Build and install ROCm via TheRock in AMD images (#7)

build_tools/install_amdgpu.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Install the AMD GPU DRM libraries and Vulkan.
+# In order to use this in the container, you must pass /dev/dri and /dev/kfd in as:
+#   docker run --device=/dev/kfd --device=/dev/dri <docker ID>
+
+set -euo pipefail
+
+AMDGPU_VERSION=$1
+
+ARCH="$(uname -m)"
+if [[ "${ARCH}" == "x86_64" ]]; then
+  apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates curl libnuma-dev gnupg \
+    && curl -sL https://repo.radeon.com/rocm/rocm.gpg.key | apt-key add - \
+    && printf "deb [arch=amd64] https://repo.radeon.com/amdgpu/${AMDGPU_VERSION}/ubuntu focal main" | tee /etc/apt/sources.list.d/amdgpu.list \
+    && apt-get update \
+    && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
+    libdrm-amdgpu-dev \
+    libdrm-dev
+else
+  echo "Installing ROCM for ${ARCH} is not supported yet."
+fi
+
+wget https://github.com/GPUOpen-Drivers/AMDVLK/releases/download/v-2023.Q3.1/amdvlk_2023.Q3.1_amd64.deb && \
+    dpkg -i amdvlk_2023.Q3.1_amd64.deb

build_tools/install_the_rock.sh

@@ -0,0 +1,71 @@
+#!/bin/bash
+# Copyright 2024 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Build and install the parts of ROCm that are required by IREE.
+
+set -euo pipefail
+
+THE_ROCK_GIT_REFSPEC=${THE_ROCK_GIT_REFSPEC:-rocm-6.1.0}
+THE_ROCK_ROCM_MANIFEST_URL=${THE_ROCK_ROCM_MANIFEST_URL:-https://github.com/ROCm/ROCm.git}
+THE_ROCK_ROCM_MANIFEST_BRANCH=${THE_ROCK_ROCM_MANIFEST_BRANCH:-refs/tags/rocm-6.1.0}
+
+ARCH="$(uname -m)"
+
+if [[ "${ARCH}" != "x86_64" ]]; then
+  echo "Installing TheRock (ROCm) for ${ARCH} is not supported yet."
+  exit 0
+fi
+
+if [ ! -d TheRock ]; then
+  git clone https://github.com/nod-ai/TheRock.git \
+    --branch $THE_ROCK_GIT_REFSPEC \
+    --depth 1
+  cd TheRock
+else
+  cd TheRock
+  git fetch --depth 1 origin $THE_ROCK_GIT_REFSPEC
+  git checkout FETCH_HEAD
+fi
+
+apt update
+apt install -y \
+  repo git-lfs libnuma-dev ninja-build g++ pkg-config libdrm-dev \
+  libelf-dev xxd libgl1-mesa-dev
+python -m pip install CppHeaderParser
+
+# Make sure git does not report Committer identity unknown errors.
+# export GIT_COMMITTER_NAME="Noname"
+# export GIT_AUTHOR_NAME="$GIT_COMMITTER_NAME"
+# export GIT_COMMITTER_EMAIL="[email protected]"
+# export GIT_AUTHOR_EMAIL="$GIT_COMMITTER_EMAIL"
+# export GIT_TERMINAL_PROMPT=0
+# The buggy repo tool does not respect the above commented vars,
+# so it asks us about our idenity.
+git config --global user.email "Noname"
+git config --global user.name "[email protected]"
+
+git config --global color.ui true
+export GIT_TERMINAL_PROMPT=0
+
+python ./build_tools/fetch_sources.py \
+  --manifest-url $THE_ROCK_ROCM_MANIFEST_URL \
+  --manifest-branch $THE_ROCK_ROCM_MANIFEST_BRANCH
+
+# The build fails with clang-19.
+export CC=gcc
+export CXX=g++
+
+cmake -B build \
+  -GNinja \
+  -DCMAKE_INSTALL_PREFIX=/usr/local \
+  .
+
+cmake --build build
+
+cmake --install build --component amdgpu-runtime
+cmake --install build --component amdgpu-runtime-dev
+cmake --install build --component amdgpu-compiler

dockerfiles/amdgpu_ubuntu_jammy_ghr_x86_64.Dockerfile

@@ -6,35 +6,33 @@ FROM docker.io/myoung34/github-runner:2.314.1-ubuntu-jammy
 # Basic packages.
 RUN apt update && \
     apt install -y \
-        wget python3.11 git unzip curl gnupg2 lsb-release vulkan-tools && \
+        wget python3.11 python3-pip git unzip curl gnupg2 lsb-release vulkan-tools && \
     update-alternatives --install /usr/local/bin/python python /usr/bin/python3.11 3 && \
     update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.11 3
 
-# Install the AMDVLK driver.
-WORKDIR /install-amdvlk
-RUN wget https://github.com/GPUOpen-Drivers/AMDVLK/releases/download/v-2023.Q3.1/amdvlk_2023.Q3.1_amd64.deb && \
-    dpkg -i amdvlk_2023.Q3.1_amd64.deb && \
-    rm -rf /install-amdvlk/*
+# CMake
+WORKDIR /install-cmake
+# Install the latest CMake version we support
+ARG CMAKE_VERSION="3.29.3"
+COPY build_tools/install_cmake.sh ./
+RUN ./install_cmake.sh "${CMAKE_VERSION}" && rm -rf /install-cmake
+
+# AMD GPU DRM & Vulkan
+WORKDIR /install-amdgpu
+ARG AMDGPU_VERSION=6.1
+COPY build_tools/install_amdgpu.sh ./
+RUN ./install_amdgpu.sh "${AMDGPU_VERSION}" && rm -rf /install-amdgpu
 WORKDIR /
-RUN rmdir /install-amdvlk
 
-# Install the ROCM driver.
-# In order to use this in the container, you must pass /dev/dri and /dev/kfd in as:
-#   docker run --device=/dev/kfd --device=/dev/dri <docker ID> rocminfo
-ARG ROCM_VERSION=5.6
-WORKDIR /install-rocm
-RUN mkdir --parents --mode=0755 /etc/apt/keyrings && wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
-    gpg --dearmor | tee /etc/apt/keyrings/rocm.gpg > /dev/null && \
-    echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu jammy main" | \
-    tee /etc/apt/sources.list.d/amdgpu.list && \
-    echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VERSION} jammy main" | \
-    tee /etc/apt/sources.list.d/rocm.list && \
-    echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | tee /etc/apt/preferences.d/rocm-pin-600 && \
-    apt update && DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends ca-certificates rocm-hip-libraries && \
-    echo "/opt/rocm/lib\n/opt/rocm/hip/lib" | tee /etc/ld.so.conf.d/24-rocm.conf && ldconfig && \
-    rm -rf /install-rocm/*
+# TheRock (ROCm)
+WORKDIR /install-the-rock
+COPY build_tools/install_the_rock.sh ./
+RUN ./install_the_rock.sh \
+  && rm -rf /install-the-rock
 WORKDIR /
-RUN rmdir /install-rocm/
+
+# Clean up.
+RUN apt clean && rm -rf /var/lib/apt/lists/*
 
 # Switch back to the working directory upstream expects.
 WORKDIR /actions-runner

dockerfiles/amdgpu_ubuntu_jammy_x86_64.Dockerfile

@@ -6,37 +6,30 @@ FROM ubuntu:jammy
 # Basic packages.
 RUN apt update && \
     apt install -y \
-        wget python3.11 git unzip curl gnupg2 lsb-release vulkan-tools && \
+        wget python3.11 python3-pip git unzip curl gnupg2 lsb-release vulkan-tools && \
     update-alternatives --install /usr/local/bin/python python /usr/bin/python3.11 3 && \
     update-alternatives --install /usr/local/bin/python3 python3 /usr/bin/python3.11 3
 
-# Install the AMDVLK driver.
-# In order to use this in the container, you must pass /dev/dri in as:
-#   --device /dev/dri
-WORKDIR /install-amdvlk
-RUN wget https://github.com/GPUOpen-Drivers/AMDVLK/releases/download/v-2023.Q3.1/amdvlk_2023.Q3.1_amd64.deb && \
-    dpkg -i amdvlk_2023.Q3.1_amd64.deb && \
-    rm -rf /install-amdvlk/*
+# CMake
+WORKDIR /install-cmake
+# Install the latest CMake version we support
+ARG CMAKE_VERSION="3.29.3"
+COPY build_tools/install_cmake.sh ./
+RUN ./install_cmake.sh "${CMAKE_VERSION}" && rm -rf /install-cmake
+
+# AMD GPU DRM & Vulkan
+WORKDIR /install-amdgpu
+ARG AMDGPU_VERSION=6.1
+COPY build_tools/install_amdgpu.sh ./
+RUN ./install_amdgpu.sh "${AMDGPU_VERSION}" && rm -rf /install-amdgpu
 WORKDIR /
-RUN rmdir /install-amdvlk
 
-# Install the ROCM driver.
-# In order to use this in the container, you must pass /dev/dri and /dev/kfd in as:
-#   docker run --device=/dev/kfd --device=/dev/dri <docker ID> rocminfo
-ARG ROCM_VERSION=5.6
-WORKDIR /install-rocm
-RUN mkdir --parents --mode=0755 /etc/apt/keyrings && wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
-    gpg --dearmor | tee /etc/apt/keyrings/rocm.gpg > /dev/null && \
-    echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/amdgpu/${ROCM_VERSION}/ubuntu jammy main" | \
-    tee /etc/apt/sources.list.d/amdgpu.list && \
-    echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VERSION} jammy main" | \
-    tee /etc/apt/sources.list.d/rocm.list && \
-    echo 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | tee /etc/apt/preferences.d/rocm-pin-600 && \
-    apt update && DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends ca-certificates rocm-hip-libraries && \
-    echo "/opt/rocm/lib\n/opt/rocm/hip/lib" | tee /etc/ld.so.conf.d/24-rocm.conf && ldconfig && \
-    rm -rf /install-rocm/*
+# TheRock (ROCm)
+WORKDIR /install-the-rock
+COPY build_tools/install_the_rock.sh ./
+RUN ./install_the_rock.sh \
+#  && rm -rf /install-the-rock
 WORKDIR /
-RUN rmdir /install-rocm/
 
 # Clean up.
 RUN apt clean && rm -rf /var/lib/apt/lists/*