[JAX SC] Extract features slice per SC in parallel.

PiperOrigin-RevId: 839887356

Author: adityagupta1089

jax_tpu_embedding/sparsecore/lib/core/BUILD

@@ -65,6 +65,7 @@ cc_library(
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@eigen_archive//:eigen3",
+        "@tsl//tsl/concurrency:async_value",
         "@tsl//tsl/profiler/lib:traceme",
         "@xla//xla:util",
     ],
@@ -153,6 +154,7 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@eigen_archive//:eigen3",
+        "@tsl//tsl/concurrency:async_value",
         "@tsl//tsl/platform:statusor",
         "@tsl//tsl/profiler/lib:traceme",
         "@xla//xla:util",
@@ -434,6 +436,7 @@ cc_test(
         "@com_google_absl//absl/types:span",
         "@com_google_benchmark//:benchmark_main",
         "@eigen_archive//:eigen3",
+        "@tsl//tsl/concurrency:async_value",
     ],
 )
 

jax_tpu_embedding/sparsecore/lib/core/abstract_input_batch.h

@@ -16,6 +16,7 @@
 #include <sys/types.h>
 
 #include <cstdint>
+#include <optional>
 
 #include "absl/base/attributes.h"  // from @com_google_absl
 #include "jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.h"
@@ -52,6 +53,13 @@ class AbstractInputBatch {
   // Returns the total number of embedding IDs across all samples.
   virtual int64_t id_count() const = 0;
 
+  // Returns number of ids in rows [start_row, end_row).
+  // If not implemented by a subclass, returns std::nullopt.
+  virtual std::optional<int64_t> GetIdsCountInSlice(int start_row,
+                                                    int end_row) const {
+    return std::nullopt;
+  }
+
   // Returns true if the input batch has variable weights.
   virtual bool HasVariableWeights() const { return true; }
 

jax_tpu_embedding/sparsecore/lib/core/coo_format.h

@@ -64,6 +64,13 @@ struct CooFormat {
   static constexpr uint32_t kDataMask = (1 << kDataBits) - 1;
   // Bit offset for rotated_col_id in grouping key.
   static constexpr uint32_t kRotatedColIdOffset = kDataBits;
+
+  // For hierarchical indexing within data bits when using multiple feature
+  // slices.
+  static constexpr uint32_t kSliceIndexBits = 6;
+  static constexpr uint32_t kItemIndexBits = kDataBits - kSliceIndexBits;
+  static constexpr uint32_t kItemIndexMask = (1UL << kItemIndexBits) - 1;
+
   // Bit offset for bucket_id in grouping key.
   static constexpr uint32_t kBucketIdOffset = kRotatedColIdOffset + 32;
 
@@ -171,6 +178,21 @@ struct CooFormat {
   static uint32_t GetBucketIdFromKey(uint64_t key) {
     return key >> kBucketIdOffset;
   }
+
+  static uint32_t EncodeHierarchicalIndex(uint32_t slice_idx,
+                                          uint32_t item_idx) {
+    DCHECK_LT(slice_idx, 1 << kSliceIndexBits);
+    DCHECK_LT(item_idx, 1 << kItemIndexBits);
+    return (slice_idx << kItemIndexBits) | item_idx;
+  }
+
+  static uint32_t GetSliceIndexFromData(uint32_t data) {
+    return data >> kItemIndexBits;
+  }
+
+  static uint32_t GetItemIndexFromData(uint32_t data) {
+    return data & kItemIndexMask;
+  }
 };
 
 }  // namespace jax_sc_embedding

jax_tpu_embedding/sparsecore/lib/core/extract_sort_and_group_benchmark.cc

@@ -14,6 +14,7 @@
 #include <cmath>
 #include <cstdint>
 #include <cstdio>
+#include <functional>
 #include <limits>
 #include <memory>
 #include <optional>
@@ -35,6 +36,8 @@
 #include "jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.h"
 #include "jax_tpu_embedding/sparsecore/lib/core/ragged_tensor_input_batch.h"
 #include "jax_tpu_embedding/sparsecore/lib/core/sort_and_group_coo_tensors_impl.h"
+#include "tsl/concurrency/async_value.h"  // from @tsl
+#include "tsl/concurrency/async_value_ref.h"  // from @tsl
 
 namespace jax_sc_embedding {
 
@@ -168,6 +171,8 @@ void BM_ExtractCooTensors(benchmark::State& state) {
   std::vector<StackedTableMetadata> stacked_table_metadata;
   stacked_table_metadata.reserve(num_features);
   for (int i = 0; i < num_features; ++i) {
+    // Set to INT_MAX to avoid ID dropping and observe the actual statistics of
+    // the generated data. This doesn't affect performance of grouping itself.
     stacked_table_metadata.push_back(StackedTableMetadata(
         absl::StrCat("table_", i), /*feature_index=*/i,
         /*max_ids_per_partition=*/std::numeric_limits<int>::max(),
@@ -188,9 +193,13 @@ void BM_ExtractCooTensors(benchmark::State& state) {
   };
 
   for (auto s : state) {
-    internal::ExtractCooTensorsForAllFeaturesPerLocalDevice(
-        stacked_table_metadata, absl::MakeSpan(input_batches),
-        /*local_device_id=*/0, options);
+    std::vector<tsl::AsyncValueRef<ExtractedCooTensors>> results_av =
+        internal::ExtractCooTensorsForAllFeaturesPerLocalDeviceAsync(
+            stacked_table_metadata, absl::MakeSpan(input_batches),
+            /*local_device_id=*/0, options);
+    for (auto& av : results_av) {
+      tsl::BlockUntilReady(av.GetAsyncValue());
+    }
   }
 }
 BENCHMARK(BM_ExtractCooTensors)
@@ -233,10 +242,20 @@ void BM_SortAndGroup_Phase1(benchmark::State& state) {
       .enable_minibatching = true,
   };
 
-  ExtractedCooTensors extracted_coo_tensors =
-      internal::ExtractCooTensorsForAllFeaturesPerLocalDevice(
+  // Extract COO tensors for all features on a single local device.
+  std::vector<tsl::AsyncValueRef<ExtractedCooTensors>> feature_results_av =
+      internal::ExtractCooTensorsForAllFeaturesPerLocalDeviceAsync(
           stacked_table_metadata_list, absl::MakeSpan(input_batches),
           /*local_device_id=*/0, options);
+  for (auto& av : feature_results_av) {
+    tsl::BlockUntilReady(av.GetAsyncValue());
+  }
+  std::vector<ExtractedCooTensors> feature_results;
+  std::vector<std::reference_wrapper<const ExtractedCooTensors>>
+      extracted_coo_tensors_list;
+  internal::GetExtractedCooTensorsFromAsyncValues(
+      absl::MakeSpan(feature_results_av), feature_results,
+      extracted_coo_tensors_list);
 
   bool minibatching_required = false;
   StatsPerHost stats_per_host(
@@ -248,8 +267,9 @@ void BM_SortAndGroup_Phase1(benchmark::State& state) {
 
   if (state.thread_index() == 0) {
     SortAndGroupCooTensorsPerLocalDevice</*kHasVariableWeights=*/false>(
-        extracted_coo_tensors, stacked_table_metadata_list[0], options,
-        stats_per_device, minibatching_required);
+        absl::MakeSpan(extracted_coo_tensors_list),
+        stacked_table_metadata_list[0], options, stats_per_device,
+        minibatching_required);
     LogStats(stats_per_device.max_ids_per_partition,
              "Max ids per partition across all global SCs");
     LogStats(stats_per_device.max_unique_ids_per_partition,
@@ -258,8 +278,9 @@ void BM_SortAndGroup_Phase1(benchmark::State& state) {
 
   for (auto s : state) {
     SortAndGroupCooTensorsPerLocalDevice</*kHasVariableWeights=*/false>(
-        extracted_coo_tensors, stacked_table_metadata_list[0], options,
-        stats_per_device, minibatching_required);
+        absl::MakeSpan(extracted_coo_tensors_list),
+        stacked_table_metadata_list[0], options, stats_per_device,
+        minibatching_required);
   }
 }
 BENCHMARK(BM_SortAndGroup_Phase1)