From f2bdfa3fa5aa7ebfb6132bd64e14203f32a3c594 Mon Sep 17 00:00:00 2001 From: Yunus Mutlu Date: Tue, 29 Sep 2026 22:44:37 +0300 Subject: [PATCH] fix(mir-importer): preserve emitted diverging device calls Signed-off-by: Yunus Mutlu --- .../crates/mir-importer/src/pipeline.rs | 13 +++-- .../mir-importer/src/translator/facts.rs | 21 +++++++ .../src/translator/terminator/helpers.rs | 57 ++++++++++++------- .../src/translator/terminator/mod.rs | 46 +++++++++------ .../examples/function_item_call/README.md | 10 ++++ .../examples/function_item_call/src/main.rs | 25 +++++++- .../function_item_call/verify-code-shape.sh | 50 ++++++++++++++++ .../rustc-codegen-cuda/src/collector.rs | 14 ++--- 8 files changed, 184 insertions(+), 52 deletions(-) create mode 100644 cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh diff --git a/cuda-oxide/crates/mir-importer/src/pipeline.rs b/cuda-oxide/crates/mir-importer/src/pipeline.rs index 5d88f4a234..d405375f40 100644 --- a/cuda-oxide/crates/mir-importer/src/pipeline.rs +++ b/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 { - // 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)?; diff --git a/cuda-oxide/crates/mir-importer/src/translator/facts.rs b/cuda-oxide/crates/mir-importer/src/translator/facts.rs index 69302b9241..6e1a262f06 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/facts.rs +++ b/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> = + const { std::cell::RefCell::new(std::collections::BTreeSet::new()) }; } /// Installs the driver-resolved lang-item ids for this pipeline run, @@ -99,6 +105,21 @@ 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) { + 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 { + let identity = instance.mangled_name().to_string(); + EMITTED_FUNCTIONS.with(|set| set.borrow().contains(&identity)) +} + /// 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. diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs index a85b623374..74a09706e3 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/helpers.rs +++ b/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,11 +428,20 @@ 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, @@ -435,14 +449,13 @@ pub fn emit_function_call( 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 diff --git a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs index 15321134d3..a7f76a5b10 100644 --- a/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs +++ b/cuda-oxide/crates/mir-importer/src/translator/terminator/mod.rs @@ -1254,25 +1254,37 @@ 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. It does not + // mean an emitted callee can be erased: a collected Rust `-> !` function + // may perform observable work before it diverges. Preserve that call and + // let `emit_function_call` terminate the caller with `mir.unreachable`. // - // 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. + // For every diverging callee that is not part of this pipeline run's + // emitted-function set, retain the pre-existing trap fallback. Keep this + // before the unsupported-intrinsic and libm checks below so non-emitted + // diverging calls still take the same path as before this fix. 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 diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md index 62efd401c6..26bd973de6 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/README.md @@ -7,3 +7,13 @@ 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` pins that contract in the pre-optimization LLVM IR, +before the middle-end can inline the diverging helper. The divergent branches +are not launched; the host run exercises only the returning path. diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs index 68b442f644..c878dbbfc9 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/src/main.rs @@ -10,8 +10,9 @@ //! 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 executes +//! before its MIR block ends with `unreachable`. //! `call_once_decoy` proves that an ordinary function whose name contains //! trait-like text is not mistaken for a callable-trait shim. @@ -57,6 +58,18 @@ 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 {} +} + #[device] fn apply_once u32>(f: F, x: u32) -> u32 { f(x) @@ -148,6 +161,14 @@ 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 let Some(slot) = out.get_mut(idx) { *slot = 0xd1ce; } diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh new file mode 100644 index 0000000000..f231d6d102 --- /dev/null +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/function_item_call/verify-code-shape.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +root="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +llvm="${root}/function_item_call.ll" + +test -s "${llvm}" + +extract_llvm_function() { + local needle="$1" + awk -v needle="${needle}" ' + /^define / && index($0, needle) != 0 { in_function = 1 } + in_function { print } + in_function && /^}/ { exit } + ' "${llvm}" +} + +kernel="$(extract_llvm_function '@diverging_callable(')" +test -n "${kernel}" + +call_count="$( + grep -Ec 'call void @[^ (]*write_then_diverge\(' <<<"${kernel}" || true +)" +if [[ "${call_count}" -ne 1 ]]; then + echo "diverging_callable must contain exactly one direct call to write_then_diverge" >&2 + exit 1 +fi + +call_line="$( + grep -nE 'call void @[^ (]*write_then_diverge\(' <<<"${kernel}" | + head -n1 | + cut -d: -f1 +)" +next_line="$(sed -n "$((call_line + 1))p" <<<"${kernel}")" +if ! grep -Eq '^[[:space:]]*unreachable[[:space:]]*$' <<<"${next_line}"; then + echo "direct diverging call must be immediately followed by LLVM unreachable" >&2 + exit 1 +fi + +helper="$(extract_llvm_function 'write_then_diverge')" +test -n "${helper}" +if ! grep -Eq 'store i32 286331153, ptr ' <<<"${helper}"; then + echo "write_then_diverge lost its observable sentinel store" >&2 + exit 1 +fi + +echo "function_item_call code shape: PASS" diff --git a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs index 0b2bb713ab..a83078e2cf 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/src/collector.rs @@ -2464,10 +2464,9 @@ impl<'tcx> DeviceCollector<'tcx> { /// "Symbol ... not found" error. /// /// Functions that are declared diverging (return type `!`) are left - /// alone: their call sites have no target block, so the translator - /// already lowers them to LLVM `unreachable` (note: NOT a trap; - /// the optimizer may delete paths that provably reach it, a known - /// gap tracked separately). + /// alone. If their body is emitted, the importer preserves the call and + /// terminates its caller with `unreachable`; if collection deliberately + /// omits the body, the importer retains the existing device-trap fallback. fn check_unreachable_callee( &self, resolved: Instance<'tcx>, @@ -2547,8 +2546,9 @@ impl<'tcx> DeviceCollector<'tcx> { // No stub marker: leave non-local functions to the existing silent // skip (this is what keeps the real `cuda_device` intrinsic - // placeholders and `core`'s cold panic wrappers working), and leave - // declared-diverging functions to the translator's `unreachable` lowering. + // placeholders and `core`'s cold panic wrappers working). Declared- + // diverging functions that are omitted here retain the translator's + // device-trap fallback rather than producing an undefined call. if !def_id.is_local() || mir.return_ty().is_never() { return; } @@ -2577,7 +2577,7 @@ impl<'tcx> DeviceCollector<'tcx> { .with_help( "annotate device helpers with `#[device]`; if the panic is \ intentional, declare the function as diverging (`-> !`) so the \ - call lowers to LLVM `unreachable`", + device pipeline can treat the path as non-returning", ) .emit() }