Chunk the D2H copy across N copy threads (#6068)

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py

@@ -4580,7 +4580,7 @@ def raw_embedding_stream(self) -> None:
                         # Lazy resize: runtime_meta shape/dtype is not known until
                         # the first data arrives from the MC module. Must use UVM
                         # (new_unified_tensor) because the C++ RawEmbeddingStreamer
-                        # reads this buffer via raw CPU pointers in tensor_copy().
+                        # reads this buffer via raw CPU pointers in tensor_copy_chunk().
                         self.register_buffer(
                             "res_runtime_meta",
                             torch.ops.fbgemm.new_unified_tensor(

fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h

@@ -13,6 +13,7 @@
 #endif
 
 #include <utility>
+#include <vector>
 
 #ifdef FBGEMM_FBCODE
 namespace facebook::aiplatform::gmpp::experimental::training_ps {
@@ -56,10 +57,12 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
 
   /// Stream out non-negative elements in <indices> and its paired embeddings
   /// from <weights> for the first <count> elements in the tensor.
-  /// It spins up a thread that will copy all 3 tensors to CPU and inject them
-  /// into the background queue which will be picked up by another set of thread
-  /// pools for streaming out to the thrift server (co-located on same host
-  /// now).
+  /// It spins up a dispatcher thread that copies the 4 tensors (indices,
+  /// weights, and the optional identities / runtime_meta) to CPU and injects
+  /// them into the background queue, which is drained by the stream thread
+  /// that streams out to the thrift server (co-located on same host
+  /// now). The copy is split into <= kChunkSize-row chunks across up to
+  /// kNumCopyThreads threads.
   ///
   /// This is used in cuda stream callback, which doesn't require to be
   /// serialized with other callbacks, thus a separate thread is used to
@@ -86,8 +89,8 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
       std::optional<at::Tensor> copy_done_flag = std::nullopt);
 
   /*
-   * Join the stream tensor copy thread, make sure the thread is properly
-   * finished before creating new.
+   * Join the pending dispatch (and the copy threads it spawned), making sure it
+   * is properly finished before creating new.
    */
   void join_stream_tensor_copy_thread();
 
@@ -97,16 +100,6 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
       const at::Tensor& weights,
       std::optional<at::Tensor> identities,
       std::optional<at::Tensor> runtime_meta);
-  /*
-   * Copy the indices, weights and count tensors and enqueue them for
-   * asynchronous stream.
-   */
-  void copy_and_enqueue_stream_tensors(
-      const at::Tensor& indices,
-      const at::Tensor& weights,
-      std::optional<at::Tensor> identities,
-      std::optional<at::Tensor> runtime_meta,
-      const at::Tensor& count);
 
   /*
    * FOR TESTING: Join the weight stream thread, make sure the thread is
@@ -130,20 +123,44 @@ class RawEmbeddingStreamer : public torch::jit::CustomClassHolder {
 #ifdef FBGEMM_FBCODE
   std::unique_ptr<std::thread> weights_stream_thread_;
   folly::UMPSCQueue<StreamQueueItem, true> weights_to_stream_queue_;
-  std::unique_ptr<std::thread> stream_tensor_copy_thread_;
+  // Copy threads for UVM cache (joined every iteration). Shared by the blocking
+  // and non-blocking stream() paths; this assumes a given table streams in a
+  // single mode at a time (blocking OR non-blocking), never concurrently.
+  std::vector<std::unique_ptr<std::thread>> chunk_copy_threads_;
+  std::unique_ptr<std::thread> dispatch_thread_;
   // OBC logger for RES silent-failure counters (res.fail.*). Emits to the
   // host-level OBC agent, so it reaches ODS from the trainer process without
   // per-process fb303 scrape config. Only constructed when streaming is on.
   std::unique_ptr<facebook::aiplatform::gmpp::experimental::training_ps::
                       TrainingPsOdsLogger>
       ods_logger_;
+
+  void join_worker_threads();
+  void chunked_copy_and_enqueue(
+      const at::Tensor& indices,
+      const at::Tensor& weights,
+      std::optional<at::Tensor> identities,
+      std::optional<at::Tensor> runtime_meta,
+      const at::Tensor& count,
+      std::vector<std::unique_ptr<std::thread>>& target_copy_threads);
 #endif
 };
 
-fbgemm_gpu::StreamQueueItem tensor_copy(
+fbgemm_gpu::StreamQueueItem tensor_copy_chunk(
     const at::Tensor& indices,
     const at::Tensor& weights,
     std::optional<at::Tensor> identities,
     std::optional<at::Tensor> runtime_meta,
-    const at::Tensor& count);
+    int64_t start_row,
+    int64_t end_row);
+
+// Tiles [0, num_rows) into per-thread groups of [start, end) chunk ranges, each
+// chunk of size <= chunk_size, contiguous and non-overlapping (union is the
+// whole range). The outer index is the thread; each inner vector is that
+// thread's contiguous band split into chunks. Empty bands produce no group.
+// Pure/build-agnostic so it is unit-testable without a GPU or FBGEMM_FBCODE.
+// num_threads bounds how the rows are pre-split before chunking, matching
+// chunked_copy_and_enqueue's tiling.
+std::vector<std::vector<std::pair<int64_t, int64_t>>>
+computeChunkRanges(int64_t num_rows, size_t chunk_size, size_t num_threads);
 } // namespace fbgemm_gpu