64 bit Image Atomics

CHANGELOG.md

@@ -165,6 +165,7 @@ By @wumpf in [#6849](https://github.com/gfx-rs/wgpu/pull/6849).
 - Add build support for Apple Vision Pro. By @guusw in [#6611](https://github.com/gfx-rs/wgpu/pull/6611).
 - Add `wgsl_language_features` for obtaining available WGSL language feature by @sagudev in [#6814](https://github.com/gfx-rs/wgpu/pull/6814)
 - Image atomic support in shaders. By @atlv24 in [#6706](https://github.com/gfx-rs/wgpu/pull/6706)
+- 64 bit image atomic support in shaders. By @atlv24 in [#5537](https://github.com/gfx-rs/wgpu/pull/5537)
 - Add `no_std` support to `wgpu-types`. By @bushrat011899 in [#6892](https://github.com/gfx-rs/wgpu/pull/6892).
 
 ##### Vulkan

naga/src/back/glsl/features.rs

@@ -408,6 +408,7 @@ impl<W> Writer<'_, W> {
                             | StorageFormat::Rgb10a2Uint
                             | StorageFormat::Rgb10a2Unorm
                             | StorageFormat::Rg11b10Ufloat
+                            | StorageFormat::R64Uint
                             | StorageFormat::Rg32Uint
                             | StorageFormat::Rg32Sint
                             | StorageFormat::Rg32Float => {

naga/src/back/glsl/mod.rs

@@ -4944,6 +4944,7 @@ fn glsl_storage_format(format: crate::StorageFormat) -> Result<&'static str, Err
         Sf::Rgb10a2Uint => "rgb10_a2ui",
         Sf::Rgb10a2Unorm => "rgb10_a2",
         Sf::Rg11b10Ufloat => "r11f_g11f_b10f",
+        Sf::R64Uint => "r64ui",
         Sf::Rg32Uint => "rg32ui",
         Sf::Rg32Sint => "rg32i",
         Sf::Rg32Float => "rg32f",

naga/src/back/hlsl/conv.rs

@@ -125,6 +125,7 @@ impl crate::StorageFormat {
             Self::R8Snorm | Self::R16Snorm => "snorm float",
             Self::R8Uint | Self::R16Uint | Self::R32Uint => "uint",
             Self::R8Sint | Self::R16Sint | Self::R32Sint => "int",
+            Self::R64Uint => "uint64_t",
 
             Self::Rg16Float | Self::Rg32Float => "float2",
             Self::Rg8Unorm | Self::Rg16Unorm => "unorm float2",

naga/src/back/msl/writer.rs

@@ -1212,7 +1212,11 @@ impl<W: Write> Writer<W> {
     ) -> BackendResult {
         write!(self.out, "{level}")?;
         self.put_expression(image, &context.expression, false)?;
-        let op = fun.to_msl();
+        let op = if context.expression.resolve_type(value).scalar_width() == Some(8) {
+            fun.to_msl_64_bit()?
+        } else {
+            fun.to_msl()
+        };
         write!(self.out, ".atomic_{}(", op)?;
         // coordinates in IR are int, but Metal expects uint
         self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;

naga/src/back/spv/image.rs

@@ -1253,6 +1253,10 @@ impl BlockContext<'_> {
             class: spirv::StorageClass::Image,
         }));
         let signed = scalar.kind == crate::ScalarKind::Sint;
+        if scalar.width == 8 {
+            self.writer
+                .require_any("64 bit image atomics", &[spirv::Capability::Int64Atomics])?;
+        }
         let pointer_id = self.gen_id();
         let coordinates = self.write_image_coordinates(coordinate, array_index, block)?;
         let sample_id = self.writer.get_constant_scalar(crate::Literal::U32(0));

naga/src/back/spv/instructions.rs

@@ -1206,6 +1206,7 @@ impl From<crate::StorageFormat> for spirv::ImageFormat {
             Sf::Rgb10a2Uint => Self::Rgb10a2ui,
             Sf::Rgb10a2Unorm => Self::Rgb10A2,
             Sf::Rg11b10Ufloat => Self::R11fG11fB10f,
+            Sf::R64Uint => Self::R64ui,
             Sf::Rg32Uint => Self::Rg32ui,
             Sf::Rg32Sint => Self::Rg32i,
             Sf::Rg32Float => Self::Rg32f,

naga/src/back/spv/writer.rs

@@ -1089,10 +1089,13 @@ impl Writer {
                 "storage image format",
                 &[spirv::Capability::StorageImageExtendedFormats],
             ),
-            If::R64ui | If::R64i => self.require_any(
-                "64-bit integer storage image format",
-                &[spirv::Capability::Int64ImageEXT],
-            ),
+            If::R64ui | If::R64i => {
+                self.use_extension("SPV_EXT_shader_image_int64");
+                self.require_any(
+                    "64-bit integer storage image format",
+                    &[spirv::Capability::Int64ImageEXT],
+                )
+            }
             If::Unknown
             | If::Rgba32f
             | If::Rgba16f

naga/src/back/wgsl/writer.rs

@@ -2079,6 +2079,7 @@ const fn storage_format_str(format: crate::StorageFormat) -> &'static str {
         Sf::Rgb10a2Uint => "rgb10a2uint",
         Sf::Rgb10a2Unorm => "rgb10a2unorm",
         Sf::Rg11b10Ufloat => "rg11b10float",
+        Sf::R64Uint => "r64uint",
         Sf::Rg32Uint => "rg32uint",
         Sf::Rg32Sint => "rg32sint",
         Sf::Rg32Float => "rg32float",

naga/src/front/glsl/parser/types.rs

@@ -430,6 +430,7 @@ fn map_image_format(word: &str) -> Option<crate::StorageFormat> {
         "rgba32ui" => Sf::Rgba32Uint,
         "rgba16ui" => Sf::Rgba16Uint,
         "rgba8ui" => Sf::Rgba8Uint,
+        "r64ui" => Sf::R64Uint,
         "rg32ui" => Sf::Rg32Uint,
         "rg16ui" => Sf::Rg16Uint,
         "rg8ui" => Sf::Rg8Uint,

naga/src/front/spv/convert.rs

@@ -105,6 +105,7 @@ pub(super) fn map_image_format(word: spirv::Word) -> Result<crate::StorageFormat
         Some(spirv::ImageFormat::Rgb10a2ui) => Ok(crate::StorageFormat::Rgb10a2Uint),
         Some(spirv::ImageFormat::Rgb10A2) => Ok(crate::StorageFormat::Rgb10a2Unorm),
         Some(spirv::ImageFormat::R11fG11fB10f) => Ok(crate::StorageFormat::Rg11b10Ufloat),
+        Some(spirv::ImageFormat::R64ui) => Ok(crate::StorageFormat::R64Uint),
         Some(spirv::ImageFormat::Rg32ui) => Ok(crate::StorageFormat::Rg32Uint),
         Some(spirv::ImageFormat::Rg32i) => Ok(crate::StorageFormat::Rg32Sint),
         Some(spirv::ImageFormat::Rg32f) => Ok(crate::StorageFormat::Rg32Float),

naga/src/front/wgsl/parse/conv.rs

@@ -95,6 +95,7 @@ pub fn map_storage_format(word: &str, span: Span) -> Result<crate::StorageFormat
         "rgb10a2uint" => Sf::Rgb10a2Uint,
         "rgb10a2unorm" => Sf::Rgb10a2Unorm,
         "rg11b10float" => Sf::Rg11b10Ufloat,
+        "r64uint" => Sf::R64Uint,
         "rg32uint" => Sf::Rg32Uint,
         "rg32sint" => Sf::Rg32Sint,
         "rg32float" => Sf::Rg32Float,

naga/src/front/wgsl/parse/mod.rs

@@ -1633,6 +1633,10 @@ impl Parser {
                 kind: Float | Sint | Uint,
                 width: 4,
             } => Ok(()),
+            Scalar {
+                kind: Uint,
+                width: 8,
+            } => Ok(()),
             _ => Err(Error::BadTextureSampleType { span, scalar }),
         }
     }

naga/src/front/wgsl/to_wgsl.rs

@@ -178,6 +178,7 @@ impl crate::StorageFormat {
             Sf::Rgb10a2Uint => "rgb10a2uint",
             Sf::Rgb10a2Unorm => "rgb10a2unorm",
             Sf::Rg11b10Ufloat => "rg11b10float",
+            Sf::R64Uint => "r64uint",
             Sf::Rg32Uint => "rg32uint",
             Sf::Rg32Sint => "rg32sint",
             Sf::Rg32Float => "rg32float",

naga/src/lib.rs

@@ -642,6 +642,7 @@ pub enum StorageFormat {
     Rg11b10Ufloat,
 
     // 64-bit formats
+    R64Uint,
     Rg32Uint,
     Rg32Sint,
     Rg32Float,

naga/src/proc/mod.rs

@@ -49,6 +49,7 @@ impl From<super::StorageFormat> for super::Scalar {
             Sf::Rgb10a2Uint => Sk::Uint,
             Sf::Rgb10a2Unorm => Sk::Float,
             Sf::Rg11b10Ufloat => Sk::Float,
+            Sf::R64Uint => Sk::Uint,
             Sf::Rg32Uint => Sk::Uint,
             Sf::Rg32Sint => Sk::Sint,
             Sf::Rg32Float => Sk::Float,
@@ -65,7 +66,11 @@ impl From<super::StorageFormat> for super::Scalar {
             Sf::Rgba16Unorm => Sk::Float,
             Sf::Rgba16Snorm => Sk::Float,
         };
-        super::Scalar { kind, width: 4 }
+        let width = match format {
+            Sf::R64Uint => 8,
+            _ => 4,
+        };
+        super::Scalar { kind, width }
     }
 }
 

naga/src/valid/function.rs

@@ -1282,6 +1282,34 @@ impl super::Validator {
                                         .with_span_handle(image, context.expressions));
                                     }
                                     match format {
+                                        crate::StorageFormat::R64Uint => {
+                                            if !self.capabilities.intersects(
+                                                super::Capabilities::TEXTURE_INT64_ATOMIC,
+                                            ) {
+                                                return Err(FunctionError::MissingCapability(
+                                                    super::Capabilities::TEXTURE_INT64_ATOMIC,
+                                                )
+                                                .with_span_static(
+                                                    span,
+                                                    "missing capability for this operation",
+                                                ));
+                                            }
+                                            match fun {
+                                                crate::AtomicFunction::Min
+                                                | crate::AtomicFunction::Max => {}
+                                                _ => {
+                                                    return Err(
+                                                        FunctionError::InvalidImageAtomicFunction(
+                                                            fun,
+                                                        )
+                                                        .with_span_handle(
+                                                            image,
+                                                            context.expressions,
+                                                        ),
+                                                    );
+                                                }
+                                            }
+                                        }
                                         crate::StorageFormat::R32Sint
                                         | crate::StorageFormat::R32Uint => {
                                             if !self

naga/src/valid/mod.rs

@@ -154,6 +154,8 @@ bitflags::bitflags! {
         const SHADER_FLOAT32_ATOMIC = 1 << 21;
         /// Support for atomic operations on images.
         const TEXTURE_ATOMIC = 1 << 22;
+        /// Support for atomic operations on 64-bit images.
+        const TEXTURE_INT64_ATOMIC = 1 << 23;
     }
 }
 

naga/tests/in/atomicTexture-int64.param.ron

@@ -0,0 +1,24 @@
+(
+	god_mode: true,
+	spv: (
+		version: (1, 0),
+		capabilities: [ Int64, Int64ImageEXT, Int64Atomics ],
+	),
+	hlsl: (
+		shader_model: V6_6,
+		binding_map: {},
+		fake_missing_bindings: true,
+		special_constants_binding: Some((space: 1, register: 0)),
+		push_constants_target: Some((space: 0, register: 0)),
+		zero_initialize_workgroup_memory: true,
+		restrict_indexing: true
+	),
+	msl: (
+		lang_version: (3, 1),
+		per_entry_point_map: {},
+		inline_samplers: [],
+		spirv_cross_compatibility: false,
+		fake_missing_bindings: true,
+		zero_initialize_workgroup_memory: true,
+	),
+)

naga/tests/in/atomicTexture-int64.wgsl

@@ -0,0 +1,12 @@
+@group(0) @binding(0)
+var image: texture_storage_2d<r64uint, atomic>;
+
+@compute
+@workgroup_size(2)
+fn cs_main(@builtin(local_invocation_id) id: vec3<u32>) {
+    textureAtomicMax(image, vec2<i32>(0, 0), 1lu);
+
+    workgroupBarrier();
+
+    textureAtomicMin(image, vec2<i32>(0, 0), 1lu);
+}

naga/tests/out/hlsl/atomicTexture-int64.hlsl

@@ -0,0 +1,17 @@
+struct NagaConstants {
+    int first_vertex;
+    int first_instance;
+    uint other;
+};
+ConstantBuffer<NagaConstants> _NagaConstants: register(b0, space1);
+
+RWTexture2D<uint64_t> image : register(u0);
+
+[numthreads(2, 1, 1)]
+void cs_main(uint3 id : SV_GroupThreadID)
+{
+    InterlockedMax(image[int2(0, 0)],1uL);
+    GroupMemoryBarrierWithGroupSync();
+    InterlockedMin(image[int2(0, 0)],1uL);
+    return;
+}

naga/tests/out/hlsl/atomicTexture-int64.ron

@@ -0,0 +1,12 @@
+(
+    vertex:[
+    ],
+    fragment:[
+    ],
+    compute:[
+        (
+            entry_point:"cs_main",
+            target_profile:"cs_6_6",
+        ),
+    ],
+)

naga/tests/out/msl/atomicTexture-int64.msl

@@ -0,0 +1,18 @@
+// language: metal3.1
+#include <metal_stdlib>
+#include <simd/simd.h>
+
+using metal::uint;
+
+
+struct cs_mainInput {
+};
+kernel void cs_main(
+  metal::uint3 id [[thread_position_in_threadgroup]]
+, metal::texture2d<ulong, metal::access::read_write> image [[user(fake0)]]
+) {
+    image.atomic_max(metal::uint2(metal::int2(0, 0)), 1uL);
+    metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
+    image.atomic_min(metal::uint2(metal::int2(0, 0)), 1uL);
+    return;
+}

naga/tests/out/spv/atomicTexture-int64.spvasm

@@ -0,0 +1,49 @@
+; SPIR-V
+; Version: 1.0
+; Generator: rspirv
+; Bound: 31
+OpCapability Shader
+OpCapability Int64ImageEXT
+OpCapability Int64
+OpCapability Int64Atomics
+OpExtension "SPV_EXT_shader_image_int64"
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %15 "cs_main" %12
+OpExecutionMode %15 LocalSize 2 1 1
+OpDecorate %9 DescriptorSet 0
+OpDecorate %9 Binding 0
+OpDecorate %12 BuiltIn LocalInvocationId
+%2 = OpTypeVoid
+%4 = OpTypeInt 64 0
+%3 = OpTypeImage %4 2D 0 0 0 2 R64ui
+%6 = OpTypeInt 32 0
+%5 = OpTypeVector %6 3
+%8 = OpTypeInt 32 1
+%7 = OpTypeVector %8 2
+%10 = OpTypePointer UniformConstant %3
+%9 = OpVariable  %10  UniformConstant
+%13 = OpTypePointer Input %5
+%12 = OpVariable  %13  Input
+%16 = OpTypeFunction %2
+%18 = OpConstant  %8  0
+%19 = OpConstantComposite  %7  %18 %18
+%20 = OpConstant  %4  1
+%22 = OpTypePointer Image %4
+%24 = OpConstant  %6  0
+%26 = OpConstant  %8  4
+%27 = OpConstant  %6  2
+%28 = OpConstant  %6  264
+%15 = OpFunction  %2  None %16
+%11 = OpLabel
+%14 = OpLoad  %5  %12
+%17 = OpLoad  %3  %9
+OpBranch %21
+%21 = OpLabel
+%23 = OpImageTexelPointer  %22  %9 %19 %24
+%25 = OpAtomicUMax  %4  %23 %26 %24 %20
+OpControlBarrier %27 %27 %28
+%29 = OpImageTexelPointer  %22  %9 %19 %24
+%30 = OpAtomicUMin  %4  %29 %26 %24 %20
+OpReturn
+OpFunctionEnd

naga/tests/out/wgsl/atomicTexture-int64.wgsl

@@ -0,0 +1,10 @@
+@group(0) @binding(0) 
+var image: texture_storage_2d<r64uint,atomic>;
+
+@compute @workgroup_size(2, 1, 1) 
+fn cs_main(@builtin(local_invocation_id) id: vec3<u32>) {
+    textureAtomicMax(image, vec2<i32>(0i, 0i), 1lu);
+    workgroupBarrier();
+    textureAtomicMin(image, vec2<i32>(0i, 0i), 1lu);
+    return;
+}

naga/tests/snapshots.rs

@@ -796,6 +796,10 @@ fn convert_wgsl() {
             "atomicOps-float32",
             Targets::SPIRV | Targets::METAL | Targets::WGSL,
         ),
+        (
+            "atomicTexture-int64",
+            Targets::SPIRV | Targets::METAL | Targets::HLSL | Targets::WGSL,
+        ),
         (
             "atomicCompareExchange-int64",
             Targets::SPIRV | Targets::WGSL,

tests/tests/image_atomics/image_64_atomics.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0)
+var image: texture_storage_2d<r64uint, atomic>;
+
+@compute
+@workgroup_size(4, 4, 1)
+fn cs_main(@builtin(local_invocation_id) id: vec3<u32>, @builtin(workgroup_id) group_id: vec3<u32>) {
+    let pixel = id + group_id * 4;
+    textureAtomicMax(image, pixel.xy, u64(pixel.x));
+
+    storageBarrier();
+
+    textureAtomicMin(image, pixel.xy, u64(pixel.y));
+}