yum-mirror/glslang

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

Wooyoung KimGLSL_QCOM_multiple_wait_queuescb30230f

main
branchesmainall refs →
1.9 KiB55 linesraw
1#version 450 core
2#extension GL_KHR_memory_scope_semantics : enable
3
4#extension GL_KHR_cooperative_matrix : enable
5#extension GL_QCOM_multiple_wait_queues : enable
6
7layout(constant_id = 0)  const uint TOTAL_K = 1;
8
9layout(constant_id = 3)  const uint TILE_M = 1;
10layout(constant_id = 4)  const uint TILE_N = 1;
11layout(constant_id = 7)  const uint TILE_K = 1;
12
13layout(set=0, binding=0) readonly buffer InputA { float x[]; } inputA;
14layout(set=0, binding=1) readonly buffer InputB { float x[]; } inputB;
15layout(set=0, binding=2) readonly buffer InputC { float x[]; } inputC;
16
17layout(local_size_x = 64, local_size_y = 2, local_size_z = 2) in;
18
19#define STRIDE_A 16
20#define STRIDE_B 16
21
22void coop_matmul(uint row, uint col)
23{
24  coopmat<float, gl_ScopeSubgroup, TILE_M, 8, gl_MatrixUseA> matA;
25  coopmat<float, gl_ScopeSubgroup, 8, TILE_N, gl_MatrixUseB> matB;
26  coopmat<float, gl_ScopeSubgroup, TILE_M, TILE_N, gl_MatrixUseAccumulator> matC;
27  uint step;
28
29  [[multiple_wait_queuesQCOM]]
30  for (step = 0; step < TOTAL_K; step += TILE_K) {
31    uint subMatrixAStart = row * STRIDE_A + step;
32    uint subMatrixBStart = col * STRIDE_B + step;
33
34    coopMatLoad(matA, inputA.x, subMatrixAStart, STRIDE_A, gl_CooperativeMatrixLayoutRowMajor); // Global to CoopMat
35    coopMatLoad(matB, inputB.x, subMatrixBStart, STRIDE_B, gl_CooperativeMatrixLayoutColumnMajor); // Global to CoopMat
36
37    matC = coopMatMulAdd(matA, matB, matC);
38  }
39
40  [[multiple_wait_queuesQCOM(2)]]
41  for (step = 0; step < TOTAL_K; step += TILE_K) {
42    uint subMatrixAStart = row * STRIDE_A + step;
43    uint subMatrixBStart = col * STRIDE_B + step;
44
45    coopMatLoad(matA, inputA.x, subMatrixAStart, STRIDE_A, gl_CooperativeMatrixLayoutRowMajor); // Global to CoopMat
46    coopMatLoad(matB, inputB.x, subMatrixBStart, STRIDE_B, gl_CooperativeMatrixLayoutColumnMajor); // Global to CoopMat
47
48    matC = coopMatMulAdd(matA, matB, matC);
49  }
50}
51
52void main()
53{
54  coop_matmul(0, 0);
55}