From 3828b5fc24b78c93a8b6a8f039042b3e6ca2df78 Mon Sep 17 00:00:00 2001 From: nihalpasham Date: Tue, 27 Oct 2026 11:41:11 +0530 Subject: [PATCH] feat(cuda-oxide): use NVPTX_DATALAYOUT_FULL and --nvptx-short-ptr in PTX backend - Set `NVPTX_DATALAYOUT_PTX = NVPTX_DATALAYOUT_FULL` (`p3:32:32:32`) and pass `--nvptx-short-ptr` to `llc` in `base_llc_args`, aligning the native PTX backend with `nvcc` and `libNVVM` on 32-bit shared-memory pointers. - Extend `mir-lower` internal function return ABI coercion (`shared_return_abi_info`) so device functions returning a bare `ptr addrspace(3)` pass a 64-bit generic `ptr` across the `.func` return slot (`func_retval0` / `retval0`) and coerce back to `ptr addrspace(3)` at the call site, avoiding LLVM 23 `NVPTXAsmPrinter`'s `.b64` vs `.b32` return-slot mismatch under `--nvptx-short-ptr`. Signed-off-by: nihalpasham --- .../crates/cuda-oxide-codegen/src/ptx.rs | 2 + .../crates/llvm-export/src/export/config.rs | 7 ++- .../src/convert/ops/aggregate/enums.rs | 4 +- .../src/convert/ops/aggregate/fields.rs | 8 +-- .../crates/mir-lower/src/convert/ops/call.rs | 6 +- .../crates/mir-lower/src/convert/ops/cast.rs | 63 ++++++++++++++----- .../mir-lower/src/convert/ops/control_flow.rs | 15 ++++- .../mir-lower/src/convert/types/func_abi.rs | 33 +++++++++- .../crates/mir-lower/src/convert/types/mod.rs | 2 +- 9 files changed, 107 insertions(+), 33 deletions(-) diff --git a/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs b/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs index 8494b51982..34adae9b92 100644 --- a/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs +++ b/cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs @@ -333,6 +333,7 @@ fn base_llc_args(target: &str) -> Vec { vec![ "-march=nvptx64".to_string(), format!("-mcpu={target}"), + "--nvptx-short-ptr".to_string(), DISABLE_BRANCH_FOLD.to_string(), DISABLE_BLOCK_PLACEMENT.to_string(), ] @@ -1222,6 +1223,7 @@ mod tests { [ "-march=nvptx64", "-mcpu=sm_90", + "--nvptx-short-ptr", "-disable-branch-fold", "-disable-block-placement", ] diff --git a/cuda-oxide/crates/llvm-export/src/export/config.rs b/cuda-oxide/crates/llvm-export/src/export/config.rs index 8cd2e86c6e..69b3f41e87 100644 --- a/cuda-oxide/crates/llvm-export/src/export/config.rs +++ b/cuda-oxide/crates/llvm-export/src/export/config.rs @@ -5,8 +5,11 @@ //! Backend configuration traits and built-in implementations. -/// Minimal data layout for PTX mode (default behavior). -pub(super) const NVPTX_DATALAYOUT_PTX: &str = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"; +/// Data layout for PTX mode (default behavior). +/// +/// Matches `nvcc` and `libNVVM` (`NVPTX_DATALAYOUT_FULL`), declaring 32-bit +/// shared-memory pointers (`p3:32:32:32`) and 64-bit generic/global pointers. +pub(super) const NVPTX_DATALAYOUT_PTX: &str = NVPTX_DATALAYOUT_FULL; /// Full NVPTX data layout for libNVVM/LTOIR mode (Blackwell+, modern dialect). /// diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs index 94aced0846..1f4d7fed22 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/enums.rs @@ -1043,8 +1043,8 @@ mod tests { let body = kernel_blocks(&ctx, module); assert_eq!( count_ops::(&ctx, &body), - 2, - "construction must genericize the pointer and extraction must restore shared space" + 3, + "construction must genericize the pointer, extraction must restore shared space, and return must genericize for .func ABI" ); assert_eq!( count_ops::(&ctx, &body), diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs index e60a2878d1..356208f099 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs @@ -505,8 +505,8 @@ mod tests { assert_eq!( count_ops::(&ctx, &body), - 2, - "union insertion must genericize AS3 and extraction must restore AS3" + 3, + "union insertion must genericize AS3, extraction must restore AS3, and return must genericize for .func ABI" ); assert_eq!( count_ops::(&ctx, &body), @@ -610,8 +610,8 @@ mod tests { assert_eq!( count_ops::(&ctx, &body), - 2, - "nested union insertion must genericize AS3 and extraction must restore AS3" + 3, + "nested union insertion must genericize AS3, extraction must restore AS3, and return must genericize for .func ABI" ); assert_eq!(count_ops::(&ctx, &body), 0); assert_eq!(count_ops::(&ctx, &body), 0); diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs index bc19168b0c..c5ae9f0325 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/call.rs @@ -64,7 +64,7 @@ use crate::convert::target_stable_storage::coerce_target_stable_value; use crate::convert::types::{ StructLayoutInfo, TransparentScalarAbiInfo, build_struct_slot_map, convert_function_type, - convert_type, is_kernel_func, is_zero_sized_type, packed_shared_internal_abi_info, + convert_type, is_kernel_func, is_zero_sized_type, shared_return_abi_info, transparent_scalar_abi_info, }; use crate::helpers; @@ -680,9 +680,7 @@ pub fn convert( let scalar_ty = abi.scalar_ty; transparent_result_abi = Some(abi); scalar_ty - } else if let Some(abi) = - packed_shared_internal_abi_info(ctx, mir_ty).map_err(anyhow_to_pliron)? - { + } else if let Some(abi) = shared_return_abi_info(ctx, mir_ty).map_err(anyhow_to_pliron)? { let storage_ty = abi.storage_ty; packed_shared_result_abi = Some(abi); storage_ty diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs index 6d9aff6b90..90c8721126 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/cast.rs @@ -1523,11 +1523,19 @@ mod tests { ); // The exposed integer is a generic address, so re-entering the shared - // space must go inttoptr → generic, then addrspacecast → shared. A - // bare inttoptr into addrspace(3) would reinterpret the generic - // address as a space-local offset. - assert_cast_lowered_to::(&ctx, module, "llvm.addrspacecast"); + // space must go inttoptr -> generic, then addrspacecast -> shared, + // followed by the .func return-boundary addrspacecast back to generic. let body = kernel_blocks(&ctx, module); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "expected cast addrspacecast plus .func return-boundary addrspacecast" + ); + assert_eq!( + count_ops::(&ctx, &body), + 0, + "mir.cast must be replaced during lowering" + ); let int_to_ptrs = find_all::(&ctx, &body); let [int_to_ptr] = int_to_ptrs.as_slice() else { panic!("integer-to-shared transmute must lower to exactly one llvm.inttoptr"); @@ -1564,8 +1572,14 @@ mod tests { let module = lower_single_cast(&mut ctx, source, destination, MirCastKindAttr::Transmute); - assert_cast_lowered_to::(&ctx, module, "llvm.addrspacecast"); let body = kernel_blocks(&ctx, module); + let expected_casts = if shared_is_source { 1 } else { 2 }; + assert_eq!( + count_ops::(&ctx, &body), + expected_casts, + "expected cast addrspacecast (plus return-boundary cast when returning shared)" + ); + assert_eq!(count_ops::(&ctx, &body), 0); assert_eq!( count_ops::(&ctx, &body), 0, @@ -1946,17 +1960,28 @@ mod tests { // The exposed integer is a generic address, so re-entry must be // inttoptr-to-generic followed by addrspacecast into the space, - // mirroring the PointerExposeAddress direction. - assert_cast_lowered_to::(&ctx, module_ptr, "llvm.addrspacecast"); + // mirroring the PointerExposeAddress direction. When returning an + // addrspace(3) pointer, the .func return boundary adds a second + // addrspacecast back to generic. + let body = kernel_blocks(&ctx, module_ptr); + let expected_casts = if address_space == llvm_export::types::address_space::SHARED { + 2 + } else { + 1 + }; assert_eq!( - count_ops::(&ctx, &kernel_blocks(&ctx, module_ptr)), + count_ops::(&ctx, &body), + expected_casts, + "expected cast addrspacecast (plus return-boundary cast for shared)" + ); + assert_eq!(count_ops::(&ctx, &body), 0); + assert_eq!( + count_ops::(&ctx, &body), 1, "address space {address_space} must re-enter through a generic inttoptr" ); - let casts = find_all::(&ctx, &kernel_blocks(&ctx, module_ptr)); - let [cast] = casts.as_slice() else { - panic!("expected exactly one llvm.addrspacecast"); - }; + let casts = find_all::(&ctx, &body); + let cast = &casts[0]; let result_ty = cast .get_operation() .deref(&ctx) @@ -1984,12 +2009,16 @@ mod tests { MirCastKindAttr::PtrToPtr, ); - assert_cast_lowered_to::(&ctx, module_ptr, "llvm.addrspacecast"); + let body = kernel_blocks(&ctx, module_ptr); + assert_eq!( + count_ops::(&ctx, &body), + 2, + "ptr -> shared cast plus .func return-boundary cast must emit two llvm.addrspacecast ops" + ); + assert_eq!(count_ops::(&ctx, &body), 0); - let casts = find_all::(&ctx, &kernel_blocks(&ctx, module_ptr)); - let [cast] = casts.as_slice() else { - panic!("ptr -> ptr addrspace cast must lower to exactly one llvm.addrspacecast"); - }; + let casts = find_all::(&ctx, &body); + let cast = &casts[0]; let result_ty = cast .get_operation() .deref(&ctx) diff --git a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs index bf31f820cf..17ae18013a 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/ops/control_flow.rs @@ -24,7 +24,9 @@ //! not copied). Successor pointers are already valid — no block map lookup needed. use crate::convert::target_stable_storage::coerce_target_stable_value; -use crate::convert::types::packed_shared_internal_abi_info; +use crate::convert::types::{ + is_kernel_func, packed_shared_internal_abi_info, shared_return_abi_info, +}; use dialect_mir::types::MirStructType; use llvm_export::ops as llvm; use pliron::basic_block::BasicBlock; @@ -81,8 +83,17 @@ pub(crate) fn convert_return( } Some(scalar) } else { + let is_kernel = op + .deref(ctx) + .get_parent_op(ctx) + .is_some_and(|parent| is_kernel_func(ctx, parent)); let packed_shared_abi = if let Some(mir_ty) = mir_ty { - match packed_shared_internal_abi_info(ctx, mir_ty) { + let abi_res = if is_kernel { + packed_shared_internal_abi_info(ctx, mir_ty) + } else { + shared_return_abi_info(ctx, mir_ty) + }; + match abi_res { Ok(abi) => abi, Err(error) => { return pliron::input_err_noloc!( diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs b/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs index 3c60bbf1b8..55c1a5112b 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/func_abi.rs @@ -349,6 +349,37 @@ pub(crate) fn packed_shared_internal_abi_info( })) } +/// Recognize an internal device-function return type that must cross the `.func` +/// boundary as a target-stable 64-bit generic pointer (or generic-pointer carrier). +/// +/// Under `--nvptx-short-ptr` (`p3:32:32:32`), LLVM 23's `NVPTXAsmPrinter` emits +/// `.func (.param .b64 func_retval0)` in the function header for a bare +/// `ptr addrspace(3)` return value while emitting `.param .b32 retval0` at the +/// call site. Returning `ptr` (`addrspace(0)`) across the `.func` boundary and +/// coercing back to `ptr addrspace(3)` at the call site matches `nvcc` and +/// avoids the PTX return-slot width mismatch. +pub(crate) fn shared_return_abi_info( + ctx: &mut Context, + mir_ty: TypeHandle, +) -> Result, anyhow::Error> { + if let Some(abi) = packed_shared_internal_abi_info(ctx, mir_ty)? { + return Ok(Some(abi)); + } + let is_shared_ptr = mir_ty + .deref(ctx) + .downcast_ref::() + .is_some_and(|ptr_ty| ptr_ty.address_space() == llvm_types::address_space::SHARED); + if !is_shared_ptr { + return Ok(None); + } + let semantic_ty = convert_type(ctx, mir_ty)?; + let storage_ty = llvm_types::PointerType::get_generic(ctx).into(); + Ok(Some(PackedSharedInternalAbiInfo { + semantic_ty, + storage_ty, + })) +} + /// Convert a type that crosses a function boundary as one LLVM value. /// /// A struct with a natural-layout divergence is legal by value only when @@ -621,7 +652,7 @@ pub fn convert_function_type( let ty = if is_transparent_scalar { transparent_scalar_llvm_type(ctx, mir_ret_ty)? } else if !is_kernel_entry { - match packed_shared_internal_abi_info(ctx, mir_ret_ty)? { + match shared_return_abi_info(ctx, mir_ret_ty)? { Some(abi) => abi.storage_ty, None => convert_by_value_abi_type(ctx, mir_ret_ty, "by-value function return")?, } diff --git a/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs b/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs index af6089f3cc..f68c950342 100644 --- a/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs +++ b/cuda-oxide/crates/mir-lower/src/convert/types/mod.rs @@ -117,7 +117,7 @@ pub(crate) use enum_layout::{ }; pub(crate) use func_abi::{ TransparentScalarAbiInfo, convert_grid_constant_storage_type, packed_shared_internal_abi_info, - transparent_scalar_abi_info, transparent_scalar_field, + shared_return_abi_info, transparent_scalar_abi_info, transparent_scalar_field, }; pub use func_abi::{convert_function_type, is_kernel_func}; // Surface parity with the pre-split `types.rs`: these four are currently