Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 31 additions & 22 deletions cuda-oxide/crates/cuda-macros/src/cuda_module/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -842,23 +842,24 @@ fn cuda_module_path_description(module_path: &[Ident]) -> String {
/// no reference at all. The backend also keeps a weak legacy alias for older
/// macro expansions in mixed-version builds.
///
/// The reference is only emitted when the module is guaranteed to produce
/// an artifact for this crate. Generic kernels are monomorphized (and
/// their PTX embedded) in the *consuming* crate, so a module with only
/// generic kernels yields no artifact here, and an anchor reference would
/// be an undefined-symbol link error. The same reasoning extends to
/// cfg-gated kernels: root `load()` emits one equivalent guarded reference per
/// concrete kernel in the complete inline tree. Each reference carries the
/// kernel's effective ancestor-plus-local availability attributes, so a module
/// containing only nested kernels is still independently loadable while no
/// anchor is referenced when every concrete kernel is absent.
/// The reference is emitted for every kernel that remains after cfg, generic
/// or concrete. A concrete kernel always embeds an artifact in this crate. A
/// generic kernel *may* also embed one here when this crate monomorphizes it
/// (issue #1365); when it does not, the codegen backend still supplies an
/// anchor-only host stub whenever PTX-merge markers are present, so the
/// reference stays defined for the #222 shape (mono only in a consumer).
/// Root `load()` emits one equivalent guarded reference per kernel in the
/// complete inline tree. Each reference carries the kernel's effective
/// ancestor-plus-local availability attributes, so a module containing only
/// nested kernels is still independently loadable while no anchor is
/// referenced when every kernel is cfg'd out.
///
/// The anchor's address also tells `load_named()` where the bundle is: the
/// linker keeps the anchor and the `.oxart` section in the same binary, so
/// the loader reads that binary rather than the process executable, which
/// differs whenever the module was compiled into a shared object. Each
/// reference records the address in `__cuda_oxide_artifact_anchor`; it stays
/// `None` when no reference is emitted or every concrete kernel is cfg'd out.
/// `None` when no reference is emitted or every kernel is cfg'd out.
fn cuda_module_artifact_anchor_statements(
kernels: &[CudaModuleKernel],
) -> syn::Result<TokenStream2> {
Expand All @@ -873,8 +874,9 @@ fn cuda_module_artifact_anchor_statements(
})
}

/// One guarded anchor reference per concrete kernel, or none when this crate
/// produces no artifact; see [`cuda_module_artifact_anchor_statements`].
/// One guarded anchor reference per kernel (generic or concrete), or none when
/// this crate is filtered out / not cargo-built; see
/// [`cuda_module_artifact_anchor_statements`].
fn cuda_module_artifact_anchor_references(
kernels: &[CudaModuleKernel],
) -> syn::Result<Vec<TokenStream2>> {
Expand All @@ -899,10 +901,6 @@ fn cuda_module_artifact_anchor_references(
return Ok(Vec::new());
}

if !kernels.iter().any(|kernel| !kernel.is_generic) {
return Ok(Vec::new());
}

let binary_name = std::env::var("CARGO_BIN_NAME").ok();
let anchor = if owner_selection.is_some() {
artifact_anchor_symbol_v2(
Expand All @@ -914,10 +912,22 @@ fn cuda_module_artifact_anchor_references(
} else {
artifact_anchor_symbol(&package_name, &package_version)
};
let anchor_name = LitStr::new(&anchor, proc_macro2::Span::call_site());
let references = kernels
Ok(artifact_anchor_reference_tokens(kernels, &anchor))
}

/// The guarded anchor references themselves, one per kernel, for an already
/// resolved anchor symbol. Split out from the env lookup above so the token
/// shape can be unit tested without cargo's per-crate environment.
pub(crate) fn artifact_anchor_reference_tokens(
kernels: &[CudaModuleKernel],
anchor: &str,
) -> Vec<TokenStream2> {
let anchor_name = LitStr::new(anchor, proc_macro2::Span::call_site());
// Emit for every kernel, including generics. Generic modules still need
// the #72 rlib keep-alive when this crate monomorphizes (issue #1365);
// when it does not, the backend's anchor-only stub keeps the link valid.
kernels
.iter()
.filter(|kernel| !kernel.is_generic)
.map(|kernel| {
let cfg_attrs = &kernel.effective_cfg_attrs;
quote! {
Expand All @@ -938,8 +948,7 @@ fn cuda_module_artifact_anchor_references(
};
}
})
.collect();
Ok(references)
.collect()
}

pub(crate) fn device_codegen_owner_selection(raw: Option<&str>, crate_name: &str) -> Option<bool> {
Expand Down
60 changes: 58 additions & 2 deletions cuda-oxide/crates/cuda-macros/src/tests/cuda_module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@

use crate::cuda_module::contract::LaunchContractArgs;
use crate::cuda_module::{
device_codegen_owner_selection, expand_cuda_module, expand_cuda_module_inner,
transform_cuda_module_items,
artifact_anchor_reference_tokens, device_codegen_owner_selection, expand_cuda_module,
expand_cuda_module_inner, transform_cuda_module_items,
};
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
Expand Down Expand Up @@ -386,6 +386,62 @@ fn cfg_gated_generic_uses_the_same_gate_for_marker_and_loader() {
);
}

/// Anchor references for every kernel in `module`, rendered compactly with a
/// fixed anchor symbol. Unit tests run without cargo's per-crate env, so the
/// expansion itself never emits references here; this exercises the token
/// builder that the expansion uses once the env is present.
fn anchor_references_compact(module: ItemMod) -> String {
let items = &module.content.expect("inline module").1;
let transformed =
transform_cuda_module_items(items, &mut Vec::new(), &[], false, true).unwrap();
let references = artifact_anchor_reference_tokens(&transformed.kernels, "test_anchor");
quote!(#(#references)*).to_string().replace(' ', "")
}

#[test]
fn all_generic_cuda_module_keeps_anchor_keepalive_and_merge_loader() {
let module: ItemMod = parse_quote! {
mod kernels {
#[kernel]
pub fn fill<T: Copy>(value: T) {}
}
};
let expanded = expand_to_compact_string(module.clone());

assert!(
expanded.contains("::cuda_host::load_all_ptx_bundles_merged(ctx)?"),
"generic modules must keep the #222 merge loader:\n{expanded}"
);
assert!(
expanded.contains("#[allow(unused_mut)]letmut__cuda_oxide_artifact_anchor:::core::option::Option<&::core::primitive::u8>=::core::option::Option::None;"),
"load_named must still initialize the optional anchor address:\n{expanded}"
);

let references = anchor_references_compact(module);
assert!(
references.contains("CUDA_OXIDE_BUNDLE_ANCHOR")
&& references.contains("black_box")
&& references.contains("#[link_name=\"test_anchor\"]"),
"all-generic modules must emit the #72/#1365 artifact-anchor keep-alive:\n{references}"
);
}

#[test]
fn cfg_gated_generic_anchor_keepalive_inherits_effective_cfg() {
let module: ItemMod = parse_quote! {
mod kernels {
#[cfg(feature = "generic")]
#[kernel]
pub fn fill<T: Copy>(value: T) {}
}
};
let references = anchor_references_compact(module);
assert!(
references.starts_with("#[cfg(feature=\"generic\")]let_={unsafeextern\"C\"{"),
"generic anchor keep-alive must inherit the kernel cfg:\n{references}"
);
}

/// The embedded-artifact loader: read the bundle from the binary that maps
/// the artifact anchor, or from the executable when no anchor was referenced.
const ANCHORED_EMBEDDED_LOADER: &str = concat!(
Expand Down
Loading