Workgroup doesn't work in MoltenVk

[redain]

Expected Behaviour

I should be able to use workgroups in compute shaders.

Example & Steps To Reproduce

This is the most simplified example I could make to reproduce.

#[spirv(compute(threads(256)))]
    pub fn test(
        #[spirv(local_invocation_id)] lid: UVec3,
        #[spirv(workgroup_id)] gid: UVec3,
        #[spirv(workgroup)] shared: &mut [u32; 32],
        #[spirv(push_constant)] pc: &PushConstants,
        #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] all_bufs: &mut RuntimeArray<
            TypedBuffer<[u32]>,
        >,
    ) {
        shared[0] = 1;
        shared[1] = 2;
        let a = Vec2::new(shared[0] as f32, shared[1] as f32);
    }

produces this runtime Vulkan error:

vkCreateShaderModule(): pCreateInfo->pCode (spirv-val produced an error):
Invalid explicit layout decorations on type for operand '17[%17]'
  %shared = OpVariable %_ptr_Workgroup__arr_uint_uint_32 Workgroup
The Vulkan spec states: All variables must have valid explicit layout decorations as described in Shader Interfaces (https://vulkan.lunarg.com/doc/view/1.4.313.0/mac/antora/spec/latest/appendices/spirvenv.html#VUID-StandaloneSpirv-None-10684)    

spir-v snippet:

 OpCapability Shader
               OpCapability Int64
               OpCapability StorageImageExtendedFormats
               OpCapability ImageQuery
               OpCapability StorageImageWriteWithoutFormat
               OpCapability RuntimeDescriptorArray
               OpExtension "SPV_EXT_descriptor_indexing"
               OpExtension "SPV_KHR_storage_buffer_storage_class"
               OpMemoryModel Logical GLSL450
               OpEntryPoint GLCompute %1 "test" %lid %gid
               OpExecutionMode %1 LocalSize 256 1 1
  OpName %lid "lid"
               OpName %gid "gid"
               OpName %shared "shared"
               OpDecorate %lid BuiltIn LocalInvocationId
               OpDecorate %gid BuiltIn WorkgroupId
               OpDecorate %_arr_uint_uint_32 ArrayStride 4
       %void = OpTypeVoid
          %8 = OpTypeFunction %void
       %uint = OpTypeInt 32 0
     %uint_0 = OpConstant %uint 0
     %uint_1 = OpConstant %uint 1
       %bool = OpTypeBool
     %v3uint = OpTypeVector %uint 3
%_ptr_Input_v3uint = OpTypePointer Input %v3uint
    %uint_32 = OpConstant %uint 32
        %lid = OpVariable %_ptr_Input_v3uint Input
        %gid = OpVariable %_ptr_Input_v3uint Input
%_ptr_Workgroup_uint = OpTypePointer Workgroup %uint
%_arr_uint_uint_32 = OpTypeArray %uint %uint_32
%_ptr_Workgroup__arr_uint_uint_32 = OpTypePointer Workgroup %_arr_uint_uint_32
     %shared = OpVariable %_ptr_Workgroup__arr_uint_uint_32 Workgroup
          %1 = OpFunction %void None %8
         %18 = OpLabel
               OpLine %4 209 8
         %19 = OpLoad %v3uint %lid
               OpLine %4 210 8
         %20 = OpLoad %v3uint %gid
               OpLine %4 217 8
         %21 = OpULessThan %bool %uint_0 %uint_32
               OpNoLine
               OpSelectionMerge %22 None
               OpBranchConditional %21 %23 %24
         %23 = OpLabel
               OpBranch %22
         %24 = OpLabel
               OpReturn
         %22 = OpLabel
               OpLine %4 217 8
         %25 = OpInBoundsAccessChain %_ptr_Workgroup_uint %shared %uint_0
               OpStore %25 %uint_1
               OpLine %4 218 26
         %26 = OpULessThan %bool %uint_0 %uint_32
               OpNoLine
               OpSelectionMerge %27 None
               OpBranchConditional %26 %28 %29
         %28 = OpLabel
               OpBranch %27
         %29 = OpLabel
               OpReturn
         %27 = OpLabel
               OpLine %4 218 26
         %30 = OpInBoundsAccessChain %_ptr_Workgroup_uint %shared %uint_0
         %31 = OpLoad %uint %30
               OpLine %4 218 44
         %32 = OpULessThan %bool %uint_1 %uint_32
               OpNoLine
               OpSelectionMerge %33 None
               OpBranchConditional %32 %34 %35
         %34 = OpLabel
               OpBranch %33
         %35 = OpLabel
               OpReturn
         %33 = OpLabel
               OpLine %4 218 44
         %36 = OpInBoundsAccessChain %_ptr_Workgroup_uint %shared %uint_1
         %37 = OpLoad %uint %36
               OpNoLine
               OpReturn
               OpFunctionEnd

System Info

• Rust: nightly-2024-11-22
• OS: MacOS 15.5
• GPU: Apple M1 (Integrated GPU) with Vulkan 1.2.296
• SPIR-V: SPIRV-Tools v2025.2 v2025.2.rc2-0-ga62abcb4
• Vulkan Loader version: 1.4.313
• rust-gpu: commit: 86fc480

[redain]

This might be related:
KhronosGroup/SPIRV-Tools#6012

[redain]

Seems like the same issue was encountered and solved in glsl compiler:
KhronosGroup/glslang#3911

[LegNeato]

Hm, I thought we had a bug for this around but can't find it. Thanks for the great report!

[Firestar99]

You mean this PR? #239

Only decorate physically sized types with ArrayStride

[LegNeato]

Ah yeah, that was what I was thinking of

[eddyb]

This is a fundamentally unsolvable issue, at the rustc_codegen_spirv level: we map Rust types to SPIR-V types, and so [u32; N] may need a ArrayStride 4 decoration, or may be disallowed from having any such decorations.
(even #239 might need to be reverted, at least if we want legalization to always have accurate layout information on the "input" side, even if it's not emitted in the "output" SPIR-V)

---

Does MoltenVK support VK_KHR_workgroup_memory_explicit_layout? AFAICT no, but https://vulkan.gpuinfo.org/ is down, so I can't check for sure (and there are other reasons to not rely on it, regardless).

In that case, the only solution involves running the spirt::{mem,qptr} transformations, which are capable of replacing typed memory with untyped, then regenerate types based on accesses, and which we might be able to enable already on main with a few minor fixes (or at least keep it limited to the problematic storage classes).

However, the current spirt::qptr::lift implementation doesn't know about (some) storage classes disallowing explicit memory layout decorations, it's just a tiny local change (explanation follow below).

---

For my example above, the same [u32; N] Rust type, and the SPIR-V OpTypeArray u32 N w/ ArrayStride 4 that it compiles into today, would result in spirt::mem::analyze most likely finding an "access pattern" like this:

#[mem.accesses(
  // for a buffer, this would also be wrapped in `buffer_data(...)`
  (0..N) × 4 => u32
)]

Every (0..N) × 4 => u32 must map to an OpTypeArray u32 N SPIR-V type, but each occurrence of it could be a different SPIR-V type if it wanted to (e.g. for excessively accurate debuginfo), there is no type equality requirement that could arise from anywhere, only the ability to use OpAccessChain with a dynamic index.

So all we need is to add another argument (e.g. spirt::AddrSpace) to the MemAccesses -> Type function, allowing it to distinguish between the cases that need Offset/ArrayStride and the ones that don't.

(NB: above names assume Rust-GPU/spirt#25 landed first)