Compute required buffer size for SC

manuadg-verily · Google-ML-Automation · commit 2d0481f09079 · 2025-05-16T08:17:45.000-07:00
PiperOrigin-RevId: 759621757
diff --git a/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing.cc b/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing.cc
@@ -250,7 +250,8 @@ void PreprocessInputForStackedTablePerLocalDevice(
     const absl::string_view stacked_table_name, const bool allow_id_dropping,
     py::array_t<int> row_pointer_buffer, py::array_t<int> embedding_id_buffer,
     py::array_t<int> sample_id_buffer, py::array_t<float> gain_buffer,
-    py::array_t<int> max_ids_buffer, py::array_t<int> max_unique_ids_buffer) {
+    py::array_t<int> max_ids_buffer, py::array_t<int> max_unique_ids_buffer,
+    py::array_t<int> required_buffer_size_per_sc_buffer) {
   const int num_scs = num_sc_per_device * num_global_devices;
   int batch_size_for_device = 0;
   int total_num_coo_tensors = 0;
@@ -300,6 +301,8 @@ void PreprocessInputForStackedTablePerLocalDevice(
   auto* gains_data = gain_buffer.mutable_data();
   auto* total_max_ids_per_sc = max_ids_buffer.mutable_data();
   auto* total_max_unique_ids_per_sc = max_unique_ids_buffer.mutable_data();
+  auto* required_buffer_size_per_sc =
+      required_buffer_size_per_sc_buffer.mutable_data();
   // The remaining section does not require GIL.
   py::gil_scoped_release release;
 
@@ -308,21 +311,15 @@ void PreprocessInputForStackedTablePerLocalDevice(
   //
   const int batch_size_per_sc =
       CeilOfRatio(batch_size_for_device, num_sc_per_device);
-  std::vector<std::vector<CooFormat>> coo_tensors_by_id;
-  coo_tensors_by_id.resize(num_sc_per_device);
 
-  const int approximate_num_coo_tensors_per_sc =
-      total_num_coo_tensors / num_sc_per_device + 1;
-  for (int i = 0; i < num_sc_per_device; ++i) {
-    // Roughly estimate the number of COO tensors for each SC.
-    coo_tensors_by_id[i].reserve(approximate_num_coo_tensors_per_sc);
-  }
-  SortAndGroupCooTensorsPerLocalDevice(
-      coo_tensors, batch_size_per_sc, num_scs, batch_size_for_device,
-      stacked_table_metadata[0].max_ids_per_partition,
-      stacked_table_metadata[0].max_unique_ids_per_partition,
-      stacked_table_name, allow_id_dropping, coo_tensors_by_id,
-      total_max_ids_per_sc, total_max_unique_ids_per_sc);
+  std::vector<std::vector<CooFormat>> coo_tensors_by_id =
+      SortAndGroupCooTensorsPerLocalDevice(
+          coo_tensors, batch_size_per_sc, num_scs, batch_size_for_device,
+          stacked_table_metadata[0].max_ids_per_partition,
+          stacked_table_metadata[0].max_unique_ids_per_partition,
+          stacked_table_name, allow_id_dropping, num_sc_per_device,
+          total_num_coo_tensors, total_max_ids_per_sc,
+          total_max_unique_ids_per_sc, required_buffer_size_per_sc);
   for (int i = 0; i < num_sc_per_device; ++i) {
     coo_tensors_by_id[i].emplace_back(batch_size_per_sc * (i + 1), 0, 0.0);
   }
@@ -384,6 +381,7 @@ py::tuple PreprocessSparseDenseMatmulInput(
   py::dict lhs_gains;
   py::dict max_ids_per_partition;
   py::dict max_unique_ids_per_partition;
+  py::dict required_buffer_sizes;
   const int num_scs = num_sc_per_device * global_device_count;
   const int row_pointers_size_per_sc = std::max(num_scs, 8);
 
@@ -446,6 +444,8 @@ py::tuple PreprocessSparseDenseMatmulInput(
             py::array_t<int>(stats_shape);
         py::array_t<int> max_unique_ids_per_partition_per_sc =
             py::array_t<int>(stats_shape);
+        py::array_t<int> required_buffer_size_per_sc =
+            py::array_t<int>(stats_shape);
         for (int local_device = 0; local_device < local_device_count;
              ++local_device) {
           // Get the tuple outputs for the current split.
@@ -466,6 +466,8 @@ py::tuple PreprocessSparseDenseMatmulInput(
               max_ids_per_partition_per_sc[stats_slice];
           auto max_unique_ids_per_partition_per_sc_buffer =
               max_unique_ids_per_partition_per_sc[stats_slice];
+          auto required_buffer_size_per_sc_buffer =
+              required_buffer_size_per_sc[stats_slice];
           PreprocessInputForStackedTablePerLocalDevice(
               stacked_table_metadata, features, feature_weights, local_device,
               local_device_count, coo_buffer_size, row_pointers_size_per_sc,
@@ -477,7 +479,8 @@ py::tuple PreprocessSparseDenseMatmulInput(
               py::cast<py::array_t<float>>(gain_buffer),
               py::cast<py::array_t<int>>(max_ids_per_partition_per_sc_buffer),
               py::cast<py::array_t<int>>(
-                  max_unique_ids_per_partition_per_sc_buffer));
+                  max_unique_ids_per_partition_per_sc_buffer),
+              py::cast<py::array_t<int>>(required_buffer_size_per_sc_buffer));
         }
         lhs_row_pointers[stacked_table_name.c_str()] =
             std::move(row_pointers_per_device);
@@ -490,6 +493,8 @@ py::tuple PreprocessSparseDenseMatmulInput(
             std::move(max_ids_per_partition_per_sc);
         max_unique_ids_per_partition[stacked_table_name.c_str()] =
             std::move(max_unique_ids_per_partition_per_sc);
+        required_buffer_sizes[stacked_table_name.c_str()] =
+            std::move(required_buffer_size_per_sc);
         counter.DecrementCount();
       });
     }
@@ -498,6 +503,8 @@ py::tuple PreprocessSparseDenseMatmulInput(
   py::dict stats;
   stats["max_ids"] = max_ids_per_partition;
   stats["max_unique_ids"] = max_unique_ids_per_partition;
+  stats["required_buffer_size"] = std::move(required_buffer_sizes);
+
   // GIL is held at this point.
   return py::make_tuple(lhs_row_pointers, lhs_embedding_ids, lhs_sample_ids,
                         lhs_gains, stats);
diff --git a/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.cc b/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.cc
@@ -34,6 +34,49 @@
 #include "tsl/profiler/lib/traceme.h"
 
 namespace jax_sc_embedding {
+namespace {
+
+void ValidateMaxIdsOrDie(const int32_t observed_max_ids_per_partition,
+                         const int32_t observed_max_unique_ids_per_partition,
+                         const int32_t max_ids_per_partition,
+                         const int32_t max_unique_ids_per_partition,
+                         const absl::string_view stacked_table_name,
+                         const bool allow_id_dropping) {
+  // If id dropping is allowed, we log a warning if the observed max ids per
+  // partition is greater than the set max ids per partition.
+  if (observed_max_ids_per_partition > max_ids_per_partition) {
+    if (allow_id_dropping) {
+      LOG(WARNING) << "Allowing ID dropping for table: " << stacked_table_name
+                   << " observed max ids per partition: "
+                   << observed_max_ids_per_partition
+                   << " is greater than the set max ids per partition: "
+                   << max_ids_per_partition;
+    } else {
+      LOG(FATAL) << "Observed max ids per partition: "
+                 << observed_max_ids_per_partition
+                 << " for table: " << stacked_table_name
+                 << " is greater than the set max ids per partition: "
+                 << max_ids_per_partition;
+    }
+  }
+  if (observed_max_unique_ids_per_partition > max_unique_ids_per_partition) {
+    if (allow_id_dropping) {
+      LOG(WARNING) << "Allowing ID dropping for table: " << stacked_table_name
+                   << " observed max unique ids per partition: "
+                   << observed_max_unique_ids_per_partition
+                   << " is greater than the set max unique ids per partition: "
+                   << max_unique_ids_per_partition;
+    } else {
+      LOG(FATAL) << "Observed max unique ids per partition: "
+                 << observed_max_unique_ids_per_partition
+                 << " for table: " << stacked_table_name
+                 << " is greater than the set max unique ids per partition: "
+                 << max_unique_ids_per_partition;
+    }
+  }
+}
+
+}  // namespace
 
 int GetColId(const int col_id, const int col_shift, const int col_offset,
              const int num_scs_mod, const int num_scs_mod_inv) {
@@ -55,19 +98,28 @@ RowCombiner GetRowCombiner(absl::string_view combiner) {
   return RowCombiner::kSum;
 }
 
-void SortAndGroupCooTensorsPerLocalDevice(
+std::vector<std::vector<CooFormat>> SortAndGroupCooTensorsPerLocalDevice(
     absl::Span<const CooFormat> coo_tensors, const int batch_size_per_sc,
     const int global_sc_count, const int32_t batch_size_for_device,
     const int32_t max_ids_per_partition,
     const int32_t max_unique_ids_per_partition,
     const absl::string_view stacked_table_name, const bool allow_id_dropping,
-    std::vector<std::vector<CooFormat>>& coo_tensors_by_id, int* max_ids_per_sc,
-    int* max_unique_ids_per_sc) {
+    const int num_sc_per_device, const int total_num_coo_tensors,
+    int max_ids_per_sc[], int max_unique_ids_per_sc[],
+    int required_buffer_size_per_sc[]) {
   tsl::profiler::TraceMe t("SortAndGroupCooTensors");
   const int local_sc_count = batch_size_for_device / batch_size_per_sc;
-  uint32_t index = 0;
+  std::vector<std::vector<CooFormat>> coo_tensors_by_id;
+  coo_tensors_by_id.resize(num_sc_per_device);
+  const int approximate_num_coo_tensors_per_sc =
+      total_num_coo_tensors / num_sc_per_device + 1;
+  for (int i = 0; i < num_sc_per_device; ++i) {
+    // Roughly estimate the number of COO tensors for each SC.
+    coo_tensors_by_id[i].reserve(approximate_num_coo_tensors_per_sc);
+  }
+
+  uint32_t coo_tensor_index = 0;
   const int32_t num_scs_bit = std::log2(global_sc_count);
-  const int total_coo_tensors = coo_tensors.size();
   // Initialize the aggregated max ids and unique ids per SC to 0.
   for (int32_t global_sc_id = 0; global_sc_id < global_sc_count;
        ++global_sc_id) {
@@ -76,29 +128,30 @@ void SortAndGroupCooTensorsPerLocalDevice(
   }
   // Loop over scs for this device.
   for (int32_t local_sc_id = 0; local_sc_id < local_sc_count; ++local_sc_id) {
-    const int num_partitions = global_sc_count;
-    std::vector<int32_t> ids_per_sc_partition(num_partitions, 0);
-    std::vector<int32_t> unique_ids_per_sc_partition(num_partitions, 0);
+    std::vector<int32_t> ids_per_sc_partition(global_sc_count, 0);
+    std::vector<int32_t> unique_ids_per_sc_partition(global_sc_count, 0);
     std::vector<uint64_t> keys;
     keys.reserve(batch_size_per_sc);
     // We take the advantage of the fact that the row_ids are already sorted
     // within each batch.
-    while (index < total_coo_tensors &&
-           coo_tensors[index].row_id < (local_sc_id + 1) * batch_size_per_sc) {
+    for (; coo_tensor_index < coo_tensors.size() &&
+           coo_tensors[coo_tensor_index].row_id <
+               (local_sc_id + 1) * batch_size_per_sc;
+         coo_tensor_index++) {
       // The key here is [col_ids % num_scs, col_ids / num_scs, index].
       // Note that this assumes `num_scs` is a power of 2.
       keys.push_back(
           (static_cast<uint64_t>(absl::rotr(
-               static_cast<uint32_t>(coo_tensors[index].col_id), num_scs_bit))
+               static_cast<uint32_t>(coo_tensors[coo_tensor_index].col_id),
+               num_scs_bit))
            << 32) +
-          index);
-      ++index;
+          coo_tensor_index);
     }
     hwy::VQSort(keys.data(), keys.size(), hwy::SortAscending());
 
     uint32_t prev_col_id = std::numeric_limits<uint32_t>::max();
     uint32_t prev_row_id = std::numeric_limits<uint32_t>::max();
-    for (const auto key : keys) {
+    for (const uint64_t key : keys) {
       const uint32_t index = static_cast<uint32_t>(key);
       const CooFormat& coo_tensor = coo_tensors[index];
       const uint32_t global_sc_id =
@@ -133,6 +186,8 @@ void SortAndGroupCooTensorsPerLocalDevice(
     for (int global_sc_id = 0; global_sc_id < global_sc_count; ++global_sc_id) {
       max_ids_per_sc[global_sc_id] = std::max(
           max_ids_per_sc[global_sc_id], ids_per_sc_partition[global_sc_id]);
+      required_buffer_size_per_sc[local_sc_id] +=
+          jax_sc_embedding::RoundUpTo(ids_per_sc_partition[global_sc_id], 8);
       max_unique_ids_per_sc[global_sc_id] =
           std::max(max_unique_ids_per_sc[global_sc_id],
                    unique_ids_per_sc_partition[global_sc_id]);
@@ -158,44 +213,14 @@ void SortAndGroupCooTensorsPerLocalDevice(
         *absl::c_max_element(ids_per_sc_partition);
     const int32_t observed_max_unique_ids_per_partition =
         *absl::c_max_element(unique_ids_per_sc_partition);
-    // If id dropping is allowed, we log a warning if the observed max ids per
-    // partition is greater than the set max ids per partition.
-    if (observed_max_ids_per_partition > max_ids_per_partition) {
-      if (allow_id_dropping) {
-        LOG(WARNING) << "Allowing ID dropping for table: " << stacked_table_name
-                     << " observed max ids per partition: "
-                     << observed_max_ids_per_partition
-                     << " is greater than the set max ids per partition: "
-                     << max_ids_per_partition;
-      } else {
-        LOG(FATAL) << "Observed max ids per partition: "
-                   << observed_max_ids_per_partition
-                   << " for table: " << stacked_table_name
-                   << " is greater than the set max ids per partition: "
-                   << max_ids_per_partition;
-      }
-    }
-    if (observed_max_unique_ids_per_partition > max_unique_ids_per_partition) {
-      if (allow_id_dropping) {
-        LOG(WARNING)
-            << "Allowing ID dropping for table: " << stacked_table_name
-            << " observed max unique ids per partition: "
-            << observed_max_unique_ids_per_partition
-            << " is greater than the set max unique ids per partition: "
-            << max_unique_ids_per_partition;
-      } else {
-        LOG(FATAL) << "Observed max unique ids per partition: "
-                   << observed_max_unique_ids_per_partition
-                   << " for table: " << stacked_table_name
-                   << " is greater than the set max unique ids per partition: "
-                   << max_unique_ids_per_partition;
-      }
-    }
+
+    ValidateMaxIdsOrDie(observed_max_ids_per_partition,
+                        observed_max_unique_ids_per_partition,
+                        max_ids_per_partition, max_unique_ids_per_partition,
+                        stacked_table_name, allow_id_dropping);
   }
+  return coo_tensors_by_id;
 }
-
-
-
 int ComputeCooBufferSize(
     const int num_scs, const int num_scs_per_device,
     absl::Span<const StackedTableMetadata> stacked_table_metadata,
@@ -251,7 +276,7 @@ void FillRowPointersPerLocalDevice(
     absl::Span<const std::vector<CooFormat>> coo_tensors_by_id,
     const int row_pointers_size_per_sc, const int coo_buffer_size_per_sc,
     const int batch_size_per_sc, const int num_scs, const int num_sc_per_device,
-    int* row_pointers, int* embedding_ids, int* sample_ids, float* gains) {
+    int row_pointers[], int embedding_ids[], int sample_ids[], float gains[]) {
   tsl::profiler::TraceMe t("FillRowPointers");
   for (int local_sc_id = 0; local_sc_id < num_sc_per_device; ++local_sc_id) {
     int lhs_row_index = 0;
diff --git a/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.h b/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.h
@@ -51,7 +51,7 @@ struct CooFormat {
   // table.
   //
   // This packing allows for efficient storage and extractions using bitwise
-  // masks (assuming `num_scs` is a power of 2).
+  // masks (assuming number of sparsecores (SC) is a power of 2).
   int col_id;
   float gain;
 
@@ -66,6 +66,8 @@ int GetColId(int col_id, int col_shift, int col_offset, int num_scs_mod,
              int num_scs_mod_inv);
 
 // Rounds up the given value to the next multiple of the given alignment.
+// This is equivalent to ceil(value / align) * align, but implemented in an
+// integer-safe way.
 template <typename T>
 static inline T RoundUpTo(T value, T align) {
   return (value + align - 1) / align * align;
@@ -116,14 +118,14 @@ struct StackedTableMetadata {
   int max_col_id;
 };
 
-void SortAndGroupCooTensorsPerLocalDevice(
+std::vector<std::vector<CooFormat>> SortAndGroupCooTensorsPerLocalDevice(
     absl::Span<const CooFormat> coo_tensors, int batch_size_per_sc,
     int global_sc_count,
     int32_t batch_size_for_device,  // Batch size for the local device.
     int32_t max_ids_per_partition, int32_t max_unique_ids_per_partition,
     absl::string_view stacked_table_name, bool allow_id_dropping,
-    std::vector<std::vector<CooFormat>>& coo_tensors_by_id, int* max_ids_per_sc,
-    int* max_unique_ids_per_sc);
+    int num_sc_per_device, int total_num_coo_tensors, int max_ids_per_sc[],
+    int max_unique_ids_per_sc[], int required_buffer_size_per_sc[]);
 
 int ComputeCooBufferSize(
     int num_scs, int num_scs_per_device,
@@ -140,7 +142,7 @@ void FillRowPointersPerLocalDevice(
     absl::Span<const std::vector<CooFormat>> coo_tensors_by_id,
     int row_pointers_size_per_sc, int coo_buffer_size_per_sc,
     int batch_size_per_sc, int num_scs, int num_sc_per_device,
-    int* row_pointers, int* embedding_ids, int* sample_ids, float* gains);
+    int row_pointers[], int embedding_ids[], int sample_ids[], float gains[]);
 
 }  // namespace jax_sc_embedding
 
diff --git a/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util_test.cc b/jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util_test.cc
