tjh.dev
1{ lib
2, stdenv
3, fetchFromGitHub
4, rocmUpdateScript
5, cmake
6, rocm-cmake
7, rocm-opencl-runtime
8, texlive
9, doxygen
10, sphinx
11, openblas
12, python3Packages
13, buildDocs ? true
14, buildTests ? false
15, buildBenchmarks ? false
16}:
17
18let
19 latex = lib.optionalAttrs buildDocs texlive.combine {
20 inherit (texlive) scheme-small
21 latexmk
22 tex-gyre
23 fncychap
24 wrapfig
25 capt-of
26 framed
27 needspace
28 tabulary
29 varwidth
30 titlesec;
31 };
32in stdenv.mkDerivation (finalAttrs: {
33 pname = "miopengemm";
34 version = "5.4.3";
35
36 outputs = [
37 "out"
38 ] ++ lib.optionals buildDocs [
39 "doc"
40 ] ++ lib.optionals buildTests [
41 "test"
42 ] ++ lib.optionals buildBenchmarks [
43 "benchmark"
44 ];
45
46 src = fetchFromGitHub {
47 owner = "ROCmSoftwarePlatform";
48 repo = "MIOpenGEMM";
49 rev = "rocm-${finalAttrs.version}";
50 hash = "sha256-AiRzOMYRA/0nbQomyq4oOEwNZdkPYWRA2W6QFlctvFc=";
51 };
52
53 nativeBuildInputs = [
54 cmake
55 rocm-cmake
56 ];
57
58 buildInputs = [
59 rocm-opencl-runtime
60 ] ++ lib.optionals buildDocs [
61 latex
62 doxygen
63 sphinx
64 python3Packages.sphinx-rtd-theme
65 python3Packages.breathe
66 ] ++ lib.optionals buildTests [
67 openblas
68 ];
69
70 cmakeFlags = [
71 # Manually define CMAKE_INSTALL_<DIR>
72 # See: https://github.com/NixOS/nixpkgs/pull/197838
73 "-DCMAKE_INSTALL_BINDIR=bin"
74 "-DCMAKE_INSTALL_LIBDIR=lib"
75 "-DCMAKE_INSTALL_INCLUDEDIR=include"
76 ] ++ lib.optionals buildTests [
77 "-DOPENBLAS=ON"
78 ] ++ lib.optionals buildBenchmarks [
79 "-DAPI_BENCH_MIOGEMM=ON"
80 # Needs https://github.com/CNugteren/CLBlast
81 # "-DAPI_BENCH_CLBLAST=ON"
82 # Needs https://github.com/openai/triton
83 # "-DAPI_BENCH_ISAAC=ON"
84 ];
85
86 # Unfortunately, it seems like we have to call make on these manually
87 postBuild = lib.optionalString buildDocs ''
88 export HOME=$(mktemp -d)
89 make doc
90 '' + lib.optionalString buildTests ''
91 make check
92 '' + lib.optionalString buildBenchmarks ''
93 make examples
94 '';
95
96 postInstall = lib.optionalString buildDocs ''
97 mv ../doc/html $out/share/doc/miopengemm
98 mv ../doc/pdf/miopengemm.pdf $out/share/doc/miopengemm
99 '' + lib.optionalString buildTests ''
100 mkdir -p $test/bin
101 find tests -executable -type f -exec mv {} $test/bin \;
102 patchelf --set-rpath ${lib.makeLibraryPath finalAttrs.buildInputs}:$out/lib $test/bin/*
103 '' + lib.optionalString buildBenchmarks ''
104 mkdir -p $benchmark/bin
105 find examples -executable -type f -exec mv {} $benchmark/bin \;
106 patchelf --set-rpath ${lib.makeLibraryPath finalAttrs.buildInputs}:$out/lib $benchmark/bin/*
107 '';
108
109 passthru.updateScript = rocmUpdateScript {
110 name = finalAttrs.pname;
111 owner = finalAttrs.src.owner;
112 repo = finalAttrs.src.repo;
113 };
114
115 meta = with lib; {
116 description = "OpenCL general matrix multiplication API for ROCm";
117 homepage = "https://github.com/ROCmSoftwarePlatform/MIOpenGEMM";
118 license = with licenses; [ mit ];
119 maintainers = teams.rocm.members;
120 platforms = platforms.linux;
121 broken = versions.minor finalAttrs.version != versions.minor stdenv.cc.version;
122 };
123})