Is it possible to index dynamically into a WebGPU storage buffer?

Question

I'm trying write a WGSL shader that reads an octree that is stored in a storage buffer. The problem is, the compiler doesn't like the dynamic index I'm calculating to access leaves in the storage buffer. wgpu produces the following validation error:

thread 'main' panicked at 'wgpu error: Validation Error

Caused by:
    In Device::create_shader_module
    Function [1] 'get_voxel' is invalid
    Expression [68] is invalid
    The expression [67] may only be indexed by a constant

The octree is structured such that I can traverse it on the GPU. The structure is outlined in this NVIDIA paper: https://developer.nvidia.com/gpugems/gpugems2/part-v-image-oriented-computing/chapter-37-octree-textures-gpu

Essentially, the octree is an array of IndirectionGrids, and each IndirectionGrid has exactly 8 GridCells colocated in memory. A grid cell can represent either a pointer to another IndirectionGrid, or some data.

Lets say at its deepest the octree is representing a 16x16x16 grid. I want to get the GridCell at 7,7,7. We know that 7,7,7 is in cell 0 of the root IndirectionGrid because each component of the coordinate is less than the midpoint. If we add up the coordinates components we can get the index of the GridCell for the current IndirectionGrid. Because I'm traversing a tree, I do this at each level. Below is the incomplete code demonstrating this.

The line in question that's causing problems is let cell = grid.cells[grid_index].data;

So ultimately my question is, are dynamic indices allowed somehow? Is there something I can change that will magically make it work? Or is there more background information I need to understand about the tradeoffs WebGPU makes?

struct GridCell {
    data: u32;
};

struct IndirectionGrid {
    cells: array<GridCell, 8>;
};

[[block]]
struct VoxelVolume {
    resolution: vec3<f32>;
    size: vec3<f32>;
    palette: array<u32, 256>;
    indirection_pool: array<IndirectionGrid>;
};

[[group(2), binding(0)]]
var<storage, read> voxel_volume: VoxelVolume;

let COLOR_RED_MASK = 0xFF000000u;
let COLOR_GREEN_MASK = 0x00FF0000u;
let COLOR_BLUE_MASK = 0x0000FF00u;
let COLOR_ALPHA_MASK = 0x000000FFu;

let CELL_TYPE_MASK: u32 = 0xFF000000u;
let CELL_DATA_MASK: u32 = 0x00FFFFFFu;

fn get_voxel(pos: vec3<f32>) -> vec4<f32> {
    let max_depth: i32 = 12;
    let color = vec4<f32>(0.0, 0.0, 0.0, 0.0);
    
    var pool_index = 0u;
    var grid_size = max(max(voxel_volume.size.x, voxel_volume.size.y), voxel_volume.size.z);
    
    for (var i: i32 = 0; i < max_depth; i = i + 1) {
        let grid = voxel_volume.indirection_pool[pool_index];
        let grid_coord_x = select(1u, 0u, pos.x / grid_size < 0.5);
        let grid_coord_y = select(1u, 0u, pos.y / grid_size < 0.5);
        let grid_coord_z = select(1u, 0u, pos.z / grid_size < 0.5);
        let grid_index = grid_coord_x + grid_coord_y * 2u + grid_coord_z * 2u * 2u;
        let cell = grid.cells[grid_index].data;
        let cell_type = cell & CELL_TYPE_MASK >> 16u;

        switch (cell_type) {
            case 1u: {
                pool_index = cell & CELL_DATA_MASK >> 8u;
            }
            case 2u: {
                let palette_index = cell & CELL_DATA_MASK >> 8u;
                let palette_color = voxel_volume.palette[palette_index];
                
                return vec4<f32>(
                    f32(palette_color & COLOR_RED_MASK >> 24u) / 255.0,
                    f32(palette_color & COLOR_GREEN_MASK >> 16u) / 255.0,
                    f32(palette_color & COLOR_BLUE_MASK >> 8u) / 255.0,
                    f32(palette_color & COLOR_ALPHA_MASK) / 255.0
                );
            }
            default: {
                // discard;
                return vec4<f32>(pos.x / 16.0, pos.y / 16.0, pos.z / 16.0, 1.0);
            }
        }
    }

    discard;
}

Answer 1

Indexing into storage buffers is totally fine. What Naga doesn't like is this line:

let cell = grid.cells[grid_index].data;

... because grid isn't a storage buffer, it's just a value "on the stack".

WGSL has recently decided to allow this, but Naga doesn't implement the necessary hacks yet. And personally, I'd not recommend anybody to rely on this. Instead, you could keep a reference to the storage buffer:

let grid = &voxel_volume.indirection_pool[pool_index];
let cell = (*grid).cells[grid_index].data;

This path doesn't require us to make any hacks/workarounds in SPIR-V output.

Also note that latest Naga (and Firefox) print out a more detailed error:

   ┌─ indexing.wgsl:25:39
   │  
25 │   let CELL_DATA_MASK: u32 = 0x00FFFFFFu;
   │ ╭──────────────────────────────────────^
26 │ │ 
27 │ │ fn get_voxel(pos: vec3<f32>) -> vec4<f32> {
28 │ │     let max_depth: i32 = 12;
   · │
40 │ │         let cell = grid.cells[grid_index].data;
   │ │                   ^^^^^^^^^^^^^^^^^^^^^^^ naga::Expression [73]
   · │
65 │ │     discard;
66 │ │ }
   │ ╰─^ naga::Function [1]

Function [1] 'get_voxel' is invalid: 
        Expression [73] is invalid
        The expression [72] may only be indexed by a constant