Update to rocm 7.1 and convert rocshmem to a sparse checkout

.github/workflows/amd-ci.yml

@@ -12,10 +12,10 @@ jobs:
     runs-on: "amd-gfx942-mi325"
     timeout-minutes: 45
     container:
-        image: rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.5.1_preview
+        image: rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1
         options: >-
           --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root
-          --volume /home/runner/.triton:/github/home/.triton
+          --volume /home/runner/.triton:/github/home/.triton --shm-size=512MB
     env:
       TRITON_BUILD_WITH_CLANG_LLD: "TRUE"
       TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE"
@@ -26,8 +26,8 @@ jobs:
       - name: Install dependencies
         run: |
           rm -f /usr/local/bin/cmake
-          apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake
-          pip3 install -i https://test.pypi.org/simple hip-python>=6.3.0 # (or whatever Rocm version you have)
+          apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake miopen-hip autoconf libtool flex ninja-build clang lld
+          pip3 install -i https://test.pypi.org/simple hip-python>=7.1 # (or whatever Rocm version you have)
           pip3 install pybind11
       - name: Checkout
         uses: actions/checkout@v4
@@ -58,10 +58,10 @@ jobs:
     runs-on: "amd-gfx942-mi325"
     timeout-minutes: 60
     container:
-        image: rocm/pytorch:rocm6.3_ubuntu22.04_py3.10_pytorch_release_2.5.1_preview
+        image: rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1
         options: >-
           --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root
-          --volume /home/runner/.triton:/github/home/.triton
+          --volume /home/runner/.triton:/github/home/.triton --shm-size=512MB
     env:
       TRITON_BUILD_WITH_CLANG_LLD: "TRUE"
       TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE"
@@ -72,9 +72,9 @@ jobs:
     steps:
       - name: Install dependencies
         run: |
-          apt-get update -y && apt install -y libopenmpi-dev
-          pip3 install -i https://test.pypi.org/simple hip-python>=6.3.0 # (or whatever Rocm version you have)
-          pip3 install pybind11
+          apt-get update -y && apt install -y libopenmpi-dev cmake miopen-hip autoconf libtool flex ninja-build clang lld
+          pip3 install -i https://test.pypi.org/simple hip-python>=7.1 # (or whatever Rocm version you have)
+          pip3 install pybind11 psutil
       - name: Checkout
         uses: actions/checkout@v4
         with:
@@ -103,10 +103,10 @@ jobs:
     runs-on: "amd-gfx942-mi325"
     timeout-minutes: 45
     container:
-        image: rocm/pytorch:rocm6.4.3_ubuntu22.04_py3.10_pytorch_release_2.5.1
+        image: rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.7.1
         options: >-
           --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root
-          --volume /home/runner/.triton:/github/home/.triton
+          --volume /home/runner/.triton:/github/home/.triton --shm-size=512MB
     env:
       TRITON_BUILD_WITH_CLANG_LLD: "TRUE"
       TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE"
@@ -118,8 +118,8 @@ jobs:
       - name: Install dependencies
         run: |
           rm -f /usr/local/bin/cmake
-          apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake
-          pip3 install -i https://test.pypi.org/simple hip-python==6.4.3.555.40
+          apt-get update -y && apt install -y libopenmpi-dev git cython3 ibverbs-utils openmpi-bin libopenmpi-dev libpci-dev libdw1 locales cmake ninja-build clang lld
+          pip3 install -i https://test.pypi.org/simple hip-python>=7.1
           pip3 install pybind11
       - name: Checkout
         uses: actions/checkout@v4

.gitignore

@@ -61,6 +61,9 @@ python/triton_dist/_C/*nvshmem
 hf_configs*
 *.whl
 
+3rdparty/rocm-systems
+3rdparty/rocshmem
+
 # conflicts
 temp_conflicts
 temp_merge_repo
@@ -82,4 +85,4 @@ tmp*
 *zst*
 
 .cursor
-*.csv
+*.csv

.gitmodules

@@ -1,6 +1,3 @@
-[submodule "3rdparty/rocshmem"]
-	path = 3rdparty/rocshmem
-	url = https://github.com/ROCm/rocSHMEM.git
 [submodule "3rdparty/triton"]
 	path = 3rdparty/triton
 	url = https://github.com/ByteDance-Seed/triton.git

python/triton_dist/test/amd/test_rocshmem_api.py

@@ -209,9 +209,8 @@ def test_rocshmem_memcpy():
 
             HIP_CHECK(cp_res)
 
-    pyrocshmem.rocshmem_barrier_all_on_stream(cur_stream.cuda_stream)
-
     torch.cuda.synchronize()
+    pyrocshmem.rocshmem_barrier_all_on_stream(cur_stream.cuda_stream)
 
     try:
         torch.testing.assert_close(comm_buffs[peer], one)

shmem/rocshmem_bind/build.sh

@@ -6,11 +6,6 @@ set -e
 SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
 PROJECT_ROOT=$(realpath ${SCRIPT_DIR})
 
-function apt_install_deps() {
-    apt update -y
-    apt-get install -y miopen-hip
-}
-
 function build_pyrocshmem_cmake() {
   pushd ${PROJECT_ROOT}/pyrocshmem
   mkdir -p build
@@ -58,8 +53,6 @@ export ROCM_INSTALL_DIR="/opt/rocm"
 export PATH="${OPENMPI_UCX_INSTALL_DIR}/bin:$PATH"
 export LD_LIBRARY_PATH="${OPENMPI_UCX_INSTALL_DIR}/lib:$LD_LIBRARY_PATH"
 
-apt_install_deps
-
 # build rocshmem
 bash -x ${PROJECT_ROOT}/build_rocshmem.sh
 # build rocshmem bitcode

shmem/rocshmem_bind/build_rocshmem.sh

@@ -7,6 +7,49 @@ SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
 PROJECT_ROOT=$(realpath ${SCRIPT_DIR})
 ROCSHMEM_SRC_DIR=${PROJECT_ROOT}/../../3rdparty/rocshmem
 
+sys_path="${PROJECT_ROOT}/../../3rdparty/rocm-systems"
+
+if [ -d "${ROCSHMEM_SRC_DIR}" ]; then
+  pushd "${PROJECT_ROOT}/../.."
+  active=$(git config submodule.3rdparty/rocshmem.active || echo "nil")
+  if [ "${active}" = "true" ]; then
+    echo "Error: Rocshmem submodule still active, please delete it"
+    quit=1
+  fi
+  popd
+  pushd "${ROCSHMEM_SRC_DIR}"
+  url=$(git remote get-url origin || echo "nil")
+  if [ "${url}" = "https://github.com/ROCm/rocSHMEM.git" ]; then
+    echo "Error: Old rocshmem checkout found, please delete it"
+    quit=1
+  fi
+  popd
+  if ! [ -z "${quit}" ]; then
+    exit $quit
+  fi
+
+  if ! [ "$(ls -A "${ROCSHMEM_SRC_DIR}")" ]; then
+    rmdir "${ROCSHMEM_SRC_DIR}"
+  fi
+fi
+
+rocm_systems_tag=hip-version_7.12.60610
+
+if ! [ -d "${ROCSHMEM_SRC_DIR}" ]; then
+  echo "Creating sparse checkout"
+  pushd "${PROJECT_ROOT}/../.."
+  git clone "https://github.com/ROCm/rocm-systems.git" -b "${rocm_systems_tag}" --depth 1 --sparse "${sys_path}"
+  popd
+  pushd "${sys_path}"
+  git config core.sparseCheckoutCone true
+  git sparse-checkout set projects/rocshmem
+  ln -s rocm-systems/projects/rocshmem ../rocshmem
+  popd
+fi
+
+pushd "${sys_path}"
+git checkout "${rocm_systems_tag}"
+popd
 
 pushd ${ROCSHMEM_SRC_DIR}
 
@@ -17,7 +60,6 @@ OMPI_INSTALL_DIR="${OMPI_INSTALL_DIR:-/opt/ompi_build}"
 # build ompi, ucx
 if [ ! -e "${OMPI_INSTALL_DIR}" ]; then
     # prepare for building ompi, ucx
-    apt-get install autoconf libtool flex -y
     BUILD_DIR=${OMPI_INSTALL_DIR} bash ${ROCSHMEM_SRC_DIR}/scripts/install_dependencies.sh
 else
     echo "ompi exists, skip building ompi and ucx"

shmem/rocshmem_bind/pyrocshmem/setup.py

@@ -108,7 +108,7 @@ def setup_pytorch_extension() -> setuptools.Extension:
     """Setup CppExtension for PyTorch support"""
     include_dirs, library_dirs, libraries = [], [], []
 
-    deps = [hip_deps(), mpi_deps(), rocshmem_deps()]
+    deps = [rocshmem_deps(), hip_deps(), mpi_deps()]
 
     for include_dir, library_dir, library in deps:
         include_dirs += include_dir

shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh

@@ -39,12 +39,11 @@ SOURCE_MAP=(
     ["${ROCSHMEM_SRC}/src/ipc/context_ipc_device_coll.cpp"]="rocshmem_context_ipc_device_coll.bc"
     ["${ROCSHMEM_SRC}/src/ipc_policy.cpp"]="rocshmem_ipc_policy.bc"
     ["${ROCSHMEM_SRC}/src/gda/context_gda_device.cpp"]="rocshmem_context_gda_device.bc"
+    ["${ROCSHMEM_SRC}/src/gda/context_gda_device_coll.cpp"]="rocshmem_context_gda_device_coll.bc"
     ["${ROCSHMEM_SRC}/src/gda/backend_gda.cpp"]="rocshmem_backend_gda.bc"
     ["${ROCSHMEM_SRC}/src/gda/queue_pair.cpp"]="rocshmem_queue_pair.bc"
     ["${ROCSHMEM_SRC}/src/gda/ionic/queue_pair_ionic.cpp"]="rocshmem_queue_pair_ionic.bc"
     ["${ROCSHMEM_SRC}/src/gda/mlx5/queue_pair_mlx5.cpp"]="rocshmem_queue_pair_mlx5.bc"
-    ["${ROCSHMEM_SRC}/src/gda/mlx5/segment_builder.cpp"]="rocshmem_segment_builder.bc"
-    ["${ROCSHMEM_SRC}/src/gda/endian.cpp"]="rocshmem_endian.bc"
     ["${ROCSHMEM_SRC}/src/team.cpp"]="rocshmem_team.bc"
     ["${ROCSHMEM_SRC}/src/sync/abql_block_mutex.cpp"]="rocshmem_abql_block_mutex.bc"
     ["${ROCSHMEM_SRC}/src/util.cpp"]="rocshmem_util.bc"