pkgs/development/libraries/composable_kernel/default.nix

   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
  14 assert 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
  19 assert buildTests == false;
  20
  21 stdenv.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 }