wehub-resource-sync
65 lines
1.9 KiB
Plaintext
65 lines
1.9 KiB
Plaintext
layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
|
|
|
|
layout(constant_id = 3) const uint ACTIVATION = 0;
|
|
|
|
layout(binding = 0) readonly buffer SrcBuffer { FLOAT4 SrcVec4[]; };
|
|
layout(binding = 1) writeonly buffer DstBuffer { FLOAT4 DstVec4[]; };
|
|
|
|
layout(binding = 2) uniform constBuffer {
|
|
uint M;
|
|
uint N;
|
|
uint padM;
|
|
uint padN;
|
|
} uConst;
|
|
|
|
shared FLOAT4 smem[32][33];
|
|
|
|
void main() {
|
|
uint tileN_out = gl_WorkGroupID.x * 32;
|
|
uint tileM = gl_WorkGroupID.y * 32;
|
|
|
|
uint flat_tid = gl_LocalInvocationIndex;
|
|
uint total_threads = gl_WorkGroupSize.x * gl_WorkGroupSize.y * gl_WorkGroupSize.z;
|
|
|
|
uint loop_count = 1024 / total_threads;
|
|
|
|
uint padN_vec4 = uConst.padN / 4;
|
|
|
|
for (uint i = 0; i < loop_count; ++i) {
|
|
uint tile_linear_idx = flat_tid + i * total_threads;
|
|
|
|
uint local_m = tile_linear_idx / 32; // Row
|
|
uint local_n = tile_linear_idx % 32; // Col
|
|
|
|
uint global_m = tileM + local_m;
|
|
uint global_n_out = tileN_out + local_n;
|
|
|
|
FLOAT4 val = FLOAT4(0.0);
|
|
if (global_m < uConst.M && (global_n_out * 4) < uConst.N) {
|
|
val = SrcVec4[global_m * padN_vec4 + global_n_out];
|
|
}
|
|
smem[local_m][local_n] = val;
|
|
}
|
|
|
|
barrier();
|
|
|
|
for (uint i = 0; i < loop_count; ++i) {
|
|
uint tile_linear_idx = flat_tid + i * total_threads;
|
|
|
|
uint local_n = tile_linear_idx / 32; // High bit -> N_out col
|
|
uint local_m = tile_linear_idx % 32; // Low bit -> M row
|
|
|
|
uint global_n_out = tileN_out + local_n;
|
|
uint global_m = tileM + local_m;
|
|
|
|
if (global_n_out * 4 < uConst.N && global_m < uConst.M) {
|
|
FLOAT4 val = smem[local_m][local_n];
|
|
if (ACTIVATION == 1) {
|
|
val = max(val, FLOAT4(0.0));
|
|
} else if (ACTIVATION == 2) {
|
|
val = clamp(val, FLOAT4(0.0), FLOAT4(6.0));
|
|
}
|
|
DstVec4[global_n_out * uConst.M + global_m] = val;
|
|
}
|
|
}
|
|
} |