Auto merge of rust-lang#118310 - scottmcm:three-way-compare, r=<try>

Add `Ord::cmp` for primitives as a `BinOp` in MIR

There are dozens of reasonable ways to implement `Ord::cmp` for integers using comparison, bit-ops, and branches.  Those differences are irrelevant at the rust level, however, so we can make things better by adding `BinOp::Cmp` at the MIR level:

1. Exactly how to implement it is left up to the backends, so LLVM can use whatever pattern its optimizer best recognizes and cranelift can use whichever pattern codegens the fastest.
2. By not inlining those details for every use of `cmp`, we drastically reduce the amount of MIR generated for `derive`d `PartialOrd`, while also making it more amenable to MIR-level optimizations.

Having extremely careful `if` ordering to μoptimize resource usage on broadwell (rust-lang#63767) is great, but it really feels to me like libcore is the wrong place to put that logic.  Similarly, using subtraction [tricks](https://graphics.stanford.edu/~seander/bithacks.html#CopyIntegerSign) (rust-lang#105840) is arguably even nicer, but depends on the optimizer understanding it (llvm/llvm-project#73417) to be practical.  Or maybe [bitor is better than add](https://discourse.llvm.org/t/representing-in-ir/67369/2?u=scottmcm)?  But maybe only on a future version that [has `or disjoint` support](https://discourse.llvm.org/t/rfc-add-or-disjoint-flag/75036?u=scottmcm)?  And just because one of those forms happens to be good for LLVM, there's no guarantee that it'd be the same form that GCC or Cranelift would rather see -- especially given their very different optimizers.  Not to mention that if LLVM gets a spaceship intrinsic -- [which it should](https://rust-lang.zulipchat.com/#narrow/stream/131828-t-compiler/topic/Suboptimal.20inlining.20in.20std.20function.20.60binary_search.60/near/404250586) -- we'll need at least a rustc intrinsic to be able to call it.

As for simplifying it in Rust, we now regularly inline `{integer}::partial_cmp`, but it's quite a large amount of IR.  The best way to see that is with rust-lang@8811efa#diff-d134c32d028fbe2bf835fef2df9aca9d13332dd82284ff21ee7ebf717bfa4765R113 -- I added a new pre-codegen MIR test for a simple 3-tuple struct, and this PR change it from 36 locals and 26 basic blocks down to 24 locals and 8 basic blocks.  Even better, as soon as the construct-`Some`-then-match-it-in-same-BB noise is cleaned up, this'll expose the `Cmp == 0` branches clearly in MIR, so that an InstCombine (rust-lang#105808) can simplify that to just a `BinOp::Eq` and thus fix some of our generated code perf issues.  (Tracking that through today's `if a < b { Less } else if a == b { Equal } else { Greater }` would be *much* harder.)

---

r? `@ghost`
But first I should check that perf is ok with this
~~...and my true nemesis, tidy.~~

Author: bors

compiler/rustc_codegen_cranelift/src/codegen_i128.rs

@@ -68,7 +68,7 @@ pub(crate) fn maybe_codegen<'tcx>(
                 Some(CValue::by_val(ret_val, lhs.layout()))
             }
         }
-        BinOp::Lt | BinOp::Le | BinOp::Eq | BinOp::Ge | BinOp::Gt | BinOp::Ne => None,
+        BinOp::Lt | BinOp::Le | BinOp::Eq | BinOp::Ge | BinOp::Gt | BinOp::Ne | BinOp::Cmp => None,
         BinOp::Shl | BinOp::ShlUnchecked | BinOp::Shr | BinOp::ShrUnchecked => None,
     }
 }
@@ -134,6 +134,7 @@ pub(crate) fn maybe_codegen_checked<'tcx>(
         BinOp::AddUnchecked | BinOp::SubUnchecked | BinOp::MulUnchecked => unreachable!(),
         BinOp::Offset => unreachable!("offset should only be used on pointers, not 128bit ints"),
         BinOp::Div | BinOp::Rem => unreachable!(),
+        BinOp::Cmp => unreachable!(),
         BinOp::Lt | BinOp::Le | BinOp::Eq | BinOp::Ge | BinOp::Gt | BinOp::Ne => unreachable!(),
         BinOp::Shl | BinOp::ShlUnchecked | BinOp::Shr | BinOp::ShrUnchecked => unreachable!(),
     }

compiler/rustc_codegen_cranelift/src/num.rs

@@ -40,13 +40,34 @@ pub(crate) fn bin_op_to_intcc(bin_op: BinOp, signed: bool) -> Option<IntCC> {
     })
 }
 
+fn codegen_three_way_compare<'tcx>(
+    fx: &mut FunctionCx<'_, '_, 'tcx>,
+    signed: bool,
+    lhs: Value,
+    rhs: Value,
+) -> CValue<'tcx> {
+    let gt_cc = crate::num::bin_op_to_intcc(BinOp::Gt, signed).unwrap();
+    let lt_cc = crate::num::bin_op_to_intcc(BinOp::Lt, signed).unwrap();
+    let gt = fx.bcx.ins().icmp(gt_cc, lhs, rhs);
+    let lt = fx.bcx.ins().icmp(lt_cc, lhs, rhs);
+    // Cranelift no longer has a single-bit type, so the comparison results
+    // are already `I8`s, which we can subtract to get the -1/0/+1 we want.
+    // See <https://github.com/bytecodealliance/wasmtime/pull/5031>.
+    let val = fx.bcx.ins().isub(gt, lt);
+    CValue::by_val(val, fx.layout_of(fx.tcx.ty_ordering_enum(Some(fx.mir.span))))
+}
+
 fn codegen_compare_bin_op<'tcx>(
     fx: &mut FunctionCx<'_, '_, 'tcx>,
     bin_op: BinOp,
     signed: bool,
     lhs: Value,
     rhs: Value,
 ) -> CValue<'tcx> {
+    if bin_op == BinOp::Cmp {
+        return codegen_three_way_compare(fx, signed, lhs, rhs);
+    }
+
     let intcc = crate::num::bin_op_to_intcc(bin_op, signed).unwrap();
     let val = fx.bcx.ins().icmp(intcc, lhs, rhs);
     CValue::by_val(val, fx.layout_of(fx.tcx.types.bool))
@@ -59,7 +80,7 @@ pub(crate) fn codegen_binop<'tcx>(
     in_rhs: CValue<'tcx>,
 ) -> CValue<'tcx> {
     match bin_op {
-        BinOp::Eq | BinOp::Lt | BinOp::Le | BinOp::Ne | BinOp::Ge | BinOp::Gt => {
+        BinOp::Eq | BinOp::Lt | BinOp::Le | BinOp::Ne | BinOp::Ge | BinOp::Gt | BinOp::Cmp => {
             match in_lhs.layout().ty.kind() {
                 ty::Bool | ty::Uint(_) | ty::Int(_) | ty::Char => {
                     let signed = type_sign(in_lhs.layout().ty);
@@ -160,7 +181,7 @@ pub(crate) fn codegen_int_binop<'tcx>(
         }
         BinOp::Offset => unreachable!("Offset is not an integer operation"),
         // Compare binops handles by `codegen_binop`.
-        BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge => {
+        BinOp::Eq | BinOp::Ne | BinOp::Lt | BinOp::Le | BinOp::Gt | BinOp::Ge | BinOp::Cmp => {
             unreachable!("{:?}({:?}, {:?})", bin_op, in_lhs.layout().ty, in_rhs.layout().ty);
         }
     };

compiler/rustc_codegen_gcc/src/common.rs

@@ -104,6 +104,10 @@ impl<'gcc, 'tcx> ConstMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         self.const_int(self.type_i32(), i as i64)
     }
 
+    fn const_i8(&self, i: i8) -> RValue<'gcc> {
+        self.const_int(self.type_i8(), i as i64)
+    }
+
     fn const_u32(&self, i: u32) -> RValue<'gcc> {
         self.const_uint(self.type_u32(), i as u64)
     }

compiler/rustc_codegen_llvm/src/common.rs

@@ -158,6 +158,10 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         self.const_int(self.type_i32(), i as i64)
     }
 
+    fn const_i8(&self, i: i8) -> &'ll Value {
+        self.const_int(self.type_i8(), i as i64)
+    }
+
     fn const_u32(&self, i: u32) -> &'ll Value {
         self.const_uint(self.type_i32(), i as u64)
     }

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -7,6 +7,7 @@ use crate::common::{self, IntPredicate};
 use crate::traits::*;
 use crate::MemFlags;
 
+use rustc_hir as hir;
 use rustc_middle::mir;
 use rustc_middle::mir::Operand;
 use rustc_middle::ty::cast::{CastTy, IntTy};
@@ -881,6 +882,27 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     bx.icmp(base::bin_op_to_icmp_predicate(op.to_hir_binop(), is_signed), lhs, rhs)
                 }
             }
+            mir::BinOp::Cmp => {
+                use std::cmp::Ordering;
+                debug_assert!(!is_float);
+                // FIXME: To avoid this PR changing behaviour, the operations used
+                // here are those from <https://github.com/rust-lang/rust/pull/63767>,
+                // as tested by `tests/codegen/integer-cmp.rs`.
+                // Something in future might want to pick different ones. For example,
+                // maybe the ones from Clang's `<=>` operator in C++20 (see
+                // <https://github.com/llvm/llvm-project/issues/60012>) or once we
+                // update to new LLVM, something to take advantage of the new folds in
+                // <https://github.com/llvm/llvm-project/issues/59666>.
+                let pred = |op| base::bin_op_to_icmp_predicate(op, is_signed);
+                let is_lt = bx.icmp(pred(hir::BinOpKind::Lt), lhs, rhs);
+                let is_ne = bx.icmp(pred(hir::BinOpKind::Ne), lhs, rhs);
+                let ge = bx.select(
+                    is_ne,
+                    bx.cx().const_i8(Ordering::Greater as i8),
+                    bx.cx().const_i8(Ordering::Equal as i8),
+                );
+                bx.select(is_lt, bx.cx().const_i8(Ordering::Less as i8), ge)
+            }
         }
     }
 

compiler/rustc_codegen_ssa/src/traits/consts.rs

@@ -19,6 +19,7 @@ pub trait ConstMethods<'tcx>: BackendTypes {
     fn const_bool(&self, val: bool) -> Self::Value;
     fn const_i16(&self, i: i16) -> Self::Value;
     fn const_i32(&self, i: i32) -> Self::Value;
+    fn const_i8(&self, i: i8) -> Self::Value;
     fn const_u32(&self, i: u32) -> Self::Value;
     fn const_u64(&self, i: u64) -> Self::Value;
     fn const_u128(&self, i: u128) -> Self::Value;

compiler/rustc_const_eval/src/interpret/operand.rs

@@ -224,6 +224,13 @@ impl<'tcx, Prov: Provenance> ImmTy<'tcx, Prov> {
         Self::from_scalar(Scalar::from_bool(b), layout)
     }
 
+    #[inline]
+    pub fn from_ordering(c: std::cmp::Ordering, tcx: TyCtxt<'tcx>) -> Self {
+        let ty = tcx.ty_ordering_enum(None);
+        let layout = tcx.layout_of(ty::ParamEnv::reveal_all().and(ty)).unwrap();
+        Self::from_scalar(Scalar::from_i8(c as i8), layout)
+    }
+
     #[inline]
     pub fn to_const_int(self) -> ConstInt {
         assert!(self.layout.ty.is_integral());

compiler/rustc_const_eval/src/interpret/operator.rs

@@ -61,6 +61,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 }
 
 impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
+    fn three_way_compare<T: Ord>(&self, lhs: T, rhs: T) -> (ImmTy<'tcx, M::Provenance>, bool) {
+        let res = Ord::cmp(&lhs, &rhs);
+        return (ImmTy::from_ordering(res, *self.tcx), false);
+    }
+
     fn binary_char_op(
         &self,
         bin_op: mir::BinOp,
@@ -69,6 +74,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ) -> (ImmTy<'tcx, M::Provenance>, bool) {
         use rustc_middle::mir::BinOp::*;
 
+        if bin_op == Cmp {
+            return self.three_way_compare(l, r);
+        }
+
         let res = match bin_op {
             Eq => l == r,
             Ne => l != r,
@@ -231,6 +240,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 let r = self.sign_extend(r, right_layout) as i128;
                 return Ok((ImmTy::from_bool(op(&l, &r), *self.tcx), false));
             }
+            if bin_op == Cmp {
+                let l = self.sign_extend(l, left_layout) as i128;
+                let r = self.sign_extend(r, right_layout) as i128;
+                return Ok(self.three_way_compare(l, r));
+            }
             let op: Option<fn(i128, i128) -> (i128, bool)> = match bin_op {
                 Div if r == 0 => throw_ub!(DivisionByZero),
                 Rem if r == 0 => throw_ub!(RemainderByZero),
@@ -270,6 +284,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             }
         }
 
+        if bin_op == Cmp {
+            return Ok(self.three_way_compare(l, r));
+        }
+
         let val = match bin_op {
             Eq => ImmTy::from_bool(l == r, *self.tcx),
             Ne => ImmTy::from_bool(l != r, *self.tcx),

compiler/rustc_const_eval/src/transform/promote_consts.rs

@@ -573,6 +573,7 @@ impl<'tcx> Validator<'_, 'tcx> {
                     | BinOp::Lt
                     | BinOp::Ge
                     | BinOp::Gt
+                    | BinOp::Cmp
                     | BinOp::Offset
                     | BinOp::Add
                     | BinOp::AddUnchecked

compiler/rustc_const_eval/src/transform/validate.rs

@@ -932,6 +932,15 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             )
                         }
                     }
+                    Cmp => {
+                        for x in [a, b] {
+                            check_kinds!(
+                                x,
+                                "Cannot three-way compare non-integer type {:?}",
+                                ty::Char | ty::Uint(..) | ty::Int(..)
+                            )
+                        }
+                    }
                     AddUnchecked | SubUnchecked | MulUnchecked | Shl | ShlUnchecked | Shr
                     | ShrUnchecked => {
                         for x in [a, b] {