From dfe5c8f72abf38c549e1a97df03f2f0394ae5d00 Mon Sep 17 00:00:00 2001 From: 25337 Date: Wed, 1 Jul 2026 11:29:31 +0000 Subject: [PATCH 01/10] perf(io): add busy-wait completion mode and numpy batch-transfer marshaling WaitBusy() spins on the completion flag instead of blocking on the condition variable, trading a CPU core for lower wakeup latency. BatchRead/BatchWrite now accept numpy arrays for offsets/sizes, avoiding per-element Python-list conversion. Benchmark gains --busy-wait, --post-batch-size, --warmup-iters flags and separate launch/transfer latency reporting. Co-Authored-By: Claude Sonnet 5 --- include/mori/io/common.hpp | 40 +++++++++ src/pybind/pybind_io.cpp | 73 ++++++++++++++-- tests/python/io/benchmark.py | 165 ++++++++++++++++++++++++++++------- 3 files changed, 240 insertions(+), 38 deletions(-) diff --git a/include/mori/io/common.hpp b/include/mori/io/common.hpp index 3ecfd3804..46eeab506 100644 --- a/include/mori/io/common.hpp +++ b/include/mori/io/common.hpp @@ -188,6 +188,34 @@ struct TransferStatus { void Wait() { (void)WaitFor(-1); } + // Busy-poll (spin) on the completion flag until the transfer finishes, + // instead of blocking on the condition variable like WaitFor(). This trades a + // CPU core for latency: it avoids the cross-thread cv.notify -> futex -> + // scheduler wakeup (~5-10us) that WaitFor() pays when a dedicated poller + // thread (NotifManager::MainLoop) delivers the completion. The waiter spins on + // the same atomic status flag that the poller thread stores via Update(), so + // it observes SUCCESS with only cross-thread atomic-visibility latency (~1us). + // + // timeoutMs < 0 spins indefinitely, == 0 polls once, > 0 spins up to the + // deadline. The returned code may still be IN_PROGRESS when a bounded wait + // times out. + StatusCode WaitBusy(int timeoutMs = -1) { + StatusCode current = code.load(std::memory_order_acquire); + if (current != StatusCode::IN_PROGRESS) return current; + + const bool bounded = timeoutMs > 0; + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::milliseconds(bounded ? timeoutMs : 0); + for (;;) { + PollProgress(); + current = code.load(std::memory_order_acquire); + if (current != StatusCode::IN_PROGRESS) return current; + if (timeoutMs == 0) return current; + if (bounded && std::chrono::steady_clock::now() >= deadline) return current; + CpuRelax(); + } + } + // timeoutMs < 0 waits indefinitely, timeoutMs == 0 polls once, timeoutMs > 0 // waits up to the requested deadline. The returned code may still be // IN_PROGRESS when a bounded wait times out. @@ -228,6 +256,18 @@ struct TransferStatus { void SetProgressCallback(std::function cb) { progressCallback = std::move(cb); } private: + // Spin-loop pause hint: relaxes the core (lowers power / frees SMT resources) + // without yielding the OS thread, keeping the waiter hot for low-latency spin. + static inline void CpuRelax() { +#if defined(__x86_64__) || defined(__i386__) + __builtin_ia32_pause(); +#elif defined(__aarch64__) + asm volatile("yield" ::: "memory"); +#else + std::atomic_thread_fence(std::memory_order_acquire); +#endif + } + StatusCode CodeWithProgress() { PollProgress(); return code.load(std::memory_order_acquire); diff --git a/src/pybind/pybind_io.cpp b/src/pybind/pybind_io.cpp index fc452ef95..5ae5a7539 100644 --- a/src/pybind/pybind_io.cpp +++ b/src/pybind/pybind_io.cpp @@ -19,6 +19,7 @@ // 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. +#include #include #include #include @@ -36,6 +37,28 @@ void* FabricMalloc(size_t size, int device); void FabricFree(void* ptr); } // namespace io +namespace { + +// Accepts either a numpy array or any Python sequence (list, tuple, ...). When the caller +// already passes a contiguous numpy array of the matching dtype, this is a single memcpy +// instead of pybind11/stl.h's generic per-element Python object iteration/conversion. +using SizeArray = py::array_t; + +mori::io::SizeVec SizeArrayToVec(const SizeArray& arr) { + py::buffer_info info = arr.request(); + const size_t* ptr = static_cast(info.ptr); + return mori::io::SizeVec(ptr, ptr + info.size); +} + +mori::io::BatchSizeVec SizeArraysToBatchVec(const std::vector& arrs) { + mori::io::BatchSizeVec out; + out.reserve(arrs.size()); + for (const auto& a : arrs) out.push_back(SizeArrayToVec(a)); + return out; +} + +} // namespace + void RegisterMoriIo(pybind11::module_& m) { m.def("set_log_level", &mori::io::SetLogLevel); @@ -133,6 +156,8 @@ void RegisterMoriIo(pybind11::module_& m) { .def("SetMessage", &mori::io::TransferStatus::SetMessage) .def("Wait", &mori::io::TransferStatus::Wait, py::call_guard()) .def("WaitFor", &mori::io::TransferStatus::WaitFor, py::arg("timeout_ms") = -1, + py::call_guard()) + .def("WaitBusy", &mori::io::TransferStatus::WaitBusy, py::arg("timeout_ms") = -1, py::call_guard()); py::class_(m, "EngineDesc") @@ -191,11 +216,27 @@ void RegisterMoriIo(pybind11::module_& m) { .def("AllocateTransferUniqueId", &mori::io::IOEngineSession::AllocateTransferUniqueId, py::call_guard()) .def("Read", &mori::io::IOEngineSession::Read, py::call_guard()) - .def("BatchRead", &mori::io::IOEngineSession::BatchRead, - py::call_guard()) + .def("BatchRead", + [](mori::io::IOEngineSession& self, const SizeArray& localOffsets, + const SizeArray& remoteOffsets, const SizeArray& sizes, + mori::io::TransferStatus* status, mori::io::TransferUniqueId id) { + mori::io::SizeVec lo = SizeArrayToVec(localOffsets); + mori::io::SizeVec ro = SizeArrayToVec(remoteOffsets); + mori::io::SizeVec sz = SizeArrayToVec(sizes); + py::gil_scoped_release release; + self.BatchRead(lo, ro, sz, status, id); + }) .def("Write", &mori::io::IOEngineSession::Write, py::call_guard()) - .def("BatchWrite", &mori::io::IOEngineSession::BatchWrite, - py::call_guard()) + .def("BatchWrite", + [](mori::io::IOEngineSession& self, const SizeArray& localOffsets, + const SizeArray& remoteOffsets, const SizeArray& sizes, + mori::io::TransferStatus* status, mori::io::TransferUniqueId id) { + mori::io::SizeVec lo = SizeArrayToVec(localOffsets); + mori::io::SizeVec ro = SizeArrayToVec(remoteOffsets); + mori::io::SizeVec sz = SizeArrayToVec(sizes); + py::gil_scoped_release release; + self.BatchWrite(lo, ro, sz, status, id); + }) .def("Alive", &mori::io::IOEngineSession::Alive); py::class_(m, "IOEngine") @@ -210,9 +251,29 @@ void RegisterMoriIo(pybind11::module_& m) { .def("AllocateTransferUniqueId", &mori::io::IOEngine::AllocateTransferUniqueId, py::call_guard()) .def("Read", &mori::io::IOEngine::Read, py::call_guard()) - .def("BatchRead", &mori::io::IOEngine::BatchRead, py::call_guard()) + .def("BatchRead", + [](mori::io::IOEngine& self, const mori::io::MemDescVec& localDest, + const std::vector& localOffsets, const mori::io::MemDescVec& remoteSrc, + const std::vector& remoteOffsets, const std::vector& sizes, + mori::io::TransferStatusPtrVec& status, mori::io::TransferUniqueIdVec& ids) { + mori::io::BatchSizeVec lo = SizeArraysToBatchVec(localOffsets); + mori::io::BatchSizeVec ro = SizeArraysToBatchVec(remoteOffsets); + mori::io::BatchSizeVec sz = SizeArraysToBatchVec(sizes); + py::gil_scoped_release release; + self.BatchRead(localDest, lo, remoteSrc, ro, sz, status, ids); + }) .def("Write", &mori::io::IOEngine::Write, py::call_guard()) - .def("BatchWrite", &mori::io::IOEngine::BatchWrite, py::call_guard()) + .def("BatchWrite", + [](mori::io::IOEngine& self, const mori::io::MemDescVec& localSrc, + const std::vector& localOffsets, const mori::io::MemDescVec& remoteDest, + const std::vector& remoteOffsets, const std::vector& sizes, + mori::io::TransferStatusPtrVec& status, mori::io::TransferUniqueIdVec& ids) { + mori::io::BatchSizeVec lo = SizeArraysToBatchVec(localOffsets); + mori::io::BatchSizeVec ro = SizeArraysToBatchVec(remoteOffsets); + mori::io::BatchSizeVec sz = SizeArraysToBatchVec(sizes); + py::gil_scoped_release release; + self.BatchWrite(localSrc, lo, remoteDest, ro, sz, status, ids); + }) .def("CreateSession", &mori::io::IOEngine::CreateSession) .def("PopInboundTransferStatus", &mori::io::IOEngine::PopInboundTransferStatus, py::call_guard()) diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 8f8bc6503..1ea858ac5 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -20,6 +20,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. from tests.python.utils import TorchDistContext +import numpy as np import torch import torch.distributed as dist from mori.io import ( @@ -229,12 +230,25 @@ def parse_args(): default=4, help="Number of QPs for a single transfer (default: 4)", ) + parser.add_argument( + "--post-batch-size", + type=int, + default=-1, + help="Number of WRs posted per ibv_post_send batch; -1 = auto (default: -1)", + ) parser.add_argument( "--num-worker-threads", type=int, default=1, help="Number of threads used for transfer", ) + parser.add_argument( + "--busy-wait", + action="store_true", + help="Busy-poll (spin) on the transfer status instead of blocking on the " + "condition variable. Avoids the cross-thread cv-wakeup latency (~5-10us) " + "per completion at the cost of burning a core while waiting.", + ) parser.add_argument( "--disable-chunking", action="store_true", @@ -281,6 +295,12 @@ def parse_args(): default=128, help="Number of iterations running test", ) + parser.add_argument( + "--warmup-iters", + type=int, + default=10, + help="Number of untimed warmup iterations before each measured run (default: 10)", + ) parser.add_argument( "--poll_cq_mode", type=str, @@ -333,6 +353,7 @@ def __init__( batch_contiguous: bool = False, enable_sess: bool = False, iters: int = 128, + warmup_iters: int = 10, sweep: bool = False, sweep_batch: bool = False, sweep_start_size: int = 8, @@ -347,7 +368,9 @@ def __init__( num_target_dev: int = 1, target_dev_offset: int = 0, num_qp_per_transfer: int = 4, + post_batch_size: int = -1, num_worker_threads: int = 1, + busy_wait: bool = False, poll_cq_mode: str = "polling", max_send_wr: int = 0, max_cqe_num: int = 0, @@ -371,6 +394,7 @@ def __init__( self.batch_contiguous = batch_contiguous self.enable_sess = enable_sess self.iters = iters + self.warmup_iters = warmup_iters self.sweep = sweep self.sweep_batch = sweep_batch self.sweep_start_size = sweep_start_size @@ -386,7 +410,9 @@ def __init__( self.num_target_dev = num_target_dev self.target_dev_offset = target_dev_offset self.num_qp_per_transfer = num_qp_per_transfer + self.post_batch_size = post_batch_size self.num_worker_threads = num_worker_threads + self.busy_wait = busy_wait self.poll_cq_mode = ( PollCqMode.POLLING if poll_cq_mode == "polling" else PollCqMode.EVENT ) @@ -564,7 +590,9 @@ def print_config(self): print(f" target_dev_offset: {self.target_dev_offset}") print(f" mem_type: {self.mem_type}") print(f" num_qp_per_transfer: {self.num_qp_per_transfer}") + print(f" post_batch_size: {self.post_batch_size}") print(f" num_worker_threads: {self.num_worker_threads}") + print(f" busy_wait: {self.busy_wait}") print(f" enable_chunking: {self.enable_chunking}") if self.enable_chunking: print(f" chunk_bytes: {self.chunk_bytes}") @@ -581,13 +609,17 @@ def print_config(self): print(f" batch_contiguous: {self.batch_contiguous}") print(f" enable_sess: {self.enable_sess}") print(f" iters: {self.iters}") + print(f" warmup_iters: {self.warmup_iters}") print() def _get_transfer_offsets(self, buffer_size, transfer_batch_size, batched): + # np.uint64 matches the C++ `size_t` used by the batch APIs, so passing + # these arrays straight into mori's batch_read/batch_write is a zero-copy + # transfer instead of the per-element Python-list conversion. if batched and not self.batch_contiguous: stride = buffer_size + 1 - return [i * stride for i in range(transfer_batch_size)] - return [i * buffer_size for i in range(transfer_batch_size)] + return np.arange(transfer_batch_size, dtype=np.uint64) * stride + return np.arange(transfer_batch_size, dtype=np.uint64) * buffer_size def _pack_tensor_segments(self, tensor, buffer_size, transfer_batch_size, batched): offsets = self._get_transfer_offsets( @@ -781,7 +813,7 @@ def _initialize_rdma(self): self.engine = IOEngine(key=f"{self.role.name}-{self.role_rank}", config=config) config = RdmaBackendConfig( qp_per_transfer=self.num_qp_per_transfer, - post_batch_size=-1, + post_batch_size=self.post_batch_size, num_worker_threads=self.num_worker_threads, poll_cq_mode=self.poll_cq_mode, enable_notification=False, @@ -926,6 +958,13 @@ def _initialize_fabric(self): "create_session returned None for fabric: peers likely not in the " "same vPOD, or remote memory is not fabric-exportable" ) + def _wait_status(self, status): + # Busy-poll spins on the completion flag (no cross-thread cv wakeup); + # otherwise block on the condition variable. + if self.busy_wait: + status.WaitBusy() + else: + status.Wait() def run_single_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size @@ -933,7 +972,7 @@ def run_single_once(self, buffer_size, transfer_batch_size): self.backend_type in ("rdma", "fabric") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: - return 0 + return 0, 0 status_list = [] transfer_uids = [] @@ -971,13 +1010,17 @@ def run_single_once(self, buffer_size, transfer_batch_size): for i in range(transfer_batch_size): status = func(*arg_list[i]) status_list.append(status) + launched = time.time() for status in status_list: - status.Wait() - duration = time.time() - st + self._wait_status(status) + done = time.time() + + launch_duration = launched - st + transfer_duration = done - launched for status in status_list: assert status.Succeeded(), f"Transfer failed: {status.Message()}" - return duration + return launch_duration, transfer_duration def run_batch_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size @@ -985,13 +1028,13 @@ def run_batch_once(self, buffer_size, transfer_batch_size): self.backend_type in ("rdma", "fabric") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: - return 0 + return 0, 0 # Strided offsets prevent merging: each transfer becomes a separate WR (to stress SQ / reproduce notify ENOMEM) offsets = self._get_transfer_offsets( buffer_size, transfer_batch_size, batched=True ) - sizes = [buffer_size for _ in range(transfer_batch_size)] + sizes = np.full(transfer_batch_size, buffer_size, dtype=np.uint64) transfer_uid = self.engine.allocate_transfer_uid() if self.enable_sess: @@ -1025,12 +1068,17 @@ def run_batch_once(self, buffer_size, transfer_batch_size): st = time.time() transfer_status = func(*args)[0] - transfer_status.Wait() - duration = time.time() - st + launched = time.time() + self._wait_status(transfer_status) + done = time.time() + + launch_duration = launched - st + transfer_duration = done - launched + assert ( transfer_status.Succeeded() ), f"Batch transfer failed: {transfer_status.Message()}" - return duration + return launch_duration, transfer_duration def run_once(self, buffer_size, transfer_batch_size): if self.enable_batch_transfer: @@ -1039,26 +1087,48 @@ def run_once(self, buffer_size, transfer_batch_size): return self.run_single_once(buffer_size, transfer_batch_size) def _run_and_compute(self, buffer_size, transfer_batch_size, iters): - latency = [] + for _ in range(self.warmup_iters): + self.run_once(buffer_size, transfer_batch_size) + + launch_latency = [] + transfer_latency = [] for _ in range(iters): - duration = self.run_once(buffer_size, transfer_batch_size) - latency.append(duration) + launch_duration, transfer_duration = self.run_once( + buffer_size, transfer_batch_size + ) + launch_latency.append(launch_duration) + transfer_latency.append(transfer_duration) if self.role is EngineRole.TARGET and ( self.backend_type in ("rdma", "fabric") or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ): - return 0, 0, 0, 0, 0 + return 0, 0, 0, 0, 0, 0, 0 + total_latency = [ + launch + transfer + for launch, transfer in zip(launch_latency, transfer_latency) + ] total_mem_mb = buffer_size * transfer_batch_size / (10**6) - avg_duration = sum(latency) / len(latency) - min_duration = min(latency) + avg_duration = sum(total_latency) / len(total_latency) + min_duration = min(total_latency) avg_duration_us = avg_duration * (10**6) min_duration_us = min_duration * (10**6) avg_bw = total_mem_mb / (10**3) / avg_duration max_bw = total_mem_mb / (10**3) / min_duration - return total_mem_mb, avg_duration_us, min_duration_us, avg_bw, max_bw + avg_launch_us = sum(launch_latency) / len(launch_latency) * (10**6) + avg_transfer_us = sum(transfer_latency) / len(transfer_latency) * (10**6) + + return ( + total_mem_mb, + avg_duration_us, + min_duration_us, + avg_bw, + max_bw, + avg_launch_us, + avg_transfer_us, + ) def _get_table_title(self): if self.backend_type == "xgmi": @@ -1083,6 +1153,8 @@ def _run_benchmark_loop(self): "Avg BW (GB/s)", "Min Lat (us)", "Avg Lat (us)", + "Avg Launch (us)", + "Avg Transfer (us)", ], title=self._get_table_title(), ) @@ -1095,11 +1167,15 @@ def _run_benchmark_loop(self): self.backend_type == "xgmi" and self.xgmi_multiprocess ): dist.barrier() - total_mem_mb, avg_duration, min_duration, avg_bw, max_bw = ( - self._run_and_compute( - cur_size, self.transfer_batch_size, self.iters - ) - ) + ( + total_mem_mb, + avg_duration, + min_duration, + avg_bw, + max_bw, + avg_launch, + avg_transfer, + ) = self._run_and_compute(cur_size, self.transfer_batch_size, self.iters) table.add_row( [ cur_size, @@ -1109,6 +1185,8 @@ def _run_benchmark_loop(self): f"{avg_bw:.2f}", f"{min_duration:.2f}", f"{avg_duration:.2f}", + f"{avg_launch:.2f}", + f"{avg_transfer:.2f}", ] ) if self.sweep_step > 0: @@ -1123,10 +1201,16 @@ def _run_benchmark_loop(self): self.backend_type == "xgmi" and self.xgmi_multiprocess ): dist.barrier() - total_mem_mb, avg_duration, min_duration, avg_bw, max_bw = ( - self._run_and_compute( - self.buffer_size, cur_transfer_batch_size, self.iters - ) + ( + total_mem_mb, + avg_duration, + min_duration, + avg_bw, + max_bw, + avg_launch, + avg_transfer, + ) = self._run_and_compute( + self.buffer_size, cur_transfer_batch_size, self.iters ) table.add_row( [ @@ -1137,14 +1221,22 @@ def _run_benchmark_loop(self): f"{avg_bw:.2f}", f"{min_duration:.2f}", f"{avg_duration:.2f}", + f"{avg_launch:.2f}", + f"{avg_transfer:.2f}", ] ) cur_transfer_batch_size *= 2 else: - total_mem_mb, avg_duration, min_duration, avg_bw, max_bw = ( - self._run_and_compute( - self.buffer_size, self.transfer_batch_size, self.iters - ) + ( + total_mem_mb, + avg_duration, + min_duration, + avg_bw, + max_bw, + avg_launch, + avg_transfer, + ) = self._run_and_compute( + self.buffer_size, self.transfer_batch_size, self.iters ) table.add_row( [ @@ -1155,6 +1247,8 @@ def _run_benchmark_loop(self): f"{avg_bw:.2f}", f"{min_duration:.2f}", f"{avg_duration:.2f}", + f"{avg_launch:.2f}", + f"{avg_transfer:.2f}", ] ) @@ -1236,6 +1330,7 @@ def benchmark_xgmi_worker(local_rank, node_rank, args): batch_contiguous=args.batch_contiguous, enable_sess=args.enable_sess, iters=args.iters, + warmup_iters=args.warmup_iters, sweep=args.all, sweep_batch=args.all_batch, sweep_start_size=args.sweep_start_size, @@ -1248,6 +1343,7 @@ def benchmark_xgmi_worker(local_rank, node_rank, args): dst_gpu=args.dst_gpu, num_streams=args.num_streams, num_events=args.num_events, + busy_wait=args.busy_wait, xgmi_multiprocess=True, ) bench.print_config() @@ -1271,6 +1367,7 @@ def benchmark_engine(local_rank, node_rank, args): batch_contiguous=args.batch_contiguous, enable_sess=args.enable_sess, iters=args.iters, + warmup_iters=args.warmup_iters, sweep=args.all, sweep_batch=args.all_batch, sweep_start_size=args.sweep_start_size, @@ -1285,7 +1382,9 @@ def benchmark_engine(local_rank, node_rank, args): num_target_dev=args.num_target_dev, target_dev_offset=args.target_dev_offset, num_qp_per_transfer=args.num_qp_per_transfer, + post_batch_size=args.post_batch_size, num_worker_threads=args.num_worker_threads, + busy_wait=args.busy_wait, poll_cq_mode=args.poll_cq_mode, max_send_wr=args.max_send_wr, max_cqe_num=args.max_cqe_num, @@ -1345,6 +1444,7 @@ def benchmark_xgmi(args): batch_contiguous=args.batch_contiguous, enable_sess=args.enable_sess, iters=args.iters, + warmup_iters=args.warmup_iters, sweep=args.all, sweep_batch=args.all_batch, sweep_start_size=args.sweep_start_size, @@ -1355,6 +1455,7 @@ def benchmark_xgmi(args): dst_gpu=args.dst_gpu, num_streams=args.num_streams, num_events=args.num_events, + busy_wait=args.busy_wait, xgmi_multiprocess=False, ) bench.print_config() From bf3078f2f380c9b70c3a747de822d1f0c1e3bcf4 Mon Sep 17 00:00:00 2001 From: 25337 Date: Wed, 1 Jul 2026 11:42:16 +0000 Subject: [PATCH 02/10] docs(io): add MORI IO tuning guide Co-Authored-By: Claude Sonnet 5 --- docs/MORI-IO-TUNING.md | 199 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 199 insertions(+) create mode 100644 docs/MORI-IO-TUNING.md diff --git a/docs/MORI-IO-TUNING.md b/docs/MORI-IO-TUNING.md new file mode 100644 index 000000000..fa41bb60a --- /dev/null +++ b/docs/MORI-IO-TUNING.md @@ -0,0 +1,199 @@ +# MORI-IO Tuning Guide + +This guide explains the knobs that most directly shape MORI-IO RDMA performance and +how to reason about them. It focuses on five flags exposed by the benchmark +(`tests/python/io/benchmark.py`) that map onto the `RdmaBackendConfig` used by every +transfer. + +## Table of Contents + +- [Where these knobs live](#where-these-knobs-live) +- [The flags](#the-flags) + - [`--num-qp-per-transfer`](#--num-qp-per-transfer) + - [`--post-batch-size`](#--post-batch-size) + - [`--busy-wait`](#--busy-wait) + - [`--disable-chunking`](#--disable-chunking) + - [`--batch-contiguous`](#--batch-contiguous) +- [Best known config (Thor2)](#best-known-config-thor2) +- [Recommended profiles](#recommended-profiles) +- [Environment variable equivalents](#environment-variable-equivalents) +- [Tuning workflow](#tuning-workflow) + +## Where these knobs live + +Three of the flags are backend config (`RdmaBackendConfig`, `include/mori/io/backend.hpp`) +and are honored by every transfer regardless of how you call the engine: + +| Flag | Config field | Env override | +|------|--------------|--------------| +| `--num-qp-per-transfer` | `qpPerTransfer` | `MORI_IO_QP_PER_TRANSFER` | +| `--post-batch-size` | `postBatchSize` | `MORI_IO_POST_BATCH_SIZE` | +| `--disable-chunking` | `enableTransferChunking` (inverted) | `MORI_IO_ENABLE_CHUNKING` | + +Two are benchmark-side only — they change *how the benchmark drives the engine*, not the +backend config: + +| Flag | Effect | +|------|--------| +| `--busy-wait` | Calls `TransferStatus::WaitBusy()` (spin) instead of `Wait()` (block on cv) | +| `--batch-contiguous` | Lays out transfer offsets contiguously so WRs can be merged | + +## The flags + +### `--num-qp-per-transfer` + +**Config `qpPerTransfer`, default `4` in the benchmark (`1` in the raw pybind default).** + +A transfer is split into per-QP post batches spread round-robin across `qpPerTransfer` queue +pairs (the start QP rotates by transfer id, so even single-WR transfers don't all land on +`eps[0]`). More QPs = more parallel send queues, which relieves single-SQ pressure under load; +fewer QPs = less state and lower overhead. `qpPerTransfer` also caps the worker-thread +executor at `min(qpPerTransfer, numWorkerThreads)` threads. + +> **AINIC:** a single QP cannot saturate the link — use **at least 2 QPs**, and **4** for full +> bandwidth. + +**Rule of thumb:** On Thor2, `1`/`2`/`4` are equivalent — default to `1`. Scale up only when +the single SQ is provably the bottleneck (host memory, multi-NIC, ionic), matching +`--num-worker-threads`. + +### `--post-batch-size` + +**Config `postBatchSize`, default `-1` (auto).** + +WRs handed to a single `ibv_post_send`. Two competing effects, and message size decides which +wins (see the [Thor2 BKC](#best-known-config-thor2)): + +- **`-1` (auto)** — posts each QP's whole share in one call + (`ceil(mergedWrCount / qpCount)`, clamped to the SQ's `maxSqDepth`). Fewest doorbell rings; + best for **small messages**, where doorbell overhead dominates. +- **`1`** — one WR per call. Overlaps posting with in-flight transfer (the CPU posts the next + WR while the NIC DMAs the current one), so it **wins on throughput for large messages** + where each DMA is long enough to hide the posting cost. +- On SQ-full / `ENOMEM`, **reduce** `postBatchSize` (or raise `MORI_IO_QP_MAX_SEND_WR` / + `qpPerTransfer`); the error reports the effective value. + +**Rule of thumb:** small messages → `-1`; large messages → `1`. + +### `--busy-wait` + +**Benchmark-side: `WaitBusy()` vs `Wait()`.** + +How the waiting thread observes completion: + +- **Default (blocking)** — blocks on a condition variable. Frees the core, but pays a + cross-thread wakeup latency of ~**5–10 µs** per completion. +- **`--busy-wait`** — spins on the completion flag (`WaitBusy`), removing that latency. + +Good for **latency-sensitive** cases, but spinning **burns a lot of CPU cycles** — avoid it +when many transfers are in flight or CPU is contended, since the wasted cycles are stolen from +the progress/completion path and can hurt throughput. + +### `--disable-chunking` + +**Config `enableTransferChunking`, chunking is ON by default.** + +Chunking splits a transfer into `--chunk-bytes` pieces (default 64 KB, up to `--max-chunks` += 64), which exposes intra-transfer parallelism across QPs and enables the GPU worker-thread +posting path (GPU local memory only; host memory falls back to inline posting). + +- **On (default)** — best for large single transfers, especially GPU memory: pipelined across + QPs/threads. +- **`--disable-chunking`** — one WR per transfer (still capped by the NIC's max message size). + Lower overhead when messages are already small (≤ chunk size) and keeps posting on the + simple inline path. + +`--chunk-bytes` and `--max-chunks` only matter while chunking is enabled. + +### `--batch-contiguous` + +**Benchmark-side: transfer buffer offset layout.** + +Whether batched transfers land on contiguous offsets: + +- **`--batch-contiguous`** — contiguous offsets let adjacent WRs **merge** + (`MergedWorkRequest`) into fewer, larger WRs: less SQ pressure, higher throughput. +- **Default (strided)** — each transfer is a separate WR, maximizing WR count. The stress + path — reproduces SQ-full / `ENOMEM` and measures the per-WR overhead floor. + +**Rule of thumb:** enable it for headline throughput; leave off to stress the SQ. + +## Best known config (Thor2) + +Best-known write-benchmark command on Thor2 (run on each node with its own `--rank`): + +```bash +tools/run_internode_io_benchmark.sh \ + --rank 0 --master-addr 10.162.224.131 --master-port 29573 --ifname enp49s0f1np1 \ + -- --op-type write --enable-batch-transfer --transfer-batch-size 128 \ + --all --sweep-start-size 1024 --sweep-max-size 1048576 --iters 80 \ + --num-qp-per-transfer 1 --post-batch-size 1 --busy-wait --disable-chunking \ + --warmup-iters 64 --batch-contiguous +``` + +> **Choosing `--post-batch-size`:** the value flips with message size. The command above uses +> `1` because it sweeps mostly large messages (≥ 1 KB up to 1 MB), where posting one WR at a +> time overlaps posting with transfer and wins on throughput. For **small messages (≤ 16 KB)** +> use `-1` (auto) instead: transfers are too short to overlap, so doorbell/CPU overhead +> dominates and posting each QP's whole share in one ring is faster. A size-adaptive client +> should pick `post_batch_size = -1 if message_bytes <= 16 * 1024 else 1`. + +Everything else is stable: **`--num-qp-per-transfer` `1`/`2`/`4` are equivalent** on Thor2, so +`1` is the default. + +## Recommended profiles + +**Latency-optimized (the command at the top of this guide):** + +```bash +--num-qp-per-transfer 1 --post-batch-size -1 --busy-wait --disable-chunking --batch-contiguous +``` + +Single QP, no chunking, merged WRs, and a spinning waiter — minimal overhead. Best for small +messages at low queue depth. + +**Throughput-optimized (large messages, > 16 KB):** + +```bash +--num-qp-per-transfer 1 --post-batch-size 1 --batch-contiguous +# chunking left ON (default), no --busy-wait +``` + +`--post-batch-size 1` overlaps posting with in-flight transfer (the throughput win on large +messages per the [Thor2 BKC](#best-known-config-thor2)); contiguous merging keeps the pipeline +full; blocking waits free the CPU for the progress path. `qp=1` is sufficient on Thor2 — raise +it only if the single SQ is provably the bottleneck. + +## Environment variable equivalents + +The config-backed flags can also be set without touching code, which is useful when MORI-IO +is embedded in another app: + +```bash +export MORI_IO_QP_PER_TRANSFER=1 +export MORI_IO_POST_BATCH_SIZE=-1 +export MORI_IO_ENABLE_CHUNKING=0 # disable chunking +export MORI_IO_CHUNK_BYTES=65536 +export MORI_IO_MAX_CHUNKS=64 +export MORI_IO_NUM_NICS_PER_TRANSFER=1 +``` + +`--busy-wait` and `--batch-contiguous` are properties of the caller's transfer loop, not the +backend config, so they have no env equivalent — replicate them in your own client code +(`WaitBusy()` and contiguous offsets, respectively). + +## Tuning workflow + +1. **Start from a profile.** Latency profile for small messages, throughput profile for + large / batched. +2. **Sweep one knob at a time.** Use `--all` (message-size sweep) and `--all-batch` + (batch-size sweep) to find where the current setting stops scaling. +3. **Leave `--num-qp-per-transfer` at `1`** (on Thor2, 1/2/4 are equivalent); scale up only + when the single SQ is provably the bottleneck, and match `--num-worker-threads`. +4. **Set `--post-batch-size` by message size:** `-1` for ≤ 16 KB, `1` for > 16 KB (Thor2 + BKC). Step to a small fixed N only if you hit SQ-full / `ENOMEM` (or raise + `MORI_IO_QP_MAX_SEND_WR`). +5. **Toggle chunking by message size:** on for large single transfers, off once messages + are at or below `--chunk-bytes`. +6. **Only spin (`--busy-wait`) with a spare core** and few outstanding transfers; otherwise + let the blocking waiter free the CPU. From 1bd5dcd3a425f1b74f29b66d78d80c5f5b789568 Mon Sep 17 00:00:00 2001 From: TianDi101 Date: Fri, 24 Jul 2026 03:17:00 +0000 Subject: [PATCH 03/10] fix(io-bench): validate RDMA transfers against the offsets run_single_once actually uses _validate_rdma() unconditionally assumed the strided (buffer_size+1) offset layout whenever --batch-contiguous was off, but that layout is only ever used by run_batch_once() (--enable-batch-transfer). The default path, run_single_once(), always issues transfers at contiguous offsets (buffer_size * i) regardless of --batch-contiguous, causing spurious AssertionErrors on batch_contiguous=False runs with batch > 1 even though the RDMA transfer itself succeeded. Gate the strided comparison on self.enable_batch_transfer and add a branch for the non-batched, non-contiguous-flag case that compares against the contiguous prefix that was actually transferred. --- tests/python/io/benchmark.py | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 1ea858ac5..7cb7a35d6 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -720,8 +720,12 @@ def _validate_rdma(self): dtype=torch.uint8, ) dist.recv(recv_tensor, src=self.num_initiator_dev + self.role_rank) - if not self.batch_contiguous: - # Received data is packed (contiguous); compare to packed view of self.tensor + if not self.batch_contiguous and self.enable_batch_transfer: + # Received data is packed (contiguous); compare to packed view of self.tensor. + # Only run_batch_once() actually honors batch_contiguous (via + # _get_transfer_offsets); run_single_once() always uses contiguous + # offsets (buffer_size * i) regardless of the flag, so validation must + # match that when enable_batch_transfer is False. stride = self.buffer_size + 1 expected = torch.empty( self.buffer_size * self.transfer_batch_size, @@ -735,13 +739,23 @@ def _validate_rdma(self): self.tensor[beg:end].view(torch.uint8) ) assert torch.equal(recv_tensor, expected) + elif not self.batch_contiguous: + # run_single_once() (enable_batch_transfer=False) always issues + # transfers at contiguous offsets (buffer_size * i), ignoring + # batch_contiguous entirely, even though self.tensor was allocated + # oversized ((buffer_size+1)*batch) for the strided layout. Compare + # against the contiguous prefix that was actually transferred. + expected = self.tensor[ + : self.buffer_size * self.transfer_batch_size + ].view(torch.uint8) + assert torch.equal(recv_tensor, expected) else: expected = self.tensor.view(torch.uint8) assert torch.equal(recv_tensor, expected) else: # Without batch_contiguous, tensor has (buffer_size+1)*transfer_batch_size # elements; Gloo send size must match initiator recv (buffer_size*transfer_batch_size). - if not self.batch_contiguous: + if not self.batch_contiguous and self.enable_batch_transfer: stride = self.buffer_size + 1 packed = torch.empty( self.buffer_size * self.transfer_batch_size, @@ -755,6 +769,11 @@ def _validate_rdma(self): self.tensor[beg:end].view(torch.uint8) ) dist.send(packed, dst=self.role_rank) + elif not self.batch_contiguous: + packed = self.tensor[ + : self.buffer_size * self.transfer_batch_size + ].view(torch.uint8) + dist.send(packed, dst=self.role_rank) else: int8_view = self.tensor.view(torch.uint8) dist.send(int8_view, dst=self.role_rank) @@ -958,6 +977,7 @@ def _initialize_fabric(self): "create_session returned None for fabric: peers likely not in the " "same vPOD, or remote memory is not fabric-exportable" ) + def _wait_status(self, status): # Busy-poll spins on the completion flag (no cross-thread cv wakeup); # otherwise block on the condition variable. @@ -1175,7 +1195,9 @@ def _run_benchmark_loop(self): max_bw, avg_launch, avg_transfer, - ) = self._run_and_compute(cur_size, self.transfer_batch_size, self.iters) + ) = self._run_and_compute( + cur_size, self.transfer_batch_size, self.iters + ) table.add_row( [ cur_size, @@ -1444,7 +1466,7 @@ def benchmark_xgmi(args): batch_contiguous=args.batch_contiguous, enable_sess=args.enable_sess, iters=args.iters, - warmup_iters=args.warmup_iters, + warmup_iters=args.warmup_iters, sweep=args.all, sweep_batch=args.all_batch, sweep_start_size=args.sweep_start_size, From 2334624abcf9d338170a1c855a8ff985d14c0ecf Mon Sep 17 00:00:00 2001 From: amirakb89 Date: Tue, 28 Jul 2026 03:45:21 +0000 Subject: [PATCH 04/10] bench(io): add nixlbench-matching C++ engine benchmark bench_engine is a pure-C++ cross-node MORI-IO benchmark whose measurement loop matches nixlbench (NVIDIA NIXL xferbench): a single whole-loop timer, inline spin-poll completion in the bench thread, GB=1e9 throughput, warmup excluded, and pipeline depth (--inflight) equivalent to nixl --pipeline_depth. This makes MORI-IO and NIXL RDMA numbers directly comparable on the same fabric. Ports the MORI Python benchmark feature set: batching (strided default + --batch-contiguous), --enable-batch-transfer, chunking (--chunk-bytes/--max-chunks), session fast-path (--enable-sess), --busy-wait, --poll_cq_mode, --mem-type gpu|cpu, --post-batch-size, and the max-send-wr/cqe/msg-sge QP knobs. Flag defaults match the Python bench (enable-sess and busy-wait off by default). Rendezvous is a tiny length-prefixed TCP socket exchanging msgpack-packed EngineDesc/MemoryDesc; rank 0 initiator drives one-sided WRITE/READ, rank 1 target only holds memory registered. Co-Authored-By: Claude Opus 4 --- tests/cpp/CMakeLists.txt | 6 + tests/cpp/io/bench_engine.cpp | 416 ++++++++++++++++++++++++++++++++++ 2 files changed, 422 insertions(+) create mode 100644 tests/cpp/io/bench_engine.cpp diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 34ba20a64..e96effab6 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -47,6 +47,12 @@ if(BUILD_IO) target_include_directories(test_engine PUBLIC ${CMAKE_SOURCE_DIR}) target_link_libraries(test_engine mori_application mori_io) + add_executable(bench_engine io/bench_engine.cpp) + + target_include_directories(bench_engine PUBLIC ${CMAKE_SOURCE_DIR}/include) + target_include_directories(bench_engine PUBLIC ${CMAKE_SOURCE_DIR}) + target_link_libraries(bench_engine mori_application mori_io) + add_executable(test_transfer_wait io/test_transfer_wait.cpp) target_include_directories(test_transfer_wait diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp new file mode 100644 index 000000000..291bc921b --- /dev/null +++ b/tests/cpp/io/bench_engine.cpp @@ -0,0 +1,416 @@ +// bench_engine.cpp +// +// A C++ benchmark for MORI-IO whose measurement methodology exactly matches +// nixlbench (NVIDIA NIXL's xferbench). The point is an apples-to-apples RDMA +// throughput/latency comparison between MORI-IO and NIXL on the same fabric, +// eliminating the gaps that exist between MORI's Python benchmark and nixlbench: +// +// 1. WHOLE-LOOP TIMER (not per-iteration sum). +// nixlbench wraps ONE timer around the entire num_iter loop +// (nixl_worker.cpp: total_timer -> total_duration = total_timer.lap()), +// then throughput = total_bytes / total_duration. MORI's Python bench +// instead times each iteration (launch+transfer) and SUMS them, which +// excludes the inter-iteration gap. Here we use a single whole-loop timer. +// +// 2. INLINE SPIN-POLL COMPLETION in the benchmark thread. +// nixlbench polls agent->getXferStatus() in a tight spin (NIXL_IN_PROG -> +// continue) directly in the bench thread. MORI's WaitBusy() spins on the +// completion flag (no cv wakeup) -> we use WaitBusy(), the closest analog. +// +// 3. PIPELINE DEPTH. +// nixlbench keeps `pipeline_depth` transfers in flight (default 1). We +// expose --inflight to match: default 1 = strict stop-and-wait, matching +// nixl --pipeline_depth 1. +// +// 4. WARMUP excluded from timing (nixl: --warmup_iter, default 100). +// +// 5. throughput_gb = total_bytes / 1e9 / (total_duration_us / 1e6), GB=10^9, +// identical unit to nixl. avg_latency = total_duration / num_iter. +// +// Rendezvous: 2 processes (one per node), TCP socket exchange of EngineDesc and +// MemoryDesc (msgpack, same as MORI's pybind pack()/unpack()). Rank 0 = +// initiator, rank 1 = target. Initiator drives all transfers (RDMA one-sided +// WRITE/READ); target only registers memory and waits. +// +// Build: see build.sh (links libmori_io + libmori_application from the editable +// install, includes 3rdparty/msgpack-c + HIP). + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "mori/io/io.hpp" + +using namespace mori::io; +using Clock = std::chrono::steady_clock; + +#define HIP_CHECK(expr) \ + do { \ + hipError_t _e = (expr); \ + if (_e != hipSuccess) { \ + std::cerr << "HIP error " << hipGetErrorString(_e) << " at " << __FILE__ << ":" \ + << __LINE__ << std::endl; \ + std::exit(1); \ + } \ + } while (0) + +// ------------------------- tiny TCP rendezvous ----------------------------- +// Rank 0 listens, rank 1 connects. Then symmetric length-prefixed blob swap. + +static int TcpListenAccept(uint16_t port) { + int lfd = socket(AF_INET, SOCK_STREAM, 0); + int opt = 1; + setsockopt(lfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = INADDR_ANY; + addr.sin_port = htons(port); + if (bind(lfd, reinterpret_cast(&addr), sizeof(addr)) != 0) { + perror("bind"); + std::exit(1); + } + listen(lfd, 1); + int fd = accept(lfd, nullptr, nullptr); + close(lfd); + return fd; +} + +static int TcpConnect(const std::string& host, uint16_t port) { + int fd = socket(AF_INET, SOCK_STREAM, 0); + sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_port = htons(port); + inet_pton(AF_INET, host.c_str(), &addr.sin_addr); + for (int retry = 0; retry < 100; ++retry) { + if (connect(fd, reinterpret_cast(&addr), sizeof(addr)) == 0) return fd; + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + } + std::cerr << "connect failed to " << host << ":" << port << std::endl; + std::exit(1); +} + +static void SendAll(int fd, const void* buf, size_t n) { + const char* p = static_cast(buf); + while (n) { + ssize_t k = send(fd, p, n, 0); + if (k <= 0) { perror("send"); std::exit(1); } + p += k; n -= k; + } +} +static void RecvAll(int fd, void* buf, size_t n) { + char* p = static_cast(buf); + while (n) { + ssize_t k = recv(fd, p, n, 0); + if (k <= 0) { perror("recv"); std::exit(1); } + p += k; n -= k; + } +} +static void SendBlob(int fd, const std::string& b) { + uint64_t len = b.size(); + SendAll(fd, &len, sizeof(len)); + SendAll(fd, b.data(), len); +} +static std::string RecvBlob(int fd) { + uint64_t len = 0; + RecvAll(fd, &len, sizeof(len)); + std::string b(len, '\0'); + RecvAll(fd, b.data(), len); + return b; +} +static void Barrier(int fd) { // symmetric 1-byte ping-pong + char c = 'x'; + SendAll(fd, &c, 1); + RecvAll(fd, &c, 1); +} + +template +static std::string Pack(const T& v) { + msgpack::sbuffer buf; + msgpack::pack(buf, v); + return std::string(buf.data(), buf.size()); +} +template +static T Unpack(const std::string& b) { + auto oh = msgpack::unpack(b.data(), b.size()); + return oh.get().as(); +} + +// ------------------------------- config ------------------------------------ +// Flags mirror MORI's Python benchmark (tests/python/io/benchmark.py) so runs +// are directly comparable. Names use the Python spelling where they exist. +struct Args { + int rank = 0; // 0=initiator, 1=target + std::string master_ip; // rank 1 connects here (rank 0's data IP) for TCP rendezvous + std::string self_ip; // this rank's own reachable IP (advertised in EngineDesc) + uint16_t port = 18515; + int gpu = 0; + std::string op = "write"; // write|read + size_t sweep_start = 1u << 20; // 1 MiB + size_t sweep_max = 32u << 20; // 32 MiB + size_t sweep_step = 1u << 20; // 1 MiB + int iters = 500; + int warmup = 50; // --warmup-iters + int inflight = 1; // == nixl pipeline_depth (transfers/requests in flight) + + // RdmaBackendConfig knobs + int qp_per_transfer = 4; // --num-qp-per-transfer + int worker_threads = 1; // --num-worker-threads + int post_batch_size = -1; // --post-batch-size + std::string poll_cq_mode = "polling"; // polling|event + bool disable_chunking = false; // --disable-chunking (default: chunking ON, like Python) + size_t chunk_bytes = 65536; // --chunk-bytes + int max_chunks = 64; // --max-chunks + int max_send_wr = 0; // --max-send-wr (0 = leave default) + int max_cqe_num = 0; // --max-cqe-num + int max_msg_sge = 0; // --max-msg-sge + + // batch / session + int batch = 1; // --transfer-batch-size (transfers per request) + bool enable_batch_transfer = false; // --enable-batch-transfer (use BatchWrite) + bool batch_contiguous = false; // --batch-contiguous (adjacent offsets → merged WR); + // default strided (each transfer a separate WR) + bool enable_sess = false; // --enable-sess (session fast-path); Python default: off + bool busy_wait = false; // --busy-wait (WaitBusy spin); Python default: off + + std::string mem_type = "gpu"; // --mem-type gpu|cpu +}; + +static Args ParseArgs(int argc, char** argv) { + Args a; + for (int i = 1; i < argc; ++i) { + std::string k = argv[i]; + auto next = [&]() { return std::string(argv[++i]); }; + if (k == "--rank") a.rank = std::stoi(next()); + else if (k == "--master-ip") a.master_ip = next(); + else if (k == "--self-ip") a.self_ip = next(); + else if (k == "--port") a.port = static_cast(std::stoi(next())); + else if (k == "--gpu") a.gpu = std::stoi(next()); + else if (k == "--op" || k == "--op-type") a.op = next(); + else if (k == "--sweep-start" || k == "--sweep-start-size") a.sweep_start = std::stoull(next()); + else if (k == "--sweep-max" || k == "--sweep-max-size") a.sweep_max = std::stoull(next()); + else if (k == "--sweep-step") a.sweep_step = std::stoull(next()); + else if (k == "--iters") a.iters = std::stoi(next()); + else if (k == "--warmup" || k == "--warmup-iters") a.warmup = std::stoi(next()); + else if (k == "--inflight") a.inflight = std::stoi(next()); + else if (k == "--qp-per-transfer" || k == "--num-qp-per-transfer") a.qp_per_transfer = std::stoi(next()); + else if (k == "--worker-threads" || k == "--num-worker-threads") a.worker_threads = std::stoi(next()); + else if (k == "--post-batch-size") a.post_batch_size = std::stoi(next()); + else if (k == "--poll_cq_mode" || k == "--poll-cq-mode") a.poll_cq_mode = next(); + else if (k == "--disable-chunking") a.disable_chunking = true; + else if (k == "--chunk-bytes") a.chunk_bytes = std::stoull(next()); + else if (k == "--max-chunks") a.max_chunks = std::stoi(next()); + else if (k == "--max-send-wr") a.max_send_wr = std::stoi(next()); + else if (k == "--max-cqe-num") a.max_cqe_num = std::stoi(next()); + else if (k == "--max-msg-sge") a.max_msg_sge = std::stoi(next()); + else if (k == "--batch" || k == "--transfer-batch-size") a.batch = std::stoi(next()); + else if (k == "--enable-batch-transfer") a.enable_batch_transfer = true; + else if (k == "--batch-contiguous") a.batch_contiguous = true; + else if (k == "--enable-sess") a.enable_sess = true; + else if (k == "--disable-sess") a.enable_sess = false; + else if (k == "--busy-wait") a.busy_wait = true; + else if (k == "--no-busy-wait") a.busy_wait = false; + else if (k == "--mem-type") a.mem_type = next(); + else { std::cerr << "unknown arg " << k << std::endl; std::exit(1); } + } + return a; +} + +// ------------------------------ main --------------------------------------- +int main(int argc, char** argv) { + Args a = ParseArgs(argc, argv); + HIP_CHECK(hipSetDevice(a.gpu)); + + // Buffer must hold the largest single REQUEST. Strided batch (default) needs + // (msg+1)*batch to keep slots non-adjacent; contiguous needs msg*batch. Round + // the strided case up to (sweep_max+1)*batch. + const bool batched = a.enable_batch_transfer && a.batch > 1; + const size_t slotStride = a.batch_contiguous ? a.sweep_max : (a.sweep_max + 1); + const size_t bufBytes = batched ? slotStride * static_cast(a.batch) : a.sweep_max; + const bool cpuMem = (a.mem_type == "cpu"); + const MemoryLocationType memLoc = cpuMem ? MemoryLocationType::CPU : MemoryLocationType::GPU; + + void* buf = nullptr; + if (cpuMem) { + HIP_CHECK(hipHostMalloc(&buf, bufBytes, 0)); + } else { + HIP_CHECK(hipMalloc(&buf, bufBytes)); + } + HIP_CHECK(hipMemset(buf, 0, bufBytes)); + + // Out-of-band control endpoint for the MORI engine. host MUST be an IP the + // peer can reach (advertised via EngineDesc for RDMA QP setup); self_ip + // defaults to master_ip on rank 0. + IOEngineConfig cfg; + cfg.host = !a.self_ip.empty() ? a.self_ip : (a.rank == 0 ? a.master_ip : a.master_ip); + cfg.port = static_cast(a.port + 1 + a.rank); // distinct per rank + std::string key = a.rank == 0 ? "initiator" : "target"; + IOEngine engine(key, cfg); + + RdmaBackendConfig rdmaCfg{}; + rdmaCfg.qpPerTransfer = a.qp_per_transfer; + rdmaCfg.postBatchSize = a.post_batch_size; + rdmaCfg.numWorkerThreads = a.worker_threads; + rdmaCfg.pollCqMode = (a.poll_cq_mode == "event") ? PollCqMode::EVENT : PollCqMode::POLLING; + rdmaCfg.enableNotification = false; // match MORI Python bench RDMA path + rdmaCfg.enableTransferChunking = !a.disable_chunking; // chunking ON by default (Python parity) + rdmaCfg.chunkBytes = a.chunk_bytes; + rdmaCfg.maxChunksPerTransfer = a.max_chunks; + if (a.max_send_wr > 0) rdmaCfg.maxSendWr = a.max_send_wr; + if (a.max_cqe_num > 0) rdmaCfg.maxCqeNum = a.max_cqe_num; + if (a.max_msg_sge > 0) rdmaCfg.maxMsgSge = a.max_msg_sge; + engine.CreateBackend(BackendType::RDMA, rdmaCfg); + + MemoryDesc localMem = engine.RegisterMemory(buf, bufBytes, a.gpu, memLoc); + + // --- rendezvous over a side TCP socket ----------------------------------- + int sock = (a.rank == 0) ? TcpListenAccept(a.port) : TcpConnect(a.master_ip, a.port); + int one = 1; + setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)); + + // Exchange EngineDesc then register the remote engine. + EngineDesc myEng = engine.GetEngineDesc(); + SendBlob(sock, Pack(myEng)); + EngineDesc peerEng = Unpack(RecvBlob(sock)); + engine.RegisterRemoteEngine(peerEng); + + // Exchange MemoryDesc. + SendBlob(sock, Pack(localMem)); + MemoryDesc peerMem = Unpack(RecvBlob(sock)); + + Barrier(sock); + + if (a.rank == 1) { + // Target: nothing to drive. RDMA one-sided ops complete without target CPU. + // Just hold memory registered until the initiator says it's done. + Barrier(sock); // wait for initiator to finish the whole sweep + if (cpuMem) HIP_CHECK(hipHostFree(buf)); else HIP_CHECK(hipFree(buf)); + close(sock); + return 0; + } + + // -------------------- initiator: run the sweep --------------------------- + // Session fast-path (matches MORI --enable-sess). When --disable-sess, we call + // the engine batch/single APIs directly with explicit MemoryDesc + uid vecs. + const bool useSess = a.enable_sess; + IOEngineSession* sessPtr = nullptr; + std::optional sessOpt; + if (useSess) { + sessOpt = engine.CreateSession(localMem, peerMem); + if (!sessOpt) { std::cerr << "CreateSession failed" << std::endl; std::exit(1); } + sessPtr = &*sessOpt; + } + const bool isRead = (a.op == "read"); + + auto waitDone = [&](TransferStatus& s) { + if (a.busy_wait) { s.WaitBusy(); } // spin on completion flag + else { s.Wait(); } // block on condition variable + }; + + std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; + + for (size_t msg = a.sweep_start; msg <= a.sweep_max; msg += a.sweep_step) { + const int depth = std::min(a.inflight, a.iters); + std::vector st(depth); + std::vector uid(depth); + + // Batch layout. Strided (default): slot i at (msg+1)*i so the N transfers + // stay SEPARATE WRs (stresses SQ, real batching) — matches Python default. + // Contiguous (--batch-contiguous): slot i at msg*i, adjacent, so MORI may + // merge them into one big WR (fast but not really batching, and hits the + // 1 GiB max_msg_sz without chunking). + const size_t stride = a.batch_contiguous ? msg : (msg + 1); + SizeVec offsets(a.batch), sizes(a.batch); + for (int i = 0; i < a.batch; ++i) { + offsets[i] = static_cast(i) * stride; + sizes[i] = msg; + } + // Engine (non-session) batch path needs vec-of-vec + desc vectors. + MemDescVec locVec{localMem}, remVec{peerMem}; + BatchSizeVec offVec{offsets}, sizeVec{sizes}; + + auto post = [&](int slot) { + st[slot].SetCode(StatusCode::INIT); + uid[slot] = useSess ? sessPtr->AllocateTransferUniqueId() : engine.AllocateTransferUniqueId(); + if (batched) { + if (useSess) { + if (isRead) sessPtr->BatchRead(offsets, offsets, sizes, &st[slot], uid[slot]); + else sessPtr->BatchWrite(offsets, offsets, sizes, &st[slot], uid[slot]); + } else { + TransferStatusPtrVec sp{&st[slot]}; + TransferUniqueIdVec ids{uid[slot]}; + if (isRead) engine.BatchRead(locVec, offVec, remVec, offVec, sizeVec, sp, ids); + else engine.BatchWrite(locVec, offVec, remVec, offVec, sizeVec, sp, ids); + } + } else { + if (useSess) { + if (isRead) sessPtr->Read(0, 0, msg, &st[slot], uid[slot]); + else sessPtr->Write(0, 0, msg, &st[slot], uid[slot]); + } else { + if (isRead) engine.Read(localMem, 0, peerMem, 0, msg, &st[slot], uid[slot]); + else engine.Write(localMem, 0, peerMem, 0, msg, &st[slot], uid[slot]); + } + } + }; + + // ---- warmup (excluded from timing) ---- + for (int w = 0; w < a.warmup; ++w) { + post(0); + waitDone(st[0]); + } + + // ---- timed region: ONE whole-loop timer, nixl-style ---- + int completed = 0, issued = 0; + auto t0 = Clock::now(); + + // prime the pipeline + for (int s = 0; s < depth; ++s) { post(s); ++issued; } + + while (completed < a.iters) { + for (int s = 0; s < depth; ++s) { + // spin-poll this slot (WaitBusy with a single poll would also work, but + // we mirror nixl's "check status, if IN_PROG continue" scan) + if (st[s].InProgress() || st[s].Init()) continue; + if (st[s].Failed()) { + std::cerr << "transfer failed: " << st[s].Message() << std::endl; + std::exit(1); + } + // completed one + ++completed; + if (issued < a.iters) { post(s); ++issued; } + } + } + auto t1 = Clock::now(); + + double total_us = + std::chrono::duration_cast>(t1 - t0).count(); + // Each of the `iters` requests moves msg*batch bytes when batching (nixl + // counts block_size*batch_size*num_iter). avg_lat is per REQUEST. + const int effBatch = batched ? a.batch : 1; + double total_bytes = static_cast(msg) * effBatch * a.iters; + double avg_bw = (total_bytes / 1e9) / (total_us / 1e6); // GB/s, GB=10^9 + double avg_lat = total_us / a.iters; // us per request + + std::printf("%-11zu %-6d %-6d %-12.2f %-11.2f %-.1f\n", msg, effBatch, a.iters, avg_bw, avg_lat, + total_us); + std::fflush(stdout); + } + + Barrier(sock); // tell target we're done + if (cpuMem) HIP_CHECK(hipHostFree(buf)); else HIP_CHECK(hipFree(buf)); + close(sock); + return 0; +} From 11cac3f2462bc81bfb9eefefd3f18b5d7fdfea0f Mon Sep 17 00:00:00 2001 From: Akbarzadeh Date: Tue, 28 Jul 2026 13:04:19 -0500 Subject: [PATCH 05/10] bench(io): add Python-parity flags to C++ engine benchmark Match benchmark.py: size/batch sweeps (--all, --all-batch, --buffer-size, --sweep-start-size/--sweep-max-size/--sweep-step), per-side memory type (--initiator-mem-type/--target-mem-type), --target-dev-offset for cross-rail transfers, and --log-level. --- tests/cpp/io/bench_engine.cpp | 96 ++++++++++++++++++++++++++++------- 1 file changed, 79 insertions(+), 17 deletions(-) diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp index 291bc921b..271ac8f58 100644 --- a/tests/cpp/io/bench_engine.cpp +++ b/tests/cpp/io/bench_engine.cpp @@ -41,6 +41,7 @@ #include #include +#include #include #include #include @@ -48,6 +49,7 @@ #include #include #include +#include #include #include @@ -158,10 +160,17 @@ struct Args { std::string self_ip; // this rank's own reachable IP (advertised in EngineDesc) uint16_t port = 18515; int gpu = 0; + int target_dev_offset = 0; // --target-dev-offset (target GPU = (gpu+offset)%ndev) std::string op = "write"; // write|read - size_t sweep_start = 1u << 20; // 1 MiB - size_t sweep_max = 32u << 20; // 32 MiB - size_t sweep_step = 1u << 20; // 1 MiB + + // Sweep control (Python parity). --all sweeps message size; --all-batch sweeps + // batch size. Neither set => single run at (buffer_size, batch). + bool sweep_all = false; // --all + bool sweep_batch = false; // --all-batch + size_t buffer_size = 32768; // --buffer-size (single message size when not sweeping) + size_t sweep_start = 8; // --sweep-start-size (Python default 8) + size_t sweep_max = 1u << 20; // --sweep-max-size (Python default 2^20) + size_t sweep_step = 0; // --sweep-step (0 = geometric x2; >0 = linear +step) int iters = 500; int warmup = 50; // --warmup-iters int inflight = 1; // == nixl pipeline_depth (transfers/requests in flight) @@ -187,6 +196,10 @@ struct Args { bool busy_wait = false; // --busy-wait (WaitBusy spin); Python default: off std::string mem_type = "gpu"; // --mem-type gpu|cpu + std::string init_mem_type; // --initiator-mem-type (empty => mem_type) + std::string target_mem_type; // --target-mem-type (empty => mem_type) + + std::string log_level = "info"; // --log-level trace|debug|info|warning|error|critical }; static Args ParseArgs(int argc, char** argv) { @@ -199,7 +212,11 @@ static Args ParseArgs(int argc, char** argv) { else if (k == "--self-ip") a.self_ip = next(); else if (k == "--port") a.port = static_cast(std::stoi(next())); else if (k == "--gpu") a.gpu = std::stoi(next()); + else if (k == "--target-dev-offset") a.target_dev_offset = std::stoi(next()); else if (k == "--op" || k == "--op-type") a.op = next(); + else if (k == "--all") a.sweep_all = true; + else if (k == "--all-batch") a.sweep_batch = true; + else if (k == "--buffer-size") a.buffer_size = std::stoull(next()); else if (k == "--sweep-start" || k == "--sweep-start-size") a.sweep_start = std::stoull(next()); else if (k == "--sweep-max" || k == "--sweep-max-size") a.sweep_max = std::stoull(next()); else if (k == "--sweep-step") a.sweep_step = std::stoull(next()); @@ -224,6 +241,9 @@ static Args ParseArgs(int argc, char** argv) { else if (k == "--busy-wait") a.busy_wait = true; else if (k == "--no-busy-wait") a.busy_wait = false; else if (k == "--mem-type") a.mem_type = next(); + else if (k == "--initiator-mem-type") a.init_mem_type = next(); + else if (k == "--target-mem-type") a.target_mem_type = next(); + else if (k == "--log-level") a.log_level = next(); else { std::cerr << "unknown arg " << k << std::endl; std::exit(1); } } return a; @@ -232,17 +252,56 @@ static Args ParseArgs(int argc, char** argv) { // ------------------------------ main --------------------------------------- int main(int argc, char** argv) { Args a = ParseArgs(argc, argv); - HIP_CHECK(hipSetDevice(a.gpu)); - - // Buffer must hold the largest single REQUEST. Strided batch (default) needs - // (msg+1)*batch to keep slots non-adjacent; contiguous needs msg*batch. Round - // the strided case up to (sweep_max+1)*batch. - const bool batched = a.enable_batch_transfer && a.batch > 1; - const size_t slotStride = a.batch_contiguous ? a.sweep_max : (a.sweep_max + 1); - const size_t bufBytes = batched ? slotStride * static_cast(a.batch) : a.sweep_max; - const bool cpuMem = (a.mem_type == "cpu"); + SetLogLevel(a.log_level); + + // Per-role memory type: initiator (rank 0) / target (rank 1) may override the + // shared --mem-type, enabling mixed CPU<->GPU transfers. Each process only + // allocates its own side, so no cross-node coupling is needed. + const std::string myMem = (a.rank == 0) + ? (!a.init_mem_type.empty() ? a.init_mem_type : a.mem_type) + : (!a.target_mem_type.empty() ? a.target_mem_type : a.mem_type); + const bool cpuMem = (myMem == "cpu"); const MemoryLocationType memLoc = cpuMem ? MemoryLocationType::CPU : MemoryLocationType::GPU; + // Target GPU shift for cross-rail pairing (GPU memory only, matches Python). + int gpu = a.gpu; + if (a.rank == 1 && !cpuMem && a.target_dev_offset != 0) { + int ndev = 0; + HIP_CHECK(hipGetDeviceCount(&ndev)); + if (ndev > 0) gpu = (a.gpu + a.target_dev_offset) % ndev; + } + HIP_CHECK(hipSetDevice(gpu)); // valid device context needed even for host mem + + // Build the run plan: a list of (msgSize, batch) points. + // --all => sweep message size (geometric x2, or linear when --sweep-step>0), + // batch fixed at --transfer-batch-size. + // --all-batch => msg fixed at --buffer-size, batch = 1,2,4,...,32768. + // neither => single point (--buffer-size, --transfer-batch-size). + std::vector> plan; + if (a.sweep_all) { + for (size_t msg = a.sweep_start; msg <= a.sweep_max; + msg = (a.sweep_step > 0) ? msg + a.sweep_step : msg * 2) { + plan.emplace_back(msg, a.batch); + } + } else if (a.sweep_batch) { + for (int b = 1; b <= 32768; b *= 2) plan.emplace_back(a.buffer_size, b); + } else { + plan.emplace_back(a.buffer_size, a.batch); + } + + // Buffer must hold the largest single REQUEST across the whole plan. Strided + // batch (default) needs (msg+1)*batch to keep slots non-adjacent; contiguous + // needs msg*batch. + size_t bufBytes = 0; + for (auto& planEntry : plan) { + const size_t msg = planEntry.first; + const int b = planEntry.second; + const bool pBatched = a.enable_batch_transfer && b > 1; + const size_t slotStride = a.batch_contiguous ? msg : (msg + 1); + const size_t need = pBatched ? slotStride * static_cast(b) : msg; + bufBytes = std::max(bufBytes, need); + } + void* buf = nullptr; if (cpuMem) { HIP_CHECK(hipHostMalloc(&buf, bufBytes, 0)); @@ -274,7 +333,7 @@ int main(int argc, char** argv) { if (a.max_msg_sge > 0) rdmaCfg.maxMsgSge = a.max_msg_sge; engine.CreateBackend(BackendType::RDMA, rdmaCfg); - MemoryDesc localMem = engine.RegisterMemory(buf, bufBytes, a.gpu, memLoc); + MemoryDesc localMem = engine.RegisterMemory(buf, bufBytes, gpu, memLoc); // --- rendezvous over a side TCP socket ----------------------------------- int sock = (a.rank == 0) ? TcpListenAccept(a.port) : TcpConnect(a.master_ip, a.port); @@ -322,7 +381,10 @@ int main(int argc, char** argv) { std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; - for (size_t msg = a.sweep_start; msg <= a.sweep_max; msg += a.sweep_step) { + for (auto& planEntry : plan) { + const size_t msg = planEntry.first; + const int curBatch = planEntry.second; + const bool batched = a.enable_batch_transfer && curBatch > 1; const int depth = std::min(a.inflight, a.iters); std::vector st(depth); std::vector uid(depth); @@ -333,8 +395,8 @@ int main(int argc, char** argv) { // merge them into one big WR (fast but not really batching, and hits the // 1 GiB max_msg_sz without chunking). const size_t stride = a.batch_contiguous ? msg : (msg + 1); - SizeVec offsets(a.batch), sizes(a.batch); - for (int i = 0; i < a.batch; ++i) { + SizeVec offsets(curBatch), sizes(curBatch); + for (int i = 0; i < curBatch; ++i) { offsets[i] = static_cast(i) * stride; sizes[i] = msg; } @@ -399,7 +461,7 @@ int main(int argc, char** argv) { std::chrono::duration_cast>(t1 - t0).count(); // Each of the `iters` requests moves msg*batch bytes when batching (nixl // counts block_size*batch_size*num_iter). avg_lat is per REQUEST. - const int effBatch = batched ? a.batch : 1; + const int effBatch = batched ? curBatch : 1; double total_bytes = static_cast(msg) * effBatch * a.iters; double avg_bw = (total_bytes / 1e9) / (total_us / 1e6); // GB/s, GB=10^9 double avg_lat = total_us / a.iters; // us per request From f690b3fcd6eef040b0e89243d63e02fed1cf0fdb Mon Sep 17 00:00:00 2001 From: Akbarzadeh Date: Tue, 28 Jul 2026 14:51:59 -0500 Subject: [PATCH 06/10] bench(io): align C++ engine bench metrics/knobs with nixlbench - latency is now per single transfer (total / (iters*batch)), matching nixl's avg_latency = total_duration/(per_thread_iter*batch_size); BW unchanged (msg*batch*iters over one whole-loop timer). - completion is always an inline spin (like nixl's getXferStatus scan); drop the --busy-wait flag, which only affected untimed warmup and would otherwise make latency non-comparable to nixl. Warmup now spins too. - batching self-activates on --transfer-batch-size > 1 (one N-descriptor batch request); drop the redundant --enable-batch-transfer gate that silently made batch sweeps a no-op by default. - simplify dead cfg.host ternary; document that --inflight > 1 reuses the same buffer window (fine for an unvalidated perf bench, matching nixl). --- tests/cpp/io/bench_engine.cpp | 68 ++++++++++++++++++++++------------- 1 file changed, 44 insertions(+), 24 deletions(-) diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp index 271ac8f58..f60276065 100644 --- a/tests/cpp/io/bench_engine.cpp +++ b/tests/cpp/io/bench_engine.cpp @@ -14,8 +14,12 @@ // // 2. INLINE SPIN-POLL COMPLETION in the benchmark thread. // nixlbench polls agent->getXferStatus() in a tight spin (NIXL_IN_PROG -> -// continue) directly in the bench thread. MORI's WaitBusy() spins on the -// completion flag (no cv wakeup) -> we use WaitBusy(), the closest analog. +// continue) directly in the bench thread. We mirror that exactly: the timed +// loop spins on the transfer status flag (stored by MORI's CQ worker +// thread). Completion is ALWAYS spin here -- never a cv-block Wait() -- so +// the measured latency matches nixl (a cv wakeup would add ~5-10us and make +// numbers non-comparable). MORI's blocking Wait() path is exercised only by +// the Python benchmark, not this nixl-parity tool. // // 3. PIPELINE DEPTH. // nixlbench keeps `pipeline_depth` transfers in flight (default 1). We @@ -25,7 +29,10 @@ // 4. WARMUP excluded from timing (nixl: --warmup_iter, default 100). // // 5. throughput_gb = total_bytes / 1e9 / (total_duration_us / 1e6), GB=10^9, -// identical unit to nixl. avg_latency = total_duration / num_iter. +// identical unit to nixl. total_bytes = msg * batch * num_iter, and +// avg_latency = total_duration / (num_iter * batch) -- i.e. PER SINGLE +// TRANSFER, matching nixl's total_duration/(per_thread_iter*batch_size) +// (nixl counts block_size*batch_size descriptors per request). // // Rendezvous: 2 processes (one per node), TCP socket exchange of EngineDesc and // MemoryDesc (msgpack, same as MORI's pybind pack()/unpack()). Rank 0 = @@ -189,11 +196,9 @@ struct Args { // batch / session int batch = 1; // --transfer-batch-size (transfers per request) - bool enable_batch_transfer = false; // --enable-batch-transfer (use BatchWrite) bool batch_contiguous = false; // --batch-contiguous (adjacent offsets → merged WR); // default strided (each transfer a separate WR) bool enable_sess = false; // --enable-sess (session fast-path); Python default: off - bool busy_wait = false; // --busy-wait (WaitBusy spin); Python default: off std::string mem_type = "gpu"; // --mem-type gpu|cpu std::string init_mem_type; // --initiator-mem-type (empty => mem_type) @@ -234,12 +239,9 @@ static Args ParseArgs(int argc, char** argv) { else if (k == "--max-cqe-num") a.max_cqe_num = std::stoi(next()); else if (k == "--max-msg-sge") a.max_msg_sge = std::stoi(next()); else if (k == "--batch" || k == "--transfer-batch-size") a.batch = std::stoi(next()); - else if (k == "--enable-batch-transfer") a.enable_batch_transfer = true; else if (k == "--batch-contiguous") a.batch_contiguous = true; else if (k == "--enable-sess") a.enable_sess = true; else if (k == "--disable-sess") a.enable_sess = false; - else if (k == "--busy-wait") a.busy_wait = true; - else if (k == "--no-busy-wait") a.busy_wait = false; else if (k == "--mem-type") a.mem_type = next(); else if (k == "--initiator-mem-type") a.init_mem_type = next(); else if (k == "--target-mem-type") a.target_mem_type = next(); @@ -296,7 +298,7 @@ int main(int argc, char** argv) { for (auto& planEntry : plan) { const size_t msg = planEntry.first; const int b = planEntry.second; - const bool pBatched = a.enable_batch_transfer && b > 1; + const bool pBatched = b > 1; const size_t slotStride = a.batch_contiguous ? msg : (msg + 1); const size_t need = pBatched ? slotStride * static_cast(b) : msg; bufBytes = std::max(bufBytes, need); @@ -311,10 +313,11 @@ int main(int argc, char** argv) { HIP_CHECK(hipMemset(buf, 0, bufBytes)); // Out-of-band control endpoint for the MORI engine. host MUST be an IP the - // peer can reach (advertised via EngineDesc for RDMA QP setup); self_ip - // defaults to master_ip on rank 0. + // peer can reach (advertised via EngineDesc for RDMA QP setup). Defaults to + // master_ip: correct for rank 0 (it IS the master); rank 1 should pass + // --self-ip when its reachable IP differs from the master's. IOEngineConfig cfg; - cfg.host = !a.self_ip.empty() ? a.self_ip : (a.rank == 0 ? a.master_ip : a.master_ip); + cfg.host = !a.self_ip.empty() ? a.self_ip : a.master_ip; cfg.port = static_cast(a.port + 1 + a.rank); // distinct per rank std::string key = a.rank == 0 ? "initiator" : "target"; IOEngine engine(key, cfg); @@ -374,9 +377,15 @@ int main(int argc, char** argv) { } const bool isRead = (a.op == "read"); - auto waitDone = [&](TransferStatus& s) { - if (a.busy_wait) { s.WaitBusy(); } // spin on completion flag - else { s.Wait(); } // block on condition variable + // Spin to completion, same mechanism as the timed loop (status flag stored by + // MORI's CQ worker thread). Used for warmup; keeps warmup identical to the + // measured path so caches/QP state are primed the same way. + auto spinDone = [&](TransferStatus& s) { + while (s.InProgress() || s.Init()) { /* spin: CQ worker thread stores status */ } + if (s.Failed()) { + std::cerr << "warmup transfer failed: " << s.Message() << std::endl; + std::exit(1); + } }; std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; @@ -384,7 +393,9 @@ int main(int argc, char** argv) { for (auto& planEntry : plan) { const size_t msg = planEntry.first; const int curBatch = planEntry.second; - const bool batched = a.enable_batch_transfer && curBatch > 1; + // nixl-style: batch>1 always issues one N-descriptor batch request (no + // separate enable flag). batch==1 is the degenerate single-transfer case. + const bool batched = curBatch > 1; const int depth = std::min(a.inflight, a.iters); std::vector st(depth); std::vector uid(depth); @@ -431,20 +442,24 @@ int main(int argc, char** argv) { // ---- warmup (excluded from timing) ---- for (int w = 0; w < a.warmup; ++w) { post(0); - waitDone(st[0]); + spinDone(st[0]); } // ---- timed region: ONE whole-loop timer, nixl-style ---- int completed = 0, issued = 0; auto t0 = Clock::now(); - // prime the pipeline + // prime the pipeline. NOTE: all `depth` in-flight slots reuse the same buffer + // window (non-batched posts all target offset 0; batched slots share the same + // offsets vector), so with --inflight > 1 concurrent transfers overlap the + // same bytes. That's fine for a throughput/latency benchmark (data is not + // validated), matching nixl which likewise reuses its IOV across slots. for (int s = 0; s < depth; ++s) { post(s); ++issued; } while (completed < a.iters) { for (int s = 0; s < depth; ++s) { - // spin-poll this slot (WaitBusy with a single poll would also work, but - // we mirror nixl's "check status, if IN_PROG continue" scan) + // spin-poll this slot, mirroring nixl's "check status, if IN_PROG + // continue" scan (status flag is stored by MORI's CQ worker thread) if (st[s].InProgress() || st[s].Init()) continue; if (st[s].Failed()) { std::cerr << "transfer failed: " << st[s].Message() << std::endl; @@ -459,12 +474,17 @@ int main(int argc, char** argv) { double total_us = std::chrono::duration_cast>(t1 - t0).count(); - // Each of the `iters` requests moves msg*batch bytes when batching (nixl - // counts block_size*batch_size*num_iter). avg_lat is per REQUEST. + // nixl parity (nixl_worker/utils.cpp): count block_size*batch_size*num_iter + // bytes over ONE whole-loop timer, and report latency PER SINGLE TRANSFER: + // total_bytes = msg * batch * iters + // avg_bw = total_bytes/1e9 / (total_us/1e6) [GB/s, GB=10^9] + // avg_latency = total_us / (iters * batch) [us per transfer] + // matching nixl's avg_latency = total_duration/(per_thread_iter*batch_size). const int effBatch = batched ? curBatch : 1; - double total_bytes = static_cast(msg) * effBatch * a.iters; + const double numXfers = static_cast(a.iters) * effBatch; + double total_bytes = static_cast(msg) * numXfers; double avg_bw = (total_bytes / 1e9) / (total_us / 1e6); // GB/s, GB=10^9 - double avg_lat = total_us / a.iters; // us per request + double avg_lat = total_us / numXfers; // us per single transfer std::printf("%-11zu %-6d %-6d %-12.2f %-11.2f %-.1f\n", msg, effBatch, a.iters, avg_bw, avg_lat, total_us); From c755fd822e33a691873f6d2a2119ccf35aa20744 Mon Sep 17 00:00:00 2001 From: Akbarzadeh Date: Tue, 28 Jul 2026 15:35:51 -0500 Subject: [PATCH 07/10] bench(io): reintroduce --enable-batch-transfer + document C++ bench Restore the Python bench's batch/single submission toggle in the C++ nixl-style engine bench, and add C++ run docs to MORI-IO-BENCHMARK.md. - --enable-batch-transfer / --disable-batch-transfer (default ON): - ON => one N-descriptor batch request (BatchWrite/BatchRead), the nixl-equivalent path (nixl always batches). - OFF => N individual single-transfer submissions per iteration at contiguous offsets i*msg, matching Python's run_single_once. Default is ON (vs Python's OFF) so the out-of-the-box run and the --all-batch sweep stay nixl-comparable; this also fixes the old sweep no-op (OFF now scales with N instead of doing one msg-byte write). - Strict stop-and-wait (nixl --pipeline_depth 1): one request outstanding at a time. perSlot statuses live in a flat status array (TransferStatus is non-copyable); a request completes only when all its sub-transfers finish, mirroring Python waiting on the whole status_list. - Metrics unchanged/consistent across modes: both move curBatch transfers of msg per iter, so total_bytes = msg*batch*iters and per-transfer latency = total_us/(iters*batch). - docs/MORI-IO-BENCHMARK.md: new "C++ Benchmark (nixlbench-matching)" section (build, 2-node run, sweep modes, Python->C++ flag mapping) and updated the --enable-batch-transfer row to describe ON vs OFF. Co-Authored-By: Claude Opus 4 --- docs/MORI-IO-BENCHMARK.md | 103 ++++++++++++++++++++++ tests/cpp/io/bench_engine.cpp | 158 +++++++++++++++++++--------------- 2 files changed, 190 insertions(+), 71 deletions(-) diff --git a/docs/MORI-IO-BENCHMARK.md b/docs/MORI-IO-BENCHMARK.md index 7d48d7be2..75f71a28a 100644 --- a/docs/MORI-IO-BENCHMARK.md +++ b/docs/MORI-IO-BENCHMARK.md @@ -3,6 +3,7 @@ ## Table of Contents - [Benchmark Commands](#benchmark-commands) +- [C++ Benchmark (nixlbench-matching)](#c-benchmark-nixlbench-matching) - [Benchmark Arguments](#benchmark-arguments) - [Results: Thor2 RDMA Read](#results-thor2-rdma-read) - [Results: Thor2 RDMA Write](#results-thor2-rdma-write) @@ -48,6 +49,108 @@ torchrun --nnodes=2 --node_rank=0 --nproc_per_node=1 \ > (`--mem-type gpu`). If the two nodes are not in the same vPOD, session creation > fails fast with a clear error. +## C++ Benchmark (nixlbench-matching) + +The commands above use the **Python** benchmark (`tests/python/io/benchmark.py`). +There is also a native **C++** benchmark, `tests/cpp/io/bench_engine.cpp`, whose +measurement loop is written to **match [nixlbench](https://github.com/ai-dynamo/nixl)** +(NVIDIA NIXL's `xferbench`) exactly, so MORI-IO and NIXL RDMA numbers are +apples-to-apples on the same fabric. Use it when you want a head-to-head +comparison against NIXL, or a benchmark with zero Python-interpreter overhead +(which biases the Python numbers by ~2 µs/iter at small message sizes). + +What "nixl-matching" means here: + +- **One whole-loop timer** around the entire iteration loop (not a per-iteration + sum), identical to nixlbench's `total_timer`. +- **Latency is reported per single transfer**: `total_us / (iters × batch)`, + matching nixl's `avg_latency = total_duration / (per_thread_iter × batch_size)`. +- **Bandwidth** `= msg × batch × iters / 1e9 / (total_us / 1e6)` (GB = 10⁹), same + units as nixl. +- **Completion is always an inline spin-poll** in the benchmark thread (mirroring + nixl's `getXferStatus` scan); there is no cv-blocking wait that would add wakeup + latency and make numbers non-comparable. +- **Strict stop-and-wait** (one request outstanding at a time, == nixl + `--pipeline_depth 1`): post one request, spin until it completes, then post the + next. +- **Warmup** iterations are excluded from timing (`--warmup-iters`, nixl + `--warmup_iter`). + +### Build + +`bench_engine` builds with the C++ tests (both options default `ON`): + +```bash +cd /path/to/mori +cmake -B build -DBUILD_IO=ON -DBUILD_TESTS=ON +cmake --build build --target bench_engine -j +# binary: build/tests/cpp/bench_engine +``` + +The MORI shared libraries must be discoverable at runtime. After an editable +install they live in `python/mori`: + +```bash +export LD_LIBRARY_PATH=/path/to/mori/python/mori:$LD_LIBRARY_PATH +``` + +### Run (two nodes) + +Unlike the Python benchmark (which uses `torchrun` for the out-of-band +rendezvous), the C++ benchmark does its own TCP rendezvous between two +processes: **rank 1 is the target and must be started first**, then **rank 0 is +the initiator** and drives all transfers. `--master-ip` is rank 0's IP; each +process passes its own peer-reachable data-plane IP via `--self-ip`. + +```bash +# On the TARGET node (rank 1) — start this FIRST: +MORI_DISABLE_AUTO_XGMI=1 build/tests/cpp/bench_engine \ + --rank 1 --master-ip 10.190.162.0 --self-ip 10.190.162.1 --port 18515 \ + --op write --all --sweep-start 1048576 --sweep-max 33554432 --sweep-step 1048576 \ + --iters 500 --warmup-iters 50 --num-qp-per-transfer 4 --enable-sess + +# On the INITIATOR node (rank 0) — start this SECOND (prints the results table): +MORI_DISABLE_AUTO_XGMI=1 build/tests/cpp/bench_engine \ + --rank 0 --master-ip 10.190.162.0 --self-ip 10.190.162.0 --port 18515 \ + --op write --all --sweep-start 1048576 --sweep-max 33554432 --sweep-step 1048576 \ + --iters 500 --warmup-iters 50 --num-qp-per-transfer 4 --enable-sess +``` + +Both ranks must be given the **same** benchmark args (op, sweep, batch, etc.). +Only rank 0 (initiator) prints the results table: + +``` +MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us) +``` + +### Sweep modes + +- `--all` — sweep message size. Geometric ×2 by default, or linear when + `--sweep-step > 0`, from `--sweep-start` to `--sweep-max`; batch stays at + `--transfer-batch-size`. +- `--all-batch` — fix message size at `--buffer-size` and sweep batch + `1,2,4,…,32768`. +- neither — a single `(--buffer-size, --transfer-batch-size)` point. + +### Differences from the Python flags + +Most flags share names with the Python benchmark (`--op-type`/`--op`, +`--num-qp-per-transfer`, `--num-worker-threads`, `--transfer-batch-size`, +`--enable-sess`, `--batch-contiguous`, `--disable-chunking`, `--chunk-bytes`, +`--poll_cq_mode`, `--mem-type`, `--all`, `--all-batch`), but note: + +| Python | C++ (nixl-style) | Notes | +|--------|------------------|-------| +| `torchrun ... --node_rank` | `--rank 0` / `--rank 1` | C++ uses its own TCP rendezvous; **start rank 1 first**. | +| `--host` | `--master-ip` + `--self-ip` | `--master-ip` = rank 0's IP; `--self-ip` = each node's peer-reachable IP. | +| `--enable-batch-transfer` (default OFF) | `--enable-batch-transfer` / `--disable-batch-transfer` (C++ default **ON**) | Same meaning as Python: **ON** = one N-descriptor batch request (the nixl-equivalent; nixl always batches); **OFF** = `--transfer-batch-size` N *individual* single-transfer submissions per iteration (Python `run_single_once`). The C++ tool defaults **ON** so the out-of-the-box run and `--all-batch` sweep are nixl-comparable; pass `--disable-batch-transfer` for the Python single-submission path. | +| *(n/a)* | `--warmup-iters` / `--iters` | Explicit warmup (untimed) and timed iteration counts. | +| `--sweep-start-size` / `--sweep-max-size` | `--sweep-start` / `--sweep-max` (aliases accepted) | Plus `--sweep-step` for a linear (vs geometric) sweep. | + +> Keep the process on the RDMA path with `MORI_DISABLE_AUTO_XGMI=1`. Use +> `--self-ip` equal to the data-plane IP each node advertises for QP setup; +> `0.0.0.0` will fail with `connect(...): Connection refused`. + ## Benchmark Arguments | Argument | Description | diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp index f60276065..7561c42bc 100644 --- a/tests/cpp/io/bench_engine.cpp +++ b/tests/cpp/io/bench_engine.cpp @@ -21,14 +21,9 @@ // numbers non-comparable). MORI's blocking Wait() path is exercised only by // the Python benchmark, not this nixl-parity tool. // -// 3. PIPELINE DEPTH. -// nixlbench keeps `pipeline_depth` transfers in flight (default 1). We -// expose --inflight to match: default 1 = strict stop-and-wait, matching -// nixl --pipeline_depth 1. +// 3. WARMUP excluded from timing (nixl: --warmup_iter, default 100). // -// 4. WARMUP excluded from timing (nixl: --warmup_iter, default 100). -// -// 5. throughput_gb = total_bytes / 1e9 / (total_duration_us / 1e6), GB=10^9, +// 4. throughput_gb = total_bytes / 1e9 / (total_duration_us / 1e6), GB=10^9, // identical unit to nixl. total_bytes = msg * batch * num_iter, and // avg_latency = total_duration / (num_iter * batch) -- i.e. PER SINGLE // TRANSFER, matching nixl's total_duration/(per_thread_iter*batch_size) @@ -180,7 +175,6 @@ struct Args { size_t sweep_step = 0; // --sweep-step (0 = geometric x2; >0 = linear +step) int iters = 500; int warmup = 50; // --warmup-iters - int inflight = 1; // == nixl pipeline_depth (transfers/requests in flight) // RdmaBackendConfig knobs int qp_per_transfer = 4; // --num-qp-per-transfer @@ -196,6 +190,10 @@ struct Args { // batch / session int batch = 1; // --transfer-batch-size (transfers per request) + bool enable_batch_transfer = true; // --enable-batch-transfer / --disable-batch-transfer. + // ON (default): batch>1 => ONE N-descriptor batch request + // (nixl-equivalent). OFF: batch>1 => N individual single + // transfers per iteration (Python run_single_once path). bool batch_contiguous = false; // --batch-contiguous (adjacent offsets → merged WR); // default strided (each transfer a separate WR) bool enable_sess = false; // --enable-sess (session fast-path); Python default: off @@ -227,7 +225,6 @@ static Args ParseArgs(int argc, char** argv) { else if (k == "--sweep-step") a.sweep_step = std::stoull(next()); else if (k == "--iters") a.iters = std::stoi(next()); else if (k == "--warmup" || k == "--warmup-iters") a.warmup = std::stoi(next()); - else if (k == "--inflight") a.inflight = std::stoi(next()); else if (k == "--qp-per-transfer" || k == "--num-qp-per-transfer") a.qp_per_transfer = std::stoi(next()); else if (k == "--worker-threads" || k == "--num-worker-threads") a.worker_threads = std::stoi(next()); else if (k == "--post-batch-size") a.post_batch_size = std::stoi(next()); @@ -239,6 +236,8 @@ static Args ParseArgs(int argc, char** argv) { else if (k == "--max-cqe-num") a.max_cqe_num = std::stoi(next()); else if (k == "--max-msg-sge") a.max_msg_sge = std::stoi(next()); else if (k == "--batch" || k == "--transfer-batch-size") a.batch = std::stoi(next()); + else if (k == "--enable-batch-transfer") a.enable_batch_transfer = true; + else if (k == "--disable-batch-transfer") a.enable_batch_transfer = false; else if (k == "--batch-contiguous") a.batch_contiguous = true; else if (k == "--enable-sess") a.enable_sess = true; else if (k == "--disable-sess") a.enable_sess = false; @@ -377,34 +376,29 @@ int main(int argc, char** argv) { } const bool isRead = (a.op == "read"); - // Spin to completion, same mechanism as the timed loop (status flag stored by - // MORI's CQ worker thread). Used for warmup; keeps warmup identical to the - // measured path so caches/QP state are primed the same way. - auto spinDone = [&](TransferStatus& s) { - while (s.InProgress() || s.Init()) { /* spin: CQ worker thread stores status */ } - if (s.Failed()) { - std::cerr << "warmup transfer failed: " << s.Message() << std::endl; - std::exit(1); - } - }; - std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; for (auto& planEntry : plan) { const size_t msg = planEntry.first; const int curBatch = planEntry.second; - // nixl-style: batch>1 always issues one N-descriptor batch request (no - // separate enable flag). batch==1 is the degenerate single-transfer case. - const bool batched = curBatch > 1; - const int depth = std::min(a.inflight, a.iters); - std::vector st(depth); - std::vector uid(depth); - - // Batch layout. Strided (default): slot i at (msg+1)*i so the N transfers - // stay SEPARATE WRs (stresses SQ, real batching) — matches Python default. - // Contiguous (--batch-contiguous): slot i at msg*i, adjacent, so MORI may - // merge them into one big WR (fast but not really batching, and hits the - // 1 GiB max_msg_sz without chunking). + // Two ways to move curBatch transfers per iteration, mirroring the Python bench: + // --enable-batch-transfer (default, batched): ONE N-descriptor batch request + // (BatchWrite/BatchRead) -- the nixl-equivalent (nixl always batches). + // --disable-batch-transfer (singles): curBatch INDIVIDUAL single-transfer + // submissions per iteration (Python run_single_once), each its own status. + // batch==1 is a single transfer either way. + const bool batched = a.enable_batch_transfer && curBatch > 1; + const int perSlot = batched ? 1 : curBatch; // statuses per request + // Strict stop-and-wait: one request outstanding at a time (nixl + // --pipeline_depth 1). [perSlot] status array (TransferStatus is + // non-copyable, so a flat vector rather than nested). + std::vector st(static_cast(perSlot)); + + // Batch layout for the batched path. Strided (default): transfer i at + // (msg+1)*i so the N transfers stay SEPARATE WRs (stresses SQ, real batching) + // -- matches Python default. Contiguous (--batch-contiguous): transfer i at + // msg*i, adjacent, so MORI may merge them into one big WR (fast but not really + // batching, and hits the 1 GiB max_msg_sz without chunking). const size_t stride = a.batch_contiguous ? msg : (msg + 1); SizeVec offsets(curBatch), sizes(curBatch); for (int i = 0; i < curBatch; ++i) { @@ -415,59 +409,80 @@ int main(int argc, char** argv) { MemDescVec locVec{localMem}, remVec{peerMem}; BatchSizeVec offVec{offsets}, sizeVec{sizes}; - auto post = [&](int slot) { - st[slot].SetCode(StatusCode::INIT); - uid[slot] = useSess ? sessPtr->AllocateTransferUniqueId() : engine.AllocateTransferUniqueId(); + auto post = [&]() { + TransferStatus* base = &st[0]; + for (int i = 0; i < perSlot; ++i) base[i].SetCode(StatusCode::INIT); if (batched) { + // ONE N-descriptor batch request (nixl / Python --enable-batch-transfer). + TransferUniqueId id = + useSess ? sessPtr->AllocateTransferUniqueId() : engine.AllocateTransferUniqueId(); if (useSess) { - if (isRead) sessPtr->BatchRead(offsets, offsets, sizes, &st[slot], uid[slot]); - else sessPtr->BatchWrite(offsets, offsets, sizes, &st[slot], uid[slot]); + if (isRead) sessPtr->BatchRead(offsets, offsets, sizes, &base[0], id); + else sessPtr->BatchWrite(offsets, offsets, sizes, &base[0], id); } else { - TransferStatusPtrVec sp{&st[slot]}; - TransferUniqueIdVec ids{uid[slot]}; + TransferStatusPtrVec sp{&base[0]}; + TransferUniqueIdVec ids{id}; if (isRead) engine.BatchRead(locVec, offVec, remVec, offVec, sizeVec, sp, ids); else engine.BatchWrite(locVec, offVec, remVec, offVec, sizeVec, sp, ids); } } else { - if (useSess) { - if (isRead) sessPtr->Read(0, 0, msg, &st[slot], uid[slot]); - else sessPtr->Write(0, 0, msg, &st[slot], uid[slot]); - } else { - if (isRead) engine.Read(localMem, 0, peerMem, 0, msg, &st[slot], uid[slot]); - else engine.Write(localMem, 0, peerMem, 0, msg, &st[slot], uid[slot]); + // curBatch individual single-transfer submissions (Python run_single_once); + // contiguous offsets i*msg, matching Python's single path. curBatch==1 is a + // single transfer at offset 0. + for (int i = 0; i < curBatch; ++i) { + const size_t off = static_cast(i) * msg; + TransferUniqueId id = + useSess ? sessPtr->AllocateTransferUniqueId() : engine.AllocateTransferUniqueId(); + if (useSess) { + if (isRead) sessPtr->Read(off, off, msg, &base[i], id); + else sessPtr->Write(off, off, msg, &base[i], id); + } else { + if (isRead) engine.Read(localMem, off, peerMem, off, msg, &base[i], id); + else engine.Write(localMem, off, peerMem, off, msg, &base[i], id); + } } } }; + // The request is complete only when ALL perSlot sub-transfers have left + // INIT/IN_PROGRESS (mirrors Python waiting on the whole status_list). Spin, + // like the timed loop. reqFailed returns the first failed status, or nullptr. + auto reqDone = [&]() { + TransferStatus* base = &st[0]; + for (int i = 0; i < perSlot; ++i) + if (base[i].InProgress() || base[i].Init()) return false; + return true; + }; + auto reqFailed = [&]() -> TransferStatus* { + TransferStatus* base = &st[0]; + for (int i = 0; i < perSlot; ++i) + if (base[i].Failed()) return &base[i]; + return nullptr; + }; + // ---- warmup (excluded from timing) ---- for (int w = 0; w < a.warmup; ++w) { - post(0); - spinDone(st[0]); + post(); + while (!reqDone()) { /* spin: CQ worker thread stores status */ } + if (TransferStatus* f = reqFailed()) { + std::cerr << "warmup transfer failed: " << f->Message() << std::endl; + std::exit(1); + } } // ---- timed region: ONE whole-loop timer, nixl-style ---- - int completed = 0, issued = 0; + // Strict stop-and-wait (nixl --pipeline_depth 1): post one request, spin + // until it completes, then post the next. auto t0 = Clock::now(); - // prime the pipeline. NOTE: all `depth` in-flight slots reuse the same buffer - // window (non-batched posts all target offset 0; batched slots share the same - // offsets vector), so with --inflight > 1 concurrent transfers overlap the - // same bytes. That's fine for a throughput/latency benchmark (data is not - // validated), matching nixl which likewise reuses its IOV across slots. - for (int s = 0; s < depth; ++s) { post(s); ++issued; } - - while (completed < a.iters) { - for (int s = 0; s < depth; ++s) { - // spin-poll this slot, mirroring nixl's "check status, if IN_PROG - // continue" scan (status flag is stored by MORI's CQ worker thread) - if (st[s].InProgress() || st[s].Init()) continue; - if (st[s].Failed()) { - std::cerr << "transfer failed: " << st[s].Message() << std::endl; - std::exit(1); - } - // completed one - ++completed; - if (issued < a.iters) { post(s); ++issued; } + for (int completed = 0; completed < a.iters; ++completed) { + post(); + // spin-poll, mirroring nixl's "check status, if IN_PROG continue" scan + // (status flags stored by MORI's CQ worker thread) + while (!reqDone()) { /* spin */ } + if (TransferStatus* f = reqFailed()) { + std::cerr << "transfer failed: " << f->Message() << std::endl; + std::exit(1); } } auto t1 = Clock::now(); @@ -475,12 +490,13 @@ int main(int argc, char** argv) { double total_us = std::chrono::duration_cast>(t1 - t0).count(); // nixl parity (nixl_worker/utils.cpp): count block_size*batch_size*num_iter - // bytes over ONE whole-loop timer, and report latency PER SINGLE TRANSFER: - // total_bytes = msg * batch * iters + // bytes over ONE whole-loop timer, and report latency PER SINGLE TRANSFER. + // Both batched and singles modes move curBatch transfers of `msg` per iter: + // total_bytes = msg * curBatch * iters // avg_bw = total_bytes/1e9 / (total_us/1e6) [GB/s, GB=10^9] - // avg_latency = total_us / (iters * batch) [us per transfer] + // avg_latency = total_us / (iters * curBatch) [us per transfer] // matching nixl's avg_latency = total_duration/(per_thread_iter*batch_size). - const int effBatch = batched ? curBatch : 1; + const int effBatch = curBatch; const double numXfers = static_cast(a.iters) * effBatch; double total_bytes = static_cast(msg) * numXfers; double avg_bw = (total_bytes / 1e9) / (total_us / 1e6); // GB/s, GB=10^9 From e0bfcfe29455f72e6bf7ccf69283fbbd92221db9 Mon Sep 17 00:00:00 2001 From: Akbarzadeh Date: Thu, 30 Jul 2026 13:28:55 -0500 Subject: [PATCH 08/10] io(rdma): add prepared transfer path to C++ and Python benchmarks Split RDMA batch transfers into a reusable build phase and a re-postable submission phase, preserving the original inline path by default. Expose the prepared session transfer through pybind and Python, then add --prepare-once to both C++ and Python benchmarks so they can remove descriptor-build work from the measured loop for NIXL-style comparisons. Co-Authored-By: Claude Opus 4 --- include/mori/io/backend.hpp | 30 +++++ include/mori/io/engine.hpp | 11 ++ python/mori/io/engine.py | 10 ++ src/io/engine.cpp | 19 +++ src/io/rdma/backend_impl.cpp | 66 ++++++++++ src/io/rdma/backend_impl.hpp | 6 + src/io/rdma/common.cpp | 240 +++++++++++++++++++++++----------- src/io/rdma/common.hpp | 69 ++++++++++ src/pybind/pybind_io.cpp | 14 ++ tests/cpp/io/bench_engine.cpp | 49 +++++++ tests/python/io/benchmark.py | 114 +++++++++++++++- 11 files changed, 551 insertions(+), 77 deletions(-) diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index 2503af73d..789059446 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -22,6 +22,7 @@ #pragma once #include +#include #include "mori/io/common.hpp" #include "mori/io/enum.hpp" @@ -114,6 +115,19 @@ inline std::ostream& operator<<(std::ostream& os, const FabricBackendConfig& c) return os << "numStreams[" << c.numStreams << "] numEvents[" << c.numEvents << "]"; } +/* ---------------------------------------------------------------------------------------------- */ +/* PreparedTransfer */ +/* ---------------------------------------------------------------------------------------------- */ +// Opaque, backend-owned handle for a transfer whose work requests were built +// once (sort/merge/chunk) and can be re-posted many times. Lets callers hoist +// the descriptor-build cost out of a hot transfer loop (the NIXL +// createXferReq/postXferReq split). Backends that do not support prepared +// transfers simply return nullptr from PrepareBatch. +struct PreparedTransfer { + PreparedTransfer() = default; + virtual ~PreparedTransfer() = default; +}; + /* ---------------------------------------------------------------------------------------------- */ /* BackendSession */ /* ---------------------------------------------------------------------------------------------- */ @@ -144,6 +158,22 @@ class BackendSession { const SizeVec& sizes, TransferStatus* status, TransferUniqueId id) { BatchReadWrite(localOffsets, remoteOffsets, sizes, status, id, true); } + + // Prepared (build-once, post-many) transfer API. PrepareBatch builds the work + // requests for the given batch and returns a reusable handle; PostPrepared + // submits it with a fresh status/id. The default implementation returns + // nullptr / reports unsupported, so backends may opt in. Offsets/sizes are + // captured at prepare time and must not change between posts. + virtual std::shared_ptr PrepareBatch(const SizeVec& localOffsets, + const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead) { + return nullptr; + } + virtual void PostPrepared(const std::shared_ptr& prepared, + TransferStatus* status, TransferUniqueId id) { + status->Update(StatusCode::ERR_BAD_STATE, "prepared transfer not supported by this backend"); + } + virtual bool Alive() const = 0; }; diff --git a/include/mori/io/engine.hpp b/include/mori/io/engine.hpp index 5155edec5..68de1f6ae 100644 --- a/include/mori/io/engine.hpp +++ b/include/mori/io/engine.hpp @@ -62,6 +62,17 @@ class IOEngineSession { TransferStatus* status, TransferUniqueId id); void BatchWrite(const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, TransferStatus* status, TransferUniqueId id); + + // Prepared (build-once, post-many) transfer API. PrepareBatch builds the work + // requests once and returns a reusable handle (nullptr if the backend does not + // support it); PostPrepared re-posts it with a fresh status/id. Offsets/sizes + // are captured at prepare time and must not change between posts. + std::shared_ptr PrepareBatch(const SizeVec& localOffsets, + const SizeVec& remoteOffsets, const SizeVec& sizes, + bool isRead); + void PostPrepared(const std::shared_ptr& prepared, TransferStatus* status, + TransferUniqueId id); + bool Alive(); friend class IOEngine; diff --git a/python/mori/io/engine.py b/python/mori/io/engine.py index e6cc76cdf..e7b8e2793 100644 --- a/python/mori/io/engine.py +++ b/python/mori/io/engine.py @@ -79,6 +79,16 @@ def batch_read(self, *args): def batch_write(self, *args): return self._batch_single_side_transfer(self._sess.BatchWrite, *args) + def prepare_batch(self, local_offsets, remote_offsets, sizes, is_read): + # Build the work requests once; returns a reusable handle (None if the + # backend does not support prepared transfers). + return self._sess.PrepareBatch(local_offsets, remote_offsets, sizes, is_read) + + def post_prepared(self, prepared, transfer_uid): + transfer_status = mori_cpp.TransferStatus() + self._sess.PostPrepared(prepared, transfer_status, transfer_uid) + return transfer_status + def alive(self): return self._sess.Alive() diff --git a/src/io/engine.cpp b/src/io/engine.cpp index a22a97b20..ca5ef26b2 100644 --- a/src/io/engine.cpp +++ b/src/io/engine.cpp @@ -160,6 +160,25 @@ void IOEngineSession::BatchWrite(const SizeVec& localOffsets, const SizeVec& rem } } +std::shared_ptr IOEngineSession::PrepareBatch(const SizeVec& localOffsets, + const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead) { + MORI_IO_FUNCTION_TIMER; + return backendSess->PrepareBatch(localOffsets, remoteOffsets, sizes, isRead); +} + +void IOEngineSession::PostPrepared(const std::shared_ptr& prepared, + TransferStatus* status, TransferUniqueId id) { + MORI_IO_FUNCTION_TIMER; + std::shared_ptr diagnostics; + internal::ScopedIoCallDiagnosticsCapture capture(&diagnostics, "Session post prepared"); + backendSess->PostPrepared(prepared, status, id); + if (status->Failed()) { + LogTransferFailure(diagnostics, "Session post prepared error {} message {}", + status->CodeUint32(), status->Message()); + } +} + bool IOEngineSession::Alive() { return backendSess->Alive(); } /* ---------------------------------------------------------------------------------------------- */ diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp index 0ba8a2b40..7de016ed3 100644 --- a/src/io/rdma/backend_impl.cpp +++ b/src/io/rdma/backend_impl.cpp @@ -1443,6 +1443,72 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV } } +// Concrete prepared-transfer handle for the RDMA backend: owns the built +// (sorted/merged/chunked) work-request list so PostPrepared can re-post it. +struct RdmaPreparedTransfer : public PreparedTransfer { + PreparedRdmaBatch batch; +}; + +std::shared_ptr RdmaBackendSession::PrepareBatch(const SizeVec& localOffsets, + const SizeVec& remoteOffsets, + const SizeVec& sizes, + bool isRead) { + MORI_IO_FUNCTION_TIMER; + if (eps.empty() || localMrPerEp.empty() || remoteMrPerEp.empty()) { + MORI_IO_WARN("PrepareBatch: session has no endpoints/memory regions"); + return nullptr; + } + // The prepared path uses the inline (non-executor) posting route, mirroring + // the RdmaBatchReadWrite `else` branch control setup. + const bool chunk = config.enableTransferChunking; + RdmaTransferControl control{}; + control.chunkBytes = chunk ? config.chunkBytes : 0; + control.maxChunks = config.maxChunksPerTransfer; + control.creditByWrCount = chunk; + + auto handle = std::make_shared(); + RdmaOpRet ret = + BuildPreparedRdmaBatch(eps, localMrPerEp.front(), remoteMrPerEp.front(), localOffsets, + remoteOffsets, sizes, isRead, control, config.postBatchSize, + handle->batch); + if (ret.Failed()) { + MORI_IO_WARN("PrepareBatch failed to build work requests: {}", ret.message); + return nullptr; + } + return handle; +} + +void RdmaBackendSession::PostPrepared(const std::shared_ptr& prepared, + TransferStatus* status, TransferUniqueId id) { + MORI_IO_FUNCTION_TIMER; + status->SetCode(StatusCode::IN_PROGRESS); + + auto* handle = dynamic_cast(prepared.get()); + if (handle == nullptr || !handle->batch.valid) { + status->Update(StatusCode::ERR_INVALID_ARGS, "invalid prepared transfer handle"); + return; + } + + // When !creditByWrCount, totalBatchSize is the original request count; the + // chunked path overwrites it inside PostMergedWorkRequests (ownsTotalBatchSize). + const int totalCredit = static_cast(handle->batch.batchSize); + auto callbackMeta = std::make_shared(status, id, totalCredit); + internal::PublishCurrentIoCallDiagnostics(callbackMeta); + + RdmaOpRet ret = + PostPreparedRdmaBatch(eps, localMrPerEp, remoteMrPerEp, handle->batch, callbackMeta, id); + assert(!ret.Init()); + if (ret.Failed() || ret.Succeeded()) { + status->Update(ret.code, ret.message); + } + if (!ret.Failed() && config.enableNotification) { + RdmaOpRet notifRet = RdmaNotifyTransfer(eps, status, id); + if (notifRet.Failed()) { + status->Update(notifRet.code, notifRet.message); + } + } +} + bool RdmaBackendSession::Alive() const { return true; } /* ---------------------------------------------------------------------------------------------- diff --git a/src/io/rdma/backend_impl.hpp b/src/io/rdma/backend_impl.hpp index 8f13acebd..98504c599 100644 --- a/src/io/rdma/backend_impl.hpp +++ b/src/io/rdma/backend_impl.hpp @@ -276,6 +276,12 @@ class RdmaBackendSession : public BackendSession { const SizeVec& sizes, TransferStatus* status, TransferUniqueId id, bool isRead); + std::shared_ptr PrepareBatch(const SizeVec& localOffsets, + const SizeVec& remoteOffsets, const SizeVec& sizes, + bool isRead) override; + void PostPrepared(const std::shared_ptr& prepared, TransferStatus* status, + TransferUniqueId id) override; + bool Alive() const; private: diff --git a/src/io/rdma/common.cpp b/src/io/rdma/common.cpp index 8e4e7f4cf..67d0d66b4 100644 --- a/src/io/rdma/common.cpp +++ b/src/io/rdma/common.cpp @@ -530,13 +530,6 @@ void PlanSgeStreamChunks(std::vector& plan, const std::vector } } -struct MergedWorkRequest { - ibv_send_wr wr{}; - std::vector sges; - size_t totalRemoteLength = 0; - size_t mergedRequests = 1; -}; - static void ResetMergedWorkRequestPointers(MergedWorkRequest* wr) { if (wr == nullptr) return; wr->wr.sg_list = wr->sges.empty() ? nullptr : wr->sges.data(); @@ -598,14 +591,19 @@ RdmaOpRet RdmaNotifyTransfer(const EpPairVec& eps, TransferStatus* status, Trans return {StatusCode::IN_PROGRESS, ""}; } -RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, - const std::vector& localMrPerEp, - const std::vector& remoteMrPerEp, - const SizeVec& localOffsets, const SizeVec& remoteOffsets, - const SizeVec& sizes, std::shared_ptr callbackMeta, - TransferUniqueId id, bool isRead, int postBatchSize, - const RdmaTransferControl& control) { - MORI_IO_FUNCTION_TIMER; +RdmaOpRet BuildMergedWorkRequests(const EpPairVec& eps, + const application::RdmaMemoryRegion& baseLocalMr, + const application::RdmaMemoryRegion& baseRemoteMr, + const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead, + const RdmaTransferControl& control, + std::vector& indices, + std::vector& mergedPool, + std::vector& chunkedPool, + std::vector& chunkPlan, + std::vector*& outWrs, size_t& outCount) { + outWrs = &mergedPool; + outCount = 0; if ((localOffsets.size() != remoteOffsets.size()) || (sizes.size() != remoteOffsets.size())) { return {StatusCode::ERR_INVALID_ARGS, @@ -621,16 +619,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, return {StatusCode::ERR_INVALID_ARGS, "no endpoints"}; } - if (localMrPerEp.size() != eps.size() || remoteMrPerEp.size() != eps.size()) { - return {StatusCode::ERR_INVALID_ARGS, "memory-region vectors must align with endpoints"}; - } - if (control.maxChunks <= 0) { return {StatusCode::ERR_INVALID_ARGS, "maxChunks must be >= 1"}; } - const application::RdmaMemoryRegion& baseLocalMr = localMrPerEp.front(); - const application::RdmaMemoryRegion& baseRemoteMr = remoteMrPerEp.front(); for (size_t i = 0; i < batchSize; i++) { if (sizes[i] > std::numeric_limits::max()) { return {StatusCode::ERR_INVALID_ARGS, "single request size " + std::to_string(sizes[i]) + @@ -642,57 +634,6 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, } } - // [tls-scratch] Per worker-thread scratch pools eliminate per-batch heap - // allocations on the RDMA hot path (slot reuse / resize() / clear() retain - // capacity across calls). The pools belong to the OUTERMOST call on a thread; - // if RdmaBatchReadWrite is ever re-entered on the same thread (e.g. from a - // completion callback), the nested call transparently falls back to local - // buffers so the outer call's pools are never clobbered. - thread_local std::vector tlIndices; - thread_local std::vector tlMergedPool; - thread_local std::vector tlChunkedPool; - thread_local std::vector tlChunkPlan; - thread_local std::vector tlEpWrsSinceSignal; - thread_local std::vector tlEpMergedSinceSignal; - thread_local int reentryDepth = 0; - - struct ReentryGuard { - int& depth; - explicit ReentryGuard(int& d) : depth(d) { ++depth; } - ~ReentryGuard() { --depth; } - } reentryGuard(reentryDepth); - const bool usePool = (reentryDepth == 1); - - // Used only on (currently non-existent) same-thread re-entry. - std::vector localIndices; - std::vector localMergedPool; - std::vector localChunkedPool; - std::vector localChunkPlan; - std::vector localEpWrsSinceSignal; - std::vector localEpMergedSinceSignal; - - std::vector& indices = usePool ? tlIndices : localIndices; - std::vector& mergedPool = usePool ? tlMergedPool : localMergedPool; - std::vector& chunkedPool = usePool ? tlChunkedPool : localChunkedPool; - std::vector& chunkPlan = usePool ? tlChunkPlan : localChunkPlan; - std::vector& epWrsSinceSignal = usePool ? tlEpWrsSinceSignal : localEpWrsSinceSignal; - std::vector& epMergedSinceSignal = - usePool ? tlEpMergedSinceSignal : localEpMergedSinceSignal; - - // Bound peak retained memory: if an earlier very large batch grew the pools far - // beyond the current need, release the excess so it doesn't stay resident. - constexpr size_t kPoolHighWater = 8192; - if (usePool && batchSize <= kPoolHighWater / 2) { - if (mergedPool.size() > kPoolHighWater) { - mergedPool.resize(kPoolHighWater); - mergedPool.shrink_to_fit(); - } - if (chunkedPool.size() > kPoolHighWater) { - chunkedPool.resize(kPoolHighWater); - chunkedPool.shrink_to_fit(); - } - } - indices.resize(batchSize); std::iota(indices.begin(), indices.end(), 0); @@ -849,9 +790,17 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, } } - std::vector& mergedWrs = useChunked ? chunkedPool : mergedPool; - size_t mergedWrCount = useChunked ? chunkedCount : wrCount; + outWrs = useChunked ? &chunkedPool : &mergedPool; + outCount = useChunked ? chunkedCount : wrCount; + return {StatusCode::SUCCESS, ""}; +} +RdmaOpRet PostMergedWorkRequests(const EpPairVec& eps, + const std::vector& localMrPerEp, + const std::vector& remoteMrPerEp, + std::vector& mergedWrs, size_t mergedWrCount, + std::shared_ptr callbackMeta, TransferUniqueId id, + int postBatchSize, const RdmaTransferControl& control) { if (control.creditByWrCount) { if (mergedWrCount > static_cast(std::numeric_limits::max())) { return {StatusCode::ERR_INVALID_ARGS, "final WR count exceeds int range"}; @@ -881,8 +830,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, if (postBatchSize <= 0) postBatchSize = 1; int numPostBatch = (mergedWrCount + postBatchSize - 1) / postBatchSize; - epWrsSinceSignal.assign(epNum, 0); - epMergedSinceSignal.assign(epNum, 0); + // Per-call scratch (was thread-local in the monolithic path; these are only + // valid for the duration of a single post and are reset here each call). + std::vector epWrsSinceSignal(epNum, 0); + std::vector epMergedSinceSignal(epNum, 0); // Rotate the starting EP by transfer id so single-segment (single WR) // transfers spread evenly across all QPs instead of always landing on eps[0]. @@ -913,6 +864,10 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, struct ibv_send_wr& wr = mergedWr.wr; wr.wr_id = 0; + // Reset signaling every post: a WR reused across posts (prepared path) may + // have been the signaled tail of a previous post; only the tail selected + // below should carry IBV_SEND_SIGNALED. No-op for the fresh-build path. + wr.send_flags = 0; wr.next = (j + 1 < end) ? &mergedWrs[j + 1].wr : nullptr; mergedReqSize += mergedWr.mergedRequests; } @@ -1028,6 +983,139 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, return {StatusCode::IN_PROGRESS, ""}; } +RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, + const std::vector& localMrPerEp, + const std::vector& remoteMrPerEp, + const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, std::shared_ptr callbackMeta, + TransferUniqueId id, bool isRead, int postBatchSize, + const RdmaTransferControl& control) { + MORI_IO_FUNCTION_TIMER; + + if (sizes.size() == 0 && localOffsets.size() == 0 && remoteOffsets.size() == 0) { + return {StatusCode::SUCCESS, ""}; + } + + if (eps.empty()) { + return {StatusCode::ERR_INVALID_ARGS, "no endpoints"}; + } + + if (localMrPerEp.size() != eps.size() || remoteMrPerEp.size() != eps.size()) { + return {StatusCode::ERR_INVALID_ARGS, "memory-region vectors must align with endpoints"}; + } + + // [tls-scratch] Per worker-thread scratch pools eliminate per-batch heap + // allocations on the RDMA hot path (slot reuse / resize() / clear() retain + // capacity across calls). The pools belong to the OUTERMOST call on a thread; + // if RdmaBatchReadWrite is ever re-entered on the same thread (e.g. from a + // completion callback), the nested call transparently falls back to local + // buffers so the outer call's pools are never clobbered. + thread_local std::vector tlIndices; + thread_local std::vector tlMergedPool; + thread_local std::vector tlChunkedPool; + thread_local std::vector tlChunkPlan; + thread_local int reentryDepth = 0; + + struct ReentryGuard { + int& depth; + explicit ReentryGuard(int& d) : depth(d) { ++depth; } + ~ReentryGuard() { --depth; } + } reentryGuard(reentryDepth); + const bool usePool = (reentryDepth == 1); + + // Used only on (currently non-existent) same-thread re-entry. + std::vector localIndices; + std::vector localMergedPool; + std::vector localChunkedPool; + std::vector localChunkPlan; + + std::vector& indices = usePool ? tlIndices : localIndices; + std::vector& mergedPool = usePool ? tlMergedPool : localMergedPool; + std::vector& chunkedPool = usePool ? tlChunkedPool : localChunkedPool; + std::vector& chunkPlan = usePool ? tlChunkPlan : localChunkPlan; + + // Bound peak retained memory: if an earlier very large batch grew the pools far + // beyond the current need, release the excess so it doesn't stay resident. + constexpr size_t kPoolHighWater = 8192; + if (usePool && sizes.size() <= kPoolHighWater / 2) { + if (mergedPool.size() > kPoolHighWater) { + mergedPool.resize(kPoolHighWater); + mergedPool.shrink_to_fit(); + } + if (chunkedPool.size() > kPoolHighWater) { + chunkedPool.resize(kPoolHighWater); + chunkedPool.shrink_to_fit(); + } + } + + const application::RdmaMemoryRegion& baseLocalMr = localMrPerEp.front(); + const application::RdmaMemoryRegion& baseRemoteMr = remoteMrPerEp.front(); + + std::vector* mergedWrs = nullptr; + size_t mergedWrCount = 0; + RdmaOpRet buildRet = + BuildMergedWorkRequests(eps, baseLocalMr, baseRemoteMr, localOffsets, remoteOffsets, sizes, + isRead, control, indices, mergedPool, chunkedPool, chunkPlan, + mergedWrs, mergedWrCount); + if (buildRet.Failed()) return buildRet; + if (mergedWrCount == 0) return {StatusCode::SUCCESS, ""}; + + return PostMergedWorkRequests(eps, localMrPerEp, remoteMrPerEp, *mergedWrs, mergedWrCount, + callbackMeta, id, postBatchSize, control); +} + +RdmaOpRet BuildPreparedRdmaBatch(const EpPairVec& eps, + const application::RdmaMemoryRegion& baseLocalMr, + const application::RdmaMemoryRegion& baseRemoteMr, + const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead, + const RdmaTransferControl& control, int postBatchSize, + PreparedRdmaBatch& out) { + out.valid = false; + std::vector indices; + std::vector mergedPool; + std::vector chunkedPool; + std::vector chunkPlan; + std::vector* outWrs = nullptr; + size_t outCount = 0; + + RdmaOpRet ret = + BuildMergedWorkRequests(eps, baseLocalMr, baseRemoteMr, localOffsets, remoteOffsets, sizes, + isRead, control, indices, mergedPool, chunkedPool, chunkPlan, outWrs, + outCount); + if (ret.Failed()) return ret; + + // Own the built WR list (moving preserves each WR's sges buffer; sg_list is + // re-armed on every post via ResetMergedWorkRequestPointers). + out.mergedWrs = std::move(*outWrs); + out.mergedWrs.resize(outCount); + out.mergedWrCount = outCount; + out.batchSize = sizes.size(); + out.control = control; + out.postBatchSize = postBatchSize; + out.isRead = isRead; + out.valid = true; + return {StatusCode::SUCCESS, ""}; +} + +RdmaOpRet PostPreparedRdmaBatch(const EpPairVec& eps, + const std::vector& localMrPerEp, + const std::vector& remoteMrPerEp, + PreparedRdmaBatch& prepared, + std::shared_ptr callbackMeta, TransferUniqueId id) { + if (!prepared.valid) { + return {StatusCode::ERR_BAD_STATE, "prepared RDMA batch is not valid"}; + } + if (localMrPerEp.size() != eps.size() || remoteMrPerEp.size() != eps.size()) { + return {StatusCode::ERR_INVALID_ARGS, "memory-region vectors must align with endpoints"}; + } + if (prepared.mergedWrCount == 0) return {StatusCode::SUCCESS, ""}; + + return PostMergedWorkRequests(eps, localMrPerEp, remoteMrPerEp, prepared.mergedWrs, + prepared.mergedWrCount, callbackMeta, id, prepared.postBatchSize, + prepared.control); +} + RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, const std::vector& localMrPerEp, const std::vector& remoteMrPerEp, diff --git a/src/io/rdma/common.hpp b/src/io/rdma/common.hpp index 471316603..9ada7d6b1 100644 --- a/src/io/rdma/common.hpp +++ b/src/io/rdma/common.hpp @@ -281,6 +281,75 @@ struct RdmaTransferControl { bool disableMerge{false}; }; +// A single (possibly merged / chunked) RDMA work request plus its scatter-gather +// list. Built once from a batch description and reused for posting. Defined here +// (rather than privately in common.cpp) so it can be cached in a PreparedRdmaBatch. +struct MergedWorkRequest { + ibv_send_wr wr{}; + std::vector sges; + size_t totalRemoteLength = 0; + size_t mergedRequests = 1; +}; + +// Result of the "build" phase of a batch transfer: the fully merged/chunked WR +// list with per-element addresses/lengths resolved, but WITHOUT the per-post +// state (lkey/rkey, wr_id, signaling, next pointers). Reused across posts by +// PostMergedWorkRequests, which re-arms that per-post state each time. This lets +// callers hoist the sort/merge/chunk cost out of a hot transfer loop (the NIXL +// createXferReq / postXferReq split), gated by the caller. +struct PreparedRdmaBatch { + std::vector mergedWrs; + size_t mergedWrCount{0}; + size_t batchSize{0}; // original request count (callbackMeta totalBatchSize when !creditByWrCount) + RdmaTransferControl control{}; + int postBatchSize{-1}; + bool isRead{false}; + bool valid{false}; +}; + +// Build phase: produce the merged/chunked WR list for the given batch into the +// caller-provided scratch pools. On success, `outWrs` points at whichever pool +// holds the final list and `outCount` is its length. Used by both the original +// RdmaBatchReadWrite (thread-local pools) and the prepared path (owned vectors). +RdmaOpRet BuildMergedWorkRequests(const EpPairVec& eps, + const application::RdmaMemoryRegion& baseLocalMr, + const application::RdmaMemoryRegion& baseRemoteMr, + const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead, + const RdmaTransferControl& control, + std::vector& indices, + std::vector& mergedPool, + std::vector& chunkedPool, + std::vector& chunkPlan, + std::vector*& outWrs, size_t& outCount); + +// Post phase: re-arm per-post state (lkey/rkey, wr_id, signaling, next) on the +// merged WR list and submit via ibv_post_send. Safe to call repeatedly on the +// same `mergedWrs` (reset-on-post semantics), which is what the prepared path does. +RdmaOpRet PostMergedWorkRequests(const EpPairVec& eps, + const std::vector& localMrPerEp, + const std::vector& remoteMrPerEp, + std::vector& mergedWrs, size_t mergedWrCount, + std::shared_ptr callbackMeta, TransferUniqueId id, + int postBatchSize, const RdmaTransferControl& control); + +// Prepared-transfer helpers: build once, post many. BuildPreparedRdmaBatch runs +// the build phase and stores the owned WR list in `out`. PostPreparedRdmaBatch +// re-posts it. Both operate on the inline (non-executor) posting path. +RdmaOpRet BuildPreparedRdmaBatch(const EpPairVec& eps, + const application::RdmaMemoryRegion& baseLocalMr, + const application::RdmaMemoryRegion& baseRemoteMr, + const SizeVec& localOffsets, const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead, + const RdmaTransferControl& control, int postBatchSize, + PreparedRdmaBatch& out); + +RdmaOpRet PostPreparedRdmaBatch(const EpPairVec& eps, + const std::vector& localMrPerEp, + const std::vector& remoteMrPerEp, + PreparedRdmaBatch& prepared, + std::shared_ptr callbackMeta, TransferUniqueId id); + RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, const std::vector& localMrPerEp, const std::vector& remoteMrPerEp, diff --git a/src/pybind/pybind_io.cpp b/src/pybind/pybind_io.cpp index 5ae5a7539..b659572eb 100644 --- a/src/pybind/pybind_io.cpp +++ b/src/pybind/pybind_io.cpp @@ -212,6 +212,9 @@ void RegisterMoriIo(pybind11::module_& m) { return out.get().as(); }); + py::class_>( + m, "PreparedTransfer"); + py::class_(m, "IOEngineSession") .def("AllocateTransferUniqueId", &mori::io::IOEngineSession::AllocateTransferUniqueId, py::call_guard()) @@ -237,6 +240,17 @@ void RegisterMoriIo(pybind11::module_& m) { py::gil_scoped_release release; self.BatchWrite(lo, ro, sz, status, id); }) + .def("PrepareBatch", + [](mori::io::IOEngineSession& self, const SizeArray& localOffsets, + const SizeArray& remoteOffsets, const SizeArray& sizes, bool isRead) { + mori::io::SizeVec lo = SizeArrayToVec(localOffsets); + mori::io::SizeVec ro = SizeArrayToVec(remoteOffsets); + mori::io::SizeVec sz = SizeArrayToVec(sizes); + py::gil_scoped_release release; + return self.PrepareBatch(lo, ro, sz, isRead); + }) + .def("PostPrepared", &mori::io::IOEngineSession::PostPrepared, + py::call_guard()) .def("Alive", &mori::io::IOEngineSession::Alive); py::class_(m, "IOEngine") diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp index 7561c42bc..d542a20fd 100644 --- a/tests/cpp/io/bench_engine.cpp +++ b/tests/cpp/io/bench_engine.cpp @@ -48,6 +48,7 @@ #include #include #include +#include #include #include #include @@ -197,6 +198,8 @@ struct Args { bool batch_contiguous = false; // --batch-contiguous (adjacent offsets → merged WR); // default strided (each transfer a separate WR) bool enable_sess = false; // --enable-sess (session fast-path); Python default: off + bool prepare_once = false; // --prepare-once (build WRs once, re-post each iter; nixl-style + // createXferReq/postXferReq split). Requires --enable-sess. std::string mem_type = "gpu"; // --mem-type gpu|cpu std::string init_mem_type; // --initiator-mem-type (empty => mem_type) @@ -241,6 +244,8 @@ static Args ParseArgs(int argc, char** argv) { else if (k == "--batch-contiguous") a.batch_contiguous = true; else if (k == "--enable-sess") a.enable_sess = true; else if (k == "--disable-sess") a.enable_sess = false; + else if (k == "--prepare-once") a.prepare_once = true; + else if (k == "--no-prepare-once") a.prepare_once = false; else if (k == "--mem-type") a.mem_type = next(); else if (k == "--initiator-mem-type") a.init_mem_type = next(); else if (k == "--target-mem-type") a.target_mem_type = next(); @@ -376,6 +381,12 @@ int main(int argc, char** argv) { } const bool isRead = (a.op == "read"); + if (a.prepare_once && !useSess) { + std::cerr << "--prepare-once requires --enable-sess (prepared handles live on a session)" + << std::endl; + std::exit(1); + } + std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; for (auto& planEntry : plan) { @@ -409,9 +420,47 @@ int main(int argc, char** argv) { MemDescVec locVec{localMem}, remVec{peerMem}; BatchSizeVec offVec{offsets}, sizeVec{sizes}; + // --prepare-once: build the merged/chunked WR list ONCE per (msg,batch) so + // the timed loop only re-posts it (nixl createXferReq / postXferReq split). + // Batched => a single N-descriptor handle; singles => one handle per + // transfer. Session-only (guarded above). + std::vector> prepared; + if (a.prepare_once) { + auto build = [&](const SizeVec& lo, const SizeVec& ro, const SizeVec& sz) { + auto h = sessPtr->PrepareBatch(lo, ro, sz, isRead); + if (!h) { + std::cerr << "PrepareBatch failed (msg=" << msg << ", batch=" << curBatch << ")" + << std::endl; + std::exit(1); + } + prepared.push_back(std::move(h)); + }; + if (batched) { + build(offsets, offsets, sizes); + } else { + for (int i = 0; i < curBatch; ++i) { + const size_t off = static_cast(i) * msg; + build(SizeVec{off}, SizeVec{off}, SizeVec{msg}); + } + } + } + auto post = [&]() { TransferStatus* base = &st[0]; for (int i = 0; i < perSlot; ++i) base[i].SetCode(StatusCode::INIT); + if (a.prepare_once) { + // Re-post the pre-built handle(s); no sort/merge/chunk on the hot path. + if (batched) { + TransferUniqueId id = sessPtr->AllocateTransferUniqueId(); + sessPtr->PostPrepared(prepared[0], &base[0], id); + } else { + for (int i = 0; i < curBatch; ++i) { + TransferUniqueId id = sessPtr->AllocateTransferUniqueId(); + sessPtr->PostPrepared(prepared[static_cast(i)], &base[i], id); + } + } + return; + } if (batched) { // ONE N-descriptor batch request (nixl / Python --enable-batch-transfer). TransferUniqueId id = diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 7cb7a35d6..63cfdbe7f 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -249,6 +249,14 @@ def parse_args(): "condition variable. Avoids the cross-thread cv-wakeup latency (~5-10us) " "per completion at the cost of burning a core while waiting.", ) + parser.add_argument( + "--prepare-once", + action="store_true", + help="Build the transfer work requests once (per size/batch) and re-post them " + "each iteration, hoisting descriptor-build cost out of the timed loop (the NIXL " + "createXferReq/postXferReq split). Falls back to the inline path if the backend " + "does not support prepared transfers.", + ) parser.add_argument( "--disable-chunking", action="store_true", @@ -371,6 +379,7 @@ def __init__( post_batch_size: int = -1, num_worker_threads: int = 1, busy_wait: bool = False, + prepare_once: bool = False, poll_cq_mode: str = "polling", max_send_wr: int = 0, max_cqe_num: int = 0, @@ -413,6 +422,16 @@ def __init__( self.post_batch_size = post_batch_size self.num_worker_threads = num_worker_threads self.busy_wait = busy_wait + # Prepared (build-once, post-many) handles live on a session, so + # --prepare-once requires --enable-sess. + if prepare_once and not self.enable_sess: + raise ValueError("--prepare-once requires --enable-sess") + self.prepare_once = prepare_once + # Cache of prepared handles keyed by (buffer_size, transfer_batch_size). + # Built lazily on first run_*_once for a size, reused across warmup+timed + # iters, and cleared between sweep points via _reset_prepared(). + self._prepared_cache = {} + self._prepare_unsupported_warned = False self.poll_cq_mode = ( PollCqMode.POLLING if poll_cq_mode == "polling" else PollCqMode.EVENT ) @@ -986,6 +1005,56 @@ def _wait_status(self, status): else: status.Wait() + def _reset_prepared(self): + # Prepared handles capture offsets/sizes at build time, so they are only + # valid for the (buffer_size, batch) they were built for. Drop them when + # moving to a new sweep point. + self._prepared_cache = {} + + def _warn_prepare_unsupported(self): + if not self._prepare_unsupported_warned: + print( + "[bench] --prepare-once: backend returned no prepared handle; " + "falling back to the inline transfer path.", + flush=True, + ) + self._prepare_unsupported_warned = True + + def _get_prepared_batch( + self, buffer_size, transfer_batch_size, offsets, sizes, is_read + ): + # Build (once) and cache the prepared handle for the batched path. Returns + # None if the backend does not support prepared transfers (caller falls + # back to inline). Prepared transfers live on a session (--enable-sess). + key = ("batch", buffer_size, transfer_batch_size, is_read) + if key in self._prepared_cache: + return self._prepared_cache[key] + prepared = self.sess.prepare_batch(offsets, offsets, sizes, is_read) + if prepared is None: + self._warn_prepare_unsupported() + self._prepared_cache[key] = prepared + return prepared + + def _get_prepared_singles(self, buffer_size, transfer_batch_size, is_read): + # Build (once) and cache one 1-element prepared handle per transfer in the + # single-transfer path. Returns None (whole list) if unsupported. + key = ("single", buffer_size, transfer_batch_size, is_read) + if key in self._prepared_cache: + return self._prepared_cache[key] + prepared_list = [] + for i in range(transfer_batch_size): + offset = buffer_size * i + off = np.array([offset], dtype=np.uint64) + sz = np.array([buffer_size], dtype=np.uint64) + p = self.sess.prepare_batch(off, off, sz, is_read) + if p is None: + self._warn_prepare_unsupported() + self._prepared_cache[key] = None + return None + prepared_list.append(p) + self._prepared_cache[key] = prepared_list + return prepared_list + def run_single_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size if ( @@ -1000,6 +1069,31 @@ def run_single_once(self, buffer_size, transfer_batch_size): for i in range(transfer_batch_size): transfer_uids.append(self.engine.allocate_transfer_uid()) + is_read = self.op_type == "read" + prepared_list = None + if self.prepare_once: + prepared_list = self._get_prepared_singles( + buffer_size, transfer_batch_size, is_read + ) + + if prepared_list is not None: + # Prepared (build-once, post-many): each single transfer's WR list was + # built outside the timed loop; only re-posting is timed here. + st = time.time() + for i in range(transfer_batch_size): + status = self.sess.post_prepared(prepared_list[i], transfer_uids[i]) + status_list.append(status) + launched = time.time() + for status in status_list: + self._wait_status(status) + done = time.time() + + launch_duration = launched - st + transfer_duration = done - launched + for status in status_list: + assert status.Succeeded(), f"Transfer failed: {status.Message()}" + return launch_duration, transfer_duration + func, arg_list = None, [] for i in range(transfer_batch_size): offset = buffer_size * i @@ -1057,7 +1151,19 @@ def run_batch_once(self, buffer_size, transfer_batch_size): sizes = np.full(transfer_batch_size, buffer_size, dtype=np.uint64) transfer_uid = self.engine.allocate_transfer_uid() - if self.enable_sess: + is_read = self.op_type == "read" + prepared = None + if self.prepare_once: + prepared = self._get_prepared_batch( + buffer_size, transfer_batch_size, offsets, sizes, is_read + ) + + if prepared is not None: + # Prepared (build-once, post-many): the WR list was built outside the + # timed loop; only the post is timed here. + st = time.time() + transfer_status = self.sess.post_prepared(prepared, transfer_uid) + elif self.enable_sess: func = ( self.sess.batch_read if self.op_type == "read" @@ -1107,6 +1213,9 @@ def run_once(self, buffer_size, transfer_batch_size): return self.run_single_once(buffer_size, transfer_batch_size) def _run_and_compute(self, buffer_size, transfer_batch_size, iters): + # Prepared handles capture offsets/sizes for a specific (size, batch); drop + # any from a previous sweep point so this point rebuilds them once. + self._reset_prepared() for _ in range(self.warmup_iters): self.run_once(buffer_size, transfer_batch_size) @@ -1366,6 +1475,7 @@ def benchmark_xgmi_worker(local_rank, node_rank, args): num_streams=args.num_streams, num_events=args.num_events, busy_wait=args.busy_wait, + prepare_once=args.prepare_once, xgmi_multiprocess=True, ) bench.print_config() @@ -1407,6 +1517,7 @@ def benchmark_engine(local_rank, node_rank, args): post_batch_size=args.post_batch_size, num_worker_threads=args.num_worker_threads, busy_wait=args.busy_wait, + prepare_once=args.prepare_once, poll_cq_mode=args.poll_cq_mode, max_send_wr=args.max_send_wr, max_cqe_num=args.max_cqe_num, @@ -1478,6 +1589,7 @@ def benchmark_xgmi(args): num_streams=args.num_streams, num_events=args.num_events, busy_wait=args.busy_wait, + prepare_once=args.prepare_once, xgmi_multiprocess=False, ) bench.print_config() From 30ed4be3258f257dcdfc0f2e2abb49809ae17379 Mon Sep 17 00:00:00 2001 From: Akbarzadeh Date: Thu, 30 Jul 2026 15:16:34 -0500 Subject: [PATCH 09/10] io(rdma): harden prepared transfers against misuse Bind each prepared handle to the session that built it and reject posts made through a different session, whose QPs and memory keys do not match the cached work requests. Serialize posts of the same handle, since every post re-arms that shared work-request list in place and concurrent callers would corrupt the chain. Refuse to prepare when the session has a worker pool: the prepared path posts inline from the calling thread, so it would bypass the pool and report a thread count that never ran. Both benchmarks now reject --prepare-once with --num-worker-threads > 1 up front, which also stops the Python bench from silently falling back to the inline path mid-sweep. QP parallelism is unaffected. Co-Authored-By: Claude Opus 4 Co-authored-by: Cursor --- include/mori/io/backend.hpp | 6 +++++ src/io/rdma/backend_impl.cpp | 41 ++++++++++++++++++++++++++++++----- tests/cpp/io/bench_engine.cpp | 9 ++++++++ tests/python/io/benchmark.py | 9 ++++++++ 4 files changed, 60 insertions(+), 5 deletions(-) diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index 789059446..f3d28d783 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -164,6 +164,12 @@ class BackendSession { // submits it with a fresh status/id. The default implementation returns // nullptr / reports unsupported, so backends may opt in. Offsets/sizes are // captured at prepare time and must not change between posts. + // + // A handle is bound to the session that created it and may only be posted + // through that session. Posting is submitted from the calling thread, so a + // backend configured with an internal worker pool may decline to prepare + // (returns nullptr) rather than silently bypass it; callers must treat a null + // handle as "use the regular batch path". virtual std::shared_ptr PrepareBatch(const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, bool isRead) { diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp index 7de016ed3..b1588a591 100644 --- a/src/io/rdma/backend_impl.cpp +++ b/src/io/rdma/backend_impl.cpp @@ -1446,7 +1446,17 @@ void RdmaBackendSession::BatchReadWrite(const SizeVec& localOffsets, const SizeV // Concrete prepared-transfer handle for the RDMA backend: owns the built // (sorted/merged/chunked) work-request list so PostPrepared can re-post it. struct RdmaPreparedTransfer : public PreparedTransfer { + explicit RdmaPreparedTransfer(const RdmaBackendSession* owner_) : owner(owner_) {} + + // The WR addresses were built against this session's memory descriptors. + // Posting through another session would combine those addresses with a + // different session's QPs/MRs and is therefore invalid. + const RdmaBackendSession* const owner; PreparedRdmaBatch batch; + // Each post re-arms the shared WR list in place (lkey/rkey, wr_id, send_flags, + // next chain). Serialize posts of the SAME handle so concurrent callers can't + // corrupt that chain; distinct handles stay fully parallel. + std::mutex postMu; }; std::shared_ptr RdmaBackendSession::PrepareBatch(const SizeVec& localOffsets, @@ -1458,15 +1468,26 @@ std::shared_ptr RdmaBackendSession::PrepareBatch(const SizeVec MORI_IO_WARN("PrepareBatch: session has no endpoints/memory regions"); return nullptr; } - // The prepared path uses the inline (non-executor) posting route, mirroring - // the RdmaBatchReadWrite `else` branch control setup. + // The prepared path always posts inline from the calling thread, so it would + // silently ignore the worker pool that BatchReadWrite uses. Rather than report + // numbers for a posting model the caller did not ask for, refuse up front. + // Prepared transfers still use every QP in the session (qpPerTransfer); only + // multi-threaded submission is unavailable. + if (executor != nullptr) { + MORI_IO_WARN( + "PrepareBatch is unsupported with numWorkerThreads>1: the prepared path posts inline " + "from the calling thread and would bypass the worker pool. Set numWorkerThreads=1 " + "(qpPerTransfer is unaffected) to use prepared transfers."); + return nullptr; + } + // Control setup mirrors the RdmaBatchReadWrite `else` (inline) branch. const bool chunk = config.enableTransferChunking; RdmaTransferControl control{}; control.chunkBytes = chunk ? config.chunkBytes : 0; control.maxChunks = config.maxChunksPerTransfer; control.creditByWrCount = chunk; - auto handle = std::make_shared(); + auto handle = std::make_shared(this); RdmaOpRet ret = BuildPreparedRdmaBatch(eps, localMrPerEp.front(), remoteMrPerEp.front(), localOffsets, remoteOffsets, sizes, isRead, control, config.postBatchSize, @@ -1488,6 +1509,11 @@ void RdmaBackendSession::PostPrepared(const std::shared_ptr& p status->Update(StatusCode::ERR_INVALID_ARGS, "invalid prepared transfer handle"); return; } + if (handle->owner != this) { + status->Update(StatusCode::ERR_INVALID_ARGS, + "prepared transfer belongs to a different RDMA session"); + return; + } // When !creditByWrCount, totalBatchSize is the original request count; the // chunked path overwrites it inside PostMergedWorkRequests (ownsTotalBatchSize). @@ -1495,8 +1521,13 @@ void RdmaBackendSession::PostPrepared(const std::shared_ptr& p auto callbackMeta = std::make_shared(status, id, totalCredit); internal::PublishCurrentIoCallDiagnostics(callbackMeta); - RdmaOpRet ret = - PostPreparedRdmaBatch(eps, localMrPerEp, remoteMrPerEp, handle->batch, callbackMeta, id); + // Guard the in-place re-arm + ibv_post_send of the shared WR list so two + // threads posting the same handle can't interleave and corrupt the chain. + RdmaOpRet ret; + { + std::lock_guard postLock(handle->postMu); + ret = PostPreparedRdmaBatch(eps, localMrPerEp, remoteMrPerEp, handle->batch, callbackMeta, id); + } assert(!ret.Init()); if (ret.Failed() || ret.Succeeded()) { status->Update(ret.code, ret.message); diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp index d542a20fd..379d64241 100644 --- a/tests/cpp/io/bench_engine.cpp +++ b/tests/cpp/io/bench_engine.cpp @@ -386,6 +386,15 @@ int main(int argc, char** argv) { << std::endl; std::exit(1); } + // The prepared path posts inline from this thread, so a worker pool would be + // ignored and the reported thread count would not match what ran. QP-level + // parallelism is unaffected: --num-qp-per-transfer still applies. + if (a.prepare_once && a.worker_threads > 1) { + std::cerr << "--prepare-once requires --num-worker-threads 1 (the prepared path posts " + "inline from the calling thread; use --num-qp-per-transfer for QP parallelism)" + << std::endl; + std::exit(1); + } std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 63cfdbe7f..417374eb0 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -426,6 +426,15 @@ def __init__( # --prepare-once requires --enable-sess. if prepare_once and not self.enable_sess: raise ValueError("--prepare-once requires --enable-sess") + # The prepared path posts inline from the calling thread, so a worker pool + # would be ignored. Reject it here instead of letting prepare_batch return + # None and silently falling back to the inline path mid-sweep. QP-level + # parallelism is unaffected (--num-qp-per-transfer still applies). + if prepare_once and self.num_worker_threads > 1: + raise ValueError( + "--prepare-once requires --num-worker-threads 1 (the prepared path posts " + "inline from the calling thread; use --num-qp-per-transfer for QP parallelism)" + ) self.prepare_once = prepare_once # Cache of prepared handles keyed by (buffer_size, transfer_batch_size). # Built lazily on first run_*_once for a size, reused across warmup+timed From 10c71ea4c7454125bd4c3b177b194deb6b996a29 Mon Sep 17 00:00:00 2001 From: amirakb89 Date: Fri, 31 Jul 2026 20:18:40 +0000 Subject: [PATCH 10/10] bench(io): worker-pool prepared transfers, XGMI backend, and validation PrepareBatch previously refused when numWorkerThreads > 1. It now builds one single-endpoint WR list per slice via SplitBatchWork -- shared with the build-and-post path, so prepared and non-prepared decompose a batch identically at a given worker count -- and MultithreadExecutor::PostPrepared fans them across workers. Slices share no mutable state; postMu still serializes posts of one handle. Add --backend rdma|xgmi; bench_engine hardcoded BackendType::RDMA and could not exercise XGMI. --prepare-once is rejected on non-rdma backends, where it would silently fall back to inline while still reporting as prepared. Validate transferred data as part of the run: both ranks seed a rank- and offset-dependent pattern and exchange one FNV-1a checksum per slot, covering every byte in O(batch) on the wire. Both ranks always call the exchange, so --skip-validate on one side reports "skipped" rather than blocking its peer in Allgather. Confine prepared transfers to the C++ benchmark: drop the Python path and its pybind bindings (those files are now byte-identical to their pre-prepared state), plus the redundant --no-prepare-once. Also from the previous review round: deregister memory before freeing its HIP allocation; hoist status/id vectors out of the timed region; replace the raw-socket rendezvous with SocketBootstrapNetwork; unwind through main instead of std::exit; declare numpy in requirements-build.txt. Co-Authored-By: Claude Opus 5 --- docs/MORI-IO-BENCHMARK.md | 10 + include/mori/io/backend.hpp | 17 +- python/mori/io/engine.py | 10 - requirements-build.txt | 1 + src/io/engine.cpp | 4 +- src/io/rdma/backend_impl.cpp | 116 +++++-- src/io/rdma/common.cpp | 45 ++- src/io/rdma/common.hpp | 10 +- src/io/rdma/executor.cpp | 96 +++-- src/io/rdma/executor.hpp | 35 +- src/pybind/pybind_io.cpp | 14 - tests/cpp/io/bench_engine.cpp | 634 +++++++++++++++++++++++----------- tests/python/io/benchmark.py | 123 +------ 13 files changed, 674 insertions(+), 441 deletions(-) diff --git a/docs/MORI-IO-BENCHMARK.md b/docs/MORI-IO-BENCHMARK.md index 75f71a28a..3d8ff823b 100644 --- a/docs/MORI-IO-BENCHMARK.md +++ b/docs/MORI-IO-BENCHMARK.md @@ -146,10 +146,20 @@ Most flags share names with the Python benchmark (`--op-type`/`--op`, | `--enable-batch-transfer` (default OFF) | `--enable-batch-transfer` / `--disable-batch-transfer` (C++ default **ON**) | Same meaning as Python: **ON** = one N-descriptor batch request (the nixl-equivalent; nixl always batches); **OFF** = `--transfer-batch-size` N *individual* single-transfer submissions per iteration (Python `run_single_once`). The C++ tool defaults **ON** so the out-of-the-box run and `--all-batch` sweep are nixl-comparable; pass `--disable-batch-transfer` for the Python single-submission path. | | *(n/a)* | `--warmup-iters` / `--iters` | Explicit warmup (untimed) and timed iteration counts. | | `--sweep-start-size` / `--sweep-max-size` | `--sweep-start` / `--sweep-max` (aliases accepted) | Plus `--sweep-step` for a linear (vs geometric) sweep. | +| `--backend rdma\|xgmi\|fabric` | `--backend rdma\|xgmi` | C++ has no `fabric` path. `--num-streams` / `--num-events` size the XGMI HIP pools, as in Python. | +| *(n/a)* | `--prepare-once` | Prepared (build-once, post-many) transfers are exercised **only** by the C++ tool; the Python benchmark always posts inline. Rejected on any backend but `rdma`, since only the RDMA backend implements prepared transfers — elsewhere the run would silently fall back to the inline path mid-sweep and report a "prepared" number that never ran prepared. | +| *(always on)* | `--skip-validate` | Both tools verify transferred data as part of the run. Python compares the target's buffer byte-for-byte over gloo; the C++ tool seeds a rank- and offset-dependent pattern and exchanges one checksum per transferred slot after the sweep (so the check is O(batch) on the wire but still covers every byte), printing `validation: OK` or naming the first mismatching slot and exiting non-zero. `--skip-validate` opts out; passing it to only one rank is safe — that rank contributes nothing and the run reports `validation: skipped` rather than hanging. | > Keep the process on the RDMA path with `MORI_DISABLE_AUTO_XGMI=1`. Use > `--self-ip` equal to the data-plane IP each node advertises for QP setup; > `0.0.0.0` will fail with `connect(...): Connection refused`. +> +> For `--backend xgmi`, set `MORI_DISABLE_AUTO_XGMI=0` instead and run both +> ranks on the **same** host with different `--gpu` (XGMI is intra-node), giving +> both the same `--master-ip`/`--self-ip`. Note this measures the *cross-process* +> XGMI path, since the C++ tool is always two processes; the Python benchmark's +> default single-process mode avoids cross-process overhead entirely, so the two +> are not directly comparable at small message sizes. ## Benchmark Arguments diff --git a/include/mori/io/backend.hpp b/include/mori/io/backend.hpp index f3d28d783..42ea4227b 100644 --- a/include/mori/io/backend.hpp +++ b/include/mori/io/backend.hpp @@ -118,11 +118,6 @@ inline std::ostream& operator<<(std::ostream& os, const FabricBackendConfig& c) /* ---------------------------------------------------------------------------------------------- */ /* PreparedTransfer */ /* ---------------------------------------------------------------------------------------------- */ -// Opaque, backend-owned handle for a transfer whose work requests were built -// once (sort/merge/chunk) and can be re-posted many times. Lets callers hoist -// the descriptor-build cost out of a hot transfer loop (the NIXL -// createXferReq/postXferReq split). Backends that do not support prepared -// transfers simply return nullptr from PrepareBatch. struct PreparedTransfer { PreparedTransfer() = default; virtual ~PreparedTransfer() = default; @@ -159,17 +154,7 @@ class BackendSession { BatchReadWrite(localOffsets, remoteOffsets, sizes, status, id, true); } - // Prepared (build-once, post-many) transfer API. PrepareBatch builds the work - // requests for the given batch and returns a reusable handle; PostPrepared - // submits it with a fresh status/id. The default implementation returns - // nullptr / reports unsupported, so backends may opt in. Offsets/sizes are - // captured at prepare time and must not change between posts. - // - // A handle is bound to the session that created it and may only be posted - // through that session. Posting is submitted from the calling thread, so a - // backend configured with an internal worker pool may decline to prepare - // (returns nullptr) rather than silently bypass it; callers must treat a null - // handle as "use the regular batch path". + // Prepared (build-once, post-many) transfer API. virtual std::shared_ptr PrepareBatch(const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, bool isRead) { diff --git a/python/mori/io/engine.py b/python/mori/io/engine.py index e7b8e2793..e6cc76cdf 100644 --- a/python/mori/io/engine.py +++ b/python/mori/io/engine.py @@ -79,16 +79,6 @@ def batch_read(self, *args): def batch_write(self, *args): return self._batch_single_side_transfer(self._sess.BatchWrite, *args) - def prepare_batch(self, local_offsets, remote_offsets, sizes, is_read): - # Build the work requests once; returns a reusable handle (None if the - # backend does not support prepared transfers). - return self._sess.PrepareBatch(local_offsets, remote_offsets, sizes, is_read) - - def post_prepared(self, prepared, transfer_uid): - transfer_status = mori_cpp.TransferStatus() - self._sess.PostPrepared(prepared, transfer_status, transfer_uid) - return transfer_status - def alive(self): return self._sess.Alive() diff --git a/requirements-build.txt b/requirements-build.txt index aa524bcf0..287d2d0ab 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -3,5 +3,6 @@ setuptools>=65 setuptools-scm>=8 pybind11 ninja +numpy prettytable pytest-assume diff --git a/src/io/engine.cpp b/src/io/engine.cpp index ca5ef26b2..581d471dc 100644 --- a/src/io/engine.cpp +++ b/src/io/engine.cpp @@ -161,8 +161,8 @@ void IOEngineSession::BatchWrite(const SizeVec& localOffsets, const SizeVec& rem } std::shared_ptr IOEngineSession::PrepareBatch(const SizeVec& localOffsets, - const SizeVec& remoteOffsets, - const SizeVec& sizes, bool isRead) { + const SizeVec& remoteOffsets, + const SizeVec& sizes, bool isRead) { MORI_IO_FUNCTION_TIMER; return backendSess->PrepareBatch(localOffsets, remoteOffsets, sizes, isRead); } diff --git a/src/io/rdma/backend_impl.cpp b/src/io/rdma/backend_impl.cpp index b1588a591..3a0655640 100644 --- a/src/io/rdma/backend_impl.cpp +++ b/src/io/rdma/backend_impl.cpp @@ -1452,10 +1452,18 @@ struct RdmaPreparedTransfer : public PreparedTransfer { // Posting through another session would combine those addresses with a // different session's QPs/MRs and is therefore invalid. const RdmaBackendSession* const owner; + // Inline posting: one WR list spanning every endpoint in the session, which the + // post phase distributes across them. Empty when posting through workers. PreparedRdmaBatch batch; - // Each post re-arms the shared WR list in place (lkey/rkey, wr_id, send_flags, - // next chain). Serialize posts of the SAME handle so concurrent callers can't - // corrupt that chain; distinct handles stay fully parallel. + // Worker-pool posting: one self-contained single-endpoint WR list per slice, so + // workers re-arm and post disjoint state. Empty when posting inline. + std::vector slices; + // Completion credit for the whole handle, fixed at build time. Used by the + // sliced path, where no single slice owns callbackMeta's total. + int totalCredit{0}; + // Each post re-arms the WR list in place (lkey/rkey, wr_id, send_flags, next + // chain). Serialize posts of the SAME handle so concurrent callers can't corrupt + // it; distinct handles stay fully parallel. std::mutex postMu; }; @@ -1468,30 +1476,69 @@ std::shared_ptr RdmaBackendSession::PrepareBatch(const SizeVec MORI_IO_WARN("PrepareBatch: session has no endpoints/memory regions"); return nullptr; } - // The prepared path always posts inline from the calling thread, so it would - // silently ignore the worker pool that BatchReadWrite uses. Rather than report - // numbers for a posting model the caller did not ask for, refuse up front. - // Prepared transfers still use every QP in the session (qpPerTransfer); only - // multi-threaded submission is unavailable. - if (executor != nullptr) { - MORI_IO_WARN( - "PrepareBatch is unsupported with numWorkerThreads>1: the prepared path posts inline " - "from the calling thread and would bypass the worker pool. Set numWorkerThreads=1 " - "(qpPerTransfer is unaffected) to use prepared transfers."); - return nullptr; + // Which posting model this handle is built for mirrors BatchReadWrite's executor + // gating, so prepared and non-prepared transfers decompose a batch identically at + // a given worker count: chunked worker posting requires GPU local memory, and + // without chunking any session that has an executor qualifies. + const bool chunk = config.enableTransferChunking; + const bool useWorkerChunking = + executor != nullptr && chunk && localLoc_ == MemoryLocationType::GPU; + const bool useWorkers = executor != nullptr && (useWorkerChunking || !chunk); + + auto handle = std::make_shared(this); + + if (useWorkers) { + // Control setup mirrors what an executor worker applies per task. + RdmaTransferControl control{}; + control.chunkBytes = useWorkerChunking ? config.chunkBytes : 0; + control.maxChunks = useWorkerChunking ? config.maxChunksPerTransfer : 1; + control.creditByWrCount = useWorkerChunking; + control.disableMerge = useWorkerChunking; + // The handle owns the completion total; a slice only knows its own share. + control.ownsTotalBatchSize = false; + + const auto splits = SplitBatchWork(static_cast(eps.size()), executor->NumWorkers(), + static_cast(sizes.size())); + handle->slices.resize(splits.size()); + size_t credit = 0; + for (size_t i = 0; i < splits.size(); ++i) { + const size_t begin = static_cast(splits[i].first); + const size_t end = static_cast(splits[i].second); + SizeVec sliceLocalOffsets(localOffsets.begin() + begin, localOffsets.begin() + end); + SizeVec sliceRemoteOffsets(remoteOffsets.begin() + begin, remoteOffsets.begin() + end); + SizeVec sliceSizes(sizes.begin() + begin, sizes.begin() + end); + + // Built against the whole endpoint set so the SGE and message-size bounds are + // the minimum across endpoints; that is what keeps a slice postable on any of + // them, and lets PostPrepared rotate slices across QPs per transfer id. + RdmaOpRet ret = BuildPreparedRdmaBatch( + eps, localMrPerEp.front(), remoteMrPerEp.front(), sliceLocalOffsets, sliceRemoteOffsets, + sliceSizes, isRead, control, config.postBatchSize, handle->slices[i]); + if (ret.Failed()) { + MORI_IO_WARN("PrepareBatch failed to build work requests: {}", ret.message); + return nullptr; + } + // Credit matches what the CQ path will award: one per WR when crediting by WR + // count, otherwise one per original descriptor. + credit += useWorkerChunking ? handle->slices[i].mergedWrCount : handle->slices[i].batchSize; + } + if (credit > static_cast(std::numeric_limits::max())) { + MORI_IO_WARN("PrepareBatch: completion credit {} exceeds int range", credit); + return nullptr; + } + handle->totalCredit = static_cast(credit); + return handle; } + // Control setup mirrors the RdmaBatchReadWrite `else` (inline) branch. - const bool chunk = config.enableTransferChunking; RdmaTransferControl control{}; control.chunkBytes = chunk ? config.chunkBytes : 0; control.maxChunks = config.maxChunksPerTransfer; control.creditByWrCount = chunk; - auto handle = std::make_shared(this); - RdmaOpRet ret = - BuildPreparedRdmaBatch(eps, localMrPerEp.front(), remoteMrPerEp.front(), localOffsets, - remoteOffsets, sizes, isRead, control, config.postBatchSize, - handle->batch); + RdmaOpRet ret = BuildPreparedRdmaBatch(eps, localMrPerEp.front(), remoteMrPerEp.front(), + localOffsets, remoteOffsets, sizes, isRead, control, + config.postBatchSize, handle->batch); if (ret.Failed()) { MORI_IO_WARN("PrepareBatch failed to build work requests: {}", ret.message); return nullptr; @@ -1505,7 +1552,8 @@ void RdmaBackendSession::PostPrepared(const std::shared_ptr& p status->SetCode(StatusCode::IN_PROGRESS); auto* handle = dynamic_cast(prepared.get()); - if (handle == nullptr || !handle->batch.valid) { + const bool sliced = handle != nullptr && !handle->slices.empty(); + if (handle == nullptr || (!sliced && !handle->batch.valid)) { status->Update(StatusCode::ERR_INVALID_ARGS, "invalid prepared transfer handle"); return; } @@ -1515,16 +1563,24 @@ void RdmaBackendSession::PostPrepared(const std::shared_ptr& p return; } - // When !creditByWrCount, totalBatchSize is the original request count; the - // chunked path overwrites it inside PostMergedWorkRequests (ownsTotalBatchSize). - const int totalCredit = static_cast(handle->batch.batchSize); - auto callbackMeta = std::make_shared(status, id, totalCredit); - internal::PublishCurrentIoCallDiagnostics(callbackMeta); - - // Guard the in-place re-arm + ibv_post_send of the shared WR list so two - // threads posting the same handle can't interleave and corrupt the chain. + // Guard the in-place re-arm + ibv_post_send of the WR list(s) so two threads + // posting the same handle can't interleave and corrupt them. Within one post the + // slices are disjoint, so workers need no further synchronization. RdmaOpRet ret; - { + if (sliced) { + auto callbackMeta = std::make_shared(status, id, handle->totalCredit); + internal::PublishCurrentIoCallDiagnostics(callbackMeta); + ExecutorPreparedReq req{ + eps, localMrPerEp.front(), remoteMrPerEp.front(), handle->slices, callbackMeta, id}; + std::lock_guard postLock(handle->postMu); + ret = executor->PostPrepared(req); + } else { + // When !creditByWrCount, totalBatchSize is the original request count; the + // chunked path overwrites it inside PostMergedWorkRequests (ownsTotalBatchSize). + const int totalCredit = static_cast(handle->batch.batchSize); + auto callbackMeta = std::make_shared(status, id, totalCredit); + internal::PublishCurrentIoCallDiagnostics(callbackMeta); + std::lock_guard postLock(handle->postMu); ret = PostPreparedRdmaBatch(eps, localMrPerEp, remoteMrPerEp, handle->batch, callbackMeta, id); } diff --git a/src/io/rdma/common.cpp b/src/io/rdma/common.cpp index 67d0d66b4..8785e4ef0 100644 --- a/src/io/rdma/common.cpp +++ b/src/io/rdma/common.cpp @@ -596,8 +596,7 @@ RdmaOpRet BuildMergedWorkRequests(const EpPairVec& eps, const application::RdmaMemoryRegion& baseRemoteMr, const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, bool isRead, - const RdmaTransferControl& control, - std::vector& indices, + const RdmaTransferControl& control, std::vector& indices, std::vector& mergedPool, std::vector& chunkedPool, std::vector& chunkPlan, @@ -830,10 +829,30 @@ RdmaOpRet PostMergedWorkRequests(const EpPairVec& eps, if (postBatchSize <= 0) postBatchSize = 1; int numPostBatch = (mergedWrCount + postBatchSize - 1) / postBatchSize; - // Per-call scratch (was thread-local in the monolithic path; these are only - // valid for the duration of a single post and are reset here each call). - std::vector epWrsSinceSignal(epNum, 0); - std::vector epMergedSinceSignal(epNum, 0); + // [tls-scratch] These counters are only meaningful for the duration of one + // post, but keeping them thread-local retains capacity across posts instead of + // heap-allocating twice per transfer on the RDMA hot path. As in + // RdmaBatchReadWrite, the pools belong to the OUTERMOST call on a thread; a + // nested post (e.g. from a completion callback reached via the failure path + // below) falls back to local vectors so the outer call's counters survive. + thread_local std::vector tlEpWrsSinceSignal; + thread_local std::vector tlEpMergedSinceSignal; + thread_local int postReentryDepth = 0; + + struct PostReentryGuard { + int& depth; + explicit PostReentryGuard(int& d) : depth(d) { ++depth; } + ~PostReentryGuard() { --depth; } + } postReentryGuard(postReentryDepth); + const bool usePool = (postReentryDepth == 1); + + std::vector localEpWrsSinceSignal; + std::vector localEpMergedSinceSignal; + std::vector& epWrsSinceSignal = usePool ? tlEpWrsSinceSignal : localEpWrsSinceSignal; + std::vector& epMergedSinceSignal = + usePool ? tlEpMergedSinceSignal : localEpMergedSinceSignal; + epWrsSinceSignal.assign(epNum, 0); + epMergedSinceSignal.assign(epNum, 0); // Rotate the starting EP by transfer id so single-segment (single WR) // transfers spread evenly across all QPs instead of always landing on eps[0]. @@ -1053,10 +1072,9 @@ RdmaOpRet RdmaBatchReadWrite(const EpPairVec& eps, std::vector* mergedWrs = nullptr; size_t mergedWrCount = 0; - RdmaOpRet buildRet = - BuildMergedWorkRequests(eps, baseLocalMr, baseRemoteMr, localOffsets, remoteOffsets, sizes, - isRead, control, indices, mergedPool, chunkedPool, chunkPlan, - mergedWrs, mergedWrCount); + RdmaOpRet buildRet = BuildMergedWorkRequests( + eps, baseLocalMr, baseRemoteMr, localOffsets, remoteOffsets, sizes, isRead, control, indices, + mergedPool, chunkedPool, chunkPlan, mergedWrs, mergedWrCount); if (buildRet.Failed()) return buildRet; if (mergedWrCount == 0) return {StatusCode::SUCCESS, ""}; @@ -1079,10 +1097,9 @@ RdmaOpRet BuildPreparedRdmaBatch(const EpPairVec& eps, std::vector* outWrs = nullptr; size_t outCount = 0; - RdmaOpRet ret = - BuildMergedWorkRequests(eps, baseLocalMr, baseRemoteMr, localOffsets, remoteOffsets, sizes, - isRead, control, indices, mergedPool, chunkedPool, chunkPlan, outWrs, - outCount); + RdmaOpRet ret = BuildMergedWorkRequests(eps, baseLocalMr, baseRemoteMr, localOffsets, + remoteOffsets, sizes, isRead, control, indices, + mergedPool, chunkedPool, chunkPlan, outWrs, outCount); if (ret.Failed()) return ret; // Own the built WR list (moving preserves each WR's sges buffer; sg_list is diff --git a/src/io/rdma/common.hpp b/src/io/rdma/common.hpp index 9ada7d6b1..5fcb1d714 100644 --- a/src/io/rdma/common.hpp +++ b/src/io/rdma/common.hpp @@ -300,7 +300,9 @@ struct MergedWorkRequest { struct PreparedRdmaBatch { std::vector mergedWrs; size_t mergedWrCount{0}; - size_t batchSize{0}; // original request count (callbackMeta totalBatchSize when !creditByWrCount) + // Original request count; becomes callbackMeta totalBatchSize when + // !creditByWrCount (the chunked path overwrites it with the final WR count). + size_t batchSize{0}; RdmaTransferControl control{}; int postBatchSize{-1}; bool isRead{false}; @@ -316,8 +318,7 @@ RdmaOpRet BuildMergedWorkRequests(const EpPairVec& eps, const application::RdmaMemoryRegion& baseRemoteMr, const SizeVec& localOffsets, const SizeVec& remoteOffsets, const SizeVec& sizes, bool isRead, - const RdmaTransferControl& control, - std::vector& indices, + const RdmaTransferControl& control, std::vector& indices, std::vector& mergedPool, std::vector& chunkedPool, std::vector& chunkPlan, @@ -335,7 +336,8 @@ RdmaOpRet PostMergedWorkRequests(const EpPairVec& eps, // Prepared-transfer helpers: build once, post many. BuildPreparedRdmaBatch runs // the build phase and stores the owned WR list in `out`. PostPreparedRdmaBatch -// re-posts it. Both operate on the inline (non-executor) posting path. +// re-posts it. Used both for inline posting (one list spanning every endpoint) and +// for worker posting, where each worker owns a single-endpoint list. RdmaOpRet BuildPreparedRdmaBatch(const EpPairVec& eps, const application::RdmaMemoryRegion& baseLocalMr, const application::RdmaMemoryRegion& baseRemoteMr, diff --git a/src/io/rdma/executor.cpp b/src/io/rdma/executor.cpp index 9b6a5ef75..ff3e86bca 100644 --- a/src/io/rdma/executor.cpp +++ b/src/io/rdma/executor.cpp @@ -124,6 +124,24 @@ void MultithreadExecutor::Worker::MainLoop() { q.pop(); } + thread_local std::vector localMrPerEp(1); + thread_local std::vector remoteMrPerEp(1); + + if (task.preq != nullptr) { + // Prepared slice: the WR list was built by PrepareBatch, so only the + // per-post state (lkey/rkey, wr_id, signaling, next chain) is re-armed here. + localMrPerEp[0] = task.preq->local; + remoteMrPerEp[0] = task.preq->remote; + + RdmaOpRet ret = mori::io::PostPreparedRdmaBatch( + {task.preq->eps[task.epId]}, localMrPerEp, remoteMrPerEp, task.preq->slices[task.sliceId], + task.preq->callbackMeta, task.preq->id); + task.ret.set_value(ret); + MORI_IO_TRACE("Worker {} post prepared task {} slice {} ep {} ret code {}", workerId, + task.preq->id, task.sliceId, task.epId, static_cast(ret.code)); + continue; + } + SizeVec tLoclOffsets(task.req->localOffsets.begin() + task.begin, task.req->localOffsets.begin() + task.end); SizeVec tRemoteOffsets(task.req->remoteOffsets.begin() + task.begin, @@ -138,8 +156,6 @@ void MultithreadExecutor::Worker::MainLoop() { control.ownsTotalBatchSize = false; control.disableMerge = chunk; - thread_local std::vector localMrPerEp(1); - thread_local std::vector remoteMrPerEp(1); localMrPerEp[0] = task.req->local; remoteMrPerEp[0] = task.req->remote; @@ -180,10 +196,7 @@ MultithreadExecutor::MultithreadExecutor(int n) : numWorker(n) { MultithreadExecutor::~MultithreadExecutor() { Shutdown(); } -std::vector> MultithreadExecutor::SplitWork(const ExecutorReq& req) { - int numEps = req.eps.size(); - int totalBatchSize = req.sizes.size(); - +std::vector> SplitBatchWork(int numEps, int numWorker, int totalBatchSize) { assert(numEps > 0); int numActiveWorkers = std::min(numEps, numWorker); @@ -200,6 +213,32 @@ std::vector> MultithreadExecutor::SplitWork(const ExecutorRe return splits; } +std::vector> MultithreadExecutor::SplitWork(const ExecutorReq& req) { + return SplitBatchWork(static_cast(req.eps.size()), numWorker, + static_cast(req.sizes.size())); +} + +RdmaOpRet MultithreadExecutor::AggregateResults(std::vector>& futs) { + bool hasFail = false; + int numSucc = 0; + RdmaOpRet failedRet; + for (auto& fut : futs) { + RdmaOpRet ret = fut.get(); + if (ret.Failed()) { + hasFail = true; + failedRet = ret; + } else if (ret.Succeeded()) { + numSucc++; + } + } + if (hasFail) return failedRet; + + if (numSucc == static_cast(futs.size())) { + return {StatusCode::SUCCESS, ""}; + } + return {StatusCode::IN_PROGRESS, ""}; +} + RdmaOpRet MultithreadExecutor::RdmaBatchReadWrite(const ExecutorReq& req) { MORI_IO_FUNCTION_TIMER; @@ -219,26 +258,37 @@ RdmaOpRet MultithreadExecutor::RdmaBatchReadWrite(const ExecutorReq& req) { pool[epId % numWorker]->Submit(std::move(task)); } - bool hasFail = false; - int numSucc = 0; - RdmaOpRet failedRet; - for (auto& fut : futs) { - RdmaOpRet ret = fut.get(); - if (ret.Failed()) { - hasFail = true; - failedRet = ret; - } else if (ret.Succeeded()) { - numSucc++; - } - } - if (hasFail) return failedRet; + RdmaOpRet ret = AggregateResults(futs); + MORI_IO_TRACE("MultithreadExecutor submit request for RdmaBatchReadWrite done"); + return ret; +} - if (numSucc == numSplits) { - return {StatusCode::SUCCESS, ""}; +RdmaOpRet MultithreadExecutor::PostPrepared(const ExecutorPreparedReq& req) { + MORI_IO_FUNCTION_TIMER; + + int numSlices = static_cast(req.slices.size()); + int numEps = static_cast(req.eps.size()); + if (numSlices == 0) return {StatusCode::SUCCESS, ""}; + if (numEps == 0) return {StatusCode::ERR_INVALID_ARGS, "no endpoints"}; + + // Slices are endpoint-agnostic (addresses come from the session's base MRs and + // lkey/rkey are re-armed per post), so the same id-based rotation the build path + // uses still spreads small transfers across every QP. + int epOffset = static_cast(req.id % static_cast(numEps)); + std::vector> futs; + futs.reserve(numSlices); + + for (int i = 0; i < numSlices; i++) { + int epId = (i + epOffset) % numEps; + Task task{&req, epId, i}; + futs.push_back(std::move(task.ret.get_future())); + // Keep each QP owned by a stable worker to preserve QP affinity. + pool[epId % numWorker]->Submit(std::move(task)); } - MORI_IO_TRACE("MultithreadExecutor submit request for RdmaBatchReadWrite done"); - return {StatusCode::IN_PROGRESS, ""}; + RdmaOpRet ret = AggregateResults(futs); + MORI_IO_TRACE("MultithreadExecutor submit request for PostPrepared done"); + return ret; } void MultithreadExecutor::Start() { diff --git a/src/io/rdma/executor.hpp b/src/io/rdma/executor.hpp index 5d5ea0381..939c6e25f 100644 --- a/src/io/rdma/executor.hpp +++ b/src/io/rdma/executor.hpp @@ -26,6 +26,8 @@ #include #include #include +#include +#include #include "mori/application/transport/rdma/rdma.hpp" #include "mori/io/common.hpp" @@ -52,6 +54,24 @@ struct ExecutorReq { int maxChunks{1}; }; +// Prepared (build-once, post-many) counterpart of ExecutorReq. PrepareBatch has +// already split the batch into one self-contained single-endpoint WR list per +// slice, so a worker only re-arms and posts its own slice: slices share no mutable +// state, and the endpoint a slice lands on is chosen per post by the caller. +struct ExecutorPreparedReq { + const EpPairVec& eps; + const application::RdmaMemoryRegion& local; + const application::RdmaMemoryRegion& remote; + std::vector& slices; + std::shared_ptr callbackMeta; + TransferUniqueId id; +}; + +// Split `totalBatchSize` descriptors into at most min(numEps, numWorker) contiguous +// [begin, end) slices. Shared by the build-and-post path and the prepared path so +// that both decompose a batch identically at a given worker count. +std::vector> SplitBatchWork(int numEps, int numWorker, int totalBatchSize); + /* ---------------------------------------------------------------------------------------------- */ /* Executor */ /* ---------------------------------------------------------------------------------------------- */ @@ -63,6 +83,10 @@ class Executor { virtual void Start() = 0; virtual void Shutdown() = 0; virtual RdmaOpRet RdmaBatchReadWrite(const ExecutorReq& req) = 0; + virtual RdmaOpRet PostPrepared(const ExecutorPreparedReq& req) = 0; + // Number of worker threads actually available, which PrepareBatch needs to + // decide how many slices to build. + virtual int NumWorkers() const = 0; }; /* ---------------------------------------------------------------------------------------------- */ @@ -74,22 +98,31 @@ class MultithreadExecutor : public Executor { ~MultithreadExecutor(); RdmaOpRet RdmaBatchReadWrite(const ExecutorReq& req); + RdmaOpRet PostPrepared(const ExecutorPreparedReq& req); + int NumWorkers() const { return numWorker; } void Start(); void Shutdown(); private: struct Task { - const ExecutorReq* req; + // Exactly one of `req` (build a descriptor range, then post it) or `preq` + // (post an already-built slice) is set. + const ExecutorReq* req{nullptr}; + const ExecutorPreparedReq* preq{nullptr}; int epId{-1}; int begin{-1}; int end{-1}; + int sliceId{-1}; std::promise ret; Task(const ExecutorReq* req_, int epId_, int begin_, int end_) : req(req_), epId(epId_), begin(begin_), end(end_) {} + Task(const ExecutorPreparedReq* preq_, int epId_, int sliceId_) + : preq(preq_), epId(epId_), sliceId(sliceId_) {} }; std::vector> SplitWork(const ExecutorReq& req); + static RdmaOpRet AggregateResults(std::vector>& futs); class Worker { public: diff --git a/src/pybind/pybind_io.cpp b/src/pybind/pybind_io.cpp index b659572eb..5ae5a7539 100644 --- a/src/pybind/pybind_io.cpp +++ b/src/pybind/pybind_io.cpp @@ -212,9 +212,6 @@ void RegisterMoriIo(pybind11::module_& m) { return out.get().as(); }); - py::class_>( - m, "PreparedTransfer"); - py::class_(m, "IOEngineSession") .def("AllocateTransferUniqueId", &mori::io::IOEngineSession::AllocateTransferUniqueId, py::call_guard()) @@ -240,17 +237,6 @@ void RegisterMoriIo(pybind11::module_& m) { py::gil_scoped_release release; self.BatchWrite(lo, ro, sz, status, id); }) - .def("PrepareBatch", - [](mori::io::IOEngineSession& self, const SizeArray& localOffsets, - const SizeArray& remoteOffsets, const SizeArray& sizes, bool isRead) { - mori::io::SizeVec lo = SizeArrayToVec(localOffsets); - mori::io::SizeVec ro = SizeArrayToVec(remoteOffsets); - mori::io::SizeVec sz = SizeArrayToVec(sizes); - py::gil_scoped_release release; - return self.PrepareBatch(lo, ro, sz, isRead); - }) - .def("PostPrepared", &mori::io::IOEngineSession::PostPrepared, - py::call_guard()) .def("Alive", &mori::io::IOEngineSession::Alive); py::class_(m, "IOEngine") diff --git a/tests/cpp/io/bench_engine.cpp b/tests/cpp/io/bench_engine.cpp index 379d64241..043da03c3 100644 --- a/tests/cpp/io/bench_engine.cpp +++ b/tests/cpp/io/bench_engine.cpp @@ -1,3 +1,24 @@ +// 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. // bench_engine.cpp // // A C++ benchmark for MORI-IO whose measurement methodology exactly matches @@ -29,118 +50,146 @@ // TRANSFER, matching nixl's total_duration/(per_thread_iter*batch_size) // (nixl counts block_size*batch_size descriptors per request). // -// Rendezvous: 2 processes (one per node), TCP socket exchange of EngineDesc and -// MemoryDesc (msgpack, same as MORI's pybind pack()/unpack()). Rank 0 = -// initiator, rank 1 = target. Initiator drives all transfers (RDMA one-sided -// WRITE/READ); target only registers memory and waits. +// Rendezvous: 2 processes (one per node) exchange EngineDesc and MemoryDesc over +// mori::application::SocketBootstrapNetwork (msgpack, same as MORI's pybind +// pack()/unpack()). +// Rank 0 = initiator, rank 1 = target. Initiator drives all transfers (RDMA +// one-sided WRITE/READ); target only registers memory and waits. // // Build: see build.sh (links libmori_io + libmori_application from the editable // install, includes 3rdparty/msgpack-c + HIP). -#include -#include -#include -#include -#include +#include #include #include #include +#include #include +#include #include #include +#include #include #include -#include #include #include -#include -#include - +#include "mori/application/bootstrap/socket_bootstrap.hpp" #include "mori/io/io.hpp" using namespace mori::io; using Clock = std::chrono::steady_clock; -#define HIP_CHECK(expr) \ - do { \ - hipError_t _e = (expr); \ - if (_e != hipSuccess) { \ - std::cerr << "HIP error " << hipGetErrorString(_e) << " at " << __FILE__ << ":" \ - << __LINE__ << std::endl; \ - std::exit(1); \ - } \ +#define HIP_CHECK(expr) \ + do { \ + hipError_t _e = (expr); \ + if (_e != hipSuccess) { \ + std::cerr << "HIP error " << hipGetErrorString(_e) << " at " << __FILE__ << ":" << __LINE__ \ + << std::endl; \ + std::exit(1); \ + } \ } while (0) -// ------------------------- tiny TCP rendezvous ----------------------------- -// Rank 0 listens, rank 1 connects. Then symmetric length-prefixed blob swap. - -static int TcpListenAccept(uint16_t port) { - int lfd = socket(AF_INET, SOCK_STREAM, 0); - int opt = 1; - setsockopt(lfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)); - sockaddr_in addr{}; - addr.sin_family = AF_INET; - addr.sin_addr.s_addr = INADDR_ANY; - addr.sin_port = htons(port); - if (bind(lfd, reinterpret_cast(&addr), sizeof(addr)) != 0) { - perror("bind"); - std::exit(1); +// ------------------------------ validation --------------------------------- +// The Python benchmark checks correctness as part of the run (benchmark.py +// _validate_rdma): it performs a transfer, ships the target's buffer back over +// the control plane and byte-compares it against the initiator's source. We do +// the same check here, but exchange a per-slot checksum instead of the payload, +// so validating a multi-hundred-MiB sweep point stays O(batch) on the wire while +// still covering every transferred byte. +// +// Both buffers are seeded with a rank-dependent, offset-dependent pattern, so an +// un-issued transfer, a short transfer, or one landing at the wrong offset all +// leave the two sides disagreeing. + +static uint64_t Fnv1a(const uint8_t* p, size_t n, uint64_t h = 1469598103934665603ull) { + for (size_t i = 0; i < n; ++i) { + h ^= p[i]; + h *= 1099511628211ull; } - listen(lfd, 1); - int fd = accept(lfd, nullptr, nullptr); - close(lfd); - return fd; + return h; } -static int TcpConnect(const std::string& host, uint16_t port) { - int fd = socket(AF_INET, SOCK_STREAM, 0); - sockaddr_in addr{}; - addr.sin_family = AF_INET; - addr.sin_port = htons(port); - inet_pton(AF_INET, host.c_str(), &addr.sin_addr); - for (int retry = 0; retry < 100; ++retry) { - if (connect(fd, reinterpret_cast(&addr), sizeof(addr)) == 0) return fd; - std::this_thread::sleep_for(std::chrono::milliseconds(200)); +// Byte at absolute offset p is a function of (p, rank), so the two ranks start +// out differing everywhere and every byte position is distinguishable. +static void FillPattern(void* buf, size_t bytes, int rank) { + constexpr size_t kTile = 1u << 20; + std::vector tile(std::min(bytes, kTile)); + for (size_t base = 0; base < bytes; base += tile.size()) { + const size_t n = std::min(tile.size(), bytes - base); + for (size_t i = 0; i < n; ++i) { + const uint64_t p = base + i; + tile[i] = static_cast((p * 2654435761ull + rank * 0x9E3779B9ull) >> 13); + } + HIP_CHECK(hipMemcpy(static_cast(buf) + base, tile.data(), n, hipMemcpyHostToDevice)); } - std::cerr << "connect failed to " << host << ":" << port << std::endl; - std::exit(1); } -static void SendAll(int fd, const void* buf, size_t n) { - const char* p = static_cast(buf); - while (n) { - ssize_t k = send(fd, p, n, 0); - if (k <= 0) { perror("send"); std::exit(1); } - p += k; n -= k; +// One checksum per transferred slot, so a mismatch names the slot that differs +// rather than just failing the whole point. +static std::vector SlotChecksums(const void* buf, const SizeVec& offsets, size_t msg) { + std::vector sums(offsets.size()); + std::vector host(msg); + for (size_t i = 0; i < offsets.size(); ++i) { + HIP_CHECK(hipMemcpy(host.data(), static_cast(buf) + offsets[i], msg, + hipMemcpyDeviceToHost)); + sums[i] = Fnv1a(host.data(), msg); } + return sums; } -static void RecvAll(int fd, void* buf, size_t n) { - char* p = static_cast(buf); - while (n) { - ssize_t k = recv(fd, p, n, 0); - if (k <= 0) { perror("recv"); std::exit(1); } - p += k; n -= k; + +// Forward declaration; defined after Rendezvous/Pack below. +struct Args; + +// ------------------------- control-plane rendezvous ------------------------- +// Adapter over MORI's SocketBootstrapNetwork -- the same bootstrap layer +// src/cco/cco_init.cpp and src/shmem/init.cpp use to bring up multi-node jobs -- +// so the benchmark inherits the library's rendezvous policy instead of carrying a +// second one: connect and accept retried over the MORI_BOOTSTRAP_TIMEOUT budget +// (300s default), a Barrier that blocks for as long as the initiator's sweep +// takes, failures raised as exceptions, and Finalize() from the destructor. +// +// Both ranks derive the same UniqueId from the master endpoint, so nothing has to +// carry it between them. The bootstrap picks its own local interface; set +// MORI_SOCKET_IFNAME when that choice has no route to the peer. +class Rendezvous { + public: + Rendezvous(int rank, const std::string& masterIp, uint16_t port) + : boot(mori::application::SocketBootstrapNetwork::GenerateUniqueId(masterIp, port), rank, + kWorldSize), + myRank(rank) { + boot.Initialize(); } -} -static void SendBlob(int fd, const std::string& b) { - uint64_t len = b.size(); - SendAll(fd, &len, sizeof(len)); - SendAll(fd, b.data(), len); -} -static std::string RecvBlob(int fd) { - uint64_t len = 0; - RecvAll(fd, &len, sizeof(len)); - std::string b(len, '\0'); - RecvAll(fd, b.data(), len); - return b; -} -static void Barrier(int fd) { // symmetric 1-byte ping-pong - char c = 'x'; - SendAll(fd, &c, 1); - RecvAll(fd, &c, 1); -} + + // Contribute this rank's blob, return the peer's. Allgather is fixed-size while + // packed descriptors are not, so sizes go first and the payload round pads every + // rank up to the larger of the two. + std::string ExchangeBlob(const std::string& mine) { + uint64_t sizes[kWorldSize] = {}; + uint64_t mySize = mine.size(); + boot.Allgather(&mySize, sizes, sizeof(mySize)); + + const size_t stride = std::max(sizes[0], sizes[1]); + if (stride == 0) return {}; + + std::vector sendBuf(stride, 0), recvBuf(stride * kWorldSize, 0); + std::memcpy(sendBuf.data(), mine.data(), mine.size()); + boot.Allgather(sendBuf.data(), recvBuf.data(), stride); + + const int peer = 1 - myRank; + return std::string(recvBuf.data() + peer * stride, sizes[peer]); + } + + void Barrier() { boot.Barrier(); } + + private: + // Rank 0 initiates, rank 1 targets. + static constexpr int kWorldSize = 2; + + mori::application::SocketBootstrapNetwork boot; + const int myRank; +}; template static std::string Pack(const T& v) { @@ -154,58 +203,116 @@ static T Unpack(const std::string& b) { return oh.get().as(); } +// Compare the transferred slots of both ranks' buffers after the sweep. Both +// sides derive the geometry from the same args, so the only thing that crosses +// the wire is one checksum per slot. Only the initiator reports. +// +// BOTH ranks must call this unconditionally, even when --skip-validate is set: +// the exchange is a collective, so a rank that returned early would leave its +// peer blocked in Allgather until the bootstrap timeout. Opting out therefore +// contributes an empty checksum vector rather than skipping the call, which also +// makes the flag safe to pass to only one side -- either empty vector downgrades +// the run to "skipped" instead of hanging or reporting a bogus failure. +static bool ValidateTransfer(Rendezvous& rdv, const void* buf, size_t msg, int batch, bool batched, + bool batchContiguous, int rank, bool wanted) { + std::vector mine; + if (wanted) { + // Mirror the offsets the sweep actually used: the batched path strides by + // msg+1 unless --batch-contiguous, while the singles path is always + // contiguous (same asymmetry as the Python bench's run_single_once). + const size_t stride = batched ? (batchContiguous ? msg : msg + 1) : msg; + SizeVec offsets(static_cast(batch)); + for (int i = 0; i < batch; ++i) offsets[i] = static_cast(i) * stride; + mine = SlotChecksums(buf, offsets, msg); + } + + const auto peer = Unpack>(rdv.ExchangeBlob(Pack(mine))); + if (rank != 0) return true; + + if (mine.empty() || peer.empty()) { + std::cout << "validation: skipped" + << (mine.empty() ? "" : " (peer opted out with --skip-validate)") << std::endl; + return true; + } + if (peer.size() != mine.size()) { + std::cerr << "VALIDATION FAILED: peer returned " << peer.size() << " checksums, expected " + << mine.size() << std::endl; + return false; + } + for (size_t i = 0; i < mine.size(); ++i) { + if (mine[i] != peer[i]) { + std::cerr << "VALIDATION FAILED: slot " << i << " of " << mine.size() + << " differs at msg=" << msg << " batch=" << batch << " (initiator " << std::hex + << mine[i] << " vs target " << peer[i] << std::dec << ")" << std::endl; + return false; + } + } + std::cout << "validation: OK (" << batch << " slot(s) x " << msg << " B byte-identical)" + << std::endl; + return true; +} + // ------------------------------- config ------------------------------------ // Flags mirror MORI's Python benchmark (tests/python/io/benchmark.py) so runs // are directly comparable. Names use the Python spelling where they exist. struct Args { - int rank = 0; // 0=initiator, 1=target - std::string master_ip; // rank 1 connects here (rank 0's data IP) for TCP rendezvous - std::string self_ip; // this rank's own reachable IP (advertised in EngineDesc) + int rank = 0; // 0=initiator, 1=target + std::string master_ip; // bootstrap root (rank 0's data IP); both ranks pass the same + std::string self_ip; // this rank's own reachable IP (advertised in EngineDesc) uint16_t port = 18515; int gpu = 0; int target_dev_offset = 0; // --target-dev-offset (target GPU = (gpu+offset)%ndev) std::string op = "write"; // write|read + std::string backend = "rdma"; // --backend rdma|xgmi (xgmi = intra-node GPU<->GPU) // Sweep control (Python parity). --all sweeps message size; --all-batch sweeps // batch size. Neither set => single run at (buffer_size, batch). - bool sweep_all = false; // --all - bool sweep_batch = false; // --all-batch - size_t buffer_size = 32768; // --buffer-size (single message size when not sweeping) - size_t sweep_start = 8; // --sweep-start-size (Python default 8) - size_t sweep_max = 1u << 20; // --sweep-max-size (Python default 2^20) - size_t sweep_step = 0; // --sweep-step (0 = geometric x2; >0 = linear +step) + bool sweep_all = false; // --all + bool sweep_batch = false; // --all-batch + size_t buffer_size = 32768; // --buffer-size (single message size when not sweeping) + size_t sweep_start = 8; // --sweep-start-size (Python default 8) + size_t sweep_max = 1u << 20; // --sweep-max-size (Python default 2^20) + size_t sweep_step = 0; // --sweep-step (0 = geometric x2; >0 = linear +step) int iters = 500; - int warmup = 50; // --warmup-iters + int warmup = 50; // --warmup-iters // RdmaBackendConfig knobs - int qp_per_transfer = 4; // --num-qp-per-transfer - int worker_threads = 1; // --num-worker-threads - int post_batch_size = -1; // --post-batch-size + int qp_per_transfer = 4; // --num-qp-per-transfer + int worker_threads = 1; // --num-worker-threads + int post_batch_size = -1; // --post-batch-size std::string poll_cq_mode = "polling"; // polling|event - bool disable_chunking = false; // --disable-chunking (default: chunking ON, like Python) - size_t chunk_bytes = 65536; // --chunk-bytes - int max_chunks = 64; // --max-chunks - int max_send_wr = 0; // --max-send-wr (0 = leave default) - int max_cqe_num = 0; // --max-cqe-num - int max_msg_sge = 0; // --max-msg-sge + bool disable_chunking = false; // --disable-chunking (default: chunking ON, like Python) + size_t chunk_bytes = 65536; // --chunk-bytes + int max_chunks = 64; // --max-chunks + int max_send_wr = 0; // --max-send-wr (0 = leave default) + int max_cqe_num = 0; // --max-cqe-num + int max_msg_sge = 0; // --max-msg-sge + + // XgmiBackendConfig knobs (--backend xgmi) + int num_streams = 64; // --num-streams + int num_events = 64; // --num-events // batch / session - int batch = 1; // --transfer-batch-size (transfers per request) + int batch = 1; // --transfer-batch-size (transfers per request) bool enable_batch_transfer = true; // --enable-batch-transfer / --disable-batch-transfer. - // ON (default): batch>1 => ONE N-descriptor batch request - // (nixl-equivalent). OFF: batch>1 => N individual single - // transfers per iteration (Python run_single_once path). - bool batch_contiguous = false; // --batch-contiguous (adjacent offsets → merged WR); - // default strided (each transfer a separate WR) - bool enable_sess = false; // --enable-sess (session fast-path); Python default: off - bool prepare_once = false; // --prepare-once (build WRs once, re-post each iter; nixl-style - // createXferReq/postXferReq split). Requires --enable-sess. + // ON (default): batch>1 => ONE N-descriptor batch request + // (nixl-equivalent). OFF: batch>1 => N individual single + // transfers per iteration (Python run_single_once path). + bool batch_contiguous = false; // --batch-contiguous (adjacent offsets → merged WR); + // default strided (each transfer a separate WR) + bool enable_sess = false; // --enable-sess (session fast-path); Python default: off + bool prepare_once = false; // --prepare-once (build WRs once, re-post each iter; nixl-style + // createXferReq/postXferReq split). Requires --enable-sess. std::string mem_type = "gpu"; // --mem-type gpu|cpu std::string init_mem_type; // --initiator-mem-type (empty => mem_type) std::string target_mem_type; // --target-mem-type (empty => mem_type) std::string log_level = "info"; // --log-level trace|debug|info|warning|error|critical + + // Correctness check on the last sweep point, on by default to match the Python + // benchmark (which always validates). --skip-validate opts out. + bool skip_validate = false; // --skip-validate }; static Args ParseArgs(int argc, char** argv) { @@ -213,59 +320,117 @@ static Args ParseArgs(int argc, char** argv) { for (int i = 1; i < argc; ++i) { std::string k = argv[i]; auto next = [&]() { return std::string(argv[++i]); }; - if (k == "--rank") a.rank = std::stoi(next()); - else if (k == "--master-ip") a.master_ip = next(); - else if (k == "--self-ip") a.self_ip = next(); - else if (k == "--port") a.port = static_cast(std::stoi(next())); - else if (k == "--gpu") a.gpu = std::stoi(next()); - else if (k == "--target-dev-offset") a.target_dev_offset = std::stoi(next()); - else if (k == "--op" || k == "--op-type") a.op = next(); - else if (k == "--all") a.sweep_all = true; - else if (k == "--all-batch") a.sweep_batch = true; - else if (k == "--buffer-size") a.buffer_size = std::stoull(next()); - else if (k == "--sweep-start" || k == "--sweep-start-size") a.sweep_start = std::stoull(next()); - else if (k == "--sweep-max" || k == "--sweep-max-size") a.sweep_max = std::stoull(next()); - else if (k == "--sweep-step") a.sweep_step = std::stoull(next()); - else if (k == "--iters") a.iters = std::stoi(next()); - else if (k == "--warmup" || k == "--warmup-iters") a.warmup = std::stoi(next()); - else if (k == "--qp-per-transfer" || k == "--num-qp-per-transfer") a.qp_per_transfer = std::stoi(next()); - else if (k == "--worker-threads" || k == "--num-worker-threads") a.worker_threads = std::stoi(next()); - else if (k == "--post-batch-size") a.post_batch_size = std::stoi(next()); - else if (k == "--poll_cq_mode" || k == "--poll-cq-mode") a.poll_cq_mode = next(); - else if (k == "--disable-chunking") a.disable_chunking = true; - else if (k == "--chunk-bytes") a.chunk_bytes = std::stoull(next()); - else if (k == "--max-chunks") a.max_chunks = std::stoi(next()); - else if (k == "--max-send-wr") a.max_send_wr = std::stoi(next()); - else if (k == "--max-cqe-num") a.max_cqe_num = std::stoi(next()); - else if (k == "--max-msg-sge") a.max_msg_sge = std::stoi(next()); - else if (k == "--batch" || k == "--transfer-batch-size") a.batch = std::stoi(next()); - else if (k == "--enable-batch-transfer") a.enable_batch_transfer = true; - else if (k == "--disable-batch-transfer") a.enable_batch_transfer = false; - else if (k == "--batch-contiguous") a.batch_contiguous = true; - else if (k == "--enable-sess") a.enable_sess = true; - else if (k == "--disable-sess") a.enable_sess = false; - else if (k == "--prepare-once") a.prepare_once = true; - else if (k == "--no-prepare-once") a.prepare_once = false; - else if (k == "--mem-type") a.mem_type = next(); - else if (k == "--initiator-mem-type") a.init_mem_type = next(); - else if (k == "--target-mem-type") a.target_mem_type = next(); - else if (k == "--log-level") a.log_level = next(); - else { std::cerr << "unknown arg " << k << std::endl; std::exit(1); } + if (k == "--rank") + a.rank = std::stoi(next()); + else if (k == "--master-ip") + a.master_ip = next(); + else if (k == "--self-ip") + a.self_ip = next(); + else if (k == "--port") + a.port = static_cast(std::stoi(next())); + else if (k == "--gpu") + a.gpu = std::stoi(next()); + else if (k == "--target-dev-offset") + a.target_dev_offset = std::stoi(next()); + else if (k == "--op" || k == "--op-type") + a.op = next(); + else if (k == "--backend") + a.backend = next(); + else if (k == "--all") + a.sweep_all = true; + else if (k == "--all-batch") + a.sweep_batch = true; + else if (k == "--buffer-size") + a.buffer_size = std::stoull(next()); + else if (k == "--sweep-start" || k == "--sweep-start-size") + a.sweep_start = std::stoull(next()); + else if (k == "--sweep-max" || k == "--sweep-max-size") + a.sweep_max = std::stoull(next()); + else if (k == "--sweep-step") + a.sweep_step = std::stoull(next()); + else if (k == "--iters") + a.iters = std::stoi(next()); + else if (k == "--warmup" || k == "--warmup-iters") + a.warmup = std::stoi(next()); + else if (k == "--qp-per-transfer" || k == "--num-qp-per-transfer") + a.qp_per_transfer = std::stoi(next()); + else if (k == "--worker-threads" || k == "--num-worker-threads") + a.worker_threads = std::stoi(next()); + else if (k == "--post-batch-size") + a.post_batch_size = std::stoi(next()); + else if (k == "--poll_cq_mode" || k == "--poll-cq-mode") + a.poll_cq_mode = next(); + else if (k == "--disable-chunking") + a.disable_chunking = true; + else if (k == "--chunk-bytes") + a.chunk_bytes = std::stoull(next()); + else if (k == "--max-chunks") + a.max_chunks = std::stoi(next()); + else if (k == "--max-send-wr") + a.max_send_wr = std::stoi(next()); + else if (k == "--max-cqe-num") + a.max_cqe_num = std::stoi(next()); + else if (k == "--max-msg-sge") + a.max_msg_sge = std::stoi(next()); + else if (k == "--num-streams") + a.num_streams = std::stoi(next()); + else if (k == "--num-events") + a.num_events = std::stoi(next()); + else if (k == "--batch" || k == "--transfer-batch-size") + a.batch = std::stoi(next()); + else if (k == "--enable-batch-transfer") + a.enable_batch_transfer = true; + else if (k == "--disable-batch-transfer") + a.enable_batch_transfer = false; + else if (k == "--batch-contiguous") + a.batch_contiguous = true; + else if (k == "--enable-sess") + a.enable_sess = true; + else if (k == "--disable-sess") + a.enable_sess = false; + else if (k == "--prepare-once") + a.prepare_once = true; + else if (k == "--mem-type") + a.mem_type = next(); + else if (k == "--initiator-mem-type") + a.init_mem_type = next(); + else if (k == "--target-mem-type") + a.target_mem_type = next(); + else if (k == "--skip-validate") + a.skip_validate = true; + else if (k == "--log-level") + a.log_level = next(); + else { + std::cerr << "unknown arg " << k << std::endl; + std::exit(1); + } } return a; } // ------------------------------ main --------------------------------------- -int main(int argc, char** argv) { +static int RunBenchmark(int argc, char** argv) { Args a = ParseArgs(argc, argv); SetLogLevel(a.log_level); + if (a.backend != "rdma" && a.backend != "xgmi") { + std::cerr << "--backend must be rdma or xgmi (got " << a.backend << ")" << std::endl; + return 1; + } + // Only the RDMA backend implements prepared transfers; elsewhere PrepareBatch + // hands back nothing and the run would quietly fall back to the inline path + // mid-sweep, reporting a prepared number that never ran prepared. + if (a.prepare_once && a.backend != "rdma") { + std::cerr << "--prepare-once requires the rdma backend (got " << a.backend << ")" << std::endl; + return 1; + } + // Per-role memory type: initiator (rank 0) / target (rank 1) may override the // shared --mem-type, enabling mixed CPU<->GPU transfers. Each process only // allocates its own side, so no cross-node coupling is needed. const std::string myMem = (a.rank == 0) - ? (!a.init_mem_type.empty() ? a.init_mem_type : a.mem_type) - : (!a.target_mem_type.empty() ? a.target_mem_type : a.mem_type); + ? (!a.init_mem_type.empty() ? a.init_mem_type : a.mem_type) + : (!a.target_mem_type.empty() ? a.target_mem_type : a.mem_type); const bool cpuMem = (myMem == "cpu"); const MemoryLocationType memLoc = cpuMem ? MemoryLocationType::CPU : MemoryLocationType::GPU; @@ -314,7 +479,13 @@ int main(int argc, char** argv) { } else { HIP_CHECK(hipMalloc(&buf, bufBytes)); } - HIP_CHECK(hipMemset(buf, 0, bufBytes)); + // Rank-dependent seed pattern rather than zeros, so the post-sweep validation + // can tell "the transfer moved my bytes" from "both buffers happened to match". + if (a.skip_validate) { + HIP_CHECK(hipMemset(buf, 0, bufBytes)); + } else { + FillPattern(buf, bufBytes, a.rank); + } // Out-of-band control endpoint for the MORI engine. host MUST be an IP the // peer can reach (advertised via EngineDesc for RDMA QP setup). Defaults to @@ -326,45 +497,61 @@ int main(int argc, char** argv) { std::string key = a.rank == 0 ? "initiator" : "target"; IOEngine engine(key, cfg); - RdmaBackendConfig rdmaCfg{}; - rdmaCfg.qpPerTransfer = a.qp_per_transfer; - rdmaCfg.postBatchSize = a.post_batch_size; - rdmaCfg.numWorkerThreads = a.worker_threads; - rdmaCfg.pollCqMode = (a.poll_cq_mode == "event") ? PollCqMode::EVENT : PollCqMode::POLLING; - rdmaCfg.enableNotification = false; // match MORI Python bench RDMA path - rdmaCfg.enableTransferChunking = !a.disable_chunking; // chunking ON by default (Python parity) - rdmaCfg.chunkBytes = a.chunk_bytes; - rdmaCfg.maxChunksPerTransfer = a.max_chunks; - if (a.max_send_wr > 0) rdmaCfg.maxSendWr = a.max_send_wr; - if (a.max_cqe_num > 0) rdmaCfg.maxCqeNum = a.max_cqe_num; - if (a.max_msg_sge > 0) rdmaCfg.maxMsgSge = a.max_msg_sge; - engine.CreateBackend(BackendType::RDMA, rdmaCfg); + // XGMI moves data over Infinity Fabric between two GPUs on the SAME host, so + // its knobs are stream/event depth rather than QPs and chunking. The RDMA-only + // flags above are simply unused on that path. + if (a.backend == "xgmi") { + XgmiBackendConfig xgmiCfg{}; + xgmiCfg.numStreams = a.num_streams; + xgmiCfg.numEvents = a.num_events; + engine.CreateBackend(BackendType::XGMI, xgmiCfg); + } else { + RdmaBackendConfig rdmaCfg{}; + rdmaCfg.qpPerTransfer = a.qp_per_transfer; + rdmaCfg.postBatchSize = a.post_batch_size; + rdmaCfg.numWorkerThreads = a.worker_threads; + rdmaCfg.pollCqMode = (a.poll_cq_mode == "event") ? PollCqMode::EVENT : PollCqMode::POLLING; + rdmaCfg.enableNotification = false; // match MORI Python bench RDMA path + rdmaCfg.enableTransferChunking = !a.disable_chunking; // chunking ON by default (Python parity) + rdmaCfg.chunkBytes = a.chunk_bytes; + rdmaCfg.maxChunksPerTransfer = a.max_chunks; + if (a.max_send_wr > 0) rdmaCfg.maxSendWr = a.max_send_wr; + if (a.max_cqe_num > 0) rdmaCfg.maxCqeNum = a.max_cqe_num; + if (a.max_msg_sge > 0) rdmaCfg.maxMsgSge = a.max_msg_sge; + engine.CreateBackend(BackendType::RDMA, rdmaCfg); + } MemoryDesc localMem = engine.RegisterMemory(buf, bufBytes, gpu, memLoc); - // --- rendezvous over a side TCP socket ----------------------------------- - int sock = (a.rank == 0) ? TcpListenAccept(a.port) : TcpConnect(a.master_ip, a.port); - int one = 1; - setsockopt(sock, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one)); + // --- rendezvous over MORI's socket bootstrap ------------------------------ + Rendezvous rdv(a.rank, a.master_ip, a.port); // Exchange EngineDesc then register the remote engine. EngineDesc myEng = engine.GetEngineDesc(); - SendBlob(sock, Pack(myEng)); - EngineDesc peerEng = Unpack(RecvBlob(sock)); + EngineDesc peerEng = Unpack(rdv.ExchangeBlob(Pack(myEng))); engine.RegisterRemoteEngine(peerEng); // Exchange MemoryDesc. - SendBlob(sock, Pack(localMem)); - MemoryDesc peerMem = Unpack(RecvBlob(sock)); + MemoryDesc peerMem = Unpack(rdv.ExchangeBlob(Pack(localMem))); - Barrier(sock); + rdv.Barrier(); if (a.rank == 1) { // Target: nothing to drive. RDMA one-sided ops complete without target CPU. // Just hold memory registered until the initiator says it's done. - Barrier(sock); // wait for initiator to finish the whole sweep - if (cpuMem) HIP_CHECK(hipHostFree(buf)); else HIP_CHECK(hipFree(buf)); - close(sock); + rdv.Barrier(); // wait for initiator to finish the whole sweep + // Collective with the initiator's call below, so it runs even under + // --skip-validate; the target only contributes checksums and does not report. + const auto& last = plan.back(); + ValidateTransfer(rdv, buf, last.first, last.second, a.enable_batch_transfer && last.second > 1, + a.batch_contiguous, a.rank, !a.skip_validate); + // Drop the NIC registration before releasing its backing pages, so the + // backend never holds an MR over freed storage. + engine.DeregisterMemory(localMem); + if (cpuMem) + HIP_CHECK(hipHostFree(buf)); + else + HIP_CHECK(hipFree(buf)); return 0; } @@ -376,7 +563,10 @@ int main(int argc, char** argv) { std::optional sessOpt; if (useSess) { sessOpt = engine.CreateSession(localMem, peerMem); - if (!sessOpt) { std::cerr << "CreateSession failed" << std::endl; std::exit(1); } + if (!sessOpt) { + std::cerr << "CreateSession failed" << std::endl; + std::exit(1); + } sessPtr = &*sessOpt; } const bool isRead = (a.op == "read"); @@ -386,15 +576,6 @@ int main(int argc, char** argv) { << std::endl; std::exit(1); } - // The prepared path posts inline from this thread, so a worker pool would be - // ignored and the reported thread count would not match what ran. QP-level - // parallelism is unaffected: --num-qp-per-transfer still applies. - if (a.prepare_once && a.worker_threads > 1) { - std::cerr << "--prepare-once requires --num-worker-threads 1 (the prepared path posts " - "inline from the calling thread; use --num-qp-per-transfer for QP parallelism)" - << std::endl; - std::exit(1); - } std::cout << "MsgSize(B) Batch Iters AvgBW(GB/s) AvgLat(us) TotalDur(us)" << std::endl; @@ -408,7 +589,7 @@ int main(int argc, char** argv) { // submissions per iteration (Python run_single_once), each its own status. // batch==1 is a single transfer either way. const bool batched = a.enable_batch_transfer && curBatch > 1; - const int perSlot = batched ? 1 : curBatch; // statuses per request + const int perSlot = batched ? 1 : curBatch; // statuses per request // Strict stop-and-wait: one request outstanding at a time (nixl // --pipeline_depth 1). [perSlot] status array (TransferStatus is // non-copyable, so a flat vector rather than nested). @@ -425,9 +606,14 @@ int main(int argc, char** argv) { offsets[i] = static_cast(i) * stride; sizes[i] = msg; } - // Engine (non-session) batch path needs vec-of-vec + desc vectors. + // Engine (non-session) batch path needs vec-of-vec + desc vectors. These and + // the status/id vectors below are built once per sweep point and reused every + // iteration: allocating them inside post() would charge two heap + // allocation/free pairs per iteration to the reported latency. MemDescVec locVec{localMem}, remVec{peerMem}; BatchSizeVec offVec{offsets}, sizeVec{sizes}; + TransferStatusPtrVec statusPtrs{&st[0]}; + TransferUniqueIdVec batchIds(1); // --prepare-once: build the merged/chunked WR list ONCE per (msg,batch) so // the timed loop only re-posts it (nixl createXferReq / postXferReq split). @@ -475,13 +661,17 @@ int main(int argc, char** argv) { TransferUniqueId id = useSess ? sessPtr->AllocateTransferUniqueId() : engine.AllocateTransferUniqueId(); if (useSess) { - if (isRead) sessPtr->BatchRead(offsets, offsets, sizes, &base[0], id); - else sessPtr->BatchWrite(offsets, offsets, sizes, &base[0], id); + if (isRead) + sessPtr->BatchRead(offsets, offsets, sizes, &base[0], id); + else + sessPtr->BatchWrite(offsets, offsets, sizes, &base[0], id); } else { - TransferStatusPtrVec sp{&base[0]}; - TransferUniqueIdVec ids{id}; - if (isRead) engine.BatchRead(locVec, offVec, remVec, offVec, sizeVec, sp, ids); - else engine.BatchWrite(locVec, offVec, remVec, offVec, sizeVec, sp, ids); + statusPtrs[0] = &base[0]; + batchIds[0] = id; + if (isRead) + engine.BatchRead(locVec, offVec, remVec, offVec, sizeVec, statusPtrs, batchIds); + else + engine.BatchWrite(locVec, offVec, remVec, offVec, sizeVec, statusPtrs, batchIds); } } else { // curBatch individual single-transfer submissions (Python run_single_once); @@ -492,11 +682,15 @@ int main(int argc, char** argv) { TransferUniqueId id = useSess ? sessPtr->AllocateTransferUniqueId() : engine.AllocateTransferUniqueId(); if (useSess) { - if (isRead) sessPtr->Read(off, off, msg, &base[i], id); - else sessPtr->Write(off, off, msg, &base[i], id); + if (isRead) + sessPtr->Read(off, off, msg, &base[i], id); + else + sessPtr->Write(off, off, msg, &base[i], id); } else { - if (isRead) engine.Read(localMem, off, peerMem, off, msg, &base[i], id); - else engine.Write(localMem, off, peerMem, off, msg, &base[i], id); + if (isRead) + engine.Read(localMem, off, peerMem, off, msg, &base[i], id); + else + engine.Write(localMem, off, peerMem, off, msg, &base[i], id); } } } @@ -521,7 +715,8 @@ int main(int argc, char** argv) { // ---- warmup (excluded from timing) ---- for (int w = 0; w < a.warmup; ++w) { post(); - while (!reqDone()) { /* spin: CQ worker thread stores status */ } + while (!reqDone()) { /* spin: CQ worker thread stores status */ + } if (TransferStatus* f = reqFailed()) { std::cerr << "warmup transfer failed: " << f->Message() << std::endl; std::exit(1); @@ -537,7 +732,8 @@ int main(int argc, char** argv) { post(); // spin-poll, mirroring nixl's "check status, if IN_PROG continue" scan // (status flags stored by MORI's CQ worker thread) - while (!reqDone()) { /* spin */ } + while (!reqDone()) { /* spin */ + } if (TransferStatus* f = reqFailed()) { std::cerr << "transfer failed: " << f->Message() << std::endl; std::exit(1); @@ -557,16 +753,44 @@ int main(int argc, char** argv) { const int effBatch = curBatch; const double numXfers = static_cast(a.iters) * effBatch; double total_bytes = static_cast(msg) * numXfers; - double avg_bw = (total_bytes / 1e9) / (total_us / 1e6); // GB/s, GB=10^9 - double avg_lat = total_us / numXfers; // us per single transfer + double avg_bw = (total_bytes / 1e9) / (total_us / 1e6); // GB/s, GB=10^9 + double avg_lat = total_us / numXfers; // us per single transfer std::printf("%-11zu %-6d %-6d %-12.2f %-11.2f %-.1f\n", msg, effBatch, a.iters, avg_bw, avg_lat, total_us); std::fflush(stdout); } - Barrier(sock); // tell target we're done - if (cpuMem) HIP_CHECK(hipHostFree(buf)); else HIP_CHECK(hipFree(buf)); - close(sock); - return 0; + rdv.Barrier(); // tell target we're done + + // Correctness check on the last sweep point, matching the Python benchmark's + // always-on validation. Runs after the timed region so it cannot perturb the + // reported numbers. + const auto& last = plan.back(); + const bool valid = ValidateTransfer(rdv, buf, last.first, last.second, + a.enable_batch_transfer && last.second > 1, + a.batch_contiguous, a.rank, !a.skip_validate); + + // Tear down in reverse order of construction: the session references the + // registered region, and the registration references the HIP allocation, so + // both must go before the pages are released. + sessOpt.reset(); + engine.DeregisterMemory(localMem); + if (cpuMem) + HIP_CHECK(hipHostFree(buf)); + else + HIP_CHECK(hipFree(buf)); + return valid ? 0 : 1; +} + +int main(int argc, char** argv) { + // Both the bootstrap layer and the IO engine report failures by throwing, so + // unwinding here runs the destructors instead of leaving teardown to an exit + // path. + try { + return RunBenchmark(argc, argv); + } catch (const std::exception& e) { + std::cerr << "bench_engine: " << e.what() << std::endl; + return 1; + } } diff --git a/tests/python/io/benchmark.py b/tests/python/io/benchmark.py index 417374eb0..7cb7a35d6 100644 --- a/tests/python/io/benchmark.py +++ b/tests/python/io/benchmark.py @@ -249,14 +249,6 @@ def parse_args(): "condition variable. Avoids the cross-thread cv-wakeup latency (~5-10us) " "per completion at the cost of burning a core while waiting.", ) - parser.add_argument( - "--prepare-once", - action="store_true", - help="Build the transfer work requests once (per size/batch) and re-post them " - "each iteration, hoisting descriptor-build cost out of the timed loop (the NIXL " - "createXferReq/postXferReq split). Falls back to the inline path if the backend " - "does not support prepared transfers.", - ) parser.add_argument( "--disable-chunking", action="store_true", @@ -379,7 +371,6 @@ def __init__( post_batch_size: int = -1, num_worker_threads: int = 1, busy_wait: bool = False, - prepare_once: bool = False, poll_cq_mode: str = "polling", max_send_wr: int = 0, max_cqe_num: int = 0, @@ -422,25 +413,6 @@ def __init__( self.post_batch_size = post_batch_size self.num_worker_threads = num_worker_threads self.busy_wait = busy_wait - # Prepared (build-once, post-many) handles live on a session, so - # --prepare-once requires --enable-sess. - if prepare_once and not self.enable_sess: - raise ValueError("--prepare-once requires --enable-sess") - # The prepared path posts inline from the calling thread, so a worker pool - # would be ignored. Reject it here instead of letting prepare_batch return - # None and silently falling back to the inline path mid-sweep. QP-level - # parallelism is unaffected (--num-qp-per-transfer still applies). - if prepare_once and self.num_worker_threads > 1: - raise ValueError( - "--prepare-once requires --num-worker-threads 1 (the prepared path posts " - "inline from the calling thread; use --num-qp-per-transfer for QP parallelism)" - ) - self.prepare_once = prepare_once - # Cache of prepared handles keyed by (buffer_size, transfer_batch_size). - # Built lazily on first run_*_once for a size, reused across warmup+timed - # iters, and cleared between sweep points via _reset_prepared(). - self._prepared_cache = {} - self._prepare_unsupported_warned = False self.poll_cq_mode = ( PollCqMode.POLLING if poll_cq_mode == "polling" else PollCqMode.EVENT ) @@ -1014,56 +986,6 @@ def _wait_status(self, status): else: status.Wait() - def _reset_prepared(self): - # Prepared handles capture offsets/sizes at build time, so they are only - # valid for the (buffer_size, batch) they were built for. Drop them when - # moving to a new sweep point. - self._prepared_cache = {} - - def _warn_prepare_unsupported(self): - if not self._prepare_unsupported_warned: - print( - "[bench] --prepare-once: backend returned no prepared handle; " - "falling back to the inline transfer path.", - flush=True, - ) - self._prepare_unsupported_warned = True - - def _get_prepared_batch( - self, buffer_size, transfer_batch_size, offsets, sizes, is_read - ): - # Build (once) and cache the prepared handle for the batched path. Returns - # None if the backend does not support prepared transfers (caller falls - # back to inline). Prepared transfers live on a session (--enable-sess). - key = ("batch", buffer_size, transfer_batch_size, is_read) - if key in self._prepared_cache: - return self._prepared_cache[key] - prepared = self.sess.prepare_batch(offsets, offsets, sizes, is_read) - if prepared is None: - self._warn_prepare_unsupported() - self._prepared_cache[key] = prepared - return prepared - - def _get_prepared_singles(self, buffer_size, transfer_batch_size, is_read): - # Build (once) and cache one 1-element prepared handle per transfer in the - # single-transfer path. Returns None (whole list) if unsupported. - key = ("single", buffer_size, transfer_batch_size, is_read) - if key in self._prepared_cache: - return self._prepared_cache[key] - prepared_list = [] - for i in range(transfer_batch_size): - offset = buffer_size * i - off = np.array([offset], dtype=np.uint64) - sz = np.array([buffer_size], dtype=np.uint64) - p = self.sess.prepare_batch(off, off, sz, is_read) - if p is None: - self._warn_prepare_unsupported() - self._prepared_cache[key] = None - return None - prepared_list.append(p) - self._prepared_cache[key] = prepared_list - return prepared_list - def run_single_once(self, buffer_size, transfer_batch_size): assert buffer_size <= self.buffer_size if ( @@ -1078,31 +1000,6 @@ def run_single_once(self, buffer_size, transfer_batch_size): for i in range(transfer_batch_size): transfer_uids.append(self.engine.allocate_transfer_uid()) - is_read = self.op_type == "read" - prepared_list = None - if self.prepare_once: - prepared_list = self._get_prepared_singles( - buffer_size, transfer_batch_size, is_read - ) - - if prepared_list is not None: - # Prepared (build-once, post-many): each single transfer's WR list was - # built outside the timed loop; only re-posting is timed here. - st = time.time() - for i in range(transfer_batch_size): - status = self.sess.post_prepared(prepared_list[i], transfer_uids[i]) - status_list.append(status) - launched = time.time() - for status in status_list: - self._wait_status(status) - done = time.time() - - launch_duration = launched - st - transfer_duration = done - launched - for status in status_list: - assert status.Succeeded(), f"Transfer failed: {status.Message()}" - return launch_duration, transfer_duration - func, arg_list = None, [] for i in range(transfer_batch_size): offset = buffer_size * i @@ -1160,19 +1057,7 @@ def run_batch_once(self, buffer_size, transfer_batch_size): sizes = np.full(transfer_batch_size, buffer_size, dtype=np.uint64) transfer_uid = self.engine.allocate_transfer_uid() - is_read = self.op_type == "read" - prepared = None - if self.prepare_once: - prepared = self._get_prepared_batch( - buffer_size, transfer_batch_size, offsets, sizes, is_read - ) - - if prepared is not None: - # Prepared (build-once, post-many): the WR list was built outside the - # timed loop; only the post is timed here. - st = time.time() - transfer_status = self.sess.post_prepared(prepared, transfer_uid) - elif self.enable_sess: + if self.enable_sess: func = ( self.sess.batch_read if self.op_type == "read" @@ -1222,9 +1107,6 @@ def run_once(self, buffer_size, transfer_batch_size): return self.run_single_once(buffer_size, transfer_batch_size) def _run_and_compute(self, buffer_size, transfer_batch_size, iters): - # Prepared handles capture offsets/sizes for a specific (size, batch); drop - # any from a previous sweep point so this point rebuilds them once. - self._reset_prepared() for _ in range(self.warmup_iters): self.run_once(buffer_size, transfer_batch_size) @@ -1484,7 +1366,6 @@ def benchmark_xgmi_worker(local_rank, node_rank, args): num_streams=args.num_streams, num_events=args.num_events, busy_wait=args.busy_wait, - prepare_once=args.prepare_once, xgmi_multiprocess=True, ) bench.print_config() @@ -1526,7 +1407,6 @@ def benchmark_engine(local_rank, node_rank, args): post_batch_size=args.post_batch_size, num_worker_threads=args.num_worker_threads, busy_wait=args.busy_wait, - prepare_once=args.prepare_once, poll_cq_mode=args.poll_cq_mode, max_send_wr=args.max_send_wr, max_cqe_num=args.max_cqe_num, @@ -1598,7 +1478,6 @@ def benchmark_xgmi(args): num_streams=args.num_streams, num_events=args.num_events, busy_wait=args.busy_wait, - prepare_once=args.prepare_once, xgmi_multiprocess=False, ) bench.print_config()