##############################################################

# This Dockerfile contains AMD compilers

# Based on https://github.com/amd/InfinityHub-CI/tree/main/base-gpu-mpi-rocm-docker

##############################################################

ARG ROCM_VERSION=7.0.2

FROM rocm/dev-ubuntu-22.04:${ROCM_VERSION}-complete AS sdk-base

ARG UCX_BRANCH="v1.17.x"

ARG OMPI_BRANCH="v5.0.x"

# Update and Install basic Linux development tools

RUN rm /etc/apt/sources.list.d/* \

&& apt-get update \

&& DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \

dh-autoreconf python3-venv python3-dev python3-pip git \

ca-certificates ssh make libtinfo* initramfs-tools libelf-dev \

wget build-essential autoconf automake libtool \

pkg-config libnuma* gfortran flex hwloc cmake

ENV ROCM_HOME=/opt/rocm\

UCX_HOME=/opt/ucx \

OMPI_HOME=/opt/ompi

# Adding ROCM

ENV PATH=$ROCM_HOME/bin:$PATH \

LD_LIBRARY_PATH=$ROCM_HOME/lib:$ROCM_HOME/lib/llvm/lib:$LD_LIBRARY_PATH

# Until rocm base has it fixed

RUN ln -s /opt/rocm/llvm/bin/offload-arch /opt/rocm/bin/offload-arch || echo "offload-arch already exists"

# Install UCX

RUN cd /tmp/ \

&& git clone --depth 1 https://github.com/openucx/ucx.git -b ${UCX_BRANCH} \

&& cd ucx \

&& ./autogen.sh \

&& mkdir build \

&& cd build \

&& ../contrib/configure-release --prefix=$UCX_HOME \

--with-rocm=$ROCM_HOME \

--without-knem \

--without-xpmem \

--without-cuda \

--without-java \

--enable-mt \

--enable-optimizations \

--disable-logging \

--disable-debug \

--disable-examples \

--disable-assertions \

--disable-params-check \

&& make -j ${nproc} \

&& make install

# Install OpenMPI

RUN cd /tmp \

&& git clone --depth 1 --recursive https://github.com/open-mpi/ompi.git -b ${OMPI_BRANCH} \

&& cd ompi \

&& ./autogen.pl \

&& mkdir build \

&& cd build \

&& ../configure CC=amdclang CXX=amdclang++ FC=amdflang F90=amdflang \

--prefix=$OMPI_HOME \

--with-ucx=$UCX_HOME \

--with-rocm=$ROCM_HOME \

--enable-mca-no-build=btl-uct \

--without-verbs \

--enable-mpi1-compatibility \

--enable-mpi-fortran=no \

--disable-debug \

--disable-oshmem \

--disable-man-pages \

&& make -j ${nproc} \

&& make install

# UCX config

ENV UCX_WARN_UNUSED_ENV_VARS=n

ENV UCX_TLS=sm,self,rocm,rocm_copy,rocm_ipc

# Cleanup

RUN rm -rf /tmp/ucx && rm -rf /tmp/ompi

# Adding OpenMPI and UCX to Environment

ENV PATH=$OMPI_HOME/bin:$UCX_HOME/bin:$PATH \

LD_LIBRARY_PATH=$OMPI_HOME/lib:$UCX_HOME/lib:$LD_LIBRARY_PATH \

C_INCLUDE_PATH=$OMPI_HOME/include:$UCX_HOME/include:$C_INCLUDE_PATH \

CPLUS_INCLUDE_PATH=$OMPI_HOME/include:$UCX_HOME/include:$CPLUS_INCLUDE_PATH \

CPATH=$OMPI_HOME/include:$UCX_HOME/include:$CPATH \

INCLUDE=$OMPI_HOME/include:$UCX_HOME/include:$INCLUDE \

PKG_CONFIG_PATH=$OMPI_HOME/lib/pkgconfig:$UCX_HOME/lib/pkgconfig:$PKG_CONFIG_PATH

# Adding environment variable for Running as ROOT

ENV OMPI_ALLOW_RUN_AS_ROOT=1

ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1

RUN apt-get clean && apt-get autoclean && apt-get autoremove -y && \

rm -rf /var/lib/apt/lists/*

EXPOSE 8888

CMD ["/bin/bash"]

########################################################################

# AOMP for GPUs (OpenMP offloading)

########################################################################

# This will only trigger if arch is aomp since the final stage depends on it

FROM sdk-base AS amdclang

# MPI env

ENV OMPI_CC="amdclang"

ENV OMPI_CXX="amdclang++"

# Devito env

ENV DEVITO_ARCH="aomp"

ENV DEVITO_PLATFORM="amdgpuX"

ENV DEVITO_LANGUAGE="openmp"

########################################################################

# HIPCC for GPUs (HIP)

########################################################################

# This will only trigger if arch is hip since the final stage depends on it

FROM sdk-base AS hip

# MPI env

ENV OMPI_CC="hipcc"

ENV OMPI_CXX="hipcc"

# Devito env

ENV DEVITO_ARCH="hip"

ENV DEVITO_PLATFORM="amdgpuX"

ENV DEVITO_LANGUAGE="hip"