feat: base complete

base.dockerfile

@@ -1,16 +1,37 @@
-FROM ubuntu:22.04
+# > base stage
-ARG ROCM_VERSION="5.7.1"
+# 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 as single target may cause error
+ARG ROCBLAS_ARCH="gfx803;gfx900"
 ARG DEBIAN_FRONTEND="noninteractive"
 
 # Install essentials
-RUN <<EOF
+RUN apt-get update && \
-apt-get update
+apt-get install -y --no-install-recommends \
-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
+    wget \
-apt-get clean
+    git \
-rm -rf /var/lib/apt/lists/*
+    ca-certificates \
-EOF
+    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
@@ -23,24 +44,11 @@ EOF
 # Install ROCm
 RUN <<EOF
 apt-get update
-apt-get install -y --no-install-recommends rocm-dev
+apt-get install -y --no-install-recommends rocm-dev rocblas-dev
 apt-get clean
 rm -rf /var/lib/apt/lists/*
 EOF
 
-# 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/
-
-# # gfx803 as single target will cause error
-# RUN sed -i -e 's/rmake.py --install/rmake.py -a="gfx803;gfx900" --install/g' install.sh
-
-# RUN ./install.sh -d --cmake_install
-# WORKDIR /root
-
-## need to fix broken version
-
 # Create non-root user
 RUN <<EOF
 echo "%sudo ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/sudo-nopasswd
@@ -52,3 +60,32 @@ ENV PATH="${PATH}:/opt/rocm/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
+RUN ./install.sh -d --cmake_install
+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"
+
+# Default to login shell
+CMD [ "bash", "-l" ]

fix-rocblas.sh

@@ -0,0 +1,25 @@
+#!/bin/bash
+
+get_version() {
+    apt show $1 2>/dev/null | grep Version | cut -d " " -f2 | tr -d '\n'
+}
+
+get_filename() {
+    ls -l $1 | awk '{print $9}' | head -n1 | tr -d '\n'
+}
+
+PKGS=("rocblas" "rocblas-dev")
+for PKG in "${PKGS[@]}"; do
+    DIR="${PKG}-deb"
+    TARGET_VER=$(get_version $PKG)
+    FILENAME=$(get_filename "${PKG}_*.deb")
+    mkdir -p $DIR
+    rm -rf ${DIR}/*
+    dpkg-deb -R $FILENAME $DIR
+    sed -i "s/^Version: .*/Version: ${TARGET_VER}/" ${DIR}/DEBIAN/control
+    chmod 755 ${DIR}/DEBIAN/postinst
+    chmod 755 ${DIR}/DEBIAN/prerm
+    OUT_NAME="${PKG}_${TARGET_VER}_amd64.deb"
+    dpkg-deb -b $DIR $OUT_NAME
+    dpkg -i $OUT_NAME
+done