vulkan: Use larger loads in scalar/coopmat1 matmul (#15729)

I think glslang will translate an access like x[i][1].z to
OpAccessChain ... x, i, 1, 2
OpLoad float16_t ...

rather than loading all of x[i] in a single OpLoad. Change the
code to explicitly load the vector/matrix.
This commit is contained in:
Jeff Bolz
2025-09-07 11:53:07 -05:00
committed by GitHub
parent 3b15924d71
commit 267e99867f
3 changed files with 57 additions and 34 deletions
+36 -24
View File
@@ -315,21 +315,23 @@ void main() {
#if LOAD_VEC_A == 8
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A;
buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx][0].x);
buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx][0].y);
buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx][0].z);
buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx][0].w);
buf_a[buf_idx + 4] = FLOAT_TYPE(data_a[idx][1].x);
buf_a[buf_idx + 5] = FLOAT_TYPE(data_a[idx][1].y);
buf_a[buf_idx + 6] = FLOAT_TYPE(data_a[idx][1].z);
buf_a[buf_idx + 7] = FLOAT_TYPE(data_a[idx][1].w);
A_TYPE32 aa = A_TYPE32(data_a[idx]);
buf_a[buf_idx ] = FLOAT_TYPE(aa[0].x);
buf_a[buf_idx + 1] = FLOAT_TYPE(aa[0].y);
buf_a[buf_idx + 2] = FLOAT_TYPE(aa[0].z);
buf_a[buf_idx + 3] = FLOAT_TYPE(aa[0].w);
buf_a[buf_idx + 4] = FLOAT_TYPE(aa[1].x);
buf_a[buf_idx + 5] = FLOAT_TYPE(aa[1].y);
buf_a[buf_idx + 6] = FLOAT_TYPE(aa[1].z);
buf_a[buf_idx + 7] = FLOAT_TYPE(aa[1].w);
#elif LOAD_VEC_A == 4
const uint idx = pos_a + (loadc_a + l) * p.stride_a / LOAD_VEC_A + loadr_a;
const uint buf_idx = (loadc_a + l) * SHMEM_STRIDE + loadr_a * LOAD_VEC_A;
buf_a[buf_idx ] = FLOAT_TYPE(data_a[idx].x);
buf_a[buf_idx + 1] = FLOAT_TYPE(data_a[idx].y);
buf_a[buf_idx + 2] = FLOAT_TYPE(data_a[idx].z);
buf_a[buf_idx + 3] = FLOAT_TYPE(data_a[idx].w);
A_TYPE32 aa = A_TYPE32(data_a[idx]);
buf_a[buf_idx ] = FLOAT_TYPE(aa.x);
buf_a[buf_idx + 1] = FLOAT_TYPE(aa.y);
buf_a[buf_idx + 2] = FLOAT_TYPE(aa.z);
buf_a[buf_idx + 3] = FLOAT_TYPE(aa.w);
#else
if (ir * BM + loadc_a + l < p.M && block + loadr_a < end_k) {
buf_a[(loadc_a + l) * SHMEM_STRIDE + loadr_a] = FLOAT_TYPE(data_a[pos_a + (loadc_a + l) * p.stride_a + loadr_a]);
@@ -808,14 +810,19 @@ void main() {
const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b;
#endif
const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B;
buf_b[buf_idx + 0] = FLOAT_TYPE(data_b[idx][0].x);
buf_b[buf_idx + 1] = FLOAT_TYPE(data_b[idx][0].y);
buf_b[buf_idx + 2] = FLOAT_TYPE(data_b[idx][0].z);
buf_b[buf_idx + 3] = FLOAT_TYPE(data_b[idx][0].w);
buf_b[buf_idx + 4] = FLOAT_TYPE(data_b[idx][1].x);
buf_b[buf_idx + 5] = FLOAT_TYPE(data_b[idx][1].y);
buf_b[buf_idx + 6] = FLOAT_TYPE(data_b[idx][1].z);
buf_b[buf_idx + 7] = FLOAT_TYPE(data_b[idx][1].w);
#if defined(DATA_B_BF16)
B_TYPE32 bb = TO_FLOAT_TYPE(data_b[idx]);
#else
B_TYPE32 bb = B_TYPE32(data_b[idx]);
#endif
buf_b[buf_idx + 0] = FLOAT_TYPE(bb[0].x);
buf_b[buf_idx + 1] = FLOAT_TYPE(bb[0].y);
buf_b[buf_idx + 2] = FLOAT_TYPE(bb[0].z);
buf_b[buf_idx + 3] = FLOAT_TYPE(bb[0].w);
buf_b[buf_idx + 4] = FLOAT_TYPE(bb[1].x);
buf_b[buf_idx + 5] = FLOAT_TYPE(bb[1].y);
buf_b[buf_idx + 6] = FLOAT_TYPE(bb[1].z);
buf_b[buf_idx + 7] = FLOAT_TYPE(bb[1].w);
#elif LOAD_VEC_B == 4
#ifdef MUL_MAT_ID
const u16vec2 row_idx = row_ids[loadc_b + l];
@@ -824,10 +831,15 @@ void main() {
const uint idx = pos_b + (loadc_b + l) * p.stride_b / LOAD_VEC_B + loadr_b;
#endif
const uint buf_idx = (loadc_b + l) * SHMEM_STRIDE + loadr_b * LOAD_VEC_B;
buf_b[buf_idx + 0] = TO_FLOAT_TYPE(data_b[idx].x);
buf_b[buf_idx + 1] = TO_FLOAT_TYPE(data_b[idx].y);
buf_b[buf_idx + 2] = TO_FLOAT_TYPE(data_b[idx].z);
buf_b[buf_idx + 3] = TO_FLOAT_TYPE(data_b[idx].w);
#if defined(DATA_B_BF16)
B_TYPE32 bb = TO_FLOAT_TYPE(data_b[idx]);
#else
B_TYPE32 bb = B_TYPE32(data_b[idx]);
#endif
buf_b[buf_idx + 0] = FLOAT_TYPE(bb.x);
buf_b[buf_idx + 1] = FLOAT_TYPE(bb.y);
buf_b[buf_idx + 2] = FLOAT_TYPE(bb.z);
buf_b[buf_idx + 3] = FLOAT_TYPE(bb.w);
#elif !MUL_MAT_ID
if (ic * BN + loadc_b + l < p.N && block + loadr_b < end_k) {
buf_b[(loadc_b + l) * SHMEM_STRIDE + loadr_b] = TO_FLOAT_TYPE(data_b[pos_b + (loadc_b + l) * p.stride_b + loadr_b]);