diff --git a/cuda-oxide/crates/llvm-export/src/export/function.rs b/cuda-oxide/crates/llvm-export/src/export/function.rs index 3fa26fa0ae..eb16ce73af 100644 --- a/cuda-oxide/crates/llvm-export/src/export/function.rs +++ b/cuda-oxide/crates/llvm-export/src/export/function.rs @@ -224,9 +224,19 @@ impl<'a> ModuleExportState<'a> { } else { "global" }; + // Host-written storage (`#[constant]`): the initializer is only the + // load-time image. Without this, whole-program LTO (libNVVM + // `-gen-lto` + nvJitLink `-lto`) internalizes the never-stored + // global, folds its loads to the initializer and drops the symbol + // the host looks up with `cuModuleGetGlobal`. + let externally_initialized = if global.is_host_written(self.ctx) { + "externally_initialized " + } else { + "" + }; write!( output, - "@{name} = addrspace({address_space}) {storage_keyword} " + "@{name} = addrspace({address_space}) {externally_initialized}{storage_keyword} " ) .unwrap(); self.export_type(ty, output)?; diff --git a/cuda-oxide/crates/llvm-export/src/lib.rs b/cuda-oxide/crates/llvm-export/src/lib.rs index 9f891f829e..fd582c514e 100644 --- a/cuda-oxide/crates/llvm-export/src/lib.rs +++ b/cuda-oxide/crates/llvm-export/src/lib.rs @@ -428,6 +428,11 @@ pub mod ops { /// reference it. The exporter adds marked globals to `@llvm.used`, keeping /// profiler metadata alive through libNVVM and nvJitLink materialization. const GLOBAL_RETAINED_KEY: &str = "cuda_oxide_global_retained"; + /// Marks a `GlobalOp` whose contents the host writes after module load + /// (`cuModuleGetGlobal` + copy). The exporter renders it + /// `externally_initialized`, so no optimizer, including libNVVM/nvJitLink + /// LTO, may treat the initializer as the value device code reads. + const GLOBAL_HOST_WRITTEN_KEY: &str = "cuda_oxide_global_host_written"; /// One pointer-width relocation inside an evaluated Rust static initializer. /// @@ -2564,6 +2569,10 @@ pub mod ops { fn mark_retained(&self, ctx: &mut Context); /// Whether this global was explicitly marked as externally consumed. fn is_retained(&self, ctx: &Context) -> bool; + /// Declare that the host writes this global after module load. + fn mark_host_written(&self, ctx: &mut Context); + /// Whether the host writes this global after module load. + fn is_host_written(&self, ctx: &Context) -> bool; } impl GlobalOpExt for GlobalOp { @@ -2685,6 +2694,25 @@ pub mod ops { .get::(&key) .is_some() } + + fn mark_host_written(&self, ctx: &mut Context) { + let key = + Identifier::try_new(GLOBAL_HOST_WRITTEN_KEY.to_string()).expect("valid identifier"); + self.get_operation() + .deref_mut(ctx) + .attributes + .set(key, pliron::builtin::attributes::UnitAttr); + } + + fn is_host_written(&self, ctx: &Context) -> bool { + let key = + Identifier::try_new(GLOBAL_HOST_WRITTEN_KEY.to_string()).expect("valid identifier"); + self.get_operation() + .deref(ctx) + .attributes + .get::(&key) + .is_some() + } } #[cfg(test)] diff --git a/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs b/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs index e6490122cf..062e44950a 100644 --- a/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs +++ b/cuda-oxide/crates/llvm-export/tests/export_test/kernels_and_roots.rs @@ -109,6 +109,38 @@ fn legacy_llvm_used_roots_retained_address_space_global() { ); } +/// Host-written storage must not look like a known constant to LTO, in either +/// IR dialect; ordinary globals keep their plain definition. +#[test] +fn host_written_global_exports_externally_initialized() { + let mut ctx = Context::new(); + let module = ModuleOp::new(&mut ctx, "host_written".try_into().unwrap()); + let module_block = module_top_block(&mut ctx, &module); + let i32_ty = IntegerType::get(&ctx, 32, Signedness::Signless); + + let params = GlobalOp::new(&mut ctx, "PARAMS".try_into().unwrap(), i32_ty.into()); + params.set_address_space(&mut ctx, 4); + params.mark_host_written(&mut ctx); + params.get_operation().insert_at_back(module_block, &ctx); + + let ordinary = GlobalOp::new(&mut ctx, "ORDINARY".try_into().unwrap(), i32_ty.into()); + ordinary.set_address_space(&mut ctx, 1); + ordinary.get_operation().insert_at_back(module_block, &ctx); + + let legacy = NvvmExportConfig::new(NvvmIrDialect::LegacyLlvm7); + for ir in [ + export_module_to_string_with_config(&ctx, &module, &PtxExportConfig), + export_module_to_string_with_config(&ctx, &module, &legacy), + ] { + let ir = ir.expect("host-written global export succeeds"); + assert!( + ir.contains("@PARAMS = addrspace(4) externally_initialized global i32"), + "{ir}" + ); + assert!(ir.contains("@ORDINARY = addrspace(1) global i32"), "{ir}"); + } +} + #[test] fn ptx_export_records_kernel_roots_for_internalization() { let mut ctx = Context::new(); diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs index 4374953b83..33e7fa85b8 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/device_global.rs @@ -257,6 +257,13 @@ fn create_device_global( }; global_op.set_address_space(ctx, spec.addr_space); global_op.set_source_global_key(ctx, spec.key); + if spec.addr_space == llvm_export::types::address_space::CONSTANT { + // `#[constant]` storage is written by the host after module load. Keep + // the symbol rooted (`@llvm.used`) even when no kernel reads it, and + // stop optimizers from treating its initializer as the value read. + global_op.mark_retained(ctx); + global_op.mark_host_written(ctx); + } if matches!( spec.addr_space, llvm_export::types::address_space::GLOBAL | llvm_export::types::address_space::CONSTANT diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs index 562db062e5..b245546638 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/memory/tests/device_global.rs @@ -200,6 +200,12 @@ fn convert_global_alloc_places_in_global_or_constant_addrspace() { .starts_with("__device_global_"), "ordinary device globals get the __device_global_ prefix" ); + // The host writes constant memory after load: it stays rooted and must not + // be folded to its initializer. Private device globals stay ordinary. + assert!(global_addr_const.is_retained(&ctx)); + assert!(global_addr_const.is_host_written(&ctx)); + assert!(!global_addr_global.is_retained(&ctx)); + assert!(!global_addr_global.is_host_written(&ctx)); } #[test]