Merge pull request #1 from rodolphito/atomic-64

Add u64 and i64 atomic support

naga/src/back/glsl/mod.rs

@@ -2451,6 +2451,9 @@ impl<'a, W: Write> Writer<'a, W> {
                     crate::Literal::I64(_) => {
                         return Err(Error::Custom("GLSL has no 64-bit integer type".into()));
                     }
+                    crate::Literal::U64(_) => {
+                        return Err(Error::Custom("GLSL has no 64-bit integer type".into()));
+                    }
                     crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => {
                         return Err(Error::Custom(
                             "Abstract types should not appear in IR presented to backends".into(),

naga/src/back/hlsl/writer.rs

@@ -2038,6 +2038,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 crate::Literal::F32(value) => write!(self.out, "{value:?}")?,
                 crate::Literal::U32(value) => write!(self.out, "{}u", value)?,
                 crate::Literal::I32(value) => write!(self.out, "{}", value)?,
+                crate::Literal::U64(value) => write!(self.out, "{}uL", value)?,
                 crate::Literal::I64(value) => write!(self.out, "{}L", value)?,
                 crate::Literal::Bool(value) => write!(self.out, "{}", value)?,
                 crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => {

naga/src/back/msl/writer.rs

@@ -1273,6 +1273,9 @@ impl<W: Write> Writer<W> {
                 crate::Literal::I32(value) => {
                     write!(self.out, "{value}")?;
                 }
+                crate::Literal::U64(value) => {
+                    write!(self.out, "{value}uL")?;
+                }
                 crate::Literal::I64(value) => {
                     write!(self.out, "{value}L")?;
                 }

naga/src/back/spv/writer.rs

@@ -1182,6 +1182,9 @@ impl Writer {
             crate::Literal::F32(value) => Instruction::constant_32bit(type_id, id, value.to_bits()),
             crate::Literal::U32(value) => Instruction::constant_32bit(type_id, id, value),
             crate::Literal::I32(value) => Instruction::constant_32bit(type_id, id, value as u32),
+            crate::Literal::U64(value) => {
+                Instruction::constant_64bit(type_id, id, value as u32, (value >> 32) as u32)
+            }
             crate::Literal::I64(value) => {
                 Instruction::constant_64bit(type_id, id, value as u32, (value >> 32) as u32)
             }

naga/src/back/wgsl/writer.rs

@@ -1105,6 +1105,9 @@ impl<W: Write> Writer<W> {
                 crate::Literal::I64(_) => {
                     return Err(Error::Custom("unsupported i64 literal".to_string()));
                 }
+                crate::Literal::U64(_) => {
+                    return Err(Error::Custom("unsupported u64 literal".to_string()));
+                }
                 crate::Literal::AbstractInt(_) | crate::Literal::AbstractFloat(_) => {
                     return Err(Error::Custom(
                         "Abstract types should not appear in IR presented to backends".into(),

naga/src/front/spv/mod.rs

@@ -4868,6 +4868,11 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                 let low = self.next()?;
                 match width {
                     4 => crate::Literal::U32(low),
+                    8 => {
+                        inst.expect(5)?;
+                        let high = self.next()?;
+                        crate::Literal::U64(u64::from(high) << 32 | u64::from(low))
+                    }
                     _ => return Err(Error::InvalidTypeWidth(width as u32)),
                 }
             }

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

@@ -1530,6 +1530,8 @@ impl<'source, 'temp> Lowerer<'source, 'temp> {
                     ast::Literal::Number(Number::F32(f)) => crate::Literal::F32(f),
                     ast::Literal::Number(Number::I32(i)) => crate::Literal::I32(i),
                     ast::Literal::Number(Number::U32(u)) => crate::Literal::U32(u),
+                    ast::Literal::Number(Number::I64(i)) => crate::Literal::I64(i),
+                    ast::Literal::Number(Number::U64(u)) => crate::Literal::U64(u),
                     ast::Literal::Number(Number::F64(f)) => crate::Literal::F64(f),
                     ast::Literal::Number(Number::AbstractInt(i)) => crate::Literal::AbstractInt(i),
                     ast::Literal::Number(Number::AbstractFloat(f)) => {

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

@@ -124,6 +124,14 @@ pub fn get_scalar_type(word: &str) -> Option<Scalar> {
             kind: Sk::Uint,
             width: 4,
         }),
+        "i64" => Some(Scalar {
+            kind: Sk::Sint,
+            width: 8,
+        }),
+        "u64" => Some(Scalar {
+            kind: Sk::Uint,
+            width: 8,
+        }),
         "bool" => Some(Scalar {
             kind: Sk::Bool,
             width: crate::BOOL_WIDTH,

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

@@ -12,6 +12,10 @@ pub enum Number {
     I32(i32),
     /// Concrete u32
     U32(u32),
+    /// Concrete i64
+    I64(i64),
+    /// Concrete u64
+    U64(u64),
     /// Concrete f32
     F32(f32),
     /// Concrete f64
@@ -31,6 +35,8 @@ enum Kind {
 enum IntKind {
     I32,
     U32,
+    I64,
+    U64,
 }
 
 #[derive(Debug)]
@@ -270,6 +276,8 @@ fn parse(input: &str) -> (Result<Number, NumberError>, &str) {
                 let kind = consume_map!(bytes, [
                     b'i' => Kind::Int(IntKind::I32),
                     b'u' => Kind::Int(IntKind::U32),
+                    b'l', b'i' => Kind::Int(IntKind::I64),
+                    b'l', b'u' => Kind::Int(IntKind::U64),
                     b'h' => Kind::Float(FloatKind::F16),
                     b'f' => Kind::Float(FloatKind::F32),
                     b'l', b'f' => Kind::Float(FloatKind::F64),
@@ -416,5 +424,13 @@ fn parse_int(input: &str, kind: Option<IntKind>, radix: u32) -> Result<Number, N
             Ok(num) => Ok(Number::U32(num)),
             Err(e) => Err(map_err(e)),
         },
+        Some(IntKind::I64) => match i64::from_str_radix(input, radix) {
+            Ok(num) => Ok(Number::I64(num)),
+            Err(e) => Err(map_err(e)),
+        },
+        Some(IntKind::U64) => match u64::from_str_radix(input, radix) {
+            Ok(num) => Ok(Number::U64(num)),
+            Err(e) => Err(map_err(e)),
+        },
     }
 }

naga/src/front/wgsl/tests.rs

@@ -17,6 +17,7 @@ fn parse_comment() {
 #[test]
 fn parse_types() {
     parse_str("const a : i32 = 2;").unwrap();
+    parse_str("const a : u64 = 2lu;").unwrap();
     assert!(parse_str("const a : x32 = 2;").is_err());
     parse_str("var t: texture_2d<f32>;").unwrap();
     parse_str("var t: texture_cube_array<i32>;").unwrap();

naga/src/keywords/wgsl.rs

@@ -14,6 +14,7 @@ pub const RESERVED: &[&str] = &[
     "f32",
     "f16",
     "i32",
+    "i64",
     "mat2x2",
     "mat2x3",
     "mat2x4",
@@ -43,6 +44,7 @@ pub const RESERVED: &[&str] = &[
     "texture_depth_cube_array",
     "texture_depth_multisampled_2d",
     "u32",
+    "u64",
     "vec2",
     "vec3",
     "vec4",

naga/src/lib.rs

@@ -886,6 +886,7 @@ pub enum Literal {
     F32(f32),
     U32(u32),
     I32(i32),
+    U64(u64),
     I64(i64),
     Bool(bool),
     AbstractInt(i64),

naga/src/proc/constant_evaluator.rs

@@ -1151,7 +1151,7 @@ impl<'a> ConstantEvaluator<'a> {
                         Literal::U32(v) => v as i32,
                         Literal::F32(v) => v as i32,
                         Literal::Bool(v) => v as i32,
-                        Literal::F64(_) | Literal::I64(_) => {
+                        Literal::F64(_) | Literal::I64(_) | Literal::U64(_) => {
                             return make_error();
                         }
                         Literal::AbstractInt(v) => i32::try_from_abstract(v)?,
@@ -1162,7 +1162,7 @@ impl<'a> ConstantEvaluator<'a> {
                         Literal::U32(v) => v,
                         Literal::F32(v) => v as u32,
                         Literal::Bool(v) => v as u32,
-                        Literal::F64(_) | Literal::I64(_) => {
+                        Literal::F64(_) | Literal::I64(_) | Literal::U64(_) => {
                             return make_error();
                         }
                         Literal::AbstractInt(v) => u32::try_from_abstract(v)?,
@@ -1173,7 +1173,7 @@ impl<'a> ConstantEvaluator<'a> {
                         Literal::U32(v) => v as f32,
                         Literal::F32(v) => v,
                         Literal::Bool(v) => v as u32 as f32,
-                        Literal::F64(_) | Literal::I64(_) => {
+                        Literal::F64(_) | Literal::I64(_) | Literal::U64(_) => {
                             return make_error();
                         }
                         Literal::AbstractInt(v) => f32::try_from_abstract(v)?,
@@ -1185,7 +1185,7 @@ impl<'a> ConstantEvaluator<'a> {
                         Literal::F32(v) => v as f64,
                         Literal::F64(v) => v,
                         Literal::Bool(v) => v as u32 as f64,
-                        Literal::I64(_) => return make_error(),
+                        Literal::I64(_) | Literal::U64(_) => return make_error(),
                         Literal::AbstractInt(v) => f64::try_from_abstract(v)?,
                         Literal::AbstractFloat(v) => f64::try_from_abstract(v)?,
                     }),
@@ -1196,6 +1196,7 @@ impl<'a> ConstantEvaluator<'a> {
                         Literal::Bool(v) => v,
                         Literal::F64(_)
                         | Literal::I64(_)
+                        | Literal::U64(_)
                         | Literal::AbstractInt(_)
                         | Literal::AbstractFloat(_) => {
                             return make_error();

naga/src/proc/mod.rs

@@ -102,6 +102,10 @@ impl super::Scalar {
         kind: crate::ScalarKind::Sint,
         width: 8,
     };
+    pub const U64: Self = Self {
+        kind: crate::ScalarKind::Uint,
+        width: 8,
+    };
     pub const BOOL: Self = Self {
         kind: crate::ScalarKind::Bool,
         width: crate::BOOL_WIDTH,
@@ -156,6 +160,7 @@ impl PartialEq for crate::Literal {
             (Self::F32(a), Self::F32(b)) => a.to_bits() == b.to_bits(),
             (Self::U32(a), Self::U32(b)) => a == b,
             (Self::I32(a), Self::I32(b)) => a == b,
+            (Self::U64(a), Self::U64(b)) => a == b,
             (Self::I64(a), Self::I64(b)) => a == b,
             (Self::Bool(a), Self::Bool(b)) => a == b,
             _ => false,
@@ -186,10 +191,18 @@ impl std::hash::Hash for crate::Literal {
                 hasher.write_u8(4);
                 v.hash(hasher);
             }
-            Self::I64(v) | Self::AbstractInt(v) => {
+            Self::I64(v) => {
                 hasher.write_u8(5);
                 v.hash(hasher);
             }
+            Self::U64(v) => {
+                hasher.write_u8(6);
+                v.hash(hasher);
+            }
+            Self::AbstractInt(v) => {
+                hasher.write_u8(7);
+                v.hash(hasher);
+            }
         }
     }
 }
@@ -201,6 +214,7 @@ impl crate::Literal {
             (value, crate::ScalarKind::Float, 4) => Some(Self::F32(value as _)),
             (value, crate::ScalarKind::Uint, 4) => Some(Self::U32(value as _)),
             (value, crate::ScalarKind::Sint, 4) => Some(Self::I32(value as _)),
+            (value, crate::ScalarKind::Uint, 8) => Some(Self::U64(value as _)),
             (value, crate::ScalarKind::Sint, 8) => Some(Self::I64(value as _)),
             (1, crate::ScalarKind::Bool, 4) => Some(Self::Bool(true)),
             (0, crate::ScalarKind::Bool, 4) => Some(Self::Bool(false)),
@@ -218,7 +232,7 @@ impl crate::Literal {
 
     pub const fn width(&self) -> crate::Bytes {
         match *self {
-            Self::F64(_) | Self::I64(_) => 8,
+            Self::F64(_) | Self::I64(_) | Self::U64(_) => 8,
             Self::F32(_) | Self::U32(_) | Self::I32(_) => 4,
             Self::Bool(_) => crate::BOOL_WIDTH,
             Self::AbstractInt(_) | Self::AbstractFloat(_) => crate::ABSTRACT_WIDTH,
@@ -230,6 +244,7 @@ impl crate::Literal {
             Self::F32(_) => crate::Scalar::F32,
             Self::U32(_) => crate::Scalar::U32,
             Self::I32(_) => crate::Scalar::I32,
+            Self::U64(_) => crate::Scalar::U64,
             Self::I64(_) => crate::Scalar::I64,
             Self::Bool(_) => crate::Scalar::BOOL,
             Self::AbstractInt(_) => crate::Scalar::ABSTRACT_INT,

naga/src/valid/type.rs

@@ -245,11 +245,31 @@ impl super::Validator {
                     scalar.width == 4
                 }
             }
-            crate::ScalarKind::Sint | crate::ScalarKind::Uint => {
+            crate::ScalarKind::Sint => {
                 if scalar.width == 8 {
-                    return Err(WidthError::Unsupported64Bit);
+                    if !self.capabilities.contains(Capabilities::SHADER_I64) {
+                        return Err(WidthError::MissingCapability {
+                            name: "i64",
+                            flag: "INT64",
+                        });
+                    }
+                    true
+                } else {
+                    scalar.width == 4
+                }
+            }
+            crate::ScalarKind::Uint => {
+                if scalar.width == 8 {
+                    if !self.capabilities.contains(Capabilities::SHADER_I64) {
+                        return Err(WidthError::MissingCapability {
+                            name: "u64",
+                            flag: "INT64",
+                        });
+                    }
+                    true
+                } else {
+                    scalar.width == 4
                 }
-                scalar.width == 4
             }
             crate::ScalarKind::AbstractInt | crate::ScalarKind::AbstractFloat => {
                 return Err(WidthError::Abstract);
@@ -335,7 +355,18 @@ impl super::Validator {
                     | crate::ScalarKind::Float
                     | crate::ScalarKind::AbstractInt
                     | crate::ScalarKind::AbstractFloat => false,
-                    crate::ScalarKind::Sint | crate::ScalarKind::Uint => width == 4,
+                    crate::ScalarKind::Sint | crate::ScalarKind::Uint => {
+                        if width == 4 {
+                            true
+                        } else if width == 8 {
+                            self.require_type_capability(
+                                Capabilities::SHADER_I64 | Capabilities::SHADER_I64_TEXTURE_ATOMIC,
+                            )?;
+                            true
+                        } else {
+                            false
+                        }
+                    }
                 };
                 if !good {
                     return Err(TypeError::InvalidAtomicWidth(kind, width));

wgpu-types/src/lib.rs

@@ -237,7 +237,7 @@ bitflags::bitflags! {
     #[repr(transparent)]
     #[derive(Default)]
     #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
-    pub struct Features: u128 {
+    pub struct Features: u64 {
         //
         // ---- Start numbering at 1 << 0 ----
         //
@@ -314,7 +314,27 @@ bitflags::bitflags! {
         /// This is a web and native feature.
         const SHADER_F16 = 1 << 8;
 
-        // 9..14 available
+        /// Allows shaders to use i64 and u64.
+        ///
+        /// Supported platforms:
+        /// - Vulkan
+        /// - DX12
+        ///
+        /// This is a native only feature.
+        const SHADER_I64 = 1 << 9;
+
+        /// Allows compute shaders to use atomic operations on i64 or u64 storage textures.
+        ///
+        /// Requires SHADER_I64 and [`Dx12Compiler::Dxc`].
+        ///
+        /// Supported platforms:
+        /// - Vulkan (with VK_EXT_shader_image_atomic_int64)
+        /// - DX12 (with Shader Model 6.6 and AtomicInt64OnTypedResourceSupported)
+        ///
+        /// This is a native only feature.
+        const SHADER_I64_TEXTURE_ATOMIC = 1 << 10;
+
+        // 11..14 available
 
         // Texture Formats:
 
@@ -853,26 +873,6 @@ bitflags::bitflags! {
         /// - Vulkan (with dualSrcBlend)
         /// - DX12
         const DUAL_SOURCE_BLENDING = 1 << 63;
-
-        /// Allows shaders to use i64 and u64.
-        ///
-        /// Supported platforms:
-        /// - Vulkan
-        /// - DX12
-        ///
-        /// This is a native only feature.
-        const SHADER_I64 = 1 << 64;
-
-        /// Allows compute shaders to use atomic operations on i64 or u64 storage textures.
-        ///
-        /// Requires SHADER_I64 and [`Dx12Compiler::Dxc`].
-        ///
-        /// Supported platforms:
-        /// - Vulkan (with VK_EXT_shader_image_atomic_int64)
-        /// - DX12 (with Shader Model 6.6 and AtomicInt64OnTypedResourceSupported)
-        ///
-        /// This is a native only feature.
-        const SHADER_I64_TEXTURE_ATOMIC = 1 << 64;
     }
 }