How do you declare an atomic input in rustgpu?

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.

Solution

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)
    };
}