tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 rocmUpdateScript,
6 cmake,
7 rocm-cmake,
8 clr,
9 libxml2,
10 libedit,
11 rocm-comgr,
12 rocm-device-libs,
13 rocm-runtime,
14 zstd,
15 zlib,
16 ncurses,
17 python3Packages,
18 buildRockCompiler ? false,
19 buildTests ? false, # `argument of type 'NoneType' is not iterable`
20}:
21
22# FIXME: rocmlir has an entire separate LLVM build in a subdirectory this is silly
23# It seems to be forked from AMD's own LLVM
24# If possible reusing the rocmPackages.llvm build would be better
25# Would have to confirm it is compatible with ROCm's tagged LLVM.
26# Fairly likely it's not given AMD's track record with forking their own software in incompatible ways
27# in subdirs
28
29# Theoretically, we could have our MLIR have an output
30# with the source and built objects so that we can just
31# use it as the external LLVM repo for this
32let
33 suffix = if buildRockCompiler then "-rock" else "";
34
35 llvmNativeTarget =
36 if stdenv.hostPlatform.isx86_64 then
37 "X86"
38 else if stdenv.hostPlatform.isAarch64 then
39 "AArch64"
40 else
41 throw "Unsupported ROCm LLVM platform";
42in
43stdenv.mkDerivation (finalAttrs: {
44 pname = "rocmlir${suffix}";
45 version = "7.1.1";
46
47 outputs = [
48 "out"
49 ]
50 ++ lib.optionals (!buildRockCompiler) [
51 "external"
52 ];
53
54 src = fetchFromGitHub {
55 owner = "ROCm";
56 repo = "rocMLIR";
57 rev = "rocm-${finalAttrs.version}";
58 hash = "sha256-A9vUvsEZrZlNEW4cscF66L48rJQ1zJYmIzwXQ2QzJ3s=";
59 };
60
61 nativeBuildInputs = [
62 clr
63 cmake
64 rocm-cmake
65 python3Packages.python
66 python3Packages.tomli
67 ];
68
69 buildInputs = [
70 libxml2
71 libedit
72 rocm-comgr
73 rocm-runtime
74 rocm-device-libs
75 ];
76
77 propagatedBuildInputs = [
78 zstd
79 zlib
80 ncurses
81 ];
82
83 cmakeFlags = [
84 (lib.cmakeFeature "LLVM_TARGETS_TO_BUILD" "AMDGPU${
85 lib.optionalString (!buildRockCompiler) ";${llvmNativeTarget}"
86 }")
87 (lib.cmakeFeature "LLVM_USE_LINKER" "lld")
88 (lib.cmakeFeature "LLVM_ENABLE_ZSTD" "FORCE_ON")
89 (lib.cmakeFeature "LLVM_ENABLE_ZLIB" "FORCE_ON")
90 (lib.cmakeBool "LLVM_ENABLE_LIBCXX" true)
91 (lib.cmakeBool "LLVM_ENABLE_TERMINFO" true)
92 (lib.cmakeFeature "ROCM_PATH" "${clr}")
93 # Manually define CMAKE_INSTALL_<DIR>
94 # See: https://github.com/NixOS/nixpkgs/pull/197838
95 (lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "bin")
96 (lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "lib")
97 (lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "include")
98 (lib.cmakeBool "MLIR_ENABLE_ROCM_RUNNER" (!buildRockCompiler))
99 (lib.cmakeBool "BUILD_FAT_LIBROCKCOMPILER" buildRockCompiler)
100 ]
101 ++ lib.optionals buildRockCompiler [
102 (lib.cmakeBool "LLVM_INCLUDE_TESTS" false)
103 ]
104 ++ lib.optionals (!buildRockCompiler) [
105 (lib.cmakeFeature "ROCM_TEST_CHIPSET" "gfx900")
106 ];
107
108 postPatch = ''
109 patchShebangs mlir
110 patchShebangs external/llvm-project/mlir/lib/Dialect/GPU/AmdDeviceLibsIncGen.py
111
112 # Fixes mlir/lib/Analysis/BufferDependencyAnalysis.cpp:41:19: error: redefinition of 'read'
113 substituteInPlace mlir/lib/Analysis/BufferDependencyAnalysis.cpp \
114 --replace-fail "enum EffectType { read, write, unknown };" "enum class EffectType { read, write, unknown };"
115
116 substituteInPlace mlir/utils/performance/common/CMakeLists.txt \
117 --replace-fail " PATHS /opt/rocm" ""
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 = {
157 description = "MLIR-based convolution and GEMM kernel generator";
158 homepage = "https://github.com/ROCm/rocMLIR";
159 license = with lib.licenses; [ asl20 ];
160 teams = [ lib.teams.rocm ];
161 platforms = lib.platforms.linux;
162 };
163})