[Feature/Performance] MetaX C500 OpenMPI backend staging experiments and follow-up items

[Dayuxiaoshui]

Problem Statement

MetaX C500 support is functional with the current OpenMPI backend, but the
backend moves accelerator buffers through host staging memory:

1. device buffer -> host buffer
2. MPI collective or point-to-point operation on host memory
3. host buffer -> device buffer

This path is portable and useful as a bootstrap implementation, but the current
performance is limited by host staging overhead. The first experiment showed
that repeated host buffer allocation was a major avoidable cost for larger
messages.

Environment

• Host: res-macloud-b01
• Accelerator: 4 x MetaX C500 (XCORE1000, reported by macainfo)
• MACA SDK: /opt/maca -> /opt/maca-3.7.0
• mxcc: mxcc version 1.0.0 (d9102a1572)
• OpenMPI: /opt/maca/ompi, mpirun (Open MPI) 4.1.6
• CMake: 3.31.10
• Backend/device configuration: WITH_METAX=ON, WITH_OMPI=ON

Manual build command used for reliable MetaX compilation:

export MACA_PATH=/opt/maca
export PATH="$MACA_PATH/ompi/bin:$PATH"
export LD_LIBRARY_PATH="$MACA_PATH/lib:$MACA_PATH/ompi/lib:$LD_LIBRARY_PATH"
export OMPI_CC="$PWD/scripts/devices/mxcc_wrapper.sh"
export OMPI_CXX="$PWD/scripts/devices/mxcc_wrapper.sh"

cmake -S . -B build-metax-fixed \
  -DCMAKE_BUILD_TYPE=Release \
  -DWITH_METAX=ON -DWITH_OMPI=ON \
  -DAUTO_DETECT_DEVICES=OFF -DAUTO_DETECT_BACKENDS=OFF \
  -DCMAKE_C_COMPILER="$MACA_PATH/ompi/bin/mpicc" \
  -DCMAKE_CXX_COMPILER="$MACA_PATH/ompi/bin/mpicxx"

cmake --build build-metax-fixed -j 8

Note: OpenMPI cannot initialize sockets inside the restricted sandbox used by
the agent, so runtime experiments were launched outside the sandbox.

Experiments

Correctness smoke tests

All internal examples passed on 2 ranks:

• all_reduce
• all_gather
• all_to_all
• broadcast
• reduce
• reduce_scatter
• send_recv

Example command:

mpirun --allow-run-as-root -np 2 ./build-metax-fixed/examples/all_reduce

Observed all_reduce result:

Correct: YES
Expect:  3.00
Actual:  3.00

Microbenchmark

A new benchmark was added at examples/metax_bench.cc to compare:

• copy_d2h: device-to-host copy
• copy_h2d: host-to-device copy
• copy_roundtrip: D2H + H2D
• host_allreduce: host-only MPI_Allreduce
• stage_reuse: manual D2H + MPI + H2D with reusable host buffers
• stage_alloc: manual D2H + MPI + H2D with per-call malloc/free
• infiniccl_allreduce: public InfiniCCL API path

Example command:

mpirun --allow-run-as-root -np 4 ./build-metax-fixed/examples/metax_bench \
  --elements 65536,262144,1048576,4194304 \
  --modes host_allreduce,stage_reuse,stage_alloc,infiniccl_allreduce \
  --warmup 5 --iters 20 --validate

Key Findings

Before host staging buffer reuse

infiniccl_allreduce was close to stage_alloc, which indicated that repeated
host buffer allocation was a major cost.

Case	stage_reuse	Original infiniccl_allreduce
2 ranks, 4 MB	3.36 ms	9.02 ms
2 ranks, 16 MB	12.78 ms	36.41 ms
4 ranks, 4 MB	4.89 ms	10.19 ms
4 ranks, 16 MB	17.57 ms	42.29 ms

After host staging buffer reuse

A reusable thread-local host staging buffer was added for the OpenMPI backend.
After the change, infiniccl_allreduce moved close to the stage_reuse
baseline.

Case	Optimized infiniccl_allreduce
2 ranks, 4 MB	3.47 ms
2 ranks, 16 MB	12.74 ms
4 ranks, 4 MB	4.45 ms
4 ranks, 16 MB	16.69 ms

This is roughly a 2x-3x improvement for large AllReduce messages compared with
the original implementation.

Changes Made During Investigation

Build and documentation

• Removed late mutation of CMAKE_C_COMPILER and CMAKE_CXX_COMPILER from the
  MetaX branch in top-level CMakeLists.txt.
• Updated scripts/build.sh to detect MetaX and use MACA OpenMPI wrappers with
  OMPI_CC/OMPI_CXX pointing to scripts/devices/mxcc_wrapper.sh.
• Made scripts/build.sh tolerate a read-only ~/.bashrc.
• Documented the MetaX + OpenMPI manual build path in README.md.
• Ignored build-* directories.

Runtime/backend

• Added src/ompi/staging_buffer.h.
• Replaced per-call host staging malloc/free in the OpenMPI backend with
  reusable thread-local staging buffers.
• Applied the staging buffer reuse to:
  • AllReduce
  • AllGather
  • AllToAll
  • Broadcast
  • Reduce
  • ReduceScatter
  • Send
  • Recv
• Added examples/metax_bench.cc for repeatable MetaX/OpenMPI experiments.

Remaining Issues

1. OpenMPI backend is still host-staging based

The current path still copies data through host memory. Buffer reuse removes a
large allocation overhead, but it does not remove the D2H/H2D copies.

Follow-up options:

• Investigate MACA/MCCL backend support for native accelerator collectives.
• Investigate whether MACA OpenMPI supports GPU-aware MPI for MetaX buffers.
• Add an optional direct device-buffer MPI path if supported and safe.

2. Staging buffers are pageable host memory

The reusable buffers currently use std::malloc. Copy bandwidth may improve if
MACA provides pinned/page-locked host allocation APIs.

Follow-up options:

• Check MACA runtime support for pinned host memory.
• Add a staging allocator abstraction with normal malloc fallback.
• Benchmark pageable vs pinned host staging.

3. ReduceScatter still needs separate analysis

Correctness passes, but reduce_scatter did not show the same clear timing
improvement as AllReduce in the example-level test. It should be investigated
separately with a benchmark mode that matches the exact ReduceScatter data
layout and counts.

Follow-up options:

• Extend metax_bench with reduce_scatter modes.
• Compare MPI_Reduce_scatter_block against alternative decompositions.
• Check if count, buffer size, or MPI algorithm selection is dominating.

4. MPI errors abort the process

INFINI_CHECK_MPI currently aborts on MPI errors. This is harsh for a library
API that otherwise returns ReturnStatus.

Follow-up options:

• Convert MPI error handling to return ReturnStatus.
• Include MPI error strings in logs.
• Keep abort behavior only under an explicit debug/fail-fast option.

5. Device binding is OpenMPI-specific

MetaX device selection currently relies on OMPI_COMM_WORLD_LOCAL_RANK.

Follow-up options:

• Add fallbacks for common MPI/PMI local-rank environment variables.
• Query MetaX device count and validate selected device IDs.
• Consider modulo mapping or explicit devlist behavior for oversubscription.

6. Reusable staging buffer retention needs policy

The new thread-local staging buffers keep the largest capacity seen by each
thread. This is good for repeated collectives but can retain large host memory
after a single large operation.

Follow-up options:

• Add an upper bound or shrink policy.
• Add an environment variable for staging buffer cache size.
• Add communicator-level or backend-level cleanup hooks if needed.

Proposed Next Steps

1. Keep the reusable host staging buffer optimization; it gives a clear
   improvement with low API risk.
2. Extend metax_bench to cover all_gather, all_to_all, broadcast,
   reduce, reduce_scatter, send, and recv.
3. Add pinned host staging support if MACA exposes a suitable API.
4. Investigate MCCL or GPU-aware MPI support as the higher-ceiling path.
5. Harden error handling and device binding before larger multi-node testing.

[Dayuxiaoshui]

cc @Ziminli @GordonYang1

[Ziminli]

Thank you for taking the time to experiment with the project and for sharing your detailed feedback and suggestions!

We greatly appreciate the insights you've provided and have carefully noted your suggestions. Many of the areas you mentioned align with our existing roadmap, and we have already planned a number of improvements and enhancements for future releases.

At the current stage, the project's primary focus is on expanding platform and backend support, as well as building out core functionalities. In the upcoming v0.2.0 release, InfiniCCL will introduce native CCL + MPI support across multiple platform combinations, including MetaX. This will significantly improve both performance and usability in real-world deployment scenarios. Beyond that, more extensive performance tuning and optimization efforts are planned starting from v0.3.0 and will continue to be rolled out incrementally as the project matures.

Thanks again for helping us improve the project. Feedback from real-world usage is invaluable to us, and we welcome any additional observations or suggestions as you continue exploring InfiniCCL.