Merge pull request bytecodealliance#41 from dhil/wasmfx-merge

Weekly merge with upstream

Cargo.toml

@@ -347,7 +347,7 @@ component-model = [
   "wasmtime-wast/component-model",
   "wasmtime-cli-flags/component-model"
 ]
-wat = ["dep:wat"]
+wat = ["dep:wat", "wasmtime/wat"]
 cache = ["dep:wasmtime-cache", "wasmtime-cli-flags/cache"]
 parallel-compilation = ["wasmtime-cli-flags/parallel-compilation"]
 logging = ["wasmtime-cli-flags/logging"]

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

@@ -833,6 +833,15 @@
         (not_taken BranchTarget)
         (kind CondBrKind))
 
+       ;; A conditional branch which tests the `bit` of `rn` and branches
+       ;; depending on `kind`.
+       (TestBitAndBranch
+        (kind TestBitAndBranchKind)
+        (taken BranchTarget)
+        (not_taken BranchTarget)
+        (rn Reg)
+        (bit u8))
+
        ;; A conditional trap: execute a `udf` if the condition is true. This is
        ;; one VCode instruction because it uses embedded control flow; it is
        ;; logically a single-in, single-out region, but needs to appear as one
@@ -1203,6 +1212,8 @@
       (enum Size32
             Size64))
 
+(type TestBitAndBranchKind (enum (Z) (NZ)))
+
 ;; Helper for calculating the `OperandSize` corresponding to a type
 (decl operand_size (Type) OperandSize)
 (rule 1 (operand_size (fits_in_32 _ty)) (OperandSize.Size32))
@@ -4079,6 +4090,21 @@
       (ConsumesFlags.ConsumesFlagsSideEffect
        (MInst.CondBr taken not_taken kind)))
 
+;; Helper for emitting `MInst.TestBitAndBranch` instructions.
+(decl test_branch (TestBitAndBranchKind BranchTarget BranchTarget Reg u8) SideEffectNoResult)
+(rule (test_branch kind taken not_taken rn bit)
+      (SideEffectNoResult.Inst (MInst.TestBitAndBranch kind taken not_taken rn bit)))
+
+;; Helper for emitting `tbnz` instructions.
+(decl tbnz (BranchTarget BranchTarget Reg u8) SideEffectNoResult)
+(rule (tbnz taken not_taken rn bit)
+      (test_branch (TestBitAndBranchKind.NZ) taken not_taken rn bit))
+
+;; Helper for emitting `tbz` instructions.
+(decl tbz (BranchTarget BranchTarget Reg u8) SideEffectNoResult)
+(rule (tbz taken not_taken rn bit)
+      (test_branch (TestBitAndBranchKind.Z) taken not_taken rn bit))
+
 ;; Helper for emitting `MInst.MovToNZCV` instructions.
 (decl mov_to_nzcv (Reg) ProducesFlags)
 (rule (mov_to_nzcv rn)

cranelift/codegen/src/isa/aarch64/inst/args.rs

@@ -340,26 +340,31 @@ impl BranchTarget {
         }
     }
 
+    /// Return the target's offset, if specified, or zero if label-based.
+    pub fn as_offset14_or_zero(self) -> u32 {
+        self.as_offset_bounded(14)
+    }
+
     /// Return the target's offset, if specified, or zero if label-based.
     pub fn as_offset19_or_zero(self) -> u32 {
-        let off = match self {
-            BranchTarget::ResolvedOffset(off) => off >> 2,
-            _ => 0,
-        };
-        assert!(off <= 0x3ffff);
-        assert!(off >= -0x40000);
-        (off as u32) & 0x7ffff
+        self.as_offset_bounded(19)
     }
 
     /// Return the target's offset, if specified, or zero if label-based.
     pub fn as_offset26_or_zero(self) -> u32 {
+        self.as_offset_bounded(26)
+    }
+
+    fn as_offset_bounded(self, bits: u32) -> u32 {
         let off = match self {
             BranchTarget::ResolvedOffset(off) => off >> 2,
             _ => 0,
         };
-        assert!(off <= 0x1ffffff);
-        assert!(off >= -0x2000000);
-        (off as u32) & 0x3ffffff
+        let hi = (1 << (bits - 1)) - 1;
+        let lo = -(1 << bits - 1);
+        assert!(off <= hi);
+        assert!(off >= lo);
+        (off as u32) & ((1 << bits) - 1)
     }
 }
 
@@ -764,3 +769,15 @@ impl APIKey {
         0xd503201f | (crm << 8) | (op2 << 5)
     }
 }
+
+pub use crate::isa::aarch64::lower::isle::generated_code::TestBitAndBranchKind;
+
+impl TestBitAndBranchKind {
+    /// Complements this branch condition to act on the opposite result.
+    pub fn complement(&self) -> TestBitAndBranchKind {
+        match self {
+            TestBitAndBranchKind::Z => TestBitAndBranchKind::NZ,
+            TestBitAndBranchKind::NZ => TestBitAndBranchKind::Z,
+        }
+    }
+}

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

@@ -172,6 +172,27 @@ fn enc_conditional_br(
     }
 }
 
+fn enc_test_bit_and_branch(
+    kind: TestBitAndBranchKind,
+    taken: BranchTarget,
+    reg: Reg,
+    bit: u8,
+) -> u32 {
+    assert!(bit < 64);
+    let op_31 = u32::from(bit >> 5);
+    let op_23_19 = u32::from(bit & 0b11111);
+    let op_30_24 = 0b0110110
+        | match kind {
+            TestBitAndBranchKind::Z => 0,
+            TestBitAndBranchKind::NZ => 1,
+        };
+    (op_31 << 31)
+        | (op_30_24 << 24)
+        | (op_23_19 << 19)
+        | (taken.as_offset14_or_zero() << 5)
+        | machreg_to_gpr(reg)
+}
+
 fn enc_move_wide(op: MoveWideOp, rd: Writable<Reg>, imm: MoveWideConst, size: OperandSize) -> u32 {
     assert!(imm.shift <= 0b11);
     let op = match op {
@@ -3224,6 +3245,32 @@ impl MachInstEmit for Inst {
                 }
                 sink.put4(enc_jump26(0b000101, not_taken.as_offset26_or_zero()));
             }
+            &Inst::TestBitAndBranch {
+                taken,
+                not_taken,
+                kind,
+                rn,
+                bit,
+            } => {
+                let rn = allocs.next(rn);
+                // Emit the conditional branch first
+                let cond_off = sink.cur_offset();
+                if let Some(l) = taken.as_label() {
+                    sink.use_label_at_offset(cond_off, l, LabelUse::Branch14);
+                    let inverted =
+                        enc_test_bit_and_branch(kind.complement(), taken, rn, bit).to_le_bytes();
+                    sink.add_cond_branch(cond_off, cond_off + 4, l, &inverted[..]);
+                }
+                sink.put4(enc_test_bit_and_branch(kind, taken, rn, bit));
+
+                // Unconditional part next.
+                let uncond_off = sink.cur_offset();
+                if let Some(l) = not_taken.as_label() {
+                    sink.use_label_at_offset(uncond_off, l, LabelUse::Branch26);
+                    sink.add_uncond_branch(uncond_off, uncond_off + 4, l);
+                }
+                sink.put4(enc_jump26(0b000101, not_taken.as_offset26_or_zero()));
+            }
             &Inst::TrapIf { kind, trap_code } => {
                 let label = sink.defer_trap(trap_code, state.take_stack_map());
                 // condbr KIND, LABEL

cranelift/codegen/src/isa/aarch64/inst/mod.rs

@@ -902,6 +902,9 @@ fn aarch64_get_operands<F: Fn(VReg) -> VReg>(inst: &Inst, collector: &mut Operan
             }
             CondBrKind::Cond(_) => {}
         },
+        &Inst::TestBitAndBranch { rn, .. } => {
+            collector.reg_use(rn);
+        }
         &Inst::IndirectBr { rn, .. } => {
             collector.reg_use(rn);
         }
@@ -1043,6 +1046,7 @@ impl MachInst for Inst {
             &Inst::ReturnCall { .. } | &Inst::ReturnCallInd { .. } => MachTerminator::RetCall,
             &Inst::Jump { .. } => MachTerminator::Uncond,
             &Inst::CondBr { .. } => MachTerminator::Cond,
+            &Inst::TestBitAndBranch { .. } => MachTerminator::Cond,
             &Inst::IndirectBr { .. } => MachTerminator::Indirect,
             &Inst::JTSequence { .. } => MachTerminator::Indirect,
             _ => MachTerminator::None,
@@ -2663,6 +2667,22 @@ impl Inst {
                     }
                 }
             }
+            &Inst::TestBitAndBranch {
+                kind,
+                ref taken,
+                ref not_taken,
+                rn,
+                bit,
+            } => {
+                let cond = match kind {
+                    TestBitAndBranchKind::Z => "z",
+                    TestBitAndBranchKind::NZ => "nz",
+                };
+                let taken = taken.pretty_print(0, allocs);
+                let not_taken = not_taken.pretty_print(0, allocs);
+                let rn = pretty_print_reg(rn, allocs);
+                format!("tb{cond} {rn}, #{bit}, {taken} ; b {not_taken}")
+            }
             &Inst::IndirectBr { rn, .. } => {
                 let rn = pretty_print_reg(rn, allocs);
                 format!("br {}", rn)
@@ -2882,6 +2902,9 @@ impl Inst {
 /// Different forms of label references for different instruction formats.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum LabelUse {
+    /// 14-bit branch offset (conditional branches). PC-rel, offset is imm <<
+    /// 2. Immediate is 14 signed bits, in bits 18:5. Used by tbz and tbnz.
+    Branch14,
     /// 19-bit branch offset (conditional branches). PC-rel, offset is imm << 2. Immediate is 19
     /// signed bits, in bits 23:5. Used by cbz, cbnz, b.cond.
     Branch19,
@@ -2908,8 +2931,10 @@ impl MachInstLabelUse for LabelUse {
     /// Maximum PC-relative range (positive), inclusive.
     fn max_pos_range(self) -> CodeOffset {
         match self {
-            // 19-bit immediate, left-shifted by 2, for 21 bits of total range. Signed, so +2^20
-            // from zero. Likewise for two other shifted cases below.
+            // N-bit immediate, left-shifted by 2, for (N+2) bits of total
+            // range. Signed, so +2^(N+1) from zero. Likewise for two other
+            // shifted cases below.
+            LabelUse::Branch14 => (1 << 15) - 1,
             LabelUse::Branch19 => (1 << 20) - 1,
             LabelUse::Branch26 => (1 << 27) - 1,
             LabelUse::Ldr19 => (1 << 20) - 1,
@@ -2941,6 +2966,7 @@ impl MachInstLabelUse for LabelUse {
         let pc_rel = pc_rel as u32;
         let insn_word = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]);
         let mask = match self {
+            LabelUse::Branch14 => 0x0007ffe0, // bits 18..5 inclusive
             LabelUse::Branch19 => 0x00ffffe0, // bits 23..5 inclusive
             LabelUse::Branch26 => 0x03ffffff, // bits 25..0 inclusive
             LabelUse::Ldr19 => 0x00ffffe0,    // bits 23..5 inclusive
@@ -2955,6 +2981,7 @@ impl MachInstLabelUse for LabelUse {
             }
         };
         let pc_rel_inserted = match self {
+            LabelUse::Branch14 => (pc_rel_shifted & 0x3fff) << 5,
             LabelUse::Branch19 | LabelUse::Ldr19 => (pc_rel_shifted & 0x7ffff) << 5,
             LabelUse::Branch26 => pc_rel_shifted & 0x3ffffff,
             LabelUse::Adr21 => (pc_rel_shifted & 0x7ffff) << 5 | (pc_rel_shifted & 0x180000) << 10,
@@ -2975,16 +3002,16 @@ impl MachInstLabelUse for LabelUse {
     /// Is a veneer supported for this label reference type?
     fn supports_veneer(self) -> bool {
         match self {
-            LabelUse::Branch19 => true, // veneer is a Branch26
-            LabelUse::Branch26 => true, // veneer is a PCRel32
+            LabelUse::Branch14 | LabelUse::Branch19 => true, // veneer is a Branch26
+            LabelUse::Branch26 => true,                      // veneer is a PCRel32
             _ => false,
         }
     }
 
     /// How large is the veneer, if supported?
     fn veneer_size(self) -> CodeOffset {
         match self {
-            LabelUse::Branch19 => 4,
+            LabelUse::Branch14 | LabelUse::Branch19 => 4,
             LabelUse::Branch26 => 20,
             _ => unreachable!(),
         }
@@ -3002,7 +3029,7 @@ impl MachInstLabelUse for LabelUse {
         veneer_offset: CodeOffset,
     ) -> (CodeOffset, LabelUse) {
         match self {
-            LabelUse::Branch19 => {
+            LabelUse::Branch14 | LabelUse::Branch19 => {
                 // veneer is a Branch26 (unconditional branch). Just encode directly here -- don't
                 // bother with constructing an Inst.
                 let insn_word = 0b000101 << 26;

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

@@ -2882,6 +2882,24 @@
         (with_flags_side_effect flags
          (cond_br taken not_taken (cond_br_not_zero rt))))))
 
+;; Special lowerings for `tbnz` - "Test bit and Branch if Nonzero"
+(rule 1 (lower_branch (brif (band x @ (value_type ty) (u64_from_iconst n)) _ _)
+                     (two_targets taken not_taken))
+  (if-let bit (test_and_compare_bit_const ty n))
+  (emit_side_effect (tbnz taken not_taken x bit)))
+
+;; Special lowering for `tbz` - "Test bit and Branch if Zero"
+(rule 1 (lower_branch (brif (icmp (IntCC.Equal)
+                                  (band x @ (value_type (fits_in_64 ty))
+                                        (u64_from_iconst n))
+                                  (u64_from_iconst 0)) _ _)
+                     (two_targets taken not_taken))
+  (if-let bit (test_and_compare_bit_const ty n))
+  (emit_side_effect (tbz taken not_taken x bit)))
+
+(decl pure partial test_and_compare_bit_const (Type u64) u8)
+(extern constructor test_and_compare_bit_const test_and_compare_bit_const)
+
 ;;; Rules for `jump` ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 
 (rule (lower_branch (jump _) (single_target label))

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

@@ -812,4 +812,15 @@ impl Context for IsleContext<'_, '_, MInst, AArch64Backend> {
     fn uimm12_scaled_from_i64(&mut self, val: i64, ty: Type) -> Option<UImm12Scaled> {
         UImm12Scaled::maybe_from_i64(val, ty)
     }
+
+    fn test_and_compare_bit_const(&mut self, ty: Type, n: u64) -> Option<u8> {
+        if n.count_ones() != 1 {
+            return None;
+        }
+        let bit = n.trailing_zeros();
+        if bit >= ty.bits() {
+            return None;
+        }
+        Some(bit as u8)
+    }
 }

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

@@ -1783,9 +1783,6 @@
 
 
 
-;; Generate a mask for the bit-width of the given type
-(decl pure shift_mask (Type) u64)
-(rule (shift_mask ty) (u64_sub (ty_bits (lane_type ty)) 1))
 
 ;; Helper for generating a i64 from a pair of Imm20 and Imm12 constants
 (decl i64_generate_imm (Imm20 Imm12) i64)

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

@@ -1193,7 +1193,7 @@
 
 ;; 8/16 bit types need a mask on the shift amount
 (rule 0 (lower (has_type (ty_int (ty_8_or_16 ty)) (ishl x y)))
-  (if-let mask (u64_to_imm12 (shift_mask ty)))
+  (if-let mask (u64_to_imm12 (ty_shift_mask ty)))
   (rv_sllw x (rv_andi (value_regs_get y 0) mask)))
 
 ;; Using the 32bit version of `sll` automatically masks the shift amount.
@@ -1206,12 +1206,12 @@
 
 ;; If the shift amount is known. We can mask it and encode it in the instruction.
 (rule 2 (lower (has_type (int_fits_in_32 ty) (ishl x (maybe_uextend (imm12_from_value y)))))
-  (rv_slliw x (imm12_and y (shift_mask ty))))
+  (rv_slliw x (imm12_and y (ty_shift_mask ty))))
 
 ;; We technically don't need to mask the shift amount here. The instruction
 ;; does the right thing. But it's neater when pretty printing it.
 (rule 3 (lower (has_type ty @ $I64 (ishl x (maybe_uextend (imm12_from_value y)))))
-  (rv_slli x (imm12_and y (shift_mask ty))))
+  (rv_slli x (imm12_and y (ty_shift_mask ty))))
 
 ;; With `Zba` we have a shift that zero extends the LHS argument.
 (rule 4 (lower (has_type $I64 (ishl (uextend x @ (value_type $I32)) (maybe_uextend (imm12_from_value y)))))
@@ -1253,7 +1253,7 @@
 ;; 8/16 bit types need a mask on the shift amount, and the LHS needs to be
 ;; zero extended.
 (rule 0 (lower (has_type (ty_int (fits_in_16 ty)) (ushr x y)))
-  (if-let mask (u64_to_imm12 (shift_mask ty)))
+  (if-let mask (u64_to_imm12 (ty_shift_mask ty)))
   (rv_srlw (zext x) (rv_andi (value_regs_get y 0) mask)))
 
 ;; Using the 32bit version of `srl` automatically masks the shift amount.
@@ -1266,7 +1266,7 @@
 
 ;; When the RHS is known we can just encode it in the instruction.
 (rule 2 (lower (has_type (ty_int (fits_in_16 ty)) (ushr x (maybe_uextend (imm12_from_value y)))))
-  (rv_srliw (zext x) (imm12_and y (shift_mask ty))))
+  (rv_srliw (zext x) (imm12_and y (ty_shift_mask ty))))
 
 (rule 3 (lower (has_type $I32 (ushr x (maybe_uextend (imm12_from_value y)))))
   (rv_srliw x y))
@@ -1308,7 +1308,7 @@
 ;; 8/16 bit types need a mask on the shift amount, and the LHS needs to be
 ;; zero extended.
 (rule 0 (lower (has_type (ty_int (fits_in_16 ty)) (sshr x y)))
-  (if-let mask (u64_to_imm12 (shift_mask ty)))
+  (if-let mask (u64_to_imm12 (ty_shift_mask ty)))
   (rv_sraw (sext x) (rv_andi (value_regs_get y 0) mask)))
 
 ;; Using the 32bit version of `sra` automatically masks the shift amount.
@@ -1321,7 +1321,7 @@
 
 ;; When the RHS is known we can just encode it in the instruction.
 (rule 2 (lower (has_type (ty_int (fits_in_16 ty)) (sshr x (maybe_uextend (imm12_from_value y)))))
-  (rv_sraiw (sext x) (imm12_and y (shift_mask ty))))
+  (rv_sraiw (sext x) (imm12_and y (ty_shift_mask ty))))
 
 (rule 3 (lower (has_type $I32 (sshr x (maybe_uextend (imm12_from_value y)))))
   (rv_sraiw x y))