是否可以动态索引到 WebGPU 存储缓冲区？

Question

我正在尝试编写一个 WGSL 着色器来读取存储在存储缓冲区中的八叉树。问题是，编译器不喜欢我正在计算的动态索引来访问存储缓冲区中的叶子。 wgpu 产生以下验证错误：

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

八叉树的结构使得我可以在 GPU 上遍历它。 NVIDIA 论文中概述了该结构：https://developer.nvidia.com/gpugems/gpugems2/part-v-image-oriented-computing/chapter-37-octree-textures-gpu

本质上，八叉树是一个 IndirectionGrid 的数组，每个 IndirectionGrid 恰好有 8 个 GridCell 在内存中并置。网格单元可以表示指向另一个 IndirectionGrid 或一些数据的指针。

假设八叉树最深处代表一个 16x16x16 的网格。我想在 7,7,7 获得 GridCell。我们知道 7,7,7 在根 IndirectionGrid 的单元格 0 中，因为坐标的每个分量都小于中点。如果我们将坐标分量相加，我们可以获得当前 IndirectionGrid 的 GridCell 的索引。因为我正在遍历一棵树，所以我在每个级别都这样做。下面是演示这一点的不完整代码。

引起问题的相关行是 let cell = grid.cells[grid_index].data;

所以最终我的问题是，是否允许以某种方式使用动态索引？有什么我可以改变的，可以神奇地让它发挥作用吗？或者是否有更多我需要了解的关于 WebGPU 权衡的背景信息？

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

索引到存储缓冲区完全没问题。 Naga 不喜欢的是这一行：

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

... 因为 grid 不是存储缓冲区，它只是“堆栈上”的一个值。

WGSL 最近决定允许这样做，但 Naga 尚未实施必要的黑客攻击。就个人而言，我不建议任何人依赖它。相反，您可以保留对存储缓冲区的引用：

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

此路径不需要我们在 SPIR-V 输出中制作任何 hacks/workarounds。

另请注意，最新的 Naga（和 Firefox）会打印出更详细的错误：

   ┌─ 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

是否可以动态索引到 WebGPU 存储缓冲区？

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

shader

rust

webgpu

wgpu-rs