pyrox.dev
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 rocmUpdateScript,
6 pkg-config,
7 cmake,
8 rocm-cmake,
9 clr,
10 openmp,
11 rocblas,
12 hipblas-common,
13 hipblas,
14 hipblaslt,
15 rocmlir,
16 miopen,
17 protobuf,
18 abseil-cpp,
19 half,
20 nlohmann_json,
21 msgpack,
22 sqlite,
23 oneDNN,
24 blaze,
25 texliveSmall,
26 doxygen,
27 sphinx,
28 docutils,
29 ghostscript,
30 python3Packages,
31 buildDocs ? false,
32 buildTests ? false,
33 gpuTargets ? clr.gpuTargets,
34}:
35
36let
37 latex = lib.optionalAttrs buildDocs (
38 texliveSmall.withPackages (
39 ps: with ps; [
40 latexmk
41 tex-gyre
42 fncychap
43 wrapfig
44 capt-of
45 framed
46 needspace
47 tabulary
48 varwidth
49 titlesec
50 epstopdf
51 ]
52 )
53 );
54 oneDNN' = oneDNN.overrideAttrs rec {
55 version = "2.7.5";
56 src = fetchFromGitHub {
57 owner = "oneapi-src";
58 repo = "oneDNN";
59 tag = "v${version}";
60 hash = "sha256-oMPBORAdL2rk2ewyUrInYVHYBRvuvNX4p4rwykO3Rhs=";
61 };
62 };
63in
64stdenv.mkDerivation (finalAttrs: {
65 pname = "migraphx";
66 version = "6.3.3";
67
68 outputs = [
69 "out"
70 ]
71 ++ lib.optionals buildDocs [
72 "doc"
73 ]
74 ++ lib.optionals buildTests [
75 "test"
76 ];
77
78 src = fetchFromGitHub {
79 owner = "ROCm";
80 repo = "AMDMIGraphX";
81 rev = "rocm-${finalAttrs.version}";
82 hash = "sha256-h9cTbrMwHeRGVJS/uHQnCXplNcrBqxbhwz2AcAEso0M=";
83 };
84
85 nativeBuildInputs = [
86 pkg-config
87 cmake
88 rocm-cmake
89 clr
90 python3Packages.python
91 ]
92 ++ lib.optionals buildDocs [
93 latex
94 doxygen
95 sphinx
96 docutils
97 ghostscript
98 python3Packages.sphinx-rtd-theme
99 python3Packages.breathe
100 ];
101
102 buildInputs = [
103 openmp
104 rocblas
105 hipblas-common
106 hipblas
107 hipblaslt
108 rocmlir
109 miopen
110 protobuf
111 half
112 nlohmann_json
113 msgpack
114 sqlite
115 oneDNN'
116 blaze
117 python3Packages.pybind11
118 python3Packages.onnx
119 ];
120
121 LDFLAGS = "-Wl,--allow-shlib-undefined";
122
123 cmakeFlags = [
124 "-DMIGRAPHX_ENABLE_GPU=ON"
125 "-DMIGRAPHX_ENABLE_CPU=ON"
126 "-DMIGRAPHX_ENABLE_FPGA=ON"
127 "-DMIGRAPHX_ENABLE_MLIR=OFF" # LLVM or rocMLIR mismatch?
128 "-DCMAKE_C_COMPILER=amdclang"
129 "-DCMAKE_CXX_COMPILER=amdclang++"
130 "-DCMAKE_VERBOSE_MAKEFILE=ON"
131 "-DEMBED_USE=CArrays" # Fixes error with lld
132 "-DDMIGRAPHX_ENABLE_PYTHON=ON"
133 "-DROCM_PATH=${clr}"
134 "-DHIP_ROOT_DIR=${clr}"
135 # migraphx relies on an incompatible fork of composable_kernel
136 # migraphxs relies on miopen which relies on current composable_kernel
137 # impossible to build with this ON; we can't link both of them even if we package both
138 "-DMIGRAPHX_USE_COMPOSABLEKERNEL=OFF"
139 "-DOpenMP_C_INCLUDE_DIR=${openmp.dev}/include"
140 "-DOpenMP_CXX_INCLUDE_DIR=${openmp.dev}/include"
141 "-DOpenMP_omp_LIBRARY=${openmp}/lib"
142 # Manually define CMAKE_INSTALL_<DIR>
143 # See: https://github.com/NixOS/nixpkgs/pull/197838
144 "-DCMAKE_INSTALL_BINDIR=bin"
145 "-DCMAKE_INSTALL_LIBDIR=lib"
146 "-DCMAKE_INSTALL_INCLUDEDIR=include"
147 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
148 ];
149
150 postPatch = ''
151 export CXXFLAGS+=" -w -isystem${rocmlir}/include/rocmlir -I${half}/include -I${abseil-cpp}/include -I${hipblas-common}/include"
152 patchShebangs tools
153
154 # `error: '__clang_hip_runtime_wrapper.h' file not found [clang-diagnostic-error]`
155 substituteInPlace CMakeLists.txt \
156 --replace "set(MIGRAPHX_TIDY_ERRORS ALL)" ""
157 ''
158 + lib.optionalString (!buildDocs) ''
159 substituteInPlace CMakeLists.txt \
160 --replace "add_subdirectory(doc)" ""
161 ''
162 + lib.optionalString (!buildTests) ''
163 substituteInPlace CMakeLists.txt \
164 --replace "add_subdirectory(test)" ""
165 '';
166
167 # Unfortunately, it seems like we have to call make on this manually
168 preInstall = lib.optionalString buildDocs ''
169 export HOME=$(mktemp -d)
170 make -j$NIX_BUILD_CORES doc
171 cd ../doc/pdf
172 make -j$NIX_BUILD_CORES
173 cd -
174 '';
175
176 postInstall =
177 lib.optionalString buildDocs ''
178 mv ../doc/html $out/share/doc/migraphx
179 mv ../doc/pdf/MIGraphX.pdf $out/share/doc/migraphx
180 ''
181 + lib.optionalString buildTests ''
182 mkdir -p $test/bin
183 mv bin/test_* $test/bin
184 patchelf $test/bin/test_* --shrink-rpath --allowed-rpath-prefixes "$NIX_STORE"
185 '';
186
187 passthru.updateScript = rocmUpdateScript {
188 name = finalAttrs.pname;
189 inherit (finalAttrs.src) owner;
190 inherit (finalAttrs.src) repo;
191 };
192
193 meta = with lib; {
194 description = "AMD's graph optimization engine";
195 homepage = "https://github.com/ROCm/AMDMIGraphX";
196 license = with licenses; [ mit ];
197 teams = [ teams.rocm ];
198 platforms = platforms.linux;
199 };
200})