How do you declare an atomic input in rustgpu?

Question

I am trying to declare an atomic input variable in rustgpu, the way it can be done in glsl.

I tried looking into both the rustgpu code base and the spirv crate documentation.

Fromt he docs I thought this would work:

#![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>;

// useful:
// https://github.com/EmbarkStudios/rust-gpu/blob/54f6978c25b7e168ded04e720b996625b3654ebd/crates/rustc_codegen_spirv/src/symbols.rs#L42
#[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(uniform, descriptor_set = 1, binding = 5)] nodes:
    // &RuntimeArray<Node<u32>>,
    #[spirv(frag_coord)] screen_pos: Vec4,
    out_color: &mut Vec4,
)
{}

However I am getting this error:

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

So I am not certain what the correct syntax is meant to be.

Based on this discussion, it seems there is no dedicated type for it, and it's just a matter of properly applying the atomic operations to a storage input?

I am not sure, but I tried declaring an SSBO with a single u32 and then copy pasted the code snippet with atomics in the prior link. It doesn't compile.

Answer 1

This seems to be the correct way of doing it. You declare your atomic as a mutable storage buffer then use the atomic semantics.

#[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(storage_buffer, descriptor_set = 1, binding = 4)] buffer: &mut [u32; 1],
    #[spirv(storage_buffer, descriptor_set = 1, binding = 5)] nodes: &mut RuntimeArray<
        Node<u32>,
    >,
    #[spirv(frag_coord)] screen_pos: Vec4,
    out_color: &mut Vec4,
)
{
    let old = unsafe {
        spirv_std::arch::atomic_i_add::<
            _,
            { Scope::Invocation as u32 },
            { Semantics::UNIFORM_MEMORY.bits() },
        >(buffer.index_unchecked_mut(0), 1)
    };
}