Fix/implement checked_{add,sub,mul} and saturating_{add,sub} (using SPIR-V instructions were possible).

crates/rustc_codegen_spirv/src/builder/builder_methods.rs

@@ -27,7 +27,7 @@ use rustc_codegen_ssa::traits::{
 };
 use rustc_middle::bug;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
-use rustc_middle::ty::layout::TyAndLayout;
+use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, AtomicOrdering, Ty};
 use rustc_span::Span;
 use rustc_target::callconv::FnAbi;
@@ -1744,30 +1744,15 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
     fn checked_binop(
         &mut self,
         oop: OverflowOp,
-        ty: Ty<'_>,
+        ty: Ty<'tcx>,
         lhs: Self::Value,
         rhs: Self::Value,
     ) -> (Self::Value, Self::Value) {
-        // adopted partially from https://github.com/ziglang/zig/blob/master/src/codegen/spirv.zig
-        let is_add = match oop {
-            OverflowOp::Add => true,
-            OverflowOp::Sub => false,
-            OverflowOp::Mul => {
-                // NOTE(eddyb) this needs to be `undef`, not `false`/`true`, because
-                // we don't want the user's boolean constants to keep the zombie alive.
-                let bool = SpirvType::Bool.def(self.span(), self);
-                let overflowed = self.undef(bool);
-
-                let result = (self.mul(lhs, rhs), overflowed);
-                self.zombie(result.1.def(self), "checked mul is not supported yet");
-                return result;
-            }
-        };
         let signed = match ty.kind() {
             ty::Int(_) => true,
             ty::Uint(_) => false,
-            other => self.fatal(format!(
-                "Unexpected {} type: {other:#?}",
+            _ => self.fatal(format!(
+                "unexpected {} type: {ty}",
                 match oop {
                     OverflowOp::Add => "checked add",
                     OverflowOp::Sub => "checked sub",
@@ -1776,48 +1761,108 @@ impl<'a, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'tcx> {
             )),
         };
 
-        let result = if is_add {
-            self.add(lhs, rhs)
-        } else {
-            self.sub(lhs, rhs)
-        };
+        // HACK(eddyb) SPIR-V `OpIAddCarry`/`OpISubBorrow` are specifically for
+        // unsigned overflow, so signed overflow still needs this custom logic.
+        if signed && let OverflowOp::Add | OverflowOp::Sub = oop {
+            let result = match oop {
+                OverflowOp::Add => self.add(lhs, rhs),
+                OverflowOp::Sub => self.sub(lhs, rhs),
+                OverflowOp::Mul => unreachable!(),
+            };
+
+            // adopted partially from https://github.com/ziglang/zig/blob/master/src/codegen/spirv.zig
+            // FIXME(eddyb) ^^ find a working a permalink, and also inform them
+            // of the difference between these two approaches:
+            // - broken: https://alive2.llvm.org/ce/z/Q3Pchi
+            // - correct: https://alive2.llvm.org/ce/z/aWvThi
 
-        let overflowed = if signed {
             // when adding, overflow could happen if
             // - rhs is positive and result < lhs; or
-            // - rhs is negative and result > lhs
-            // this is equivalent to (rhs < 0) == (result > lhs)
+            // - rhs is negative and result >= lhs
+            //   (`result == lhs` impossible, `>=` used as it's `!(result < lhs)`)
+            // this is equivalent to (rhs < 0) == (result >= lhs)
             //
             // when subtracting, overflow happens if
             // - rhs is positive and result > lhs; or
-            // - rhs is negative and result < lhs
-            // this is equivalent to (rhs < 0) == (result < lhs)
+            // - rhs is negative and result <= lhs
+            //   (`result == lhs` impossible, `<=` used as it's `!(result > lhs)`)
+            // this is equivalent to (rhs < 0) == (result <= lhs)
             let rhs_lt_zero = self.icmp(IntPredicate::IntSLT, rhs, self.constant_int(rhs.ty, 0));
-            let result_gt_lhs = self.icmp(
-                if is_add {
-                    IntPredicate::IntSGT
-                } else {
-                    IntPredicate::IntSLT
+            let result_ge_lhs = self.icmp(
+                match oop {
+                    OverflowOp::Add => IntPredicate::IntSGE,
+                    OverflowOp::Sub => IntPredicate::IntSLE,
+                    OverflowOp::Mul => unreachable!(),
                 },
                 result,
                 lhs,
             );
-            self.icmp(IntPredicate::IntEQ, rhs_lt_zero, result_gt_lhs)
-        } else {
-            // for unsigned addition, overflow occurred if the result is less than any of the operands.
-            // for subtraction, overflow occurred if the result is greater.
-            self.icmp(
-                if is_add {
-                    IntPredicate::IntULT
+
+            let overflowed = self.icmp(IntPredicate::IntEQ, rhs_lt_zero, result_ge_lhs);
+
+            return (result, overflowed);
+        }
+
+        let result_type = self.layout_of(ty).spirv_type(self.span(), self);
+        let pair_result_type = {
+            let field_types = [result_type, result_type];
+            let (field_offsets, size, align) = crate::abi::auto_struct_layout(self, &field_types);
+            SpirvType::Adt {
+                def_id: None,
+                size,
+                align,
+                field_types: &field_types,
+                field_offsets: &field_offsets,
+                field_names: None,
+            }
+            .def(self.span(), self)
+        };
+
+        let lhs = lhs.def(self);
+        let rhs = rhs.def(self);
+        let pair_result = match oop {
+            OverflowOp::Add => self
+                .emit()
+                .i_add_carry(pair_result_type, None, lhs, rhs)
+                .unwrap(),
+            OverflowOp::Sub => self
+                .emit()
+                .i_sub_borrow(pair_result_type, None, lhs, rhs)
+                .unwrap(),
+            OverflowOp::Mul => {
+                if signed {
+                    self.emit()
+                        .s_mul_extended(pair_result_type, None, lhs, rhs)
+                        .unwrap()
                 } else {
-                    IntPredicate::IntUGT
-                },
-                result,
-                lhs,
-            )
+                    self.emit()
+                        .u_mul_extended(pair_result_type, None, lhs, rhs)
+                        .unwrap()
+                }
+            }
+        }
+        .with_type(pair_result_type);
+        let result_lo = self.extract_value(pair_result, 0);
+        let result_hi = self.extract_value(pair_result, 1);
+
+        // HACK(eddyb) SPIR-V lacks any `(T, T) -> (T, bool)` instructions,
+        // so instead `result_hi` is compared with the value expected in the
+        // non-overflow case (`0`, or `-1` for negative signed multiply result).
+        let expected_nonoverflowing_hi = match (oop, signed) {
+            (OverflowOp::Add | OverflowOp::Sub, _) | (OverflowOp::Mul, false) => {
+                self.const_uint(result_type, 0)
+            }
+            (OverflowOp::Mul, true) => {
+                // HACK(eddyb) `(x: iN) >> (N - 1)` will spread the sign bit
+                // across all `N` bits of `iN`, and should be equivalent to
+                // `if x < 0 { -1 } else { 0 }`, without needing compare+select).
+                let result_width = u32::try_from(self.int_width(result_type)).unwrap();
+                self.ashr(result_lo, self.const_u32(result_width - 1))
+            }
         };
+        let overflowed = self.icmp(IntPredicate::IntNE, result_hi, expected_nonoverflowing_hi);
 
-        (result, overflowed)
+        (result_lo, overflowed)
     }
 
     // rustc has the concept of an immediate vs. memory type - bools are compiled to LLVM bools as

crates/rustc_codegen_spirv/src/builder/intrinsics.rs

@@ -119,33 +119,39 @@ impl<'a, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'a, 'tcx> {
 
             sym::saturating_add => {
                 assert_eq!(arg_tys[0], arg_tys[1]);
-                let result = match arg_tys[0].kind() {
-                    TyKind::Int(_) | TyKind::Uint(_) => {
-                        self.add(args[0].immediate(), args[1].immediate())
-                    }
-                    TyKind::Float(_) => self.fadd(args[0].immediate(), args[1].immediate()),
+                match arg_tys[0].kind() {
+                    TyKind::Int(_) | TyKind::Uint(_) => self
+                        .emit()
+                        .i_add_sat_intel(
+                            ret_ty,
+                            None,
+                            args[0].immediate().def(self),
+                            args[1].immediate().def(self),
+                        )
+                        .unwrap()
+                        .with_type(ret_ty),
                     other => self.fatal(format!(
-                        "Unimplemented saturating_add intrinsic type: {other:#?}"
+                        "unimplemented saturating_add intrinsic type: {other:#?}"
                     )),
-                };
-                // TODO: Implement this
-                self.zombie(result.def(self), "saturating_add is not implemented yet");
-                result
+                }
             }
             sym::saturating_sub => {
                 assert_eq!(arg_tys[0], arg_tys[1]);
-                let result = match &arg_tys[0].kind() {
-                    TyKind::Int(_) | TyKind::Uint(_) => {
-                        self.sub(args[0].immediate(), args[1].immediate())
-                    }
-                    TyKind::Float(_) => self.fsub(args[0].immediate(), args[1].immediate()),
+                match &arg_tys[0].kind() {
+                    TyKind::Int(_) | TyKind::Uint(_) => self
+                        .emit()
+                        .i_sub_sat_intel(
+                            ret_ty,
+                            None,
+                            args[0].immediate().def(self),
+                            args[1].immediate().def(self),
+                        )
+                        .unwrap()
+                        .with_type(ret_ty),
                     other => self.fatal(format!(
-                        "Unimplemented saturating_sub intrinsic type: {other:#?}"
+                        "unimplemented saturating_sub intrinsic type: {other:#?}"
                     )),
-                };
-                // TODO: Implement this
-                self.zombie(result.def(self), "saturating_sub is not implemented yet");
-                result
+                }
             }
 
             sym::sqrtf32 | sym::sqrtf64 | sym::sqrtf128 => {

tests/difftests/tests/Cargo.toml

@@ -23,6 +23,8 @@ members = [
     "lang/abi/vector_layout_cuda/rust-gpu",
     "lang/abi/vector_layout_scalar_math/cpu",
     "lang/abi/vector_layout_scalar_math/rust-gpu",
+    "lang/math/checked_saturating/cpu",
+    "lang/math/checked_saturating/rust-gpu",
     "lang/control_flow/control_flow-rust",
     "lang/control_flow/control_flow-wgsl",
     "lang/control_flow_complex/control_flow_complex-rust",

tests/difftests/tests/lang/math/checked_saturating/cpu/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "math-checked-saturating-cpu"
+edition.workspace = true
+
+[lints]
+workspace = true
+
+# GPU deps
+[dependencies]
+spirv-std.workspace = true
+glam.workspace = true
+bytemuck.workspace = true
+
+# CPU deps (for the test harness)
+[target.'cfg(not(target_arch = "spirv"))'.dependencies]
+difftest.workspace = true

tests/difftests/tests/lang/math/checked_saturating/cpu/src/lib.rs

@@ -0,0 +1,103 @@
+#![cfg_attr(target_arch = "spirv", no_std)]
+
+use spirv_std::spirv;
+
+pub trait ToU32 {
+    fn to_u32(&self) -> u32;
+}
+
+impl ToU32 for u32 {
+    fn to_u32(&self) -> u32 {
+        *self
+    }
+}
+
+impl ToU32 for i32 {
+    fn to_u32(&self) -> u32 {
+        *self as u32
+    }
+}
+
+#[inline]
+pub fn ch<T: ToU32>(opt: Option<T>) -> [u32; 2] {
+    match opt {
+        None => [u32::default(), 1],
+        Some(e) => [e.to_u32(), 0],
+    }
+}
+
+#[inline]
+pub fn ov<T: ToU32>((value, overflow): (T, bool)) -> [u32; 2] {
+    [value.to_u32(), overflow.into()]
+}
+
+pub const LEN: usize = 24;
+pub type Output = [[u32; 2]; LEN];
+
+pub fn eval() -> Output {
+    [
+        // u32
+        ov(1u32.overflowing_add(1)),
+        ch(1u32.checked_add(1)),
+        ov(u32::MAX.overflowing_add(1)),
+        ch(u32::MAX.checked_add(1)),
+        ch(0u32.checked_add(0)),
+        ov(0u32.overflowing_add(0)),
+        ov(1u32.overflowing_sub(1)),
+        ch(1u32.checked_sub(1)),
+        ov(u32::MIN.overflowing_sub(1)),
+        ch(u32::MIN.checked_sub(1)),
+        ch(0u32.checked_sub(0)),
+        ov(0u32.overflowing_sub(0)),
+        // i32
+        ov(1i32.overflowing_add(1)),
+        ch(1i32.checked_add(1)),
+        ov(i32::MAX.overflowing_add(1)),
+        ch(i32::MAX.checked_add(1)),
+        ch(0i32.checked_add(0)),
+        ov(0i32.overflowing_add(0)),
+        ov(1i32.overflowing_sub(1)),
+        ch(1i32.checked_sub(1)),
+        ov(i32::MIN.overflowing_sub(1)),
+        ch(i32::MIN.checked_sub(1)),
+        ch(0i32.checked_sub(0)),
+        ov(0i32.overflowing_sub(0)),
+    ]
+}
+
+#[spirv(compute(threads(1)))]
+pub fn main_cs(#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] out: &mut Output) {
+    *out = eval();
+}
+
+#[cfg(not(target_arch = "spirv"))]
+mod cpu {
+    use super::*;
+    use difftest::config::Config;
+    use difftest::scaffold::compute::{BufferConfig, WgpuComputeTest};
+    use difftest::scaffold::shader::RustComputeShader;
+
+    pub fn cpu_driver() {
+        let config = Config::from_path(std::env::args().nth(1).unwrap()).unwrap();
+        config.write_result(&eval()).unwrap()
+    }
+
+    pub fn shader_driver() {
+        let config = Config::from_path(std::env::args().nth(1).unwrap()).unwrap();
+        if cfg!(target_os = "macos") {
+            let skip = difftest::scaffold::Skip::new("no Vulkan on MacOS");
+            skip.run_test(&config).unwrap();
+            return;
+        }
+
+        let test = WgpuComputeTest::new(
+            RustComputeShader::default().passthrough(),
+            [1, 1, 1],
+            Vec::from(&[BufferConfig::writeback(size_of::<Output>())]),
+        );
+        test.run_test(&config).unwrap();
+    }
+}
+
+#[cfg(not(target_arch = "spirv"))]
+pub use cpu::*;

tests/difftests/tests/lang/math/checked_saturating/cpu/src/main.rs

@@ -0,0 +1,3 @@
+fn main() {
+    math_checked_saturating_cpu::cpu_driver();
+}