tjh.dev
Clone of https://github.com/NixOS/nixpkgs.git (to stress-test knotserver)
1{ lib
2, stdenv
3, fetchFromGitHub
4, unstableGitUpdater
5, cmake
6, rocm-cmake
7, hip
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 = "unstable-2023-01-16";
19
20 outputs = [
21 "out"
22 ] ++ lib.optionals buildTests [
23 "test"
24 ] ++ lib.optionals buildExamples [
25 "example"
26 ];
27
28 # ROCm 5.6 should release composable_kernel as stable with a tag in the future
29 src = fetchFromGitHub {
30 owner = "ROCmSoftwarePlatform";
31 repo = "composable_kernel";
32 rev = "80e05267417f948e4f7e63c0fe807106d9a0c0ef";
33 hash = "sha256-+c0E2UtlG/abweLwCWWjNHDO5ZvSIVKwwwettT9mqR4=";
34 };
35
36 nativeBuildInputs = [
37 cmake
38 rocm-cmake
39 hip
40 clang-tools-extra
41 ];
42
43 buildInputs = [
44 openmp
45 ];
46
47 cmakeFlags = [
48 "-DCMAKE_C_COMPILER=hipcc"
49 "-DCMAKE_CXX_COMPILER=hipcc"
50 ] ++ lib.optionals (gpuTargets != [ ]) [
51 "-DGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
52 "-DAMDGPU_TARGETS=${lib.concatStringsSep ";" gpuTargets}"
53 ] ++ lib.optionals buildTests [
54 "-DGOOGLETEST_DIR=${gtest.src}" # Custom linker names
55 ];
56
57 # No flags to build selectively it seems...
58 postPatch = lib.optionalString (!buildTests) ''
59 substituteInPlace CMakeLists.txt \
60 --replace "add_subdirectory(test)" ""
61 '' + lib.optionalString (!buildExamples) ''
62 substituteInPlace CMakeLists.txt \
63 --replace "add_subdirectory(example)" ""
64 '';
65
66 postInstall = lib.optionalString buildTests ''
67 mkdir -p $test/bin
68 mv $out/bin/test_* $test/bin
69 '' + lib.optionalString buildExamples ''
70 mkdir -p $example/bin
71 mv $out/bin/example_* $example/bin
72 '';
73
74 passthru.updateScript = unstableGitUpdater { };
75
76 # Times out otherwise
77 requiredSystemFeatures = [ "big-parallel" ];
78
79 meta = with lib; {
80 description = "Performance portable programming model for machine learning tensor operators";
81 homepage = "https://github.com/ROCmSoftwarePlatform/composable_kernel";
82 license = with licenses; [ mit ];
83 maintainers = teams.rocm.members;
84 platforms = platforms.linux;
85 };
86})