wehub-resource-sync
69 lines
1.9 KiB
Plaintext
69 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 TILE_M = 64;
|
|
layout(constant_id = 4) const uint TILE_K = 16;
|
|
|
|
// Input: A_c4 [K/4, M, 4]
|
|
layout(binding = 0) readonly buffer SrcBuffer {
|
|
FLOAT Src[];
|
|
};
|
|
|
|
// Output: A_coop [M_pad/TILE_M, K_pad/TILE_K, TILE_M, TILE_K] (Block-Linear)
|
|
layout(binding = 1) writeonly buffer DstBuffer {
|
|
FLOAT Dst[];
|
|
};
|
|
|
|
layout(binding = 2) uniform constBuffer {
|
|
uint M;
|
|
uint K;
|
|
uint padM;
|
|
uint padK;
|
|
} uConst;
|
|
|
|
shared FLOAT smem[TILE_M * TILE_K];
|
|
|
|
void main() {
|
|
uint tileID_K = gl_WorkGroupID.x;
|
|
uint tileID_M = gl_WorkGroupID.y;
|
|
uint tid = gl_LocalInvocationID.x;
|
|
uint groupSize = gl_WorkGroupSize.x;
|
|
|
|
uint global_m_start = tileID_M * TILE_M;
|
|
uint base_global_k = tileID_K * TILE_K;
|
|
|
|
uint num_blocks_k = TILE_K / 4;
|
|
|
|
for (uint c = 0; c < num_blocks_k; ++c) {
|
|
uint k_out = (base_global_k >> 2) + c;
|
|
uint src_base = k_out * (uConst.M * 4) + (global_m_start * 4);
|
|
|
|
uint elements_in_chunk = TILE_M * 4;
|
|
|
|
for (uint i = tid; i < elements_in_chunk; i += groupSize) {
|
|
uint m_in_tile = i >> 2; // i / 4
|
|
uint k_in = i & 3; // i % 4
|
|
|
|
uint global_m = global_m_start + m_in_tile;
|
|
uint current_k = base_global_k + c * 4 + k_in;
|
|
|
|
FLOAT val = FLOAT(0.0);
|
|
|
|
if (global_m < uConst.M && current_k < uConst.K) {
|
|
val = Src[src_base + i];
|
|
}
|
|
|
|
uint smem_idx = m_in_tile * TILE_K + (c * 4) + k_in;
|
|
smem[smem_idx] = val;
|
|
}
|
|
}
|
|
|
|
barrier();
|
|
|
|
uint tiles_per_row = uConst.padK / TILE_K;
|
|
uint dst_tile_offset = (tileID_M * tiles_per_row + tileID_K) * (TILE_M * TILE_K);
|
|
|
|
uint total_elements = TILE_M * TILE_K;
|
|
for (uint i = tid; i < total_elements; i += groupSize) {
|
|
Dst[dst_tile_offset + i] = smem[i];
|
|
}
|
|
} |