yum-mirror/glslang

git clone https://git.yummers.dev/yum-mirror/glslang

Jeff BolzImplement GL_EXT_cooperative_matrix_maintenance124ae85ba

main
branchesmainall refs →
5.4 KiB128 linesraw
1#version 450 core
2#extension GL_KHR_memory_scope_semantics : enable
3#extension GL_KHR_cooperative_matrix : enable
4#extension GL_EXT_shader_explicit_arithmetic_types : enable
5#extension GL_EXT_cooperative_matrix_maintenance1 : enable
6#extension GL_EXT_float_e5m2 : enable
7
8layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
9
10float16_t relu(const in uint32_t row, const in uint32_t col, const in float16_t x) { return max(x, float16_t(0)); }
11
12float16_t add(const in uint32_t row, const in uint32_t col, const in float16_t x, const in float16_t y) { return x+y; }
13
14float16_t combineSum(const in float16_t a, const in float16_t b) { return a + b; }
15float16_t combineMax(const in float16_t a, const in float16_t b) { return max(a, b); }
16
17layout(constant_id = 0) const uint32_t Dim = 32;
18
19void main()
20{
21    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA> A;
22    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseB> B;
23    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Acc;
24
25    A = coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Acc);
26    B = coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseB>(Acc);
27
28    coopmat<float16_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> tr;
29
30    coopMatTransposeEXT(tr, Acc);
31
32    coopMatReduceEXT(Acc, Acc, gl_CooperativeMatrixReduceRowEXT, combineSum);
33    coopMatReduceEXT(Acc, Acc, gl_CooperativeMatrixReduceColumnEXT, combineSum);
34    coopMatReduceEXT(Acc, Acc, gl_CooperativeMatrixReduceRowAndColumnEXT, combineSum);
35
36    coopmat<float16_t, gl_ScopeWorkgroup, 32, 16, gl_MatrixUseAccumulator> Acc2x2;
37    coopMatReduceEXT(Acc2x2, Acc, gl_CooperativeMatrixReduce2x2EXT, combineMax);
38
39    coopMatPerElementEXT(Acc, Acc, relu);
40    coopMatPerElementEXT(Acc, Acc, add, float16_t(1.0));
41
42    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accf16;
43    coopmat<float32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accf32;
44    coopmat<uint32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accu32;
45    coopmat<int32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accs32;
46
47    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
48    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
49    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
50
51    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
52    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
53    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
54
55    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf16);
56    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf16);
57    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf16);
58
59    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf32);
60    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf32);
61    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf32);
62
63    coopmat<uint32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
64    coopmat<int32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
65
66    coopmat<float, gl_ScopeWorkgroup, Dim, Dim, gl_MatrixUseAccumulator> li, mijm1;
67    mijm1 = coopmat<float, gl_ScopeWorkgroup, Dim, Dim, gl_MatrixUseAccumulator>(-1.0/0.0);
68
69    // New conversions in GL_EXT_cooperative_matrix_maintenance1:
70
71    Acc = coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator>(A);
72    Acc = coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator>(B);
73
74    uint index = 1;
75    uvec2 coord;
76    coord = coopMatGetCoordinateEXT(A, index);
77    coord = coopMatGetCoordinateEXT(B, index);
78    coord = coopMatGetCoordinateEXT(Acc, index);
79
80    coopmat<uint8_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseA> Au8t;
81    coopmat<int8_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseA> As8t;
82    coopmat<float16_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseA> Af16t;
83    coopmat<int32_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> As32t;
84
85    coopmat<uint8_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> Bu8t;
86    coopmat<int8_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> Bs8t;
87    coopmat<float16_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> Bf16t;
88    coopmat<int32_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> Bs32t;
89
90    coopMatTransposeEXT(Bu8t, Accf16);
91    coopMatTransposeEXT(Bs8t, Accf16);
92    coopMatTransposeEXT(Bf16t, Accf16);
93
94    coopMatTransposeEXT(Bu8t, Accf32);
95    coopMatTransposeEXT(Bs8t, Accf32);
96    coopMatTransposeEXT(Bf16t, Accf32);
97
98    coopMatTransposeEXT(Bu8t, Accs32);
99    coopMatTransposeEXT(Bs8t, Accs32);
100    coopMatTransposeEXT(Bf16t, Accs32);
101
102    coopMatTransposeEXT(Bu8t, Accu32);
103    coopMatTransposeEXT(Bs8t, Accu32);
104    coopMatTransposeEXT(Bf16t, Accu32);
105
106    coopMatTransposeEXT(Bs32t, Accu32);
107
108    coopmat<floate5m2_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA> matAe5m2;
109    saturatedConvertEXT(matAe5m2, Accf16);
110
111    coopMatTransposeEXT(Au8t, Accf16);
112    coopMatTransposeEXT(As8t, Accf16);
113    coopMatTransposeEXT(Af16t, Accf16);
114
115    coopMatTransposeEXT(Au8t, Accf32);
116    coopMatTransposeEXT(As8t, Accf32);
117    coopMatTransposeEXT(Af16t, Accf32);
118
119    coopMatTransposeEXT(Au8t, Accs32);
120    coopMatTransposeEXT(As8t, Accs32);
121    coopMatTransposeEXT(Af16t, Accs32);
122
123    coopMatTransposeEXT(Au8t, Accu32);
124    coopMatTransposeEXT(As8t, Accu32);
125    coopMatTransposeEXT(Af16t, Accu32);
126
127    coopMatTransposeEXT(As32t, Accu32);
128}