tjh.dev
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5 fetchpatch,
6 rocmUpdateScript,
7 cmake,
8 rocm-cmake,
9 clr,
10 diffutils,
11 python3,
12 tensile,
13 boost,
14 msgpack-cxx,
15 libxml2,
16 gtest,
17 gfortran,
18 openmp,
19 gitMinimal,
20 amd-blis,
21 zstd,
22 roctracer,
23 hipblas-common,
24 hipblaslt,
25 python3Packages,
26 rocm-smi,
27 pkg-config,
28 removeReferencesTo,
29 buildTensile ? true,
30 buildTests ? true,
31 buildBenchmarks ? true,
32 tensileSepArch ? true,
33 tensileLazyLib ? true,
34 withHipBlasLt ? true,
35 gpuTargets ? (clr.localGpuTargets or clr.gpuTargets),
36}:
37
38let
39 gpuTargets' = lib.concatStringsSep ";" gpuTargets;
40in
41stdenv.mkDerivation (finalAttrs: {
42 pname = "rocblas${clr.gpuArchSuffix}";
43 version = "7.1.1";
44
45 src = fetchFromGitHub {
46 owner = "ROCm";
47 repo = "rocBLAS";
48 rev = "rocm-${finalAttrs.version}";
49 hash = "sha256-obKypbYmnSeOtOr7g0pxOz02YfzZ0bGugTtznkeHz14=";
50 };
51
52 outputs = [ "out" ] ++ lib.optional buildBenchmarks "benchmark" ++ lib.optional buildTests "test";
53
54 nativeBuildInputs = [
55 cmake
56 # no ninja, it buffers console output and nix times out long periods of no output
57 rocm-cmake
58 clr
59 gitMinimal
60 pkg-config
61 removeReferencesTo
62 ]
63 ++ lib.optionals buildTensile [
64 tensile
65 ];
66
67 buildInputs = [
68 python3
69 hipblas-common
70 roctracer
71 openmp
72 amd-blis
73 ]
74 ++ lib.optionals withHipBlasLt [
75 hipblaslt
76 ]
77 ++ lib.optionals buildTensile [
78 zstd
79 msgpack-cxx
80 libxml2
81 python3Packages.msgpack
82 python3Packages.zstandard
83 ]
84 ++ lib.optionals (buildTests || buildBenchmarks) [
85 gtest
86 gfortran
87 rocm-smi
88 ]
89 ++ lib.optionals (buildTensile || buildTests || buildBenchmarks) [
90 python3Packages.pyyaml
91 ];
92
93 env.CXXFLAGS = "-fopenmp -I${lib.getDev boost}/include -I${hipblas-common}/include -I${roctracer}/include";
94 # Fails to link tests with undefined symbol: cblas_*
95 env.LDFLAGS =
96 "-Wl,--as-needed -lzstd" + lib.optionalString (buildTests || buildBenchmarks) " -lcblas";
97 env.TENSILE_ROCM_ASSEMBLER_PATH = "${stdenv.cc}/bin/clang++";
98
99 cmakeFlags = [
100 (lib.cmakeFeature "Boost_INCLUDE_DIR" "${lib.getDev boost}/include") # msgpack FindBoost fails to find boost
101 (lib.cmakeFeature "CMAKE_EXECUTE_PROCESS_COMMAND_ECHO" "STDERR")
102 (lib.cmakeFeature "CMAKE_Fortran_COMPILER" "${lib.getBin gfortran}/bin/gfortran")
103 (lib.cmakeFeature "CMAKE_Fortran_COMPILER_AR" "${lib.getBin gfortran}/bin/ar")
104 (lib.cmakeFeature "CMAKE_Fortran_COMPILER_RANLIB" "${lib.getBin gfortran}/bin/ranlib")
105 (lib.cmakeFeature "python" "python3")
106 (lib.cmakeFeature "SUPPORTED_TARGETS" gpuTargets')
107 (lib.cmakeFeature "AMDGPU_TARGETS" gpuTargets')
108 (lib.cmakeFeature "GPU_TARGETS" gpuTargets')
109 (lib.cmakeBool "BUILD_WITH_TENSILE" buildTensile)
110 (lib.cmakeBool "ROCM_SYMLINK_LIBS" false)
111 (lib.cmakeFeature "ROCBLAS_TENSILE_LIBRARY_DIR" "lib/rocblas")
112 (lib.cmakeBool "BUILD_WITH_HIPBLASLT" withHipBlasLt)
113 (lib.cmakeBool "BUILD_CLIENTS_TESTS" buildTests)
114 (lib.cmakeBool "BUILD_CLIENTS_BENCHMARKS" buildBenchmarks)
115 (lib.cmakeBool "BUILD_CLIENTS_SAMPLES" buildBenchmarks)
116 (lib.cmakeBool "BUILD_OFFLOAD_COMPRESS" true)
117 # # Temporarily set variables to work around upstream CMakeLists issue
118 # # Can be removed once https://github.com/ROCm/rocm-cmake/issues/121 is fixed
119 "-DCMAKE_INSTALL_BINDIR=bin"
120 "-DCMAKE_INSTALL_INCLUDEDIR=include"
121 "-DCMAKE_INSTALL_LIBDIR=lib"
122 ]
123 ++ lib.optionals buildTensile [
124 "-DCPACK_SET_DESTDIR=OFF"
125 "-DLINK_BLIS=ON"
126 "-DBLAS_LIBRARY=${amd-blis}/lib/libblis-mt.so"
127 "-DBLIS_INCLUDE_DIR=${amd-blis}/include/blis/"
128 "-DBLA_PREFER_PKGCONFIG=ON"
129 "-DTensile_CODE_OBJECT_VERSION=default"
130 "-DTensile_LOGIC=asm_full"
131 "-DTensile_LIBRARY_FORMAT=msgpack"
132 (lib.cmakeBool "BUILD_WITH_PIP" false)
133 (lib.cmakeBool "Tensile_SEPARATE_ARCHITECTURES" tensileSepArch)
134 (lib.cmakeBool "Tensile_LAZY_LIBRARY_LOADING" tensileLazyLib)
135 ];
136
137 patches = [
138 (fetchpatch {
139 name = "Extend-rocBLAS-HIP-ISA-compatibility.patch";
140 url = "https://github.com/GZGavinZhao/rocm-libraries/commit/2850f22f80f90c9e498f520608a82989a4932ec3.patch";
141 hash = "sha256-SPsdEwGe+r8bQudkChRzBDAgu3tPQWFweZCgzh+4nOE=";
142 stripLen = 2;
143 })
144 ];
145
146 # Pass $NIX_BUILD_CORES to Tensile
147 postPatch = ''
148 substituteInPlace cmake/build-options.cmake \
149 --replace-fail 'Tensile_CPU_THREADS ""' 'Tensile_CPU_THREADS "$ENV{NIX_BUILD_CORES}"'
150 ''
151 # Workaround: libblis detection uses broken absolute paths
152 # TODO: upstream a proper fix
153 + ''
154 substituteInPlace clients/CMakeLists.txt \
155 --replace-fail "if ( NOT WIN32 )" "if(OFF)" \
156 --replace-fail "else() # WIN32" "elseif(OFF)"
157 ''
158 # Fixes sh: line 1: /usr/bin/diff: No such file or directory
159 # /build/source/clients/gtest/../include/testing_logging.hpp:1117: Failure
160 + lib.optionalString buildTests ''
161 substituteInPlace clients/include/testing_logging.hpp \
162 --replace-fail "/usr/bin/diff" "${lib.getExe' diffutils "diff"}"
163 '';
164
165 postInstall =
166 # tensile isn't needed at runtime and pulls in ~400MB of python deps
167 ''
168 remove-references-to -t ${tensile} \
169 "$out/lib/librocblas.so."*
170 ''
171 + lib.optionalString buildBenchmarks ''
172 moveToOutput "bin/*-tune" "$benchmark"
173 moveToOutput "bin/*-bench" "$benchmark"
174 moveToOutput "bin/*example*" "$benchmark"
175 cp "$out/bin/"*.{yaml,txt} "$benchmark/bin"
176 ''
177 + lib.optionalString buildTests ''
178 moveToOutput "bin/*test*" "$test"
179 cp "$out/bin/"*.{yaml,txt} "$test/bin"
180 ''
181 + ''
182 if [ -d $out/bin ]; then
183 rm $out/bin/*.{yaml,txt} || true
184 rmdir $out/bin
185 fi
186 '';
187
188 passthru = {
189 amdgpu_targets = gpuTargets';
190 updateScript = rocmUpdateScript {
191 name = finalAttrs.pname;
192 inherit (finalAttrs.src) owner;
193 inherit (finalAttrs.src) repo;
194 };
195 };
196
197 enableParallelBuilding = true;
198 requiredSystemFeatures = [ "big-parallel" ];
199
200 meta = {
201 description = "BLAS implementation for ROCm platform";
202 homepage = "https://github.com/ROCm/rocBLAS";
203 license = with lib.licenses; [ mit ];
204 teams = [ lib.teams.rocm ];
205 platforms = lib.platforms.linux;
206 };
207})