python3Packages.warp-lang: init at 1.7.2.post1

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

@@ -0,0 +1,39 @@
+diff --git a/warp/build_dll.py b/warp/build_dll.py
+index 4d411e1b..a9304c6a 100644
+--- a/warp/build_dll.py
++++ b/warp/build_dll.py
+@@ -316,6 +316,8 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
+         cuda_includes = f' -I"{cuda_home}/include"' if cu_path else ""
+         includes = cpp_includes + cuda_includes
+ 
++        includes += " -isystem @LIBCXX_DEV@/include/c++/v1"
++
+         if sys.platform == "darwin":
+             version = f"--target={arch}-apple-macos11"
+         else:
+@@ -345,6 +347,8 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
+                 build_cmd = f'g++ {cpp_flags} -c "{cpp_path}" -o "{cpp_out}"'
+                 run_cmd(build_cmd)
+ 
++        ld_inputs.append('-L"@LIBCXX_LIB@/lib" -lc++')
++
+         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/darwin-single-target.patch

@@ -0,0 +1,50 @@
+diff --git a/build_llvm.py b/build_llvm.py
+index 9d5a26a7..0be02a89 100644
+--- a/build_llvm.py
++++ b/build_llvm.py
+@@ -389,15 +389,4 @@ def build_warp_clang_for_arch(args, lib_name, arch):
+ 
+ 
+ def build_warp_clang(args, lib_name):
+-    if sys.platform == "darwin":
+-        # create a universal binary by combining x86-64 and AArch64 builds
+-        build_warp_clang_for_arch(args, lib_name + "-x86_64", "x86_64")
+-        build_warp_clang_for_arch(args, lib_name + "-aarch64", "aarch64")
+-
+-        dylib_path = os.path.join(build_path, f"bin/{lib_name}")
+-        run_cmd(f"lipo -create -output {dylib_path} {dylib_path}-x86_64 {dylib_path}-aarch64")
+-        os.remove(f"{dylib_path}-x86_64")
+-        os.remove(f"{dylib_path}-aarch64")
+-
+-    else:
+-        build_warp_clang_for_arch(args, lib_name, machine_architecture())
++    build_warp_clang_for_arch(args, lib_name, machine_architecture())
+diff --git a/warp/build_dll.py b/warp/build_dll.py
+index 4d411e1b..4cf4a6c2 100644
+--- a/warp/build_dll.py
++++ b/warp/build_dll.py
+@@ -317,7 +317,7 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
+         includes = cpp_includes + cuda_includes
+ 
+         if sys.platform == "darwin":
+-            version = f"--target={arch}-apple-macos11"
++            version = ""
+         else:
+             version = "-fabi-version=13"  # GCC 8.2+
+ 
+@@ -392,14 +392,4 @@ def build_dll(args, dll_path, cpp_paths, cu_path, libs=None):
+     if libs is None:
+         libs = []
+ 
+-    if sys.platform == "darwin":
+-        # create a universal binary by combining x86-64 and AArch64 builds
+-        build_dll_for_arch(args, dll_path + "-x86_64", cpp_paths, cu_path, libs, "x86_64")
+-        build_dll_for_arch(args, dll_path + "-aarch64", cpp_paths, cu_path, libs, "aarch64")
+-
+-        run_cmd(f"lipo -create -output {dll_path} {dll_path}-x86_64 {dll_path}-aarch64")
+-        os.remove(f"{dll_path}-x86_64")
+-        os.remove(f"{dll_path}-aarch64")
+-
+-    else:
+-        build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, machine_architecture())
++    build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, machine_architecture())

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

@@ -0,0 +1,251 @@
+{
+  config,
+  lib,
+  stdenv,
+  buildPythonPackage,
+  fetchurl,
+  fetchFromGitHub,
+  replaceVars,
+  build,
+  setuptools,
+  numpy,
+  llvmPackages,
+  cudaPackages,
+  unittestCheckHook,
+  jax,
+  torch,
+  nix-update-script,
+
+  # Use standalone LLVM-based JIT compiler and CPU device support
+  standaloneSupport ? true,
+
+  # Use CUDA toolchain and GPU device support
+  cudaSupport ? config.cudaSupport,
+
+  # Build Warp with MathDx support (requires CUDA support)
+  # Most linear-algebra tile operations like tile_cholesky(), tile_fft(),
+  # and tile_matmul() require Warp to be built with the MathDx library.
+  libmathdxSupport ? cudaSupport && stdenv.hostPlatform.isLinux,
+}:
+
+let
+  version = "1.7.2.post1";
+
+  libmathdx = stdenv.mkDerivation (finalAttrs: {
+    pname = "libmathdx";
+    version = "0.2.0";
+
+    src =
+      let
+        inherit (stdenv.hostPlatform) system;
+        selectSystem = attrs: attrs.${system} or (throw "Unsupported system: ${system}");
+
+        suffix = selectSystem {
+          x86_64-linux = "Linux-x86_64";
+          aarch64-linux = "Linux-aarch64";
+          x86_64-windows = "win32-x86_64";
+        };
+
+        # nix-hash --type sha256 --to-sri $(nix-prefetch-url "https://...")
+        hash = selectSystem {
+          x86_64-linux = "sha256-Lk+PxWFvyQGRClFdmyuo4y7HBdR7pigOhMyEzajqbmg=";
+          aarch64-linux = "sha256-6tH9YH98kSvDiut9rQEU5potEpeKqma/QtrCHLxwRLo=";
+          x86_64-windows = "sha256-B8qwj7UzOXEDZh2oT3ip1qW0uqtygMsyfcbhh5Dgc8U=";
+        };
+      in
+      fetchurl {
+        url = "https://developer.nvidia.com/downloads/compute/cublasdx/redist/cublasdx/libmathdx-${suffix}-${finalAttrs.version}.tar.gz";
+        inherit hash;
+      };
+
+    unpackPhase = ''
+      runHook preUnpack
+
+      mkdir unpacked
+      cd unpacked
+      tar -xzf $src
+      export sourceRoot=$(pwd)
+
+      runHook postUnpack
+    '';
+
+    dontConfigure = true;
+    dontBuild = true;
+
+    installPhase = ''
+      runHook preInstall
+
+      cp -rT "$sourceRoot" "$out"
+
+      runHook postInstall
+    '';
+
+    meta = {
+      description = "library used to integrate cuBLASDx and cuFFTDx into Warp";
+      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.
+        (
+          nvidiaCudaRedist
+          // {
+            url = "https://developer.download.nvidia.cn/compute/mathdx/License.txt";
+          }
+        )
+
+        # 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.
+        bsd
+      ];
+      platforms = with lib.platforms; linux ++ [ "x86_64-windows" ];
+      maintainers = with lib.maintainers; [ yzx9 ];
+    };
+  });
+in
+buildPythonPackage {
+  pname = "warp-lang";
+  inherit version;
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "NVIDIA";
+    repo = "warp";
+    tag = "v${version}";
+    hash = "sha256-cT0CrD71nNZnQMimGrmnSQl6RQx4MiUv2xBFPWNI/0s=";
+  };
+
+  patches =
+    lib.optionals stdenv.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
+    ];
+
+  postPatch =
+    lib.optionalString (!stdenv.cc.isGNU) ''
+      substituteInPlace warp/build_dll.py \
+        --replace-fail "g++" "${lib.getExe stdenv.cc}"
+    ''
+    # Broken tests on aarch64. Since unittest doesn't support disabling a
+    # single test, and pytest isn't compatible, we patch the test file directly
+    # instead.
+    #
+    # See: https://github.com/NVIDIA/warp/issues/552
+    + lib.optionalString stdenv.hostPlatform.isAarch64 ''
+      substituteInPlace warp/tests/test_fem.py \
+        --replace-fail "add_function_test(TestFem, \"test_integrate_gradient\", test_integrate_gradient, devices=devices)" ""
+    '';
+
+  build-system = [
+    build
+    setuptools
+  ];
+
+  dependencies = [
+    numpy
+  ];
+
+  nativeBuildInputs = lib.optionals libmathdxSupport [
+    libmathdx
+    cudaPackages.libcublas
+    cudaPackages.libcufft
+    cudaPackages.libnvjitlink
+  ];
+
+  buildInputs =
+    lib.optionals standaloneSupport [
+      llvmPackages.llvm
+      llvmPackages.clang
+      llvmPackages.libcxx
+    ]
+    ++ lib.optionals cudaSupport [
+      cudaPackages.cudatoolkit
+      cudaPackages.cuda_cudart
+      cudaPackages.cuda_nvcc
+      cudaPackages.cuda_nvrtc
+    ];
+
+  preBuild =
+    let
+      buildOptions =
+        lib.optionals (!standaloneSupport) [
+          "--no_standalone"
+        ]
+        ++ lib.optionals cudaSupport [
+          "--cuda_path=${cudaPackages.cudatoolkit}"
+        ]
+        ++ lib.optionals libmathdxSupport [
+          "--libmathdx"
+          "--libmathdx_path=${libmathdx}"
+        ]
+        ++ lib.optionals (!libmathdxSupport) [
+          "--no_libmathdx"
+        ];
+
+      buildOptionString = lib.concatStringsSep " " buildOptions;
+    in
+    ''
+      python build_lib.py ${buildOptionString}
+    '';
+
+  pythonImportsCheck = [
+    "warp"
+  ];
+
+  # Many unit tests fail with segfaults on aarch64-linux, especially in the sim
+  # and grad modules. However, other functionality generally works, so we don't
+  # mark the package as broken.
+  #
+  # See: https://www.github.com/NVIDIA/warp/issues/{356,372,552}
+  doCheck = !(stdenv.hostPlatform.isAarch64 && stdenv.hostPlatform.isLinux);
+
+  nativeCheckInputs = [
+    unittestCheckHook
+    (jax.override { inherit cudaSupport; })
+    (torch.override { inherit cudaSupport; })
+
+    # # Disable paddlepaddle interop tests: malloc(): unaligned tcache chunk detected
+    #  (paddlepaddle.override { inherit cudaSupport; })
+  ];
+
+  preCheck = ''
+    export WARP_CACHE_PATH=$(mktemp -d) # warp.config.kernel_cache_dir
+  '';
+
+  passthru.updateScript = nix-update-script { };
+
+  meta = {
+    description = "Python framework for high performance GPU simulation and graphics";
+    longDescription = ''
+      Warp is a Python framework for writing high-performance simulation
+      and graphics code. Warp takes regular Python functions and JIT
+      compiles them to efficient kernel code that can run on the CPU or
+      GPU.
+
+      Warp is designed for spatial computing and comes with a rich set
+      of primitives that make it easy to write programs for physics
+      simulation, perception, robotics, and geometry processing. In
+      addition, Warp kernels are differentiable and can be used as part
+      of machine-learning pipelines with frameworks such as PyTorch,
+      JAX and Paddle.
+    '';
+    homepage = "https://github.com/NVIDIA/warp";
+    changelog = "https://github.com/NVIDIA/warp/blob/v${version}/CHANGELOG.md";
+    license = lib.licenses.asl20;
+    platforms = with lib.platforms; linux ++ darwin;
+    maintainers = with lib.maintainers; [ yzx9 ];
+  };
+}

pkgs/development/python-modules/warp-lang/standalone-cxx11-abi.patch

@@ -0,0 +1,25 @@
+diff --git a/build_llvm.py b/build_llvm.py
+index 9d5a26a7..839909ad 100644
+--- a/build_llvm.py
++++ b/build_llvm.py
+@@ -161,7 +161,6 @@ def build_from_source_for_arch(args, arch, llvm_source):
+         "-D", "LLVM_INCLUDE_TESTS=FALSE",
+         "-D", "LLVM_INCLUDE_TOOLS=TRUE",  # Needed by Clang
+         "-D", "LLVM_INCLUDE_UTILS=FALSE",
+-        "-D", f"CMAKE_CXX_FLAGS=-D_GLIBCXX_USE_CXX11_ABI=0 {abi_version}",  # The pre-C++11 ABI is still the default on the CentOS 7 toolchain
+         "-D", f"CMAKE_INSTALL_PREFIX={install_path}",
+         "-D", f"LLVM_HOST_TRIPLE={host_triple}",
+         "-D", f"CMAKE_OSX_ARCHITECTURES={osx_architectures}",
+diff --git a/warp/build_dll.py b/warp/build_dll.py
+index 4d411e1b..4177725b 100644
+--- a/warp/build_dll.py
++++ b/warp/build_dll.py
+@@ -321,7 +321,7 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
+         else:
+             version = "-fabi-version=13"  # GCC 8.2+
+
+-        cpp_flags = f'{version} --std=c++17 -fno-rtti -D{cuda_enabled} -D{mathdx_enabled} -D{cuda_compat_enabled} -fPIC -fvisibility=hidden -D_GLIBCXX_USE_CXX11_ABI=0 -I"{native_dir}" {includes} '
++        cpp_flags = f'{version} --std=c++17 -fno-rtti -D{cuda_enabled} -D{mathdx_enabled} -D{cuda_compat_enabled} -fPIC -fvisibility=hidden -I"{native_dir}" {includes} '
+
+         if mode == "debug":
+             cpp_flags += "-O0 -g -D_DEBUG -DWP_ENABLE_DEBUG=1 -fkeep-inline-functions"

pkgs/development/python-modules/warp-lang/standalone-llvm.patch

@@ -0,0 +1,62 @@
+diff --git a/build_llvm.py b/build_llvm.py
+index 9d5a26a7..3663e9c9 100644
+--- a/build_llvm.py
++++ b/build_llvm.py
+@@ -338,25 +338,19 @@ def build_warp_clang_for_arch(args, lib_name, arch):
+ 
+         clang_dll_path = os.path.join(build_path, f"bin/{lib_name}")
+ 
+-        if args.build_llvm:
+-            # obtain Clang and LLVM libraries from the local build
+-            install_path = os.path.join(llvm_install_path, f"{args.mode}-{arch}")
+-            libpath = os.path.join(install_path, "lib")
+-        else:
+-            # obtain Clang and LLVM libraries from packman
+-            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 = []
+-
+-        for _, _, libraries in os.walk(libpath):
+-            libs.extend(libraries)
+-            break  # just the top level contains library files
++        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")
+             libs.append("Ws2_32.lib")
+-            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"]
+             if sys.platform == "darwin":
+@@ -364,7 +358,8 @@ def build_warp_clang_for_arch(args, lib_name, arch):
+             else:
+                 libs.insert(0, "-Wl,--start-group")
+                 libs.append("-Wl,--end-group")
+-            libs.append(f"-L{libpath}")
++            libs.extend(f"-L{libpath}" for libpath in libpaths)
++            libs.append("-lz")
+             libs.append("-lpthread")
+             libs.append("-ldl")
+             if sys.platform != "darwin":
+diff --git a/warp/build_dll.py b/warp/build_dll.py
+index 4d411e1b..95fb7eaf 100644
+--- a/warp/build_dll.py
++++ b/warp/build_dll.py
+@@ -311,8 +311,8 @@ def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, libs, arch, mode=None
+             run_cmd(link_cmd)
+ 
+     else:
+-        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/top-level/python-packages.nix

@@ -19133,6 +19133,8 @@ self: super: with self; {
 
   warlock = callPackage ../development/python-modules/warlock { };
 
+  warp-lang = callPackage ../development/python-modules/warp-lang { };
+
   warrant = callPackage ../development/python-modules/warrant { };
 
   warrant-lite = callPackage ../development/python-modules/warrant-lite { };