diff --git a/pkgs/development/python-modules/warp-lang/clang-libs.patch b/pkgs/development/python-modules/warp-lang/clang-libs.patch new file mode 100644 index 000000000000..3c8a5b87a0b8 --- /dev/null +++ b/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 + diff --git a/pkgs/development/python-modules/warp-lang/clang-path.patch b/pkgs/development/python-modules/warp-lang/clang-path.patch new file mode 100644 index 000000000000..ad87f9b3ab15 --- /dev/null +++ b/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) + diff --git a/pkgs/development/python-modules/warp-lang/standalone-cxx11-abi.patch b/pkgs/development/python-modules/warp-lang/cxx11-abi.patch similarity index 100% rename from pkgs/development/python-modules/warp-lang/standalone-cxx11-abi.patch rename to pkgs/development/python-modules/warp-lang/cxx11-abi.patch diff --git a/pkgs/development/python-modules/warp-lang/darwin-libcxx.patch b/pkgs/development/python-modules/warp-lang/darwin-libcxx.patch index 64bc5c58c7e5..beda9f58e6de 100644 --- a/pkgs/development/python-modules/warp-lang/darwin-libcxx.patch +++ b/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 - #include - -+#undef isfinite -+#undef isinf -+#undef isnan -+ -+#include -+ - #else - - // These definitions are taken from Jitify: https://github.com/NVIDIA/jitify diff --git a/pkgs/development/python-modules/warp-lang/default.nix b/pkgs/development/python-modules/warp-lang/default.nix index 9de805812305..0f90b01b2801 100644 --- a/pkgs/development/python-modules/warp-lang/default.nix +++ b/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"; + version = "1.9.1"; - 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 - - 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 diff --git a/pkgs/development/python-modules/warp-lang/libmathdx.nix b/pkgs/development/python-modules/warp-lang/libmathdx.nix new file mode 100644 index 000000000000..1de70997fd66 --- /dev/null +++ b/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 ]; + }; +}) diff --git a/pkgs/development/python-modules/warp-lang/standalone-llvm.patch b/pkgs/development/python-modules/warp-lang/llvm-libs.patch similarity index 53% rename from pkgs/development/python-modules/warp-lang/standalone-llvm.patch rename to pkgs/development/python-modules/warp-lang/llvm-libs.patch index 83860f49fbba..60c967004001 100644 --- a/pkgs/development/python-modules/warp-lang/standalone-llvm.patch +++ b/pkgs/development/python-modules/warp-lang/llvm-libs.patch @@ -1,8 +1,8 @@ diff --git a/build_llvm.py b/build_llvm.py -index aee0a9f1..6b9806c9 100644 +index aee0a9f1..20cfe8f3 100644 --- a/build_llvm.py +++ b/build_llvm.py -@@ -350,25 +350,19 @@ def build_warp_clang_for_arch(args, lib_name: str, arch: str) -> None: +@@ -350,33 +350,23 @@ def build_warp_clang_for_arch(args, lib_name: str, arch: str) -> None: clang_dll_path = os.path.join(build_path, f"bin/{lib_name}") @@ -15,18 +15,14 @@ index aee0a9f1..6b9806c9 100644 - fetch_prebuilt_libraries(arch) - libpath = os.path.join(base_path, f"_build/host-deps/llvm-project/release-{arch}/lib") - -+ # obtain Clang and LLVM libraries from the local build - libs = [] +- libs = [] - - for _, _, libraries in os.walk(libpath): - libs.extend(libraries) - break # just the top level contains library files ++ libs = ["LLVM", "clang-cpp"] + install_paths = ["@LLVM_LIB@", "@LIBCLANG_LIB@"] + libpaths = [os.path.join(install_path, "lib") for install_path in install_paths] -+ for libpath in libpaths: -+ for _, _, libraries in os.walk(libpath): -+ libs.extend(libraries) -+ break # just the top level contains library files if os.name == "nt": libs.append("Version.lib") @@ -34,9 +30,10 @@ index aee0a9f1..6b9806c9 100644 - libs.append(f'/LIBPATH:"{libpath}"') + libs.extend(f'/LIBPATH:"{libpath}"' for libpath in libpaths) else: - libs = [f"-l{lib[3:-2]}" for lib in libs if os.path.splitext(lib)[1] == ".a"] +- libs = [f"-l{lib[3:-2]}" for lib in libs if os.path.splitext(lib)[1] == ".a"] ++ libs = [f"-l{lib}" for lib in libs] if sys.platform == "darwin": -@@ -376,7 +370,8 @@ def build_warp_clang_for_arch(args, lib_name: str, arch: str) -> None: + libs += libs # prevents unresolved symbols due to link order else: libs.insert(0, "-Wl,--start-group") libs.append("-Wl,--end-group") @@ -46,17 +43,4 @@ index aee0a9f1..6b9806c9 100644 libs.append("-lpthread") libs.append("-ldl") if sys.platform != "darwin": -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 + diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index dc429cf6c55c..06688aca99be 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -20048,7 +20048,10 @@ self: super: with self; { warlock = callPackage ../development/python-modules/warlock { }; - warp-lang = callPackage ../development/python-modules/warp-lang { }; + warp-lang = callPackage ../development/python-modules/warp-lang { + stdenv = if stdenv.hostPlatform.isDarwin then pkgs.llvmPackages_19.stdenv else pkgs.stdenv; + llvmPackages = pkgs.llvmPackages_19; + }; warrant = callPackage ../development/python-modules/warrant { };