Implement dynamic polymorphism in CPUHashMap

pyg_lib/csrc/classes/cpu/hash_map.cpp

@@ -7,7 +7,7 @@ namespace pyg {
 namespace classes {
 
 template <typename KeyType>
-CPUHashMap<KeyType>::CPUHashMap(const at::Tensor& key) {
+CPUHashMapImpl<KeyType>::CPUHashMapImpl(const at::Tensor& key) {
   at::TensorArg key_arg{key, "key", 0};
   at::CheckedFrom c{"HashMap.init"};
   at::checkDeviceType(c, key, at::DeviceType::CPU);
@@ -30,7 +30,7 @@ CPUHashMap<KeyType>::CPUHashMap(const at::Tensor& key) {
 };
 
 template <typename KeyType>
-at::Tensor CPUHashMap<KeyType>::get(const at::Tensor& query) {
+at::Tensor CPUHashMapImpl<KeyType>::get(const at::Tensor& query) {
   at::TensorArg query_arg{query, "query", 0};
   at::CheckedFrom c{"HashMap.get"};
   at::checkDeviceType(c, query, at::DeviceType::CPU);
@@ -57,10 +57,25 @@ at::Tensor CPUHashMap<KeyType>::get(const at::Tensor& query) {
   return out;
 }
 
+CPUHashMap::CPUHashMap(const at::Tensor& key) {
+  // clang-format off
+  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Bool,
+  key.scalar_type(),
+  "cpu_hash_map_init",
+  [&] {
+    map_ = std::make_unique<CPUHashMapImpl<scalar_t>>(key);
+  });
+  // clang-format on
+}
+
+at::Tensor CPUHashMap::get(const at::Tensor& query) {
+  return map_->get(query);
+}
+
 TORCH_LIBRARY(pyg, m) {
-  m.class_<CPUHashMap<int64_t>>("CPUHashMap")
+  m.class_<CPUHashMap>("CPUHashMap")
       .def(torch::init<at::Tensor&>())
-      .def("get", &CPUHashMap<int64_t>::get);
+      .def("get", &CPUHashMap::get);
 }
 
 }  // namespace classes

pyg_lib/csrc/classes/cpu/hash_map.h

@@ -6,13 +6,18 @@
 namespace pyg {
 namespace classes {
 
+struct IHashMap {
+  virtual ~IHashMap() = default;
+  virtual at::Tensor get(const at::Tensor& query) = 0;
+};
+
 template <typename KeyType>
-struct CPUHashMap : torch::CustomClassHolder {
+struct CPUHashMapImpl : IHashMap {
  public:
   using ValueType = int64_t;
 
-  CPUHashMap(const at::Tensor& key);
-  at::Tensor get(const at::Tensor& query);
+  CPUHashMapImpl(const at::Tensor& key);
+  at::Tensor get(const at::Tensor& query) override;
 
  private:
   phmap::parallel_flat_hash_map<
@@ -21,10 +26,19 @@ struct CPUHashMap : torch::CustomClassHolder {
       phmap::priv::hash_default_hash<KeyType>,
       phmap::priv::hash_default_eq<KeyType>,
       phmap::priv::Allocator<std::pair<const KeyType, ValueType>>,
-      8,
-      phmap::NullMutex>
+      12,
+      std::mutex>
       map_;
 };
 
+struct CPUHashMap : torch::CustomClassHolder {
+ public:
+  CPUHashMap(const at::Tensor& key);
+  at::Tensor get(const at::Tensor& query);
+
+ private:
+  std::unique_ptr<IHashMap> map_;
+};
+
 }  // namespace classes
 }  // namespace pyg

test/csrc/classes/test_hash_map.cpp

@@ -7,7 +7,7 @@ TEST(CPUHashMapTest, BasicAssertions) {
   auto options = at::TensorOptions().dtype(at::kLong);
   auto key = at::tensor({0, 10, 30, 20}, options);
 
-  auto map = pyg::classes::CPUHashMap<int64_t>(key);
+  auto map = pyg::classes::CPUHashMap(key);
 
   auto query = at::tensor({30, 10, 20, 40}, options);
   auto expected = at::tensor({2, 1, 3, -1}, options);