diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py index e76e5d03cd..f63553811e 100644 --- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py +++ b/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( diff --git a/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h b/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h index 9b0f7be6ee..432ad5a615 100644 --- a/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h +++ b/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h @@ -13,6 +13,7 @@ #endif #include +#include #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 and its paired embeddings /// from for the first 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 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 identities, std::optional 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 identities, - std::optional 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 weights_stream_thread_; folly::UMPSCQueue weights_to_stream_queue_; - std::unique_ptr 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> chunk_copy_threads_; + std::unique_ptr 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 ods_logger_; + + void join_worker_threads(); + void chunked_copy_and_enqueue( + const at::Tensor& indices, + const at::Tensor& weights, + std::optional identities, + std::optional runtime_meta, + const at::Tensor& count, + std::vector>& 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 identities, std::optional 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>> +computeChunkRanges(int64_t num_rows, size_t chunk_size, size_t num_threads); } // namespace fbgemm_gpu diff --git a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp index 0a887a4193..bd9b0f5b94 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/raw_embedding_streamer.cpp @@ -29,7 +29,12 @@ namespace { #ifdef FBGEMM_FBCODE // Timeout for copy_done_flag polling loop (microseconds). -constexpr int64_t kCopyDonePollTimeoutUs = 10'000'000; // 10 seconds +constexpr int64_t kCopyDonePollTimeoutUs = 10'000'000; + +// Max rows copied into one enqueued chunk. +constexpr size_t kChunkSize = 500000; +// Parallel chunk-copy threads spawned per non-blocking stream() call. +constexpr size_t kNumCopyThreads = 4; /* * Get the thrift client to the training parameter server service @@ -64,86 +69,114 @@ inline int64_t get_maybe_uvm_scalar(const at::Tensor& tensor) { } // namespace -fbgemm_gpu::StreamQueueItem tensor_copy( +fbgemm_gpu::StreamQueueItem tensor_copy_chunk( const at::Tensor& indices, const at::Tensor& weights, std::optional identities, std::optional runtime_meta, - const at::Tensor& count) { - auto num_sets = get_maybe_uvm_scalar(count); - auto new_indices = at::empty( - num_sets, at::TensorOptions().device(at::kCPU).dtype(indices.dtype())); + int64_t start_row, + int64_t end_row) { + int64_t n = end_row - start_row; + auto new_indices = + at::empty(n, at::TensorOptions().device(at::kCPU).dtype(indices.dtype())); auto new_weights = at::empty( - {num_sets, weights.size(1)}, + {n, weights.size(1)}, at::TensorOptions().device(at::kCPU).dtype(weights.dtype())); std::optional new_identities = std::nullopt; if (identities.has_value()) { new_identities = at::empty( - {num_sets, identities->size(1)}, + {n, identities->size(1)}, at::TensorOptions().device(at::kCPU).dtype(identities->dtype())); } std::optional new_runtime_meta = std::nullopt; if (runtime_meta.has_value()) { new_runtime_meta = at::empty( - {num_sets, runtime_meta->size(1)}, + {n, runtime_meta->size(1)}, at::TensorOptions().device(at::kCPU).dtype(runtime_meta->dtype())); } auto new_count = at::empty({1}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE( - weights.scalar_type(), "tensor_copy", [&] { + weights.scalar_type(), "tensor_copy_chunk", [&] { using value_t = scalar_t; FBGEMM_DISPATCH_INTEGRAL_TYPES( - indices.scalar_type(), "tensor_copy", [&] { + indices.scalar_type(), "tensor_copy_chunk", [&] { using index_t = scalar_t; - auto indices_addr = indices.const_data_ptr(); - auto new_indices_addr = new_indices.mutable_data_ptr(); std::copy( - indices_addr, - indices_addr + num_sets, - new_indices_addr); // dst_start - - auto weights_addr = weights.const_data_ptr(); - auto new_weights_addr = new_weights.mutable_data_ptr(); + indices.const_data_ptr() + start_row, + indices.const_data_ptr() + end_row, + new_indices.mutable_data_ptr()); std::copy( - weights_addr, - weights_addr + num_sets * weights.size(1), - new_weights_addr); // dst_start + weights.const_data_ptr() + + start_row * weights.size(1), + weights.const_data_ptr() + end_row * weights.size(1), + new_weights.mutable_data_ptr()); if (identities.has_value()) { FBGEMM_DISPATCH_INTEGRAL_TYPES( - identities->scalar_type(), "tensor_copy", [&] { - using identities_t = scalar_t; - const auto identities_addr = - identities->const_data_ptr(); - auto new_identities_addr = - new_identities->mutable_data_ptr(); + identities->scalar_type(), "tensor_copy_chunk", [&] { + using id_t = scalar_t; std::copy( - identities_addr, - identities_addr + num_sets * identities->size(1), - new_identities_addr); // dst_start + identities->const_data_ptr() + + start_row * identities->size(1), + identities->const_data_ptr() + + end_row * identities->size(1), + new_identities->mutable_data_ptr()); }); } if (runtime_meta.has_value()) { FBGEMM_DISPATCH_ALL_TYPES( - runtime_meta->scalar_type(), "tensor_copy", [&] { - using runtime_meta_t = scalar_t; - auto runtime_meta_addr = - runtime_meta->const_data_ptr(); - auto new_runtime_meta_addr = - new_runtime_meta->mutable_data_ptr(); + runtime_meta->scalar_type(), "tensor_copy_chunk", [&] { + using rm_t = scalar_t; std::copy( - runtime_meta_addr, - runtime_meta_addr + num_sets * runtime_meta->size(1), - new_runtime_meta_addr); // dst_start + runtime_meta->const_data_ptr() + + start_row * runtime_meta->size(1), + runtime_meta->const_data_ptr() + + end_row * runtime_meta->size(1), + new_runtime_meta->mutable_data_ptr()); }); } }); }); - *new_count.mutable_data_ptr() = num_sets; + *new_count.mutable_data_ptr() = n; return fbgemm_gpu::StreamQueueItem{ new_indices, new_weights, new_identities, new_runtime_meta, new_count}; } +std::vector>> +computeChunkRanges(int64_t num_rows, size_t chunk_size, size_t num_threads) { + // Split [0, num_rows) across up to num_threads contiguous per-thread bands, + // then split each band into <= chunk_size chunks. Returns one inner vector of + // [start, end) chunk ranges per thread (outer index = thread); ranges are + // contiguous, non-overlapping, and their union is the whole range. Empty + // bands produce no group. + std::vector>> thread_chunks; + if (num_rows <= 0 || chunk_size == 0 || num_threads == 0) { + return thread_chunks; + } + // ceil-div (a + b - 1) / b: rounds up so a partial final chunk/band counts. + const size_t n_chunks = + (static_cast(num_rows) + chunk_size - 1) / chunk_size; + const size_t n_threads = std::min(n_chunks, num_threads); + const size_t rows_per_thread = + (static_cast(num_rows) + n_threads - 1) / n_threads; + for (size_t ti = 0; ti < n_threads; ++ti) { + const int64_t thread_start = static_cast(ti * rows_per_thread); + const int64_t thread_end = + std::min(static_cast((ti + 1) * rows_per_thread), num_rows); + std::vector> chunks; + for (int64_t s = thread_start; s < thread_end; + s += static_cast(chunk_size)) { + const int64_t e = + std::min(s + static_cast(chunk_size), thread_end); + chunks.emplace_back(s, e); + } + if (!chunks.empty()) { + thread_chunks.push_back(std::move(chunks)); + } + } + return thread_chunks; +} + RawEmbeddingStreamer::RawEmbeddingStreamer( std::string unique_id, bool enable_raw_embedding_streaming, @@ -211,6 +244,7 @@ RawEmbeddingStreamer::~RawEmbeddingStreamer() { #ifdef FBGEMM_FBCODE if (enable_raw_embedding_streaming_) { join_stream_tensor_copy_thread(); + join_worker_threads(); join_weights_stream_thread(); } #endif @@ -241,7 +275,7 @@ void RawEmbeddingStreamer::stream( weights_to_stream_queue_.enqueue(stream_item); return; } - if (blocking_tensor_copy) { + auto poll_flag = [this, copy_done_flag]() { if (copy_done_flag.has_value()) { auto* ptr = static_cast(copy_done_flag->data_ptr()); folly::stop_watch poll_watch; @@ -249,59 +283,58 @@ void RawEmbeddingStreamer::stream( std::this_thread::yield(); if (poll_watch.elapsed().count() > kCopyDonePollTimeoutUs) { LOG(ERROR) << "[TBE_ID" << unique_id_ - << "] copy_done_flag blocking: poll timed out after " + << "] copy_done_flag poll timed out after " << kCopyDonePollTimeoutUs / 1'000'000 << "s"; - return; + return false; } } - *ptr = 0; // Reset for next iteration - } else { - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] copy_done_flag not provided, skipping wait (blocking)"; + *ptr = 0; + } + return true; + }; + + if (blocking_tensor_copy) { + if (!poll_flag()) { + return; } - copy_and_enqueue_stream_tensors( + chunked_copy_and_enqueue( indices, weights, std::move(identities), std::move(runtime_meta), - count); + count, + chunk_copy_threads_); + join_worker_threads(); return; } - // Make sure the previous thread is done before starting a new one + // Non-blocking: join the previous dispatch + copy threads, then spawn new + // ones. The join is the serializer: it guarantees iter i's copy finished + // reading the source cache rows before iter i+1 overwrites them. join_stream_tensor_copy_thread(); - // Cuda dispatches the host callbacks all in the same CPU thread. But the - // callbacks don't need to be serialized. - // So, We need to spin up a new thread to unblock the CUDA stream, so the CUDA - // can continue executing other host callbacks, eg. get/evict. - stream_tensor_copy_thread_ = std::make_unique([this, - copy_done_flag, - indices, - weights, - identities, - runtime_meta, - count]() { - if (copy_done_flag.has_value()) { - auto* ptr = static_cast(copy_done_flag->data_ptr()); - folly::stop_watch poll_watch; - while (*ptr == 0) { - std::this_thread::yield(); - if (poll_watch.elapsed().count() > kCopyDonePollTimeoutUs) { - LOG(ERROR) << "[TBE_ID" << unique_id_ - << "] copy_done_flag non-blocking: poll timed out after " - << kCopyDonePollTimeoutUs / 1'000'000 << "s"; - return; + dispatch_thread_ = std::make_unique( + [this, poll_flag, indices, weights, identities, runtime_meta, count]() { + // Guard the dispatcher body so a copy/enqueue failure logs instead of + // escaping the std::thread and calling std::terminate. + try { + if (!poll_flag()) { + return; + } + chunked_copy_and_enqueue( + indices, + weights, + identities, + runtime_meta, + count, + chunk_copy_threads_); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] stream dispatcher thread caught exception: " + << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ + << "] stream dispatcher thread caught unknown exception"; } - } - *ptr = 0; // Reset for next iteration - } else { - XLOG_EVERY_MS(INFO, 60000) - << "[TBE_ID" << unique_id_ - << "] copy_done_flag not provided, skipping wait (non-blocking)"; - } - copy_and_enqueue_stream_tensors( - indices, weights, identities, runtime_meta, count); - }); + }); rec->record.end(); #endif } @@ -310,15 +343,91 @@ void RawEmbeddingStreamer::join_stream_tensor_copy_thread() { #ifdef FBGEMM_FBCODE auto rec = torch::autograd::profiler::record_function_enter_new( "## RawEmbeddingStreamer::join_stream_tensor_copy_thread ##"); - if (stream_tensor_copy_thread_ != nullptr && - stream_tensor_copy_thread_->joinable()) { - stream_tensor_copy_thread_->join(); + if (dispatch_thread_ != nullptr && dispatch_thread_->joinable()) { + dispatch_thread_->join(); } + join_worker_threads(); rec->record.end(); #endif } #ifdef FBGEMM_FBCODE +void RawEmbeddingStreamer::join_worker_threads() { + for (auto& t : chunk_copy_threads_) { + if (t && t->joinable()) { + t->join(); + } + } + chunk_copy_threads_.clear(); +} + +void RawEmbeddingStreamer::chunked_copy_and_enqueue( + const at::Tensor& indices, + const at::Tensor& weights, + std::optional identities, + std::optional runtime_meta, + const at::Tensor& count, + std::vector>& target_copy_threads) { + auto rec = torch::autograd::profiler::record_function_enter_new( + "## RawEmbeddingStreamer::chunked_copy_and_enqueue ##"); + const auto num_rows = get_maybe_uvm_scalar(count); + const auto thread_chunks = + computeChunkRanges(num_rows, kChunkSize, kNumCopyThreads); + + for (auto& t : target_copy_threads) { // join+clear the previous batch + if (t && t->joinable()) { + t->join(); + } + } + target_copy_threads.clear(); + if (thread_chunks.empty()) { + rec->record.end(); + return; + } + + // One copy thread per pre-computed group. Chunk boundaries and per-thread + // grouping live entirely in computeChunkRanges, so the enqueued row set is + // identical regardless of how threads are laid out here. + target_copy_threads.reserve(thread_chunks.size()); + for (size_t ti = 0; ti < thread_chunks.size(); ++ti) { + target_copy_threads.push_back( + std::make_unique([this, + indices, + weights, + identities, + runtime_meta, + chunks = thread_chunks[ti], + ti]() { + // Guard the copy body so a per-chunk failure logs instead of escaping + // the std::thread and calling std::terminate. + try { + folly::stop_watch thread_watch; + int64_t rows_done = 0; + for (const auto& [s, e] : chunks) { + auto chunk_item = tensor_copy_chunk( + indices, weights, identities, runtime_meta, s, e); + weights_to_stream_queue_.enqueue(std::move(chunk_item)); + rows_done += (e - s); + } + XLOG_EVERY_MS(INFO, 15000) + << "[TBE_ID" << unique_id_ << "] copy_thread tid=" << ti + << " rows=" << rows_done << " chunks=" << chunks.size() + << " copy_ms=" << thread_watch.elapsed().count(); + } catch (const std::exception& e) { + XLOG(ERR) << "[TBE_ID" << unique_id_ << "] copy_thread tid=" << ti + << " caught exception: " << e.what(); + } catch (...) { + XLOG(ERR) << "[TBE_ID" << unique_id_ << "] copy_thread tid=" << ti + << " caught unknown exception"; + } + })); + } + XLOG_EVERY_MS(INFO, 15000) + << "[RES] chunked_copy tbe=" << unique_id_ << " rows=" << num_rows + << " threads=" << thread_chunks.size(); + rec->record.end(); +} + folly::coro::Task RawEmbeddingStreamer::tensor_stream( const at::Tensor& indices, const at::Tensor& weights, @@ -451,20 +560,6 @@ folly::coro::Task RawEmbeddingStreamer::tensor_stream( co_return; } -void RawEmbeddingStreamer::copy_and_enqueue_stream_tensors( - const at::Tensor& indices, - const at::Tensor& weights, - std::optional identities, - std::optional runtime_meta, - const at::Tensor& count) { - auto rec = torch::autograd::profiler::record_function_enter_new( - "## RawEmbeddingStreamer::copy_and_enqueue_stream_tensors ##"); - auto stream_item = tensor_copy( - indices, weights, std::move(identities), std::move(runtime_meta), count); - weights_to_stream_queue_.enqueue(stream_item); - rec->record.end(); -} - void RawEmbeddingStreamer::join_weights_stream_thread() { if (weights_stream_thread_ != nullptr && weights_stream_thread_->joinable()) { stop_ = true; diff --git a/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp b/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp index 50b6f28ada..2c17b9d262 100644 --- a/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp +++ b/fbgemm_gpu/src/split_embeddings_cache/tests/raw_embedding_streamer_test.cpp @@ -8,7 +8,7 @@ #include #include -#include "deeplearning/fbgemm/fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h" // @manual=//deeplearning/fbgemm/fbgemm_gpu/src/split_embeddings_cache:raw_embedding_streamer +#include "fbgemm_gpu/split_embeddings_cache/raw_embedding_streamer.h" #ifdef FBGEMM_FBCODE #include #include "aiplatform/gmpp/experimental/training_ps/gen-cpp2/TrainingParameterServerService.h" @@ -35,6 +35,11 @@ class MockTrainingParameterServerService }; #endif +// These object-constructing tests run in BOTH targets: the fbcode target links +// the flag-on :raw_embedding_streamer, and the no-fbcode target links the +// flag-off :raw_embedding_streamer_no_fbcode -- so each test's flag view of the +// header layout matches its library, avoiding the sizeof mismatch (an all-off +// test linked against the always-flag-on library would overrun the object). static std::unique_ptr getRawEmbeddingStreamer( const std::string& unique_id, @@ -84,6 +89,312 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithoutStreaming) { indices, weights, std::nullopt, std::nullopt, count, true, true); } +namespace { +// Row-major tensor with distinct, predictable values so a sliced copy is +// unambiguous: value at [r, c] == r * dim + c. +at::Tensor makeRowMajor(int64_t num_rows, int64_t dim, at::ScalarType dtype) { + return at::arange( + num_rows * dim, at::TensorOptions().device(at::kCPU).dtype(dtype)) + .reshape({num_rows, dim}); +} +} // namespace + +// tensor_copy_chunk is build-agnostic (defined outside FBGEMM_FBCODE). Expected +// tensors are constructed independently via at::slice, never copied from impl +// output. +TEST(RawEmbeddingStreamerTest, TensorCopyChunkFullRange) { + constexpr int64_t kNumRows = 5; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, /*start_row=*/0, kNumRows); + + EXPECT_TRUE(at::equal(item.indices, indices)); + EXPECT_TRUE(at::equal(item.weights, weights)); + EXPECT_FALSE(item.identities.has_value()); + EXPECT_FALSE(item.runtime_meta.has_value()); + const auto expected_count = at::tensor( + {kNumRows}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + EXPECT_TRUE(at::equal(item.count, expected_count)); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkNonZeroStartSlicesCorrectSlice) { + // Guards the start_row*dim / end_row*dim offset arithmetic in the copy. + constexpr int64_t kNumRows = 6; + constexpr int64_t kStart = 2; + constexpr int64_t kEnd = 5; // n == 3 + auto indices = at::tensor( + {10, 20, 30, 40, 50, 60}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_TRUE(at::equal(item.indices, indices.slice(0, kStart, kEnd))); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); + const auto expected_count = at::tensor( + {kEnd - kStart}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + EXPECT_TRUE(at::equal(item.count, expected_count)); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkCopiesIdentitiesAndRuntimeMeta) { + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; // n == 3 + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + auto identities = makeRowMajor(kNumRows, /*dim=*/2, at::kLong); + auto runtime_meta = makeRowMajor(kNumRows, /*dim=*/1, at::kLong); + + auto item = tensor_copy_chunk( + indices, weights, identities, runtime_meta, kStart, kEnd); + + ASSERT_TRUE(item.identities.has_value()); + ASSERT_TRUE(item.runtime_meta.has_value()); + EXPECT_TRUE(at::equal(*item.identities, identities.slice(0, kStart, kEnd))); + EXPECT_TRUE( + at::equal(*item.runtime_meta, runtime_meta.slice(0, kStart, kEnd))); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkAbsentOptionalsStayNullopt) { + constexpr int64_t kNumRows = 4; + auto indices = at::tensor( + {10, 20, 30, 40}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, /*start_row=*/0, kNumRows); + + EXPECT_FALSE(item.identities.has_value()); + EXPECT_FALSE(item.runtime_meta.has_value()); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkInt32IndicesDtype) { + // Coverage for the integral-index dispatch branch with int32 indices. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kInt)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_EQ(item.indices.scalar_type(), at::kInt); + EXPECT_TRUE(at::equal(item.indices, indices.slice(0, kStart, kEnd))); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkFloatRuntimeMeta) { + // tensor_copy_chunk must dispatch runtime_meta over all dtypes + // (FBGEMM_DISPATCH_ALL_TYPES): a float runtime_meta would throw if it were + // dispatched integral-only. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kFloat); + auto runtime_meta = makeRowMajor(kNumRows, /*dim=*/2, c10::kFloat); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, runtime_meta, kStart, kEnd); + + ASSERT_TRUE(item.runtime_meta.has_value()); + EXPECT_TRUE( + at::equal(*item.runtime_meta, runtime_meta.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkHalfWeights) { + // Coverage for the half (fp16) weights dispatch branch of + // FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE -- a common quantized serving dtype. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kHalf); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_EQ(item.weights.scalar_type(), at::kHalf); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +TEST(RawEmbeddingStreamerTest, TensorCopyChunkByteWeights) { + // Coverage for the byte (int8) weights dispatch branch of + // FBGEMM_DISPATCH_FLOAT_HALF_AND_BYTE -- the int8-quantized serving dtype. + constexpr int64_t kNumRows = 5; + constexpr int64_t kStart = 1; + constexpr int64_t kEnd = 4; + auto indices = at::tensor( + {10, 20, 30, 40, 50}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = makeRowMajor(kNumRows, EMBEDDING_DIMENSION, c10::kByte); + + auto item = tensor_copy_chunk( + indices, weights, std::nullopt, std::nullopt, kStart, kEnd); + + EXPECT_EQ(item.weights.scalar_type(), at::kByte); + EXPECT_TRUE(at::equal(item.weights, weights.slice(0, kStart, kEnd))); +} + +namespace { +// computeChunkRanges groups chunks per thread (outer index = thread). Flatten +// to the in-order chunk list so the coverage/contiguity invariants can be +// checked across the whole range regardless of the thread grouping. +std::vector> flatten( + const std::vector>>& + thread_chunks) { + std::vector> ranges; + for (const auto& chunks : thread_chunks) { + ranges.insert(ranges.end(), chunks.begin(), chunks.end()); + } + return ranges; +} + +// Structural invariants computeChunkRanges must always satisfy: ranges are +// contiguous + non-overlapping starting at 0, cover exactly [0, num_rows), and +// every chunk is non-empty and no larger than chunk_size. An off-by-one in the +// tiling arithmetic breaks at least one of these. +void expectValidChunkRanges( + const std::vector>& ranges, + int64_t num_rows, + int64_t chunk_size) { + int64_t cursor = 0; + for (const auto& [start, end] : ranges) { + EXPECT_EQ(start, cursor) << "ranges must be contiguous and non-overlapping"; + EXPECT_GT(end, start) << "no empty ranges"; + EXPECT_LE(end - start, chunk_size) << "each chunk must be <= chunk_size"; + cursor = end; + } + EXPECT_EQ(cursor, num_rows) << "ranges must cover exactly [0, num_rows)"; +} +} // namespace + +// computeChunkRanges (like tensor_copy_chunk) is build-agnostic. Expected +// ranges are constructed independently. +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesExactMultiple) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/8, /*chunk_size=*/4, /*num_threads=*/2)); + const std::vector> expected = {{0, 4}, {4, 8}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/8, /*chunk_size=*/4); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesRemainderChunk) { + // Single thread => pure chunking; last chunk carries the remainder. + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/10, /*chunk_size=*/4, /*num_threads=*/1)); + const std::vector> expected = { + {0, 4}, {4, 8}, {8, 10}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/10, /*chunk_size=*/4); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesCountLessThanChunkSize) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/3, /*chunk_size=*/10, /*num_threads=*/4)); + const std::vector> expected = {{0, 3}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/3, /*chunk_size=*/10); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesSingleChunk) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/5, /*num_threads=*/4)); + const std::vector> expected = {{0, 5}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/5); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesChunkSizeOne) { + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/4, /*chunk_size=*/1, /*num_threads=*/1)); + const std::vector> expected = { + {0, 1}, {1, 2}, {2, 3}, {3, 4}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/4, /*chunk_size=*/1); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesZeroRowsIsEmpty) { + EXPECT_TRUE( + computeChunkRanges(/*num_rows=*/0, /*chunk_size=*/4, /*num_threads=*/4) + .empty()); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesNumThreadsExceedsNumChunks) { + // n_threads is clamped to n_chunks, so exactly n_chunks groups are emitted, + // one chunk each, with no empty group. + const auto thread_chunks = + computeChunkRanges(/*num_rows=*/6, /*chunk_size=*/3, /*num_threads=*/10); + EXPECT_EQ(thread_chunks.size(), 2u) << "one group per chunk"; + const auto ranges = flatten(thread_chunks); + const std::vector> expected = {{0, 3}, {3, 6}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/6, /*chunk_size=*/3); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesThreadSplitThenChunk) { + // num_threads < num_chunks and rows_per_thread not a multiple of chunk_size: + // rows are pre-split into 2 per-thread bands ([0,50), [50,100)) and each band + // is then chunked by 30, so boundaries land at the thread split (50), not at + // 60. This locks the tiling to the original inline behavior. + const auto thread_chunks = computeChunkRanges( + /*num_rows=*/100, /*chunk_size=*/30, /*num_threads=*/2); + EXPECT_EQ(thread_chunks.size(), 2u) << "one group per thread band"; + const std::vector>> expected_groups = + {{{0, 30}, {30, 50}}, {{50, 80}, {80, 100}}}; + EXPECT_EQ(thread_chunks, expected_groups); + expectValidChunkRanges( + flatten(thread_chunks), /*num_rows=*/100, /*chunk_size=*/30); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesNoEmptyTailRange) { + // rows_per_thread rounds up (ceil(5/4)=2) so the 4th thread's band would be + // [6,5); that empty band must be dropped, leaving 3 non-empty groups. + const auto thread_chunks = + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/1, /*num_threads=*/4); + EXPECT_EQ(thread_chunks.size(), 3u) << "empty trailing band is dropped"; + const auto ranges = flatten(thread_chunks); + const std::vector> expected = { + {0, 1}, {1, 2}, {2, 3}, {3, 4}, {4, 5}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/1); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesOneOverChunk) { + // num_rows == chunk_size + 1: the +1 spills into a second, single-row chunk. + const auto ranges = flatten( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/4, /*num_threads=*/1)); + const std::vector> expected = {{0, 4}, {4, 5}}; + EXPECT_EQ(ranges, expected); + expectValidChunkRanges(ranges, /*num_rows=*/5, /*chunk_size=*/4); +} + +TEST(RawEmbeddingStreamerTest, ComputeChunkRangesZeroThreadsOrChunkSizeEmpty) { + // Defensive guards: chunk_size==0 and num_threads==0 would divide by zero in + // the ceil-div tiling, so both must short-circuit to an empty result. + EXPECT_TRUE( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/0, /*num_threads=*/4) + .empty()); + EXPECT_TRUE( + computeChunkRanges(/*num_rows=*/5, /*chunk_size=*/4, /*num_threads=*/0) + .empty()); +} + #ifdef FBGEMM_FBCODE TEST(RawEmbeddingStreamerTest, TestTensorStream) { std::vector table_names = {"tb1", "tb2", "tb3"}; @@ -178,14 +489,60 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopy) { indices, weights, std::nullopt, std::nullopt, count, true, true); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); - // Test non-blocking tensor copy + // Test non-blocking tensor copy. The copy runs on the dispatcher, so we must + // join_stream_tensor_copy_thread() before checking the queue -- asserting the + // size before the join would race the background copy. streamer->stream( indices, weights, std::nullopt, std::nullopt, count, true, false); - EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); streamer->join_stream_tensor_copy_thread(); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 2); } +TEST(RawEmbeddingStreamerTest, TestStreamWithCopyZeroCountEnqueuesNothing) { + // count <= 0 drives num_rows == 0, so chunked_copy_and_enqueue early-returns + // and nothing is enqueued. + std::vector table_names = {"tb1", "tb2", "tb3"}; + std::vector table_offsets = {0, 100, 300}; + std::vector table_sizes = {0, 50, 200, 300}; + + auto streamer = getRawEmbeddingStreamer( + "test_zero_count", true, table_names, table_offsets, table_sizes); + + auto mock_service = std::make_shared(); + auto mock_server = + std::make_shared( + mock_service, + "::1", + 0, + facebook::services::TLSConfig::applyDefaultsToThriftServer); + auto& mock_client_factory = + facebook::servicerouter::getMockSRClientFactory(false /* strict */); + mock_client_factory.registerMockService( + "realtime.delta.publish.esr", mock_server); + + auto indices = at::tensor( + {10, 2, 1, 150, 170, 230, 280}, + at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + auto weights = at::randn( + {indices.size(0), EMBEDDING_DIMENSION}, + at::TensorOptions().device(at::kCPU).dtype(c10::kFloat)); + auto count = + at::tensor({0}, at::TensorOptions().device(at::kCPU).dtype(at::kLong)); + + // Stop the consumer threads so the queue size is stable to read. + streamer->join_weights_stream_thread(); + + streamer->stream( + indices, + weights, + std::nullopt, + std::nullopt, + count, + /*require_tensor_copy=*/true, + /*blocking_tensor_copy=*/true); + EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 0); +} + TEST(RawEmbeddingStreamerTest, TestStreamE2E) { std::vector table_names = {"tb1", "tb2", "tb3"}; std::vector table_offsets = {0, 100, 300}; @@ -232,8 +589,13 @@ TEST(RawEmbeddingStreamerTest, TestStreamE2E) { streamer->stream( indices, weights, std::nullopt, std::nullopt, count, true, true); - // Make sure dequeue finished - std::this_thread::sleep_for(std::chrono::seconds(1)); + // Bounded wait for the consumer to drain the enqueued item (so + // co_setEmbeddings has run) before stopping the thread -- avoids a + // fixed-sleep flake and the stop_-between-peek-and-process race. + for (int i = 0; i < 1000 && streamer->get_weights_to_stream_queue_size() > 0; + ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } streamer->join_weights_stream_thread(); } @@ -431,5 +793,8 @@ TEST(RawEmbeddingStreamerTest, TestStreamWithCopyDoneFlagNonBlockingCopy) { // Wait for the async thread to complete streamer->join_stream_tensor_copy_thread(); EXPECT_EQ(streamer->get_weights_to_stream_queue_size(), 1); + // poll_flag() must have observed the flag (1) and reset it to 0; without the + // reset the next iteration would stream before the D2H copy finished. + EXPECT_EQ(copy_done_flag.item(), 0); } #endif