tjh.dev
1{
2 lib,
3 config,
4 callPackage,
5 newScope,
6 recurseIntoAttrs,
7 symlinkJoin,
8 fetchFromGitHub,
9 ffmpeg_4,
10 boost179,
11 opencv,
12 libjpeg_turbo,
13 python3Packages,
14 triton-llvm,
15 openmpi,
16 rocmGpuArches ? [ ],
17}:
18
19let
20 outer = lib.makeScope newScope (
21 self:
22 let
23 inherit (self) llvm;
24 pyPackages = python3Packages;
25 openmpi-orig = openmpi;
26 in
27 {
28 inherit rocmGpuArches;
29 buildTests = false;
30 buildBenchmarks = false;
31 stdenv = llvm.rocmClangStdenv;
32
33 rocmPath = self.callPackage ./rocm-path { };
34 rocmUpdateScript = self.callPackage ./update.nix { };
35
36 ## ROCm ##
37 llvm = recurseIntoAttrs (
38 callPackage ./llvm/default.nix {
39 inherit (self) rocm-device-libs rocm-runtime;
40 }
41 );
42 inherit (self.llvm) rocm-merged-llvm clang openmp;
43
44 rocm-core = self.callPackage ./rocm-core { };
45 amdsmi = pyPackages.callPackage ./amdsmi {
46 inherit (self) rocmUpdateScript;
47 };
48
49 rocm-cmake = self.callPackage ./rocm-cmake { };
50
51 rocm-smi = pyPackages.callPackage ./rocm-smi {
52 inherit (self) rocmUpdateScript;
53 };
54
55 rocm-device-libs = self.callPackage ./rocm-device-libs {
56 inherit (llvm) rocm-merged-llvm;
57 };
58
59 rocm-runtime = self.callPackage ./rocm-runtime {
60 inherit (llvm) rocm-merged-llvm;
61 };
62
63 rocm-comgr = self.callPackage ./rocm-comgr {
64 inherit (llvm) rocm-merged-llvm;
65 };
66
67 rocminfo = self.callPackage ./rocminfo { };
68
69 # Unfree
70 hsa-amd-aqlprofile-bin = self.callPackage ./hsa-amd-aqlprofile-bin { };
71
72 rdc = self.callPackage ./rdc { };
73
74 rocm-docs-core = python3Packages.callPackage ./rocm-docs-core { };
75
76 hip-common = self.callPackage ./hip-common { };
77
78 # Eventually will be in the LLVM repo
79 hipcc = self.callPackage ./hipcc {
80 inherit (llvm) rocm-merged-llvm;
81 };
82
83 # Replaces hip, opencl-runtime, and rocclr
84 clr = self.callPackage ./clr { };
85
86 aotriton = self.callPackage ./aotriton { };
87
88 hipify = self.callPackage ./hipify {
89 inherit (llvm)
90 clang
91 rocm-merged-llvm
92 ;
93 };
94
95 # hsakmt was merged into rocm-runtime
96 hsakmt = self.rocm-runtime;
97
98 rocprofiler = self.callPackage ./rocprofiler {
99 inherit (llvm) clang;
100 };
101 rocprofiler-register = self.callPackage ./rocprofiler-register {
102 inherit (llvm) clang;
103 };
104
105 # Needs GCC
106 roctracer = self.callPackage ./roctracer { };
107
108 rocgdb = self.callPackage ./rocgdb { };
109
110 rocdbgapi = self.callPackage ./rocdbgapi { };
111
112 rocr-debug-agent = self.callPackage ./rocr-debug-agent { };
113
114 rocprim = self.callPackage ./rocprim { };
115
116 rocsparse = self.callPackage ./rocsparse { };
117
118 rocthrust = self.callPackage ./rocthrust { };
119
120 rocrand = self.callPackage ./rocrand { };
121
122 hiprand = self.callPackage ./hiprand { };
123
124 rocfft = self.callPackage ./rocfft { };
125
126 mscclpp = self.callPackage ./mscclpp { };
127
128 rccl = self.callPackage ./rccl { };
129
130 # RCCL with sanitizers and tests
131 # Can't have with sanitizer build as dep of other packages without
132 # runtime crashes due to ASAN not loading first
133 rccl-tests = self.callPackage ./rccl {
134 buildTests = true;
135 };
136
137 hipcub = self.callPackage ./hipcub { };
138
139 hipsparse = self.callPackage ./hipsparse { };
140
141 hipfort = self.callPackage ./hipfort { };
142
143 hipfft = self.callPackage ./hipfft { };
144
145 tensile = pyPackages.callPackage ./tensile {
146 inherit (self)
147 rocmUpdateScript
148 clr
149 ;
150 };
151
152 rocblas = self.callPackage ./rocblas {
153 buildTests = true;
154 buildBenchmarks = true;
155 };
156
157 rocsolver = self.callPackage ./rocsolver { };
158
159 rocwmma = self.callPackage ./rocwmma { };
160
161 rocalution = self.callPackage ./rocalution { };
162
163 rocmlir-rock = self.callPackage ./rocmlir {
164 buildRockCompiler = true;
165 };
166 rocmlir = self.rocmlir-rock;
167
168 hipsolver = self.callPackage ./hipsolver { };
169
170 hipblas-common = self.callPackage ./hipblas-common { };
171
172 hipblas = self.callPackage ./hipblas { };
173
174 hipblaslt = self.callPackage ./hipblaslt { };
175
176 # hipTensor - Only supports GFX9
177
178 composable_kernel_base = self.callPackage ./composable_kernel/base.nix { };
179 composable_kernel = self.callPackage ./composable_kernel { };
180
181 ck4inductor = pyPackages.callPackage ./composable_kernel/ck4inductor.nix {
182 inherit (self) composable_kernel;
183 inherit (llvm) rocm-merged-llvm;
184 };
185
186 half = self.callPackage ./half { };
187
188 miopen = self.callPackage ./miopen {
189 boost = boost179.override { enableStatic = true; };
190 };
191
192 miopen-hip = self.miopen;
193
194 migraphx = self.callPackage ./migraphx { };
195
196 rpp = self.callPackage ./rpp { };
197
198 rpp-hip = self.rpp.override {
199 useOpenCL = false;
200 useCPU = false;
201 };
202
203 rpp-opencl = self.rpp.override {
204 useOpenCL = true;
205 useCPU = false;
206 };
207
208 rpp-cpu = self.rpp.override {
209 useOpenCL = false;
210 useCPU = true;
211 };
212
213 mivisionx = self.callPackage ./mivisionx {
214 opencv = opencv.override { enablePython = true; };
215 # TODO: Remove this pin in ROCm 6.4+
216 # FFMPEG support was improved in https://github.com/ROCm/MIVisionX/pull/1460
217 ffmpeg = ffmpeg_4;
218 # Unfortunately, rocAL needs a custom libjpeg-turbo until further notice
219 # See: https://github.com/ROCm/MIVisionX/issues/1051
220 libjpeg_turbo = libjpeg_turbo.overrideAttrs {
221 version = "2.0.6.1";
222 src = fetchFromGitHub {
223 owner = "rrawther";
224 repo = "libjpeg-turbo";
225 rev = "640d7ee1917fcd3b6a5271aa6cf4576bccc7c5fb";
226 sha256 = "sha256-T52whJ7nZi8jerJaZtYInC2YDN0QM+9tUDqiNr6IsNY=";
227 };
228 # overwrite all patches, since patches for newer version do not apply
229 patches = [ ./0001-Compile-transupp.c-as-part-of-the-library.patch ];
230 };
231 };
232
233 mivisionx-hip = self.mivisionx.override {
234 rpp = self.rpp-hip;
235 useOpenCL = false;
236 useCPU = false;
237 };
238
239 mivisionx-cpu = self.mivisionx.override {
240 rpp = self.rpp-cpu;
241 useOpenCL = false;
242 useCPU = true;
243 };
244
245 # Even if config.rocmSupport is false we need rocmSupport true
246 # version of ucc/ucx in openmpi in this package set
247 openmpi = openmpi-orig.override (
248 prev:
249 let
250 ucx = prev.ucx.override {
251 enableCuda = false;
252 enableRocm = true;
253 };
254 in
255 {
256 inherit ucx;
257 ucc = prev.ucc.override {
258 enableCuda = false;
259 inherit ucx;
260 };
261 }
262 );
263 mpi = self.openmpi;
264
265 triton-llvm = triton-llvm.overrideAttrs {
266 src = fetchFromGitHub {
267 owner = "llvm";
268 repo = "llvm-project";
269 # make sure this matches triton llvm rel branch hash for now
270 # https://github.com/triton-lang/triton/blob/release/3.2.x/cmake/llvm-hash.txt
271 rev = "86b69c31642e98f8357df62c09d118ad1da4e16a";
272 hash = "sha256-W/mQwaLGx6/rIBjdzUTIbWrvGjdh7m4s15f70fQ1/hE=";
273 };
274 pname = "triton-llvm-rocm";
275 patches = [ ]; # FIXME: https://github.com/llvm/llvm-project//commit/84837e3cc1cf17ed71580e3ea38299ed2bfaa5f6.patch doesn't apply, may need to rebase
276 };
277
278 triton = pyPackages.callPackage ./triton { rocmPackages = self; };
279
280 ## Meta ##
281 # Emulate common ROCm meta layout
282 # These are mainly for users. I strongly suggest NOT using these in nixpkgs derivations
283 # Don't put these into `propagatedBuildInputs` unless you want PATH/PYTHONPATH issues!
284 # See: https://rocm.docs.amd.com/en/docs-5.7.1/_images/image.004.png
285 # See: https://rocm.docs.amd.com/en/docs-5.7.1/deploy/linux/os-native/package_manager_integration.html
286 meta = with self; rec {
287 rocm-developer-tools = symlinkJoin {
288 name = "rocm-developer-tools-meta";
289 paths = [
290 hsa-amd-aqlprofile-bin
291 rocm-core
292 rocr-debug-agent
293 roctracer
294 rocdbgapi
295 rocprofiler
296 rocgdb
297 rocm-language-runtime
298 ];
299 };
300 rocm-ml-sdk = symlinkJoin {
301 name = "rocm-ml-sdk-meta";
302 paths = [
303 rocm-core
304 miopen-hip
305 rocm-hip-sdk
306 rocm-ml-libraries
307 ];
308 };
309 rocm-ml-libraries = symlinkJoin {
310 name = "rocm-ml-libraries-meta";
311 paths = [
312 llvm.clang
313 llvm.mlir
314 llvm.openmp
315 rocm-core
316 miopen-hip
317 rocm-hip-libraries
318 ];
319 };
320 rocm-hip-sdk = symlinkJoin {
321 name = "rocm-hip-sdk-meta";
322 paths = [
323 rocprim
324 rocalution
325 hipfft
326 rocm-core
327 hipcub
328 hipblas
329 hipblaslt
330 rocrand
331 rocfft
332 rocsparse
333 rccl
334 rocthrust
335 rocblas
336 hipsparse
337 hipfort
338 rocwmma
339 hipsolver
340 rocsolver
341 rocm-hip-libraries
342 rocm-hip-runtime-devel
343 ];
344 };
345 rocm-hip-libraries = symlinkJoin {
346 name = "rocm-hip-libraries-meta";
347 paths = [
348 rocblas
349 hipfort
350 rocm-core
351 rocsolver
352 rocalution
353 rocrand
354 hipblas
355 hipblaslt
356 rocfft
357 hipfft
358 rccl
359 rocsparse
360 hipsparse
361 hipsolver
362 rocm-hip-runtime
363 ];
364 };
365 rocm-openmp-sdk = symlinkJoin {
366 name = "rocm-openmp-sdk-meta";
367 paths = [
368 rocm-core
369 llvm.clang
370 llvm.mlir
371 llvm.openmp # openmp-extras-devel (https://github.com/ROCm/aomp)
372 rocm-language-runtime
373 ];
374 };
375 rocm-opencl-sdk = symlinkJoin {
376 name = "rocm-opencl-sdk-meta";
377 paths = [
378 rocm-core
379 rocm-runtime
380 clr
381 clr.icd
382 rocm-opencl-runtime
383 ];
384 };
385 rocm-opencl-runtime = symlinkJoin {
386 name = "rocm-opencl-runtime-meta";
387 paths = [
388 rocm-core
389 clr
390 clr.icd
391 rocm-language-runtime
392 ];
393 };
394 rocm-hip-runtime-devel = symlinkJoin {
395 name = "rocm-hip-runtime-devel-meta";
396 paths = [
397 clr
398 rocm-core
399 hipify
400 rocm-cmake
401 llvm.clang
402 llvm.mlir
403 llvm.openmp
404 rocm-runtime
405 rocm-hip-runtime
406 ];
407 };
408 rocm-hip-runtime = symlinkJoin {
409 name = "rocm-hip-runtime-meta";
410 paths = [
411 rocm-core
412 rocminfo
413 clr
414 rocm-language-runtime
415 ];
416 };
417 rocm-language-runtime = symlinkJoin {
418 name = "rocm-language-runtime-meta";
419 paths = [
420 rocm-runtime
421 rocm-core
422 rocm-comgr
423 llvm.openmp # openmp-extras-runtime (https://github.com/ROCm/aomp)
424 ];
425 };
426 rocm-all = symlinkJoin {
427 name = "rocm-all-meta";
428 paths = [
429 rocm-developer-tools
430 rocm-ml-sdk
431 rocm-ml-libraries
432 rocm-hip-sdk
433 rocm-hip-libraries
434 rocm-openmp-sdk
435 rocm-opencl-sdk
436 rocm-opencl-runtime
437 rocm-hip-runtime-devel
438 rocm-hip-runtime
439 rocm-language-runtime
440 ];
441 };
442 };
443
444 rocm-tests = self.callPackage ./rocm-tests {
445 rocmPackages = self;
446 };
447 }
448 // lib.optionalAttrs config.allowAliases {
449 rocm-thunk = throw ''
450 'rocm-thunk' has been removed. It's now part of the ROCm runtime.
451 ''; # Added 2025-3-16
452
453 clang-ocl = throw ''
454 'clang-ocl' has been deprecated upstream. Use ROCm's clang directly.
455 ''; # Added 2025-3-16
456
457 miopengemm = throw ''
458 'miopengemm' has been deprecated.
459 ''; # Added 2024-3-3
460
461 miopen-opencl = throw ''
462 'miopen-opencl' has been deprecated.
463 ''; # Added 2024-3-3
464
465 mivisionx-opencl = throw ''
466 'mivisionx-opencl' has been deprecated.
467 Other versions of mivisionx are still available.
468 ''; # Added 2024-3-24
469 }
470 );
471 scopeForArches =
472 arches:
473 outer.overrideScope (
474 _final: prev: {
475 clr = prev.clr.override {
476 localGpuTargets = arches;
477 };
478 }
479 );
480in
481outer
482// builtins.listToAttrs (
483 builtins.map (arch: {
484 name = arch;
485 value = scopeForArches [ arch ];
486 }) outer.clr.gpuTargets
487)
488// {
489 gfx9 = scopeForArches [
490 "gfx906"
491 "gfx908"
492 "gfx90a"
493 "gfx942"
494 ];
495 gfx10 = scopeForArches [
496 "gfx1010"
497 "gfx1030"
498 ];
499 gfx11 = scopeForArches [
500 "gfx1100"
501 "gfx1101"
502 "gfx1102"
503 ];
504}