majit: #1287 review follow-ups, an inline-trial geometry diagnostic, and one wasm cost removal

majit/majit-backend-wasm/src/codegen.rs

@@ -1035,16 +1035,38 @@ impl LabelResumeData {
         self.capture_by_id.get(r.raw() as usize).copied().flatten()
     }
 
+    fn shortage(&self, frame: FrameGeometry) -> Option<super::FrameShortage> {
+        if self.ref_slots > frame.label_ref_slots {
+            return Some(super::FrameShortage::new(
+                super::FrameShortageKind::LabelResumeRefSlots,
+                self.ref_slots,
+                frame.label_ref_slots,
+            ));
+        }
+        for storage in self.capture_by_id.iter().flatten() {
+            match storage {
+                LabelCaptureStorage::ValueSlot(slot) if *slot >= frame.value_slots => {
+                    return Some(super::FrameShortage::new(
+                        super::FrameShortageKind::LabelResumeCaptureSlots,
+                        slot + 1,
+                        frame.value_slots,
+                    ));
+                }
+                LabelCaptureStorage::RefSlot(slot) if *slot >= frame.label_ref_slots => {
+                    return Some(super::FrameShortage::new(
+                        super::FrameShortageKind::LabelResumeCaptureSlots,
+                        slot + 1,
+                        frame.label_ref_slots,
+                    ));
+                }
+                LabelCaptureStorage::ValueSlot(_) | LabelCaptureStorage::RefSlot(_) => {}
+            }
+        }
+        None
+    }
+
     fn supported_by(&self, frame: FrameGeometry) -> bool {
-        self.ref_slots <= frame.label_ref_slots
-            && self
-                .capture_by_id
-                .iter()
-                .flatten()
-                .all(|storage| match storage {
-                    LabelCaptureStorage::ValueSlot(slot) => *slot < frame.value_slots,
-                    LabelCaptureStorage::RefSlot(slot) => *slot < frame.label_ref_slots,
-                })
+        self.shortage(frame).is_none()
     }
 
     fn frame_offset(&self, storage: LabelCaptureStorage, frame: FrameGeometry) -> u64 {
@@ -2365,13 +2387,17 @@ pub fn build_wasm_module(
     let max_value_slots =
         normal_frame_value_slots(&analysis_inputargs, &analysis_ops) + label_resume.scalar_slots;
     if max_value_slots > frame.value_slots {
+        let shortage = super::FrameShortage::new(
+            super::FrameShortageKind::FrameValueSlots,
+            max_value_slots,
+            frame.value_slots,
+        );
         if !inlined_bridges.is_empty() {
-            super::record_inline_geometry(max_value_slots, frame.value_slots);
+            super::record_inline_geometry(shortage.kind, shortage.needed, shortage.available);
         }
         return Err(BackendError::Unsupported(format!(
-            "wasm backend: {max_value_slots} frame value slots exceed frozen frame layout \
-             ({})",
-            frame.value_slots,
+            "wasm backend: {} frame value slots exceed frozen frame layout ({})",
+            shortage.needed, shortage.available,
         )));
     }
 
@@ -2398,16 +2424,37 @@ pub fn build_wasm_module(
         &label_resume.captured_refs,
     );
     let num_ref_homes = ref_homes.len();
-    if num_ref_homes > frame.ordinary_home_slots() || !label_resume.supported_by(*frame) {
+    let shortage = if num_ref_homes > frame.ordinary_home_slots() {
+        Some(super::FrameShortage::new(
+            super::FrameShortageKind::OrdinaryRefHomes,
+            num_ref_homes,
+            frame.ordinary_home_slots(),
+        ))
+    } else {
+        label_resume.shortage(*frame)
+    };
+    if let Some(shortage) = shortage {
         if !inlined_bridges.is_empty() {
-            super::record_inline_geometry(num_ref_homes, frame.ordinary_home_slots());
+            super::record_inline_geometry(shortage.kind, shortage.needed, shortage.available);
         }
-        return Err(BackendError::Unsupported(format!(
-            "wasm backend: {num_ref_homes} ordinary ref homes and {} LABEL ref captures exceed frozen frame layout ({}, {})",
-            label_resume.ref_slots,
-            frame.ordinary_home_slots(),
-            frame.label_ref_slots,
-        )));
+        let reason = match shortage.kind {
+            super::FrameShortageKind::OrdinaryRefHomes => format!(
+                "wasm backend: {} ordinary ref homes exceed frozen frame layout ({})",
+                shortage.needed, shortage.available,
+            ),
+            super::FrameShortageKind::LabelResumeRefSlots => format!(
+                "wasm backend: {} LABEL ref captures exceed label resume layout ({} label ref slots)",
+                shortage.needed, shortage.available,
+            ),
+            super::FrameShortageKind::LabelResumeCaptureSlots => format!(
+                "wasm backend: {} LABEL capture slots exceed label resume layout ({})",
+                shortage.needed, shortage.available,
+            ),
+            super::FrameShortageKind::FrameValueSlots => {
+                unreachable!("value-slot shortage was checked above")
+            }
+        };
+        return Err(BackendError::Unsupported(reason));
     }
 
     // Self-recursive CALL_ASSEMBLER arm (`PYRE_WASM_CA`): `bridge_finish_fi` is

majit/majit-backend-wasm/src/lib.rs

@@ -83,31 +83,69 @@ use std::sync::{Arc, Mutex};
 /// because the source module has frame-only dispatch; 46 = parameter entry
 /// declined because the source guard and bridge input arities disagree; 47 =
 /// LABEL publication suppressed because the bridge entry has nonzero parameters.
-pub static BRIDGE_DIAG: [AtomicU64; 48] = [const { AtomicU64::new(0) }; 48];
+/// 48 = an inline trial's LABEL-resume storage exceeds the frozen frame.
+pub static BRIDGE_DIAG: [AtomicU64; 49] = [const { AtomicU64::new(0) }; 49];
 
-/// The first three inline geometry failures, packed as `(needed, available)`.
-/// They expose a frozen-layout shortage without changing the compile result.
+#[repr(u8)]
+#[derive(Clone, Copy)]
+pub(crate) enum FrameShortageKind {
+    FrameValueSlots = 1,
+    OrdinaryRefHomes = 2,
+    LabelResumeRefSlots = 3,
+    LabelResumeCaptureSlots = 4,
+}
+
+#[derive(Clone, Copy)]
+pub(crate) struct FrameShortage {
+    pub(crate) kind: FrameShortageKind,
+    pub(crate) needed: usize,
+    pub(crate) available: usize,
+}
+
+impl FrameShortage {
+    pub(crate) const fn new(kind: FrameShortageKind, needed: usize, available: usize) -> Self {
+        Self {
+            kind,
+            needed,
+            available,
+        }
+    }
+}
+
+/// The first three inline geometry failures, packed as
+/// `(kind: u8, needed: u24, available: u24)`. They expose a frozen-layout
+/// shortage without changing the compile result.
 static INLINE_GEOMETRY: [AtomicU64; 3] = [const { AtomicU64::new(0) }; 3];
 static INLINE_GEOMETRY_COUNT: AtomicU64 = AtomicU64::new(0);
 static INLINE_TRIAL_ERRORS: Mutex<Vec<String>> = Mutex::new(Vec::new());
 
-pub(crate) fn record_inline_geometry(needed: usize, available: usize) {
+pub(crate) fn record_inline_geometry(kind: FrameShortageKind, needed: usize, available: usize) {
+    const FIELD_MASK: u64 = (1 << 24) - 1;
+
     let index = INLINE_GEOMETRY_COUNT.fetch_add(1, Ordering::Relaxed) as usize;
     if let Some(slot) = INLINE_GEOMETRY.get(index) {
         slot.store(
-            ((needed as u64) << 32) | available as u64,
+            ((kind as u64) << 48)
+                | ((needed as u64).min(FIELD_MASK) << 24)
+                | (available as u64).min(FIELD_MASK),
             Ordering::Relaxed,
         );
     }
 }
 
-/// Read a packed `(needed, available)` inline geometry failure.
+/// Read a packed `(kind, needed, available)` inline geometry failure.
 pub fn inline_geometry_diag(index: usize) -> u64 {
     INLINE_GEOMETRY
         .get(index)
         .map_or(0, |slot| slot.load(Ordering::Relaxed))
 }
 
+/// Number of inline geometry failures, including records beyond the three
+/// diagnostics retained in [`INLINE_GEOMETRY`].
+pub fn inline_geometry_count() -> u64 {
+    INLINE_GEOMETRY_COUNT.load(Ordering::Relaxed)
+}
+
 pub fn inline_trial_errors() -> String {
     INLINE_TRIAL_ERRORS.lock().unwrap().join(" | ")
 }
@@ -228,11 +266,9 @@ fn reemit_enabled() -> bool {
     REEMIT_ENABLED.load(Ordering::Relaxed)
 }
 
-/// Arm loop-closing bridge inlining. Inlining rebuilds the owning loop, so it
-/// also enables the replacement path.
+/// Arm loop-closing bridge inlining from the host before guest execution starts.
 pub fn inline_bridge_enable() {
     INLINE_BRIDGE_ENABLED.store(true, Ordering::Relaxed);
-    reemit_enable();
 }
 
 fn inline_bridge_enabled() -> bool {
@@ -3301,6 +3337,8 @@ impl majit_backend::Backend for WasmBackend {
                                 diag_bump(40);
                             } else if reason.contains("ordinary ref homes") {
                                 diag_bump(41);
+                            } else if reason.contains("label resume layout") {
+                                diag_bump(48);
                             } else if reason
                                 .contains("inlined bridge stream has no local loop LABEL")
                             {
@@ -3606,10 +3644,9 @@ impl majit_backend::Backend for WasmBackend {
             unsafe {
                 core::ptr::write(cell, bridge_slot);
             }
-            // Only retained module replacement needs to restore this cell
-            // after allocating a fresh dispatch array.  Without replacement,
-            // the live cell is already the sole dispatch state.
-            if is_direct && reemit_enabled() {
+            // Retained module replacement and loop-closing bridge inlining
+            // restore this cell after allocating a fresh dispatch array.
+            if is_direct && (reemit_enabled() || inline_bridge_enabled()) {
                 if let Some(source_loop) = original_token
                     .compiled
                     .get()

majit/majit-backend/src/lib.rs

@@ -1304,12 +1304,13 @@ pub struct JitCellToken {
     /// (`LoopTargetDescr` Arc; `TargetToken IS-A AbstractDescr` in
     /// PyPy, so a `DescrRef` is the matching identity).  Each
     /// successful loop / retrace populates this through
-    /// `record_target_token`.  Its one reader is
-    /// [`Self::first_target_token`], the descr a bridge closes onto:
-    /// neither `has_compiled_loop` (token presence) nor pyre's
-    /// `has_compiled_targets` (the `compiled_loops` side table) reads
-    /// this list, so it is not pyre's `has_compiled_targets` signal
-    /// despite mirroring what upstream's reads.  The metainterp-side
+    /// `record_target_token`.  Its one reader with a caller is
+    /// [`Self::first_target_token`], the descr a bridge closes onto;
+    /// [`Self::has_target_tokens`] reads it too but is called from
+    /// nowhere.  Neither `has_compiled_loop` (token presence) nor
+    /// pyre's `has_compiled_targets` (the `compiled_loops` side table)
+    /// reads this list, so it is not pyre's `has_compiled_targets`
+    /// signal despite mirroring what upstream's reads.  The metainterp-side
     /// `TargetToken` value (with `virtual_state` / `short_preamble`)
     /// stays on the `CompiledEntry::front_target_tokens` list per
     /// the F.6 retirement plan — the per-target descr identity is the

majit/majit-metainterp/src/optimizeopt/mod.rs

@@ -4347,6 +4347,17 @@ impl OptContext {
         self.active_short_preamble_producer.as_mut()
     }
 
+    /// Address of the active producer's storage for the MetaInterp root
+    /// walker. This field has the same
+    /// `Option<ExtendedShortPreambleBuilder>` type as
+    /// `Optimizer.short_preamble_producer`, so either address is valid for
+    /// the walker's cast while the builder is moved between them.
+    pub(crate) fn active_short_preamble_producer_slot_addr(&mut self) -> usize {
+        (&mut self.active_short_preamble_producer
+            as *mut Option<crate::optimizeopt::shortpreamble::ExtendedShortPreambleBuilder>)
+            as usize
+    }
+
     pub fn build_active_short_preamble(
         &self,
     ) -> Option<crate::optimizeopt::shortpreamble::ShortPreamble> {

majit/majit-metainterp/src/optimizeopt/optimizer.rs

@@ -408,6 +408,11 @@ pub struct Optimizer {
     /// the extended builder's home without merging those stages.
     pub short_preamble_producer:
         Option<crate::optimizeopt::shortpreamble::ExtendedShortPreambleBuilder>,
+    /// MetaInterp's `Option<usize>` publication slot for this producer. While
+    /// the builder is on loan to OptContext, the slot is re-pointed at
+    /// `OptContext.active_short_preamble_producer` so the root walker always
+    /// follows the builder's current home.
+    pub(crate) published_short_preamble_producer_slot: Option<usize>,
     /// RPython unroll.py: `label_args = import_state(...)`.
     /// The peeled loop's LABEL must use these args, not the phase-1 end_args.
     pub imported_label_args: Option<Vec<OpRef>>,
@@ -1490,6 +1495,7 @@ impl Optimizer {
             imported_short_preamble: None,
             imported_short_preamble_builder: None,
             short_preamble_producer: None,
+            published_short_preamble_producer_slot: None,
             imported_label_args: None,
             patchguardop: None,
             skip_flush: false,