fix(naga): Detect overflowing const shift amounts at compile time

Author: andyleiserson

CHANGELOG.md

@@ -99,6 +99,7 @@ Bottom level categories:
 - Fixed `workgroupUniformLoad` incorrectly returning an atomic when called on an atomic, it now returns the inner `T` as per the spec. By @cryvosh in [#8791](https://github.com/gfx-rs/wgpu/pull/8791).
 - Fixed constant evaluation for `sign()` builtin to return zero when the argument is zero. By @mandryskowski in [#8942](https://github.com/gfx-rs/wgpu/pull/8942).
 - Allow array generation to compile with the macOS 10.12 Metal compiler. By @madsmtm in [#8953](https://github.com/gfx-rs/wgpu/pull/8953)
+- Naga now detects bitwise shifts by a constant exceeding the operand bit width at compile time, and disallows scalar-by-vector and vector-by-scalar shifts in constant evaluation. By @andyleiserson in [#8907](https://github.com/gfx-rs/wgpu/pull/8907).
 
 #### Validation
 

cts_runner/fail.lst

@@ -91,7 +91,7 @@ webgpu:shader,validation,expression,access,matrix:* // 93%, runtime OOB matrix a
 webgpu:shader,validation,expression,access,vector:* // 52%, https://github.com/gfx-rs/wgpu/issues/4390, and missing swizzle validation
 webgpu:shader,validation,expression,binary,add_sub_mul:* // 95%, u32 const-eval overflow incorrectly rejected, f16 const-eval overflow not rejected, atomics #5474
 webgpu:shader,validation,expression,binary,and_or_xor:* // 96%, https://github.com/gfx-rs/wgpu/issues/5474
-webgpu:shader,validation,expression,binary,bitwise_shift:* // 97%, atomics https://github.com/gfx-rs/wgpu/issues/5474, partial eval errors
+webgpu:shader,validation,expression,binary,bitwise_shift:invalid_types:* // 93%, atomics #5474
 webgpu:shader,validation,expression,binary,comparison:* // 74%, https://github.com/gfx-rs/wgpu/issues/5474
 webgpu:shader,validation,expression,binary,div_rem:* // 86%, https://github.com/gfx-rs/wgpu/issues/5474
 webgpu:shader,validation,expression,binary,short_circuiting_and_or:* // 92%, https://github.com/gfx-rs/wgpu/issues/8440

cts_runner/test.lst

@@ -326,7 +326,12 @@ webgpu:shader,validation,expression,access,array:early_eval_errors:case="overrid
 webgpu:shader,validation,expression,access,structure:*
 webgpu:shader,validation,expression,binary,add_sub_mul:scalar_vector_out_of_range:lhs="i32";*
 webgpu:shader,validation,expression,binary,add_sub_mul:scalar_vector_out_of_range:lhs="u32";*
+webgpu:shader,validation,expression,binary,bitwise_shift:partial_eval_errors:*
 webgpu:shader,validation,expression,binary,bitwise_shift:scalar_vector:*
+webgpu:shader,validation,expression,binary,bitwise_shift:shift_left_abstract:*
+webgpu:shader,validation,expression,binary,bitwise_shift:shift_left_concrete:*
+webgpu:shader,validation,expression,binary,bitwise_shift:shift_right_abstract:*
+webgpu:shader,validation,expression,binary,bitwise_shift:shift_right_concrete:*
 webgpu:shader,validation,expression,binary,parse:*
 webgpu:shader,validation,expression,binary,short_circuiting_and_or:array_override:op="%26%26";a_val=1;b_val=1
 webgpu:shader,validation,expression,binary,short_circuiting_and_or:invalid_types:*

naga/src/proc/constant_evaluator.rs

@@ -3040,6 +3040,10 @@ impl<'a> ConstantEvaluator<'a> {
         h
     }
 
+    /// Resolve the type of `expr` if it is a constant expression.
+    ///
+    /// If `expr` was evaluated to a constant, returns its type.
+    /// Otherwise, returns an error.
     fn resolve_type(
         &self,
         expr: Handle<Expression>,

naga/src/proc/mod.rs

@@ -475,7 +475,12 @@ pub struct GlobalCtx<'a> {
 }
 
 impl GlobalCtx<'_> {
-    /// Try to evaluate the expression in `self.global_expressions` using its `handle` and return it as a `u32`.
+    /// Try to evaluate the expression in `self.global_expressions` using its `handle`
+    /// and return it as a `T: TryFrom<ir::Literal>`.
+    ///
+    /// This currently only evaluates scalar expressions. If adding support for vectors,
+    /// consider changing `valid::expression::validate_constant_shift_amounts` to use that
+    /// support.
     #[cfg_attr(
         not(any(
             feature = "glsl-in",

naga/src/valid/expression.rs

@@ -148,6 +148,11 @@ pub enum ExpressionError {
     UnsupportedWidth(crate::MathFunction, crate::ScalarKind, crate::Bytes),
     #[error("Invalid operand for cooperative op")]
     InvalidCooperativeOperand(Handle<crate::Expression>),
+    #[error("Shift amount exceeds the bit width of {lhs_type:?}")]
+    ShiftAmountTooLarge {
+        lhs_type: crate::TypeInner,
+        rhs_expr: Handle<crate::Expression>,
+    },
 }
 
 #[derive(Clone, Debug, thiserror::Error)]
@@ -243,6 +248,74 @@ impl super::Validator {
         Ok(())
     }
 
+    /// Return an error if a constant shift amount in `right` exceeds the bit
+    /// width of `left_ty`.
+    ///
+    /// This function promises to return an error in cases where (1) the
+    /// expression is well-typed, (2) `left_ty` is a concrete integer, and
+    /// (3) the shift will overflow. It does not return an error in cases where
+    /// the expression is not well-typed (e.g. vector dimension mismatch),
+    /// because those will be rejected elsewhere.
+    fn validate_constant_shift_amounts(
+        left_ty: &crate::TypeInner,
+        right: Handle<crate::Expression>,
+        module: &crate::Module,
+        function: &crate::Function,
+    ) -> Result<(), ExpressionError> {
+        fn is_overflowing_shift(
+            left_ty: &crate::TypeInner,
+            right: Handle<crate::Expression>,
+            module: &crate::Module,
+            function: &crate::Function,
+        ) -> bool {
+            let Some((vec_size, scalar)) = left_ty.vector_size_and_scalar() else {
+                return false;
+            };
+            if !matches!(
+                scalar.kind,
+                crate::ScalarKind::Sint | crate::ScalarKind::Uint
+            ) {
+                return false;
+            }
+            let lhs_bits = u32::from(8 * scalar.width);
+            if vec_size.is_none() {
+                let shift_amount = module
+                    .to_ctx()
+                    .get_const_val_from::<u32, _>(right, &function.expressions);
+                shift_amount.ok().is_some_and(|s| s >= lhs_bits)
+            } else {
+                match function.expressions[right] {
+                    crate::Expression::ZeroValue(_) => false, // zero shift does not overflow
+                    crate::Expression::Splat { value, .. } => module
+                        .to_ctx()
+                        .get_const_val_from::<u32, _>(value, &function.expressions)
+                        .ok()
+                        .is_some_and(|s| s >= lhs_bits),
+                    crate::Expression::Compose {
+                        ty: _,
+                        ref components,
+                    } => components.iter().any(|comp| {
+                        module
+                            .to_ctx()
+                            .get_const_val_from::<u32, _>(*comp, &function.expressions)
+                            .ok()
+                            .is_some_and(|s| s >= lhs_bits)
+                    }),
+                    _ => false,
+                }
+            }
+        }
+
+        if is_overflowing_shift(left_ty, right, module, function) {
+            Err(ExpressionError::ShiftAmountTooLarge {
+                lhs_type: left_ty.clone(),
+                rhs_expr: right,
+            })
+        } else {
+            Ok(())
+        }
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub(super) fn validate_expression(
         &self,
@@ -984,6 +1057,10 @@ impl super::Validator {
                         rhs_type: right_inner.clone(),
                     });
                 }
+                // For shift operations, check if the constant shift amount exceeds the bit width
+                if matches!(op, Bo::ShiftLeft | Bo::ShiftRight) {
+                    Self::validate_constant_shift_amounts(left_inner, right, module, function)?;
+                }
                 ShaderStages::all()
             }
             E::Select {

naga/tests/naga/wgsl_errors.rs

@@ -31,6 +31,18 @@ fn check(input: &str, snapshot: &str) {
     }
 }
 
+#[track_caller]
+fn check_error_matches(input: &str, expected_substring: &str) {
+    let result = naga::front::wgsl::parse_str(input);
+    let Err(ref err) = result else {
+        panic!("expected ParseError, got {result:#?}");
+    };
+    let message = err.message();
+    if !message.contains(expected_substring) {
+        panic!("expected error containing '{expected_substring}', got '{message}'",);
+    }
+}
+
 #[track_caller]
 fn check_success(input: &str) {
     match naga::front::wgsl::parse_str(input) {
@@ -3816,18 +3828,10 @@ fn inconsistent_type() {
 fn more_inconsistent_type() {
     #[track_caller]
     fn variant(call: &str) {
-        let input = format!(
-            r#"
-            fn f() {{ var x = {call}; }}
-        "#
+        check_error_matches(
+            &format!("fn f() {{ var x = {call}; }}"),
+            "inconsistent type",
         );
-        let result = naga::front::wgsl::parse_str(&input);
-        let Err(ref err) = result else {
-            panic!("expected ParseError, got {result:#?}");
-        };
-        if !err.message().contains("inconsistent type") {
-            panic!("expected 'inconsistent type' error, got {result:#?}");
-        }
     }
 
     variant("min(1.0, 1i)");
@@ -4986,3 +4990,71 @@ fn enable_without_capability() {
         );
     }
 }
+
+#[test]
+fn bitwise_shift_errors() {
+    // 32-bit const by const >= bitwidth
+    check_error_matches(
+        "const N: u32 = 1u >> 32;",
+        "RHS of shift operation is greater than or equal to 32",
+    );
+    check_error_matches(
+        "const N: i32 = 1i >> 32;",
+        "RHS of shift operation is greater than or equal to 32",
+    );
+
+    // 32-bit const by const overflow
+    check_error_matches("const N: u32 = 0xFFFFFFFFu << 1;", "overflowed");
+    check_error_matches("const N: i32 = 1i << 31;", "overflowed");
+
+    // 32-bit const by const negative shift
+    check_error_matches("const N: u32 = 1u << -1;", "cannot represent");
+    check_error_matches("const N: i32 = 1i << -1;", "cannot represent");
+
+    // 32-bit runtime by const < bitwidth
+    check_success("fn foo() { var x: u32; var n = x << 31; }");
+    check_success("fn foo() { var x: i32; var n = x << 31; }");
+    check_success("fn foo() { var x: u32; var n = x >> 31; }");
+    check_success("fn foo() { var x: i32; var n = x >> 31; }");
+
+    // 32-bit runtime by const >= bitwidth
+    check_validation! {
+        "fn foo() { var x: u32; var n = x >> 32; }",
+        "fn foo() { var x: i32; var n = x >> 32; }",
+        "fn foo() { var x: u32; var n = x << 32; }",
+        "fn foo() { var x: i32; var n = x << 32; }":
+        Err(naga::valid::ValidationError::Function {
+            source: naga::valid::FunctionError::Expression {
+                source: naga::valid::ExpressionError::ShiftAmountTooLarge { .. },
+                ..
+            },
+            ..
+        })
+    }
+
+    // (CTS has more 32-bit test cases)
+
+    // Const evaluation of `i64` and `u64` is not implemented, https://github.com/gfx-rs/wgpu/issues/8972
+
+    // 64-bit runtime by const < bitwidth
+    check_success("fn foo() { var x: u64; var n = x << 63; }");
+    check_success("fn foo() { var x: i64; var n = x << 63; }");
+    check_success("fn foo() { var x: u64; var n = x >> 63; }");
+    check_success("fn foo() { var x: i64; var n = x >> 63; }");
+
+    // 64-bit runtime by const >= bitwidth
+    check_validation! {
+        "fn foo() { var x: u64; var n = x << 64; }",
+        "fn foo() { var x: i64; var n = x << 64; }",
+        "fn foo() { var x: u64; var n = x >> 64; }",
+        "fn foo() { var x: i64; var n = x >> 64; }":
+        Err(naga::valid::ValidationError::Function {
+            source: naga::valid::FunctionError::Expression {
+                source: naga::valid::ExpressionError::ShiftAmountTooLarge { .. },
+                ..
+            },
+            ..
+        }),
+        naga::valid::Capabilities::SHADER_INT64
+    }
+}