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 rocm-smi,
9 clr,
10 openmp,
11 gtest,
12 rocblas,
13 buildTests ? false, # Will likely fail building because wavefront shifts are not supported for certain archs
14 buildExtendedTests ? false,
15 buildBenchmarks ? false,
16 buildSamples ? false,
17 gpuTargets ? clr.localGpuTargets or [ ],
18}:
19
20stdenv.mkDerivation (finalAttrs: {
21 pname = "rocwmma";
22 version = "7.1.1";
23
24 outputs = [
25 "out"
26 ]
27 ++ lib.optionals (buildTests || buildBenchmarks) [
28 "test"
29 ]
30 ++ lib.optionals buildBenchmarks [
31 "benchmark"
32 ]
33 ++ lib.optionals buildSamples [
34 "sample"
35 ];
36
37 src = fetchFromGitHub {
38 owner = "ROCm";
39 repo = "rocWMMA";
40 rev = "rocm-${finalAttrs.version}";
41 hash = "sha256-6hvdq3J8UhldcyNqXxi9UEF+P0RgG8o5YdIZLzvpgXw=";
42 };
43
44 patches = lib.optionals (buildTests || buildBenchmarks) [
45 ./0000-dont-fetch-googletest.patch
46 ];
47
48 nativeBuildInputs = [
49 cmake
50 rocm-cmake
51 clr
52 ];
53
54 buildInputs = [
55 openmp
56 ]
57 ++ lib.optionals (buildTests || buildBenchmarks) [
58 rocm-smi
59 gtest
60 rocblas
61 ];
62
63 cmakeFlags = [
64 "-DROCWMMA_BUILD_TESTS=${if buildTests || buildBenchmarks then "ON" else "OFF"}"
65 "-DROCWMMA_BUILD_SAMPLES=${if buildSamples then "ON" else "OFF"}"
66 # Manually define CMAKE_INSTALL_<DIR>
67 # See: https://github.com/NixOS/nixpkgs/pull/197838
68 "-DCMAKE_INSTALL_BINDIR=bin"
69 "-DCMAKE_INSTALL_LIBDIR=lib"
70 "-DCMAKE_INSTALL_INCLUDEDIR=include"
71 ]
72 ++ lib.optionals (gpuTargets != [ ]) [
73 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
74 ]
75 ++ lib.optionals buildExtendedTests [
76 "-DROCWMMA_BUILD_EXTENDED_TESTS=ON"
77 ]
78 ++ lib.optionals buildBenchmarks [
79 "-DROCWMMA_BUILD_BENCHMARK_TESTS=ON"
80 "-DROCWMMA_BENCHMARK_WITH_ROCBLAS=ON"
81 ];
82
83 postInstall =
84 lib.optionalString (buildTests || buildBenchmarks) ''
85 mkdir -p $test/bin
86 mv $out/bin/{*_test,*-validate} $test/bin
87 ''
88 + lib.optionalString buildBenchmarks ''
89 mkdir -p $benchmark/bin
90 mv $out/bin/*-bench $benchmark/bin
91 ''
92 + lib.optionalString buildSamples ''
93 mkdir -p $sample/bin
94 mv $out/bin/sgemmv $sample/bin
95 mv $out/bin/simple_gemm $sample/bin
96 mv $out/bin/simple_dlrm $sample/bin
97 ''
98 + lib.optionalString (buildTests || buildBenchmarks || buildSamples) ''
99 rm -rf $out/bin
100 '';
101
102 passthru.updateScript = rocmUpdateScript {
103 name = finalAttrs.pname;
104 inherit (finalAttrs.src) owner;
105 inherit (finalAttrs.src) repo;
106 };
107
108 meta = {
109 description = "Mixed precision matrix multiplication and accumulation";
110 homepage = "https://github.com/ROCm/rocWMMA";
111 license = with lib.licenses; [ mit ];
112 teams = [ lib.teams.rocm ];
113 platforms = lib.platforms.linux;
114 };
115})