How to use VK_KHR_cooperative_matrix in Vulkan

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_id fed by subgroupSize (64 on RADV wave64).
  • Layout constant is gl_CooperativeMatrixLayoutRowMajor — not gl_MatrixLayoutRowMajor; the shorter name doesn’t exist.
  • int8_t/int32_t in SSBOs need GL_EXT_shader_8bit_storage + GL_EXT_shader_explicit_arithmetic_types_int8/int32 plus 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).


Check out similar posts by category: C++ Vulkan