yum-mirror/glslang

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

Jeff BolzImplement GLSL_NV_cooperative_matrix2d0e67c58

main
branchesmainall refs →
3.4 KiB77 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_NV_cooperative_matrix2 : enable
6
7layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
8
9buffer BufType {
10   float16_t x[];
11} Buf;
12
13uint32_t addr(const in uint32_t x, const in uint32_t y) {
14    return y*64+x;
15}
16
17uint32_t foo() { return 124; }
18
19float16_t relu(const in uint32_t row, const in uint32_t col, const in float16_t x) { return max(x, float16_t(0)); }
20
21float16_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; }
22
23float16_t combineSum(const in float16_t a, const in float16_t b) { return a + b; }
24float16_t combineMax(const in float16_t a, const in float16_t b) { return max(a, b); }
25
26layout(constant_id = 0) const uint32_t Dim = 32;
27
28void main()
29{
30    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA> A;
31    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseB> B;
32    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Acc;
33
34    A = coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Acc);
35    B = coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseB>(Acc);
36
37    coopmat<float16_t, gl_ScopeWorkgroup, 32, 64, gl_MatrixUseB> tr;
38
39    coopMatTransposeNV(tr, Acc);
40
41    coopMatReduceNV(Acc, Acc, gl_CooperativeMatrixReduceRowNV, combineSum);
42    coopMatReduceNV(Acc, Acc, gl_CooperativeMatrixReduceColumnNV, combineSum);
43    coopMatReduceNV(Acc, Acc, gl_CooperativeMatrixReduceRowAndColumnNV, combineSum);
44
45    coopmat<float16_t, gl_ScopeWorkgroup, 32, 16, gl_MatrixUseAccumulator> Acc2x2;
46    coopMatReduceNV(Acc2x2, Acc, gl_CooperativeMatrixReduce2x2NV, combineMax);
47
48    //coopMatLoadTensorNV(A, Buf.x, foo(), addr);
49    //coopMatStoreTensorNV(A, Buf.x, foo(), addr);
50
51    coopMatPerElementNV(Acc, Acc, relu);
52    coopMatPerElementNV(Acc, Acc, add, float16_t(1.0));
53
54    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accf16;
55    coopmat<float32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accf32;
56    coopmat<uint32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accu32;
57    coopmat<int32_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseAccumulator> Accs32;
58
59    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
60    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
61    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accu32);
62
63    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
64    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
65    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accs32);
66
67    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf16);
68    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf16);
69    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf16);
70
71    coopmat<uint8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf32);
72    coopmat<int8_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf32);
73    coopmat<float16_t, gl_ScopeWorkgroup, 64, 32, gl_MatrixUseA>(Accf32);
74
75    coopmat<float, gl_ScopeWorkgroup, Dim, Dim, gl_MatrixUseAccumulator> li, mijm1;
76    mijm1 = coopmat<float, gl_ScopeWorkgroup, Dim, Dim, gl_MatrixUseAccumulator>(-1.0/0.0);
77}