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
|
{
lib,
stdenv,
fetchFromGitHub,
rocmUpdateScript,
cmake,
rocm-cmake,
clr,
libxml2,
libedit,
rocm-comgr,
rocm-device-libs,
rocm-runtime,
zstd,
zlib,
ncurses,
python3Packages,
buildRockCompiler ? false,
buildTests ? false, # `argument of type 'NoneType' is not iterable`
}:
# FIXME: rocmlir has an entire separate LLVM build in a subdirectory this is silly
# It seems to be forked from AMD's own LLVM
# If possible reusing the rocmPackages.llvm build would be better
# Would have to confirm it is compatible with ROCm's tagged LLVM.
# Fairly likely it's not given AMD's track record with forking their own software in incompatible ways
# in subdirs
# Theoretically, we could have our MLIR have an output
# with the source and built objects so that we can just
# use it as the external LLVM repo for this
let
suffix = if buildRockCompiler then "-rock" else "";
llvmNativeTarget =
if stdenv.hostPlatform.isx86_64 then
"X86"
else if stdenv.hostPlatform.isAarch64 then
"AArch64"
else
throw "Unsupported ROCm LLVM platform";
in
stdenv.mkDerivation (finalAttrs: {
pname = "rocmlir${suffix}";
version = "7.2.3";
outputs = [
"out"
]
++ lib.optionals (!buildRockCompiler) [
"external"
];
src = fetchFromGitHub {
owner = "ROCm";
repo = "rocMLIR";
rev = "rocm-${finalAttrs.version}";
hash = "sha256-0OvQT8pX6GbEqUwuauKGI66IHw8dsnt5mIijnzYyiRc=";
};
nativeBuildInputs = [
clr
cmake
rocm-cmake
python3Packages.python
python3Packages.tomli
];
buildInputs = [
libxml2
libedit
rocm-comgr
rocm-runtime
rocm-device-libs
];
propagatedBuildInputs = [
zstd
zlib
ncurses
];
cmakeFlags = [
(lib.cmakeFeature "LLVM_TARGETS_TO_BUILD" "AMDGPU${
lib.optionalString (!buildRockCompiler) ";${llvmNativeTarget}"
}")
(lib.cmakeFeature "LLVM_USE_LINKER" "lld")
(lib.cmakeFeature "LLVM_ENABLE_ZSTD" "FORCE_ON")
(lib.cmakeFeature "LLVM_ENABLE_ZLIB" "FORCE_ON")
(lib.cmakeBool "LLVM_ENABLE_LIBCXX" true)
(lib.cmakeBool "LLVM_ENABLE_TERMINFO" true)
(lib.cmakeFeature "ROCM_PATH" "${clr}")
# Manually define CMAKE_INSTALL_<DIR>
# See: https://github.com/NixOS/nixpkgs/pull/197838
(lib.cmakeFeature "CMAKE_INSTALL_BINDIR" "bin")
(lib.cmakeFeature "CMAKE_INSTALL_LIBDIR" "lib")
(lib.cmakeFeature "CMAKE_INSTALL_INCLUDEDIR" "include")
(lib.cmakeBool "MLIR_ENABLE_ROCM_RUNNER" (!buildRockCompiler))
(lib.cmakeBool "BUILD_FAT_LIBROCKCOMPILER" buildRockCompiler)
]
++ lib.optionals buildRockCompiler [
(lib.cmakeBool "LLVM_INCLUDE_TESTS" false)
]
++ lib.optionals (!buildRockCompiler) [
(lib.cmakeFeature "ROCM_TEST_CHIPSET" "gfx900")
];
postPatch = ''
patchShebangs mlir
patchShebangs external/llvm-project/mlir/lib/Dialect/GPU/AmdDeviceLibsIncGen.py
# Fixes mlir/lib/Analysis/BufferDependencyAnalysis.cpp:41:19: error: redefinition of 'read'
substituteInPlace mlir/lib/Analysis/BufferDependencyAnalysis.cpp \
--replace-fail "enum EffectType { read, write, unknown };" "enum class EffectType { read, write, unknown };"
substituteInPlace mlir/utils/performance/common/CMakeLists.txt \
--replace-fail " PATHS /opt/rocm" ""
'';
dontBuild = true;
doCheck = true;
# Certain libs aren't being generated, try enabling tests next update
checkTarget =
if buildRockCompiler then
"librockCompiler"
else if buildTests then
"check-rocmlir"
else
"check-rocmlir-build-only";
postInstall =
let
libPath = lib.makeLibraryPath [
zstd
zlib
ncurses
clr
stdenv.cc.cc
];
in
lib.optionals (!buildRockCompiler) ''
mkdir -p $external/lib
cp -a external/llvm-project/llvm/lib/{*.a*,*.so*} $external/lib
patchelf --set-rpath $external/lib:$out/lib:${libPath} $external/lib/*.so*
patchelf --set-rpath $out/lib:$external/lib:${libPath} $out/{bin/*,lib/*.so*}
'';
passthru.updateScript = rocmUpdateScript {
inherit finalAttrs;
page = "tags";
};
meta = {
description = "MLIR-based convolution and GEMM kernel generator";
homepage = "https://github.com/ROCm/rocMLIR";
license = with lib.licenses; [ asl20 ];
teams = [ lib.teams.rocm ];
platforms = lib.platforms.linux;
};
})
|