feat(cco): merge dev/cco — Collective Communication Object (CCO) module

.github/workflows/ci_cco.yml

@@ -0,0 +1,72 @@
+name: CCO CI test
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  IMAGE: rocm/mori:ci-cco
+  BASE_IMAGE: rocm/pytorch:rocm7.2.4_ubuntu24.04_py3.12_pytorch_release_2.8.0
+  CONTAINER: mori_cco_ci_${{ github.run_id }}
+  CT: docker
+
+jobs:
+  cco-unit-test:
+    name: CCO unit test (${{ matrix.platform }})
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - platform: MI325X_BNXT
+            runner: [self-hosted, 325_bnxt_108]
+            rdma_devices: bnxt_re0,bnxt_re1,bnxt_re2,bnxt_re3,bnxt_re4,bnxt_re5,bnxt_re7,bnxt_re8
+            rdma_sl: 3
+            rdma_tc: 104
+    env:
+      MORI_RDMA_DEVICES: ${{ matrix.rdma_devices }}
+      MORI_RDMA_SL: ${{ matrix.rdma_sl }}
+      MORI_RDMA_TC: ${{ matrix.rdma_tc }}
+    timeout-minutes: 30
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          submodules: true
+
+      - name: Build CI image
+        run: $CT build --network=host --build-arg BASE_IMAGE=$BASE_IMAGE -t $IMAGE -f docker/Dockerfile.cco .
+
+      - name: Start container
+        run: |
+          $CT rm -f $CONTAINER 2>/dev/null || true
+          CONTAINER_RUNTIME=$CT ./docker/ci_run.sh --name $CONTAINER \
+            -e MORI_RDMA_DEVICES=$MORI_RDMA_DEVICES \
+            -e MORI_RDMA_SL=$MORI_RDMA_SL \
+            -e MORI_RDMA_TC=$MORI_RDMA_TC \
+            -v $GITHUB_WORKSPACE:$GITHUB_WORKSPACE \
+            -w $GITHUB_WORKSPACE \
+            $IMAGE sleep infinity
+          $CT exec $CONTAINER \
+            git config --global --add safe.directory $GITHUB_WORKSPACE
+
+      - name: Build mori with tests
+        run: |
+          $CT exec $CONTAINER bash -c \
+            "cd $GITHUB_WORKSPACE && BUILD_TESTS=ON MORI_WITH_MPI=ON pip install ."
+
+      - name: Run CCO unit tests (fork mode, 4 ranks)
+        run: |
+          $CT exec $CONTAINER bash $GITHUB_WORKSPACE/tools/run_cco_tests.sh $GITHUB_WORKSPACE/build 4 2>&1
+
+      - name: Cleanup
+        if: always()
+        run: |
+          $CT exec $CONTAINER chown -R $(id -u):$(id -g) $GITHUB_WORKSPACE 2>/dev/null || true
+          $CT rm -f $CONTAINER || true

.gitignore

@@ -13,8 +13,15 @@ __pycache__
 python/mori/_version.py
 python/mori/_jit-sources/
 python/mori/*.so
+python/mori/cco/*.so
+python/mori/cco/cco.cpp          # Cython-generated from cco.pyx (build artifact)
 python/mori/umbp_master
 python/mori/spdk_proxy
+python/mori/examples/          # BUILD_EXAMPLES=ON copies example binaries here
+python/mori/benchmarks/        # BUILD_BENCHMARK=ON copies benchmark binaries here
+
+# local cco experiments (not for commit)
+examples/cco/**/_dump_*.py
 
 .vscode
 *.log

.pre-commit-config.yaml

@@ -1,7 +1,7 @@
 # Exclude all 3rd-party libraries
 exclude: |
   (?x)^(
-    third_party/.+
+    3rdparty/.+
   )$
 repos:
   # Common hooks
@@ -67,8 +67,10 @@ repos:
     rev: v0.6.13
     hooks:
       - id: cmake-format
+        exclude: '.*CMakeLists\.txt$'
   - repo: https://github.com/PFCCLab/cmake-lint-paddle
     rev: v1.5.2
     hooks:
       - id: cmakelint
         args: [--config=./tools/codestyle/.cmakelintrc]
+        exclude: '.*CMakeLists\.txt$'

CMakeLists.txt

@@ -23,11 +23,21 @@ option(USE_ROCM "Whether to use rocm" ON)
 option(USE_BNXT "Whether to use BNXT NIC" OFF)
 option(USE_IONIC "Whether to use IONIC" OFF)
 option(BUILD_EXAMPLES "Whether to build examples" ON)
+option(BUILD_BENCHMARK "Whether to build benchmark targets" OFF)
+# MPI is required by mpi_bootstrap, the C++ examples, and every benchmark
+# target (all use the MPI multi-rank launcher), so default it on whenever
+# either is enabled.
+if(BUILD_EXAMPLES OR BUILD_BENCHMARK)
+  set(_with_mpi_default ON)
+else()
+  set(_with_mpi_default OFF)
+endif()
 option(WITH_MPI
-       "Enable MPI support (required for mpi_bootstrap and C++ examples)"
-       ${BUILD_EXAMPLES})
+       "Enable MPI support (required for mpi_bootstrap, C++ examples, benchmarks)"
+       ${_with_mpi_default})
 option(BUILD_APPLICATION "Whether to build application library" ON)
 option(BUILD_SHMEM "Whether to build shmem library" ON)
+option(BUILD_CCO "Whether to build cco library" ON)
 option(BUILD_OPS "Whether to build mori operation kernels" ON)
 option(BUILD_IO "Whether to build mori io library" ON)
 option(BUILD_COLLECTIVE "Whether to build mori collective library" ON)
@@ -36,7 +46,6 @@ option(BUILD_XLA_FFI_OPS "Whether to build mori xla ffi python bindings" OFF)
 option(BUILD_UMBP "Whether to build UMBP (Unified Memory/Bandwidth Pool)" OFF)
 option(BUILD_METRICS "Whether to build the Prometheus metrics server module" ON)
 option(BUILD_TESTS "Whether to build mori CPP tests" ON)
-option(BUILD_BENCHMARK "Whether to build benchmark targets" OFF)
 option(ENABLE_PROFILER "Enable kernel profiling" OFF)
 option(ENABLE_DEBUG_PRINTF "Enable debug printf in device kernels" OFF)
 option(ENABLE_STANDARD_MOE_ADAPT "Enable standard moe adapt" OFF)
@@ -243,6 +252,10 @@ if(BUILD_SHMEM)
   add_subdirectory(src/shmem)
 endif()
 
+if(BUILD_CCO)
+  add_subdirectory(src/cco)
+endif()
+
 if(BUILD_OPS)
   add_subdirectory(src/ops)
 endif()
@@ -282,6 +295,9 @@ endif()
 if(BUILD_SHMEM)
   install(TARGETS mori_shmem LIBRARY DESTINATION lib)
 endif()
+if(BUILD_CCO)
+  install(TARGETS mori_cco LIBRARY DESTINATION lib)
+endif()
 if(BUILD_OPS)
   install(TARGETS mori_ops LIBRARY DESTINATION lib)
 endif()

benchmark/CMakeLists.txt

@@ -44,6 +44,38 @@ function(add_mori_benchmark_app name)
   endif()
 endfunction()
 
+# --- CCO benchmark: HIP kernel TU + a shared host-CXX control-plane lib ---
+# CCO makes `ccoComm` host-only (opaque under __HIPCC__). util.cpp touches its
+# members, so it must compile as plain CXX, never as HIP. Putting it in its own
+# CXX static library guarantees that regardless of the per-target HIP language of
+# the kernel TUs. Each benchmark = one HIP kernel .cpp + this lib.
+function(add_mori_benchmark_cco name)
+  cmake_parse_arguments(ARG "" "" "SOURCES;LIBS;INCLUDES" ${ARGN})
+  if(NOT ARG_SOURCES)
+    message(FATAL_ERROR "add_mori_benchmark_cco(${name}): SOURCES required")
+  endif()
+  add_executable(${name} ${ARG_SOURCES})
+  set_source_files_properties(${ARG_SOURCES} PROPERTIES LANGUAGE HIP)
+  # libmori_cco.so is self-contained (absorbs the application layer).
+  target_link_libraries(${name} PRIVATE cco_bench_util mori_cco hip::host
+                                        hip::device ${ARG_LIBS})
+  target_include_directories(${name} PRIVATE ${CMAKE_SOURCE_DIR}/include
+                                             ${CMAKE_CURRENT_SOURCE_DIR}/cco)
+  # $ORIGIN-relative rpath (kept alongside the build-tree rpath) so the binaries
+  # resolve libmori_*.so after `BUILD_BENCHMARK=ON pip install .` copies them to
+  # site-packages/mori/benchmarks/cco/ (the libs live at site-packages/mori/).
+  set_target_properties(${name} PROPERTIES BUILD_RPATH "$ORIGIN/../..")
+  if(DEFINED GPU_TARGETS)
+    set_target_properties(${name} PROPERTIES HIP_ARCHITECTURES "${GPU_TARGETS}")
+  endif()
+  if(MORI_DEVICE_NIC_DEFINE)
+    target_compile_definitions(${name} PRIVATE ${MORI_DEVICE_NIC_DEFINE})
+  endif()
+  if(ARG_INCLUDES)
+    target_include_directories(${name} PRIVATE ${ARG_INCLUDES})
+  endif()
+endfunction()
+
 # ---------------------------------------------------------------------------
 # Registered benchmark targets
 # ---------------------------------------------------------------------------
@@ -62,3 +94,36 @@ if(WITH_MPI)
   add_mori_benchmark_shmem(p2p_get_latency SOURCES shmem/p2p_get_latency.cpp
                            shmem/util.cpp LIBS stdc++fs)
 endif()
+
+# --- CCO p2p benchmarks (LSA + IGBDA transports, MPI 2-rank launcher) ---
+if(WITH_MPI AND BUILD_CCO)
+  # MPI is a direct dependency of the cco benchmarks; previously it came
+  # transitively via mori_application, which the benchmarks no longer link.
+  find_package(MPI REQUIRED)
+  # Host control-plane (MPI + ccoComm lifecycle) as a plain CXX static lib, so it
+  # never gets the kernel TUs' HIP language (which would make ccoComm opaque).
+  add_library(cco_bench_util STATIC cco/util.cpp)
+  set_source_files_properties(cco/util.cpp PROPERTIES LANGUAGE CXX)
+  target_include_directories(cco_bench_util PRIVATE ${CMAKE_SOURCE_DIR}/include
+                                                    ${CMAKE_CURRENT_SOURCE_DIR}/cco
+                                                    ${MPI_CXX_INCLUDE_DIRS})
+  target_link_libraries(cco_bench_util PRIVATE mori_cco hip::host
+                                               ${MPI_CXX_LIBRARIES})
+  set_target_properties(cco_bench_util PROPERTIES POSITION_INDEPENDENT_CODE ON)
+
+  add_mori_benchmark_cco(cco_p2p_put_bw SOURCES cco/p2p_put_bw.cpp
+                         LIBS stdc++fs ${MPI_CXX_LIBRARIES}
+                         INCLUDES ${MPI_CXX_INCLUDE_DIRS})
+  add_mori_benchmark_cco(cco_p2p_put_latency SOURCES cco/p2p_put_latency.cpp
+                         LIBS stdc++fs ${MPI_CXX_LIBRARIES}
+                         INCLUDES ${MPI_CXX_INCLUDE_DIRS})
+  add_mori_benchmark_cco(cco_p2p_get_bw SOURCES cco/p2p_get_bw.cpp
+                         LIBS stdc++fs ${MPI_CXX_LIBRARIES}
+                         INCLUDES ${MPI_CXX_INCLUDE_DIRS})
+  add_mori_benchmark_cco(cco_p2p_get_latency SOURCES cco/p2p_get_latency.cpp
+                         LIBS stdc++fs ${MPI_CXX_LIBRARIES}
+                         INCLUDES ${MPI_CXX_INCLUDE_DIRS})
+elseif(WITH_MPI AND NOT BUILD_CCO)
+  message(WARNING
+    "BUILD_BENCHMARK=ON with BUILD_CCO=OFF: CCO p2p benchmark targets skipped")
+endif()

benchmark/cco/device_utils.hpp

@@ -0,0 +1,71 @@
+// Copyright © Advanced Micro Devices, Inc. All rights reserved.
+//
+// MIT License
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+// Device-side helpers shared across the CCO p2p benchmark kernels. Include only
+// from HIP translation units.
+
+#pragma once
+
+#include "mori/cco/cco_scale_out.hpp"
+#include "util.hpp"
+
+namespace mori::cco::benchmark {
+
+// Intra-block linear thread index.
+__device__ inline int linear_tid() {
+  return threadIdx.x * blockDim.y * blockDim.z + threadIdx.y * blockDim.z + threadIdx.z;
+}
+
+// Strided element copy across [lane, nlanes): dst[i] = src[i] for i in [0, n).
+template <typename T>
+__device__ inline void lsa_copy_strided(T* __restrict__ dst, const T* __restrict__ src, size_t n,
+                                        int lane, int nlanes) {
+  for (size_t i = lane; i < n; i += static_cast<size_t>(nlanes)) {
+    dst[i] = src[i];
+  }
+}
+
+// GPU-internal all-block barrier (single GPU, not cross-rank), matching shmem's
+// bw_cross_block_barrier_round so the LSA bw timed window includes the same
+// per-round sync. counter_d[0]=arrivals, counter_d[1]=phase; call from all
+// threads of all blocks with the same (counter_d, nblocks, i).
+__device__ inline void bw_cross_block_barrier_round(volatile unsigned int* counter_d, int nblocks,
+                                                    int i) {
+  __syncthreads();
+  if (linear_tid() == 0) {
+    __threadfence();
+    unsigned int c = atomicInc((unsigned int*)counter_d, 0xffffffffu);
+    if (c == static_cast<unsigned int>(nblocks * (i + 1) - 1)) {
+      counter_d[1] += 1u;
+    }
+    while (counter_d[1] != static_cast<unsigned int>(i + 1)) {
+    }
+  }
+  __syncthreads();
+}
+
+}  // namespace mori::cco::benchmark
+
+// CCO_GDA_DISPATCH is provided by mori/cco/cco_scale_out.hpp: GDA provider is
+// fixed at build time (per-NIC, from MORI_DEVICE_NIC_*), so the macro just binds
+// `constexpr auto P = CCO_GDA_BUILD_PROVIDER` and runs the statement — no
+// runtime provider argument.