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 ? [ ],
18}:
19
20stdenv.mkDerivation (finalAttrs: {
21 pname = "rocwmma";
22 version = "6.3.3";
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-kih3hn6QhcMmyj9n8f8eO+RIgKQgWKIuzg8fb0eoRPE=";
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 "-DOpenMP_C_INCLUDE_DIR=${openmp.dev}/include"
65 "-DOpenMP_CXX_INCLUDE_DIR=${openmp.dev}/include"
66 "-DOpenMP_omp_LIBRARY=${openmp}/lib"
67 "-DROCWMMA_BUILD_TESTS=${if buildTests || buildBenchmarks then "ON" else "OFF"}"
68 "-DROCWMMA_BUILD_SAMPLES=${if buildSamples then "ON" else "OFF"}"
69 # Manually define CMAKE_INSTALL_<DIR>
70 # See: https://github.com/NixOS/nixpkgs/pull/197838
71 "-DCMAKE_INSTALL_BINDIR=bin"
72 "-DCMAKE_INSTALL_LIBDIR=lib"
73 "-DCMAKE_INSTALL_INCLUDEDIR=include"
74 ]
75 ++ lib.optionals (gpuTargets != [ ]) [
76 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
77 ]
78 ++ lib.optionals buildExtendedTests [
79 "-DROCWMMA_BUILD_EXTENDED_TESTS=ON"
80 ]
81 ++ lib.optionals buildBenchmarks [
82 "-DROCWMMA_BUILD_BENCHMARK_TESTS=ON"
83 "-DROCWMMA_BENCHMARK_WITH_ROCBLAS=ON"
84 ];
85
86 postInstall =
87 lib.optionalString (buildTests || buildBenchmarks) ''
88 mkdir -p $test/bin
89 mv $out/bin/{*_test,*-validate} $test/bin
90 ''
91 + lib.optionalString buildBenchmarks ''
92 mkdir -p $benchmark/bin
93 mv $out/bin/*-bench $benchmark/bin
94 ''
95 + lib.optionalString buildSamples ''
96 mkdir -p $sample/bin
97 mv $out/bin/sgemmv $sample/bin
98 mv $out/bin/simple_gemm $sample/bin
99 mv $out/bin/simple_dlrm $sample/bin
100 ''
101 + lib.optionalString (buildTests || buildBenchmarks || buildSamples) ''
102 rm -rf $out/bin
103 '';
104
105 passthru.updateScript = rocmUpdateScript {
106 name = finalAttrs.pname;
107 inherit (finalAttrs.src) owner;
108 inherit (finalAttrs.src) repo;
109 };
110
111 meta = with lib; {
112 description = "Mixed precision matrix multiplication and accumulation";
113 homepage = "https://github.com/ROCm/rocWMMA";
114 license = with licenses; [ mit ];
115 teams = [ teams.rocm ];
116 platforms = platforms.linux;
117 };
118})