Optimize Krea Vulkan - MUL_MAT - #27493
Conversation
Assisted-by: OpenAI Codex
| uint elem_row[CM_ELEMS]; | ||
| uint elem_col[CM_ELEMS]; | ||
| if (WARP == 32) { | ||
| [[unroll]] for (uint e = 0; e < CM_ELEMS; ++e) { | ||
| elem_row[e] = gl_SubgroupInvocationID / CM_N + 2 * e; | ||
| elem_col[e] = gl_SubgroupInvocationID % CM_N; | ||
| } | ||
| } else { | ||
| for (uint i = gl_LocalInvocationID.x; i < CM_M * CM_N; i += BLOCK_SIZE) { | ||
| cm_layout_probe[i] = int32_t(i); | ||
| } | ||
| barrier(); | ||
|
|
||
| coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator> probe; | ||
| coopMatLoad(probe, cm_layout_probe, 0, CM_N, gl_CooperativeMatrixLayoutRowMajor); | ||
|
|
||
| [[unroll]] for (uint e = 0; e < CM_ELEMS; ++e) { | ||
| elem_row[e] = uint(probe[e]) / CM_N; | ||
| elem_col[e] = uint(probe[e]) % CM_N; | ||
| } | ||
| } |
There was a problem hiding this comment.
This is a huge hack and explains why you got it to work at all over my (naively spec-adhering) attempts.
There was a problem hiding this comment.
Usually you have to write a coopmat to shared memory to be able to work with its values directly, that is why my previous attempts to write an int8-based coopmat1 shader failed, it was too inefficient. I hadn't thought of probing the order the driver stores the values. It's hacky, but not a bad idea.
There was a problem hiding this comment.
There are better ways to get the row/col now with https://github.com/KhronosGroup/GLSL/blob/main/extensions/ext/GLSL_EXT_cooperative_matrix_maintenance1.txt.
There was a problem hiding this comment.
Do you know when we can expect driver support?
There was a problem hiding this comment.
For NVIDIA it's already in the https://developer.nvidia.com/vulkan-driver, can't say for the mainline driver (could be a few months)
86110ec to
fcf82c2
Compare
|
This PR needs a description and some perf data. |
Assisted-by: OpenAI Codex
|
Here are the numbers:
Useful-operation throughput:
This is all for AMD gfx1151 (Strix Halo) |
Overview
The MUL_MAT part of the Krea Vulkan optimization
Requirements