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

sbryngelson · web-flow · commit f1fd862cc661 · 2026-05-11T18:58:19.000-04:00
diff --git a/.github/Dockerfile b/.github/Dockerfile
@@ -7,41 +7,82 @@ ARG CXX_COMPILER
 ARG FC_COMPILER
 ARG COMPILER_PATH
 ARG COMPILER_LD_LIBRARY_PATH
+ARG AFAR_VERSION
+ARG OLCF_AFAR_ROOT=""
 
 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 HYDRA_LAUNCHER=fork
 ENV PATH="/opt/MFC:$PATH"
 
 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=${OLCF_AFAR_ROOT}
+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 +95,17 @@ RUN echo "TARGET=$TARGET CC=$CC_COMPILER FC=$FC_COMPILER" && \
     cd /opt/MFC && \
     if [ "$TARGET" = "gpu" ]; then \
       ./mfc.sh build --gpu acc -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 acc -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,10 +24,11 @@ jobs:
     strategy:
       matrix:
         config:
-            - { name: 'cpu', runner: 'ubuntu-22.04',     base_image: 'ubuntu:22.04' }
-            - { name: 'cpu', runner: 'ubuntu-22.04-arm', base_image: 'ubuntu:22.04' }
-            - { name: 'gpu', runner: 'ubuntu-22.04',     base_image: 'nvcr.io/nvidia/nvhpc:24.5-devel-cuda_multi-ubuntu22.04', compiler_arch: 'Linux_x86_64' }
-            - { name: 'gpu', runner: 'ubuntu-22.04-arm', base_image: 'nvcr.io/nvidia/nvhpc:24.5-devel-cuda_multi-ubuntu22.04', compiler_arch: 'Linux_aarch64' }
+            - { name: 'cpu',  tag: 'cpu',       runner: 'ubuntu-22.04',     base_image: 'ubuntu:22.04' }
+            - { name: 'cpu',  tag: 'cpu',       runner: 'ubuntu-22.04-arm', base_image: 'ubuntu:22.04' }
+            - { name: 'gpu',  tag: 'gpu-nvidia', runner: 'ubuntu-22.04',     base_image: 'nvcr.io/nvidia/nvhpc:24.5-devel-cuda_multi-ubuntu22.04', compiler_arch: 'Linux_x86_64' }
+            - { name: 'gpu',  tag: 'gpu-nvidia', runner: 'ubuntu-22.04-arm', base_image: 'nvcr.io/nvidia/nvhpc:24.5-devel-cuda_multi-ubuntu22.04', compiler_arch: 'Linux_aarch64' }
+            - { name: 'amd',  tag: 'gpu-amd',   runner: 'ubuntu-22.04',     base_image: 'ubuntu:22.04' }
     runs-on: ${{ matrix.config.runner }}
     outputs:
       tag: ${{ steps.clone.outputs.tag }}
@@ -108,8 +109,8 @@ jobs:
           labels: |
             org.opencontainers.image.source=https://github.com/${{ github.repository }}
           tags: |
-            ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-${{ matrix.config.name }}-${{ matrix.config.runner }}
-            ${{ env.GH_REGISTRY }}:${{ env.TAG }}-${{ matrix.config.name }}-${{ matrix.config.runner }}
+            ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-${{ matrix.config.tag }}-${{ matrix.config.runner }}
+            ${{ env.GH_REGISTRY }}:${{ env.TAG }}-${{ matrix.config.tag }}-${{ matrix.config.runner }}
           push: true
 
       - name: Build and push image (gpu)
@@ -130,8 +131,39 @@ jobs:
           labels: |
             org.opencontainers.image.source=https://github.com/${{ github.repository }}
           tags: |
-            ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-${{ matrix.config.name }}-${{ matrix.config.runner }}
-            ${{ env.GH_REGISTRY }}:${{ env.TAG }}-${{ matrix.config.name }}-${{ matrix.config.runner }}
+            ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-${{ matrix.config.tag }}-${{ matrix.config.runner }}
+            ${{ env.GH_REGISTRY }}:${{ env.TAG }}-${{ matrix.config.tag }}-${{ matrix.config.runner }}
+          push: true
+
+      - name: Set AMD AFAR vars
+        if: ${{ matrix.config.name == 'amd' }}
+        run: |
+          AFAR=rocm-afar-8873-drop-22.2.0
+          ROOT="/opt/${AFAR}"
+          echo "AFAR_VERSION=${AFAR}"      >> $GITHUB_ENV
+          echo "AFAR_ROOT=${ROOT}"         >> $GITHUB_ENV
+
+      - 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=${{ env.AFAR_VERSION }}
+            OLCF_AFAR_ROOT=${{ env.AFAR_ROOT }}
+            CC_COMPILER=gcc
+            CXX_COMPILER=g++
+            FC_COMPILER=${{ env.AFAR_ROOT }}/bin/amdflang
+            COMPILER_PATH=${{ env.AFAR_ROOT }}/lib/llvm/bin:${{ env.AFAR_ROOT }}/bin
+            COMPILER_LD_LIBRARY_PATH=${{ env.AFAR_ROOT }}/lib:${{ env.AFAR_ROOT }}/lib/llvm/lib
+          labels: |
+            org.opencontainers.image.source=https://github.com/${{ github.repository }}
+          tags: |
+            ${{ secrets.DOCKERHUB_USERNAME }}/mfc:${{ env.TAG }}-gpu-amd-ubuntu-22.04
+            ${{ env.GH_REGISTRY }}:${{ env.TAG }}-gpu-amd-ubuntu-22.04
           push: true
 
   manifests:
@@ -169,8 +201,9 @@ jobs:
         run: |
           GH="${{ env.GH_REGISTRY }}"
           for R in "$DH" "$GH"; do
-            docker buildx imagetools create -t $R:$TAG-cpu $R:$TAG-cpu-ubuntu-22.04 $R:$TAG-cpu-ubuntu-22.04-arm
-            docker buildx imagetools create -t $R:$TAG-gpu $R:$TAG-gpu-ubuntu-22.04 $R:$TAG-gpu-ubuntu-22.04-arm
+            docker buildx imagetools create -t $R:$TAG-cpu       $R:$TAG-cpu-ubuntu-22.04 $R:$TAG-cpu-ubuntu-22.04-arm
+            docker buildx imagetools create -t $R:$TAG-gpu-nvidia $R:$TAG-gpu-nvidia-ubuntu-22.04 $R:$TAG-gpu-nvidia-ubuntu-22.04-arm
+            docker buildx imagetools create -t $R:$TAG-gpu-amd   $R:$TAG-gpu-amd-ubuntu-22.04
           done
 
       - name: Update latest tags
@@ -181,6 +214,7 @@ jobs:
         run: |
           GH="${{ env.GH_REGISTRY }}"
           for R in "$DH" "$GH"; do
-            docker buildx imagetools create -t $R:latest-cpu $R:$TAG-cpu-ubuntu-22.04 $R:$TAG-cpu-ubuntu-22.04-arm
-            docker buildx imagetools create -t $R:latest-gpu $R:$TAG-gpu-ubuntu-22.04 $R:$TAG-gpu-ubuntu-22.04-arm
+            docker buildx imagetools create -t $R:latest-cpu       $R:$TAG-cpu-ubuntu-22.04 $R:$TAG-cpu-ubuntu-22.04-arm
+            docker buildx imagetools create -t $R:latest-gpu-nvidia $R:$TAG-gpu-nvidia-ubuntu-22.04 $R:$TAG-gpu-nvidia-ubuntu-22.04-arm
+            docker buildx imagetools create -t $R:latest-gpu-amd   $R:$TAG-gpu-amd-ubuntu-22.04
           done
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,7 +46,7 @@ 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
