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:
@@ -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]);
|
||||
|
||||
Reference in New Issue
Block a user