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