fix(mir-importer): preserve emitted diverging device calls

cuda-oxide/crates/mir-importer/src/pipeline.rs

@@ -494,11 +494,16 @@ pub fn run_pipeline(
     config: &PipelineConfig,
     known_defs: crate::translator::facts::KnownDefs,
 ) -> Result<CompilationResult, PipelineError> {
-    // Install the driver-resolved lang-item ids for this run. Set (not
-    // merged) every entry: the ids are only valid inside the caller's
-    // `rustc_internal::run` context, so a stale set from a previous run on
-    // this thread must never survive.
+    // Install the driver-resolved lang-item ids and emitted-function identities
+    // for this run. Replace (rather than merge) both thread-local sets so
+    // stale state from a previous run on this thread cannot survive.
     crate::translator::facts::set_known_defs(known_defs);
+    crate::translator::facts::set_emitted_functions(
+        functions
+            .iter()
+            .map(|func| func.instance.mangled_name().to_string())
+            .collect(),
+    );
 
     prepare_output_dir(&config.output_dir)?;
 

cuda-oxide/crates/mir-importer/src/translator/facts.rs

@@ -85,6 +85,12 @@ thread_local! {
             index_mut_trait: None,
         })
     };
+    // Exact monomorphized Rust functions that this pipeline run will emit.
+    //
+    // Store rustc's linkage identity rather than the legalised module symbol:
+    // call classification must not depend on Legaliser ordering or spelling.
+    static EMITTED_FUNCTIONS: std::cell::RefCell<std::collections::BTreeSet<String>> =
+        const { std::cell::RefCell::new(std::collections::BTreeSet::new()) };
 }
 
 /// Installs the driver-resolved lang-item ids for this pipeline run,
@@ -99,6 +105,28 @@ pub(crate) fn known_defs() -> KnownDefs {
     KNOWN_DEFS.with(|cell| cell.get())
 }
 
+/// Replace the emitted-function set for the current pipeline run.
+///
+/// Like [`set_known_defs`], this replaces rather than extends previous
+/// thread-local state so one rustc invocation cannot leak into the next.
+pub(crate) fn set_emitted_functions(names: std::collections::BTreeSet<String>) {
+    EMITTED_FUNCTIONS.with(|set| *set.borrow_mut() = names);
+}
+
+/// Whether this exact monomorphized Rust instance will have a `mir.func`
+/// definition in the module being translated.
+pub(crate) fn is_emitted_function(instance: &mir::mono::Instance) -> bool {
+    is_emitted_name(&instance.mangled_name())
+}
+
+/// Whether `mangled_name`, rustc's linkage identity of one monomorphized
+/// function, is in the current run's emitted-function set. The name-based
+/// form lets the set's replace-not-merge contract be tested without a rustc
+/// session.
+pub(crate) fn is_emitted_name(mangled_name: &str) -> bool {
+    EMITTED_FUNCTIONS.with(|set| set.borrow().contains(mangled_name))
+}
+
 /// Validate the source reference and the fully specialized read-only storage
 /// contract using rustc's type system. MIR locals have erased regions, so the
 /// outer lifetime is read from the declaration before instantiation instead.
@@ -1104,4 +1132,27 @@ pub fn reference_validity(
             assert!(fact.pointee_alignment.is_power_of_two());
         }
     }
+
+    /// `set_emitted_functions` replaces the previous run's set instead of
+    /// merging into it. A name left over from an earlier rustc invocation on
+    /// the same thread would keep a diverging call to a function this run
+    /// never emits, and the module would reference an undefined symbol.
+    #[test]
+    fn emitted_function_set_is_replaced_per_run() {
+        let first_run = std::collections::BTreeSet::from(["first_run_die".to_string()]);
+        set_emitted_functions(first_run);
+        assert!(is_emitted_name("first_run_die"));
+        assert!(!is_emitted_name("second_run_spin"));
+
+        let second_run = std::collections::BTreeSet::from(["second_run_spin".to_string()]);
+        set_emitted_functions(second_run);
+        assert!(is_emitted_name("second_run_spin"));
+        assert!(
+            !is_emitted_name("first_run_die"),
+            "an entry from the previous run must not survive the replacement"
+        );
+
+        set_emitted_functions(std::collections::BTreeSet::new());
+        assert!(!is_emitted_name("second_run_spin"));
+    }
 }

cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs

@@ -316,10 +316,10 @@ pub fn bundle_generated_u32_results_as_array(
 /// 1. Translate all MIR arguments to Pliron IR values
 /// 2. At a foreign ABI boundary, adapt pointer address spaces to the exact
 ///    declared parameter types without changing pointer kind or mutability
-/// 3. Prepare the call destination using the current projection state
+/// 3. For returning calls, prepare the destination using the current projection state
 /// 4. Create a `mir.call` operation carrying the callee's name attribute
-/// 5. Store the result through the prepared destination
-/// 6. Emit a zero-operand goto to the call's success target
+/// 5. For returning calls, store the result and goto the success target
+/// 6. For diverging calls, emit `mir.unreachable` after the call
 ///
 /// Reference arguments (`&mut local`) are handed the local's alloca slot
 /// pointer directly, so callee writes through the reference are observed by
@@ -383,16 +383,21 @@ pub fn emit_function_call(
         }
     }
 
-    let (prepared_destination, prepared_last_op) = prepare_destination_write(
-        ctx,
-        body,
-        destination,
-        value_map,
-        block_ptr,
-        last_op,
-        loc.clone(),
-    )?;
-    last_op = prepared_last_op;
+    let prepared_destination = if target.is_some() {
+        let (prepared, prepared_last_op) = prepare_destination_write(
+            ctx,
+            body,
+            destination,
+            value_map,
+            block_ptr,
+            last_op,
+            loc.clone(),
+        )?;
+        last_op = prepared_last_op;
+        Some(prepared)
+    } else {
+        None
+    };
 
     use pliron::builtin::attributes::StringAttr;
 
@@ -423,26 +428,34 @@ pub fn emit_function_call(
         call_op
     };
 
+    if target.is_none() {
+        return Ok(super::emit_unreachable_after(
+            ctx,
+            block_ptr,
+            Some(call_op),
+            loc,
+        ));
+    }
+
     let result_value = call_op.deref(ctx).get_result(0);
 
     let goto_prev = finish_destination_write(
         ctx,
-        prepared_destination,
+        prepared_destination.expect("returning call prepared its destination"),
         result_value,
         value_map,
         block_ptr,
         call_op,
         loc.clone(),
     )?;
 
-    if let Some(target_idx) = target {
-        Ok(emit_goto(ctx, *target_idx, goto_prev, block_map, loc))
-    } else {
-        input_err!(
-            loc.clone(),
-            TranslationErr::unsupported("Call terminator without target not supported".to_string(),)
-        )
-    }
+    Ok(emit_goto(
+        ctx,
+        target.expect("returning call has a success target"),
+        goto_prev,
+        block_map,
+        loc,
+    ))
 }
 
 /// Adapt a foreign-call pointer argument only when the address space is the

cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs

@@ -1254,25 +1254,47 @@ fn translate_call(
         return Ok(result);
     }
 
-    // Handle diverging calls (calls that never return, like unwrap_failed, panic, etc.)
-    // These have no target block because the function never returns.
+    // `Call { target: None }` means there is no normal successor. rustc
+    // clears the target when the destination type is visibly uninhabited:
+    // `!`, `core::convert::Infallible`, an empty enum. That does not mean the
+    // callee can be erased. An emitted Rust function with such a return type
+    // may do observable work before it diverges. When the resolved instance
+    // is in this pipeline run's emitted-function set, keep the call and let
+    // `emit_function_call` end the block with `mir.unreachable`, which is
+    // what rustc's own backend emits for a call without a destination.
     //
-    // The callee is dropped and replaced by an immediate trap. That applies
-    // to every `-> !` callee reaching this point, including a user
-    // `#[device]` fn that legitimately diverges (e.g. `loop {}`); full Rust
-    // fidelity would emit the call for resolvable callees the way
-    // `translate_function_item_call` does. Dropping the call is a
-    // pre-existing semantic: the trap only makes it safe, where the bare
-    // `unreachable` previously emitted here was UB that let `opt` delete
-    // the whole panic path.
-    //
-    // Panic entry points additionally never get here with any statements
-    // translated ahead of them: `block::translate_block` recognizes the same
-    // shape via [`is_dropped_panic_call`] and emits the trap directly.
+    // Every other diverging call keeps the pre-existing trap fallback. That
+    // covers callees the collector skipped, such as core helpers without MIR
+    // and local `-> !` bodies that are nothing but a panic, and a call
+    // through a function pointer, whose `func` operand is not a constant and
+    // so resolves no instance. Panic entry points never reach this arm: for
+    // a callee whose path matches `is_panic_entry_path`,
+    // `block::translate_block` lowers the whole block to the trap through
+    // `is_dropped_panic_call` before any call is classified. Keep this arm
+    // before the unsupported-intrinsic and libm checks below so the trapped
+    // cases take the same path as before.
     if target_usize.is_none() {
-        // This is a diverging call (returns !) - emit trap + unreachable
-        // Examples: unwrap_failed(), panic!(), abort()
-        return Ok(emit_trap_unreachable_after(ctx, block_ptr, prev_op, loc));
+        let resolved_instance = match func {
+            mir::Operand::Constant(constant)
+                if *constant.const_.kind() == ConstantKind::ZeroSized =>
+            {
+                match constant.const_.ty().kind() {
+                    rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::FnDef(
+                        fn_def,
+                        substs,
+                    )) => rustc_public::mir::mono::Instance::resolve(fn_def, &substs).ok(),
+                    _ => None,
+                }
+            }
+            _ => None,
+        };
+
+        if !resolved_instance
+            .as_ref()
+            .is_some_and(crate::translator::facts::is_emitted_function)
+        {
+            return Ok(emit_trap_unreachable_after(ctx, block_ptr, prev_op, loc));
+        }
     }
 
     // A call to a rustc intrinsic that no dispatch arm above recognized can

cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md

@@ -7,3 +7,20 @@ Rust lowers `apply_once(plus_seven, x)` through a callable trait method such as
 function-item body, and MIR import must emit a direct call name that matches the
 collector/export naming policy. Otherwise the generated device IR references a
 callee symbol that was never emitted.
+
+The diverging-call regression also covers an ordinary direct `#[device]`
+function returning `!`. Its body performs an observable store before looping
+forever. A MIR `Call` with no normal successor must therefore keep the call and
+end the caller block with `unreachable`; replacing the call itself with a trap
+would erase the store.
+
+`verify-code-shape.sh` checks that contract in the pre-optimization LLVM IR,
+before the middle-end can inline the diverging helper. It covers the three
+call sites in `diverging_callable` that used to lower to a trap: two through
+`apply_never` (a function item and a closure) and the direct call to
+`write_then_diverge`. Each kept call must be followed by `unreachable`. When
+the build also wrote PTX, the kernel entry must keep the sentinel store or a
+call to the helper. The control is `panic_only`, a `-> !` function whose body
+is nothing but a panic: the collector does not emit it, so its call site must
+still lower to the device trap, in the IR and in the PTX. The divergent
+branches are not launched; the host run exercises only the returning path.

cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs

@@ -10,8 +10,11 @@
 //! back to the concrete function body instead of emitting a dangling trait-shim
 //! callee symbol. The regression also covers `Fn`/`FnMut`, references,
 //! multiple and nested-tuple arguments, tuple returns, and type/const-generic
-//! function items. Diverging function items and closures also verify that a
-//! direct Rust call returning `!` ends its MIR block with `unreachable`.
+//! function items. Diverging function items, closures, and an ordinary direct
+//! `-> !` device helper also verify that a non-returning Rust call is emitted
+//! before its MIR block ends with `unreachable`. `panic_only` is the control:
+//! its `-> !` body is nothing but a panic, so the collector never emits it and
+//! its call site must still lower to the device trap.
 //! `call_once_decoy` proves that an ordinary function whose name contains
 //! trait-like text is not mistaken for a callable-trait shim.
 
@@ -57,6 +60,30 @@ fn spin_forever(_value: u32) -> ! {
     loop {}
 }
 
+// Unlike `spin_forever`, this is reached by an ordinary direct MIR `Call`.
+// The store makes replacing that call with an immediate trap observably wrong.
+#[allow(clippy::empty_loop)]
+#[inline(never)]
+#[device]
+unsafe fn write_then_diverge(ptr: *mut u32) -> ! {
+    unsafe {
+        ptr.write(0x1111_1111);
+    }
+    loop {}
+}
+
+// Control for the retained trap. A `-> !` body that is nothing but a panic is
+// a placeholder to the collector, so this function is never emitted and its
+// call site must keep the device trap. `#[inline(never)]` keeps the call in
+// the kernel's MIR, so the importer classifies this callee rather than the
+// panic entry point inside it. The name must not contain a `panicking` module
+// segment: that would trap the whole block before the call is classified.
+#[inline(never)]
+#[device]
+fn panic_only() -> ! {
+    panic!("not collected")
+}
+
 #[device]
 fn apply_once<F: FnOnce(u32) -> u32>(f: F, x: u32) -> u32 {
     f(x)
@@ -133,9 +160,10 @@ mod kernels {
         }
     }
 
-    /// Compile both diverging callable paths without entering them on-device.
-    /// A Rust call returning `!` has no successor block, so the importer must
-    /// emit the direct call followed by `mir.unreachable`.
+    /// Compile the diverging call paths without entering them on-device.
+    /// A Rust call returning `!` has no successor block. When the callee's
+    /// body is emitted, the importer must keep the direct call followed by
+    /// `mir.unreachable`; when it is not, the call site keeps the device trap.
     #[kernel]
     pub fn diverging_callable(flag: u32, mut out: DisjointSlice<u32>) {
         let idx = thread::index_1d();
@@ -148,6 +176,19 @@ mod kernels {
             #[allow(clippy::empty_loop)]
             apply_never(|_| loop {}, flag);
         }
+        if flag == 3 {
+            // The importer must preserve this ordinary direct call. The callee
+            // stores first and then diverges, so replacing it with a trap loses
+            // an observable side effect.
+            unsafe {
+                write_then_diverge(out.as_mut_ptr());
+            }
+        }
+        if flag == 4 {
+            // Retained-trap control: the callee's body is only a panic, so the
+            // collector does not emit it and this call site must still trap.
+            panic_only();
+        }
         if let Some(slot) = out.get_mut(idx) {
             *slot = 0xd1ce;
         }