tjh.dev
fork
Configure Feed
Select the types of activity you want to include in your feed.
nixpkgs mirror (for testing)
github.com/NixOS/nixpkgs
nix
fork
Configure Feed
Select the types of activity you want to include in your feed.
1{ lib
2, stdenv
3, fetchFromGitHub
4, rocmUpdateScript
5, cmake
6, rocm-cmake
7, hip
8, openmp
9, gtest
10, rocblas
11, texlive
12, doxygen
13, sphinx
14, python3Packages
15, buildDocs ? true
16, buildTests ? false
17, buildExtendedTests ? false
18, buildBenchmarks ? false
19, buildSamples ? false
20, gpuTargets ? [ ] # gpuTargets = [ "gfx908:xnack-" "gfx90a:xnack-" "gfx90a:xnack+" ... ]
21}:
22
23let
24 latex = lib.optionalAttrs buildDocs texlive.combine {
25 inherit (texlive) scheme-small
26 latexmk
27 tex-gyre
28 fncychap
29 wrapfig
30 capt-of
31 framed
32 needspace
33 tabulary
34 varwidth
35 titlesec;
36 };
37in stdenv.mkDerivation (finalAttrs: {
38 pname = "rocwmma";
39 version = "5.4.3";
40
41 outputs = [
42 "out"
43 ] ++ lib.optionals buildDocs [
44 "doc"
45 ] ++ lib.optionals (buildTests || buildBenchmarks) [
46 "test"
47 ] ++ lib.optionals buildBenchmarks [
48 "benchmark"
49 ] ++ lib.optionals buildSamples [
50 "sample"
51 ];
52
53 src = fetchFromGitHub {
54 owner = "ROCmSoftwarePlatform";
55 repo = "rocWMMA";
56 rev = "rocm-${finalAttrs.version}";
57 hash = "sha256-HUJPb6IahBgl/v+W4kXludBTNAjRm8k6v0jxKAX+qZM=";
58 };
59
60 patches = lib.optionals (buildTests || buildBenchmarks) [
61 ./0000-dont-fetch-googletest.patch
62 ];
63
64 nativeBuildInputs = [
65 cmake
66 rocm-cmake
67 hip
68 ];
69
70 buildInputs = [
71 openmp
72 ] ++ lib.optionals (buildTests || buildBenchmarks) [
73 gtest
74 rocblas
75 ] ++ lib.optionals buildDocs [
76 latex
77 doxygen
78 sphinx
79 python3Packages.sphinx-rtd-theme
80 python3Packages.breathe
81 ];
82
83 cmakeFlags = [
84 "-DCMAKE_CXX_COMPILER=hipcc"
85 "-DROCWMMA_BUILD_TESTS=${if buildTests || buildBenchmarks then "ON" else "OFF"}"
86 "-DROCWMMA_BUILD_VALIDATION_TESTS=ON"
87 "-DROCWMMA_BUILD_SAMPLES=${if buildSamples then "ON" else "OFF"}"
88 "-DROCWMMA_VALIDATE_WITH_ROCBLAS=ON"
89 # Manually define CMAKE_INSTALL_<DIR>
90 # See: https://github.com/NixOS/nixpkgs/pull/197838
91 "-DCMAKE_INSTALL_BINDIR=bin"
92 "-DCMAKE_INSTALL_LIBDIR=lib"
93 "-DCMAKE_INSTALL_INCLUDEDIR=include"
94 ] ++ lib.optionals (gpuTargets != [ ]) [
95 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
96 ] ++ lib.optionals buildExtendedTests [
97 "-DROCWMMA_BUILD_EXTENDED_TESTS=ON"
98 ] ++ lib.optionals buildBenchmarks [
99 "-DROCWMMA_BUILD_BENCHMARK_TESTS=ON"
100 "-DROCWMMA_BENCHMARK_WITH_ROCBLAS=ON"
101 ];
102
103 postPatch = lib.optionalString buildDocs ''
104 patchShebangs docs/*.sh
105 '';
106
107 # Unfortunately, it seems like we have to call make on this manually
108 # -DROCWMMA_BUILD_DOCS=ON is invalid, despite being on the README
109 postBuild = lib.optionalString buildDocs ''
110 export HOME=$(mktemp -d)
111 ../docs/run_doc.sh
112 '';
113
114 postInstall = lib.optionalString buildDocs ''
115 mv ../docs/source/_build/html $out/share/doc/rocwmma
116 mv ../docs/source/_build/latex/rocWMMA.pdf $out/share/doc/rocwmma
117 '' + lib.optionalString (buildTests || buildBenchmarks) ''
118 mkdir -p $test/bin
119 mv $out/bin/{*_test,*-validate} $test/bin
120 '' + lib.optionalString buildBenchmarks ''
121 mkdir -p $benchmark/bin
122 mv $out/bin/*-bench $benchmark/bin
123 '' + lib.optionalString buildSamples ''
124 mkdir -p $sample/bin
125 mv $out/bin/sgemmv $sample/bin
126 mv $out/bin/simple_gemm $sample/bin
127 mv $out/bin/simple_dlrm $sample/bin
128 '' + lib.optionalString (buildTests || buildBenchmarks || buildSamples) ''
129 rm -rf $out/bin
130 '';
131
132 passthru.updateScript = rocmUpdateScript {
133 name = finalAttrs.pname;
134 owner = finalAttrs.src.owner;
135 repo = finalAttrs.src.repo;
136 };
137
138 meta = with lib; {
139 description = "Mixed precision matrix multiplication and accumulation";
140 homepage = "https://github.com/ROCmSoftwarePlatform/rocWMMA";
141 license = with licenses; [ mit ];
142 maintainers = teams.rocm.members;
143 platforms = platforms.linux;
144 broken = versions.minor finalAttrs.version != versions.minor hip.version;
145 };
146})