ptxjitcompiler.so

Author: brandonros

examples/vecadd/src/main.rs

@@ -78,7 +78,7 @@ fn load_ptx_with_log(ptx: &str) -> Result<Module, Box<dyn Error>> {
 
     // The driver accepted the PTX; drop our raw handle and re-load via cust so the
     // caller gets a typed Module with cust's lifetime/drop machinery.
-    unsafe { driver_sys::cuModuleUnload(module_ptr) };
+    let _ = unsafe { driver_sys::cuModuleUnload(module_ptr) };
     Module::from_ptx(ptx, &[]).map_err(|e| e.into())
 }
 

flake.nix

@@ -84,12 +84,18 @@
         CUDA_LIBRARY_PATH =
           "${cudaRoot}/targets/x86_64-linux/lib:${cudaRoot}/lib64:${cudaRoot}/lib64/stubs";
       };
+      # Symlink every NVIDIA-shipped driver library (libcuda, libnvidia-*) into a
+      # single shim dir that we then stick on LD_LIBRARY_PATH. libcuda alone is not
+      # enough: the driver will dlopen companions like libnvidia-ptxjitcompiler.so.1
+      # when JITing PTX, and failing to find them surfaces as
+      # CUDA_ERROR_JIT_COMPILER_NOT_FOUND from cuModuleLoadDataEx.
       driverShimHook = ''
         driver_shim_dir="$PWD/.nix-driver-libs"
         mkdir -p "$driver_shim_dir"
-        if [ -e "${driverLibPath}" ]; then
-          ln -sf "${driverLibPath}" "$driver_shim_dir/libcuda.so.1"
-        fi
+        for src in "${driverLibDir}"/libcuda.so* "${driverLibDir}"/libnvidia-*.so*; do
+          [ -e "$src" ] || continue
+          ln -sf "$src" "$driver_shim_dir/$(basename "$src")"
+        done
       '';
 
       # ---- LLVM 7-only shell ----