tjh.dev
1{
2 lib,
3 stdenv,
4 addOpenGLRunpath,
5 cudaPackages,
6 buildPythonPackage,
7 fetchurl,
8 isPy38,
9 isPy39,
10 isPy310,
11 isPy311,
12 python,
13 autoPatchelfHook,
14 filelock,
15 lit,
16 pythonRelaxDepsHook,
17 zlib,
18}:
19
20buildPythonPackage rec {
21 pname = "triton";
22 version = "2.1.0";
23 format = "wheel";
24
25 src =
26 let
27 pyVerNoDot = lib.replaceStrings [ "." ] [ "" ] python.pythonVersion;
28 unsupported = throw "Unsupported system";
29 srcs = (import ./binary-hashes.nix version)."${stdenv.system}-${pyVerNoDot}" or unsupported;
30 in
31 fetchurl srcs;
32
33 disabled = !(isPy38 || isPy39 || isPy310 || isPy311);
34
35 pythonRemoveDeps = [
36 "cmake"
37 "torch"
38 ];
39
40 buildInputs = [ zlib ];
41
42 nativeBuildInputs = [
43 pythonRelaxDepsHook # torch and triton refer to each other so this hook is included to mitigate that.
44 autoPatchelfHook
45 ];
46
47 propagatedBuildInputs = [
48 filelock
49 lit
50 zlib
51 ];
52
53 dontStrip = true;
54
55 # If this breaks, consider replacing with "${cuda_nvcc}/bin/ptxas"
56 postFixup =
57 ''
58 chmod +x "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas"
59 ''
60 + (
61 let
62 # Bash was getting weird without linting,
63 # but basically upstream contains [cc, ..., "-lcuda", ...]
64 # and we replace it with [..., "-lcuda", "-L/run/opengl-driver/lib", "-L$stubs", ...]
65 old = [ "-lcuda" ];
66 new = [
67 "-lcuda"
68 "-L${addOpenGLRunpath.driverLink}"
69 "-L${cudaPackages.cuda_cudart}/lib/stubs/"
70 ];
71
72 quote = x: ''"${x}"'';
73 oldStr = lib.concatMapStringsSep ", " quote old;
74 newStr = lib.concatMapStringsSep ", " quote new;
75 in
76 ''
77 substituteInPlace $out/${python.sitePackages}/triton/common/build.py \
78 --replace '${oldStr}' '${newStr}'
79 ''
80 );
81
82 meta = with lib; {
83 description = "A language and compiler for custom Deep Learning operations";
84 homepage = "https://github.com/openai/triton/";
85 changelog = "https://github.com/openai/triton/releases/tag/v${version}";
86 # Includes NVIDIA's ptxas, but redistributions of the binary are not limited.
87 # https://docs.nvidia.com/cuda/eula/index.html
88 # triton's license is MIT.
89 # openai-triton-bin includes ptxas binary, therefore unfreeRedistributable is set.
90 license = with licenses; [
91 unfreeRedistributable
92 mit
93 ];
94 sourceProvenance = with sourceTypes; [ binaryNativeCode ];
95 platforms = [ "x86_64-linux" ];
96 maintainers = with maintainers; [ junjihashimoto ];
97 };
98}