nix flake

Author: brandonros

crates/cust/src/memory/unified.rs

@@ -19,6 +19,29 @@ use crate::memory::UnifiedPointer;
 use crate::memory::malloc::{cuda_free_unified, cuda_malloc_unified};
 use crate::prelude::Stream;
 
+#[cfg(any(cuMemPrefetchAsync_v2, cuMemAdvise_v2))]
+unsafe fn cu_mem_location(
+    type_: driver_sys::CUmemLocationType,
+    id: std::os::raw::c_int,
+) -> driver_sys::CUmemLocation {
+    let mut location = std::mem::MaybeUninit::<driver_sys::CUmemLocation>::zeroed();
+    let location_ptr = location.as_mut_ptr();
+
+    // Support both older bindgen output (`{ type_, id }`) and the newer
+    // anonymous-union layout emitted from CUDA 13.2 headers.
+    unsafe {
+        (*location_ptr).type_ = type_;
+        std::ptr::write(
+            (location_ptr.cast::<u8>())
+                .add(std::mem::size_of::<driver_sys::CUmemLocationType>())
+                .cast::<std::os::raw::c_int>(),
+            id,
+        );
+
+        location.assume_init()
+    }
+}
+
 /// A pointer type for heap-allocation in CUDA unified memory.
 ///
 /// See the [`module-level documentation`](../memory/index.html) for more information on unified
@@ -640,17 +663,13 @@ pub trait MemoryAdvise<T: DeviceCopy>: private::Sealed {
         let mem_size = std::mem::size_of_val(slice);
 
         unsafe {
-            let id = -1; // -1 is CU_DEVICE_CPU
             driver_sys::cuMemPrefetchAsync(
                 slice.as_ptr() as driver_sys::CUdeviceptr,
                 mem_size,
                 #[cfg(cuMemPrefetchAsync_v2)]
-                driver_sys::CUmemLocation {
-                    type_: driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
-                    id,
-                },
+                cu_mem_location(driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_HOST, 0),
                 #[cfg(not(cuMemPrefetchAsync_v2))]
-                id,
+                -1, // -1 is CU_DEVICE_CPU
                 #[cfg(cuMemPrefetchAsync_v2)]
                 0, // flags for future use, must be 0 as of CUDA 13.0
                 stream.as_inner(),
@@ -691,10 +710,7 @@ pub trait MemoryAdvise<T: DeviceCopy>: private::Sealed {
                 slice.as_ptr() as driver_sys::CUdeviceptr,
                 mem_size,
                 #[cfg(cuMemPrefetchAsync_v2)]
-                driver_sys::CUmemLocation {
-                    type_: driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
-                    id,
-                },
+                cu_mem_location(driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE, id),
                 #[cfg(not(cuMemPrefetchAsync_v2))]
                 id,
                 #[cfg(cuMemPrefetchAsync_v2)]
@@ -727,18 +743,14 @@ pub trait MemoryAdvise<T: DeviceCopy>: private::Sealed {
         };
 
         unsafe {
-            let id = 0;
             driver_sys::cuMemAdvise(
                 slice.as_ptr() as driver_sys::CUdeviceptr,
                 mem_size,
                 advice,
                 #[cfg(cuMemAdvise_v2)]
-                driver_sys::CUmemLocation {
-                    type_: driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
-                    id,
-                },
+                cu_mem_location(driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_HOST, 0),
                 #[cfg(not(cuMemAdvise_v2))]
-                id,
+                0,
             )
             .to_result()?;
         }
@@ -775,9 +787,11 @@ pub trait MemoryAdvise<T: DeviceCopy>: private::Sealed {
                 mem_size,
                 driver_sys::CUmem_advise::CU_MEM_ADVISE_SET_PREFERRED_LOCATION,
                 #[cfg(cuMemAdvise_v2)]
-                driver_sys::CUmemLocation {
-                    type_: driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
-                    id,
+                match preferred_location {
+                    Some(_) => {
+                        cu_mem_location(driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE, id)
+                    }
+                    None => cu_mem_location(driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_HOST, 0),
                 },
                 #[cfg(not(cuMemAdvise_v2))]
                 id,
@@ -793,18 +807,14 @@ pub trait MemoryAdvise<T: DeviceCopy>: private::Sealed {
         let mem_size = std::mem::size_of_val(slice);
 
         unsafe {
-            let id = 0;
             driver_sys::cuMemAdvise(
                 slice.as_ptr() as driver_sys::CUdeviceptr,
                 mem_size,
                 driver_sys::CUmem_advise::CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION,
                 #[cfg(cuMemAdvise_v2)]
-                driver_sys::CUmemLocation {
-                    type_: driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE,
-                    id,
-                },
+                cu_mem_location(driver_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_HOST, 0),
                 #[cfg(not(cuMemAdvise_v2))]
-                id,
+                0,
             )
             .to_result()?;
         }

flake.lock

@@ -0,0 +1,55 @@
+{
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixpkgs-unstable";
+    rust-overlay.url = "github:oxalica/rust-overlay";
+    rust-overlay.inputs.nixpkgs.follows = "nixpkgs";
+  };
+
+  outputs = { nixpkgs, rust-overlay, ... }:
+    let
+      system = "x86_64-linux";
+      pkgs = import nixpkgs {
+        inherit system;
+        overlays = [ rust-overlay.overlays.default ];
+      };
+      lib = pkgs.lib;
+      cudaRoot = "/usr/local/cuda-13.2";
+      llvm19 = pkgs.llvmPackages_19;
+      toolchain = pkgs.rust-bin.fromRustupToolchainFile ./rust-toolchain.toml;
+    in
+    {
+      devShells.${system}.default = pkgs.mkShell {
+        packages = [
+          toolchain
+          pkgs.pkg-config
+          pkgs.openssl
+          pkgs.cmake
+          pkgs.ninja
+          pkgs.ncurses
+          llvm19.clang
+          llvm19.libclang
+          (lib.getDev llvm19.llvm)
+          pkgs.stdenv.cc.cc.lib
+        ];
+
+        CUDA_HOME = cudaRoot;
+        CUDA_ROOT = cudaRoot;
+        CUDA_PATH = cudaRoot;
+        CUDA_TOOLKIT_ROOT_DIR = cudaRoot;
+        CUDA_LIBRARY_PATH =
+          "${cudaRoot}/targets/x86_64-linux/lib:${cudaRoot}/lib64:${cudaRoot}/lib64/stubs";
+
+        LLVM_CONFIG_19 = "${lib.getDev llvm19.llvm}/bin/llvm-config";
+        LIBCLANG_PATH = "${lib.getLib llvm19.libclang}/lib";
+
+        shellHook = ''
+          export PATH="${cudaRoot}/bin:${cudaRoot}/nvvm/bin:$PATH"
+          export LD_LIBRARY_PATH="${cudaRoot}/nvvm/lib64:${cudaRoot}/lib64:${pkgs.ncurses.out}/lib:${pkgs.stdenv.cc.cc.lib}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+
+          echo "rust-cuda llvm19 shell"
+          echo "  CUDA_HOME=$CUDA_HOME"
+          echo "  LLVM_CONFIG_19=$LLVM_CONFIG_19"
+        '';
+      };
+    };
+}