tjh.dev
Clone of https://github.com/NixOS/nixpkgs.git (to stress-test knotserver)
1{ lib
2, stdenv
3, fetchFromGitHub
4, rocmUpdateScript
5, cmake
6, rocm-cmake
7, clr
8, openmp
9, clang-tools-extra
10, git
11, gtest
12, zstd
13, buildTests ? false
14, buildExamples ? false
15, gpuTargets ? [ ] # gpuTargets = [ "gfx803" "gfx900" "gfx1030" ... ]
16}:
17
18stdenv.mkDerivation (finalAttrs: {
19 pname = "composable_kernel";
20 version = "6.0.2";
21
22 outputs = [
23 "out"
24 ] ++ lib.optionals buildTests [
25 "test"
26 ] ++ lib.optionals buildExamples [
27 "example"
28 ];
29
30 src = fetchFromGitHub {
31 owner = "ROCm";
32 repo = "composable_kernel";
33 rev = "rocm-${finalAttrs.version}";
34 hash = "sha256-NCqMganmNyQfz3X+KQOrfrimnrgd3HbAGK5DeC4+J+o=";
35 };
36
37 nativeBuildInputs = [
38 git
39 cmake
40 rocm-cmake
41 clr
42 clang-tools-extra
43 zstd
44 ];
45
46 buildInputs = [ openmp ];
47
48 cmakeFlags = [
49 "-DCMAKE_C_COMPILER=hipcc"
50 "-DCMAKE_CXX_COMPILER=hipcc"
51 ] ++ lib.optionals (gpuTargets != [ ]) [
52 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
53 "-DAMDGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
54 ] ++ lib.optionals buildTests [
55 "-DGOOGLETEST_DIR=${gtest.src}" # Custom linker names
56 ];
57
58 # No flags to build selectively it seems...
59 postPatch = lib.optionalString (!buildTests) ''
60 substituteInPlace CMakeLists.txt \
61 --replace "add_subdirectory(test)" ""
62 '' + lib.optionalString (!buildExamples) ''
63 substituteInPlace CMakeLists.txt \
64 --replace "add_subdirectory(example)" ""
65 '' + ''
66 substituteInPlace CMakeLists.txt \
67 --replace "add_subdirectory(profiler)" ""
68 ''
69 ;
70
71 postInstall = ''
72 zstd --rm $out/lib/libdevice_operations.a
73 '' + lib.optionalString buildTests ''
74 mkdir -p $test/bin
75 mv $out/bin/test_* $test/bin
76 '' + lib.optionalString buildExamples ''
77 mkdir -p $example/bin
78 mv $out/bin/example_* $example/bin
79 '';
80
81 passthru.updateScript = rocmUpdateScript {
82 name = finalAttrs.pname;
83 owner = finalAttrs.src.owner;
84 repo = finalAttrs.src.repo;
85 };
86
87 # Times out otherwise
88 requiredSystemFeatures = [ "big-parallel" ];
89
90 meta = with lib; {
91 description = "Performance portable programming model for machine learning tensor operators";
92 homepage = "https://github.com/ROCm/composable_kernel";
93 license = with licenses; [ mit ];
94 maintainers = teams.rocm.members;
95 platforms = platforms.linux;
96 broken = versions.minor finalAttrs.version != versions.minor stdenv.cc.version || versionAtLeast finalAttrs.version "7.0.0";
97 };
98})