Skip to content

Commit be95178

Browse files
atlv24atlv24
authored
64 bit image atomics (#5537)
1 parent 6c10e0b commit be95178

39 files changed

+382
-20
lines changed

Diff for: ‎CHANGELOG.md

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

170171
##### Vulkan

Diff for: ‎naga/src/back/glsl/features.rs

+1
Original file line numberDiff line numberDiff line change
@@ -408,6 +408,7 @@ impl<W> Writer<'_, W> {
408408
| StorageFormat::Rgb10a2Uint
409409
| StorageFormat::Rgb10a2Unorm
410410
| StorageFormat::Rg11b10Ufloat
411+
| StorageFormat::R64Uint
411412
| StorageFormat::Rg32Uint
412413
| StorageFormat::Rg32Sint
413414
| StorageFormat::Rg32Float => {

Diff for: ‎naga/src/back/glsl/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -4944,6 +4944,7 @@ fn glsl_storage_format(format: crate::StorageFormat) -> Result<&'static str, Err
49444944
Sf::Rgb10a2Uint => "rgb10_a2ui",
49454945
Sf::Rgb10a2Unorm => "rgb10_a2",
49464946
Sf::Rg11b10Ufloat => "r11f_g11f_b10f",
4947+
Sf::R64Uint => "r64ui",
49474948
Sf::Rg32Uint => "rg32ui",
49484949
Sf::Rg32Sint => "rg32i",
49494950
Sf::Rg32Float => "rg32f",

Diff for: ‎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",

Diff for: ‎naga/src/back/msl/writer.rs

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

Diff for: ‎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));

Diff for: ‎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,

Diff for: ‎naga/src/back/spv/writer.rs

+7-4
Original file line numberDiff line numberDiff line change
@@ -1089,10 +1089,13 @@ impl Writer {
10891089
"storage image format",
10901090
&[spirv::Capability::StorageImageExtendedFormats],
10911091
),
1092-
If::R64ui | If::R64i => self.require_any(
1093-
"64-bit integer storage image format",
1094-
&[spirv::Capability::Int64ImageEXT],
1095-
),
1092+
If::R64ui | If::R64i => {
1093+
self.use_extension("SPV_EXT_shader_image_int64");
1094+
self.require_any(
1095+
"64-bit integer storage image format",
1096+
&[spirv::Capability::Int64ImageEXT],
1097+
)
1098+
}
10961099
If::Unknown
10971100
| If::Rgba32f
10981101
| If::Rgba16f

Diff for: ‎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",

Diff for: ‎naga/src/front/glsl/parser/types.rs

+1
Original file line numberDiff line numberDiff line change
@@ -430,6 +430,7 @@ fn map_image_format(word: &str) -> Option<crate::StorageFormat> {
430430
"rgba32ui" => Sf::Rgba32Uint,
431431
"rgba16ui" => Sf::Rgba16Uint,
432432
"rgba8ui" => Sf::Rgba8Uint,
433+
"r64ui" => Sf::R64Uint,
433434
"rg32ui" => Sf::Rg32Uint,
434435
"rg16ui" => Sf::Rg16Uint,
435436
"rg8ui" => Sf::Rg8Uint,

Diff for: ‎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),

Diff for: ‎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,

Diff for: ‎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
}

Diff for: ‎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",

Diff for: ‎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,

Diff for: ‎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

Diff for: ‎naga/src/valid/function.rs

+28
Original file line numberDiff line numberDiff line change
@@ -1282,6 +1282,34 @@ impl super::Validator {
12821282
.with_span_handle(image, context.expressions));
12831283
}
12841284
match format {
1285+
crate::StorageFormat::R64Uint => {
1286+
if !self.capabilities.intersects(
1287+
super::Capabilities::TEXTURE_INT64_ATOMIC,
1288+
) {
1289+
return Err(FunctionError::MissingCapability(
1290+
super::Capabilities::TEXTURE_INT64_ATOMIC,
1291+
)
1292+
.with_span_static(
1293+
span,
1294+
"missing capability for this operation",
1295+
));
1296+
}
1297+
match fun {
1298+
crate::AtomicFunction::Min
1299+
| crate::AtomicFunction::Max => {}
1300+
_ => {
1301+
return Err(
1302+
FunctionError::InvalidImageAtomicFunction(
1303+
fun,
1304+
)
1305+
.with_span_handle(
1306+
image,
1307+
context.expressions,
1308+
),
1309+
);
1310+
}
1311+
}
1312+
}
12851313
crate::StorageFormat::R32Sint
12861314
| crate::StorageFormat::R32Uint => {
12871315
if !self

Diff for: ‎naga/src/valid/mod.rs

+2
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,8 @@ bitflags::bitflags! {
154154
const SHADER_FLOAT32_ATOMIC = 1 << 21;
155155
/// Support for atomic operations on images.
156156
const TEXTURE_ATOMIC = 1 << 22;
157+
/// Support for atomic operations on 64-bit images.
158+
const TEXTURE_INT64_ATOMIC = 1 << 23;
157159
}
158160
}
159161

Diff for: ‎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+
)

Diff for: ‎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+
}

Diff for: ‎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+
}

Diff for: ‎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+
)

Diff for: ‎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+
}

Diff for: ‎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

Diff for: ‎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+
}

Diff for: ‎naga/tests/snapshots.rs

+4
Original file line numberDiff line numberDiff line change
@@ -796,6 +796,10 @@ fn convert_wgsl() {
796796
"atomicOps-float32",
797797
Targets::SPIRV | Targets::METAL | Targets::WGSL,
798798
),
799+
(
800+
"atomicTexture-int64",
801+
Targets::SPIRV | Targets::METAL | Targets::HLSL | Targets::WGSL,
802+
),
799803
(
800804
"atomicCompareExchange-int64",
801805
Targets::SPIRV | Targets::WGSL,

Diff for: ‎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)