Introduce FdoStats struct for tracking FDO info.

PiperOrigin-RevId: 757849206

Author: manuadg-verily

jax_tpu_embedding/sparsecore/lib/core/BUILD

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 load("//jax_tpu_embedding/sparsecore:jax_tpu_embedding.bzl", "CORE_USERS")
-load("//third_party/bazel/python:pybind11.bzl", "pybind_extension")
+load("//third_party/bazel/python:pybind11.bzl", "pybind_extension", "pybind_library")
 load("//third_party/bazel/python:pypi.bzl", "pypi_requirement")
 load("//third_party/bazel/python:pytype.bzl", "pytype_strict_contrib_test", "pytype_strict_library")
 
@@ -78,6 +78,7 @@ pybind_extension(
     name = "input_preprocessing_cc",
     srcs = ["input_preprocessing.cc"],
     deps = [
+        ":fdo_types",
         ":input_preprocessing_threads",
         ":input_preprocessing_util",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -90,6 +91,21 @@ pybind_extension(
     ],
 )
 
+pybind_library(
+    name = "fdo_types",
+    hdrs = ["fdo_types.h"],
+    deps = ["@com_google_absl//absl/container:flat_hash_map"],
+)
+
+pybind_extension(
+    name = "fdo_types_cc",
+    srcs = [
+        "fdo_types.cc",
+        "fdo_types.h",
+    ],
+    deps = ["@com_google_absl//absl/container:flat_hash_map"],
+)
+
 pytype_strict_library(
     name = "input_preprocessing",
     srcs = [
@@ -161,12 +177,17 @@ pytype_strict_library(
     srcs = ["__init__.py"],
     # C++ dependencies must go in "data".
     data = [
+        ":fdo_types",  # buildcleaner: keep
         ":input_preprocessing_threads",  # buildcleaner: keep
         ":input_preprocessing_util",  # buildcleaner: keep
     ],
-    visibility = ["//jax_tpu_embedding/sparsecore/lib:__pkg__"],
+    visibility = [
+        "//jax_tpu_embedding/sparsecore/lib:__pkg__",
+        "//jax_tpu_embedding/sparsecore/lib/extensions:__pkg__",
+    ],
     deps = [
         ":constants",  # buildcleaner: keep
+        ":fdo_types_cc",  # buildcleaner: keep
         ":input_preprocessing",  # buildcleaner: keep
         ":input_preprocessing_cc",  # buildcleaner: keep
         "//jax_tpu_embedding/sparsecore/lib/core/primitives",  # buildcleaner: keep

jax_tpu_embedding/sparsecore/lib/core/fdo_types.cc

@@ -0,0 +1,33 @@
+// Copyright 2024 The JAX SC Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#include "jax_tpu_embedding/sparsecore/lib/core/fdo_types.h"
+
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/numpy.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+
+namespace jax_sc_embedding {
+
+namespace py = ::pybind11;
+
+PYBIND11_MODULE(fdo_types_cc, m) {
+  py::class_<PyFdoStats>(m, "FdoStats")
+      .def_readwrite("max_ids_per_partition",
+                     &PyFdoStats::max_ids_per_partition)
+      .def_readwrite("max_unique_ids_per_partition",
+                     &PyFdoStats::max_unique_ids_per_partition);
+}
+
+}  // namespace jax_sc_embedding

jax_tpu_embedding/sparsecore/lib/core/fdo_types.h

@@ -0,0 +1,71 @@
+// Copyright 2024 The JAX SC Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#ifndef JAX_TPU_EMBEDDING_SPARSECORE_LIB_CORE_FDO_TYPES_H_
+#define JAX_TPU_EMBEDDING_SPARSECORE_LIB_CORE_FDO_TYPES_H_
+
+#include <string>
+
+#include "absl/container/flat_hash_map.h"  // from @com_google_absl
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/numpy.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+
+namespace jax_sc_embedding {
+
+namespace py = ::pybind11;
+
+template <typename ArrayType, typename MapType>
+struct is_valid_map_type {
+  // is the value type of map same as T?
+  static constexpr bool value = false;
+};
+
+template <typename ArrayType, typename ValueType>
+struct is_valid_map_type<ArrayType,
+                         absl::flat_hash_map<std::string, ValueType>> {
+  static constexpr bool value = true;
+};
+
+template <typename ArrayType>
+struct is_valid_map_type<ArrayType, py::dict> {
+  // dict does not support type annotations
+  static constexpr bool value = true;
+};
+
+template <typename ArrayType,  // std::vector<int> or py::array_t<int>
+          typename MapType     // absl::flat_hash_map<std::string,
+                               // std::vector<int>> or py::dict
+          >
+struct FdoStats {
+  static_assert(is_valid_map_type<ArrayType, MapType>::value,
+                "MapType must be a valid map type");
+
+  using FdoStatsPerSparseCore = ArrayType;
+
+  using StackedTableName = std::string;
+
+  // <StackedTableName, FdoStatsPerSparseCore>
+  using FdoStatsPerStackedTable = MapType;
+
+  FdoStatsPerStackedTable max_ids_per_partition;
+  FdoStatsPerStackedTable max_unique_ids_per_partition;
+  FdoStatsPerStackedTable id_drop_counters;
+};
+
+using PyFdoStats = FdoStats<py::array_t<int>, py::dict>;
+
+}  // namespace jax_sc_embedding
+
+#endif  // JAX_TPU_EMBEDDING_SPARSECORE_LIB_CORE_FDO_TYPES_H_

jax_tpu_embedding/sparsecore/lib/core/input_preprocessing.cc

@@ -13,6 +13,7 @@
 // limitations under the License.
 #include <algorithm>
 #include <cmath>
+#include <cstdint>
 #include <optional>
 #include <string>
 #include <utility>
@@ -24,6 +25,7 @@
 #include "absl/strings/string_view.h"  // from @com_google_absl
 #include "absl/synchronization/blocking_counter.h"  // from @com_google_absl
 #include "absl/types/span.h"  // from @com_google_absl
+#include "jax_tpu_embedding/sparsecore/lib/core/fdo_types.h"
 #include "jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_threads.h"
 #include "jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_util.h"
 #include "pybind11/cast.h"  // from @pybind11
@@ -382,8 +384,8 @@ py::tuple PreprocessSparseDenseMatmulInput(
   py::dict lhs_embedding_ids;
   py::dict lhs_sample_ids;
   py::dict lhs_gains;
-  py::dict max_ids_per_partition;
-  py::dict max_unique_ids_per_partition;
+  PyFdoStats::FdoStatsPerStackedTable max_ids_per_partition;
+  PyFdoStats::FdoStatsPerStackedTable max_unique_ids_per_partition;
   const int num_scs = num_sc_per_device * global_device_count;
   const int row_pointers_size_per_sc = std::max(num_scs, 8);
 
@@ -495,9 +497,10 @@ py::tuple PreprocessSparseDenseMatmulInput(
     }
     counter.Wait();
   }
-  py::dict stats;
-  stats["max_ids"] = max_ids_per_partition;
-  stats["max_unique_ids"] = max_unique_ids_per_partition;
+  PyFdoStats stats{
+      .max_ids_per_partition = max_ids_per_partition,
+      .max_unique_ids_per_partition = max_unique_ids_per_partition,
+  };
   // GIL is held at this point.
   return py::make_tuple(lhs_row_pointers, lhs_embedding_ids, lhs_sample_ids,
                         lhs_gains, stats);

jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_cc_test.py

@@ -763,7 +763,7 @@ def test_multi_process_fdo(self, has_leading_dimension):
             allow_id_dropping=False,
         )
     )
-    stats = embedding.SparseDenseMatmulInputStats.from_dict(stats)
+    stats = embedding.SparseDenseMatmulInputStats.from_cc(stats)
     fdo_client.record(stats)
     fdo_client.publish()
     # Duplicated ids on row 0 and 6 are combined.

jax_tpu_embedding/sparsecore/lib/nn/BUILD

@@ -40,6 +40,7 @@ pytype_strict_library(
     deps = [
         ":embedding_spec",
         ":table_stacking",
+        "//jax_tpu_embedding/sparsecore/lib/core:fdo_types_cc",
         "//jax_tpu_embedding/sparsecore/lib/core:input_preprocessing_cc",
         "//jax_tpu_embedding/sparsecore/lib/core/primitives:sparse_dense_matmul_csr",
         "//jax_tpu_embedding/sparsecore/lib/proto:embedding_spec_py_pb2",

jax_tpu_embedding/sparsecore/lib/nn/embedding.py

@@ -24,6 +24,7 @@
 from jax.experimental.layout import DeviceLocalLayout as DLL
 from jax.experimental.layout import Layout
 import jax.numpy as jnp
+from jax_tpu_embedding.sparsecore.lib.core import fdo_types_cc
 from jax_tpu_embedding.sparsecore.lib.core import input_preprocessing_cc
 from jax_tpu_embedding.sparsecore.lib.core.primitives import sparse_dense_matmul_csr
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
@@ -69,12 +70,12 @@ class SparseDenseMatmulInputStats:
   max_unique_ids_per_partition: Mapping[str, np.ndarray]
 
   @classmethod
-  def from_dict(
-      cls, stats: Mapping[str, Mapping[str, np.ndarray]]
+  def from_cc(
+      cls, stats: fdo_types_cc.FdoStats
   ) -> "SparseDenseMatmulInputStats":
     return cls(
-        max_ids_per_partition=stats["max_ids"],
-        max_unique_ids_per_partition=stats["max_unique_ids"],
+        max_ids_per_partition=stats.max_ids_per_partition,
+        max_unique_ids_per_partition=stats.max_unique_ids_per_partition,
     )
 
 
@@ -380,7 +381,7 @@ def preprocess_sparse_dense_matmul_input(
 
   return SparseDenseMatmulInput(
       *preprocessed_inputs
-  ), SparseDenseMatmulInputStats.from_dict(stats)
+  ), SparseDenseMatmulInputStats.from_cc(stats)
 
 
 def _get_activation_for_feature(