Auto merge of #135674 - scottmcm:assume-better, r=estebank

Update our range `assume`s to the format that LLVM prefers

I found out in https://github.com/llvm/llvm-project/issues/123278#issuecomment-2597440158 that the way I started emitting the `assume`s in #109993 was suboptimal, and as seen in that LLVM issue the way we're doing it -- with two `assume`s sometimes -- can at times lead to CVP/SCCP not realize what's happening because one of them turns into a `ne` instead of conveying a range.

So this updates how it's emitted from
```
assume( x >= LOW );
assume( x <= HIGH );
```
or
```
// (for ranges that wrap the range)
assume( (x <= LOW) | (x >= HIGH) );
```
to
```
assume( (x - LOW) <= (HIGH - LOW) );
```
so that we don't need multiple `icmp`s nor multiple `assume`s for a single value, and both wrappping and non-wrapping ranges emit the same shape.

(And we don't bother emitting the subtraction if `LOW` is zero, since that's trivial for us to check too.)

2 files changed, 39 insertions, 29 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index cb4862663fd..e775d219c7b 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -387,10 +387,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         use abi::Primitive::*;
         imm = bx.from_immediate(imm);
 
-        // When scalars are passed by value, there's no metadata recording their
-        // valid ranges. For example, `char`s are passed as just `i32`, with no
-        // way for LLVM to know that they're 0x10FFFF at most. Thus we assume
-        // the range of the input value too, not just the output range.
+        // If we have a scalar, we must already know its range. Either
+        //
+        // 1) It's a parameter with `range` parameter metadata,
+        // 2) It's something we `load`ed with `!range` metadata, or
+        // 3) After a transmute we `assume`d the range (see below).
+        //
+        // That said, last time we tried removing this, it didn't actually help
+        // the rustc-perf results, so might as well keep doing it
+        // <https://github.com/rust-lang/rust/pull/135610#issuecomment-2599275182>
         self.assume_scalar_range(bx, imm, from_scalar, from_backend_ty);
 
         imm = match (from_scalar.primitive(), to_scalar.primitive()) {
@@ -411,7 +416,14 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 bx.bitcast(int_imm, to_backend_ty)
             }
         };
+
+        // This `assume` remains important for cases like (a conceptual)
+        //    transmute::<u32, NonZeroU32>(x) == 0
+        // since it's never passed to something with parameter metadata (especially
+        // after MIR inlining) so the only way to tell the backend about the
+        // constraint that the `transmute` introduced is to `assume` it.
         self.assume_scalar_range(bx, imm, to_scalar, to_backend_ty);
+
         imm = bx.to_immediate_scalar(imm, to_scalar);
         imm
     }
@@ -433,31 +445,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             return;
         }
 
-        let abi::WrappingRange { start, end } = scalar.valid_range(self.cx);
-
-        if start <= end {
-            if start > 0 {
-                let low = bx.const_uint_big(backend_ty, start);
-                let cmp = bx.icmp(IntPredicate::IntUGE, imm, low);
-                bx.assume(cmp);
-            }
-
-            let type_max = scalar.size(self.cx).unsigned_int_max();
-            if end < type_max {
-                let high = bx.const_uint_big(backend_ty, end);
-                let cmp = bx.icmp(IntPredicate::IntULE, imm, high);
-                bx.assume(cmp);
-            }
-        } else {
-            let low = bx.const_uint_big(backend_ty, start);
-            let cmp_low = bx.icmp(IntPredicate::IntUGE, imm, low);
-
-            let high = bx.const_uint_big(backend_ty, end);
-            let cmp_high = bx.icmp(IntPredicate::IntULE, imm, high);
-
-            let or = bx.or(cmp_low, cmp_high);
-            bx.assume(or);
-        }
+        let range = scalar.valid_range(self.cx);
+        bx.assume_integer_range(imm, backend_ty, range);
     }
 
     pub(crate) fn codegen_rvalue_unsized(
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
index 3ee13b19f66..bbf87a59942 100644
--- a/compiler/rustc_codegen_ssa/src/traits/builder.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -217,6 +217,27 @@ pub trait BuilderMethods<'a, 'tcx>:
         dest: PlaceRef<'tcx, Self::Value>,
     );
 
+    /// Emits an `assume` that the integer value `imm` of type `ty` is contained in `range`.
+    ///
+    /// This *always* emits the assumption, so you probably want to check the
+    /// optimization level and `Scalar::is_always_valid` before calling it.
+    fn assume_integer_range(&mut self, imm: Self::Value, ty: Self::Type, range: WrappingRange) {
+        let WrappingRange { start, end } = range;
+
+        // Perhaps one day we'll be able to use assume operand bundles for this,
+        // but for now this encoding with a single icmp+assume is best per
+        // <https://github.com/llvm/llvm-project/issues/123278#issuecomment-2597440158>
+        let shifted = if start == 0 {
+            imm
+        } else {
+            let low = self.const_uint_big(ty, start);
+            self.sub(imm, low)
+        };
+        let width = self.const_uint_big(ty, u128::wrapping_sub(end, start));
+        let cmp = self.icmp(IntPredicate::IntULE, shifted, width);
+        self.assume(cmp);
+    }
+
     fn range_metadata(&mut self, load: Self::Value, range: WrappingRange);
     fn nonnull_metadata(&mut self, load: Self::Value);