From 659babe816692bc60588a205608599fad5ef5767 Mon Sep 17 00:00:00 2001 From: Stephen Huan Date: Sun, 19 Jan 2025 17:35:36 -0500 Subject: [PATCH 1/4] python312Packages.jax-cuda12-pjrt: add missing cuda libraries --- .../python-modules/jax-cuda12-pjrt/default.nix | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix index 67f5774a2f49..62513250b62d 100644 --- a/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-pjrt/default.nix @@ -18,10 +18,16 @@ let cudaLibPath = lib.makeLibraryPath ( with cudaPackages; [ + (lib.getLib libcublas) # libcublas.so + (lib.getLib cuda_cupti) # libcupti.so (lib.getLib cuda_cudart) # libcudart.so (lib.getLib cudnn) # libcudnn.so - (lib.getLib libcublas) # libcublas.so - addDriverRunpath.driverLink # libcuda.so + (lib.getLib libcufft) # libcufft.so + (lib.getLib libcusolver) # libcusolver.so + (lib.getLib libcusparse) # libcusparse.so + (lib.getLib nccl) # libnccl.so + (lib.getLib libnvjitlink) # libnvJitLink.so + (lib.getLib addDriverRunpath.driverLink) # libcuda.so ] ); @@ -83,6 +89,8 @@ buildPythonPackage { pythonImportsCheck = [ "jax_plugins" ]; + inherit cudaLibPath; + meta = { description = "JAX XLA PJRT Plugin for NVIDIA GPUs"; homepage = "https://github.com/jax-ml/jax/tree/main/jax_plugins/cuda"; From 9bd08f040ed9ce929f90314e4461e93b97872e51 Mon Sep 17 00:00:00 2001 From: Stephen Huan Date: Sun, 19 Jan 2025 17:35:53 -0500 Subject: [PATCH 2/4] python3Packages.jax-cuda12-plugin: patch like jax-cuda12-pjrt --- .../jax-cuda12-plugin/default.nix | 29 +++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index 0b63fac213a1..782b3260459d 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -12,8 +12,9 @@ jax-cuda12-pjrt, }: let - inherit (cudaPackages) cudaVersion; inherit (jaxlib) version; + inherit (cudaPackages) cudaVersion; + inherit (jax-cuda12-pjrt) cudaLibPath; getSrcFromPypi = { @@ -94,12 +95,34 @@ buildPythonPackage { wheelUnpackHook ]; + # jax-cuda12-plugin looks for ptxas at runtime, e.g. with a xla custom call. + # Linking into $out is the least bad solution. See + # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 + # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211 + # * https://github.com/NixOS/nixpkgs/pull/375186 + # for more info. + postInstall = '' + mkdir -p $out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin + ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin + ln -s ${lib.getExe' cudaPackages.cuda_nvcc "nvlink"} $out/${python.sitePackages}/jax_cuda12_plugin/cuda/bin + ''; + + # jax-cuda12-plugin contains shared libraries that open other shared libraries via dlopen + # and these implicit dependencies are not recognized by ldd or + # autoPatchelfHook. That means we need to sneak them into rpath. This step + # must be done after autoPatchelfHook and the automatic stripping of + # artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the + # patchPhase. + preInstallCheck = '' + patchelf --add-rpath "${cudaLibPath}" $out/${python.sitePackages}/jax_cuda12_plugin/*.so + ''; + dependencies = [ jax-cuda12-pjrt ]; pythonImportsCheck = [ "jax_cuda12_plugin" ]; - # no tests - doCheck = false; + # FIXME: there are no tests, but we need to run preInstallCheck above + doCheck = true; meta = { description = "JAX Plugin for CUDA12"; From 17b79646e2648e36d234021233ce8c8bc342cd75 Mon Sep 17 00:00:00 2001 From: Stephen Huan Date: Sun, 19 Jan 2025 17:36:51 -0500 Subject: [PATCH 3/4] python3Packages.jax: remove outdated longDescription --- pkgs/development/python-modules/jax/default.nix | 4 ---- 1 file changed, 4 deletions(-) diff --git a/pkgs/development/python-modules/jax/default.nix b/pkgs/development/python-modules/jax/default.nix index 99154a8ea86a..25e39e72d125 100644 --- a/pkgs/development/python-modules/jax/default.nix +++ b/pkgs/development/python-modules/jax/default.nix @@ -194,10 +194,6 @@ buildPythonPackage rec { meta = { description = "Source-built JAX frontend: differentiate, compile, and transform Numpy code"; - longDescription = '' - This is the JAX frontend package, it's meant to be used together with one of the jaxlib implementations, - e.g. `python3Packages.jaxlib`, `python3Packages.jaxlib-bin`, or `python3Packages.jaxlibWithCuda`. - ''; homepage = "https://github.com/google/jax"; license = lib.licenses.asl20; maintainers = with lib.maintainers; [ samuela ]; From f9c16aa4768850e24021d1ae684125177c8ebfe3 Mon Sep 17 00:00:00 2001 From: Stephen Huan Date: Sat, 1 Feb 2025 22:32:51 -0500 Subject: [PATCH 4/4] python3Packages.jax: add operations to cuda test --- .../python-modules/jax-cuda12-plugin/default.nix | 2 +- pkgs/development/python-modules/jax/test-cuda.nix | 11 ++++++++--- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix index 782b3260459d..eadeb4a161dc 100644 --- a/pkgs/development/python-modules/jax-cuda12-plugin/default.nix +++ b/pkgs/development/python-modules/jax-cuda12-plugin/default.nix @@ -95,7 +95,7 @@ buildPythonPackage { wheelUnpackHook ]; - # jax-cuda12-plugin looks for ptxas at runtime, e.g. with a xla custom call. + # jax-cuda12-plugin looks for ptxas at runtime, e.g. with a triton kernel. # Linking into $out is the least bad solution. See # * https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 # * https://github.com/NixOS/nixpkgs/pull/288829#discussion_r1493852211 diff --git a/pkgs/development/python-modules/jax/test-cuda.nix b/pkgs/development/python-modules/jax/test-cuda.nix index ebeebb7ed88c..fefde55687d3 100644 --- a/pkgs/development/python-modules/jax/test-cuda.nix +++ b/pkgs/development/python-modules/jax/test-cuda.nix @@ -11,13 +11,18 @@ pkgs.writers.writePython3Bin "jax-test-cuda" } '' import jax + import jax.numpy as jnp from jax import random + from jax.experimental import sparse - assert jax.devices()[0].platform == "gpu" + assert jax.devices()[0].platform == "gpu" # libcuda.so - rng = random.PRNGKey(0) + rng = random.key(0) # libcudart.so, libcudnn.so x = random.normal(rng, (100, 100)) - x @ x + x @ x # libcublas.so + jnp.fft.fft(x) # libcufft.so + jnp.linalg.inv(x) # libcusolver.so + sparse.CSR.fromdense(x) @ x # libcusparse.so print("success!") ''