riscv64: Refactor constant emission

cranelift/codegen/src/isa/riscv64/abi.rs

@@ -285,7 +285,6 @@ impl ABIMachineSpec for Riscv64MachineDeps {
             insts.extend(Inst::load_constant_u32(
                 writable_spilltmp_reg2(),
                 imm as u64,
-                &mut |_| writable_spilltmp_reg2(),
             ));
             insts.push(Inst::AluRRR {
                 alu_op: AluOPRRR::Add,
@@ -391,11 +390,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
     }
 
     fn gen_probestack(insts: &mut SmallInstVec<Self::I>, frame_size: u32) {
-        insts.extend(Inst::load_constant_u32(
-            writable_a0(),
-            frame_size as u64,
-            &mut |_| writable_a0(),
-        ));
+        insts.extend(Inst::load_constant_u32(writable_a0(), frame_size as u64));
         insts.push(Inst::Call {
             info: Box::new(CallInfo {
                 dest: ExternalName::LibCall(LibCall::Probestack),
@@ -578,7 +573,7 @@ impl ABIMachineSpec for Riscv64MachineDeps {
         let arg1 = Writable::from_reg(x_reg(11));
         let arg2 = Writable::from_reg(x_reg(12));
         let tmp = alloc_tmp(Self::word_type());
-        insts.extend(Inst::load_constant_u64(tmp, size as u64, &mut alloc_tmp).into_iter());
+        insts.extend(Inst::load_constant_u64(tmp, size as u64).into_iter());
         insts.push(Inst::Call {
             info: Box::new(CallInfo {
                 dest: ExternalName::LibCall(LibCall::Memcpy),

cranelift/codegen/src/isa/riscv64/inst.isle

@@ -1517,16 +1517,51 @@
 (decl pure shift_mask (Type) u64)
 (rule (shift_mask ty) (u64_sub (ty_bits (lane_type ty)) 1))
 
-;; for load immediate
+;; Helper for generating a i64 from a pair of Imm20 and Imm12 constants
+(decl i64_generate_imm (Imm20 Imm12) i64)
+(extern extractor i64_generate_imm i64_generate_imm)
+
+
+;; Immediate Loading rules
+;; TODO: Loading the zero reg directly causes a bunch of regalloc errors, we should look into it.
+;; TODO: We can load constants like 0x3ff0000000000000 by loading the bits then doing a shift left.
 (decl imm (Type u64) Reg)
-(extern constructor imm imm)
+
+;; Refs get loaded as integers.
+(rule 4 (imm $R32 c) (imm $I32 c))
+(rule 4 (imm $R64 c) (imm $I64 c))
+
+;; Floats get loaded as integers and then moved into an F register.
+(rule 4 (imm $F32 c) (gen_bitcast (imm $I32 c) $I32 $F32))
+(rule 4 (imm $F64 c) (gen_bitcast (imm $I64 c) $I64 $F64))
+
+;; Try to match just an imm12
+(rule 3 (imm (ty_int ty) c)
+  (if-let (i64_generate_imm (imm20_is_zero) imm12) (i64_sextend_u64 ty c))
+  (rv_addi (zero_reg) imm12))
+
+;; We can also try to load using a single LUI.
+;; LUI takes a 20 bit immediate, places it on bits 13 to 32 of the register.
+;; In RV64 this value is then sign extended to 64bits.
+(rule 2 (imm (ty_int ty) c)
+  (if-let (i64_generate_imm imm20 (imm12_is_zero)) (i64_sextend_u64 ty c))
+  (rv_lui imm20))
+
+;; We can combo addi + lui to represent all 32-bit immediates
+;; And some 64-bit immediates as well.
+(rule 1 (imm (ty_int ty) c)
+  (if-let (i64_generate_imm imm20 imm12) (i64_sextend_u64 ty c))
+  (rv_addi (rv_lui imm20) imm12))
+
+;; Otherwise we fall back to loading the immediate from memory
+;; TODO: This doesen't yet emit the constant into the memory pool, but it would
+;; be pretty neat if it did.
+(rule 0 (imm (ty_int ty) c) (emit_load_const_64 c))
 
 ;; Imm12 Rules
 
 (decl pure imm12_zero () Imm12)
-(rule
-  (imm12_zero)
-  (imm12_const 0))
+(rule (imm12_zero) (imm12_const 0))
 
 (decl pure imm12_const (i32) Imm12)
 (extern constructor imm12_const imm12_const)
@@ -1549,39 +1584,48 @@
 (decl imm12_and (Imm12 u64) Imm12)
 (extern constructor imm12_and imm12_and)
 
-;; Helper for get negative of Imm12
-(decl neg_imm12 (Imm12) Imm12)
-(extern constructor neg_imm12 neg_imm12)
-
 ;; Imm12 Extractors
 
 ;; Helper to go directly from a `Value`, when it's an `iconst`, to an `Imm12`.
 (decl imm12_from_value (Imm12) Value)
-(extractor
-  (imm12_from_value n)
-  (def_inst (iconst (u64_from_imm64 (imm12_from_u64 n)))))
+(extractor (imm12_from_value n) (i64_from_iconst (imm12_from_i64 n)))
 
 (decl imm12_from_u64 (Imm12) u64)
 (extern extractor imm12_from_u64 imm12_from_u64)
 
+(decl imm12_from_i64 (Imm12) i64)
+(extern extractor imm12_from_i64 imm12_from_i64)
+
 (decl pure partial u64_to_imm12 (u64) Imm12)
 (rule (u64_to_imm12 (imm12_from_u64 n)) n)
 
+(decl pure imm12_is_zero () Imm12)
+(extern extractor imm12_is_zero imm12_is_zero)
+
+;; Imm20
+
+;; Extractor that matches if a Imm20 is zero
+(decl pure imm20_is_zero () Imm20)
+(extern extractor imm20_is_zero imm20_is_zero)
+
 
 ;; Imm5 Extractors
 
 (decl imm5_from_u64 (Imm5) u64)
 (extern extractor imm5_from_u64 imm5_from_u64)
 
+(decl imm5_from_i64 (Imm5) i64)
+(extern extractor imm5_from_i64 imm5_from_i64)
+
 ;; Construct a Imm5 from an i8
-(decl pure partial imm5_from_i8 (i8) Imm5)
-(extern constructor imm5_from_i8 imm5_from_i8)
+(decl pure partial i8_to_imm5 (i8) Imm5)
+(extern constructor i8_to_imm5 i8_to_imm5)
 
 ;; Extractor that matches a `Value` equivalent to a replicated Imm5 on all lanes.
 ;; TODO(#6527): Try matching vconst here as well
 (decl replicated_imm5 (Imm5) Value)
 (extractor (replicated_imm5 n)
-  (def_inst (splat (iconst (u64_from_imm64 (imm5_from_u64 n))))))
+  (def_inst (splat (i64_from_iconst (imm5_from_i64 n)))))
 
 ;; UImm5 Helpers
 
@@ -1670,6 +1714,24 @@
             (_ Unit (emit (MInst.AluRRImm12 op dst src (imm12_zero)))))
         dst))
 
+;; Helper for emitting the `LoadConst64` instruction.
+(decl emit_load_const_64 (u64) XReg)
+(rule (emit_load_const_64 imm)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.LoadConst64 dst imm))))
+        dst))
+
+;; Helper for emitting the `Lui` instruction.
+;; TODO: This should be something like `emit_u_type`. And should share the
+;; `MInst` with `auipc` since these instructions share the U-Type format.
+(decl rv_lui (Imm20) XReg)
+(rule (rv_lui imm)
+      (let ((dst WritableXReg (temp_writable_xreg))
+            (_ Unit (emit (MInst.Lui dst imm))))
+        dst))
+
+
+
 (decl select_addi (Type) AluOPRRI)
 (rule 1 (select_addi (fits_in_32 ty)) (AluOPRRI.Addiw))
 (rule (select_addi (fits_in_64 ty)) (AluOPRRI.Addi))
@@ -1989,7 +2051,7 @@
 (rule 3 (extend val (ExtendOp.Zero) (fits_in_64 from_ty) $I128)
   (let ((val XReg (value_regs_get val 0))
         (low XReg (zext val from_ty $I64))
-        (high XReg (load_u64_constant 0)))
+        (high XReg (imm $I64 0)))
     (value_regs low high)))
 
 ;; Catch all rule for ignoring extensions of the same type.
@@ -2127,7 +2189,7 @@
 (rule (gen_bseti val bit)
   (if-let $false (has_zbs))
   (if-let $false (u64_le bit 12))
-  (let ((const XReg (load_u64_constant (u64_shl 1 bit))))
+  (let ((const XReg (imm $I64 (u64_shl 1 bit))))
     (rv_or val const)))
 
 (rule (gen_bseti val bit)
@@ -2169,7 +2231,7 @@
       (high XReg (lower_popcnt (value_regs_get a 1) $I64))
       ;; add toghter.
       (result XReg (rv_add low high)))
-    (value_regs result (load_u64_constant 0))))
+    (value_regs result (imm $I64 0))))
 
 (decl lower_i128_rotl (ValueRegs ValueRegs) ValueRegs)
 (rule
@@ -2190,7 +2252,7 @@
       (high_part3 XReg (gen_select_reg (IntCC.Equal) shamt (zero_reg) (zero_reg) high_part2))
       (high XReg (rv_or high_part1 high_part3))
       ;;
-      (const64 XReg (load_u64_constant 64))
+      (const64 XReg (imm $I64 64))
       (shamt_128 XReg (rv_andi (value_regs_get y 0) (imm12_const 127))))
     ;; right now we only rotate less than 64 bits.
     ;; if shamt is greater than or equal 64 , we should switch low and high.
@@ -2220,7 +2282,7 @@
       (high XReg (rv_or high_part1 high_part3))
 
       ;;
-      (const64 XReg (load_u64_constant 64))
+      (const64 XReg (imm $I64 64))
       (shamt_128 XReg (rv_andi (value_regs_get y 0) (imm12_const 127))))
     ;; right now we only rotate less than 64 bits.
     ;; if shamt is greater than or equal 64 , we should switch low and high.
@@ -2402,10 +2464,6 @@
 (decl gen_select_reg (IntCC XReg XReg Reg Reg) Reg)
 (extern constructor gen_select_reg gen_select_reg)
 
-;; load a constant into reg.
-(decl load_u64_constant (u64) Reg)
-(extern constructor load_u64_constant load_u64_constant)
-
 ;;; clone WritableReg
 ;;; if not rust compiler will complain about use moved value.
 (decl vec_writable_clone (VecWritableReg) VecWritableReg)
@@ -2733,8 +2791,8 @@
 (rule
   (gen_div_overflow rs1 rs2 ty)
   (let
-    ((r_const_neg_1 XReg (load_imm12 -1))
-      (r_const_min XReg (rv_slli (load_imm12 1) (imm12_const 63)))
+    ((r_const_neg_1 XReg (imm $I64 (i64_as_u64 -1)))
+      (r_const_min XReg (rv_slli (imm $I64 1) (imm12_const 63)))
       (tmp_rs1 XReg (shift_int_to_most_significant rs1 ty))
       (t1 XReg (gen_icmp (IntCC.Equal) r_const_neg_1 rs2 ty))
       (t2 XReg (gen_icmp (IntCC.Equal) r_const_min tmp_rs1 ty))

cranelift/codegen/src/isa/riscv64/inst/emit.rs

@@ -177,18 +177,6 @@ impl MachInstEmitState<Inst> for EmitState {
 }
 
 impl Inst {
-    /// construct a "imm - rs".
-    pub(crate) fn construct_imm_sub_rs(rd: Writable<Reg>, imm: u64, rs: Reg) -> SmallInstVec<Inst> {
-        let mut insts = Inst::load_constant_u64(rd, imm, &mut |_| rd);
-        insts.push(Inst::AluRRR {
-            alu_op: AluOPRRR::Sub,
-            rd,
-            rs1: rd.to_reg(),
-            rs2: rs,
-        });
-        insts
-    }
-
     /// Load int mask.
     /// If ty is int then 0xff in rd.
     pub(crate) fn load_int_mask(rd: Writable<Reg>, ty: Type) -> SmallInstVec<Inst> {
@@ -676,7 +664,7 @@ impl MachInstEmit for Inst {
 
                 let base = from.get_base_register();
                 let offset = from.get_offset_with_state(state);
-                let offset_imm12 = Imm12::maybe_from_u64(offset as u64);
+                let offset_imm12 = Imm12::maybe_from_i64(offset);
 
                 let (addr, imm12) = match (base, offset_imm12) {
                     // If the offset fits into an imm12 we can directly encode it.
@@ -703,7 +691,7 @@ impl MachInstEmit for Inst {
 
                 let base = to.get_base_register();
                 let offset = to.get_offset_with_state(state);
-                let offset_imm12 = Imm12::maybe_from_u64(offset as u64);
+                let offset_imm12 = Imm12::maybe_from_i64(offset);
 
                 let (addr, imm12) = match (base, offset_imm12) {
                     // If the offset fits into an imm12 we can directly encode it.
@@ -785,7 +773,7 @@ impl MachInstEmit for Inst {
                     .for_each(|i| i.emit(&[], sink, emit_info, state));
             }
             &Inst::AdjustSp { amount } => {
-                if let Some(imm) = Imm12::maybe_from_u64(amount as u64) {
+                if let Some(imm) = Imm12::maybe_from_i64(amount) {
                     Inst::AluRRImm12 {
                         alu_op: AluOPRRI::Addi,
                         rd: writable_stack_reg(),
@@ -795,7 +783,7 @@ impl MachInstEmit for Inst {
                     .emit(&[], sink, emit_info, state);
                 } else {
                     let tmp = writable_spilltmp_reg();
-                    let mut insts = Inst::load_constant_u64(tmp, amount as u64, &mut |_| tmp);
+                    let mut insts = Inst::load_constant_u64(tmp, amount as u64);
                     insts.push(Inst::AluRRR {
                         alu_op: AluOPRRR::Add,
                         rd: writable_stack_reg(),
@@ -1119,7 +1107,7 @@ impl MachInstEmit for Inst {
                 // Check if the index passed in is larger than the number of jumptable
                 // entries that we have. If it is, we fallthrough to a jump into the
                 // default block.
-                Inst::load_constant_u32(tmp2, targets.len() as u64, &mut |_| tmp2)
+                Inst::load_constant_u32(tmp2, targets.len() as u64)
                     .iter()
                     .for_each(|i| i.emit(&[], sink, emit_info, state));
                 Inst::CondBr {
@@ -1264,7 +1252,7 @@ impl MachInstEmit for Inst {
 
                 let base = mem.get_base_register();
                 let offset = mem.get_offset_with_state(state);
-                let offset_imm12 = Imm12::maybe_from_u64(offset as u64);
+                let offset_imm12 = Imm12::maybe_from_i64(offset);
 
                 match (mem, base, offset_imm12) {
                     (_, Some(rs), Some(imm12)) => {
@@ -1914,7 +1902,6 @@ impl MachInstEmit for Inst {
                             // I8
                             (u8::MAX >> 1) as u64
                         },
-                        &mut |_| writable_spilltmp_reg2(),
                     )
                     .into_iter()
                     .for_each(|x| x.emit(&[], sink, emit_info, state));
@@ -2791,14 +2778,10 @@ impl MachInstEmit for Inst {
                 tmp: guard_size_tmp,
             } => {
                 let step = writable_spilltmp_reg();
-                Inst::load_constant_u64(
-                    step,
-                    (guard_size as u64) * (probe_count as u64),
-                    &mut |_| step,
-                )
-                .iter()
-                .for_each(|i| i.emit(&[], sink, emit_info, state));
-                Inst::load_constant_u64(guard_size_tmp, guard_size as u64, &mut |_| guard_size_tmp)
+                Inst::load_constant_u64(step, (guard_size as u64) * (probe_count as u64))
+                    .iter()
+                    .for_each(|i| i.emit(&[], sink, emit_info, state));
+                Inst::load_constant_u64(guard_size_tmp, guard_size as u64)
                     .iter()
                     .for_each(|i| i.emit(&[], sink, emit_info, state));
 
@@ -3205,7 +3188,7 @@ fn emit_return_call_common_sequence(
         alu_op: AluOPRRI::Addi,
         rd: regs::writable_stack_reg(),
         rs: regs::fp_reg(),
-        imm12: Imm12::maybe_from_u64(fp_to_callee_sp as u64).unwrap(),
+        imm12: Imm12::maybe_from_i64(fp_to_callee_sp).unwrap(),
     }
     .emit(&[], sink, emit_info, state);