pyrox.dev
1{ lib
2, stdenv
3, fetchFromGitHub
4, cmake
5, rocm-cmake
6, hip
7, openmp
8, gtest ? null
9, buildTests ? false
10, buildExamples ? false
11, gpuTargets ? null # gpuTargets = [ "gfx803" "gfx900" "gfx1030" ... ]
12}:
13
14assert buildTests -> gtest != null;
15
16# Several tests seem to either not compile or have a race condition
17# Undefined reference to symbol '_ZTIN7testing4TestE'
18# Try removing this next update
19assert buildTests == false;
20
21stdenv.mkDerivation rec {
22 pname = "composable_kernel";
23 version = "unstable-2022-11-02";
24
25 outputs = [
26 "out"
27 ] ++ lib.optionals buildTests [
28 "test"
29 ] ++ lib.optionals buildExamples [
30 "example"
31 ];
32
33 src = fetchFromGitHub {
34 owner = "ROCmSoftwarePlatform";
35 repo = "composable_kernel";
36 rev = "79aa3fb1793c265c59d392e916baa851a55521c8";
37 hash = "sha256-vIfMdvRYCTqrjMGSb7gQfodzLw2wf3tGoCAa5jtfbvw=";
38 };
39
40 nativeBuildInputs = [
41 cmake
42 rocm-cmake
43 hip
44 ];
45
46 buildInputs = [
47 openmp
48 ] ++ lib.optionals buildTests [
49 gtest
50 ];
51
52 cmakeFlags = [
53 "-DCMAKE_C_COMPILER=hipcc"
54 "-DCMAKE_CXX_COMPILER=hipcc"
55 ] ++ lib.optionals (gpuTargets != null) [
56 "-DGPU_TARGETS=${lib.strings.concatStringsSep ";" gpuTargets}"
57 ];
58
59 # No flags to build selectively it seems...
60 postPatch = ''
61 substituteInPlace test/CMakeLists.txt \
62 --replace "include(googletest)" ""
63
64 substituteInPlace CMakeLists.txt \
65 --replace "enable_testing()" ""
66 '' + lib.optionalString (!buildTests) ''
67 substituteInPlace CMakeLists.txt \
68 --replace "add_subdirectory(test)" ""
69 '' + lib.optionalString (!buildExamples) ''
70 substituteInPlace CMakeLists.txt \
71 --replace "add_subdirectory(example)" ""
72 '';
73
74 postInstall = ''
75 mkdir -p $out/bin
76 mv bin/ckProfiler $out/bin
77 '' + lib.optionalString buildTests ''
78 mkdir -p $test/bin
79 mv bin/test_* $test/bin
80 '' + lib.optionalString buildExamples ''
81 mkdir -p $example/bin
82 mv bin/example_* $example/bin
83 '';
84
85 meta = with lib; {
86 description = "Performance portable programming model for machine learning tensor operators";
87 homepage = "https://github.com/ROCmSoftwarePlatform/composable_kernel";
88 license = with licenses; [ mit ];
89 maintainers = with maintainers; [ Madouura ];
90 };
91}