pyrox.dev
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 rocmUpdateScript,
6 cmake,
7 rocm-cmake,
8 rocminfo,
9 clr,
10 git,
11 libxml2,
12 libedit,
13 zstd,
14 zlib,
15 ncurses,
16 python3Packages,
17 buildRockCompiler ? false,
18 buildTests ? false, # `argument of type 'NoneType' is not iterable`
19}:
20
21# FIXME: rocmlir has an entire separate LLVM build in a subdirectory this is silly
22# It seems to be forked from AMD's own LLVM
23# If possible reusing the rocmPackages.llvm build would be better
24# Would have to confirm it is compatible with ROCm's tagged LLVM.
25# Fairly likely it's not given AMD's track record with forking their own software in incompatible ways
26# in subdirs
27
28# Theoretically, we could have our MLIR have an output
29# with the source and built objects so that we can just
30# use it as the external LLVM repo for this
31let
32 suffix = if buildRockCompiler then "-rock" else "";
33
34 llvmNativeTarget =
35 if stdenv.hostPlatform.isx86_64 then
36 "X86"
37 else if stdenv.hostPlatform.isAarch64 then
38 "AArch64"
39 else
40 throw "Unsupported ROCm LLVM platform";
41in
42stdenv.mkDerivation (finalAttrs: {
43 pname = "rocmlir${suffix}";
44 version = "6.3.3";
45
46 outputs = [
47 "out"
48 ]
49 ++ lib.optionals (!buildRockCompiler) [
50 "external"
51 ];
52
53 src = fetchFromGitHub {
54 owner = "ROCm";
55 repo = "rocMLIR";
56 rev = "rocm-${finalAttrs.version}";
57 hash = "sha256-0SQ6uLDRfVfdCX+8a7D6pu6dYlFvX0HFzCDEvlKYfak=";
58 };
59
60 nativeBuildInputs = [
61 cmake
62 rocm-cmake
63 clr
64 python3Packages.python
65 python3Packages.tomli
66 ];
67
68 buildInputs = [
69 git
70 libxml2
71 libedit
72 ];
73
74 propagatedBuildInputs = [
75 zstd
76 zlib
77 ncurses
78 ];
79
80 patches = [
81 ./initparamdata-sort-const.patch
82 ];
83
84 cmakeFlags = [
85 "-DLLVM_TARGETS_TO_BUILD=AMDGPU;${llvmNativeTarget}"
86 "-DCMAKE_BUILD_TYPE=Release"
87 "-DLLVM_USE_LINKER=lld"
88 "-DLLVM_ENABLE_ZSTD=FORCE_ON"
89 "-DLLVM_ENABLE_ZLIB=FORCE_ON"
90 "-DLLVM_ENABLE_LIBCXX=ON"
91 "-DLLVM_ENABLE_TERMINFO=ON"
92 "-DROCM_PATH=${clr}"
93 # Manually define CMAKE_INSTALL_<DIR>
94 # See: https://github.com/NixOS/nixpkgs/pull/197838
95 "-DCMAKE_INSTALL_BINDIR=bin"
96 "-DCMAKE_INSTALL_LIBDIR=lib"
97 "-DCMAKE_INSTALL_INCLUDEDIR=include"
98 (lib.cmakeBool "BUILD_FAT_LIBROCKCOMPILER" buildRockCompiler)
99 ]
100 ++ lib.optionals (!buildRockCompiler) [
101 "-DROCM_TEST_CHIPSET=gfx000"
102 ];
103
104 postPatch = ''
105 patchShebangs mlir
106 patchShebangs external/llvm-project/mlir/lib/Dialect/GPU/AmdDeviceLibsIncGen.py
107
108 # Fixes mlir/lib/Analysis/BufferDependencyAnalysis.cpp:41:19: error: redefinition of 'read'
109 substituteInPlace mlir/lib/Analysis/BufferDependencyAnalysis.cpp \
110 --replace-fail "enum EffectType { read, write, unknown };" "enum class EffectType { read, write, unknown };"
111
112 # remove when no longer required
113 substituteInPlace mlir/test/{e2e/generateE2ETest.py,fusion/e2e/generate-fusion-tests.py} \
114 --replace-fail "\"/opt/rocm/bin" "\"${rocminfo}/bin"
115
116 substituteInPlace mlir/utils/performance/common/CMakeLists.txt \
117 --replace-fail "/opt/rocm" "${clr}"
118 '';
119
120 dontBuild = true;
121 doCheck = true;
122
123 # Certain libs aren't being generated, try enabling tests next update
124 checkTarget =
125 if buildRockCompiler then
126 "librockCompiler"
127 else if buildTests then
128 "check-rocmlir"
129 else
130 "check-rocmlir-build-only";
131
132 postInstall =
133 let
134 libPath = lib.makeLibraryPath [
135 zstd
136 zlib
137 ncurses
138 clr
139 stdenv.cc.cc
140 ];
141 in
142 lib.optionals (!buildRockCompiler) ''
143 mkdir -p $external/lib
144 cp -a external/llvm-project/llvm/lib/{*.a*,*.so*} $external/lib
145 patchelf --set-rpath $external/lib:$out/lib:${libPath} $external/lib/*.so*
146 patchelf --set-rpath $out/lib:$external/lib:${libPath} $out/{bin/*,lib/*.so*}
147 '';
148
149 passthru.updateScript = rocmUpdateScript {
150 name = finalAttrs.pname;
151 inherit (finalAttrs.src) owner;
152 inherit (finalAttrs.src) repo;
153 page = "tags?per_page=4";
154 };
155
156 meta = with lib; {
157 description = "MLIR-based convolution and GEMM kernel generator";
158 homepage = "https://github.com/ROCm/rocMLIR";
159 license = with licenses; [ asl20 ];
160 teams = [ teams.rocm ];
161 platforms = platforms.linux;
162 };
163})