perf(umbp): build transfer plans incrementally

src/umbp/distributed/pool_client.cpp

@@ -3468,19 +3468,18 @@ std::vector<bool> PoolClient::BatchGetRanges(const std::vector<std::string>& key
   // key and one bad key does not fail the batch.  This is why the local half
   // needs no per-key threading the way whole-object BatchGet does: the copies
   // are already batched into a single plan set.
-  std::vector<TransferItem> local_items;
+  // Preserve the batch vector's original reservation for buffered engines.
+  // HBM builds plans directly and does not need this hint.
+  size_t range_total = 0;
+  for (const auto& per_key : sizes) range_total += per_key.size();
+  auto plan_builder = transfer_engine_->CreatePlanBuilder(range_total);
+  // Buffer one object until its complete page walk succeeds. This preserves
+  // the existing rollback: a bad object cannot leak partial work into a plan.
+  std::vector<TransferItem> key_items;
+  size_t local_item_count = 0;
   std::vector<size_t> local_keys;
   std::vector<size_t> missed;
   missed.reserve(n);
-  // A TransferItem carries two TransferRefs, each with a MemoryDesc, so growing
-  // this vector by doubling relocates a lot of bytes.  One range is at least one
-  // item -- more only when it straddles a page -- so the range count is a tight
-  // lower bound and covers the common case exactly.
-  {
-    size_t range_total = 0;
-    for (const auto& per_key : sizes) range_total += per_key.size();
-    local_items.reserve(range_total);
-  }
   std::vector<ObjectRange> key_ranges;
 
   // A key that asked for no ranges is never resolved -- there is nothing to
@@ -3526,7 +3525,7 @@ std::vector<bool> PoolClient::BatchGetRanges(const std::vector<std::string>& key
                       keys[i], resolved.size);
       continue;  // a caller bug, not a miss — do not go looking for a replica
     }
-    const size_t before = local_items.size();
+    key_items.clear();
     bool built = false;
     {
       PhaseTimer build_timer(dbg ? &dbg->build : nullptr);
@@ -3535,34 +3534,46 @@ std::vector<bool> PoolClient::BatchGetRanges(const std::vector<std::string>& key
       FillReadRanges(dsts[i], sizes[i], src_offsets[i], &key_ranges);
       built =
           resolved.file_ref.IsFile()
-              ? BuildLocalFileRangeTransfers(resolved.file_ref, key_ranges, /*tag=*/i, &local_items)
+              ? BuildLocalFileRangeTransfers(resolved.file_ref, key_ranges, /*tag=*/i, &key_items)
               : BuildLocalRangeTransfers(holder, resolved.pages, resolved.page_size, resolved.size,
                                          key_ranges, /*to_backend=*/false,
-                                         /*tag=*/i, &local_items, dbg ? &dbg->classify : nullptr);
+                                         /*tag=*/i, &key_items, dbg ? &dbg->classify : nullptr);
     }
     if (!built) {
       // The medium holds the key but cannot be read in-process.  Same call as
       // Same disposition the local get path takes on a size mismatch: look
       // reporting a miss the caller cannot distinguish from absence.
-      local_items.resize(before);
+      key_items.clear();
       missed.push_back(i);
       continue;
     }
+    {
+      PhaseTimer xfer_timer(dbg ? &dbg->xfer : nullptr);
+      PhaseTimer plan_timer(dbg ? &dbg->xfer_plan : nullptr);
+      for (const auto& item : key_items) plan_builder->Add(item);
+    }
+    local_item_count += key_items.size();
     local_keys.push_back(i);
   }
   if (dbg != nullptr) {
     // Classification is timed inside the build window, so subtract it out to
     // keep the phases disjoint and `other` meaningful.
     dbg->build -= dbg->classify;
-    dbg->items += local_items.size();
+    dbg->items += local_item_count;
   }
-  if (!local_items.empty()) {
+  if (local_item_count != 0) {
     std::vector<size_t> failed_tags;
     {
       PhaseTimer xfer_timer(dbg ? &dbg->xfer : nullptr);
       TransferEngine::StepTiming steps;
       if (dbg != nullptr) steps = {&dbg->xfer_plan, &dbg->xfer_submit, &dbg->xfer_wait};
-      transfer_engine_->Transfer(local_items, &failed_tags, dbg != nullptr ? &steps : nullptr);
+      TransferPlanSet planned;
+      {
+        PhaseTimer plan_timer(dbg ? &dbg->xfer_plan : nullptr);
+        planned = plan_builder->Finish();
+      }
+      transfer_engine_->Transfer(std::move(planned), &failed_tags,
+                                 dbg != nullptr ? &steps : nullptr);
     }
     const std::unordered_set<size_t> failed(failed_tags.begin(), failed_tags.end());
     for (size_t i : local_keys) {

src/umbp/distributed/transfer/composite_transfer_engine.cpp

@@ -193,6 +193,52 @@ bool CompositeTransferEngine::CanHandle(const TransferRef& src, const TransferRe
   return SelectEngine(src, dst) != nullptr;
 }
 
+class CompositeTransferEngine::PlanBuilder final : public TransferPlanBuilder {
+ public:
+  PlanBuilder(const CompositeTransferEngine& owner, size_t expected_items)
+      : owner_(owner), expected_items_(expected_items) {}
+  void Add(const TransferItem& item) override {
+    TransferEngine* engine = owner_.SelectEngine(item.src, item.dst);
+    if (engine == nullptr) {
+      rejected_.push_back(item.tag);
+      owner_.rejected_items_.fetch_add(1, std::memory_order_relaxed);
+      return;
+    }
+    auto it = builders_.find(engine);
+    if (it == builders_.end()) {
+      order_.push_back(engine);
+      it = builders_.emplace(engine, engine->CreatePlanBuilder(expected_items_)).first;
+    }
+    it->second->Add(item);
+  }
+  TransferPlanSet Finish() override {
+    TransferPlanSet out;
+    out.rejected_tags = std::move(rejected_);
+    for (TransferEngine* engine : order_) {
+      TransferPlanSet sub = builders_.at(engine)->Finish();
+      for (auto& plan : sub.plans) {
+        plan.engine = engine;
+        out.plans.push_back(std::move(plan));
+      }
+      out.rejected_tags.insert(out.rejected_tags.end(), sub.rejected_tags.begin(),
+                               sub.rejected_tags.end());
+    }
+    return out;
+  }
+
+ private:
+  const CompositeTransferEngine& owner_;
+  const size_t expected_items_;
+  std::unordered_map<TransferEngine*, std::unique_ptr<TransferPlanBuilder>> builders_;
+  std::vector<TransferEngine*> order_;
+  std::vector<size_t> rejected_;
+};
+
+std::unique_ptr<TransferPlanBuilder> CompositeTransferEngine::CreatePlanBuilder(
+    size_t expected_items) const {
+  return std::make_unique<PlanBuilder>(*this, expected_items);
+}
+
 TransferPlanSet CompositeTransferEngine::Plan(const std::vector<TransferItem>& items) const {
   TransferPlanSet out;
   if (items.empty()) return out;

src/umbp/distributed/transfer/hbm_copy_engine.cpp

@@ -358,19 +358,16 @@ bool HbmCopyEngine::CanHandle(const TransferRef& src, const TransferRef& dst) co
   return IsGpu(src) || IsGpu(dst);
 }
 
-TransferPlanSet HbmCopyEngine::Plan(const std::vector<TransferItem>& items) const {
-  TransferPlanSet out;
-  std::unordered_map<PtrPairKey, size_t, PtrPairKeyHash> pair_to_plan;
-  pair_to_plan.reserve(items.size() * 2);
-  // No segment may cross a host registration piece boundary on its host side;
-  // see HostRegistrationChunkBytes().
-  const size_t piece = HostRegistrationChunkBytes();
-
-  for (const auto& item : items) {
-    if (item.size == 0) continue;
-    if (!CanHandle(item.src, item.dst)) {
+namespace {
+class HbmPlanBuilder final : public TransferPlanBuilder {
+ public:
+  explicit HbmPlanBuilder(const HbmCopyEngine& engine)
+      : engine_(engine), host_registration_piece_(HostRegistrationChunkBytes()) {}
+  void Add(const TransferItem& item) override {
+    if (item.size == 0) return;
+    if (!engine_.CanHandle(item.src, item.dst)) {
       out.rejected_tags.push_back(item.tag);
-      continue;
+      return;
     }
     // Bounds are checked here, once, rather than per copy: a page index that
     // does not fit its buffer is a bug in the backend's slot bookkeeping, and
@@ -380,7 +377,7 @@ TransferPlanSet HbmCopyEngine::Plan(const std::vector<TransferItem>& items) cons
     if (item.size > item.src.size || item.src_offset > item.src.size - item.size ||
         item.size > item.dst.size || item.dst_offset > item.dst.size - item.size) {
       out.rejected_tags.push_back(item.tag);
-      continue;
+      return;
     }
 
     const PtrPairKey key{item.src.host_ptr, item.dst.host_ptr};
@@ -405,7 +402,7 @@ TransferPlanSet HbmCopyEngine::Plan(const std::vector<TransferItem>& items) cons
                                 : host_is_dst ? reinterpret_cast<uintptr_t>(item.dst.host_ptr)
                                               : 0;
     const auto host_piece = [&](size_t src_offset, size_t dst_offset) {
-      return (host_base + (host_is_src ? src_offset : dst_offset)) / piece;
+      return (host_base + (host_is_src ? src_offset : dst_offset)) / host_registration_piece_;
     };
     size_t so = static_cast<size_t>(item.src_offset);
     size_t dof = static_cast<size_t>(item.dst_offset);
@@ -414,7 +411,7 @@ TransferPlanSet HbmCopyEngine::Plan(const std::vector<TransferItem>& items) cons
       size_t sz = remaining;
       if (host_base != 0) {
         const uintptr_t host = host_base + (host_is_src ? so : dof);
-        sz = std::min<size_t>(sz, piece - host % piece);
+        sz = std::min<size_t>(sz, host_registration_piece_ - host % host_registration_piece_);
       }
       const bool can_coalesce =
           !plan.sizes.empty() &&
@@ -435,7 +432,29 @@ TransferPlanSet HbmCopyEngine::Plan(const std::vector<TransferItem>& items) cons
     }
     if (plan.tags.empty() || plan.tags.back() != item.tag) plan.tags.push_back(item.tag);
   }
-  return out;
+  TransferPlanSet Finish() override { return std::move(out); }
+
+ private:
+  const HbmCopyEngine& engine_;
+  const size_t host_registration_piece_;
+  TransferPlanSet out;
+  // Grows with distinct endpoint pairs, rather than reserving two hash
+  // buckets for every page fragment in the batch.
+  std::unordered_map<PtrPairKey, size_t, PtrPairKeyHash> pair_to_plan;
+};
+}  // namespace
+
+std::unique_ptr<TransferPlanBuilder> HbmCopyEngine::CreatePlanBuilder(
+    size_t /*expected_items*/) const {
+  return std::make_unique<HbmPlanBuilder>(*this);
+}
+
+TransferPlanSet HbmCopyEngine::Plan(const std::vector<TransferItem>& items) const {
+  auto builder = CreatePlanBuilder();
+  for (const auto& item : items) {
+    builder->Add(item);
+  }
+  return builder->Finish();
 }
 
 void HbmCopyEngine::AddHostGatherRegion(void* base, size_t bytes) {

src/umbp/distributed/transfer/transfer_engine.cpp

@@ -53,6 +53,26 @@ class StepTimer {
 
 }  // namespace
 
+namespace {
+class BufferedPlanBuilder final : public TransferPlanBuilder {
+ public:
+  BufferedPlanBuilder(const TransferEngine& engine, size_t expected_items) : engine_(engine) {
+    items_.reserve(expected_items);
+  }
+  void Add(const TransferItem& item) override { items_.push_back(item); }
+  TransferPlanSet Finish() override { return engine_.Plan(items_); }
+
+ private:
+  const TransferEngine& engine_;
+  std::vector<TransferItem> items_;
+};
+}  // namespace
+
+std::unique_ptr<TransferPlanBuilder> TransferEngine::CreatePlanBuilder(
+    size_t expected_items) const {
+  return std::make_unique<BufferedPlanBuilder>(*this, expected_items);
+}
+
 bool TransferEngine::Transfer(const std::vector<TransferItem>& items,
                               std::vector<size_t>* failed_tags, const StepTiming* timing) {
   if (items.empty()) return true;
@@ -62,6 +82,11 @@ bool TransferEngine::Transfer(const std::vector<TransferItem>& items,
     StepTimer plan_timer(timing ? timing->plan : nullptr);
     planned = Plan(items);
   }
+  return Transfer(std::move(planned), failed_tags, timing);
+}
+
+bool TransferEngine::Transfer(TransferPlanSet planned, std::vector<size_t>* failed_tags,
+                              const StepTiming* timing) {
   bool ok = planned.rejected_tags.empty();
   if (failed_tags != nullptr) {
     failed_tags->insert(failed_tags->end(), planned.rejected_tags.begin(),

src/umbp/include/umbp/distributed/transfer/composite_transfer_engine.h

@@ -79,6 +79,7 @@ class CompositeTransferEngine final : public TransferEngine {
   // Partition items by the engine that will carry them, delegate, and tag each
   // returned plan with its engine so Submit can route it back.
   TransferPlanSet Plan(const std::vector<TransferItem>& items) const override;
+  std::unique_ptr<TransferPlanBuilder> CreatePlanBuilder(size_t expected_items = 0) const override;
 
   std::unique_ptr<TransferHandle> Submit(std::vector<TransferPlan> plans) override;
 
@@ -95,6 +96,7 @@ class CompositeTransferEngine final : public TransferEngine {
 
  private:
   class FanOutHandle;
+  class PlanBuilder;
 
   // Per (engine, direction) accumulators.  Indexed by the engine's position in
   // engines_, which is fixed after composition, so the hot path is an array

src/umbp/include/umbp/distributed/transfer/hbm_copy_engine.h

@@ -112,6 +112,7 @@ class HbmCopyEngine final : public TransferEngine {
   // rather than just loop iterations.  This is the concrete instance of the
   // reason Plan() is virtual rather than a base-class helper.
   TransferPlanSet Plan(const std::vector<TransferItem>& items) const override;
+  std::unique_ptr<TransferPlanBuilder> CreatePlanBuilder(size_t expected_items = 0) const override;
 
   // Performs the copies inline; the returned handle is already settled.
   //

src/umbp/include/umbp/distributed/transfer/transfer_engine.h

@@ -326,6 +326,15 @@ struct TransferPlanSet {
   std::vector<size_t> rejected_tags;
 };
 
+// A call owns its builder. Adding fragments only plans work; it never
+// submits IO. Finish consumes the accumulated plans once.
+class TransferPlanBuilder {
+ public:
+  virtual ~TransferPlanBuilder() = default;
+  virtual void Add(const TransferItem& item) = 0;
+  virtual TransferPlanSet Finish() = 0;
+};
+
 // One failed plan, mapped back to the keys that rode in it.  Carries enough to
 // reproduce the pre-refactor per-transfer error log.
 struct TransferFailure {
@@ -405,6 +414,15 @@ class TransferEngine : public MemoryRegistrar, public MetricSource {
   // and fills `failed_tags` if anything was rejected or failed.
   bool Transfer(const std::vector<TransferItem>& items, std::vector<size_t>* failed_tags,
                 const StepTiming* timing = nullptr);
+
+  // Submit and wait for work that has already been planned.
+  bool Transfer(TransferPlanSet planned, std::vector<size_t>* failed_tags,
+                const StepTiming* timing = nullptr);
+
+  // Engines may build incrementally; the default preserves their existing
+  // batch planner by buffering items until Finish. expected_items is a
+  // reservation hint for buffered engines; incremental engines may ignore it.
+  virtual std::unique_ptr<TransferPlanBuilder> CreatePlanBuilder(size_t expected_items = 0) const;
 };
 
 }  // namespace mori::umbp