From a146283b2683eef3479db5b7bc7a4813d856fc42 Mon Sep 17 00:00:00 2001 From: Mikhail Borisov Date: Sat, 10 Oct 2026 12:48:05 +0500 Subject: [PATCH] fix(cuda-macros): obey #[inline] on generic device functions `#[device]` added `#[inline(never)]` before the other attributes of a generic device function. rustc uses only the first `inline` attribute. Because of this, rustc did not use an `#[inline(always)]` or `#[inline]` attribute on the function. Add `#[inline(never)]` only if there is no `inline` attribute on the function. `#[kernel]` uses the same check for the attributes of generic kernels. rustc expands `cfg_attr` before it calls the macro. Because of this, the check also finds an `inline` attribute in `cfg_attr`. Move the expansion to `device_function_tokens`, as for generic kernels, so that unit tests can examine it. Fixes #1471 Signed-off-by: Mikhail Borisov Co-Authored-By: Claude Opus 5.5 --- cuda-oxide/crates/cuda-macros/src/device.rs | 39 ++++++---- .../crates/cuda-macros/src/tests/device.rs | 75 +++++++++++++++++++ .../crates/cuda-macros/src/tests/mod.rs | 1 + .../crates/cuda-macros/tests/macro_guard.rs | 1 + .../tests/pass/device_inline_policy.rs | 22 ++++++ .../examples/standalone_device_fn/src/main.rs | 2 +- 6 files changed, 126 insertions(+), 14 deletions(-) create mode 100644 cuda-oxide/crates/cuda-macros/src/tests/device.rs create mode 100644 cuda-oxide/crates/cuda-macros/tests/pass/device_inline_policy.rs diff --git a/cuda-oxide/crates/cuda-macros/src/device.rs b/cuda-oxide/crates/cuda-macros/src/device.rs index bda0e02b18..5e98bd9d4e 100644 --- a/cuda-oxide/crates/cuda-macros/src/device.rs +++ b/cuda-oxide/crates/cuda-macros/src/device.rs @@ -5,12 +5,16 @@ //! `#[device]` functions and extern blocks, plus `#[constant]` statics. -use crate::common::{impl_trait_parameter_error, reject_reserved_name, track_codegen_environment}; +use crate::common::{ + attr_path_ends_with, impl_trait_parameter_error, reject_reserved_name, + track_codegen_environment, +}; use crate::cuda_module::constants::extract_constant_inner_ty; use crate::cuda_module::launchers::{codegen_generic_arguments, has_codegen_generics}; use crate::kernel::scope::{forwarding_inputs, inject_device_thread_index_scope}; use crate::launch_attrs::rewrite_loop_unroll_attrs; use proc_macro::TokenStream; +use proc_macro2::TokenStream as TokenStream2; use quote::{format_ident, quote}; use reserved_oxide_symbols::{DEVICE_EXTERN_PREFIX, DEVICE_PREFIX, constant_symbol}; use syn::{ @@ -115,7 +119,7 @@ pub(crate) fn device_entry(_attr: TokenStream, item: TokenStream) -> TokenStream /// Handles both non-generic and generic device functions: /// - **Non-generic**: `#[no_mangle]` on the prefixed function, `#[inline(always)]` wrapper. /// - **Generic**: No `#[no_mangle]` (generics use mangled symbols), `#[inline(never)]` on -/// the prefixed function (so monomorphizations appear in CGUs for the collector), +/// the prefixed function by default (so monomorphizations appear in CGUs for the collector), /// `#[inline(always)]` wrapper with generics + turbofish forwarding. /// /// This mirrors the pattern used by `#[kernel]` for generic kernels @@ -131,7 +135,12 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { return err.to_compile_error().into(); } inject_device_thread_index_scope(&mut input); + device_function_tokens(input).into() +} +/// Expansion body of [`generate_device_function`], split out so unit tests +/// can inspect the generated items without a live proc-macro bridge. +pub(crate) fn device_function_tokens(mut input: ItemFn) -> TokenStream2 { let fn_name = input.sig.ident.clone(); let vis = input.vis.clone(); let new_name = format_ident!("{}{}", DEVICE_PREFIX, fn_name); @@ -148,7 +157,7 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { let (wrapper_inputs, params) = match forwarding_inputs(&input.sig.inputs) { Ok(forwarding) => forwarding, - Err(err) => return err.to_compile_error().into(), + Err(err) => return err.to_compile_error(), }; // Rename the original function with the prefix @@ -161,7 +170,7 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { // monomorphization (e.g., `cuda_oxide_device__add::` gets a // unique mangled name). #[no_mangle] requires a single concrete symbol. // - // - #[inline(never)] on the prefixed function — ensures each monomorphization + // - #[inline(never)] on the prefixed function by default — ensures each monomorphization // appears as a distinct CGU item so the collector can find it. If it were // inlined, the function would disappear from the CGU. // @@ -177,8 +186,16 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { call }; - let expanded = quote! { - #[inline(never)] + // A user `inline` attribute replaces the default. rustc expands + // `cfg_attr` before it calls this macro. + let has_inline_attr = input + .attrs + .iter() + .any(|attr| attr_path_ends_with(attr, "inline")); + let default_inline = (!has_inline_attr).then(|| quote! { #[inline(never)] }); + + quote! { + #default_inline #input /// Wrapper for the generic device function with the original name. @@ -186,9 +203,7 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { #vis #constness #unsafety #abi fn #fn_name #generics (#(#wrapper_inputs),*) #return_type #where_clause { #call } - }; - - TokenStream::from(expanded) + } } else { let call = quote! { #new_name(#(#params),*) }; let call = if unsafety.is_some() { @@ -197,7 +212,7 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { call }; // Non-generic device function: simple case. - let expanded = quote! { + quote! { #[unsafe(no_mangle)] #input @@ -206,9 +221,7 @@ fn generate_device_function(mut input: ItemFn) -> TokenStream { #vis #constness #unsafety #abi fn #fn_name #generics (#(#wrapper_inputs),*) #return_type #where_clause { #call } - }; - - TokenStream::from(expanded) + } } } diff --git a/cuda-oxide/crates/cuda-macros/src/tests/device.rs b/cuda-oxide/crates/cuda-macros/src/tests/device.rs new file mode 100644 index 0000000000..54c089c42e --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/src/tests/device.rs @@ -0,0 +1,75 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +use crate::device::device_function_tokens; +use quote::{format_ident, quote}; +use reserved_oxide_symbols::DEVICE_PREFIX; +use syn::{Attribute, ItemFn, parse_quote}; + +/// The `inline` attributes of the renamed function in the expansion of +/// `function`. +fn renamed_inline_attrs(function: ItemFn) -> Vec { + let name = format_ident!("{}{}", DEVICE_PREFIX, function.sig.ident); + let expanded = device_function_tokens(function); + let file: syn::File = + syn::parse2(expanded.clone()).expect("generated tokens must parse as items"); + let renamed = file + .items + .iter() + .find_map(|item| match item { + syn::Item::Fn(item_fn) if item_fn.sig.ident == name => Some(item_fn), + _ => None, + }) + .unwrap_or_else(|| panic!("expansion has no fn `{name}`:\n{expanded}")); + renamed + .attrs + .iter() + .filter(|attr| attr.path().is_ident("inline")) + .map(|attr| quote!(#attr).to_string().replace(' ', "")) + .collect() +} + +#[test] +fn generic_device_function_keeps_the_inline_attribute_of_the_user() { + let attrs: [Attribute; 3] = [ + parse_quote!(#[inline]), + parse_quote!(#[inline(always)]), + parse_quote!(#[inline(never)]), + ]; + for attr in attrs { + let expected = quote!(#attr).to_string().replace(' ', ""); + let function: ItemFn = parse_quote! { + #attr + fn scale(value: T) -> T { value } + }; + assert_eq!(renamed_inline_attrs(function), [expected]); + } +} + +#[test] +fn generic_device_function_is_inline_never_by_default() { + let function: ItemFn = parse_quote! { + fn scale(value: T) -> T { value } + }; + assert_eq!(renamed_inline_attrs(function), ["#[inline(never)]"]); +} + +#[test] +fn other_attributes_do_not_replace_the_default() { + let function: ItemFn = parse_quote! { + #[cold] + #[doc = "inline"] + fn scale(value: T) -> T { value } + }; + assert_eq!(renamed_inline_attrs(function), ["#[inline(never)]"]); +} + +#[test] +fn non_generic_device_function_gets_no_inline_attribute() { + let function: ItemFn = parse_quote! { + fn scale(value: f32) -> f32 { value } + }; + assert!(renamed_inline_attrs(function).is_empty()); +} diff --git a/cuda-oxide/crates/cuda-macros/src/tests/mod.rs b/cuda-oxide/crates/cuda-macros/src/tests/mod.rs index 8e7851e3b5..27801d5790 100644 --- a/cuda-oxide/crates/cuda-macros/src/tests/mod.rs +++ b/cuda-oxide/crates/cuda-macros/src/tests/mod.rs @@ -4,5 +4,6 @@ */ mod cuda_module; +mod device; mod kernel; mod launch; diff --git a/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs b/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs index 46683828d3..a4f3993870 100644 --- a/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs +++ b/cuda-oxide/crates/cuda-macros/tests/macro_guard.rs @@ -10,6 +10,7 @@ fn macro_guards() { let t = trybuild::TestCases::new(); t.pass("tests/pass/const_generic_hygiene.rs"); t.pass("tests/pass/cuda_module_inline_namespaces.rs"); + t.pass("tests/pass/device_inline_policy.rs"); t.compile_fail("tests/compile_fail/kernel_reserved_name.rs"); t.compile_fail("tests/compile_fail/device_reserved_name.rs"); t.compile_fail("tests/compile_fail/device_extern_reserved_name.rs"); diff --git a/cuda-oxide/crates/cuda-macros/tests/pass/device_inline_policy.rs b/cuda-oxide/crates/cuda-macros/tests/pass/device_inline_policy.rs new file mode 100644 index 0000000000..8338bb86fb --- /dev/null +++ b/cuda-oxide/crates/cuda-macros/tests/pass/device_inline_policy.rs @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// rustc gives `unused_attributes` for a second `inline` attribute. +#![deny(unused_attributes)] +#![allow(dead_code)] + +use cuda_macros::device; + +#[device] +#[inline(always)] +fn always(value: T) -> T { + value +} + +#[device] +#[cfg_attr(all(), inline)] +fn conditional(value: T) -> T { + value +} + +fn main() {} diff --git a/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/src/main.rs b/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/src/main.rs index 6fec7faf30..9484d86470 100644 --- a/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/src/main.rs +++ b/cuda-oxide/crates/rustc-codegen-cuda/examples/standalone_device_fn/src/main.rs @@ -79,7 +79,7 @@ pub fn safe_sqrt(x: f32) -> f32 { // // #[device] now supports generics, mirroring how #[kernel] handles them: // - No #[no_mangle] (generics use mangled symbol names) -// - #[inline(never)] on the prefixed function (so monomorphizations appear in CGUs) +// - #[inline(never)] on the prefixed function by default (so monomorphizations appear in CGUs) // - Wrapper forwards type params via turbofish // // Monomorphization is triggered by concrete call sites (other #[device] fns