Example: VulkanExamples/p78_coopmat
VK_KHR_cooperative_matrix is the tensor-core-shaped op in Vulkan: a
coopmat is a matrix distributed across an entire subgroup — no
lane owns it, coopMatLoad/coopMatStore scatter/gather it, and
coopMatMulAdd(A, B, C) computes A*B + C as one collective op.
example.txt
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
coopmat<int8_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA> A;
coopmat<int8_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB> B;
coopmat<int32_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator> C;
coopMatLoad(A, a, 0, 16, gl_CooperativeMatrixLayoutRowMajor);
C = coopMatMulAdd(A, B, C);
coopMatStore(C, c, 0, 16, gl_CooperativeMatrixLayoutRowMajor);The supported shapes are a hardware query —
vkGetPhysicalDeviceCooperativeMatrixPropertiesKHR — and on
NAVI32/RADV there are exactly 14 entries, all 16×16×16 at
VK_SCOPE_SUBGROUP, with uint8/sint8/float16 inputs and
uint32/sint32/f32 accumulators. The demo picks s8 × s8 → s32: int
math, so the 16×16 matmul-add verifies bit-exactly against the
CPU, bad=0.
Practical notes:
- The workgroup must effectively be the cooperating scope — use
local_size_x_idfed bysubgroupSize(64 on RADV wave64). - Layout constant is
gl_CooperativeMatrixLayoutRowMajor— notgl_MatrixLayoutRowMajor; the shorter name doesn’t exist. int8_t/int32_tin SSBOs needGL_EXT_shader_8bit_storage+GL_EXT_shader_explicit_arithmetic_types_int8/int32plus the matching Vulkan feature bits (storageBuffer8BitAccess,shaderInt8).- The whole point vs. doing it lane-wise: the fill of the matrix is hidden — the driver maps it onto WMMA/dot-product hardware without you caring which lane holds which element.
All source code in this series is released under CC0 1.0 Universal (CC0-1.0).
If this post helped you, please consider buying me a coffee or donating via PayPal to support research & publishing of new posts on TechOverflow