cranelift: Choose memory trap code based on memflags (#8134)

Ideally we'd allow the frontend to specify what trap code a memory
access instruction can produce, but we don't have room to store that
information.

We do have a few bits to spare in `MemFlags`, though, so add a new
`tabletrap` bit to indicate that the trap code should be
TableOutOfBounds instead of the default HeapOutOfBounds.

Nothing yet sets the `tabletrap` flag, so the TableOutOfBounds case is
never hit, but I wanted to separate this change out from the
cranelift-wasm changes which will use it.

The original version of this patch reused the `table` flag instead of
introducing a new `tabletrap` flag. But this usage changes the semantics
of an instruction, while the `table` flag is meant to be a hint for
alias analysis in the egraph optimization pass.

cranelift/codegen/src/ir/memflags.rs

@@ -1,5 +1,6 @@
 //! Memory operation flags.
 
+use super::TrapCode;
 use core::fmt;
 
 #[cfg(feature = "enable-serde")]
@@ -9,6 +10,9 @@ enum FlagBit {
  /// Guaranteed not to trap. This may enable additional
  /// optimizations to be performed.
  Notrap,
+ /// If this memory access traps, use [TrapCode::TableOutOfBounds]
+ /// instead of [TrapCode::HeapOutOfBounds].
+ TableTrap,
  /// Guaranteed to use "natural alignment" for the given type. This
  /// may enable better instruction selection.
  Aligned,
@@ -36,8 +40,17 @@ enum FlagBit {
  Checked,
 }
 
-const NAMES: [&str; 9] = [
- "notrap", "aligned", "readonly", "little", "big", "heap", "table", "vmctx", "checked",
+const NAMES: [&str; 10] = [
+ "notrap",
+ "tabletrap",
+ "aligned",
+ "readonly",
+ "little",
+ "big",
+ "heap",
+ "table",
+ "vmctx",
+ "checked",
 ];
 
 /// Endianness of a memory access.
@@ -162,6 +175,25 @@ impl MemFlags {
  self.with(FlagBit::Notrap)
  }
 
+ /// Test if the `tabletrap` bit is set.
+ ///
+ /// By default, if this memory access traps, the associated trap
+ /// code will be `HeapOutOfBounds`. With this flag set, the trap
+ /// code will instead be `TableOutOfBounds`.
+ pub const fn tabletrap(self) -> bool {
+ self.read(FlagBit::TableTrap)
+ }
+
+ /// Set the `tabletrap` bit.
+ pub fn set_tabletrap(&mut self) {
+ *self = self.with_tabletrap();
+ }
+
+ /// Set the `tabletrap` bit, returning new flags.
+ pub const fn with_tabletrap(self) -> Self {
+ self.with(FlagBit::TableTrap)
+ }
+
  /// Test if the `aligned` flag is set.
  ///
  /// By default, Cranelift memory instructions work with any unaligned effective address. If the
@@ -296,6 +328,21 @@ impl MemFlags {
  pub const fn with_checked(self) -> Self {
  self.with(FlagBit::Checked)
  }
+
+ /// Get the trap code to report if this memory access traps. If the
+ /// `notrap` flag is set, then return `None`. Otherwise, the choice
+ /// of trap code depends on the `tabletrap` flag. The default trap
+ /// code is `HeapOutOfBounds`, but with `tabletrap` set, the trap
+ /// code is instead `TableOutOfBounds`.
+ pub const fn trap_code(self) -> Option<TrapCode> {
+ if self.notrap() {
+ None
+ } else if self.tabletrap() {
+ Some(TrapCode::TableOutOfBounds)
+ } else {
+ Some(TrapCode::HeapOutOfBounds)
+ }
+ }
 }
 
 impl fmt::Display for MemFlags {

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

@@ -4,7 +4,7 @@ use cranelift_control::ControlPlane;
 use regalloc2::Allocation;
 
 use crate::binemit::StackMap;
-use crate::ir::{self, types::*, TrapCode};
+use crate::ir::{self, types::*};
 use crate::isa::aarch64::inst::*;
 use crate::trace;
 
@@ -1005,9 +1005,9 @@ impl MachInstEmit for Inst {
  _ => unreachable!(),
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  match &mem {
@@ -1140,9 +1140,9 @@ impl MachInstEmit for Inst {
  _ => unreachable!(),
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual store instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  match &mem {
@@ -1219,9 +1219,9 @@ impl MachInstEmit for Inst {
  let rt = allocs.next(rt);
  let rt2 = allocs.next(rt2);
  let mem = mem.with_allocs(&mut allocs);
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual store instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
  match &mem {
  &PairAMode::SignedOffset { reg, simm7 } => {
@@ -1250,9 +1250,9 @@ impl MachInstEmit for Inst {
  let rt = allocs.next(rt.to_reg());
  let rt2 = allocs.next(rt2.to_reg());
  let mem = mem.with_allocs(&mut allocs);
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  match &mem {
@@ -1289,9 +1289,9 @@ impl MachInstEmit for Inst {
  let rt2 = allocs.next(rt2.to_reg());
  let mem = mem.with_allocs(&mut allocs);
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  let opc = match self {
@@ -1334,9 +1334,9 @@ impl MachInstEmit for Inst {
  let rt2 = allocs.next(rt2);
  let mem = mem.with_allocs(&mut allocs);
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual store instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  let opc = match self {
@@ -1494,8 +1494,8 @@ impl MachInstEmit for Inst {
  let rt = allocs.next_writable(rt);
  let rn = allocs.next(rn);
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_acq_rel(ty, op, rs, rt, rn));
@@ -1534,8 +1534,8 @@ impl MachInstEmit for Inst {
  // again:
  sink.bind_label(again_label, &mut state.ctrl_plane);
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_ldaxr(ty, x27wr, x25)); // ldaxr x27, [x25]
@@ -1658,8 +1658,8 @@ impl MachInstEmit for Inst {
  }
  }
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
  if op == AtomicRMWLoopOp::Xchg {
  sink.put4(enc_stlxr(ty, x24wr, x26, x25)); // stlxr w24, x26, [x25]
@@ -1699,8 +1699,8 @@ impl MachInstEmit for Inst {
  _ => panic!("Unsupported type: {}", ty),
  };
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_cas(size, rd, rt, rn));
@@ -1733,8 +1733,8 @@ impl MachInstEmit for Inst {
  // again:
  sink.bind_label(again_label, &mut state.ctrl_plane);
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  // ldaxr x27, [x25]
@@ -1760,8 +1760,8 @@ impl MachInstEmit for Inst {
  ));
  sink.use_label_at_offset(br_out_offset, out_label, LabelUse::Branch19);
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_stlxr(ty, x24wr, x28, x25)); // stlxr w24, x28, [x25]
@@ -1789,8 +1789,8 @@ impl MachInstEmit for Inst {
  let rn = allocs.next(rn);
  let rt = allocs.next_writable(rt);
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_ldar(access_ty, rt, rn));
@@ -1804,8 +1804,8 @@ impl MachInstEmit for Inst {
  let rn = allocs.next(rn);
  let rt = allocs.next(rt);
 
- if !flags.notrap() {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_stlr(access_ty, rt, rn));
@@ -2977,9 +2977,9 @@ impl MachInstEmit for Inst {
  let rn = allocs.next(rn);
  let (q, size) = size.enc_size();
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  sink.put4(enc_ldst_vec(q, size, rn, rd));

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

@@ -630,9 +630,9 @@ impl Inst {
  _ => return None,
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
  sink.put2(encode_ci_sp_load(op, rd, imm6));
  }
@@ -701,9 +701,9 @@ impl Inst {
  encode_cl_type(op, rd, base, imm5)
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
  sink.put2(encoded);
  }
@@ -732,9 +732,9 @@ impl Inst {
  _ => return None,
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
  sink.put2(encode_css_type(op, src, imm6));
  }
@@ -799,9 +799,9 @@ impl Inst {
  encode_cs_type(op, src, base, imm5)
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
  sink.put2(encoded);
  }
@@ -1077,9 +1077,9 @@ impl Inst {
  }
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  sink.put4(encode_i_type(op.op_code(), rd, op.funct3(), addr, imm12));
@@ -1100,9 +1100,9 @@ impl Inst {
  }
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  sink.put4(encode_s_type(op.op_code(), op.funct3(), addr, src, imm12));
@@ -1492,7 +1492,11 @@ impl Inst {
  src,
  amo,
  } => {
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ // TODO: get flags from original CLIF atomic instruction
+ let flags = MemFlags::new();
+ if let Some(trap_code) = flags.trap_code() {
+ sink.add_trap(trap_code);
+ }
  let x = op.op_code()
  | reg_to_gpr_num(rd.to_reg()) << 7
  | op.funct3() << 12
@@ -2718,9 +2722,9 @@ impl Inst {
  }
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  sink.put4(encode_vmem_load(
@@ -2765,9 +2769,9 @@ impl Inst {
  }
  };
 
- if !flags.notrap() {
+ if let Some(trap_code) = flags.trap_code() {
  // Register the offset at which the actual load instruction starts.
- sink.add_trap(TrapCode::HeapOutOfBounds);
+ sink.add_trap(trap_code);
  }
 
  sink.put4(encode_vmem_store(

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

@@ -102,10 +102,6 @@ impl MemArg {
  }
  }
 
- pub(crate) fn can_trap(&self) -> bool {
- !self.get_flags().notrap()
- }
-
  /// Edit registers with allocations.
  pub fn with_allocs(&self, allocs: &mut AllocationConsumer<'_>) -> Self {
  match self {
@@ -187,10 +183,6 @@ impl MemArgPair {
  }
  }
 
- pub(crate) fn can_trap(&self) -> bool {
- !self.flags.notrap()
- }
-
  /// Edit registers with allocations.
  pub fn with_allocs(&self, allocs: &mut AllocationConsumer<'_>) -> Self {
  MemArgPair {