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