feat(x64): fuse *_overflow + brif to jo/jb

cranelift/codegen/src/isa/x64/lower.isle

@@ -193,6 +193,84 @@
                    (result MultiReg (with_flags_chained lo_inst hi_inst of_inst)))
                   (multi_reg_to_pair_and_single result)))))
 
+;; Match a value that is the overflow/carry flag (result 1) of a fusable
+;; `u/sadd/sub_overflow` instruction. Binds the arithmetic result, operands,
+;; condition code (`CC.O` or `CC.B`), ALU op, and defining instruction.
+;;
+;; Fails when fusion is unsafe (flag already materialized, multiple uses of the
+;; flag, intervening same-block use of the sum, i128, already absorbed, etc.).
+(decl overflow_alu_flag (Value Value Value CC ProduceFlagsOp Inst) Value)
+(extern extractor overflow_alu_flag overflow_alu_flag)
+
+;; Like `overflow_alu_flag`, but for `u/smul_overflow`. `signed` selects imul vs mul.
+(decl overflow_mul_flag (Value Value Value bool Inst) Value)
+(extern extractor overflow_mul_flag overflow_mul_flag)
+
+;; Alias the arithmetic result of an absorbed `*_overflow` and skip lowering it.
+(decl absorb_overflow_inst (Inst ValueRegs) Unit)
+(extern constructor absorb_overflow_inst absorb_overflow_inst)
+
+;; Emit `alu` + `jcc` for a fused overflow+branch, defining `sum` and absorbing
+;; the overflow instruction.
+(decl fuse_overflow_alu_brif (Type Value Value Value CC ProduceFlagsOp Inst MachLabel MachLabel) Unit)
+(rule (fuse_overflow_alu_brif ty sum x y cc op inst then else)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_produce_flags op ty x y)
+              (jmp_cond cc then else)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (unit)))
+
+(decl fuse_overflow_mul_brif (Type Value Value Value bool Inst MachLabel MachLabel) Unit)
+(rule 1 (fuse_overflow_mul_brif $I8 sum x y signed inst then else)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_mul8_with_flags_paired signed x y)
+              (jmp_cond (CC.O) then else)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (unit)))
+(rule 0 (fuse_overflow_mul_brif (ty_int_ref_16_to_64 ty) sum x y signed inst then else)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_mul_lo_with_flags_paired ty signed x y)
+              (jmp_cond (CC.O) then else)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (unit)))
+
+;; Emit `alu` + `trap_if` for a fused overflow+trapnz.
+(decl fuse_overflow_alu_trap (Type Value Value Value CC ProduceFlagsOp Inst TrapCode) InstOutput)
+(rule (fuse_overflow_alu_trap ty sum x y cc op inst code)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_produce_flags op ty x y)
+              (trap_if cc code)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (output_none)))
+
+(decl fuse_overflow_mul_trap (Type Value Value Value bool Inst TrapCode) InstOutput)
+(rule 1 (fuse_overflow_mul_trap $I8 sum x y signed inst code)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_mul8_with_flags_paired signed x y)
+              (trap_if (CC.O) code)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (output_none)))
+(rule 0 (fuse_overflow_mul_trap (ty_int_ref_16_to_64 ty) sum x y signed inst code)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_mul_lo_with_flags_paired ty signed x y)
+              (trap_if (CC.O) code)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (output_none)))
+
+(decl fuse_overflow_mul_trap_inv (Type Value Value Value bool Inst TrapCode) InstOutput)
+(rule 1 (fuse_overflow_mul_trap_inv $I8 sum x y signed inst code)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_mul8_with_flags_paired signed x y)
+              (trap_if (CC.NO) code)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (output_none)))
+(rule 0 (fuse_overflow_mul_trap_inv (ty_int_ref_16_to_64 ty) sum x y signed inst code)
+      (let ((sum_regs ValueRegs (with_flags
+              (x64_mul_lo_with_flags_paired ty signed x y)
+              (trap_if (CC.NO) code)))
+            (_ Unit (absorb_overflow_inst inst sum_regs)))
+        (output_none)))
+
 ;;;; Rules for `uadd_overflow` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (rule 1 (lower (uadd_overflow _ x y @ (value_type (fits_in_64 ty))))
@@ -1897,6 +1975,16 @@
 
 ;;;; Rules for `trapz` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
+;; Fuse `*_overflow` flag + `trapz` into alu + inverted `trap_if`.
+(rule 2 (lower (trapz of code))
+      (if-let (overflow_alu_flag sum x y cc op inst) of)
+      (if-let (value_type ty) sum)
+      (fuse_overflow_alu_trap ty sum x y (cc_invert cc) op inst code))
+(rule 1 (lower (trapz of code))
+      (if-let (overflow_mul_flag sum x y signed inst) of)
+      (if-let (value_type ty) sum)
+      (fuse_overflow_mul_trap_inv ty sum x y signed inst code))
+
 (rule (lower (trapz val code))
   (side_effect (trap_if_cond (cond_invert (is_nonzero_cmp val)) code)))
 
@@ -1910,6 +1998,17 @@
 
 ;;;; Rules for `trapnz` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
+;; Fuse `*_overflow` flag + `trapnz` into alu + `trap_if` (jo/jb), absorbing the
+;; overflow instruction so it is not also lowered to `setcc`.
+(rule 2 (lower (trapnz of code))
+      (if-let (overflow_alu_flag sum x y cc op inst) of)
+      (if-let (value_type ty) sum)
+      (fuse_overflow_alu_trap ty sum x y cc op inst code))
+(rule 1 (lower (trapnz of code))
+      (if-let (overflow_mul_flag sum x y signed inst) of)
+      (if-let (value_type ty) sum)
+      (fuse_overflow_mul_trap ty sum x y signed inst code))
+
 (rule (lower (trapnz val code))
   (side_effect (trap_if_cond (is_nonzero_cmp val) code)))
 
@@ -3600,6 +3699,17 @@
 
 ;; Rules for `brif` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
+;; Fuse `*_overflow` + `brif` into alu + `jo`/`jb`, defining the arithmetic
+;; result and absorbing the overflow instruction.
+(rule 2 (lower_branch (brif of _ _) (two_targets then else))
+      (if-let (overflow_alu_flag sum x y cc op inst) of)
+      (if-let (value_type ty) sum)
+      (fuse_overflow_alu_brif ty sum x y cc op inst then else))
+(rule 1 (lower_branch (brif of _ _) (two_targets then else))
+      (if-let (overflow_mul_flag sum x y signed inst) of)
+      (if-let (value_type ty) sum)
+      (fuse_overflow_mul_brif ty sum x y signed inst then else))
+
 (rule (lower_branch (brif val _ _) (two_targets then else))
       (emit_side_effect (jmp_cond_result (is_nonzero_cmp val) then else)))
 

cranelift/codegen/src/isa/x64/lower/isle.rs

@@ -3,21 +3,23 @@
 // Pull in the ISLE generated code.
 pub(crate) mod generated_code;
 use crate::{ir::AtomicRmwOp, ir::types};
-use generated_code::{AssemblerOutputs, Context, MInst, RegisterClass};
+use generated_code::{AssemblerOutputs, Context, MInst, ProduceFlagsOp, RegisterClass};
 
 // Types that the generated ISLE code uses via `use super::*`.
 use super::external::{CraneliftRegisters, PairedGpr, PairedXmm, isle_assembler_methods};
 use super::{MergeableLoadSize, is_int_or_ref_ty, is_mergeable_load, lower_to_amode};
 use crate::ir::condcodes::{FloatCC, IntCC};
 use crate::ir::immediates::*;
 use crate::ir::types::*;
-use crate::ir::{BlockCall, Inst, InstructionData, LibCall, Opcode, TrapCode, Value, ValueList};
+use crate::ir::{
+    BlockCall, Inst, InstructionData, LibCall, Opcode, TrapCode, Value, ValueDef, ValueList,
+};
 use crate::isa::x64::X64Backend;
 use crate::isa::x64::inst::{ReturnCallInfo, args::*, regs};
 use crate::isa::x64::lower::{InsnInput, emit_vm_call};
 use crate::machinst::isle::*;
 use crate::machinst::{
-    ArgPair, CallArgList, CallInfo, CallRetList, InstOutput, MachInst, VCodeConstant,
+    ArgPair, CallArgList, CallInfo, CallRetList, InstOutput, Lower, MachInst, VCodeConstant,
     VCodeConstantData,
 };
 use alloc::boxed::Box;
@@ -1926,6 +1928,36 @@ impl Context for IsleContext<'_, '_, MInst, X64Backend> {
             None
         }
     }
+
+    fn overflow_alu_flag(
+        &mut self,
+        of: Value,
+    ) -> Option<(Value, Value, Value, CC, ProduceFlagsOp, Inst)> {
+        let (inst, sum, x, y, opcode) = overflow_flag_common(self.lower_ctx, of)?;
+        let (cc, op) = match opcode {
+            Opcode::UaddOverflow => (CC::B, ProduceFlagsOp::Add),
+            Opcode::SaddOverflow => (CC::O, ProduceFlagsOp::Add),
+            Opcode::UsubOverflow => (CC::B, ProduceFlagsOp::Sub),
+            Opcode::SsubOverflow => (CC::O, ProduceFlagsOp::Sub),
+            _ => return None,
+        };
+        Some((sum, x, y, cc, op, inst))
+    }
+
+    fn overflow_mul_flag(&mut self, of: Value) -> Option<(Value, Value, Value, bool, Inst)> {
+        let (inst, sum, x, y, opcode) = overflow_flag_common(self.lower_ctx, of)?;
+        let signed = match opcode {
+            Opcode::UmulOverflow => false,
+            Opcode::SmulOverflow => true,
+            _ => return None,
+        };
+        Some((sum, x, y, signed, inst))
+    }
+
+    fn absorb_overflow_inst(&mut self, inst: Inst, sum_regs: ValueRegs) -> Unit {
+        use crate::machinst::ValueRegs as VR;
+        self.lower_ctx.absorb_inst(inst, &[sum_regs, VR::invalid()]);
+    }
 }
 
 impl IsleContext<'_, '_, MInst, X64Backend> {
@@ -2052,6 +2084,45 @@ impl IsleContext<'_, '_, MInst, X64Backend> {
     }
 }
 
+/// Shared fuse-safety checks for matching an overflow-flag value.
+fn overflow_flag_common(
+    lower_ctx: &Lower<MInst>,
+    of: Value,
+) -> Option<(Inst, Value, Value, Value, Opcode)> {
+    let ValueDef::Result(inst, 1) = lower_ctx.dfg().value_def(of) else {
+        return None;
+    };
+    if lower_ctx.inst_is_sunk(inst) {
+        return None;
+    }
+    // Flag must not already be forced into a register, and must have a single
+    // IR use (this consumer) so we can skip materializing `setcc`.
+    if lower_ctx.value_has_lowered_uses(of) || !lower_ctx.value_is_once_used(of) {
+        return None;
+    }
+    let cur = lower_ctx.cur_inst()?;
+    let results = lower_ctx.dfg().inst_results(inst);
+    if results.len() != 2 {
+        return None;
+    }
+    let sum = results[0];
+    let ty = lower_ctx.dfg().value_type(sum);
+    if ty.bits() > 64 || !ty.is_int() {
+        return None;
+    }
+    // Do not fuse if the arithmetic result is used in the same block between
+    // the overflow op and this consumer (flags would be clobbered / order wrong).
+    if lower_ctx.has_same_block_use_between(sum, inst, cur) {
+        return None;
+    }
+    let opcode = lower_ctx.data(inst).opcode();
+    let args = lower_ctx.dfg().inst_args(inst);
+    if args.len() != 2 {
+        return None;
+    }
+    Some((inst, sum, args[0], args[1], opcode))
+}
+
 // Since x64 doesn't have 8x16 shifts and we must use a 16x8 shift instead, we
 // need to fix up the bits that migrate from one half of the lane to the
 // other. Each 16-byte mask is indexed by the shift amount: e.g. if we shift

cranelift/codegen/src/machinst/lower.rs

@@ -1645,8 +1645,16 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         debug_assert!(self.f.dfg.value_is_real(val));
         trace!("put_value_in_regs: val {}", val);
 
+        // Absorbed multi-result instructions (see `absorb_inst`) are marked
+        // sunk even though their results remain live; those results already
+        // have vreg aliases set, so reading them here is fine.
         if let Some(inst) = self.f.dfg.value_def(val).inst() {
-            assert!(!self.inst_sunk.contains(&inst));
+            if self.inst_sunk.contains(&inst) {
+                let regs = self.value_regs[val];
+                assert!(regs.is_valid());
+                self.value_lowered_uses[val] += 1;
+                return regs;
+            }
         }
 
         let regs = self.value_regs[val];
@@ -1657,6 +1665,50 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
 
         regs
     }
+
+    /// Returns true if this value has already been forced into registers during
+    /// lowering (i.e., `put_value_in_regs` has been called on it).
+    pub fn value_has_lowered_uses(&self, val: Value) -> bool {
+        self.value_lowered_uses[val] > 0
+    }
+
+    /// Returns true if this value has exactly one use in the IR use-state
+    /// analysis (see `ValueUseState::Once`).
+    pub fn value_is_once_used(&self, val: Value) -> bool {
+        self.value_ir_uses[val] == ValueUseState::Once
+    }
+
+    /// Current CLIF instruction being lowered, if any.
+    pub fn cur_inst(&self) -> Option<Inst> {
+        self.cur_inst
+    }
+
+    /// Has this instruction been sunk / absorbed away from its original site?
+    pub fn inst_is_sunk(&self, ir_inst: Inst) -> bool {
+        self.inst_sunk.contains(&ir_inst)
+    }
+
+    /// Returns true if `val` is used by any instruction strictly between
+    /// `from_inst` and `to_inst` in layout order within the same block.
+    ///
+    /// Uses on `to_inst` itself (e.g. a branch blockparam) are not counted.
+    pub fn has_same_block_use_between(&self, val: Value, from_inst: Inst, to_inst: Inst) -> bool {
+        let layout = &self.f.layout;
+        if layout.inst_block(from_inst) != layout.inst_block(to_inst) {
+            return true;
+        }
+        let mut cursor = layout.next_inst(from_inst);
+        while let Some(inst) = cursor {
+            if inst == to_inst {
+                break;
+            }
+            if self.f.dfg.inst_values(inst).any(|arg| arg == val) {
+                return true;
+            }
+            cursor = layout.next_inst(inst);
+        }
+        false
+    }
 }
 
 /// Codegen primitives: allocate temps, emit instructions, set result registers,
@@ -1697,6 +1749,37 @@ impl<'func, I: VCodeInst> Lower<'func, I> {
         self.inst_sunk.insert(ir_inst);
     }
 
+    /// Absorb a multi-result instruction into the current lowering site.
+    ///
+    /// Used when fusing e.g. `*_overflow` into a following `brif` so the ALU op
+    /// is emitted next to a flag-consuming branch (`jo`/`jb`). `results` must
+    /// provide a `ValueRegs` for each CLIF result that should be defined
+    /// (typically just the arithmetic result; unused overflow-flag results can
+    /// be `ValueRegs::invalid()`).
+    ///
+    /// Unlike [`Self::sink_inst`], this does not require a lowering side-effect
+    /// and is allowed to leave results live via vreg aliases.
+    pub fn absorb_inst(&mut self, ir_inst: Inst, results: &[ValueRegs<Reg>]) {
+        assert!(!self.inst_sunk.contains(&ir_inst));
+        let inst_results = self.f.dfg.inst_results(ir_inst);
+        debug_assert_eq!(results.len(), inst_results.len());
+
+        for (regs, &result) in results.iter().zip(inst_results) {
+            if regs.is_invalid() {
+                continue;
+            }
+            let dsts = self.value_regs[result];
+            let mut regs_iter = regs.regs().iter();
+            for &dst in dsts.regs().iter() {
+                let temp = regs_iter.next().copied().unwrap_or(Reg::invalid_sentinel());
+                trace!("absorb_inst: set vreg alias for {result:?}: {dst:?} -> {temp:?}");
+                self.vregs.set_vreg_alias(dst, temp);
+            }
+        }
+
+        self.inst_sunk.insert(ir_inst);
+    }
+
     /// Retrieve immediate data given a handle.
     pub fn get_immediate_data(&self, imm: Immediate) -> &ConstantData {
         self.f.dfg.immediates.get(imm).unwrap()