python3Packages.warp-lang: 1.9.0 -> 1.9.1, fix build

pkgs/development/python-modules/warp-lang/clang-libs.patch

@@ -0,0 +1,15 @@
+diff --git a/warp/build_dll.py b/warp/build_dll.py
+index 2218ff13..3fcf5796 100644
+--- a/warp/build_dll.py
++++ b/warp/build_dll.py
+@@ -404,8 +404,8 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: Optional[
+         cuda_compiler = "clang++" if getattr(args, "clang_build_toolchain", False) else "nvcc"
+         cpp_compiler = "clang++" if getattr(args, "clang_build_toolchain", False) else "g++"
+
+-        cpp_includes = f' -I"{warp_home_path.parent}/external/llvm-project/out/install/{mode}-{arch}/include"'
+-        cpp_includes += f' -I"{warp_home_path.parent}/_build/host-deps/llvm-project/release-{arch}/include"'
++        cpp_includes = ' -I"@LLVM_DEV@/include"'
++        cpp_includes += ' -I"@LIBCLANG_DEV@/include"'
+         cuda_includes = f' -I"{cuda_home}/include"' if cu_path else ""
+         includes = cpp_includes + cuda_includes
+

pkgs/development/python-modules/warp-lang/clang-path.patch

@@ -0,0 +1,28 @@
+diff --git a/warp/build_dll.py b/warp/build_dll.py
+index 11d2a898..33fb2a76 100644
+--- a/warp/build_dll.py
++++ b/warp/build_dll.py
+@@ -636,8 +636,8 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: list[str]
+
+     else:
+         # Unix compilation
+-        cuda_compiler = "clang++" if getattr(args, "clang_build_toolchain", False) else "nvcc"
+-        cpp_compiler = "clang++" if getattr(args, "clang_build_toolchain", False) else "g++"
++        cuda_compiler = "@CLANG@" if getattr(args, "clang_build_toolchain", False) else "nvcc"
++        cpp_compiler = "@CLANG@" if getattr(args, "clang_build_toolchain", False) else "g++"
+
+         cpp_includes = f' -I"{warp_home_path.parent}/external/llvm-project/out/install/{mode}-{arch}/include"'
+         cpp_includes += f' -I"{warp_home_path.parent}/_build/host-deps/llvm-project/release-{arch}/include"'
+@@ -687,9 +687,9 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: list[str]
+             else:
+                 # Use Clang compiler
+                 if mode == "debug":
+-                    cuda_cmd = f'clang++ -Werror -Wuninitialized -Wno-unknown-cuda-version {" ".join(clang_opts)} -g -O0 -fPIC -fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
++                    cuda_cmd = f'@CLANG@ -Werror -Wuninitialized -Wno-unknown-cuda-version {" ".join(clang_opts)} -g -O0 -fPIC -fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+                 elif mode == "release":
+-                    cuda_cmd = f'clang++ -Werror -Wuninitialized -Wno-unknown-cuda-version {" ".join(clang_opts)} -O3 -fPIC -fvisibility=hidden -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
++                    cuda_cmd = f'@CLANG@ -Werror -Wuninitialized -Wno-unknown-cuda-version {" ".join(clang_opts)} -O3 -fPIC -fvisibility=hidden -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
+
+             with ScopedTimer("build_cuda", active=args.verbose):
+                 run_cmd(cuda_cmd)
+

pkgs/development/python-modules/warp-lang/darwin-libcxx.patch

@@ -18,20 +18,3 @@ index 2218ff13..53786017 100644
          if cu_path:
              cu_out = cu_path + ".o"
 
-diff --git a/warp/native/crt.h b/warp/native/crt.h
-index 47ef9983..89ae289b 100644
---- a/warp/native/crt.h
-+++ b/warp/native/crt.h
-@@ -65,6 +65,12 @@ extern "C" WP_API int _wp_isinf(double);
- #include <float.h>
- #include <string.h>
-
-+#undef isfinite
-+#undef isinf
-+#undef isnan
-+
-+#include <cmath>
-+
- #else
-
- // These definitions are taken from Jitify: https://github.com/NVIDIA/jitify

pkgs/development/python-modules/warp-lang/default.nix

@@ -1,14 +1,13 @@
 {
-  _cuda,
   autoAddDriverRunpath,
   buildPythonPackage,
   config,
   cudaPackages,
+  callPackage,
   fetchFromGitHub,
-  fetchurl,
   jax,
   lib,
-  llvmPackages,
+  llvmPackages, # TODO: use llvm 21 in 1.10, see python-packages.nix
   numpy,
   pkgsBuildHost,
   python,
@@ -33,100 +32,16 @@
   # libmathdxSupport ? cudaSupport && stdenv.hostPlatform.isLinux,
   libmathdxSupport ? cudaSupport,
 }@args:
+
 assert libmathdxSupport -> cudaSupport;
+
 let
   effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else args.stdenv;
   stdenv = throw "Use effectiveStdenv instead of stdenv directly, as it may be replaced by cudaPackages.backendStdenv";
 
-  version = "1.9.0";
-
-  libmathdx = effectiveStdenv.mkDerivation (finalAttrs: {
-    # NOTE: The version used should match the version Warp requires:
-    # https://github.com/NVIDIA/warp/blob/${version}/deps/libmathdx-deps.packman.xml
-    pname = "libmathdx";
-    version = "0.2.3";
-
-    outputs = [
-      "out"
-      "static"
-    ];
-
-    src =
-      let
-        baseURL = "https://developer.nvidia.com/downloads/compute/cublasdx/redist/cublasdx";
-        cudaMajorVersion = cudaPackages.cudaMajorVersion; # only 12, 13 supported
-        cudaVersion = "${cudaMajorVersion}.0"; # URL example: ${baseURL}/cuda12/${name}-${version}-cuda12.0.zip
-        name = lib.concatStringsSep "-" [
-          finalAttrs.pname
-          "Linux"
-          effectiveStdenv.hostPlatform.parsed.cpu.name
-          finalAttrs.version
-          "cuda${cudaVersion}"
-        ];
-
-        # nix-hash --type sha256 --to-sri $(nix-prefetch-url "https://...")
-        hashes = {
-          "12" = {
-            aarch64-linux = "sha256-d/aBC+zU2ciaw3isv33iuviXYaLGLdVDdzynGk9SFck=";
-            x86_64-linux = "sha256-CHIH0s4SnA67COtHBkwVCajW/3f0VxNBmuDLXy4LFIg=";
-          };
-          "13" = {
-            aarch64-linux = "sha256-TetJbMts8tpmj5PV4+jpnUHMcooDrXUEKL3aGWqilKI=";
-            x86_64-linux = "sha256-wLJLbRpQWa6QEm8ibm1gxt3mXvkWvu0vEzpnqTIvE1M=";
-          };
-        };
-      in
-      lib.mapNullable (
-        hash:
-        fetchurl {
-          inherit hash name;
-          url = "${baseURL}/cuda${cudaMajorVersion}/${name}.tar.gz";
-        }
-      ) (hashes.${cudaMajorVersion}.${effectiveStdenv.hostPlatform.system} or null);
-
-    dontUnpack = true;
-    dontConfigure = true;
-    dontBuild = true;
-
-    installPhase = ''
-      runHook preInstall
+  version = "1.9.1";
 
-      mkdir -p "$out"
-      tar -xzf "$src" -C "$out"
-
-      mkdir -p "$static"
-      moveToOutput "lib/libmathdx_static.a" "$static"
-
-      runHook postInstall
-    '';
-
-    meta = {
-      description = "Library used to integrate cuBLASDx and cuFFTDx into Warp";
-      homepage = "https://developer.nvidia.com/cublasdx-downloads";
-      sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
-      license = with lib.licenses; [
-        # By downloading and using the software, you agree to fully
-        # comply with the terms and conditions of the NVIDIA Software
-        # License Agreement.
-        _cuda.lib.licenses.math_sdk_sla
-
-        # Some of the libmathdx routines were written by or derived
-        # from code written by Meta Platforms, Inc. and affiliates and
-        # are subject to the BSD License.
-        bsd3
-
-        # Some of the libmathdx routines were written by or derived from
-        # code written by Victor Zverovich and are subject to the following
-        # license:
-        mit
-      ];
-      platforms = [
-        "aarch64-linux"
-        "x86_64-linux"
-      ];
-      maintainers = with lib.maintainers; [ yzx9 ];
-    };
-  });
+  libmathdx = callPackage ./libmathdx.nix { };
 in
 buildPythonPackage {
   pname = "warp-lang";
@@ -144,44 +59,54 @@ buildPythonPackage {
     owner = "NVIDIA";
     repo = "warp";
     tag = "v${version}";
-    hash = "sha256-OEg2mUsEdRKhgx0fIraqme4moKNh1RSdN7/yCT1V5+g=";
+    hash = "sha256-Atp3WyxQ7GYwWLmQIUgoPULyVlNjduh4/9CBixNWFwc=";
   };
 
-  patches =
-    lib.optionals effectiveStdenv.hostPlatform.isDarwin [
-      (replaceVars ./darwin-libcxx.patch {
-        LIBCXX_DEV = llvmPackages.libcxx.dev;
-        LIBCXX_LIB = llvmPackages.libcxx;
-      })
-      ./darwin-single-target.patch
-    ]
-    ++ lib.optionals standaloneSupport [
-      (replaceVars ./standalone-llvm.patch {
-        LLVM_DEV = llvmPackages.llvm.dev;
-        LLVM_LIB = llvmPackages.llvm.lib;
-        LIBCLANG_DEV = llvmPackages.libclang.dev;
-        LIBCLANG_LIB = llvmPackages.libclang.lib;
-      })
-      ./standalone-cxx11-abi.patch
-    ];
+  patches = [
+    ./cxx11-abi.patch
+  ]
+  ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [
+    (replaceVars ./darwin-libcxx.patch {
+      LIBCXX_DEV = llvmPackages.libcxx.dev;
+      LIBCXX_LIB = llvmPackages.libcxx;
+    })
+
+    ./darwin-single-target.patch
+  ]
+  ++ lib.optionals (effectiveStdenv.cc.isClang || standaloneSupport) [
+    (replaceVars ./clang-path.patch {
+      CLANG = "${effectiveStdenv.cc}/bin/cc";
+    })
+
+    (replaceVars ./clang-libs.patch {
+      LLVM_DEV = llvmPackages.llvm.dev;
+      LIBCLANG_DEV = llvmPackages.libclang.dev;
+    })
+  ]
+  ++ lib.optionals standaloneSupport [
+    (replaceVars ./llvm-libs.patch {
+      LLVM_LIB = llvmPackages.llvm.lib;
+      LIBCLANG_LIB = llvmPackages.libclang.lib;
+    })
+  ];
 
   postPatch =
     # Patch build_dll.py to use our gencode flags rather than NVIDIA's very broad defaults.
     lib.optionalString cudaSupport ''
       nixLog "patching $PWD/warp/build_dll.py to use our gencode flags"
       substituteInPlace "$PWD/warp/build_dll.py" \
-          --replace-fail \
-            '*gencode_opts,' \
-            '${
-              lib.concatMapStringsSep ", " (gencodeString: ''"${gencodeString}"'') cudaPackages.flags.gencode
-            },' \
-          --replace-fail \
-            '*clang_arch_flags,' \
-            '${
-              lib.concatMapStringsSep ", " (
-                realArch: ''"--cuda-gpu-arch=${realArch}"''
-              ) cudaPackages.flags.realArches
-            },'
+        --replace-fail \
+          '*gencode_opts,' \
+          '${
+            lib.concatMapStringsSep ", " (gencodeString: ''"${gencodeString}"'') cudaPackages.flags.gencode
+          },' \
+        --replace-fail \
+          '*clang_arch_flags,' \
+          '${
+            lib.concatMapStringsSep ", " (
+              realArch: ''"--cuda-gpu-arch=${realArch}"''
+            ) cudaPackages.flags.realArches
+          },'
     ''
     # Patch build_dll.py to use dynamic libraries rather than static ones.
     # NOTE: We do not patch the `nvptxcompiler_static` path because it is not available as a dynamic library.
@@ -232,11 +157,10 @@ buildPythonPackage {
     numpy
   ];
 
-  # NOTE: While normally we wouldn't include autoAddDriverRunpath for packages built from source, since Warp
-  # will be loading GPU drivers at runtime, we need to inject the path to our video drivers.
   nativeBuildInputs = lib.optionals cudaSupport [
+    # NOTE: While normally we wouldn't include autoAddDriverRunpath for packages built from source, since Warp
+    # will be loading GPU drivers at runtime, we need to inject the path to our video drivers.
     autoAddDriverRunpath
-    cudaPackages.cuda_nvcc
   ];
 
   buildInputs =
@@ -246,7 +170,7 @@ buildPythonPackage {
       llvmPackages.libcxx
     ]
     ++ lib.optionals cudaSupport [
-      (lib.getOutput "static" cudaPackages.cuda_nvcc) # dependency on nvptxcompiler_static; no dynamic version available
+      (lib.getStatic cudaPackages.cuda_nvcc) # dependency on nvptxcompiler_static; no dynamic version available
       cudaPackages.cuda_cccl
       cudaPackages.cuda_cudart
       cudaPackages.cuda_nvcc

pkgs/development/python-modules/warp-lang/libmathdx.nix

@@ -0,0 +1,95 @@
+{
+  _cuda,
+  cudaPackages,
+  fetchurl,
+  lib,
+  stdenvNoCC,
+}:
+
+stdenvNoCC.mkDerivation (finalAttrs: {
+  # NOTE: The version used should match the version Warp requires:
+  # https://github.com/NVIDIA/warp/blob/${version}/deps/libmathdx-deps.packman.xml
+  pname = "libmathdx";
+  version = "0.2.3";
+
+  outputs = [
+    "out"
+    "static"
+  ];
+
+  src =
+    let
+      baseURL = "https://developer.nvidia.com/downloads/compute/cublasdx/redist/cublasdx";
+      cudaMajorVersion = cudaPackages.cudaMajorVersion; # only 12, 13 supported
+      cudaVersion = "${cudaMajorVersion}.0"; # URL example: ${baseURL}/cuda12/${name}-${version}-cuda12.0.zip
+      name = lib.concatStringsSep "-" [
+        finalAttrs.pname
+        "Linux"
+        stdenvNoCC.hostPlatform.parsed.cpu.name
+        finalAttrs.version
+        "cuda${cudaVersion}"
+      ];
+
+      # nix-hash --type sha256 --to-sri $(nix-prefetch-url "https://...")
+      hashes = {
+        "12" = {
+          aarch64-linux = "sha256-d/aBC+zU2ciaw3isv33iuviXYaLGLdVDdzynGk9SFck=";
+          x86_64-linux = "sha256-CHIH0s4SnA67COtHBkwVCajW/3f0VxNBmuDLXy4LFIg=";
+        };
+        "13" = {
+          aarch64-linux = "sha256-TetJbMts8tpmj5PV4+jpnUHMcooDrXUEKL3aGWqilKI=";
+          x86_64-linux = "sha256-wLJLbRpQWa6QEm8ibm1gxt3mXvkWvu0vEzpnqTIvE1M=";
+        };
+      };
+    in
+    lib.mapNullable (
+      hash:
+      fetchurl {
+        inherit hash name;
+        url = "${baseURL}/cuda${cudaMajorVersion}/${name}.tar.gz";
+      }
+    ) (hashes.${cudaMajorVersion}.${stdenvNoCC.hostPlatform.system} or null);
+
+  dontUnpack = true;
+  dontConfigure = true;
+  dontBuild = true;
+
+  installPhase = ''
+    runHook preInstall
+
+    mkdir -p "$out"
+    tar -xzf "$src" -C "$out"
+
+    mkdir -p "$static"
+    moveToOutput "lib/libmathdx_static.a" "$static"
+
+    runHook postInstall
+  '';
+
+  meta = {
+    description = "Library used to integrate cuBLASDx and cuFFTDx into Warp";
+    homepage = "https://developer.nvidia.com/cublasdx-downloads";
+    sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
+    license = with lib.licenses; [
+      # By downloading and using the software, you agree to fully
+      # comply with the terms and conditions of the NVIDIA Software
+      # License Agreement.
+      _cuda.lib.licenses.math_sdk_sla
+
+      # Some of the libmathdx routines were written by or derived
+      # from code written by Meta Platforms, Inc. and affiliates and
+      # are subject to the BSD License.
+      bsd3
+
+      # Some of the libmathdx routines were written by or derived from
+      # code written by Victor Zverovich and are subject to the following
+      # license:
+      mit
+    ];
+    platforms = [
+      "aarch64-linux"
+      "x86_64-linux"
+    ];
+    maintainers = with lib.maintainers; [ yzx9 ];
+  };
+})