Migrate naive id transformer (pytorch#661)

Summary:
This PR is trying to migrate `NaiveIDTransformer` to the top folder. The idea behind the naive transformer is quite simple, we use a hash map (`ska::flat_hash_map`) to store the global id to cache id map, and use a bitmap to record the empty slot in the cache.

cc reyoung

Pull Request resolved: pytorch#661

Reviewed By: zyan0

Differential Revision: D39857301

Pulled By: colin2328

fbshipit-source-id: a7681826827d84d51dcf5d91b3a0eca7583f391e

Author: zhuzilin

.github/workflows/unittest_ci_cpu.yml

@@ -121,7 +121,9 @@ jobs:
       working-directory: ./cpp-build
       shell: bash
       run: |
-        conda run -n build_binary cmake -DBUILD_TEST=ON -DCMAKE_CXX_COMPILER=$(which x86_64-conda-linux-gnu-g++) ..
+        conda run -n build_binary cmake \
+            -DBUILD_TEST=ON \
+            -DCMAKE_PREFIX_PATH=/home/ec2-user/miniconda/envs/build_binary/lib/python3.7/site-packages/torch/share/cmake ..
         conda run -n build_binary make -j
     - name: Test
       working-directory: ./cpp-build

CMakeLists.txt

@@ -9,12 +9,16 @@ cmake_minimum_required(VERSION 3.11.0 FATAL_ERROR)
 project(TorchRec
         LANGUAGES CXX C)
 
+find_package(Torch REQUIRED)
+
 set(CMAKE_CXX_STANDARD 20)
 
 include(FetchContent)
 
 option(BUILD_TEST "Build C++ test binaries (need gtest and gbenchmark)" OFF)
 
+add_definitions("-D_GLIBCXX_USE_CXX11_ABI=0")
+
 add_subdirectory(torchrec/csrc)
 
 if (BUILD_TEST)
@@ -34,5 +38,6 @@ if (BUILD_TEST)
     FetchContent_MakeAvailable(googletest google_benchmark)
 
     enable_testing()
+    add_subdirectory(benchmarks/cpp)
     add_subdirectory(test/cpp)
 endif()

benchmarks/cpp/CMakeLists.txt

@@ -0,0 +1,7 @@
+# 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.
+
+add_subdirectory(dynamic_embedding)

benchmarks/cpp/dynamic_embedding/CMakeLists.txt

@@ -0,0 +1,12 @@
+# 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.
+
+function(add_tde_benchmark NAME)
+    add_executable(${NAME} ${ARGN})
+    target_link_libraries(${NAME} tde_cpp_objs benchmark::benchmark_main benchmark::benchmark)
+endfunction()
+
+add_tde_benchmark(naive_id_transformer_benchmark naive_id_transformer_benchmark.cpp)

benchmarks/cpp/dynamic_embedding/naive_id_transformer_benchmark.cpp

@@ -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.
+ */
+
+#include <benchmark/benchmark.h>
+#include <torch/torch.h>
+#include <torchrec/csrc/dynamic_embedding/details/naive_id_transformer.h>
+
+namespace torchrec {
+
+static void BM_NaiveIDTransformer(benchmark::State& state) {
+  using Tag = int32_t;
+  NaiveIDTransformer<Tag> transformer(2e8);
+  torch::Tensor global_ids = torch::empty({1024, 1024}, torch::kLong);
+  torch::Tensor cache_ids = torch::empty_like(global_ids);
+  for (auto _ : state) {
+    state.PauseTiming();
+    global_ids.random_(state.range(0), state.range(1));
+    state.ResumeTiming();
+    transformer.transform(
+        std::span{
+            global_ids.template data_ptr<int64_t>(),
+            static_cast<size_t>(global_ids.numel())},
+        std::span{
+            cache_ids.template data_ptr<int64_t>(),
+            static_cast<size_t>(cache_ids.numel())});
+  }
+}
+
+BENCHMARK(BM_NaiveIDTransformer)
+    ->Iterations(100)
+    ->Unit(benchmark::kMillisecond)
+    ->ArgNames({"rand_from", "rand_to"})
+    ->Args({static_cast<long long>(1e10), static_cast<long long>(2e10)})
+    ->Args({static_cast<long long>(1e6), static_cast<long long>(2e6)});
+
+} // namespace torchrec

test/cpp/dynamic_embedding/CMakeLists.txt

@@ -11,3 +11,4 @@ function(add_tde_test NAME)
 endfunction()
 
 add_tde_test(bits_op_test bits_op_test.cpp)
+add_tde_test(naive_id_transformer_test naive_id_transformer_test.cpp)

test/cpp/dynamic_embedding/naive_id_transformer_test.cpp

@@ -0,0 +1,77 @@
+/*
+ * 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 <gtest/gtest.h>
+#include <torchrec/csrc/dynamic_embedding/details/naive_id_transformer.h>
+
+namespace torchrec {
+
+TEST(tde, NaiveThreadedIDTransformer_NoFilter) {
+  using Tag = int32_t;
+  NaiveIDTransformer<Tag, 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};
+  ASSERT_TRUE(transformer.transform(global_ids, cache_ids));
+  for (size_t i = 0; i < 5; i++) {
+    ASSERT_EQ(expected_cache_ids[i], cache_ids[i]);
+  }
+}
+
+TEST(tde, NaiveThreadedIDTransformer_Full) {
+  using Tag = int32_t;
+  NaiveIDTransformer<Tag, 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};
+
+  ASSERT_FALSE(transformer.transform(global_ids, cache_ids));
+  for (size_t i = 0; i < 4; i++) {
+    EXPECT_EQ(expected_cache_ids[i], cache_ids[i]);
+  }
+}
+
+TEST(tde, NaiveThreadedIDTransformer_Evict) {
+  using Tag = int32_t;
+  NaiveIDTransformer<Tag, Bitmap<uint8_t>> transformer(4);
+  const int64_t global_ids[5] = {100, 101, 102, 103, 104};
+  int64_t cache_ids[5];
+
+  ASSERT_FALSE(transformer.transform(global_ids, cache_ids));
+
+  const int64_t evict_global_ids[2] = {100, 102};
+  transformer.evict(evict_global_ids);
+
+  const int64_t new_global_ids[4] = {101, 102, 103, 104};
+  int64_t new_cache_ids[4];
+
+  ASSERT_TRUE(transformer.transform(new_global_ids, new_cache_ids));
+
+  int64_t expected_cache_ids[4] = {1, 0, 3, 2};
+
+  for (size_t i = 0; i < 4; i++) {
+    EXPECT_EQ(expected_cache_ids[i], new_cache_ids[i]);
+  }
+}
+
+TEST(tde, NaiveThreadedIDTransformer_Iterator) {
+  using Tag = int32_t;
+  NaiveIDTransformer<Tag, 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};
+  ASSERT_TRUE(transformer.transform(global_ids, cache_ids));
+
+  auto iterator = transformer.iterator();
+  for (size_t i = 0; i < 3; i++) {
+    EXPECT_TRUE(iterator().has_value());
+  }
+  EXPECT_TRUE(!iterator().has_value());
+}
+
+} // namespace torchrec

torchrec/csrc/dynamic_embedding/CMakeLists.txt

@@ -10,3 +10,6 @@ add_library(tde_cpp_objs
             details/ctz_impl.cpp)
 
 target_include_directories(tde_cpp_objs PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/../../../)
+target_include_directories(tde_cpp_objs PUBLIC ${TORCH_INCLUDE_DIRS})
+target_link_libraries(tde_cpp_objs PUBLIC ${TORCH_LIBRARIES})
+target_compile_options(tde_cpp_objs PUBLIC -fPIC)

torchrec/csrc/dynamic_embedding/details/bitmap.h

@@ -0,0 +1,53 @@
+/*
+ * 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 <stdint.h>
+
+namespace torchrec {
+
+/**
+ * Bitmap
+ *
+ * A bitmap for recording whether num_bits of slots are
+ *  occupied or free.
+ */
+template <typename T = uint32_t>
+struct Bitmap {
+  explicit Bitmap(int64_t num_bits);
+  Bitmap(const Bitmap&) = delete;
+  Bitmap(Bitmap&&) noexcept = default;
+
+  /**
+   * Returns the position of the next free slot.
+   * If the bitmap is full, return `num_total_bits_`.
+   */
+  int64_t next_free_bit();
+
+  /**
+   * Set the slot of position `offset` to free.
+   */
+  void free_bit(int64_t offset);
+
+  /**
+   * Returns if all slots in the bitmap is occupied.
+   */
+  bool full() const;
+
+  static constexpr int64_t num_bits_per_value = sizeof(T) * 8;
+
+  const int64_t num_total_bits_;
+  const int64_t num_values_;
+  std::unique_ptr<T[]> values_;
+
+  int64_t next_free_bit_;
+};
+
+} // namespace torchrec
+
+#include <torchrec/csrc/dynamic_embedding/details/bitmap_impl.h>

torchrec/csrc/dynamic_embedding/details/bitmap_impl.h

@@ -0,0 +1,56 @@
+/*
+ * 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 <stdint.h>
+#include <torchrec/csrc/dynamic_embedding/details/bits_op.h>
+
+namespace torchrec {
+
+template <typename T>
+inline Bitmap<T>::Bitmap(int64_t num_bits)
+    : num_total_bits_(num_bits),
+      num_values_((num_bits + num_bits_per_value - 1) / num_bits_per_value),
+      values_(new T[num_values_]),
+      next_free_bit_(0) {
+  std::fill(values_.get(), values_.get() + num_values_, -1);
+}
+
+template <typename T>
+inline int64_t Bitmap<T>::next_free_bit() {
+  int64_t result = next_free_bit_;
+  int64_t offset = result / num_bits_per_value;
+  T value = values_[offset];
+  // set the last 1 bit to zero
+  values_[offset] = value & (value - 1);
+  while (values_[offset] == 0 && offset < num_values_) {
+    offset++;
+  }
+  value = values_[offset];
+  if (C10_LIKELY(value)) {
+    next_free_bit_ = offset * num_bits_per_value + ctz(value);
+  } else {
+    next_free_bit_ = num_total_bits_;
+  }
+
+  return result;
+}
+
+template <typename T>
+inline void Bitmap<T>::free_bit(int64_t offset) {
+  int64_t mask_offset = offset / num_bits_per_value;
+  int64_t bit_offset = offset % num_bits_per_value;
+  values_[mask_offset] |= 1 << bit_offset;
+  next_free_bit_ = std::min(offset, next_free_bit_);
+}
+template <typename T>
+inline bool Bitmap<T>::full() const {
+  return next_free_bit_ >= num_total_bits_;
+}
+
+} // namespace torchrec

torchrec/csrc/dynamic_embedding/details/clz_impl.cpp

@@ -20,8 +20,6 @@ template <typename T>
 struct ClzImpl {
   /**
    * Naive implementation for no __builtin_clz
-   * @param v
-   * @return
    */
   int operator()(T v) const {
     int result = 0;

torchrec/csrc/dynamic_embedding/details/ctz_impl.cpp

@@ -14,8 +14,6 @@ template <typename T>
 struct CtzImpl {
   /**
    * Naive implementation for no __builtin_ctz
-   * @param v
-   * @return
    */
   int operator()(T v) const {
     if (v == 0)