diff --git a/pkgs/development/python-modules/nvidia-cutlass-dsl-libs-base/default.nix b/pkgs/development/python-modules/nvidia-cutlass-dsl-libs-base/default.nix new file mode 100644 index 000000000000..cd5c4659f1ac --- /dev/null +++ b/pkgs/development/python-modules/nvidia-cutlass-dsl-libs-base/default.nix @@ -0,0 +1,101 @@ +{ + lib, + stdenv, + buildPythonPackage, + fetchPypi, + python, + + # nativeBuildInputs + autoAddDriverRunpath, + autoPatchelfHook, + + # dependencies + cuda-bindings, + numpy, + typing-extensions, +}: + +let + platform = + { + x86_64-linux = "manylinux_2_28_x86_64"; + aarch64-linux = "manylinux_2_28_aarch64"; + } + .${stdenv.hostPlatform.system} + or (throw "nvidia-cutlass-dsl-libs-base is not supported on ${stdenv.hostPlatform.system}"); + + pyShortVersion = "cp${builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion}"; + + hashes = { + x86_64-linux = { + cp311 = "sha256-kReQDLpT08Iajay6a789bl8mnkJ6UmwyD7RHB6DVc2M="; + cp312 = "sha256-Fe9qWRk2Z+Zjk070hz+MytN0Vem3w8QZwwchE7iu32E="; + cp313 = "sha256-5Z2n2J5eT4UUxlMIQ/kQ+dhzTYBC3KoHnJ2cUGPrNRQ="; + cp314 = "sha256-EsKffB8fgoUQkro4aSZNr6+wNSKMDZgnqNsIuIT7gMo="; + }; + aarch64-linux = { + cp311 = "sha256-y7VVqVxwEeSzyjKL5AcpnHfSiWYK2+oi7VFdRAbmlJw="; + cp312 = "sha256-0qPEEih+NW++SP6fhF1tM8013qXiDX5PYowglXlnys0="; + cp313 = "sha256-OVvXfPZCru8xExNFPmWC8RyTV6S4H+Yg6j2szR/Mq5s="; + cp314 = "sha256-IW7uaqgQfTVWn5RRtmsDo8UxZ4QdGvm2MLlm742Wbhk="; + }; + }; +in +buildPythonPackage (finalAttrs: { + pname = "nvidia-cutlass-dsl-libs-base"; + version = "4.5.2"; + format = "wheel"; + + src = fetchPypi { + pname = "nvidia_cutlass_dsl_libs_base"; + inherit (finalAttrs) version; + format = "wheel"; + inherit platform; + dist = pyShortVersion; + python = pyShortVersion; + abi = pyShortVersion; + hash = + hashes.${stdenv.hostPlatform.system}.${pyShortVersion} + or (throw "No hash specified for '${stdenv.hostPlatform.system}.${pyShortVersion}'"); + }; + + pythonRemoveDeps = [ + # Only cuda-bindings is needed + "cuda-python" + ]; + dependencies = [ + cuda-bindings + numpy + typing-extensions + ]; + + nativeBuildInputs = [ + autoAddDriverRunpath + autoPatchelfHook + ]; + + autoPatchelfIgnoreMissingDeps = [ + # libmlir_cuda_runtime.so links libcuda.so.1 + # autoAddDriverRunpath bakes the driver path into the runpath; tell autoPatchelfHook not to fail + # on it. + "libcuda.so.1" + ]; + + # This wheel ships the `cutlass` module via `nvidia_cutlass_dsl.pth`. + # Python only processes `.pth` files in dirs that `site.py` registers as site-packages, not in + # PYTHONPATH entries so pythonImportsCheck (which uses PYTHONPATH) can't see `cutlass`. + dontUsePythonImportsCheck = true; + + # No tests in the Pypi archive + doCheck = false; + + meta = { + description = "Bundled MLIR/CUDA runtime libraries and Python sources for the NVIDIA CUTLASS DSL"; + homepage = "https://github.com/NVIDIA/cutlass"; + changelog = "https://github.com/NVIDIA/cutlass/blob/v${finalAttrs.version}/CHANGELOG.md"; + sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ]; + license = lib.licenses.unfreeRedistributable; # NVIDIA Proprietary + maintainers = with lib.maintainers; [ GaetanLepage ]; + platforms = lib.platforms.linux; + }; +}) diff --git a/pkgs/development/python-modules/nvidia-cutlass-dsl-libs-base/prefetch.sh b/pkgs/development/python-modules/nvidia-cutlass-dsl-libs-base/prefetch.sh new file mode 100755 index 000000000000..ce0200c446c2 --- /dev/null +++ b/pkgs/development/python-modules/nvidia-cutlass-dsl-libs-base/prefetch.sh @@ -0,0 +1,29 @@ +#! /usr/bin/env nix-shell +#! nix-shell -i sh -p jq + +pname="nvidia-cutlass-dsl-libs-base" +outfile="${pname}-hashes.nix" +# Clear file +rm -f $outfile + +prefetch() { + package_attr="python${1}Packages.${pname}" + echo "Fetching hash for $package_attr on $2" + + expr="(import { system = \"$2\"; }).$package_attr.src.url" + url=$(NIX_PATH=.. nix-instantiate --eval -E "$expr" | jq -r) + + sha256=$(nix-prefetch-url "$url") + hash=$(nix --extra-experimental-features nix-command hash convert --to sri --hash-algo sha256 "$sha256") + + echo -e " cp${1} = \"${hash}\";" >>$outfile + echo +} + +for system in "x86_64-linux" "aarch64-linux"; do + echo "${system} = {" >>$outfile + for python_version in "311" "312" "313" "314"; do + prefetch "$python_version" "$system" + done + echo "};" >>$outfile +done diff --git a/pkgs/development/python-modules/nvidia-cutlass-dsl/default.nix b/pkgs/development/python-modules/nvidia-cutlass-dsl/default.nix new file mode 100644 index 000000000000..59f4cd405157 --- /dev/null +++ b/pkgs/development/python-modules/nvidia-cutlass-dsl/default.nix @@ -0,0 +1,45 @@ +{ + lib, + buildPythonPackage, + fetchPypi, + + # dependencies + nvidia-cutlass-dsl-libs-base, +}: + +buildPythonPackage (finalAttrs: { + pname = "nvidia-cutlass-dsl"; + inherit (nvidia-cutlass-dsl-libs-base) version; + format = "wheel"; + + # Universal metadata-only wheel that just pulls in `nvidia-cutlass-dsl-libs-base` + # (which actually ships the Python code and the bundled MLIR/CUDA runtime libs). + src = fetchPypi { + pname = "nvidia_cutlass_dsl"; + inherit (finalAttrs) version; + format = "wheel"; + python = "py3"; + dist = "py3"; + hash = "sha256-aO0bY8p0quh5VQEtqd/X/arkcTKdACiyKbhBxxksz1I="; + }; + + dependencies = [ + nvidia-cutlass-dsl-libs-base + ]; + + # The `cutlass` module is provided by `nvidia-cutlass-dsl-libs-base` via a `.pth` file. + # pythonImportsCheck uses PYTHONPATH, which Python's site.py doesn't scan for `.pth` files. + dontUsePythonImportsCheck = true; + + # No tests in the Pypi archive + doCheck = false; + + meta = { + description = "NVIDIA CUTLASS Python DSL"; + homepage = "https://github.com/NVIDIA/cutlass"; + changelog = "https://github.com/NVIDIA/cutlass/blob/v${finalAttrs.version}/CHANGELOG.md"; + license = lib.licenses.unfreeRedistributable; # NVIDIA Proprietary + maintainers = with lib.maintainers; [ GaetanLepage ]; + platforms = lib.platforms.linux; + }; +}) diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 03ef1a894184..ab9c91f5b506 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11453,6 +11453,12 @@ self: super: with self; { nvidia-cutlass = callPackage ../development/python-modules/nvidia-cutlass { }; + nvidia-cutlass-dsl = callPackage ../development/python-modules/nvidia-cutlass-dsl { }; + + nvidia-cutlass-dsl-libs-base = + callPackage ../development/python-modules/nvidia-cutlass-dsl-libs-base + { }; + nvidia-dlprof-pytorch-nvtx = callPackage ../development/python-modules/nvidia-dlprof-pytorch-nvtx { };