Files
rocm/base.dockerfile
T
2024-09-13 10:28:18 -07:00

104 lines
2.6 KiB
Docker

# > base stage
# Uses official ROCm packages
FROM ubuntu:22.04 AS base
ARG ROCM_VERSION="5.7.3"
ARG ROCM_APT_CODE_NAME="jammy"
ARG ROCBLAS_BRANCH="release/rocm-rel-5.7"
# "gfx803" or "gfx803;gfx900"
ARG ROCBLAS_ARCH="gfx803"
ARG DEBIAN_FRONTEND="noninteractive"
# Install essentials
RUN apt-get update && \
apt-get install -y --no-install-recommends \
wget \
git \
ca-certificates \
curl \
gnupg \
sudo \
libelf1 \
libnuma-dev \
build-essential \
vim-nox \
cmake-curses-gui \
kmod \
file \
make \
python3 \
python3-yaml \
python3-venv \
python3-joblib \
python3-pip \
libmsgpack-dev \
&& apt-get clean && rm -rf /var/lib/apt/lists/*
# Add ROCm apt repo
RUN <<EOF
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
bash -c 'echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VERSION} ${ROCM_APT_CODE_NAME} main" > "/etc/apt/sources.list.d/rocm.list"'
bash -c 'echo -e "Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600" | tee /etc/apt/preferences.d/rocm-pin-600'
EOF
# Install ROCm
RUN <<EOF
apt-get update
apt-get install -y --no-install-recommends rocm-dev rocm-libs
apt-get clean
rm -rf /var/lib/apt/lists/*
EOF
# For ROCm 4.5+
RUN <<EOF
sudo apt-get update
sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk || true
sudo apt-get clean
sudo rm -rf /var/lib/apt/lists/*
EOF
# Create non-root user
RUN <<EOF
echo "%sudo ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/sudo-nopasswd
useradd --create-home -G sudo,video --shell /bin/bash rocm
EOF
USER rocm
WORKDIR /home/rocm
ENV PATH="${PATH}:/opt/rocm/bin:/home/rocm/.local/bin"
# Default to login shell
CMD [ "bash", "-l" ]
# > patched stage
# Uses patched rocblas and rocblas-dev
# (Others are official ROCm packages)
FROM base AS patched
# Compile and install patched rocBLAS
USER root
WORKDIR /root
RUN git clone -b ${ROCBLAS_BRANCH} https://github.com/ROCmSoftwarePlatform/rocBLAS.git
WORKDIR /root/rocBLAS
RUN rm -rf library/src/blas3/Tensile/Logic/asm_full/r9nano/
RUN sed -i -e 's/rmake.py --install/rmake.py -a="${ROCBLAS_ARCH}" --install/g' install.sh
# > https://github.com/ROCm/Tensile/pull/1897
RUN ./install.sh -a "${ROCBLAS_ARCH}" -d --cmake_install --merge-architectures --no-lazy-library-loading
WORKDIR /root/rocBLAS/build/release
COPY fix-rocblas.sh .
RUN <<EOF
make package
chmod +x fix-rocblas.sh
./fix-rocblas.sh
EOF
# Change to non-root user
USER rocm
WORKDIR /home/rocm
ENV PATH="${PATH}:/opt/rocm/bin:/home/rocm/.local/bin"
# Default to login shell
CMD [ "bash", "-l" ]