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