pkgs/development/rocm-modules/5/rocmlir/default.nix

   1 { lib
   2 , stdenv
   3 , fetchFromGitHub
   4 , rocmUpdateScript
   5 , cmake
   6 , rocm-cmake
   7 , ninja
   8 , clr
   9 , git
  10 , libxml2
  11 , libedit
  12 , zstd
  13 , zlib
  14 , ncurses
  15 , python3Packages
  16 , buildRockCompiler ? false
  17 , buildTests ? false # `argument of type 'NoneType' is not iterable`
  18 }:
  19
  20 # Theoretically, we could have our MLIR have an output
  21 # with the source and built objects so that we can just
  22 # use it as the external LLVM repo for this
  23 let
  24   suffix =
  25     if buildRockCompiler
  26     then "-rock"
  27     else "";
  28
  29   llvmNativeTarget =
  30     if stdenv.hostPlatform.isx86_64 then "X86"
  31     else if stdenv.hostPlatform.isAarch64 then "AArch64"
  32     else throw "Unsupported ROCm LLVM platform";
  33 in stdenv.mkDerivation (finalAttrs: {
  34   pname = "rocmlir${suffix}";
  35   version = "5.7.1";
  36
  37   outputs = [
  38     "out"
  39   ] ++ lib.optionals (!buildRockCompiler) [
  40     "external"
  41   ];
  42
  43   src = fetchFromGitHub {
  44     owner = "ROCm";
  45     repo = "rocMLIR";
  46     rev = "rocm-${finalAttrs.version}";
  47     hash = "sha256-vPi4UVljohVAfnwDVQqeOVaJPa6v8aV5uBOtqLddTtc=";
  48   };
  49
  50   nativeBuildInputs = [
  51     cmake
  52     rocm-cmake
  53     ninja
  54     clr
  55     python3Packages.python
  56     python3Packages.tomli
  57   ];
  58
  59   buildInputs = [
  60     git
  61     libxml2
  62     libedit
  63   ];
  64
  65   propagatedBuildInputs = [
  66     zstd
  67     zlib
  68     ncurses
  69   ];
  70
  71   cmakeFlags = [
  72     "-DLLVM_TARGETS_TO_BUILD=AMDGPU;${llvmNativeTarget}"
  73     "-DLLVM_ENABLE_ZSTD=ON"
  74     "-DLLVM_ENABLE_ZLIB=ON"
  75     "-DLLVM_ENABLE_TERMINFO=ON"
  76     "-DROCM_PATH=${clr}"
  77     # Manually define CMAKE_INSTALL_<DIR>
  78     # See: https://github.com/NixOS/nixpkgs/pull/197838
  79     "-DCMAKE_INSTALL_BINDIR=bin"
  80     "-DCMAKE_INSTALL_LIBDIR=lib"
  81     "-DCMAKE_INSTALL_INCLUDEDIR=include"
  82   ] ++ lib.optionals buildRockCompiler [
  83     "-DBUILD_FAT_LIBROCKCOMPILER=ON"
  84   ] ++ lib.optionals (!buildRockCompiler) [
  85     "-DROCM_TEST_CHIPSET=gfx000"
  86   ];
  87
  88   postPatch = ''
  89     patchShebangs mlir
  90
  91     substituteInPlace mlir/utils/performance/common/CMakeLists.txt \
  92       --replace "/opt/rocm" "${clr}"
  93   '';
  94
  95   dontBuild = true;
  96   doCheck = true;
  97
  98   # Certain libs aren't being generated, try enabling tests next update
  99   checkTarget = if buildRockCompiler
 100                 then "librockCompiler"
 101                 else if buildTests
 102                 then "check-rocmlir"
 103                 else "check-rocmlir-build-only";
 104
 105   postInstall = let
 106     libPath = lib.makeLibraryPath [ zstd zlib ncurses clr stdenv.cc.cc ];
 107   in lib.optionals (!buildRockCompiler) ''
 108     mkdir -p $external/lib
 109     cp -a external/llvm-project/llvm/lib/{*.a*,*.so*} $external/lib
 110     patchelf --set-rpath $external/lib:$out/lib:${libPath} $external/lib/*.so*
 111     patchelf --set-rpath $out/lib:$external/lib:${libPath} $out/{bin/*,lib/*.so*}
 112   '';
 113
 114   passthru.updateScript = rocmUpdateScript {
 115     name = finalAttrs.pname;
 116     owner = finalAttrs.src.owner;
 117     repo = finalAttrs.src.repo;
 118     page = "tags?per_page=2";
 119     filter = ".[1].name | split(\"-\") | .[1]";
 120   };
 121
 122   meta = with lib; {
 123     description = "MLIR-based convolution and GEMM kernel generator";
 124     homepage = "https://github.com/ROCm/rocMLIR";
 125     license = with licenses; [ asl20 ];
 126     maintainers = teams.rocm.members;
 127     platforms = platforms.linux;
 128     broken = versions.minor finalAttrs.version != versions.minor stdenv.cc.version || versionAtLeast finalAttrs.version "6.0.0";
 129   };
 130 })