yum-mirror/glslang

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

Wooyoung KimAdd support for GLSL_QCOM_cooperative_matrix_conversionf3a7f2c2

main
branchesmainall refs →
9.2 KiB200 linesraw
1#version 450 core
2#pragma use_vulkan_memory_model
3#extension GL_KHR_shader_subgroup_basic : enable
4#extension GL_EXT_scalar_block_layout : enable
5#extension GL_KHR_memory_scope_semantics : enable
6#extension GL_KHR_cooperative_matrix : enable
7#extension GL_EXT_buffer_reference : enable
8#extension GL_EXT_control_flow_attributes : enable
9#extension GL_KHR_shader_subgroup_basic : enable
10
11#extension GL_EXT_shader_explicit_arithmetic_types_float32 : enable
12#extension GL_EXT_shader_explicit_arithmetic_types_float16 : enable
13#extension GL_EXT_shader_explicit_arithmetic_types_int32   : enable
14#extension GL_EXT_shader_explicit_arithmetic_types_int8    : enable
15#extension GL_QCOM_cooperative_matrix_conversion : enable
16
17layout(constant_id = 3)  const uint TILE_M = 1;
18layout(constant_id = 4)  const uint TILE_N = 1;
19layout(constant_id = 5)  const uint TILE_N2 = 2;
20layout(constant_id = 6)  const uint TILE_N4 = 4;
21layout(constant_id = 7)  const uint TILE_K = 1;
22layout(constant_id = 8)  const uint TILE_K2 = 2;
23layout(constant_id = 9)  const uint TILE_K4 = 4;
24
25layout(set=0, binding=0) readonly buffer InputA { float x[]; } inputA;
26layout(set=0, binding=1) readonly buffer InputB { float x[]; } inputB;
27layout(set=0, binding=2) readonly buffer InputC { float x[]; } inputC;
28layout(set=0, binding=3) buffer Output { float x[]; } outputO;
29
30layout(local_size_x = 64, local_size_y = 2, local_size_z = 2) in;
31
32void main()
33{
34    uint32_t uvec8[8];
35    float     vec8[8];
36    uint     uvec64Acc[64];
37
38    {
39      uint32_t uvecSC[TILE_K];
40      float    f32vecSC[TILE_K];
41      bitcastQCOM(uvec8,  f32vecSC);
42      bitcastQCOM(uvecSC, vec8);
43      bitcastQCOM(uvecSC, f32vecSC);
44
45      uint32_t uvecSC2[TILE_K/2];
46      float    f32vecSC2[TILE_K/2];
47      bitcastQCOM(uvecSC2, f32vecSC2);
48    }
49
50    // extractSubArrayQCOM
51    {
52      uint32_t uvecSC[TILE_K];
53      uint32_t uvec64AccSC[TILE_N];
54      extractSubArrayQCOM(uvec64Acc, 0, uvecSC);
55      extractSubArrayQCOM(uvec64AccSC, 4, uvec8);
56      extractSubArrayQCOM(uvec64AccSC, 8, uvecSC);
57
58      uint32_t uvecSC2[TILE_K/2];
59      uint32_t uvec64AccSC2[TILE_N/2];
60      extractSubArrayQCOM(uvec64Acc, 2, uvecSC2);
61      extractSubArrayQCOM(uvec64AccSC2, 2, uvec8);
62    }
63
64    // vectorToCoopmatQCOM gl_MatrixUseA
65    {
66      coopmat<float, gl_ScopeSubgroup, TILE_M, 8, gl_MatrixUseA> V2C_matA_f32_SC_M8;
67      vectorToCoopmatQCOM(uvec8, V2C_matA_f32_SC_M8);
68
69      uint32_t uvecSC_K11[TILE_K];
70      uint32_t uvecSC_K12[TILE_K2];
71      coopmat<float, gl_ScopeSubgroup, TILE_M, TILE_K, gl_MatrixUseA> V2C_matA_f32_SC_MK;
72      coopmat<float, gl_ScopeSubgroup, 64, TILE_K, gl_MatrixUseA> V2C_matA_f32_SC_64K;
73      coopmat<float, gl_ScopeSubgroup, 64, 8, gl_MatrixUseA> V2C_matA_f32_SC_64_8;
74      vectorToCoopmatQCOM(uvecSC_K11, V2C_matA_f32_SC_MK);
75      vectorToCoopmatQCOM(uvecSC_K12, V2C_matA_f32_SC_64_8);
76
77      coopmat<float16_t, gl_ScopeSubgroup, TILE_M, TILE_K2, gl_MatrixUseA> V2C_matA_f16_SC_MK;
78      coopmat< int8_t, gl_ScopeSubgroup, 64, TILE_K4, gl_MatrixUseA> V2C_matA_i8_SC_64K;
79      coopmat<uint8_t, gl_ScopeSubgroup, 64, TILE_K4, gl_MatrixUseA> V2C_matA_u8_SC_64K;
80
81      vectorToCoopmatQCOM(uvecSC_K11, V2C_matA_f16_SC_MK);
82      vectorToCoopmatQCOM(uvecSC_K11, V2C_matA_i8_SC_64K);
83      vectorToCoopmatQCOM(uvecSC_K11, V2C_matA_u8_SC_64K);
84    }
85
86    // vectorToCoopmatQCOM gl_MatrixUseB
87    {
88      coopmat<float, gl_ScopeSubgroup, 8, TILE_N, gl_MatrixUseB> V2C_matB_f32_SC_8N;
89      vectorToCoopmatQCOM(uvec8, V2C_matB_f32_SC_8N);
90
91      uint32_t uvecSC_K21[TILE_K];
92      uint32_t uvecSC_K22[TILE_K2];
93      coopmat<float, gl_ScopeSubgroup, TILE_K, TILE_N, gl_MatrixUseB> V2C_matB_f32_SC_KN;
94      coopmat<float, gl_ScopeSubgroup, TILE_K, 64, gl_MatrixUseB> V2C_matB_f32_SC_K64;
95      coopmat<float, gl_ScopeSubgroup, 8, 64, gl_MatrixUseB> V2C_matB_f32_SC_8_64;
96      vectorToCoopmatQCOM(uvecSC_K21, V2C_matB_f32_SC_KN);
97      vectorToCoopmatQCOM(uvecSC_K22, V2C_matB_f32_SC_8_64);
98
99      coopmat<float16_t, gl_ScopeSubgroup, TILE_K2, TILE_N, gl_MatrixUseB> V2C_matB_f16_SC_KN;
100      coopmat< int8_t, gl_ScopeSubgroup, TILE_K4, 64, gl_MatrixUseB> V2C_matB_i8_SC_K64;
101      coopmat<uint8_t, gl_ScopeSubgroup, TILE_K4, 64, gl_MatrixUseB> V2C_matB_u8_SC_K64;
102      coopmat<float, gl_ScopeSubgroup, TILE_K, 32, gl_MatrixUseB> V2C_matB_f32_SC_K32;
103      coopmat<float, gl_ScopeSubgroup, TILE_K, 16, gl_MatrixUseB> V2C_matB_f32_SC_K16;
104
105      vectorToCoopmatQCOM(uvecSC_K21, V2C_matB_f16_SC_KN);
106      vectorToCoopmatQCOM(uvecSC_K21, V2C_matB_i8_SC_K64);
107      vectorToCoopmatQCOM(uvecSC_K21, V2C_matB_u8_SC_K64);
108      vectorToCoopmatQCOM(uvecSC_K21, V2C_matB_f32_SC_K32);
109      vectorToCoopmatQCOM(uvecSC_K21, V2C_matB_f32_SC_K16);
110    }
111
112    // vectorToCoopmatQCOM gl_MatrixUseAccumulator
113    {
114      uint32_t uvecSC_N31[TILE_N];
115      uint32_t uvecSC_N32[TILE_N2];
116
117      coopmat<float, gl_ScopeSubgroup, TILE_M, TILE_N, gl_MatrixUseAccumulator> V2C_matAcc_f32_SC_MN;
118      coopmat<float, gl_ScopeSubgroup, 64, TILE_N, gl_MatrixUseAccumulator> V2C_matAcc_f32_SC_64N;
119      coopmat<float, gl_ScopeSubgroup, 64, 64, gl_MatrixUseAccumulator> V2C_matAcc_f32_SC_64_64;
120      vectorToCoopmatQCOM(uvecSC_N31, V2C_matAcc_f32_SC_MN);
121      vectorToCoopmatQCOM(uvecSC_N32, V2C_matAcc_f32_SC_64_64);
122
123      coopmat<float16_t, gl_ScopeSubgroup, TILE_M, TILE_N2, gl_MatrixUseAccumulator> V2C_matAcc_f16_SC_MN;
124      coopmat< int32_t, gl_ScopeSubgroup, 64, TILE_N4, gl_MatrixUseAccumulator> V2C_matAcc_i32_SC_64N;
125      coopmat<uint32_t, gl_ScopeSubgroup, 64, TILE_N4, gl_MatrixUseAccumulator> V2C_matAcc_u32_SC_64N;
126
127      vectorToCoopmatQCOM(uvecSC_N31, V2C_matAcc_f16_SC_MN);
128      vectorToCoopmatQCOM(uvecSC_N31, V2C_matAcc_i32_SC_64N);
129      vectorToCoopmatQCOM(uvecSC_N31, V2C_matAcc_u32_SC_64N);
130    }
131
132    // coopmatToVectorQCOM gl_MatrixUseA
133    {
134      coopmat<float, gl_ScopeSubgroup, TILE_M, 8, gl_MatrixUseA> C2V_matA_f32_SC_M8;
135      coopmatToVectorQCOM(C2V_matA_f32_SC_M8, uvec8);
136
137      uint32_t uvecSC_K11[TILE_K];
138      uint32_t uvecSC_K12[TILE_K2];
139      coopmat<float, gl_ScopeSubgroup, TILE_M, TILE_K, gl_MatrixUseA> C2V_matA_f32_SC_MK;
140      coopmat<float, gl_ScopeSubgroup, 64, TILE_K, gl_MatrixUseA> C2V_matA_f32_SC_64K;
141      coopmat<float, gl_ScopeSubgroup, 64, 8, gl_MatrixUseA> C2V_matA_f32_SC_64_8;
142      coopmatToVectorQCOM(C2V_matA_f32_SC_MK, uvecSC_K11);
143      coopmatToVectorQCOM(C2V_matA_f32_SC_64_8, uvecSC_K12);
144
145      coopmat<float16_t, gl_ScopeSubgroup, TILE_M, TILE_K2, gl_MatrixUseA> C2V_matA_f16_SC_MK;
146      coopmat< int8_t, gl_ScopeSubgroup, 64, TILE_K4, gl_MatrixUseA> C2V_matA_i8_SC_64K;
147      coopmat<uint8_t, gl_ScopeSubgroup, 64, TILE_K4, gl_MatrixUseA> C2V_matA_u8_SC_64K;
148
149      coopmatToVectorQCOM(C2V_matA_f16_SC_MK, uvecSC_K11);
150      coopmatToVectorQCOM(C2V_matA_i8_SC_64K, uvecSC_K11);
151      coopmatToVectorQCOM(C2V_matA_u8_SC_64K, uvecSC_K11);
152    }
153
154    // coopmatToVectorQCOM gl_MatrixUseB
155    {
156      coopmat<float, gl_ScopeSubgroup, 8, TILE_N, gl_MatrixUseB> C2V_matB_f32_SC_8N;
157      coopmatToVectorQCOM(C2V_matB_f32_SC_8N, uvec8);
158
159      uint32_t uvecSC_K21[TILE_K];
160      uint32_t uvecSC_K22[TILE_K2];
161      coopmat<float, gl_ScopeSubgroup, TILE_K, TILE_N, gl_MatrixUseB> C2V_matB_f32_SC_KN;
162      coopmat<float, gl_ScopeSubgroup, TILE_K, 64, gl_MatrixUseB> C2V_matB_f32_SC_K64;
163      coopmat<float, gl_ScopeSubgroup, 8, 64, gl_MatrixUseB> C2V_matB_f32_SC_8_64;
164      coopmatToVectorQCOM(C2V_matB_f32_SC_KN, uvecSC_K21);
165      coopmatToVectorQCOM(C2V_matB_f32_SC_8_64, uvecSC_K22);
166
167      coopmat<float16_t, gl_ScopeSubgroup, TILE_K2, TILE_N, gl_MatrixUseB> C2V_matB_f16_SC_KN;
168      coopmat< int8_t, gl_ScopeSubgroup, TILE_K4, 64, gl_MatrixUseB> C2V_matB_i8_SC_K64;
169      coopmat<uint8_t, gl_ScopeSubgroup, TILE_K4, 64, gl_MatrixUseB> C2V_matB_u8_SC_K64;
170      coopmat<float, gl_ScopeSubgroup, TILE_K, 32, gl_MatrixUseB> C2V_matB_f32_SC_K32;
171      coopmat<float, gl_ScopeSubgroup, TILE_K, 16, gl_MatrixUseB> C2V_matB_f32_SC_K16;
172
173      coopmatToVectorQCOM(C2V_matB_f16_SC_KN, uvecSC_K21);
174      coopmatToVectorQCOM(C2V_matB_i8_SC_K64, uvecSC_K21);
175      coopmatToVectorQCOM(C2V_matB_u8_SC_K64, uvecSC_K21);
176      coopmatToVectorQCOM(C2V_matB_f32_SC_K32, uvecSC_K21);
177      coopmatToVectorQCOM(C2V_matB_f32_SC_K16, uvecSC_K21);
178    }
179
180    // coopmatToVectorQCOM gl_MatrixUseAccumulator
181    {
182      uint32_t uvecSC_N31[TILE_N];
183      uint32_t uvecSC_N32[TILE_N2];
184
185      coopmat<float, gl_ScopeSubgroup, TILE_M, TILE_N, gl_MatrixUseAccumulator> C2V_matAcc_f32_SC_MN;
186      coopmat<float, gl_ScopeSubgroup, 64, TILE_N, gl_MatrixUseAccumulator> C2V_matAcc_f32_SC_64N;
187      coopmat<float, gl_ScopeSubgroup, 64, 64, gl_MatrixUseAccumulator> C2V_matAcc_f32_SC_64_64;
188      coopmatToVectorQCOM(C2V_matAcc_f32_SC_MN, uvecSC_N31);
189      coopmatToVectorQCOM(C2V_matAcc_f32_SC_64N, uvecSC_N31);
190      coopmatToVectorQCOM(C2V_matAcc_f32_SC_64_64, uvecSC_N32);
191
192      coopmat<float16_t, gl_ScopeSubgroup, TILE_M, TILE_N2, gl_MatrixUseAccumulator> C2V_matAcc_f16_SC_MN;
193      coopmat< int32_t, gl_ScopeSubgroup, 64, TILE_N4, gl_MatrixUseAccumulator> C2V_matAcc_i32_SC_64N;
194      coopmat<uint32_t, gl_ScopeSubgroup, 64, TILE_N4, gl_MatrixUseAccumulator> C2V_matAcc_u32_SC_64N;
195
196      coopmatToVectorQCOM(C2V_matAcc_f16_SC_MN, uvecSC_N31);
197      coopmatToVectorQCOM(C2V_matAcc_i32_SC_64N, uvecSC_N31);
198      coopmatToVectorQCOM(C2V_matAcc_u32_SC_64N, uvecSC_N31);
199    }
200}