Skip to content

Commit 17a304c

Browse files
atlv24atlv24
committed
64 bit image atomics
1 parent 86c378d commit 17a304c

39 files changed

+370
-9
lines changed

‎CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ By @ErichDonGubler in [#6456](https://github.com/gfx-rs/wgpu/pull/6456), [#6148]
106106
### New Features
107107

108108
Image atomic support in shaders. By @atlv24 in [#6706](https://github.com/gfx-rs/wgpu/pull/6706)
109+
64 bit image atomic support in shaders. By @atlv24 in [#5537](https://github.com/gfx-rs/wgpu/pull/5537)
109110

110111
#### Naga
111112

‎naga/src/back/glsl/features.rs

+1
Original file line numberDiff line numberDiff line change
@@ -400,6 +400,7 @@ impl<W> Writer<'_, W> {
400400
| StorageFormat::Rgb10a2Uint
401401
| StorageFormat::Rgb10a2Unorm
402402
| StorageFormat::Rg11b10Ufloat
403+
| StorageFormat::R64Uint
403404
| StorageFormat::Rg32Uint
404405
| StorageFormat::Rg32Sint
405406
| StorageFormat::Rg32Float => {

‎naga/src/back/glsl/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -4941,6 +4941,7 @@ fn glsl_storage_format(format: crate::StorageFormat) -> Result<&'static str, Err
49414941
Sf::Rgb10a2Uint => "rgb10_a2ui",
49424942
Sf::Rgb10a2Unorm => "rgb10_a2",
49434943
Sf::Rg11b10Ufloat => "r11f_g11f_b10f",
4944+
Sf::R64Uint => "r64ui",
49444945
Sf::Rg32Uint => "rg32ui",
49454946
Sf::Rg32Sint => "rg32i",
49464947
Sf::Rg32Float => "rg32f",

‎naga/src/back/hlsl/conv.rs

+1
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ impl crate::StorageFormat {
125125
Self::R8Snorm | Self::R16Snorm => "snorm float",
126126
Self::R8Uint | Self::R16Uint | Self::R32Uint => "uint",
127127
Self::R8Sint | Self::R16Sint | Self::R32Sint => "int",
128+
Self::R64Uint => "uint64_t",
128129

129130
Self::Rg16Float | Self::Rg32Float => "float2",
130131
Self::Rg8Unorm | Self::Rg16Unorm => "unorm float2",

‎naga/src/back/msl/writer.rs

+5-1
Original file line numberDiff line numberDiff line change
@@ -1211,7 +1211,11 @@ impl<W: Write> Writer<W> {
12111211
) -> BackendResult {
12121212
write!(self.out, "{level}")?;
12131213
self.put_expression(image, &context.expression, false)?;
1214-
let op = fun.to_msl();
1214+
let op = if context.expression.resolve_type(value).scalar_width() == Some(8) {
1215+
fun.to_msl_64_bit()?
1216+
} else {
1217+
fun.to_msl()
1218+
};
12151219
write!(self.out, ".atomic_{}(", op)?;
12161220
// coordinates in IR are int, but Metal expects uint
12171221
self.put_cast_to_uint_scalar_or_vector(address.coordinate, &context.expression)?;

‎naga/src/back/spv/image.rs

+4
Original file line numberDiff line numberDiff line change
@@ -1253,6 +1253,10 @@ impl BlockContext<'_> {
12531253
class: spirv::StorageClass::Image,
12541254
}));
12551255
let signed = scalar.kind == crate::ScalarKind::Sint;
1256+
if scalar.width == 8 {
1257+
self.writer
1258+
.require_any("64 bit image atomics", &[spirv::Capability::Int64Atomics])?;
1259+
}
12561260
let pointer_id = self.gen_id();
12571261
let coordinates = self.write_image_coordinates(coordinate, array_index, block)?;
12581262
let sample_id = self.writer.get_constant_scalar(crate::Literal::U32(0));

‎naga/src/back/spv/instructions.rs

+1
Original file line numberDiff line numberDiff line change
@@ -1206,6 +1206,7 @@ impl From<crate::StorageFormat> for spirv::ImageFormat {
12061206
Sf::Rgb10a2Uint => Self::Rgb10a2ui,
12071207
Sf::Rgb10a2Unorm => Self::Rgb10A2,
12081208
Sf::Rg11b10Ufloat => Self::R11fG11fB10f,
1209+
Sf::R64Uint => Self::R64ui,
12091210
Sf::Rg32Uint => Self::Rg32ui,
12101211
Sf::Rg32Sint => Self::Rg32i,
12111212
Sf::Rg32Float => Self::Rg32f,

‎naga/src/back/spv/writer.rs

+7-4
Original file line numberDiff line numberDiff line change
@@ -1079,10 +1079,13 @@ impl Writer {
10791079
"storage image format",
10801080
&[spirv::Capability::StorageImageExtendedFormats],
10811081
),
1082-
If::R64ui | If::R64i => self.require_any(
1083-
"64-bit integer storage image format",
1084-
&[spirv::Capability::Int64ImageEXT],
1085-
),
1082+
If::R64ui | If::R64i => {
1083+
self.use_extension("SPV_EXT_shader_image_int64");
1084+
self.require_any(
1085+
"64-bit integer storage image format",
1086+
&[spirv::Capability::Int64ImageEXT],
1087+
)
1088+
}
10861089
If::Unknown
10871090
| If::Rgba32f
10881091
| If::Rgba16f

‎naga/src/back/wgsl/writer.rs

+1
Original file line numberDiff line numberDiff line change
@@ -2079,6 +2079,7 @@ const fn storage_format_str(format: crate::StorageFormat) -> &'static str {
20792079
Sf::Rgb10a2Uint => "rgb10a2uint",
20802080
Sf::Rgb10a2Unorm => "rgb10a2unorm",
20812081
Sf::Rg11b10Ufloat => "rg11b10float",
2082+
Sf::R64Uint => "r64uint",
20822083
Sf::Rg32Uint => "rg32uint",
20832084
Sf::Rg32Sint => "rg32sint",
20842085
Sf::Rg32Float => "rg32float",

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

+1
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,7 @@ fn map_image_format(word: &str) -> Option<crate::StorageFormat> {
420420
"rgba32ui" => Sf::Rgba32Uint,
421421
"rgba16ui" => Sf::Rgba16Uint,
422422
"rgba8ui" => Sf::Rgba8Uint,
423+
"r64ui" => Sf::R64Uint,
423424
"rg32ui" => Sf::Rg32Uint,
424425
"rg16ui" => Sf::Rg16Uint,
425426
"rg8ui" => Sf::Rg8Uint,

‎naga/src/front/spv/convert.rs

+1
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ pub(super) fn map_image_format(word: spirv::Word) -> Result<crate::StorageFormat
105105
Some(spirv::ImageFormat::Rgb10a2ui) => Ok(crate::StorageFormat::Rgb10a2Uint),
106106
Some(spirv::ImageFormat::Rgb10A2) => Ok(crate::StorageFormat::Rgb10a2Unorm),
107107
Some(spirv::ImageFormat::R11fG11fB10f) => Ok(crate::StorageFormat::Rg11b10Ufloat),
108+
Some(spirv::ImageFormat::R64ui) => Ok(crate::StorageFormat::R64Uint),
108109
Some(spirv::ImageFormat::Rg32ui) => Ok(crate::StorageFormat::Rg32Uint),
109110
Some(spirv::ImageFormat::Rg32i) => Ok(crate::StorageFormat::Rg32Sint),
110111
Some(spirv::ImageFormat::Rg32f) => Ok(crate::StorageFormat::Rg32Float),

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

+1
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,7 @@ pub fn map_storage_format(word: &str, span: Span) -> Result<crate::StorageFormat
9595
"rgb10a2uint" => Sf::Rgb10a2Uint,
9696
"rgb10a2unorm" => Sf::Rgb10a2Unorm,
9797
"rg11b10float" => Sf::Rg11b10Ufloat,
98+
"r64uint" => Sf::R64Uint,
9899
"rg32uint" => Sf::Rg32Uint,
99100
"rg32sint" => Sf::Rg32Sint,
100101
"rg32float" => Sf::Rg32Float,

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

+4
Original file line numberDiff line numberDiff line change
@@ -1633,6 +1633,10 @@ impl Parser {
16331633
kind: Float | Sint | Uint,
16341634
width: 4,
16351635
} => Ok(()),
1636+
Scalar {
1637+
kind: Uint,
1638+
width: 8,
1639+
} => Ok(()),
16361640
_ => Err(Error::BadTextureSampleType { span, scalar }),
16371641
}
16381642
}

‎naga/src/front/wgsl/to_wgsl.rs

+1
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,7 @@ impl crate::StorageFormat {
178178
Sf::Rgb10a2Uint => "rgb10a2uint",
179179
Sf::Rgb10a2Unorm => "rgb10a2unorm",
180180
Sf::Rg11b10Ufloat => "rg11b10float",
181+
Sf::R64Uint => "r64uint",
181182
Sf::Rg32Uint => "rg32uint",
182183
Sf::Rg32Sint => "rg32sint",
183184
Sf::Rg32Float => "rg32float",

‎naga/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -642,6 +642,7 @@ pub enum StorageFormat {
642642
Rg11b10Ufloat,
643643

644644
// 64-bit formats
645+
R64Uint,
645646
Rg32Uint,
646647
Rg32Sint,
647648
Rg32Float,

‎naga/src/proc/mod.rs

+6-1
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ impl From<super::StorageFormat> for super::Scalar {
4949
Sf::Rgb10a2Uint => Sk::Uint,
5050
Sf::Rgb10a2Unorm => Sk::Float,
5151
Sf::Rg11b10Ufloat => Sk::Float,
52+
Sf::R64Uint => Sk::Uint,
5253
Sf::Rg32Uint => Sk::Uint,
5354
Sf::Rg32Sint => Sk::Sint,
5455
Sf::Rg32Float => Sk::Float,
@@ -65,7 +66,11 @@ impl From<super::StorageFormat> for super::Scalar {
6566
Sf::Rgba16Unorm => Sk::Float,
6667
Sf::Rgba16Snorm => Sk::Float,
6768
};
68-
super::Scalar { kind, width: 4 }
69+
let width = match format {
70+
Sf::R64Uint => 8,
71+
_ => 4,
72+
};
73+
super::Scalar { kind, width }
6974
}
7075
}
7176

‎naga/src/valid/function.rs

+28
Original file line numberDiff line numberDiff line change
@@ -1231,6 +1231,34 @@ impl super::Validator {
12311231
.with_span_handle(image, context.expressions));
12321232
}
12331233
match format {
1234+
crate::StorageFormat::R64Uint => {
1235+
if !self.capabilities.intersects(
1236+
super::Capabilities::TEXTURE_INT64_ATOMIC,
1237+
) {
1238+
return Err(FunctionError::MissingCapability(
1239+
super::Capabilities::TEXTURE_INT64_ATOMIC,
1240+
)
1241+
.with_span_static(
1242+
span,
1243+
"missing capability for this operation",
1244+
));
1245+
}
1246+
match fun {
1247+
crate::AtomicFunction::Min
1248+
| crate::AtomicFunction::Max => {}
1249+
_ => {
1250+
return Err(
1251+
FunctionError::InvalidImageAtomicFunction(
1252+
fun,
1253+
)
1254+
.with_span_handle(
1255+
image,
1256+
context.expressions,
1257+
),
1258+
);
1259+
}
1260+
}
1261+
}
12341262
crate::StorageFormat::R32Sint
12351263
| crate::StorageFormat::R32Uint => {
12361264
if !self

‎naga/src/valid/mod.rs

+2
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,8 @@ bitflags::bitflags! {
145145
const SHADER_INT64_ATOMIC_ALL_OPS = 0x100000;
146146
/// Support for atomic operations on images.
147147
const TEXTURE_ATOMIC = 0x200000;
148+
/// Support for atomic operations on 64-bit images.
149+
const TEXTURE_INT64_ATOMIC = 0x400000;
148150
}
149151
}
150152

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

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
(
2+
god_mode: true,
3+
spv: (
4+
version: (1, 0),
5+
capabilities: [ Int64, Int64ImageEXT, Int64Atomics ],
6+
),
7+
hlsl: (
8+
shader_model: V6_6,
9+
binding_map: {},
10+
fake_missing_bindings: true,
11+
special_constants_binding: Some((space: 1, register: 0)),
12+
push_constants_target: Some((space: 0, register: 0)),
13+
zero_initialize_workgroup_memory: true,
14+
restrict_indexing: true
15+
),
16+
msl: (
17+
lang_version: (3, 1),
18+
per_entry_point_map: {},
19+
inline_samplers: [],
20+
spirv_cross_compatibility: false,
21+
fake_missing_bindings: true,
22+
zero_initialize_workgroup_memory: true,
23+
),
24+
)

‎naga/tests/in/atomicTexture-int64.wgsl

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
@group(0) @binding(0)
2+
var image: texture_storage_2d<r64uint, atomic>;
3+
4+
@compute
5+
@workgroup_size(2)
6+
fn cs_main(@builtin(local_invocation_id) id: vec3<u32>) {
7+
textureAtomicMax(image, vec2<i32>(0, 0), 1lu);
8+
9+
workgroupBarrier();
10+
11+
textureAtomicMin(image, vec2<i32>(0, 0), 1lu);
12+
}

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

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
struct NagaConstants {
2+
int first_vertex;
3+
int first_instance;
4+
uint other;
5+
};
6+
ConstantBuffer<NagaConstants> _NagaConstants: register(b0, space1);
7+
8+
RWTexture2D<uint64_t> image : register(u0);
9+
10+
[numthreads(2, 1, 1)]
11+
void cs_main(uint3 id : SV_GroupThreadID)
12+
{
13+
InterlockedMax(image[int2(0, 0)],1uL);
14+
GroupMemoryBarrierWithGroupSync();
15+
InterlockedMin(image[int2(0, 0)],1uL);
16+
return;
17+
}

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

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
(
2+
vertex:[
3+
],
4+
fragment:[
5+
],
6+
compute:[
7+
(
8+
entry_point:"cs_main",
9+
target_profile:"cs_6_6",
10+
),
11+
],
12+
)

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

+18
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
// language: metal3.1
2+
#include <metal_stdlib>
3+
#include <simd/simd.h>
4+
5+
using metal::uint;
6+
7+
8+
struct cs_mainInput {
9+
};
10+
kernel void cs_main(
11+
metal::uint3 id [[thread_position_in_threadgroup]]
12+
, metal::texture2d<ulong, metal::access::read_write> image [[user(fake0)]]
13+
) {
14+
image.atomic_max(metal::uint2(metal::int2(0, 0)), 1uL);
15+
metal::threadgroup_barrier(metal::mem_flags::mem_threadgroup);
16+
image.atomic_min(metal::uint2(metal::int2(0, 0)), 1uL);
17+
return;
18+
}

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

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

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

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
@group(0) @binding(0)
2+
var image: texture_storage_2d<r64uint,atomic>;
3+
4+
@compute @workgroup_size(2, 1, 1)
5+
fn cs_main(@builtin(local_invocation_id) id: vec3<u32>) {
6+
textureAtomicMax(image, vec2<i32>(0i, 0i), 1lu);
7+
workgroupBarrier();
8+
textureAtomicMin(image, vec2<i32>(0i, 0i), 1lu);
9+
return;
10+
}

‎naga/tests/snapshots.rs

+4
Original file line numberDiff line numberDiff line change
@@ -792,6 +792,10 @@ fn convert_wgsl() {
792792
"atomicTexture",
793793
Targets::SPIRV | Targets::METAL | Targets::HLSL | Targets::WGSL,
794794
),
795+
(
796+
"atomicTexture-int64",
797+
Targets::SPIRV | Targets::METAL | Targets::HLSL | Targets::WGSL,
798+
),
795799
(
796800
"atomicCompareExchange-int64",
797801
Targets::SPIRV | Targets::WGSL,

‎tests/tests/image_atomics/image_64_atomics.wgsl

+13
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
@group(0) @binding(0)
2+
var image: texture_storage_2d<r64uint, atomic>;
3+
4+
@compute
5+
@workgroup_size(4, 4, 1)
6+
fn cs_main(@builtin(local_invocation_id) id: vec3<u32>, @builtin(workgroup_id) group_id: vec3<u32>) {
7+
let pixel = id + group_id * 4;
8+
textureAtomicMax(image, pixel.xy, u64(pixel.x));
9+
10+
storageBarrier();
11+
12+
textureAtomicMin(image, pixel.xy, u64(pixel.y));
13+
}

0 commit comments

Comments
 (0)