From 3b9ac2f44807fd37be46a936aa52bac142435549 Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Mon, 9 Mar 2026 21:27:12 -0500 Subject: [PATCH 1/7] Switch rocshmem to a sparse checkout of rocm-systems Signed-off-by: Eric Eaton --- .gitmodules | 3 -- 3rdparty/rocshmem | 1 - .../scripts/build_rshm_ipc_single.sh | 41 ++++++++++++++++++- 3 files changed, 40 insertions(+), 5 deletions(-) delete mode 160000 3rdparty/rocshmem diff --git a/.gitmodules b/.gitmodules index 11e11c8cc..c61680680 100644 --- a/.gitmodules +++ b/.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 diff --git a/3rdparty/rocshmem b/3rdparty/rocshmem deleted file mode 160000 index 7b5765ec0..000000000 --- a/3rdparty/rocshmem +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 7b5765ec0e137db04b50d657ed83ea239be8dd58 diff --git a/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh b/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh index 0915bb267..207a23a93 100755 --- a/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh +++ b/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh @@ -9,7 +9,46 @@ else install_path=$1 fi -src_path=$(dirname "$(realpath $0)")/../../../3rdparty/rocshmem/ +root_path="$(dirname "$(realpath $0)")/../../.." +src_path="${root_path}/3rdparty/rocshmem" +sys_path="${root_path}/3rdparty/rocm-systems" + +if [ -d "${src_path}" ]; then + pushd "${root_path}" + 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 "${src_path}" + 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 "${src_path}")" ]; then + rmdir "${src_path}" + fi +fi + +if ! [ -d "${src_path}" ]; then + echo "Creating sparse checkout" + tag=hip-version_7.12.60610 + pushd "${root_path}" + git clone "https://github.com/ROCm/rocm-systems.git" -b "${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 cmake \ -DCMAKE_BUILD_TYPE=Release \ From 002a8338993b83752507a181648385bc7ddadf5a Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Wed, 18 Mar 2026 17:08:30 -0500 Subject: [PATCH 2/7] Fix build with sparse checkout Signed-off-by: Eric Eaton --- shmem/rocshmem_bind/build_rocshmem.sh | 38 +++++++++++++++++ shmem/rocshmem_bind/pyrocshmem/setup.py | 2 +- .../scripts/build_rocshmem_device_bc.sh | 3 +- .../scripts/build_rshm_ipc_single.sh | 41 +------------------ 4 files changed, 41 insertions(+), 43 deletions(-) diff --git a/shmem/rocshmem_bind/build_rocshmem.sh b/shmem/rocshmem_bind/build_rocshmem.sh index 9bb464288..0c614e1e3 100755 --- a/shmem/rocshmem_bind/build_rocshmem.sh +++ b/shmem/rocshmem_bind/build_rocshmem.sh @@ -7,6 +7,44 @@ 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 + +if ! [ -d "${ROCSHMEM_SRC_DIR}" ]; then + echo "Creating sparse checkout" + tag=hip-version_7.12.60610 + pushd "${PROJECT_ROOT}/../.." + git clone "https://github.com/ROCm/rocm-systems.git" -b "${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 ${ROCSHMEM_SRC_DIR} diff --git a/shmem/rocshmem_bind/pyrocshmem/setup.py b/shmem/rocshmem_bind/pyrocshmem/setup.py index 250a30566..1f242a224 100644 --- a/shmem/rocshmem_bind/pyrocshmem/setup.py +++ b/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 diff --git a/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh b/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh index d1797c564..1dc3b56cf 100755 --- a/shmem/rocshmem_bind/scripts/build_rocshmem_device_bc.sh +++ b/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" diff --git a/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh b/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh index 207a23a93..0915bb267 100755 --- a/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh +++ b/shmem/rocshmem_bind/scripts/build_rshm_ipc_single.sh @@ -9,46 +9,7 @@ else install_path=$1 fi -root_path="$(dirname "$(realpath $0)")/../../.." -src_path="${root_path}/3rdparty/rocshmem" -sys_path="${root_path}/3rdparty/rocm-systems" - -if [ -d "${src_path}" ]; then - pushd "${root_path}" - 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 "${src_path}" - 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 "${src_path}")" ]; then - rmdir "${src_path}" - fi -fi - -if ! [ -d "${src_path}" ]; then - echo "Creating sparse checkout" - tag=hip-version_7.12.60610 - pushd "${root_path}" - git clone "https://github.com/ROCm/rocm-systems.git" -b "${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 +src_path=$(dirname "$(realpath $0)")/../../../3rdparty/rocshmem/ cmake \ -DCMAKE_BUILD_TYPE=Release \ From b95838c58eb34f23ae797e84a05039786c374832 Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Thu, 26 Mar 2026 18:14:59 -0500 Subject: [PATCH 3/7] Update mori --- 3rdparty/mori | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/mori b/3rdparty/mori index f07a21c92..cef7ecb92 160000 --- a/3rdparty/mori +++ b/3rdparty/mori @@ -1 +1 @@ -Subproject commit f07a21c926e827bf1a5c24bb5338b14be88b6702 +Subproject commit cef7ecb921e8abfbfd74e376e15d2fd283ab0004 From 974dfc1f8947748c2d82c9c786916d1ebd29ee30 Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Wed, 1 Apr 2026 20:04:30 -0500 Subject: [PATCH 4/7] Update CI to rocm 7.1 Signed-off-by: Eric Eaton --- .github/workflows/amd-ci.yml | 20 ++++++++++---------- shmem/rocshmem_bind/build.sh | 7 ------- shmem/rocshmem_bind/build_rocshmem.sh | 10 +++++++--- 3 files changed, 17 insertions(+), 20 deletions(-) diff --git a/.github/workflows/amd-ci.yml b/.github/workflows/amd-ci.yml index 643172d7a..ff23c92de 100644 --- a/.github/workflows/amd-ci.yml +++ b/.github/workflows/amd-ci.yml @@ -12,7 +12,7 @@ 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 @@ -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,7 +58,7 @@ 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 @@ -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,7 +103,7 @@ 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 @@ -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 diff --git a/shmem/rocshmem_bind/build.sh b/shmem/rocshmem_bind/build.sh index f6497dab8..f90944a48 100755 --- a/shmem/rocshmem_bind/build.sh +++ b/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 diff --git a/shmem/rocshmem_bind/build_rocshmem.sh b/shmem/rocshmem_bind/build_rocshmem.sh index 0c614e1e3..91becc603 100755 --- a/shmem/rocshmem_bind/build_rocshmem.sh +++ b/shmem/rocshmem_bind/build_rocshmem.sh @@ -33,11 +33,12 @@ if [ -d "${ROCSHMEM_SRC_DIR}" ]; then fi fi +rocm_systems_tag=hip-version_7.12.60610 + if ! [ -d "${ROCSHMEM_SRC_DIR}" ]; then echo "Creating sparse checkout" - tag=hip-version_7.12.60610 pushd "${PROJECT_ROOT}/../.." - git clone "https://github.com/ROCm/rocm-systems.git" -b "${tag}" --depth 1 --sparse "${sys_path}" + 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 @@ -46,6 +47,10 @@ if ! [ -d "${ROCSHMEM_SRC_DIR}" ]; then popd fi +pushd "${sys_path}" +git checkout "${rocm_systems_tag}" +popd + pushd ${ROCSHMEM_SRC_DIR} ROCSHMEM_BUILD_DIR=${PROJECT_ROOT}/rocshmem_build @@ -55,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" From dcdea9eec3dd3256b223473efb2b734b9d18a965 Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Fri, 3 Apr 2026 12:49:53 -0500 Subject: [PATCH 5/7] Add sparse checkout to .gitignore --- .gitignore | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index ffb38f58a..a840320d1 100644 --- a/.gitignore +++ b/.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 \ No newline at end of file +*.csv From f945f012442060f13976efd8e9f1333327c4fa72 Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Tue, 7 Apr 2026 16:41:25 -0500 Subject: [PATCH 6/7] Set the shared memory size for docker Signed-off-by: Eric Eaton --- .github/workflows/amd-ci.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/amd-ci.yml b/.github/workflows/amd-ci.yml index ff23c92de..26e53ed18 100644 --- a/.github/workflows/amd-ci.yml +++ b/.github/workflows/amd-ci.yml @@ -15,7 +15,7 @@ jobs: 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" @@ -61,7 +61,7 @@ jobs: 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" @@ -106,7 +106,7 @@ jobs: 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" From fd41ad340d3abef6363f7da8a1595cc57e27f372 Mon Sep 17 00:00:00 2001 From: Eric Eaton Date: Thu, 9 Apr 2026 13:58:38 -0500 Subject: [PATCH 7/7] Swap barrier order Signed-off-by: Eric Eaton --- python/triton_dist/test/amd/test_rocshmem_api.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python/triton_dist/test/amd/test_rocshmem_api.py b/python/triton_dist/test/amd/test_rocshmem_api.py index 20c7c7d95..9c747dbe9 100644 --- a/python/triton_dist/test/amd/test_rocshmem_api.py +++ b/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)