1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
|
{ lib
, buildPythonPackage
, python
, fetchpatch
, fetchFromGitHub
, addOpenGLRunpath
, cmake
, cudaPackages
, llvmPackages
, pybind11
, gtest
, zlib
, ncurses
, libxml2
, lit
, filelock
, torchWithRocm
, pytest
, pytestCheckHook
, pythonRelaxDepsHook
, pkgsTargetTarget
}:
let
pname = "triton";
version = "2.0.0";
inherit (cudaPackages) cuda_cudart backendStdenv;
# A time may come we'll want to be cross-friendly
#
# Short explanation: we need pkgsTargetTarget, because we use string
# interpolation instead of buildInputs.
#
# Long explanation: OpenAI/triton downloads and vendors a copy of NVidia's
# ptxas compiler. We're not running this ptxas on the build machine, but on
# the user's machine, i.e. our Target platform. The second "Target" in
# pkgsTargetTarget maybe doesn't matter, because ptxas compiles programs to
# be executed on the GPU.
# Cf. https://nixos.org/manual/nixpkgs/unstable/#sec-cross-infra
ptxas = "${pkgsTargetTarget.cudaPackages.cuda_nvcc}/bin/ptxas";
llvm = (llvmPackages.llvm.override {
llvmTargetsToBuild = [ "NATIVE" "NVPTX" ];
# Upstream CI sets these too:
# targetProjects = [ "mlir" ];
extraCMakeFlags = [
"-DLLVM_INSTALL_UTILS=ON"
];
});
in
buildPythonPackage {
inherit pname version;
format = "setuptools";
src = fetchFromGitHub {
owner = "openai";
repo = pname;
rev = "v${version}";
hash = "sha256-9GZzugab+Pdt74Dj6zjlEzjj4BcJ69rzMJmqcVMxsKU=";
};
patches = [
# Prerequisite for llvm15 patch
(fetchpatch {
url = "https://github.com/openai/triton/commit/2aba985daaa70234823ea8f1161da938477d3e02.patch";
hash = "sha256-LGv0+Ut2WYPC4Ksi4803Hwmhi3FyQOF9zElJc/JCobk=";
})
(fetchpatch {
url = "https://github.com/openai/triton/commit/e3941f9d09cdd31529ba4a41018cfc0096aafea6.patch";
hash = "sha256-A+Gor6qzFlGQhVVhiaaYOzqqx8yO2MdssnQS6TIfUWg=";
})
# Source: https://github.com/openai/triton/commit/fc7a8e35819bda632bdcf1cf75fd9abe4d4e077a.patch
# The original patch adds ptxas binary, so we include our own clean copy
# Drop with the next update
./llvm15.patch
# TODO: there have been commits upstream aimed at removing the "torch"
# circular dependency, but the patches fail to apply on the release
# revision. Keeping the link for future reference
# Also cf. https://github.com/openai/triton/issues/1374
# (fetchpatch {
# url = "https://github.com/openai/triton/commit/fc7c0b0e437a191e421faa61494b2ff4870850f1.patch";
# hash = "sha256-f0shIqHJkVvuil2Yku7vuqWFn7VCRKFSFjYRlwx25ig=";
# })
];
postPatch = ''
substituteInPlace python/setup.py \
--replace \
'= get_thirdparty_packages(triton_cache_path)' \
'= os.environ["cmakeFlags"].split()'
''
# Wiring triton=2.0.0 with llcmPackages_rocm.llvm=5.4.3
# Revisit when updating either triton or llvm
+ ''
substituteInPlace CMakeLists.txt \
--replace "nvptx" "NVPTX" \
--replace "LLVM 11" "LLVM"
sed -i '/AddMLIR/a set(MLIR_TABLEGEN_EXE "${llvmPackages.mlir}/bin/mlir-tblgen")' CMakeLists.txt
sed -i '/AddMLIR/a set(MLIR_INCLUDE_DIR ''${MLIR_INCLUDE_DIRS})' CMakeLists.txt
find -iname '*.td' -exec \
sed -i \
-e '\|include "mlir/IR/OpBase.td"|a include "mlir/IR/AttrTypeBase.td"' \
-e 's|include "mlir/Dialect/StandardOps/IR/Ops.td"|include "mlir/Dialect/Func/IR/FuncOps.td"|' \
'{}' ';'
substituteInPlace unittest/CMakeLists.txt --replace "include(GoogleTest)" "find_package(GTest REQUIRED)"
sed -i 's/^include.*$//' unittest/CMakeLists.txt
sed -i '/LINK_LIBS/i NVPTXInfo' lib/Target/PTX/CMakeLists.txt
sed -i '/LINK_LIBS/i NVPTXCodeGen' lib/Target/PTX/CMakeLists.txt
''
# TritonMLIRIR already links MLIRIR. Not transitive?
# + ''
# echo "target_link_libraries(TritonPTX PUBLIC MLIRIR)" >> lib/Target/PTX/CMakeLists.txt
# ''
# Already defined in llvm, when built with -DLLVM_INSTALL_UTILS
+ ''
substituteInPlace bin/CMakeLists.txt \
--replace "add_subdirectory(FileCheck)" ""
rm cmake/FindLLVM.cmake
''
+
(
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${cuda_cudart}/lib/stubs/" ];
quote = x: ''"${x}"'';
oldStr = lib.concatMapStringsSep ", " quote old;
newStr = lib.concatMapStringsSep ", " quote new;
in
''
substituteInPlace python/triton/compiler.py \
--replace '${oldStr}' '${newStr}'
''
)
# Triton seems to be looking up cuda.h
+ ''
sed -i 's|cu_include_dir = os.path.join.*$|cu_include_dir = "${cuda_cudart}/include"|' python/triton/compiler.py
'';
nativeBuildInputs = [
cmake
pythonRelaxDepsHook
# Requires torch (circular dependency) and probably needs GPUs:
# pytestCheckHook
# Note for future:
# These *probably* should go in depsTargetTarget
# ...but we cannot test cross right now anyway
# because we only support cudaPackages on x86_64-linux atm
lit
llvm
llvmPackages.mlir
];
buildInputs = [
gtest
libxml2.dev
ncurses
pybind11
zlib
];
propagatedBuildInputs = [
filelock
];
# Avoid GLIBCXX mismatch with other cuda-enabled python packages
preConfigure = ''
export CC="${backendStdenv.cc}/bin/cc";
export CXX="${backendStdenv.cc}/bin/c++";
# Upstream's setup.py tries to write cache somewhere in ~/
export HOME=$TMPDIR
# Upstream's github actions patch setup.cfg to write base-dir. May be redundant
echo "
[build_ext]
base-dir=$PWD" >> python/setup.cfg
# The rest (including buildPhase) is relative to ./python/
cd python/
# Work around download_and_copy_ptxas()
dst_cuda="$PWD/triton/third_party/cuda/bin"
mkdir -p "$dst_cuda"
ln -s "${ptxas}" "$dst_cuda/"
'';
# CMake is run by setup.py instead
dontUseCmakeConfigure = true;
cmakeFlags = [
"-DMLIR_DIR=${llvmPackages.mlir}/lib/cmake/mlir"
];
postFixup =
let
ptxasDestination = "$out/${python.sitePackages}/triton/third_party/cuda/bin/ptxas";
in
# Setuptools (?) strips runpath and +x flags. Let's just restore the symlink
''
rm -f ${ptxasDestination}
ln -s ${ptxas} ${ptxasDestination}
'';
checkInputs = [
cmake # ctest
];
dontUseSetuptoolsCheck = true;
preCheck =
# build/temp* refers to build_ext.build_temp (looked up in the build logs)
''
(cd /build/source/python/build/temp* ; ctest)
'' # For pytestCheckHook
+ ''
cd test/unit
'';
pythonImportsCheck = [
# Circular dependency on torch
# "triton"
# "triton.language"
];
# Ultimately, torch is our test suite:
passthru.tests = {
inherit torchWithRocm;
};
pythonRemoveDeps = [
# Circular dependency, cf. https://github.com/openai/triton/issues/1374
"torch"
# CLI tools without dist-info
"cmake"
"lit"
];
meta = with lib; {
description = "Development repository for the Triton language and compiler";
homepage = "https://github.com/openai/triton/";
platforms = lib.platforms.unix;
license = licenses.mit;
maintainers = with maintainers; [ SomeoneSerge ];
};
}
|