如何在rust-gpu中正确声明原子输入变量？

问题描述 #

我想在rust-gpu中声明类似GLSL原子计数器的原子输入变量，查阅相关资源后编写了代码，却出现编译错误：

error: Operand 2 of TypePointer requires one of these capabilities: AtomicStorage 
         %_ptr_AtomicCounter__struct_16 = OpTypePointer AtomicCounter %_struct_16

根据相关讨论，rust-gpu没有专用的原子计数器类型，需要对存储输入应用原子操作，但我尝试声明单个u32的SSBO后仍无法编译，求正确语法。

尝试的代码 #

#![cfg_attr(target_arch = "spirv", no_std, feature(lang_items))]
#![allow(internal_features)]

extern crate bytemuck;
extern crate spirv_std;

use core::sync::atomic::AtomicU32;

use shader_utils::ray_tracing::*;
use shader_utils::voxel_oct_tree::*;
use spirv_std::{arch::IndexUnchecked, glam::*, image::*, spirv, RuntimeArray};

type Texture2D = Image!(2D, type=f32, sampled);
type SampledTexture2D = SampledImage<Texture2D>;

#[spirv(fragment)]
pub fn main_fs(
    #[spirv(uniform, descriptor_set = 1, binding = 0)] volume_data: &VoxelVolumeMeta,
    #[spirv(uniform, descriptor_set = 1, binding = 1)] window_data: &WindowData,
    #[spirv(uniform, descriptor_set = 1, binding = 2)] voctree_levels: &u32,
    #[spirv(descriptor_set = 1, binding = 3)] volume: &Storage3D,
    #[spirv(atomic_counter, descriptor_set = 1, binding = 4)] node_count: &AtomicU32,
    #[spirv(frag_coord)] screen_pos: Vec4,
    out_color: &mut Vec4,
)
{}

正确实现方式 #

rust-gpu不支持GLSL风格的atomic_counter属性，需改用**存储缓冲区（SSBO）**结合原子类型实现，具体步骤如下：

1. 移除atomic_counter属性：该属性在rust-gpu中无对应实现，会触发能力缺失错误。
2. 用结构体包裹原子变量：SPIR-V要求存储缓冲区绑定聚合类型（结构体/数组），不能直接绑定单个原子变量，需创建结构体并实现bytemuck::Pod和bytemuck::Zeroable以确保内存安全。
3. 通过storage_buffer声明绑定：使用该属性标记SSBO，并通过core::sync::atomic提供的方法执行原子操作。

修正后的代码示例 #

#![cfg_attr(target_arch = "spirv", no_std, feature(lang_items))]
#![allow(internal_features)]

extern crate bytemuck;
extern crate spirv_std;

use core::sync::atomic::{AtomicU32, Ordering};
use bytemuck::{Pod, Zeroable};

use shader_utils::ray_tracing::*;
use shader_utils::voxel_oct_tree::*;
use spirv_std::{arch::IndexUnchecked, glam::*, image::*, spirv, RuntimeArray};

type Texture2D = Image!(2D, type=f32, sampled);
type SampledTexture2D = SampledImage<Texture2D>;

// 定义SSBO结构体，包裹原子计数器变量
#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct AtomicCounterBuffer {
    node_count: AtomicU32,
}

#[spirv(fragment)]
pub fn main_fs(
    #[spirv(uniform, descriptor_set = 1, binding = 0)] volume_data: &VoxelVolumeMeta,
    #[spirv(uniform, descriptor_set = 1, binding = 1)] window_data: &WindowData,
    #[spirv(uniform, descriptor_set = 1, binding = 2)] voctree_levels: &u32,
    #[spirv(descriptor_set = 1, binding = 3)] volume: &Storage3D,
    // 使用storage_buffer属性声明SSBO绑定
    #[spirv(storage_buffer, descriptor_set = 1, binding = 4)] counter_buffer: &AtomicCounterBuffer,
    #[spirv(frag_coord)] screen_pos: Vec4,
    out_color: &mut Vec4,
) {
    // 示例：原子自增操作
    let _ = counter_buffer.node_count.fetch_add(1, Ordering::Relaxed);
    // 示例：加载当前计数器值
    let current_count = counter_buffer.node_count.load(Ordering::Relaxed);
}

关键说明 #

• 聚合类型要求：SPIR-V规范强制存储缓冲区绑定聚合类型，因此必须用结构体包裹原子变量。
• 内存顺序选择：原子操作需指定内存顺序（如Ordering::Relaxed、Ordering::AcqRel），需根据同步需求合理选择。
• 自动能力启用：使用storage_buffer结合原子类型会自动启用AtomicStorage能力，解决编译时的能力缺失错误。

内容的提问来源于stack exchange，提问作者Makogan