Cranelift: Take user stack maps through lowering and emission

cranelift/codegen/src/ir/mod.rs

@@ -65,7 +65,7 @@ pub use crate::ir::stackslot::{
 };
 pub use crate::ir::trapcode::TrapCode;
 pub use crate::ir::types::Type;
-pub use crate::ir::user_stack_maps::UserStackMapEntry;
+pub use crate::ir::user_stack_maps::{UserStackMap, UserStackMapEntry};
 
 use crate::entity::{entity_impl, PrimaryMap, SecondaryMap};
 

cranelift/codegen/src/ir/user_stack_maps.rs

@@ -29,13 +29,19 @@
 //! contrast to the old system and its `r64` values).
 
 use crate::ir;
+use cranelift_bitset::CompoundBitSet;
+use cranelift_entity::PrimaryMap;
 use smallvec::SmallVec;
 
 pub(crate) type UserStackMapEntryVec = SmallVec<[UserStackMapEntry; 4]>;
 
-/// A stack map entry describes a GC-managed value and its location at a
-/// particular instruction.
-#[derive(Clone, PartialEq, Hash)]
+/// A stack map entry describes a single GC-managed value and its location on
+/// the stack.
+///
+/// A stack map entry is associated with a particular instruction, and that
+/// instruction must be a safepoint. The GC-managed value must be stored in the
+/// described location across this entry's instruction.
+#[derive(Clone, Debug, PartialEq, Hash)]
 #[cfg_attr(
     feature = "enable-serde",
     derive(serde_derive::Serialize, serde_derive::Deserialize)
@@ -50,3 +56,46 @@ pub struct UserStackMapEntry {
     /// The offset within the stack slot where this entry's value can be found.
     pub offset: u32,
 }
+
+/// A compiled stack map, describing the location of many GC-managed values.
+///
+/// A stack map is associated with a particular instruction, and that
+/// instruction is a safepoint.
+#[derive(Clone, Debug, PartialEq)]
+#[cfg_attr(
+    feature = "enable-serde",
+    derive(serde_derive::Deserialize, serde_derive::Serialize)
+)]
+pub struct UserStackMap {
+    by_type: SmallVec<[(ir::Type, CompoundBitSet); 1]>,
+}
+
+impl UserStackMap {
+    /// Coalesce the given entries into a new `UserStackMap`.
+    pub fn new(
+        entries: &[UserStackMapEntry],
+        stack_slot_offsets: &PrimaryMap<ir::StackSlot, u32>,
+    ) -> Self {
+        let mut by_type = SmallVec::<[(ir::Type, CompoundBitSet); 1]>::default();
+
+        for entry in entries {
+            let offset = stack_slot_offsets[entry.slot] + entry.offset;
+            let offset = usize::try_from(offset).unwrap();
+
+            // Don't bother trying to avoid an `O(n)` search here: `n` is
+            // basically always one in practice; even if it isn't, there aren't
+            // that many different CLIF types.
+            let index = by_type
+                .iter()
+                .position(|(ty, _)| *ty == entry.ty)
+                .unwrap_or_else(|| {
+                    by_type.push((entry.ty, CompoundBitSet::with_capacity(offset + 1)));
+                    by_type.len() - 1
+                });
+
+            by_type[index].1.insert(offset);
+        }
+
+        UserStackMap { by_type }
+    }
+}

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

@@ -651,25 +651,38 @@ fn enc_asimd_mod_imm(rd: Writable<Reg>, q_op: u32, cmode: u32, imm: u8) -> u32 {
 /// State carried between emissions of a sequence of instructions.
 #[derive(Default, Clone, Debug)]
 pub struct EmitState {
-    /// Safepoint stack map for upcoming instruction, as provided to `pre_safepoint()`.
+    /// Safepoint stack map for upcoming instruction, as provided to
+    /// `pre_safepoint()`.
     stack_map: Option<StackMap>,
+
+    /// The user stack map for the upcoming instruction, as provided to
+    /// `pre_safepoint()`.
+    user_stack_map: Option<ir::UserStackMap>,
+
     /// Only used during fuzz-testing. Otherwise, it is a zero-sized struct and
     /// optimized away at compiletime. See [cranelift_control].
     ctrl_plane: ControlPlane,
+
     frame_layout: FrameLayout,
 }
 
 impl MachInstEmitState<Inst> for EmitState {
     fn new(abi: &Callee<AArch64MachineDeps>, ctrl_plane: ControlPlane) -> Self {
         EmitState {
             stack_map: None,
+            user_stack_map: None,
             ctrl_plane,
             frame_layout: abi.frame_layout().clone(),
         }
     }
 
-    fn pre_safepoint(&mut self, stack_map: StackMap) {
-        self.stack_map = Some(stack_map);
+    fn pre_safepoint(
+        &mut self,
+        stack_map: Option<StackMap>,
+        user_stack_map: Option<ir::UserStackMap>,
+    ) {
+        self.stack_map = stack_map;
+        self.user_stack_map = user_stack_map;
     }
 
     fn ctrl_plane_mut(&mut self) -> &mut ControlPlane {
@@ -686,8 +699,8 @@ impl MachInstEmitState<Inst> for EmitState {
 }
 
 impl EmitState {
-    fn take_stack_map(&mut self) -> Option<StackMap> {
-        self.stack_map.take()
+    fn take_stack_map(&mut self) -> (Option<StackMap>, Option<ir::UserStackMap>) {
+        (self.stack_map.take(), self.user_stack_map.take())
     }
 
     fn clear_post_insn(&mut self) {
@@ -2921,11 +2934,16 @@ impl MachInstEmit for Inst {
                 }
             }
             &Inst::Call { ref info } => {
-                if let Some(s) = state.take_stack_map() {
+                let (stack_map, user_stack_map) = state.take_stack_map();
+                if let Some(s) = stack_map {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(4), s);
                 }
                 sink.add_reloc(Reloc::Arm64Call, &info.dest, 0);
                 sink.put4(enc_jump26(0b100101, 0));
+                if let Some(s) = user_stack_map {
+                    let offset = sink.cur_offset();
+                    sink.push_user_stack_map(state, offset, s);
+                }
                 if info.opcode.is_call() {
                     sink.add_call_site(info.opcode);
                 }
@@ -2939,11 +2957,16 @@ impl MachInstEmit for Inst {
                 }
             }
             &Inst::CallInd { ref info } => {
-                if let Some(s) = state.take_stack_map() {
+                let (stack_map, user_stack_map) = state.take_stack_map();
+                if let Some(s) = stack_map {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(4), s);
                 }
                 let rn = info.rn;
                 sink.put4(0b1101011_0001_11111_000000_00000_00000 | (machreg_to_gpr(rn) << 5));
+                if let Some(s) = user_stack_map {
+                    let offset = sink.cur_offset();
+                    sink.push_user_stack_map(state, offset, s);
+                }
                 if info.opcode.is_call() {
                     sink.add_call_site(info.opcode);
                 }

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

@@ -46,20 +46,28 @@ pub enum EmitVState {
 /// State carried between emissions of a sequence of instructions.
 #[derive(Default, Clone, Debug)]
 pub struct EmitState {
-    /// Safepoint stack map for upcoming instruction, as provided to `pre_safepoint()`.
+    /// Safepoint stack map for upcoming instruction, as provided to
+    /// `pre_safepoint()`.
     stack_map: Option<StackMap>,
+
+    /// The user stack map for the upcoming instruction, as provided to
+    /// `pre_safepoint()`.
+    user_stack_map: Option<ir::UserStackMap>,
+
     /// Only used during fuzz-testing. Otherwise, it is a zero-sized struct and
     /// optimized away at compiletime. See [cranelift_control].
     ctrl_plane: ControlPlane,
+
     /// Vector State
     /// Controls the current state of the vector unit at the emission point.
     vstate: EmitVState,
+
     frame_layout: FrameLayout,
 }
 
 impl EmitState {
-    fn take_stack_map(&mut self) -> Option<StackMap> {
-        self.stack_map.take()
+    fn take_stack_map(&mut self) -> (Option<StackMap>, Option<ir::UserStackMap>) {
+        (self.stack_map.take(), self.user_stack_map.take())
     }
 }
 
@@ -70,14 +78,20 @@ impl MachInstEmitState<Inst> for EmitState {
     ) -> Self {
         EmitState {
             stack_map: None,
+            user_stack_map: None,
             ctrl_plane,
             vstate: EmitVState::Unknown,
             frame_layout: abi.frame_layout().clone(),
         }
     }
 
-    fn pre_safepoint(&mut self, stack_map: StackMap) {
-        self.stack_map = Some(stack_map);
+    fn pre_safepoint(
+        &mut self,
+        stack_map: Option<StackMap>,
+        user_stack_map: Option<ir::UserStackMap>,
+    ) {
+        self.stack_map = stack_map;
+        self.user_stack_map = user_stack_map;
     }
 
     fn ctrl_plane_mut(&mut self) -> &mut ControlPlane {
@@ -1127,13 +1141,21 @@ impl Inst {
                     sink.add_call_site(info.opcode);
                 }
                 sink.add_reloc(Reloc::RiscvCallPlt, &info.dest, 0);
-                if let Some(s) = state.take_stack_map() {
+
+                let (stack_map, user_stack_map) = state.take_stack_map();
+                if let Some(s) = stack_map {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(8), s);
                 }
+
                 Inst::construct_auipc_and_jalr(Some(writable_link_reg()), writable_link_reg(), 0)
                     .into_iter()
                     .for_each(|i| i.emit_uncompressed(sink, emit_info, state, start_off));
 
+                if let Some(s) = user_stack_map {
+                    let offset = sink.cur_offset();
+                    sink.push_user_stack_map(state, offset, s);
+                }
+
                 let callee_pop_size = i32::try_from(info.callee_pop_size).unwrap();
                 if callee_pop_size > 0 {
                     for inst in Riscv64MachineDeps::gen_sp_reg_adjust(-callee_pop_size) {
@@ -1151,9 +1173,14 @@ impl Inst {
                 }
                 .emit(sink, emit_info, state);
 
-                if let Some(s) = state.take_stack_map() {
+                let (stack_map, user_stack_map) = state.take_stack_map();
+                if let Some(s) = stack_map {
                     sink.add_stack_map(StackMapExtent::StartedAtOffset(start_offset), s);
                 }
+                if let Some(s) = user_stack_map {
+                    let offset = sink.cur_offset();
+                    sink.push_user_stack_map(state, offset, s);
+                }
 
                 if info.opcode.is_call() {
                     sink.add_call_site(info.opcode);

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

@@ -1,7 +1,7 @@
 //! S390x ISA: binary code emission.
 
 use crate::binemit::StackMap;
-use crate::ir::{MemFlags, TrapCode};
+use crate::ir::{self, MemFlags, TrapCode};
 use crate::isa::s390x::inst::*;
 use crate::isa::s390x::settings as s390x_settings;
 use cranelift_control::ControlPlane;
@@ -1306,11 +1306,19 @@ fn put_with_trap(sink: &mut MachBuffer<Inst>, enc: &[u8], trap_code: TrapCode) {
 #[derive(Default, Clone, Debug)]
 pub struct EmitState {
     pub(crate) initial_sp_offset: i64,
-    /// Safepoint stack map for upcoming instruction, as provided to `pre_safepoint()`.
+
+    /// Safepoint stack map for upcoming instruction, as provided to
+    /// `pre_safepoint()`.
     stack_map: Option<StackMap>,
+
+    /// The user stack map for the upcoming instruction, as provided to
+    /// `pre_safepoint()`.
+    user_stack_map: Option<ir::UserStackMap>,
+
     /// Only used during fuzz-testing. Otherwise, it is a zero-sized struct and
     /// optimized away at compiletime. See [cranelift_control].
     ctrl_plane: ControlPlane,
+
     frame_layout: FrameLayout,
 }
 
@@ -1319,13 +1327,19 @@ impl MachInstEmitState<Inst> for EmitState {
         EmitState {
             initial_sp_offset: abi.frame_size() as i64,
             stack_map: None,
+            user_stack_map: None,
             ctrl_plane,
             frame_layout: abi.frame_layout().clone(),
         }
     }
 
-    fn pre_safepoint(&mut self, stack_map: StackMap) {
-        self.stack_map = Some(stack_map);
+    fn pre_safepoint(
+        &mut self,
+        stack_map: Option<StackMap>,
+        user_stack_map: Option<ir::UserStackMap>,
+    ) {
+        self.stack_map = stack_map;
+        self.user_stack_map = user_stack_map;
     }
 
     fn ctrl_plane_mut(&mut self) -> &mut ControlPlane {
@@ -1342,8 +1356,8 @@ impl MachInstEmitState<Inst> for EmitState {
 }
 
 impl EmitState {
-    fn take_stack_map(&mut self) -> Option<StackMap> {
-        self.stack_map.take()
+    fn take_stack_map(&mut self) -> (Option<StackMap>, Option<ir::UserStackMap>) {
+        (self.stack_map.take(), self.user_stack_map.take())
     }
 
     fn clear_post_insn(&mut self) {
@@ -3243,9 +3257,15 @@ impl Inst {
                     _ => unreachable!(),
                 }
 
-                if let Some(s) = state.take_stack_map() {
+                let (stack_map, user_stack_map) = state.take_stack_map();
+                if let Some(s) = stack_map {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(6), s);
                 }
+                if let Some(s) = user_stack_map {
+                    let offset = sink.cur_offset() + 6;
+                    sink.push_user_stack_map(state, offset, s);
+                }
+
                 put(sink, &enc_ril_b(opcode, link.to_reg(), 0));
                 if info.opcode.is_call() {
                     sink.add_call_site(info.opcode);
@@ -3255,10 +3275,16 @@ impl Inst {
                 debug_assert_eq!(link.to_reg(), gpr(14));
                 let rn = info.rn;
 
-                let opcode = 0x0d; // BASR
-                if let Some(s) = state.take_stack_map() {
+                let (stack_map, user_stack_map) = state.take_stack_map();
+                if let Some(s) = stack_map {
                     sink.add_stack_map(StackMapExtent::UpcomingBytes(2), s);
                 }
+                if let Some(s) = user_stack_map {
+                    let offset = sink.cur_offset() + 2;
+                    sink.push_user_stack_map(state, offset, s);
+                }
+
+                let opcode = 0x0d; // BASR
                 put(sink, &enc_rr(opcode, link.to_reg(), rn));
                 if info.opcode.is_call() {
                     sink.add_call_site(info.opcode);

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

@@ -1599,9 +1599,15 @@ pub(crate) fn emit(
             opcode,
             info: call_info,
         } => {
-            if let Some(s) = state.take_stack_map() {
+            let (stack_map, user_stack_map) = state.take_stack_map();
+            if let Some(s) = stack_map {
                 sink.add_stack_map(StackMapExtent::UpcomingBytes(5), s);
             }
+            if let Some(s) = user_stack_map {
+                let offset = sink.cur_offset() + 5;
+                sink.push_user_stack_map(state, offset, s);
+            }
+
             sink.put1(0xE8);
             // The addend adjusts for the difference between the end of the instruction and the
             // beginning of the immediate field.
@@ -1696,9 +1702,16 @@ pub(crate) fn emit(
                     );
                 }
             }
-            if let Some(s) = state.take_stack_map() {
+
+            let (stack_map, user_stack_map) = state.take_stack_map();
+            if let Some(s) = stack_map {
                 sink.add_stack_map(StackMapExtent::StartedAtOffset(start_offset), s);
             }
+            if let Some(s) = user_stack_map {
+                let offset = sink.cur_offset();
+                sink.push_user_stack_map(state, offset, s);
+            }
+
             if opcode.is_call() {
                 sink.add_call_site(*opcode);
             }