Merge pull request #227420 from GaetanLepage/torchbin-update

    # Should appear BEFORE lit, since lit passes through python rather

    # than the python environment we make.

    kleePython

    (lit.override { python3 = kleePython; })

    (lit.override { python = kleePython; })

  ];

  cmakeFlags = let

{ lib, python3 }:

{ lib

, buildPythonPackage

, fetchPypi

, python

}:

python3.pkgs.buildPythonApplication rec {

buildPythonPackage rec {

  pname = "lit";

  version = "15.0.6";

  src = python3.pkgs.fetchPypi {

  src = fetchPypi {

    inherit pname version;

    hash = "sha256-S06OQfDmDyutls21HxyQ016ku3FTTsDOP8Di67d9f+k=";

  };

  passthru = {

    python = python3;

    inherit python;

  };

  # Non-standard test suite. Needs custom checkPhase.

{ lib

, stdenv

, addOpenGLRunpath

, cudaPackages

, buildPythonPackage

, fetchurl

, isPy38

, isPy39

, isPy310

, isPy311

, python

, autoPatchelfHook

, filelock

, lit

, pythonRelaxDepsHook

, zlib

}:

buildPythonPackage rec {

  pname = "triton";

  version = "2.0.0";

  format = "wheel";

  src =

    let pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion;

        unsupported = throw "Unsupported system";

        srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported;

    in fetchurl srcs;

  disabled = !(isPy38 || isPy39 || isPy310 || isPy311);

  pythonRemoveDeps = [ "cmake" "torch" ];

  buildInputs = [ zlib ];

  nativeBuildInputs = [

    pythonRelaxDepsHook # torch and triton refer to each other so this hook is included to mitigate that.

    autoPatchelfHook

  ];

  propagatedBuildInputs = [

    filelock

    lit

    zlib

  ];

  dontStrip = true;

  # If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"

  postFixup = ''

    chmod +x "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas"

  '' +

  (let

    # Bash was getting weird without linting,

    # but basically upstream contains [cc, ..., "-lcuda", ...]

    # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]

    old = [ "-lcuda" ];

    new = [ "-lcuda" "-L${addOpenGLRunpath.driverLink}" "-L${cudaPackages.cuda_cudart}/lib/stubs/" ];

    quote = x: ''"${x}"'';

    oldStr = lib.concatMapStringsSep ", " quote old;

    newStr = lib.concatMapStringsSep ", " quote new;

  in

    ''

      substituteInPlace $out/${python.sitePackages}/triton/compiler.py \

        --replace '${oldStr}' '${newStr}'

    '');

  meta = with lib; {

    description = "A language and compiler for custom Deep Learning operations";

    homepage = "https://github.com/openai/triton/";

    changelog = "https://github.com/openai/triton/releases/tag/v${version}";

    # Includes NVIDIA's ptxas, but redistributions of the binary are not limited.

    # https://docs.nvidia.com/cuda/eula/index.html

    # triton's license is MIT.

    # openai-triton-bin includes ptxas binary, therefore unfreeRedistributable is set.

    license = with licenses; [ unfreeRedistributable mit ];

    sourceProvenance = with sourceTypes; [ binaryNativeCode ];

    platforms = [ "x86_64-linux" ];

    maintainers = with maintainers; [ junjihashimoto ];

  };

}

# Warning: Need to update at the same time as torch-bin

#

# Precompiled wheels can be found at:

# https://download.pytorch.org/whl/torch_stable.html

# To add a new version, run "prefetch.sh 'new-version'" to paste the generated file as follows.

version : builtins.getAttr version {

  "2.0.0" = {

    x86_64-linux-38 = {

      name = "triton-2.0.0-1-cp38-cp38-linux_x86_64.whl";

      url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";

      hash = "sha256-nUl4KYt0/PWadf5x5TXAkrAjCIkzsvHfkz7DJhXkvu8=";

    };

    x86_64-linux-39 = {

      name = "triton-2.0.0-1-cp39-cp39-linux_x86_64.whl";

      url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";

      hash = "sha256-dPEYwStDf7LKJeGgR1kXO1F1gvz0x74RkTMWx2QhNlY=";

    };

    x86_64-linux-310 = {

      name = "triton-2.0.0-1-cp310-cp310-linux_x86_64.whl";

      url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";

      hash = "sha256-OIBu6WY/Sw981keQ6WxXk3QInlj0mqxKZggSGqVeJQU=";

    };

    x86_64-linux-311 = {

      name = "triton-2.0.0-1-cp311-cp311-linux_x86_64.whl";

      url = "https://download.pytorch.org/whl/triton-2.0.0-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl";

      hash = "sha256-ImlBx7hZUhnd71mh/bgh6MdEKJoTJBXd1YT6zt60dbE=";

    };

  };

}

#!/usr/bin/env nix-shell

#!nix-shell -i bash -p nix-prefetch-scripts

set -eou pipefail

version=$1

linux_bucket="https://download.pytorch.org/whl"

url_and_key_list=(

  "x86_64-linux-38 $linux_bucket/triton-${version}-1-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp38-cp38-linux_x86_64.whl"

  "x86_64-linux-39 $linux_bucket/triton-${version}-1-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp39-cp39-linux_x86_64.whl"

  "x86_64-linux-310 $linux_bucket/triton-${version}-1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp310-cp310-linux_x86_64.whl"

  "x86_64-linux-311 $linux_bucket/triton-${version}-1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl triton-${version}-cp311-cp311-linux_x86_64.whl"

)

hashfile=binary-hashes-"$version".nix

echo "  \"$version\" = {" >> $hashfile

for url_and_key in "${url_and_key_list[@]}"; do

  key=$(echo "$url_and_key" | cut -d' ' -f1)

  url=$(echo "$url_and_key" | cut -d' ' -f2)

  name=$(echo "$url_and_key" | cut -d' ' -f3)

  echo "prefetching ${url}..."

  hash=$(nix hash to-sri --type sha256 `nix-prefetch-url "$url" --name "$name"`)

  cat << EOF >> $hashfile

    $key = {

      name = "$name";

      url = "$url";

      hash = "$hash";

    };

EOF

  echo

done

echo "  };" >> $hashfile

echo "done."