python3Packages.nvidia-cutlass-dsl: init at 4.5.0 (#522219)

2026-05-26 13:10:51 +00:00
parent 5ff973b4c8 1a4f8bb9de
commit 97581bc7ab
4 changed files with 181 additions and 0 deletions
@@ -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;
+  };
+})
@@ -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 <nixpkgs> { 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
@@ -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;
+  };
+})
@@ -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
      { };