feat(cuda-oxide): use 32-bit shared-memory pointers (p3:32:32:32 + --nvptx-short-ptr) in PTX backend

cuda-oxide/crates/cuda-oxide-codegen/src/ptx.rs

@@ -333,6 +333,7 @@ fn base_llc_args(target: &str) -> Vec<String> {
     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",
             ]

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).
 ///

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::<llvm::AddrSpaceCastOp>(&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::<llvm::PtrToIntOp>(&ctx, &body),

cuda-oxide/crates/mir-lower/src/convert/ops/aggregate/fields.rs

@@ -505,8 +505,8 @@ mod tests {
 
         assert_eq!(
             count_ops::<llvm::AddrSpaceCastOp>(&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::<MirInsertFieldOp>(&ctx, &body),
@@ -610,8 +610,8 @@ mod tests {
 
         assert_eq!(
             count_ops::<llvm::AddrSpaceCastOp>(&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::<MirInsertFieldOp>(&ctx, &body), 0);
         assert_eq!(count_ops::<MirExtractFieldOp>(&ctx, &body), 0);

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

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::<llvm::AddrSpaceCastOp>(&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::<llvm::AddrSpaceCastOp>(&ctx, &body),
+            2,
+            "expected cast addrspacecast plus .func return-boundary addrspacecast"
+        );
+        assert_eq!(
+            count_ops::<mir::MirCastOp>(&ctx, &body),
+            0,
+            "mir.cast must be replaced during lowering"
+        );
         let int_to_ptrs = find_all::<llvm::IntToPtrOp>(&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::<llvm::AddrSpaceCastOp>(&ctx, module, "llvm.addrspacecast");
             let body = kernel_blocks(&ctx, module);
+            let expected_casts = if shared_is_source { 1 } else { 2 };
+            assert_eq!(
+                count_ops::<llvm::AddrSpaceCastOp>(&ctx, &body),
+                expected_casts,
+                "expected cast addrspacecast (plus return-boundary cast when returning shared)"
+            );
+            assert_eq!(count_ops::<mir::MirCastOp>(&ctx, &body), 0);
             assert_eq!(
                 count_ops::<llvm::AllocaOp>(&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::<llvm::AddrSpaceCastOp>(&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::<llvm::IntToPtrOp>(&ctx, &kernel_blocks(&ctx, module_ptr)),
+                count_ops::<llvm::AddrSpaceCastOp>(&ctx, &body),
+                expected_casts,
+                "expected cast addrspacecast (plus return-boundary cast for shared)"
+            );
+            assert_eq!(count_ops::<mir::MirCastOp>(&ctx, &body), 0);
+            assert_eq!(
+                count_ops::<llvm::IntToPtrOp>(&ctx, &body),
                 1,
                 "address space {address_space} must re-enter through a generic inttoptr"
             );
-            let casts = find_all::<llvm::AddrSpaceCastOp>(&ctx, &kernel_blocks(&ctx, module_ptr));
-            let [cast] = casts.as_slice() else {
-                panic!("expected exactly one llvm.addrspacecast");
-            };
+            let casts = find_all::<llvm::AddrSpaceCastOp>(&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::<llvm::AddrSpaceCastOp>(&ctx, module_ptr, "llvm.addrspacecast");
+        let body = kernel_blocks(&ctx, module_ptr);
+        assert_eq!(
+            count_ops::<llvm::AddrSpaceCastOp>(&ctx, &body),
+            2,
+            "ptr -> shared cast plus .func return-boundary cast must emit two llvm.addrspacecast ops"
+        );
+        assert_eq!(count_ops::<mir::MirCastOp>(&ctx, &body), 0);
 
-        let casts = find_all::<llvm::AddrSpaceCastOp>(&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::<llvm::AddrSpaceCastOp>(&ctx, &body);
+        let cast = &casts[0];
         let result_ty = cast
             .get_operation()
             .deref(&ctx)

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!(

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<Option<PackedSharedInternalAbiInfo>, 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::<MirPtrType>()
+        .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")?,
             }

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