Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Implement dynamic polymorphism in CPUHashMap #381

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Jan 13, 2025
Merged

Implement dynamic polymorphism in CPUHashMap #381

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
9ca8fc5
update
rusty1s Jan 11, 2025
9284bc3
update
rusty1s Jan 11, 2025
5376a11
update
rusty1s Jan 13, 2025
9cd78e1
Merge branch 'master' into hash_map_benchmark2
rusty1s Jan 13, 2025
7bec40f
update
rusty1s Jan 13, 2025
809ab3d
update
rusty1s Jan 13, 2025
9701e15
update
rusty1s Jan 13, 2025
a32c794
update
rusty1s Jan 13, 2025
540f03e
update
rusty1s Jan 13, 2025
c2b5ace
Merge branch 'master' into hash_map_benchmark3
rusty1s Jan 13, 2025
cc52803
update
rusty1s Jan 13, 2025
7d496d7
Merge branch 'hash_map_benchmark3' of github.com:pyg-team/pyg-lib int…
rusty1s Jan 13, 2025
300cbd1
update
rusty1s Jan 13, 2025
94bba2d
update
rusty1s Jan 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 19 additions & 4 deletions pyg_lib/csrc/classes/cpu/hash_map.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
Expand All @@ -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
Expand Down
24 changes: 19 additions & 5 deletions pyg_lib/csrc/classes/cpu/hash_map.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<
Expand All @@ -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
2 changes: 1 addition & 1 deletion test/csrc/classes/test_hash_map.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Loading