[spv-in] Add support for structs that contain pointers

[tombh]

I'm writing a physics compute shader with rust-gpu. Recently I tried adding Workgroup support, but I ran into this issue with naga validation:

Type [34] '' is invalid:
        Expected data type, found [17]

I'm using Bevy 0.14.1 that depends on naga 0.20.0. But I can reproduce the error with naga-cli on trunk.

I've managed to get a minimal reproducible rust-gpu shader that cross-compiles to a small GLSL shader:

#![cfg_attr(target_arch = "spirv", no_std)]

use spirv_std::{glam::Vec2, spirv};

type Workgroup = [Vec2; 1];

pub struct Cell<'world> {
    pub positions: &'world mut [Vec2],
    pub velocities: &'world [Vec2],
    pub workgroup: &'world mut Workgroup,
}

impl Cell<'_> {
    pub fn physics_for_cell(&mut self) {
        self.workgroup[0] = self.velocities[0];
        self.positions[0] = self.workgroup[0];
    }
}

#[spirv(compute(threads(1)))]
pub fn main(
    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] positions: &mut [Vec2],
    #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] velocities: &[Vec2],
    #[spirv(workgroup)] workgroup: &mut Workgroup,
) {
    let mut cell = Cell {
        positions,
        velocities,
        workgroup,
    };

    cell.physics_for_cell();
}

And so the validation error now is:

Type [7] '' is invalid:
        Expected data type, found [3]

How do I dump the Naga representation? Would that be an internal WGSL version? Here's the GLSL version from spirv-cross (I don't know if the variable/type numbers match up with Naga's?). Anyway, looking at this version, type 7 could be the workgroup type: shared vec2 _7[1];. But for some reason my hunch is that the problem is the struct in struct in type 6, therefore: struct _6 { _5 _m0 ...

GLSL:

#version 450
layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;

struct _5
{
    vec2 _m0[];
    uint _m1;
};

struct _6
{
    _5 _m0;
    _5 _m1;
    vec2 _m2[1];
};

layout(binding = 0, std430) buffer _9_3
{
    vec2 _m0[];
} _3;

layout(binding = 1, std430) readonly buffer _9_4
{
    vec2 _m0[];
} _4;

shared vec2 _7[1];

void main()
{
    if (0u < uint(_4._m0.length()))
    {
        _7[0u] = _4._m0[0u];
        if (0u < uint(_3._m0.length()))
        {
            _3._m0[0u] = _7[0u];
        }
        else
        {
        }
    }
    else
    {
    }
}

And here's the SPIR-V:

; SPIR-V
; Version: 1.3
; Generator: Google rspirv; 0
; Bound: 56
; Schema: 0
               OpCapability Shader
               OpMemoryModel Logical Simple
               OpEntryPoint GLCompute %1 "main"
               OpExecutionMode %1 LocalSize 1 1 1
               OpDecorate %_runtimearr_v2float ArrayStride 8
               OpDecorate %_struct_9 Block
               OpMemberDecorate %_struct_9 0 Offset 0
               OpDecorate %3 Binding 0
               OpDecorate %3 DescriptorSet 0
               OpDecorate %4 NonWritable
               OpDecorate %4 Binding 1
               OpDecorate %4 DescriptorSet 0
               OpMemberDecorate %_struct_5 0 Offset 0
               OpMemberDecorate %_struct_5 1 Offset 4
               OpDecorate %_arr_v2float_uint_1 ArrayStride 8
               OpMemberDecorate %_struct_6 0 Offset 0
               OpMemberDecorate %_struct_6 1 Offset 8
               OpMemberDecorate %_struct_6 2 Offset 16
       %void = OpTypeVoid
         %12 = OpTypeFunction %void
      %float = OpTypeFloat 32
    %v2float = OpTypeVector %float 2
%_runtimearr_v2float = OpTypeRuntimeArray %v2float
%_ptr_StorageBuffer__runtimearr_v2float = OpTypePointer StorageBuffer %_runtimearr_v2float
  %_struct_9 = OpTypeStruct %_runtimearr_v2float
%_ptr_StorageBuffer__struct_9 = OpTypePointer StorageBuffer %_struct_9
          %3 = OpVariable %_ptr_StorageBuffer__struct_9 StorageBuffer
       %uint = OpTypeInt 32 0
     %uint_0 = OpConstant %uint 0
          %4 = OpVariable %_ptr_StorageBuffer__struct_9 StorageBuffer
  %_struct_5 = OpTypeStruct %_ptr_StorageBuffer__runtimearr_v2float %uint
     %uint_1 = OpConstant %uint 1
%_arr_v2float_uint_1 = OpTypeArray %v2float %uint_1
%_ptr_Workgroup__arr_v2float_uint_1 = OpTypePointer Workgroup %_arr_v2float_uint_1
  %_struct_6 = OpTypeStruct %_struct_5 %_struct_5 %_ptr_Workgroup__arr_v2float_uint_1
          %7 = OpVariable %_ptr_Workgroup__arr_v2float_uint_1 Workgroup
       %bool = OpTypeBool
%_ptr_StorageBuffer_v2float = OpTypePointer StorageBuffer %v2float
%_ptr_Workgroup_v2float = OpTypePointer Workgroup %v2float
          %1 = OpFunction %void None %12
         %25 = OpLabel
         %26 = OpAccessChain %_ptr_StorageBuffer__runtimearr_v2float %3 %uint_0
         %27 = OpArrayLength %uint %3 0
         %28 = OpAccessChain %_ptr_StorageBuffer__runtimearr_v2float %4 %uint_0
         %29 = OpArrayLength %uint %4 0
         %53 = OpCompositeConstruct %_struct_5 %26 %27
         %54 = OpCompositeConstruct %_struct_5 %28 %29
         %55 = OpCompositeConstruct %_struct_6 %53 %54 %7
         %36 = OpULessThan %bool %uint_0 %29
               OpSelectionMerge %37 None
               OpBranchConditional %36 %38 %39
         %38 = OpLabel
         %41 = OpInBoundsAccessChain %_ptr_StorageBuffer_v2float %28 %uint_0
         %42 = OpLoad %v2float %41
         %43 = OpInBoundsAccessChain %_ptr_Workgroup_v2float %7 %uint_0
               OpStore %43 %42
         %45 = OpLoad %v2float %43
         %47 = OpULessThan %bool %uint_0 %27
               OpSelectionMerge %48 None
               OpBranchConditional %47 %49 %50
         %49 = OpLabel
         %52 = OpInBoundsAccessChain %_ptr_StorageBuffer_v2float %26 %uint_0
               OpStore %52 %45
               OpBranch %48
         %50 = OpLabel
               OpBranch %48
         %48 = OpLabel
               OpBranch %37
         %39 = OpLabel
               OpBranch %37
         %37 = OpLabel
               OpReturn
               OpFunctionEnd

I originally posted this as a discussion at #6361

[teoxoy]

This is because naga doesn't allow structs to contain pointers.

[tombh]

Thank you. Do you think a good first step would just be to improve the error?

Just out of curiosity, this does validate, even though the Rust code has a pointer in a struct:

#![cfg_attr(target_arch = "spirv", no_std)]

use spirv_std::{glam::Vec2, spirv};

pub struct Cell<'world> {
    pub positions: &'world mut [Vec2],
}

impl Cell<'_> {
    pub fn physics_for_cell(&mut self) {
        self.positions[0] = Vec2::new(0.1, 0.1);
    }
}

#[spirv(compute(threads(1)))]
pub fn main(#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] positions: &mut [Vec2]) {
    let mut cell = Cell { positions };

    cell.physics_for_cell();
}

But I suspect that's because Rust compiles down to SPIR-V without a pointer in a struct?

; SPIR-V
; Version: 1.3
; Generator: Google rspirv; 0
; Bound: 28
; Schema: 0
               OpCapability Shader
               OpMemoryModel Logical Simple
               OpEntryPoint GLCompute %1 "main"
               OpExecutionMode %1 LocalSize 1 1 1
               OpDecorate %_runtimearr_v2float ArrayStride 8
               OpDecorate %_struct_6 Block
               OpMemberDecorate %_struct_6 0 Offset 0
               OpDecorate %4 Binding 0
               OpDecorate %4 DescriptorSet 0
       %void = OpTypeVoid
          %8 = OpTypeFunction %void
      %float = OpTypeFloat 32
    %v2float = OpTypeVector %float 2
%_runtimearr_v2float = OpTypeRuntimeArray %v2float
  %_struct_6 = OpTypeStruct %_runtimearr_v2float
%_ptr_StorageBuffer__struct_6 = OpTypePointer StorageBuffer %_struct_6
          %4 = OpVariable %_ptr_StorageBuffer__struct_6 StorageBuffer
       %uint = OpTypeInt 32 0
     %uint_0 = OpConstant %uint 0
%float_0_100000001 = OpConstant %float 0.100000001
       %bool = OpTypeBool
%_ptr_StorageBuffer_v2float = OpTypePointer StorageBuffer %v2float
         %27 = OpConstantComposite %v2float %float_0_100000001 %float_0_100000001
          %1 = OpFunction %void None %8
         %18 = OpLabel
         %20 = OpArrayLength %uint %4 0
         %22 = OpULessThan %bool %uint_0 %20
               OpSelectionMerge %23 None
               OpBranchConditional %22 %24 %25
         %24 = OpLabel
         %26 = OpAccessChain %_ptr_StorageBuffer_v2float %4 %uint_0 %uint_0
               OpStore %26 %27
               OpBranch %23
         %25 = OpLabel
               OpBranch %23
         %23 = OpLabel
               OpReturn
               OpFunctionEnd

So I suspect it's not enough to just not use pointers in structs in rust-gpu code.

[teoxoy]

Do you think a good first step would just be to improve the error?

For sure, some of the validator errors are lacking.

Just out of curiosity, this does validate, even though the Rust code has a pointer in a struct:

But I suspect that's because Rust compiles down to SPIR-V without a pointer in a struct?

In that case the codegen is different, which avoids the issue.

So I suspect it's not enough to just not use pointers in structs in rust-gpu code.

Maybe, but not using pointers in structs is a lot more likely to work (with naga).