pyrox.dev
1{
2 lib,
3 stdenv,
4 addDriverRunpath,
5 cudaPackages,
6 buildPythonPackage,
7 fetchurl,
8 python,
9 pythonOlder,
10 autoPatchelfHook,
11 filelock,
12 lit,
13 zlib,
14}:
15
16buildPythonPackage rec {
17 pname = "triton";
18 version = "3.1.0";
19 format = "wheel";
20
21 src =
22 let
23 pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion;
24 unsupported = throw "Unsupported system";
25 srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported;
26 in
27 fetchurl srcs;
28
29 disabled = pythonOlder "3.8";
30
31 pythonRemoveDeps = [
32 "cmake"
33 # torch and triton refer to each other so this hook is included to mitigate that.
34 "torch"
35 ];
36
37 buildInputs = [ zlib ];
38
39 nativeBuildInputs = [
40 autoPatchelfHook
41 ];
42
43 propagatedBuildInputs = [
44 filelock
45 lit
46 zlib
47 ];
48
49 dontStrip = true;
50
51 # If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"
52 postFixup = ''
53 mkdir -p $out/${python.sitePackages}/triton/third_party/cuda/bin/
54 ln -s ${cudaPackages.cuda_nvcc}/bin/ptxas $out/${python.sitePackages}/triton/third_party/cuda/bin/
55 '';
56
57 meta = with lib; {
58 description = "Language and compiler for custom Deep Learning operations";
59 homepage = "https://github.com/triton-lang/triton/";
60 changelog = "https://github.com/triton-lang/triton/releases/tag/v${version}";
61 # Includes NVIDIA's ptxas, but redistributions of the binary are not limited.
62 # https://docs.nvidia.com/cuda/eula/index.html
63 # triton's license is MIT.
64 # triton-bin includes ptxas binary, therefore unfreeRedistributable is set.
65 license = with licenses; [
66 unfreeRedistributable
67 mit
68 ];
69 sourceProvenance = with sourceTypes; [ binaryNativeCode ];
70 platforms = [ "x86_64-linux" ];
71 maintainers = with maintainers; [ junjihashimoto ];
72 };
73}