yum-mirror/glslang

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

Matthew MoultonAdd checks to coopMatMulAdd (#4022)8dd7112d

main
branchesmainall refs →
2.5 KiB61 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_buffer_reference : enable
6
7layout(constant_id = 0) const int M = 3;
8layout(constant_id = 1) const int N = 3;
9
10layout(local_size_x = 4, local_size_y = 1, local_size_z = 1) in;
11
12layout(set = 0, binding = 0) coherent buffer Block16 {
13    float16_t x[];
14} block;
15
16void main() {
17    coopmat<float16_t, gl_ScopeSubgroup, 3, 4, gl_MatrixUseA> a;
18    coopmat<float16_t, gl_ScopeSubgroup, 3, 4, gl_MatrixUseB> b;
19    coopMatLoad(a, block.x, 0, 3, gl_CooperativeMatrixLayoutRowMajor);
20    coopMatLoad(b, block.x, 0, 3, gl_CooperativeMatrixLayoutRowMajor);
21
22    coopmat<float16_t, gl_ScopeSubgroup, 3, 3, gl_MatrixUseAccumulator> c =
23        coopmat<float16_t, gl_ScopeSubgroup, 3, 3, gl_MatrixUseAccumulator>(1.0);
24
25    // Incompatible dimensions (inside doesn't match): (3, 4) x (3, 4)
26    c = coopMatMulAdd(a, b, c);
27
28    coopmat<float16_t, gl_ScopeSubgroup, M, N, gl_MatrixUseA> i;
29    coopmat<float16_t, gl_ScopeSubgroup, M, N, gl_MatrixUseB> j;
30    coopMatLoad(a, block.x, 0, M, gl_CooperativeMatrixLayoutRowMajor);
31    coopMatLoad(b, block.x, 0, M, gl_CooperativeMatrixLayoutRowMajor);
32
33    coopmat<float16_t, gl_ScopeSubgroup, M, N, gl_MatrixUseAccumulator> k =
34        coopmat<float16_t, gl_ScopeSubgroup, M, N, gl_MatrixUseAccumulator>(2.0);
35
36    // Incompatible dimensions (inside doesn't match): (M, N) x (M, N)
37    k = coopMatMulAdd(i, j, k);
38
39    coopmat<float16_t, gl_ScopeSubgroup, 4, 3, gl_MatrixUseB> d;
40    coopMatLoad(d, block.x, 0, 4, gl_CooperativeMatrixLayoutRowMajor);
41    coopmat<float16_t, gl_ScopeSubgroup, 4, 4, gl_MatrixUseAccumulator> e =
42        coopmat<float16_t, gl_ScopeSubgroup, 4, 4, gl_MatrixUseAccumulator>(3.0);
43
44    // Incompatible dimensions (product doesn't match sum): (3, 4) * (4, 3) = (3, 3) + (4, 4)
45    e = coopMatMulAdd(a, d, e);
46
47    // Improper usage type: A
48    coopmat<float16_t, gl_ScopeSubgroup, 3, 4, gl_MatrixUseB> f;
49    coopMatLoad(f, block.x, 0, 3, gl_CooperativeMatrixLayoutRowMajor);
50    c = coopMatMulAdd(f, d, c);
51
52    // Improper usage type: B
53    coopmat<float16_t, gl_ScopeSubgroup, 4, 3, gl_MatrixUseAccumulator> g;
54    coopMatLoad(g, block.x, 0, 4, gl_CooperativeMatrixLayoutRowMajor);
55    c = coopMatMulAdd(a, g, c);
56
57    // Improper usage type: Accum
58    coopmat<float16_t, gl_ScopeSubgroup, 3, 3, gl_MatrixUseA> h =
59        coopmat<float16_t, gl_ScopeSubgroup, 3, 3, gl_MatrixUseA>(1.0);
60    e = coopMatMulAdd(a, d, h);
61}