wehub-resource-sync
66 lines
2.0 KiB
Plaintext
66 lines
2.0 KiB
Plaintext
#extension GL_KHR_cooperative_matrix : require
|
|
#extension GL_KHR_memory_scope_semantics : require
|
|
|
|
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
|
|
|
|
// A is Packed: (M/COOP_M) x (K/COOP_K) tiles of COOP_M x COOP_K
|
|
layout(binding = 0) readonly buffer MatrixA {
|
|
FLOAT A[];
|
|
};
|
|
|
|
// B is Packed: (K/COOP_K) x (N/COOP_N) tiles of COOP_K x COOP_N
|
|
layout(binding = 1) readonly buffer MatrixB {
|
|
FLOAT B[];
|
|
};
|
|
|
|
// Bias is Linear: (N)
|
|
layout(binding = 2) readonly buffer BiasBuffer {
|
|
FLOAT Bias[];
|
|
};
|
|
|
|
// Output C is Linear (Row-Major) for this stage
|
|
layout(binding = 3) writeonly buffer MatrixC {
|
|
FLOAT C[];
|
|
};
|
|
|
|
layout(binding = 4) uniform constBuffer {
|
|
uint M; // Padded M
|
|
uint N; // Padded N
|
|
uint K; // Padded K
|
|
uint pad;
|
|
} uConst;
|
|
|
|
layout(constant_id = 3) const uint COOP_M = 64;
|
|
layout(constant_id = 4) const uint COOP_N = 64;
|
|
layout(constant_id = 5) const uint COOP_K = 16;
|
|
|
|
void main() {
|
|
uint wgRow = gl_WorkGroupID.y;
|
|
uint wgCol = gl_WorkGroupID.x;
|
|
|
|
coopmat<FLOAT, gl_ScopeSubgroup, COOP_M, COOP_K, gl_MatrixUseA> matA;
|
|
coopmat<FLOAT, gl_ScopeSubgroup, COOP_K, COOP_N, gl_MatrixUseB> matB;
|
|
coopmat<FLOAT, gl_ScopeSubgroup, COOP_M, COOP_N, gl_MatrixUseAccumulator> matC;
|
|
|
|
uint biasOffset = wgCol * COOP_N;
|
|
|
|
coopMatLoad(matC, Bias, biasOffset, 0, gl_CooperativeMatrixLayoutRowMajor);
|
|
|
|
uint tilesK = uConst.K / COOP_K;
|
|
uint tilesN = uConst.N / COOP_N;
|
|
|
|
for (uint t = 0; t < tilesK; ++t) {
|
|
uint tileIdxA = wgRow * tilesK + t;
|
|
uint offsetA = tileIdxA * (COOP_M * COOP_K);
|
|
coopMatLoad(matA, A, offsetA, COOP_K, gl_CooperativeMatrixLayoutRowMajor);
|
|
|
|
uint tileIdxB = t * tilesN + wgCol;
|
|
uint offsetB = tileIdxB * (COOP_K * COOP_N);
|
|
coopMatLoad(matB, B, offsetB, COOP_N, gl_CooperativeMatrixLayoutRowMajor);
|
|
|
|
matC = coopMatMulAdd(matA, matB, matC);
|
|
}
|
|
|
|
uint offsetC = (wgRow * COOP_M) * uConst.N + (wgCol * COOP_N);
|
|
coopMatStore(matC, C, offsetC, uConst.N, gl_CooperativeMatrixLayoutRowMajor);
|
|
} |