feat: add AMD GPU (amdflang/OpenMP offload) container support

sbryngelson · sbryngelson · commit 1557fc43cb78 · 2026-05-10T22:15:51.000-05:00
- Dockerfile: add TARGET=amd branch — downloads AFAR drop from repo.radeon.com,
  installs cmake 3.28 (3.22 doesn't recognise LLVMFlang), builds MPICH 3.4.3
  with amdflang so mpi.mod is compiler-compatible; runtime libs libnuma1/libdrm2
  added so only --rocm is needed at apptainer runtime
- docker.yml: add amd matrix entry + build/push/manifest steps; fix cpu to run
  natively on amd64/arm64 instead of QEMU cross-build; add weekly nightly cron
- CMakeLists.txt: make Cray-specific MPI/hipfft paths conditional on
  CRAY_MPICH_INC/CRAY_HIPFORT_LIB being set; fall back to standard
  find_package(MPI) and find_library(hipfft/amdhip64) so the self-contained
  container image works without any OLCF env vars loaded
- toolchain: add amd90a cluster profile (HPCFund gfx90a / MI250); fix module
  variable export loop so vars that reference previously exported vars expand correctly
diff --git a/.github/Dockerfile b/.github/Dockerfile
@@ -7,30 +7,68 @@ ARG CXX_COMPILER
 ARG FC_COMPILER
 ARG COMPILER_PATH
 ARG COMPILER_LD_LIBRARY_PATH
+ARG AFAR_VERSION
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV TZ=UTC
 
 RUN apt-get update -y && \
-    apt-get install -y software-properties-common ca-certificates gnupg && \
+    apt-get install -y software-properties-common ca-certificates gnupg wget && \
     add-apt-repository ppa:deadsnakes/ppa && \
     apt-get update -y && \
-    if [ "$TARGET" != "gpu" ]; then \
+    if [ "$TARGET" = "cpu" ]; then \
         apt-get install -y \
             build-essential git make cmake gcc g++ gfortran bc \
             python3.12 python3.12-venv python3-pip \
             openmpi-bin libopenmpi-dev libfftw3-dev \
             mpich libmpich-dev; \
-    else \
+    elif [ "$TARGET" = "gpu" ]; then \
         apt-get install -y \
             build-essential git make cmake bc \
             python3.12 python3.12-venv python3-pip \
             libfftw3-dev \
             openmpi-bin libopenmpi-dev; \
+    elif [ "$TARGET" = "amd" ]; then \
+        apt-get install -y \
+            build-essential git make gcc g++ gfortran bc \
+            python3.12 python3.12-venv python3-pip \
+            libfftw3-dev libnuma1 libdrm2 libdrm-amdgpu1; \
     fi && \
     update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 2 && \
     rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
 
+# AMD: download AFAR, install cmake 3.28 (Ubuntu 22.04 ships 3.22 which doesn't
+# recognize amdflang as LLVMFlang), then build MPICH with amdflang so the
+# generated mpi.mod is compiler-compatible.
+RUN if [ "$TARGET" = "amd" ] && [ -n "$AFAR_VERSION" ]; then \
+        OLCF_AFAR_ROOT="/opt/${AFAR_VERSION}" && \
+        wget -q "https://repo.radeon.com/rocm/misc/flang/${AFAR_VERSION}-ubuntu.tar.bz2" -O /tmp/afar.tar.bz2 && \
+        tar -xjf /tmp/afar.tar.bz2 -C /opt/ && \
+        rm /tmp/afar.tar.bz2 && \
+        CMAKE_VER=3.28.6 && \
+        wget -q "https://github.com/Kitware/CMake/releases/download/v${CMAKE_VER}/cmake-${CMAKE_VER}-linux-x86_64.sh" \
+            -O /tmp/cmake-install.sh && \
+        chmod +x /tmp/cmake-install.sh && \
+        /tmp/cmake-install.sh --prefix=/usr/local --skip-license --exclude-subdir && \
+        rm /tmp/cmake-install.sh && \
+        printf '#!/bin/bash\nargs=()\nwhile [ "$#" -gt 0 ]; do\n    if [ "$1" = "-soname" ]; then\n        args+=("-Wl,-soname,$2"); shift 2\n    else\n        args+=("$1"); shift\n    fi\ndone\nexec '"${OLCF_AFAR_ROOT}"'/bin/amdflang "${args[@]}"\n' \
+            > /usr/local/bin/amdflang-ld-wrap && \
+        chmod +x /usr/local/bin/amdflang-ld-wrap && \
+        MPICH_VER=3.4.3 && \
+        wget -q "https://www.mpich.org/static/downloads/${MPICH_VER}/mpich-${MPICH_VER}.tar.gz" \
+            -O /tmp/mpich.tar.gz && \
+        mkdir -p /tmp/mpich-src && \
+        tar -xzf /tmp/mpich.tar.gz -C /tmp/mpich-src --strip-components=1 && \
+        cd /tmp/mpich-src && \
+        FC=/usr/local/bin/amdflang-ld-wrap CC=gcc CXX=g++ \
+            ./configure --prefix=/opt/mpich --enable-shared --disable-static \
+            --with-device=ch3 2>&1 && \
+        make -j$(nproc) 2>&1 && \
+        make install 2>&1 && \
+        cd / && \
+        rm -rf /tmp/mpich-src /tmp/mpich.tar.gz; \
+    fi
+
 ENV OMPI_ALLOW_RUN_AS_ROOT=1
 ENV OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1
 ENV PATH="/opt/MFC:$PATH"
@@ -40,8 +78,9 @@ COPY ../ /opt/MFC
 ENV CC=${CC_COMPILER}
 ENV CXX=${CXX_COMPILER}
 ENV FC=${FC_COMPILER}
-ENV PATH="${COMPILER_PATH}:$PATH"
-ENV LD_LIBRARY_PATH="${COMPILER_LD_LIBRARY_PATH}:${LD_LIBRARY_PATH:-}"
+ENV OLCF_AFAR_ROOT=/opt/${AFAR_VERSION}
+ENV PATH="${COMPILER_PATH}:/opt/mpich/bin:$PATH"
+ENV LD_LIBRARY_PATH="${COMPILER_LD_LIBRARY_PATH}:/opt/mpich/lib:${LD_LIBRARY_PATH:-}"
 
 # Pre-install numpy into the venv before mfc.sh runs, as it's required at
 # build time by several dependencies (pandas, cantera, matplotlib, etc.) that
@@ -54,13 +93,17 @@ RUN echo "TARGET=$TARGET CC=$CC_COMPILER FC=$FC_COMPILER" && \
     cd /opt/MFC && \
     if [ "$TARGET" = "gpu" ]; then \
       ./mfc.sh build --gpu -j $(nproc); \
+    elif [ "$TARGET" = "amd" ]; then \
+      ./mfc.sh build --gpu mp -j $(nproc); \
     else \
       ./mfc.sh build -j $(nproc); \
     fi
 
 RUN cd /opt/MFC && \
     if [ "$TARGET" = "gpu" ]; then \
       ./mfc.sh test -a --dry-run --gpu -j $(nproc); \
+    elif [ "$TARGET" = "amd" ]; then \
+      ./mfc.sh test -a --dry-run --gpu mp -j $(nproc); \
     else \
       ./mfc.sh test -a --dry-run -j $(nproc); \
     fi
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
@@ -24,6 +24,7 @@ jobs:
             - { name: 'cpu', runner: 'ubuntu-22.04-arm', base_image: 'ubuntu:22.04' }
             - { name: 'gpu', runner: 'ubuntu-22.04',     base_image: 'nvcr.io/nvidia/nvhpc:23.11-devel-cuda_multi-ubuntu22.04' }
             - { name: 'gpu', runner: 'ubuntu-22.04-arm', base_image: 'nvcr.io/nvidia/nvhpc:23.11-devel-cuda_multi-ubuntu22.04' }
+            - { name: 'amd', runner: 'ubuntu-22.04',     base_image: 'ubuntu:22.04' }
     runs-on: ${{ matrix.config.runner }}
     outputs:
       tag: ${{ steps.clone.outputs.tag }}
@@ -113,6 +114,24 @@ jobs:
           tags: ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-${{ matrix.config.name }}-${{ matrix.config.runner}}
           push: true
 
+      - name: Build and push image (amd)
+        if: ${{ matrix.config.name == 'amd' }}
+        uses: docker/build-push-action@v6
+        with:
+          context: /mnt/share
+          file: /mnt/share/Dockerfile
+          build-args: |
+            BASE_IMAGE=${{ matrix.config.base_image }}
+            TARGET=amd
+            AFAR_VERSION=rocm-afar-8873-drop-22.2.0
+            CC_COMPILER=gcc
+            CXX_COMPILER=g++
+            FC_COMPILER=/opt/rocm-afar-8873-drop-22.2.0/bin/amdflang
+            COMPILER_PATH=/opt/rocm-afar-8873-drop-22.2.0/lib/llvm/bin:/opt/rocm-afar-8873-drop-22.2.0/bin
+            COMPILER_LD_LIBRARY_PATH=/opt/rocm-afar-8873-drop-22.2.0/lib:/opt/rocm-afar-8873-drop-22.2.0/lib/llvm/lib
+          tags: ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-amd-ubuntu-22.04
+          push: true
+
   manifests:
     runs-on: ubuntu-latest
     needs: Container
@@ -133,6 +152,7 @@ jobs:
         run: |
           docker buildx imagetools create -t $REGISTRY:$TAG-cpu  $REGISTRY:$TAG-cpu-ubuntu-22.04 $REGISTRY:$TAG-cpu-ubuntu-22.04-arm
           docker buildx imagetools create -t $REGISTRY:$TAG-gpu  $REGISTRY:$TAG-gpu-ubuntu-22.04 $REGISTRY:$TAG-gpu-ubuntu-22.04-arm
+          docker buildx imagetools create -t $REGISTRY:$TAG-amd  $REGISTRY:$TAG-amd-ubuntu-22.04
 
       - name: Update latest tags
         if: github.event_name == 'release'
@@ -142,3 +162,4 @@ jobs:
         run: |
           docker buildx imagetools create -t $REGISTRY:latest-cpu $REGISTRY:$TAG-cpu-ubuntu-22.04 $REGISTRY:$TAG-cpu-ubuntu-22.04-arm
           docker buildx imagetools create -t $REGISTRY:latest-gpu $REGISTRY:$TAG-gpu-ubuntu-22.04 $REGISTRY:$TAG-gpu-ubuntu-22.04-arm
+          docker buildx imagetools create -t $REGISTRY:latest-amd $REGISTRY:$TAG-amd-ubuntu-22.04
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -563,10 +563,11 @@ exit 0
             find_package(MPI COMPONENTS Fortran REQUIRED)
 	    
             target_compile_definitions(${a_target} PRIVATE MFC_MPI)
-            if(CMAKE_Fortran_COMPILER_ID STREQUAL "LLVMFlang")
+            if(CMAKE_Fortran_COMPILER_ID STREQUAL "LLVMFlang" AND
+               DEFINED ENV{CRAY_MPICH_INC} AND NOT "$ENV{CRAY_MPICH_INC}" STREQUAL "")
                 target_compile_options(${a_target} PRIVATE "$ENV{CRAY_MPICH_INC}")
                 target_link_libraries(${a_target} PRIVATE $ENV{CRAY_MPICH_LIB})
-            else() 
+            else()
                 target_link_libraries(${a_target} PRIVATE MPI::MPI_Fortran)
             endif()
         endif()
@@ -587,10 +588,16 @@ exit 0
                     find_package(CUDAToolkit REQUIRED)
                     target_link_libraries(${a_target} PRIVATE CUDA::cudart CUDA::cufft)
                 elseif(CMAKE_Fortran_COMPILER_ID STREQUAL "LLVMFlang")
-                    target_link_libraries(${a_target} PRIVATE $ENV{CRAY_HIPFORT_LIB})
-                else() 
+                    if(DEFINED ENV{CRAY_HIPFORT_LIB} AND NOT "$ENV{CRAY_HIPFORT_LIB}" STREQUAL "")
+                        target_link_libraries(${a_target} PRIVATE $ENV{CRAY_HIPFORT_LIB})
+                    else()
+                        find_library(HIPFFT_LIB hipfft
+                            HINTS "$ENV{OLCF_AFAR_ROOT}/lib" REQUIRED)
+                        target_link_libraries(${a_target} PRIVATE ${HIPFFT_LIB})
+                    endif()
+                else()
                     find_package(hipfort COMPONENTS hipfft CONFIG REQUIRED)
-                    target_link_libraries(${a_target} PRIVATE hipfort::hipfft)                    
+                    target_link_libraries(${a_target} PRIVATE hipfort::hipfft)
                 endif()
             else()
                 find_package(FFTW REQUIRED)
@@ -703,8 +710,14 @@ exit 0
                         PRIVATE -DFRONTIER_UNIFIED)
                 endif()
 
-                find_package(hipfort COMPONENTS hip CONFIG REQUIRED)
-                target_link_libraries(${a_target} PRIVATE hipfort::hip hipfort::hipfort-amdgcn flang_rt.hostdevice)
+                find_library(HIP_LIB amdhip64
+                    HINTS "$ENV{OLCF_AFAR_ROOT}/lib" REQUIRED)
+                find_library(HIPFORT_AMDGCN_LIB hipfort-amdgcn
+                    HINTS "$ENV{OLCF_AFAR_ROOT}/lib" REQUIRED)
+                target_include_directories(${a_target} PRIVATE
+                    "$ENV{OLCF_AFAR_ROOT}/include/hipfort/amdgcn")
+                target_link_libraries(${a_target} PRIVATE
+                    ${HIP_LIB} ${HIPFORT_AMDGCN_LIB} flang_rt.hostdevice)
             endif()
         elseif (CMAKE_Fortran_COMPILER_ID STREQUAL "Cray")
             target_compile_options(${a_target} PRIVATE "SHELL:-h noacc" "SHELL:-x acc")
diff --git a/toolchain/bootstrap/modules.sh b/toolchain/bootstrap/modules.sh
@@ -105,7 +105,6 @@ fi
 
 ELEMENTS="$(__extract "$u_c-all") $(__extract "$u_c-$cg")"
 MODULES=`echo "$ELEMENTS" | tr ' ' '\n' | grep -v = | xargs`
-VARIABLES=`echo "$ELEMENTS" | tr ' ' '\n' | grep = | xargs`
 
 log " $ module load $MODULES"
 if ! module load $MODULES; then
@@ -114,10 +113,17 @@ if ! module load $MODULES; then
     return
 fi
 
-if [ $(echo "$VARIABLES" | grep = | wc -c) -gt 0 ]; then
-    log " $ export $(eval "echo $VARIABLES")"
-    export $(eval "echo $VARIABLES") > /dev/null
-fi
+# Export variables one line at a time so each can reference previously exported vars
+# (e.g. PATH="${OLCF_AFAR_ROOT}/..." requires OLCF_AFAR_ROOT to already be set)
+for _suffix in "all" "$cg"; do
+    while IFS= read -r _entry; do
+        if echo "$_entry" | grep -q '='; then
+            log " $ export $(eval "echo \"$_entry\"")"
+            eval "export $_entry"
+        fi
+    done < <(grep -E "^$u_c-$_suffix\s+" toolchain/modules | sed "s/^$u_c-$_suffix\s\+//")
+done
+unset _suffix _entry
 
 UNLOAD_MODULES="$(__extract "$u_c-all-unload") $(__extract "$u_c-$cg-unload")"
 UNLOAD_MODULES=$(echo "$UNLOAD_MODULES" | xargs)
diff --git a/toolchain/modules b/toolchain/modules
@@ -46,11 +46,24 @@ f-all cpe/25.03 rocm/6.3.1
 f-all cray-fftw cray-hdf5 python cmake
 f-gpu python craype-accel-amd-gfx90a rocprofiler-compute/3.0.0
 
-famd    OLCF Frontier AMD 
+famd    OLCF Frontier AMD
 famd-all python cmake
 famd-all cpe/25.09
 famd-all PrgEnv-amd
 
+amd90a  HPCFund gfx90a compiler (MI250)
+amd90a-gpu rocm/7.2.0
+amd90a-gpu OLCF_AFAR_ROOT="/work1/spencerbryngelson/sbryngelson/software/therock-afar-23.2.1-gfx90a-7.13.0-7357b5084b"
+amd90a-gpu PATH="${OLCF_AFAR_ROOT}/lib/llvm/bin:${PATH}"
+amd90a-gpu LD_LIBRARY_PATH="${OLCF_AFAR_ROOT}/lib:${OLCF_AFAR_ROOT}/lib/llvm/lib:${LD_LIBRARY_PATH}"
+amd90a-gpu CRAY_HIPFORT_INC="-I${OLCF_AFAR_ROOT}/include/hipfort/amdgcn"
+amd90a-gpu CRAY_HIPFORT_LIB="-L${OLCF_AFAR_ROOT}/lib -lhipfft"
+amd90a-gpu CRAY_HIP_INC="-I${OLCF_AFAR_ROOT}/include/hip"
+amd90a-gpu CMAKE_PREFIX_PATH="${OLCF_AFAR_ROOT}:${CMAKE_PREFIX_PATH}"
+amd90a-gpu CC="${OLCF_AFAR_ROOT}/lib/llvm/bin/clang"
+amd90a-gpu CXX="${OLCF_AFAR_ROOT}/lib/llvm/bin/clang++"
+amd90a-gpu FC="${OLCF_AFAR_ROOT}/bin/amdflang"
+
 tuo     OLCF Tuolumne
 tuo-all cpe/25.03 rocm/6.3.1
 tuo-all cray-fftw/3.3.10.9 cray-hdf5 python/3.12.2 cmake
