Migrate id transformer binding (#835)

Summary:
This PR integrated the migrated id transformer and lxu strategy to create a python binding. There are several point that may need to be mentioned:
- We use `IDTransformerVariant` and `LXUStrategyVariant` type to support adding more kinds of id transformers and lxu strategies.
- The origin design used [nlohmann::json](https://github.com/nlohmann/json) to parse the configs. To remove this dependency, we list all configs of all kinds of id transformers and lxu strategies as the parameter of the constructor of the variant.
- We choose to use pytorch native binding (`CustomClassHolder`) instead of pybind to be coherent with fbgemm_gpu.

Thank you for your time on reviewing this PR :)

gently ping divchenko colin2328 reyoung

Pull Request resolved: #835

Reviewed By: s4ayub

Differential Revision: D41507250

Pulled By: colin2328

fbshipit-source-id: f6e8bc6234f287dfba333e1fe7bbe7de64df88b6

Author: zhuzilin

benchmarks/cpp/dynamic_embedding/mixed_lfu_lru_strategy_benchmark.cpp

@@ -12,7 +12,7 @@
 namespace torchrec {
 void BM_MixedLFULRUStrategy(benchmark::State& state) {
   size_t num_ext_values = state.range(0);
-  std::vector<MixedLFULRUStrategy::lxu_record_t> ext_values(num_ext_values);
+  std::vector<lxu_record_t> ext_values(num_ext_values);
 
   MixedLFULRUStrategy strategy;
   for (auto& v : ext_values) {

benchmarks/cpp/dynamic_embedding/mixed_lfu_lru_strategy_evict_benchmark.cpp

@@ -18,13 +18,13 @@ class RecordIterator {
  public:
   RecordIterator(Container::const_iterator begin, Container::const_iterator end)
       : begin_(begin), end_(end) {}
-  std::optional<TransformerRecord<uint32_t>> operator()() {
+  std::optional<record_t> operator()() {
     if (begin_ == end_) {
       return std::nullopt;
     }
-    TransformerRecord<uint32_t> record{};
-    record.global_id_ = next_global_id_++;
-    record.lxu_record_ = *reinterpret_cast<const uint32_t*>(&(*begin_++));
+    record_t record{};
+    record.global_id = next_global_id_++;
+    record.lxu_record = *reinterpret_cast<const uint32_t*>(&(*begin_++));
     return record;
   }
 
@@ -52,8 +52,8 @@ class RandomizeMixedLXUSet {
     std::uniform_int_distribution<uint32_t> time_dist(0, max_time - 1);
     for (size_t i = 0; i < n; ++i) {
       MixedLFULRUStrategy::Record record{};
-      record.freq_power_ = freq_dist(engine) + min_freq;
-      record.time_ = time_dist(engine);
+      record.freq_power = freq_dist(engine) + min_freq;
+      record.time = time_dist(engine);
       records_.emplace_back(record);
     }
   }
@@ -68,8 +68,9 @@ class RandomizeMixedLXUSet {
 
 void BM_MixedLFULRUStrategyEvict(benchmark::State& state) {
   RandomizeMixedLXUSet lxuSet(state.range(0), state.range(1), state.range(2));
+  MixedLFULRUStrategy strategy;
   for (auto _ : state) {
-    MixedLFULRUStrategy::evict(lxuSet.Iterator(), state.range(3));
+    strategy.evict(lxuSet.Iterator(), state.range(3));
   }
 }
 

benchmarks/cpp/dynamic_embedding/naive_id_transformer_benchmark.cpp

@@ -13,8 +13,7 @@
 namespace torchrec {
 
 static void BM_NaiveIDTransformer(benchmark::State& state) {
-  using Tag = int32_t;
-  NaiveIDTransformer<Tag> transformer(2e8);
+  NaiveIDTransformer transformer(2e8);
   torch::Tensor global_ids = torch::empty({1024, 1024}, torch::kLong);
   torch::Tensor cache_ids = torch::empty_like(global_ids);
   for (auto _ : state) {

test/cpp/dynamic_embedding/mixed_lfu_lru_strategy_test.cpp

@@ -12,8 +12,8 @@
 namespace torchrec {
 TEST(TDE, order) {
   MixedLFULRUStrategy::Record a;
-  a.time_ = 1;
-  a.freq_power_ = 31;
+  a.time = 1;
+  a.freq_power = 31;
   uint32_t i32 = a.ToUint32();
   ASSERT_EQ(0xF8000001, i32);
 }
@@ -23,42 +23,41 @@ TEST(TDE, MixedLFULRUStrategy_Evict) {
   {
     records.emplace_back();
     records.back().first = 1;
-    records.back().second.time_ = 100;
-    records.back().second.freq_power_ = 2;
+    records.back().second.time = 100;
+    records.back().second.freq_power = 2;
   }
   {
     records.emplace_back();
     records.back().first = 2;
-    records.back().second.time_ = 10;
-    records.back().second.freq_power_ = 2;
+    records.back().second.time = 10;
+    records.back().second.freq_power = 2;
   }
   {
     records.emplace_back();
     records.back().first = 3;
-    records.back().second.time_ = 100;
-    records.back().second.freq_power_ = 1;
+    records.back().second.time = 100;
+    records.back().second.freq_power = 1;
   }
   {
     records.emplace_back();
     records.back().first = 4;
-    records.back().second.time_ = 150;
-    records.back().second.freq_power_ = 2;
+    records.back().second.time = 150;
+    records.back().second.freq_power = 2;
   }
   size_t offset_{0};
-  auto ids = MixedLFULRUStrategy::evict(
-      [&offset_,
-       &records]() -> std::optional<MixedLFULRUStrategy::transformer_record_t> {
+  MixedLFULRUStrategy strategy;
+  auto ids = strategy.evict(
+      [&offset_, &records]() -> std::optional<record_t> {
         if (offset_ == records.size()) {
           return std::nullopt;
         }
         auto record = records[offset_++];
-        MixedLFULRUStrategy::lxu_record_t ext_type =
-            *reinterpret_cast<MixedLFULRUStrategy::lxu_record_t*>(
-                &record.second);
-        return MixedLFULRUStrategy::transformer_record_t{
-            .global_id_ = record.first,
-            .cache_id_ = 0,
-            .lxu_record_ = ext_type,
+        lxu_record_t ext_type =
+            *reinterpret_cast<lxu_record_t*>(&record.second);
+        return record_t{
+            .global_id = record.first,
+            .cache_id = 0,
+            .lxu_record = ext_type,
         };
       },
       3);
@@ -73,12 +72,12 @@ TEST(TDE, MixedLFULRUStrategy_Transform) {
   constexpr static size_t n_iter = 1000000;
   MixedLFULRUStrategy strategy;
   strategy.update_time(10);
-  MixedLFULRUStrategy::lxu_record_t val;
+  lxu_record_t val;
   {
     val = strategy.update(0, 0, std::nullopt);
     auto record = reinterpret_cast<MixedLFULRUStrategy::Record*>(&val);
-    ASSERT_EQ(record->freq_power_, 5);
-    ASSERT_EQ(record->time_, 10);
+    ASSERT_EQ(record->freq_power, 5);
+    ASSERT_EQ(record->time, 10);
   }
 
   uint32_t freq_power_5_cnt = 0;
@@ -87,13 +86,13 @@ TEST(TDE, MixedLFULRUStrategy_Transform) {
   for (size_t i = 0; i < n_iter; ++i) {
     auto tmp = strategy.update(0, 0, val);
     auto record = reinterpret_cast<MixedLFULRUStrategy::Record*>(&tmp);
-    ASSERT_EQ(record->time_, 10);
-    if (record->freq_power_ == 5) {
+    ASSERT_EQ(record->time, 10);
+    if (record->freq_power == 5) {
       ++freq_power_5_cnt;
-    } else if (record->freq_power_ == 6) {
+    } else if (record->freq_power == 6) {
       ++freq_power_6_cnt;
     } else {
-      ASSERT_TRUE(record->freq_power_ == 5 || record->freq_power_ == 6);
+      ASSERT_TRUE(record->freq_power == 5 || record->freq_power == 6);
     }
   }
 

test/cpp/dynamic_embedding/naive_id_transformer_test.cpp

@@ -12,8 +12,7 @@
 namespace torchrec {
 
 TEST(tde, NaiveThreadedIDTransformer_NoFilter) {
-  using Tag = int32_t;
-  NaiveIDTransformer<Tag, Bitmap<uint8_t>> transformer(16);
+  NaiveIDTransformer<Bitmap<uint8_t>> transformer(16);
   const int64_t global_ids[5] = {100, 101, 100, 102, 101};
   int64_t cache_ids[5];
   int64_t expected_cache_ids[5] = {0, 1, 0, 2, 1};
@@ -24,8 +23,7 @@ TEST(tde, NaiveThreadedIDTransformer_NoFilter) {
 }
 
 TEST(tde, NaiveThreadedIDTransformer_Full) {
-  using Tag = int32_t;
-  NaiveIDTransformer<Tag, Bitmap<uint8_t>> transformer(4);
+  NaiveIDTransformer<Bitmap<uint8_t>> transformer(4);
   const int64_t global_ids[5] = {100, 101, 102, 103, 104};
   int64_t cache_ids[5];
   int64_t expected_cache_ids[5] = {0, 1, 2, 3, -1};
@@ -37,8 +35,7 @@ TEST(tde, NaiveThreadedIDTransformer_Full) {
 }
 
 TEST(tde, NaiveThreadedIDTransformer_Evict) {
-  using Tag = int32_t;
-  NaiveIDTransformer<Tag, Bitmap<uint8_t>> transformer(4);
+  NaiveIDTransformer<Bitmap<uint8_t>> transformer(4);
   const int64_t global_ids[5] = {100, 101, 102, 103, 104};
   int64_t cache_ids[5];
 
@@ -60,8 +57,7 @@ TEST(tde, NaiveThreadedIDTransformer_Evict) {
 }
 
 TEST(tde, NaiveThreadedIDTransformer_Iterator) {
-  using Tag = int32_t;
-  NaiveIDTransformer<Tag, Bitmap<uint8_t>> transformer(16);
+  NaiveIDTransformer<Bitmap<uint8_t>> transformer(16);
   const int64_t global_ids[5] = {100, 101, 100, 102, 101};
   int64_t cache_ids[5];
   int64_t expected_cache_ids[5] = {3, 4, 3, 5, 4};

torchrec/csrc/dynamic_embedding/CMakeLists.txt

@@ -6,10 +6,11 @@
 
 add_library(tde_cpp_objs
             OBJECT
+            bind.cpp
+            id_transformer_wrapper.cpp
             details/clz_impl.cpp
             details/ctz_impl.cpp
             details/random_bits_generator.cpp
-            details/mixed_lfu_lru_strategy.cpp
             details/io_registry.cpp
             details/io.cpp
             details/notification.cpp)

torchrec/csrc/dynamic_embedding/bind.cpp

@@ -0,0 +1,30 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <torch/torch.h>
+
+#include <torchrec/csrc/dynamic_embedding/details/io_registry.h>
+#include <torchrec/csrc/dynamic_embedding/id_transformer_wrapper.h>
+
+namespace torchrec {
+TORCH_LIBRARY(tde, m) {
+  m.def("register_io", [](const std::string& name) {
+    IORegistry::Instance().register_plugin(name.c_str());
+  });
+
+  m.class_<TransformResult>("TransformResult")
+      .def_readonly("success", &TransformResult::success)
+      .def_readonly("ids_to_fetch", &TransformResult::ids_to_fetch);
+
+  m.class_<IDTransformerWrapper>("IDTransformer")
+      .def(torch::init<int64_t, std::string, std::string, int64_t>())
+      .def("transform", &IDTransformerWrapper::transform)
+      .def("evict", &IDTransformerWrapper::evict)
+      .def("save", &IDTransformerWrapper::save);
+}
+} // namespace torchrec

torchrec/csrc/dynamic_embedding/details/id_transformer.h

@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+#include <torchrec/csrc/dynamic_embedding/details/types.h>
+#include <functional>
+#include <span>
+#include <type_traits>
+
+namespace torchrec {
+
+namespace transform_default {
+inline lxu_record_t no_update(
+    int64_t global_id,
+    int64_t cache_id,
+    std::optional<lxu_record_t> record) {
+  return record.value_or(lxu_record_t{});
+};
+
+inline void no_fetch(int64_t global_id, int64_t cache_id) {}
+} // namespace transform_default
+
+class IDTransformer {
+ public:
+  /**
+   * Transform global ids to cache ids
+   *
+   * @tparam Update Update the eviction strategy tag type. Update LXU Record
+   * @tparam Fetch Fetch the not existing global-id/cache-id pair. It is used
+   * by dynamic embedding parameter server.
+   *
+   * @param global_ids Global ID vector
+   * @param cache_ids [out] Cache ID vector
+   * @param update update lambda. See `Update` doc.
+   * @param fetch fetch lambda. See `Fetch` doc.
+   * @return true if all transformed, otherwise need eviction.
+   */
+  virtual bool transform(
+      std::span<const int64_t> global_ids,
+      std::span<int64_t> cache_ids,
+      update_t update = transform_default::no_update,
+      fetch_t fetch = transform_default::no_fetch) = 0;
+
+  /**
+   * Evict global ids from the transformer
+   *
+   * @param global_ids Global IDs to evict.
+   */
+  virtual void evict(std::span<const int64_t> global_ids) = 0;
+
+  /**
+   * Create an iterator of the id transformer, a possible usecase is:
+   *
+   *   auto iterator = transformer.iterator();
+   *   auto record = iterator();
+   *   while (record.has_value()) {
+   *     // do sth with the record
+   *     // ...
+   *     // get next record
+   *     auto record = iterator();
+   *   }
+   *
+   * @return the iterator created.
+   */
+  virtual iterator_t iterator() const = 0;
+};
+
+} // namespace torchrec

torchrec/csrc/dynamic_embedding/details/lxu_strategy.h

@@ -0,0 +1,41 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+#include <torchrec/csrc/dynamic_embedding/details/types.h>
+#include <optional>
+
+namespace torchrec {
+
+class LXUStrategy {
+ public:
+  LXUStrategy() = default;
+  LXUStrategy(const LXUStrategy&) = delete;
+  LXUStrategy(LXUStrategy&& o) noexcept = default;
+
+  virtual void update_time(uint32_t time) = 0;
+  virtual int64_t time(lxu_record_t record) = 0;
+
+  virtual lxu_record_t update(
+      int64_t global_id,
+      int64_t cache_id,
+      std::optional<lxu_record_t> val) = 0;
+
+  /**
+   * Analysis all ids and returns the num_elems that are most need to evict.
+   * @param iterator Returns each global_id to ExtValue pair. Returns nullopt
+   * when at ends.
+   * @param num_to_evict
+   * @return
+   */
+  virtual std::vector<int64_t> evict(
+      iterator_t iterator,
+      uint64_t num_to_evict) = 0;
+};
+
+} // namespace torchrec