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. diff --git a/include/mori/io/common.hpp b/include/mori/io/common.hpp index 670f6c274..99e0e5d3c 100644 --- a/include/mori/io/common.hpp +++ b/include/mori/io/common.hpp @@ -165,6 +165,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. @@ -205,6 +233,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 f1ba2d269..e30ebde63 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 @@ -29,6 +30,28 @@ namespace py = pybind11; namespace mori { +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); @@ -108,6 +131,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") @@ -166,11 +191,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") @@ -185,9 +226,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 654bf0415..0637bcfd3 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 ( @@ -165,12 +166,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", @@ -201,6 +215,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, @@ -253,6 +273,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, @@ -266,7 +287,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, @@ -288,6 +311,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 @@ -302,7 +326,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 ) @@ -422,7 +448,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}") @@ -439,13 +467,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( @@ -582,7 +614,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_transfer_chunking=self.enable_chunking, @@ -677,13 +709,21 @@ def _initialize_xgmi(self): if self.enable_sess: self.sess = self.engine.create_session(self.mem, self.target_mem) + 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 if ( self.backend_type == "rdma" or (self.backend_type == "xgmi" and self.xgmi_multiprocess) ) and self.role is EngineRole.TARGET: - return 0 + return 0, 0 status_list = [] transfer_uids = [] @@ -721,13 +761,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 @@ -735,13 +779,13 @@ def run_batch_once(self, buffer_size, transfer_batch_size): self.backend_type == "rdma" 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: @@ -775,12 +819,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: @@ -789,26 +838,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 == "rdma" 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": @@ -831,6 +902,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(), ) @@ -843,11 +916,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, @@ -857,6 +934,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}", ] ) cur_size *= 2 @@ -868,10 +947,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( [ @@ -882,14 +967,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( [ @@ -900,6 +993,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}", ] ) @@ -980,6 +1075,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, @@ -991,6 +1087,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() @@ -1014,6 +1111,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, @@ -1027,7 +1125,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, @@ -1083,6 +1183,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, @@ -1092,6 +1193,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()