NixOS ROCm update: 6.0.2->6.2.2

Jachin Shen
Jachin Shen

Categories

Tags

背景

尝试在 NixOS 系统中编译 ONNX Runtime 时，我发现 NixOS 自带的 ROCm 版本较低，仅为 6.0.2，这导致编译过程中会出现奇奇怪怪的问题。因此，需要将 ROCm 升级至 6.2.2 版本，顺便可以进一步了解 ROCm 的细节。由于 NixOS 是一个独特的 Linux 发行版，采用函数式包定义方式，而 ROCm 本身又相当复杂，在这样一个小众发行版上打包一个小众软件，自然会遇到诸多挑战。本文将记录升级 ROCm 的全过程，代码存放在 rocm-6.2.2 中。

NixOS 中 ROCm 的组织结构

NixOS 通过 GitHub 管理所有软件包。在官方仓库中，我们可以找到 ROCm 的函数式定义。ROCm 的编译顺序如下：

  1. 定制化的 LLVM
  2. rocm-core/cmake/thunk/smi/device-libs/runtime/comgr/info 等基础组件
  3. hip-common/hipcc/hipify 等基础 HIP 工具链
  4. CLR（Common Language Runtime）运行库
  5. roc/hip 系列库
  6. miopen/migraphx/mivisionx 等深度学习库

为了顺利完成 ROCm 的升级，主要参考了以下资源：

  1. NixOS 的 ROCm 维护团队提供的[tracking] ROCm packages，其中包含了源代码升级脚本。
  2. Archlinux 的新版 6.2.2 ROCm 库，例如 rocm-llvm中的 PKGBUILD 文件。

升级 ROCm

升级源代码

首先，我尝试使用了 ROCm 维护团队提供的源代码升级脚本，的确可以自动更新源代码以及校验值。然而，我发现 device-libs/comgr/hipcc 等库已经整合进了 LLVM，因此，我需要首先解决 LLVM 的编译问题。

升级 LLVM

旧版本 LLVM 的编译

旧版本的 Nix 定义位于 llvm/default.nix:

{ # stdenv FIXME: Try changing back to this with a new ROCm release https://github.com/NixOS/nixpkgs/issues/271943
  gcc12Stdenv
, callPackage
, rocmUpdateScript
, wrapBintoolsWith
, overrideCC
, rocm-device-libs
, rocm-runtime
, rocm-thunk
, clr
}:

let
  ## Stage 1 ##
  # Projects
  llvm = callPackage ./stage-1/llvm.nix { inherit rocmUpdateScript; stdenv = gcc12Stdenv; };
  clang-unwrapped = callPackage ./stage-1/clang-unwrapped.nix { inherit rocmUpdateScript llvm; stdenv = gcc12Stdenv; };
  lld = callPackage ./stage-1/lld.nix { inherit rocmUpdateScript llvm; stdenv = gcc12Stdenv; };

  # Runtimes
  runtimes = callPackage ./stage-1/runtimes.nix { inherit rocmUpdateScript llvm; stdenv = gcc12Stdenv; };

  ## Stage 2 ##
  # Helpers
  bintools-unwrapped = callPackage ./stage-2/bintools-unwrapped.nix { inherit llvm lld; };
  bintools = wrapBintoolsWith { bintools = bintools-unwrapped; };
  rStdenv = callPackage ./stage-2/rstdenv.nix { inherit llvm clang-unwrapped lld runtimes bintools; stdenv = gcc12Stdenv; };
in rec {
  inherit
  llvm
  clang-unwrapped
  lld
  bintools;

  # Runtimes
  libc = callPackage ./stage-2/libc.nix { inherit rocmUpdateScript; stdenv = rStdenv; };
  libunwind = callPackage ./stage-2/libunwind.nix { inherit rocmUpdateScript; stdenv = rStdenv; };
  libcxxabi = callPackage ./stage-2/libcxxabi.nix { inherit rocmUpdateScript; stdenv = rStdenv; };
  libcxx = callPackage ./stage-2/libcxx.nix { inherit rocmUpdateScript; stdenv = rStdenv; };
  compiler-rt = callPackage ./stage-2/compiler-rt.nix { inherit rocmUpdateScript llvm; stdenv = rStdenv; };

  ## Stage 3 ##
  # Helpers
  clang = callPackage ./stage-3/clang.nix { inherit llvm lld clang-unwrapped bintools libc libunwind libcxxabi libcxx compiler-rt; stdenv = gcc12Stdenv; };
  rocmClangStdenv = overrideCC gcc12Stdenv clang;

  # Projects
  clang-tools-extra = callPackage ./stage-3/clang-tools-extra.nix { inherit rocmUpdateScript llvm clang-unwrapped; stdenv = rocmClangStdenv; };
  libclc = callPackage ./stage-3/libclc.nix { inherit rocmUpdateScript llvm clang; stdenv = rocmClangStdenv; };
  lldb = callPackage ./stage-3/lldb.nix { inherit rocmUpdateScript clang; stdenv = rocmClangStdenv; };
  mlir = callPackage ./stage-3/mlir.nix { inherit rocmUpdateScript clr; stdenv = rocmClangStdenv; };
  polly = callPackage ./stage-3/polly.nix { inherit rocmUpdateScript; stdenv = rocmClangStdenv; };
  flang = callPackage ./stage-3/flang.nix { inherit rocmUpdateScript clang-unwrapped mlir; stdenv = rocmClangStdenv; };
  openmp = callPackage ./stage-3/openmp.nix { inherit rocmUpdateScript llvm clang-unwrapped clang rocm-device-libs rocm-runtime rocm-thunk; stdenv = rocmClangStdenv; };

  # Runtimes
  pstl = callPackage ./stage-3/pstl.nix { inherit rocmUpdateScript; stdenv = rocmClangStdenv; };
}

这里分成了三个阶段：

  1. 阶段一编译核心的 LLVM、 CLANG 和 LLD 二进制，分别是 LLVM 项目的后端，C前端和链接器。不过此时只有孤零零的二进制，参考 C-NixOS Wiki，需要 wrapper 设置好环境变量才能正常使用
  2. 阶段二编译核心的库，比如 libc、libunwind、libcxxabi、libcxx 等
  3. 阶段三构建类似 NixOS 中 stdenv 的编译环境，用来编译剩下的库

在每个阶段内部，都会调用 base.nix，比如 llvm.nix：

{ stdenv
, callPackage
, rocmUpdateScript
}:

callPackage ../base.nix {
  inherit stdenv rocmUpdateScript;
  requiredSystemFeatures = [ "big-parallel" ];
  isBroken = stdenv.isAarch64; # https://github.com/ROCm/ROCm/issues/1831#issuecomment-1278205344
}

而 base.nix 的内容如下：

{ lib
, stdenv
, gcc12Stdenv
, fetchFromGitHub
, rocmUpdateScript
, pkg-config
, cmake
, ninja
, git
, doxygen
, sphinx
, lit
, libxml2
, libxcrypt
, libedit
, libffi
, mpfr
, zlib
, ncurses
, python3Packages
, buildDocs ? true
, buildMan ? true
, buildTests ? true
, targetName ? "llvm"
, targetDir ? "llvm"
, targetProjects ? [ ]
, targetRuntimes ? [ ]
, llvmTargetsToBuild ? [ "NATIVE" ] # "NATIVE" resolves into x86 or aarch64 depending on stdenv
, extraPatches ? [ ]
, extraNativeBuildInputs ? [ ]
, extraBuildInputs ? [ ]
, extraCMakeFlags ? [ ]
, extraPostPatch ? ""
, checkTargets ? [(
  lib.optionalString buildTests (
    if targetDir == "runtimes"
    then "check-runtimes"
    else "check-all"
  )
)]
, extraPostInstall ? ""
, hardeningDisable ? [ ]
, requiredSystemFeatures ? [ ]
, extraLicenses ? [ ]
, isBroken ? false
}:

let stdenv' = stdenv; in
let stdenv =
      if stdenv'.cc.cc.isGNU or false && lib.versionAtLeast stdenv'.cc.cc.version "13.0"
      then gcc12Stdenv
      else stdenv';
in

let
  llvmNativeTarget =
    if stdenv.isx86_64 then "X86"
    else if stdenv.isAarch64 then "AArch64"
    else throw "Unsupported ROCm LLVM platform";
  inferNativeTarget = t: if t == "NATIVE" then llvmNativeTarget else t;
  llvmTargetsToBuild' = [ "AMDGPU" ] ++ builtins.map inferNativeTarget llvmTargetsToBuild;
in stdenv.mkDerivation (finalAttrs: {
  pname = "rocm-llvm-${targetName}";
  version = "6.0.2";

  outputs = [
    "out"
  ] ++ lib.optionals buildDocs [
    "doc"
  ] ++ lib.optionals buildMan [
    "man"
    "info" # Avoid `attribute 'info' missing` when using with wrapCC
  ];

  patches = extraPatches;

  src = fetchFromGitHub {
    owner = "ROCm";
    repo = "llvm-project";
    rev = "rocm-${finalAttrs.version}";
    hash = "sha256-uGxalrwMNCOSqSFVrYUBi3ijkMEFFTrzFImmvZKQf6I=";
  };

  nativeBuildInputs = [
    pkg-config
    cmake
    ninja
    git
    python3Packages.python
  ] ++ lib.optionals (buildDocs || buildMan) [
    doxygen
    sphinx
    python3Packages.recommonmark
  ] ++ lib.optionals (buildTests && !finalAttrs.passthru.isLLVM) [
    lit
  ] ++ extraNativeBuildInputs;

  buildInputs = [
    libxml2
    libxcrypt
    libedit
    libffi
    mpfr
  ] ++ extraBuildInputs;

  propagatedBuildInputs = lib.optionals finalAttrs.passthru.isLLVM [
    zlib
    ncurses
  ];

  sourceRoot = "${finalAttrs.src.name}/${targetDir}";

  cmakeFlags = [
    "-DLLVM_TARGETS_TO_BUILD=${builtins.concatStringsSep ";" llvmTargetsToBuild'}"
  ] ++ lib.optionals (finalAttrs.passthru.isLLVM && targetProjects != [ ]) [
    "-DLLVM_ENABLE_PROJECTS=${lib.concatStringsSep ";" targetProjects}"
  ] ++ lib.optionals ((finalAttrs.passthru.isLLVM || targetDir == "runtimes") && targetRuntimes != [ ]) [
    "-DLLVM_ENABLE_RUNTIMES=${lib.concatStringsSep ";" targetRuntimes}"
  ] ++ lib.optionals finalAttrs.passthru.isLLVM [
    "-DLLVM_INSTALL_UTILS=ON"
    "-DLLVM_INSTALL_GTEST=ON"
  ] ++ lib.optionals (buildDocs || buildMan) [
    "-DLLVM_INCLUDE_DOCS=ON"
    "-DLLVM_BUILD_DOCS=ON"
    # "-DLLVM_ENABLE_DOXYGEN=ON" Way too slow, only uses one core
    "-DLLVM_ENABLE_SPHINX=ON"
    "-DSPHINX_OUTPUT_HTML=ON"
    "-DSPHINX_OUTPUT_MAN=ON"
    "-DSPHINX_WARNINGS_AS_ERRORS=OFF"
  ] ++ lib.optionals buildTests [
    "-DLLVM_INCLUDE_TESTS=ON"
    "-DLLVM_BUILD_TESTS=ON"
    "-DLLVM_EXTERNAL_LIT=${lit}/bin/.lit-wrapped"
  ] ++ extraCMakeFlags;

  prePatch = ''
    cd ../
    chmod -R u+w .
  '';

  postPatch = ''
    cd ${targetDir}
  '' + lib.optionalString finalAttrs.passthru.isLLVM ''
    patchShebangs lib/OffloadArch/make_generated_offload_arch_h.sh
  '' + lib.optionalString (buildTests && finalAttrs.passthru.isLLVM) ''
    # FileSystem permissions tests fail with various special bits
    rm test/tools/llvm-objcopy/ELF/mirror-permissions-unix.test
    rm unittests/Support/Path.cpp

    substituteInPlace unittests/Support/CMakeLists.txt \
      --replace-fail "Path.cpp" ""
  '' + extraPostPatch;

  doCheck = buildTests;
  checkTarget = lib.concatStringsSep " " checkTargets;

  postInstall = lib.optionalString buildMan ''
    mkdir -p $info
  '' + extraPostInstall;

  passthru = {
    isLLVM = targetDir == "llvm";
    isClang = targetDir == "clang" || builtins.elem "clang" targetProjects;

    updateScript = rocmUpdateScript {
      name = finalAttrs.pname;
      owner = finalAttrs.src.owner;
      repo = finalAttrs.src.repo;
    };
  };

  inherit hardeningDisable requiredSystemFeatures;

  meta = with lib; {
    description = "ROCm fork of the LLVM compiler infrastructure";
    homepage = "https://github.com/ROCm/llvm-project";
    license = with licenses; [ ncsa ] ++ extraLicenses;
    maintainers = with maintainers; [ acowley lovesegfault ] ++ teams.rocm.members;
    platforms = platforms.linux;
    broken = isBroken || versionAtLeast finalAttrs.version "7.0.0";
  };
})

这里是最难以理解的地方，因为普通的 Nix 包一般都是一次编译一整个项目，但这里既区分了三个阶段，又复用了 base.nix。参考 Archlinux 中的 llvm 包，我们可以知道，llvm-project 是一个包含多个子项目的集合，所以需要尽量模块化，方便管理。另外由于 NixOS 的特殊性，需要给不同阶段配置不同的环境，所以需要区分不同的阶段。

在 base.nix 中，这段代码：

let stdenv' = stdenv; in
let stdenv =
      if stdenv'.cc.cc.isGNU or false && lib.versionAtLeast stdenv'.cc.cc.version "13.0"
      then gcc12Stdenv
      else stdenv';
in

就是用来区分不同环境的。如果输入的环境是 GNU 编译器，意味着在阶段一，那么就统一使用 gcc12Stdenv，否则意味着在阶段二，使用了阶段一产生的 llvm/clang 环境，那么就保持输入的环境。

需要编译的项目就通过targetProjects和targetRuntimes来指定，另外通过extraCMakeFlags等参数配置每个项目特殊的编译选项，这样就可以保证 base.nix 的通用性，修改新项目的参数不会影响已经编译好的项目，避免不必要的重新编译。

阶段一

  1. 只更新源码，不更改编译环境，编译过程中会提示需要 GCC 13 及以上版本。
  2. 将gcc12Stdenv改为gcc13Stdenv/gcc14Stdenv，编译过程中出现 gcc: error: unrecognized command-line option ‘-nostdlib++’ 错误，具体原因不清，应该是 CMake 中检测到环境支持 -nostdlib++ 选项会启用，虽然 g++ 支持但是 gcc 不支持。
  3. 将编译环境更改为llvmPackages.stdenv，改用 llvm 来编译这个定制化的 llvm-project，可以避免 GNU 奇奇怪怪的魔法，编译成功。

这样阶段一的 llvm、clang、lld 都可以正常编译。

阶段二

阶段二原本是使用阶段一编译好的 llvm、clang、lld 组建新环境 rStdenv 来编译libc等库，但是会出现找不到头文件的错误。考虑到后面并没有用到 rStdenv，所以直接使用 llvmPackages.stdenv 来编译 libc，编译成功。这里不太明白为什么要使用阶段一编译好的 llvm、clang、lld，可能原作者有自己的考虑，不过我们这么做也没有遇到问题，使用现成的 llvm 环境编译 llvm-project 从道理上也没有问题。

阶段三

阶段三首先会使用阶段一二编译好的二进制组建新环境 rocmClangStdenv，这就是 ROCm 定制的 llvm 环境，也是之后都会用到的环境。这个环境的构建没有问题，但是使用这个环境构建剩余项目时，首先出现了环境不对的问题。还记得 base.nix 使用 isGNU 区分现成环境和我们自己编译的环境吗？由于我们现成环境也是使用的 llvm，所以这里的判断失效了，导致后续项目使用的是 gcc12Stdenv，而不是我们期望的 rocmClangStdenv。既然如此，我直接去除了这里的判断，直接使用输入的 stdenv 环境，具体使用哪个环境由调用者自己决定。

现在我们成功的进入了 rocmClangStdenv 环境，可以开始编译后续项目了。但是尝试了几个项目，都出现了找不到标准库头文件的错误。这就很奇怪了，毕竟在阶段二中，我们已经编译了标准库，为什么阶段三找不到呢？

于是我回头编译了旧版本的 rocmClangStdenv 进行对比，最后问题在于路径！旧版 Clang 的安装路径中使用了完整的版本号，所以 nix 中使用

clang_version=`${clang-unwrapped}/bin/clang -v 2>&1 | grep "clang version " | grep -E -o "[0-9.-]+"`

获取版本号，继而组装出路径 mkdir -p $out/{bin,include/c++/v1,lib/{cmake,clang/$clang_version/{include,lib}},libexec,share} 而新版 Clang 的安装路径中只使用了主版本号，所以后续指向了错误的路径，导致找不到头文件。知道了问题，改成使用主版本号即可。

clang_version=`${clang-unwrapped}/bin/clang -v 2>&1 | grep "clang version " | grep -E -o "[0-9.-]+" | grep -E -o "^[0-9]+"`

这样就可以找到头文件了，成功地编译了后续项目。

小结

俗话说，万事开头难。笔者在这第一步卡了很长时间，中途还搁置了很长时间。为了搞清楚编译流程，先是详细了解了 llvm-project 的整体架构，然后又在 Archlinux 中尝试编译，又回头补课了 NixOS 中的 C 编译环境。其实主要难点在于 llvm-project 的 nix 打包比较复杂，对 base.nix 进行了多次调用，导致细节都被隐藏起来了。测试编译环境的时候也一直在 GCC 和 LLVM 之间犹豫，一方面旧版本使用的是 GCC，总觉得问题会更少一点，另一方面阶段一改用 LLVM 就能成功编译了，但是后续编译又有问题。后来才知道，后续编译的问题不在于 LLVM，而是路径问题。整个过程中，有一个成功的案例作为参考还是很重要的，否则很容易陷入迷茫，定位不到问题。

升级 rocMLIR

在 [Tracking] ROCm packages 中提到了两个 patch，打上就可以解决找不到头文件的问题，正常编译。

升级 composable_kernel

升级完源代码后，在 CMake 阶段就会出现问题：

-- The HIP compiler identification is unknown
-- Detecting HIP compiler ABI info
-- Detecting HIP compiler ABI info - failed
-- Check for working HIP compiler: /nix/store/h0aa4jal4731bhrc5k865y927qndsjwa-rocm-llvm-clang-wrapper-6.2.2/bin/clang++
-- Check for working HIP compiler: /nix/store/h0aa4jal4731bhrc5k865y927qndsjwa-rocm-llvm-clang-wrapper-6.2.2/bin/clang++ - broken
...
    clang++: error: cannot find ROCm device library; provide its path via '--rocm-path' or '--rocm-device-lib-path', or pass '-nogpulib' to build without ROCm device library

检查 CMakeLists.txt，发现问题出在 project(composable_kernel VERSION ${version} LANGUAGES CXX HIP) 这一行，在 commit 7b027d5 中进行了重构，使用了 CMake 新的 HIP language support 。在旧版本 6.0.2 中并没有这个功能，所以那时的编译环境没有问题。

找到了问题，但是解决没有那么简单。翻看 CMake 的源码，和 HIP 相关的部分主要在：

  1. https://gitlab.kitware.com/cmake/cmake/-/blob/v3.30.0/Modules/CMakeDetermineHIPCompiler.cmake?ref_type=tags
  2. https://gitlab.kitware.com/cmake/cmake/-/blob/v3.30.0/Modules/CMakeTestHIPCompiler.cmake?ref_type=tags

整体流程是，CMake 会从环境中寻找 HIP，然后检查 ABI，如果失败就尝试编译一个简单的 HIP 程序。这里 CMake 找到了 ROCm llvm 的 clang++，但是这里的 clang++ 并没有配套的环境，所以编译失败。我们希望 CMake 找到的应该是已经包装好的 hipcc 。

那我们直接给 CMake 指定 HIP 的编译器行不行？在参数中增加 -DCMAKE_HIP_COMPILER=hipcc，结果如下：

CMake Error at /nix/store/irwcgpm8csj7br49jih6kdpqkaqya3vf-cmake-3.30.4/share/cmake-3.30/Modules/CMakeDetermineHIPCompiler.cmake:73 (message):
  CMAKE_HIP_COMPILER is set to the hipcc wrapper:

   hipcc

  This is not supported.  Use Clang directly, or let CMake pick a default.
Call Stack (most recent call first):
  CMakeLists.txt:26 (project)

很遗憾不行， CMake 不支持使用包装好的 hipcc ，必须使用 clang++ 。

那么现在就有矛盾的地方了， CMake 只接受光秃秃的 clang++，这在一般的发行版中不是问题，因为 ROCm 库已经在系统中了， clang++ 可以执行。然而，在 NixOS 中，安装的库并不会直接暴露在系统环境中，二进制可执行文件需要通过环境变量到对应的 nix store 中寻找库。这就导致了 clang++ 找不到库，从而编译失败。那么这里的一个 workaround 就是在这个编译环境中，将 ROCm 库的路径添加到环境变量中，让 clang++ 找到库。

  rocmtoolkit-merged = symlinkJoin {
    name = "rocmtoolkit-merged";

    paths = [
      rocm-core
      rocm-thunk
      rocm-device-libs
      roctracer
      rocdbgapi
      rocm-smi
      hsa-amd-aqlprofile-bin
      clr
    ];

    postBuild = ''
      rm -rf $out/nix-support
    '';
  };

  propagatedBuildInputs = [ rocmtoolkit-merged ];

  ROCM_PATH = "${rocmtoolkit-merged}";
  DEVICE_LIB_PATH = "${rocmtoolkit-merged}/amdgpu/bitcode";

这里使用了 NixOS 的 symlinkJoin 来将多个库合并成一个，然后通过环境变量来传递给编译环境。笔者尝试过只将基础库 ${clr} 添加到环境变量 ROCM_PATH 中，但发现 CMake 还是无法识别 ABI 信息，所以最后还是将所有相关库都添加到环境变量中。

LinkedIn Reddit