perf(io): add busy-wait completion mode and numpy batch-transfer mars…

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.

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<void()> 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);

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 <pybind11/numpy.h>
 #include <pybind11/operators.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
@@ -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<size_t, py::array::c_style | py::array::forcecast>;
+
+mori::io::SizeVec SizeArrayToVec(const SizeArray& arr) {
+  py::buffer_info info = arr.request();
+  const size_t* ptr = static_cast<const size_t*>(info.ptr);
+  return mori::io::SizeVec(ptr, ptr + info.size);
+}
+
+mori::io::BatchSizeVec SizeArraysToBatchVec(const std::vector<SizeArray>& 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<py::gil_scoped_release>())
       .def("WaitFor", &mori::io::TransferStatus::WaitFor, py::arg("timeout_ms") = -1,
+           py::call_guard<py::gil_scoped_release>())
+      .def("WaitBusy", &mori::io::TransferStatus::WaitBusy, py::arg("timeout_ms") = -1,
            py::call_guard<py::gil_scoped_release>());
 
   py::class_<mori::io::EngineDesc>(m, "EngineDesc")
@@ -166,11 +191,27 @@ void RegisterMoriIo(pybind11::module_& m) {
       .def("AllocateTransferUniqueId", &mori::io::IOEngineSession::AllocateTransferUniqueId,
            py::call_guard<py::gil_scoped_release>())
       .def("Read", &mori::io::IOEngineSession::Read, py::call_guard<py::gil_scoped_release>())
-      .def("BatchRead", &mori::io::IOEngineSession::BatchRead,
-           py::call_guard<py::gil_scoped_release>())
+      .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<py::gil_scoped_release>())
-      .def("BatchWrite", &mori::io::IOEngineSession::BatchWrite,
-           py::call_guard<py::gil_scoped_release>())
+      .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_<mori::io::IOEngine>(m, "IOEngine")
@@ -185,9 +226,29 @@ void RegisterMoriIo(pybind11::module_& m) {
       .def("AllocateTransferUniqueId", &mori::io::IOEngine::AllocateTransferUniqueId,
            py::call_guard<py::gil_scoped_release>())
       .def("Read", &mori::io::IOEngine::Read, py::call_guard<py::gil_scoped_release>())
-      .def("BatchRead", &mori::io::IOEngine::BatchRead, py::call_guard<py::gil_scoped_release>())
+      .def("BatchRead",
+           [](mori::io::IOEngine& self, const mori::io::MemDescVec& localDest,
+              const std::vector<SizeArray>& localOffsets, const mori::io::MemDescVec& remoteSrc,
+              const std::vector<SizeArray>& remoteOffsets, const std::vector<SizeArray>& 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<py::gil_scoped_release>())
-      .def("BatchWrite", &mori::io::IOEngine::BatchWrite, py::call_guard<py::gil_scoped_release>())
+      .def("BatchWrite",
+           [](mori::io::IOEngine& self, const mori::io::MemDescVec& localSrc,
+              const std::vector<SizeArray>& localOffsets, const mori::io::MemDescVec& remoteDest,
+              const std::vector<SizeArray>& remoteOffsets, const std::vector<SizeArray>& 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<py::gil_scoped_release>())