Anchor GPU supported GPU archs in dockerfile.

Author: Hongtao Zhang

dockerfile/rocm6.3.x.dockerfile

@@ -7,6 +7,7 @@ FROM ${BASE_IMAGE}
 #   - Ubuntu: 22.04
 #   - Python: 3.10
 #   - ROCm: 6.3.4
+#   - openmpi: 4.0.7rc2
 #   - torch: 2.8.0a0+git7d205b2
 #   - rccl: 2.21.5.60304-76
 #   - hipblaslt: 0.15.0-8c69191d
@@ -17,9 +18,8 @@ FROM ${BASE_IMAGE}
 #   - amd-smi: 25.1.0+8dc45db
 # Added by this Dockerfile:
 #   - Docker Client: 27.5.1
-#   - openmpi: pre-installed at /opt/ompi
 #   - mlc: v3.12
-#   - OFED: 24.10-1.1.4.0 LTS (if not present)
+#   - OFED: 24.10-1.1.4.0 LTS
 
 # Fix base image botocore/urllib3 incompatibility:
 # Base image ships botocore 1.22.12 (expects urllib3 1.x) with urllib3 2.6.3,
@@ -104,6 +104,9 @@ RUN if ! command -v ofed_info >/dev/null 2>&1; then \
 
 ENV ROCM_PATH=/opt/rocm
 
+# Target GPU architectures for ROCm builds (space-separated)
+ENV AMDGPU_TARGETS="gfx908 gfx90a gfx942"
+
 # Use pre-installed OpenMPI from base image at /opt/ompi
 ENV MPI_HOME=/opt/ompi
 

third_party/Makefile

@@ -134,9 +134,17 @@ ifneq (,$(wildcard fio/Makefile))
 endif
 
 # Build rccl-tests from commit 46375b1 of default branch.
+# If AMDGPU_TARGETS env var is set (space-separated, e.g. "gfx908 gfx90a gfx942"),
+# explicit --offload-arch flags and include paths are used.
+# Otherwise, the original build command is used (relies on hipcc auto-detection).
+ROCM_OFFLOAD_ARCH_FLAGS := $(foreach arch,$(AMDGPU_TARGETS),--offload-arch=$(arch))
 rocm_rccl_tests: sb_micro_path
 ifneq (, $(wildcard rccl-tests/Makefile))
+ifdef AMDGPU_TARGETS
+	cd ./rccl-tests && make MPI=1 MPI_HOME=$(MPI_HOME) HIPCUFLAGS="-std=c++14 -O3 $(ROCM_OFFLOAD_ARCH_FLAGS) -I$(ROCM_PATH)/include -I$(ROCM_PATH)/include/rccl -I$(ROCM_PATH)/include/hip -DMPI_SUPPORT -I$(MPI_HOME)/include -I$(MPI_HOME)/include/mpi" -j
+else
 	cd ./rccl-tests && make MPI=1 MPI_HOME=$(MPI_HOME) -j
+endif
 	cp -v -r ./rccl-tests/build/* $(SB_MICRO_PATH)/bin/
 endif
 
@@ -168,6 +176,7 @@ rocm_hipblaslt: sb_micro_path
 # Build hipBusBandwidth.
 # HIP is released with rocm, like rocm-4.2.0 and so on.
 # The version we use is the released tag which is consistent with the rocm version in the environment or docker.
+
 rocm_bandwidthTest: sb_micro_path
 	git clone -b ${ROCM_VER} https://github.com/ROCm-Developer-Tools/HIP.git
 	cd ./HIP/samples/1_Utils/hipBusBandwidth/ && mkdir -p build && cd build && cmake .. && make