Skip to content
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

Update benchmark script to respect num_submaps #402

Merged
merged 3 commits into from
Feb 6, 2025
Merged

Update benchmark script to respect num_submaps #402

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8412b57
update
rusty1s Feb 6, 2025
2ce9dc3
update
rusty1s Feb 6, 2025
8b6313b
update
rusty1s Feb 6, 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
65 changes: 39 additions & 26 deletions benchmark/classes/hash_map.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,39 +26,52 @@
max_value = torch.iinfo(dtype).max

key1 = torch.randint(0, max_value, (args.num_keys, ), dtype=dtype,
device=args.device)
device=args.device).unique()
query1 = key1[torch.randperm(key1.size(0), device=args.device)]
query1 = query1[:args.num_queries]

key2 = torch.randperm(args.num_keys, dtype=dtype, device=args.device)
query2 = torch.randperm(args.num_queries, dtype=dtype, device=args.device)
query2 = query2[:args.num_queries]

if key1.is_cpu:
HashMap = torch.classes.pyg.CPUHashMap
elif key1.is_cuda:
HashMap = torch.classes.pyg.CUDAHashMap
if key1.is_cuda:
t_init = t_get = 0
for i in range(num_warmups + num_steps):
torch.cuda.synchronize()
t_start = time.perf_counter()
hash_map = torch.classes.pyg.CUDAHashMap(key1, 0.5)
torch.cuda.synchronize()
if i >= num_warmups:
t_init += time.perf_counter() - t_start

t_start = time.perf_counter()
out1 = hash_map.get(query1)
torch.cuda.synchronize()
if i >= num_warmups:
t_get += time.perf_counter() - t_start

print(f'HashMap Init: {t_init:.4f}s')
print(f'HashMap Get: {t_get:.4f}s')
print('=====================')
else:
raise NotImplementedError(f"Unsupported device '{args.device}'")
for num_submaps in [-1, 0, 16, 256]:
t_init = t_get = 0
for i in range(num_warmups + num_steps):
t_start = time.perf_counter()
hash_map = torch.classes.pyg.CPUHashMap(key1, num_submaps)
if i >= num_warmups:
t_init += time.perf_counter() - t_start

t_init = t_get = 0
for i in range(num_warmups + num_steps):
torch.cuda.synchronize()
t_start = time.perf_counter()
hash_map = HashMap(key1)
torch.cuda.synchronize()
if i >= num_warmups:
t_init += time.perf_counter() - t_start
t_start = time.perf_counter()
out1 = hash_map.get(query1)
if i >= num_warmups:
t_get += time.perf_counter() - t_start

t_start = time.perf_counter()
out1 = hash_map.get(query1)
torch.cuda.synchronize()
if i >= num_warmups:
t_get += time.perf_counter() - t_start
print(f'HashMap[{num_submaps}] Init: {t_init:.4f}s')
print(f'HashMap[{num_submaps}] Get: {t_get:.4f}s')
print('=====================')

print(f'HashMap Init: {t_init / num_steps:.4f}s')
print(f'HashMap Get: {t_get / num_steps:.4f}s')
print('=====================')
quit()

t_init = t_get = 0
for i in range(num_warmups + num_steps):
Expand All @@ -78,8 +91,8 @@
if i >= num_warmups:
t_get += time.perf_counter() - t_start

print(f' Memory Init: {t_init / num_steps:.4f}s')
print(f' Memory Get: {t_get / num_steps:.4f}s')
print(f' Memory Init: {t_init:.4f}s')
print(f' Memory Get: {t_get:.4f}s')
print('=====================')

if key1.is_cpu:
Expand All @@ -97,7 +110,7 @@
if i >= num_warmups:
t_get += time.perf_counter() - t_start

print(f' Pandas Init: {t_init / num_steps:.4f}s')
print(f' Pandas Get: {t_get / num_steps:.4f}s')
print(f' Pandas Init: {t_init:.4f}s')
print(f' Pandas Get: {t_get:.4f}s')

assert out1.equal(torch.tensor(out3))
64 changes: 27 additions & 37 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 @@ -27,9 +27,8 @@ struct CPUHashMapImpl : HashMapImpl {
public:
using ValueType = int64_t;

CPUHashMapImpl(const at::Tensor& key, double load_factor) {
size_t capacity = std::ceil(key.numel() / load_factor);
map_.reserve(capacity);
CPUHashMapImpl(const at::Tensor& key) {
map_.reserve(key.numel());

const auto key_data = key.data_ptr<KeyType>();
for (int64_t i = 0; i < key.numel(); ++i) {
Expand Down Expand Up @@ -89,9 +88,8 @@ struct ParallelCPUHashMapImpl : HashMapImpl {
public:
using ValueType = int64_t;

ParallelCPUHashMapImpl(const at::Tensor& key, double load_factor) {
size_t capacity = std::ceil(key.numel() / load_factor);
map_.reserve(capacity);
ParallelCPUHashMapImpl(const at::Tensor& key) {
map_.reserve(key.numel());

const auto key_data = key.data_ptr<KeyType>();

Expand Down Expand Up @@ -157,16 +155,13 @@ struct ParallelCPUHashMapImpl : HashMapImpl {
phmap::priv::hash_default_hash<KeyType>,
phmap::priv::hash_default_eq<KeyType>,
phmap::priv::Allocator<std::pair<const KeyType, ValueType>>,
num_submaps,
std::mutex>
num_submaps>
map_;
};

struct CPUHashMap : torch::CustomClassHolder {
public:
CPUHashMap(const at::Tensor& key,
int64_t num_submaps = 0,
double load_factor = 0.5) {
CPUHashMap(const at::Tensor& key, int64_t num_submaps = 0) {
at::TensorArg key_arg{key, "key", 0};
at::CheckedFrom c{"CPUHashMap.init"};
at::checkDeviceType(c, key, at::DeviceType::CPU);
Expand All @@ -175,56 +170,51 @@ struct CPUHashMap : torch::CustomClassHolder {

DISPATCH_KEY(key.scalar_type(), "cpu_hash_map_init", [&] {
switch (num_submaps) {
case -1: // Auto-infer:
if (key.numel() < 200'000) {
map_ = std::make_unique<CPUHashMapImpl<scalar_t>>(key);
} else {
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 8>>(key);
}
break;
case 0:
map_ = std::make_unique<CPUHashMapImpl<scalar_t>>(key, load_factor);
map_ = std::make_unique<CPUHashMapImpl<scalar_t>>(key);
break;
case 2:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 1>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 1>>(key);
break;
case 4:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 2>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 2>>(key);
break;
case 8:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 3>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 3>>(key);
break;
case 16:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 4>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 4>>(key);
break;
case 32:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 5>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 5>>(key);
break;
case 64:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 6>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 6>>(key);
break;
case 128:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 7>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 7>>(key);
break;
case 256:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 8>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 8>>(key);
break;
case 512:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 9>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 9>>(key);
break;
case 1024:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 10>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 10>>(key);
break;
case 2048:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 11>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 11>>(key);
break;
case 4096:
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 12>>(
key, load_factor);
map_ = std::make_unique<ParallelCPUHashMapImpl<scalar_t, 12>>(key);
break;
default:
TORCH_CHECK(false, "'num_submaps' needs to be a power of 2");
Expand Down Expand Up @@ -252,7 +242,7 @@ struct CPUHashMap : torch::CustomClassHolder {

TORCH_LIBRARY_FRAGMENT(pyg, m) {
m.class_<CPUHashMap>("CPUHashMap")
.def(torch::init<at::Tensor&, int64_t, double>())
.def(torch::init<at::Tensor&, int64_t>())
.def("get", &CPUHashMap::get)
.def("keys", &CPUHashMap::keys)
.def_pickle(
Expand Down
2 changes: 1 addition & 1 deletion pyg_lib/csrc/classes/cuda/hash_map.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ struct CUDAHashMapImpl : HashMapImpl {
map_->retrieve_all(key_data, value_data);

const auto perm = at::empty_like(value);
perm[value] = at::arange(value.numel(), value.options());
perm.scatter_(0, value, at::arange(value.numel(), value.options()));

return key.index_select(0, perm);
}
Expand Down
20 changes: 8 additions & 12 deletions test/classes/test_hash_map.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,31 +8,30 @@


@withCUDA
@pytest.mark.parametrize('load_factor', [0.5, 0.25])
@pytest.mark.parametrize('dtype', [torch.short, torch.int, torch.long])
def test_hash_map(load_factor, dtype, device):
def test_hash_map(dtype, device):
key = torch.tensor([0, 10, 30, 20], device=device, dtype=dtype)
query = torch.tensor([30, 10, 20, 40], device=device, dtype=dtype)

if key.is_cpu:
HashMap = torch.classes.pyg.CPUHashMap
hash_map = HashMap(key, 0, load_factor)
hash_map = HashMap(key, 0)
elif key.is_cuda:
HashMap = torch.classes.pyg.CUDAHashMap
hash_map = HashMap(key, load_factor)
hash_map = HashMap(key, 0.5)
else:
raise NotImplementedError(f"Unsupported device '{device}'")

assert hash_map.keys().equal(key)
assert hash_map.keys().equal(key)
assert hash_map.keys().dtype == dtype
expected = torch.tensor([2, 1, 3, -1], device=device)
assert hash_map.get(query).equal(expected)
assert hash_map.get(query).dtype == torch.long

if key.is_cpu:
hash_map = HashMap(key, 16, load_factor)
assert hash_map.keys().dtype == dtype
if key.is_cpu: # Test parallel hash map:
hash_map = HashMap(key, 16)
assert hash_map.keys().equal(key)
assert hash_map.keys().dtype == dtype
assert hash_map.get(query).equal(expected)
assert hash_map.get(query).dtype == torch.long

Expand All @@ -42,14 +41,11 @@ def __init__(self, key: Tensor):
super().__init__()
if key.is_cpu:
HashMap = torch.classes.pyg.CPUHashMap
self.map = HashMap(key, 0, 0.5)
self.map = HashMap(key, 0)
elif key.is_cuda:
HashMap = torch.classes.pyg.CUDAHashMap
self.map = HashMap(key, 0.5)

def forward(self, query: Tensor) -> Tensor:
return self.map.get(query)


@withCUDA
def test_serialization(device, tmp_path):
Expand Down
2 changes: 1 addition & 1 deletion third_party/parallel-hashmap
Open in desktop
Submodule parallel-hashmap updated 47 files
+0 −57 .appveyor_old.yml
+19 −7 .github/workflows/linux.yml
+19 −7 .github/workflows/macos.yml
+18 −8 .github/workflows/windows.yml
+5 −0 .gitignore
+10 −0 CITATION.cff
+133 −54 CMakeLists.txt
+41 −27 README.md
+6 −0 cmake/DownloadGTest.cmake
+1 −2 cmake/helpers.cmake
+36 −0 doc/new_release.md
+64 −0 examples/custom_pointer.cc
+6 −6 examples/dump_nested.cc
+48 −0 examples/hash.cc
+127 −0 examples/hash_bench.cc
+1 −10 examples/lazy_emplace_l.cc
+453 −0 examples/llil.cc
+507 −0 examples/llil4map.cc
+52 −0 examples/llil_utils/gen-llil.pl
+34 −0 examples/llil_utils/gen_files
+4 −0 examples/llil_utils/run_llil4map
+6 −0 examples/llil_utils/shuffle.pl
+3 −2 examples/matt.cc
+84 −0 examples/mt_word_counter.cc
+190 −0 examples/p_bench.cc
+1 −1 examples/pmr.cc
+1 −1 examples/serialize.cc
+39 −43 parallel_hashmap/btree.h
+0 −37 parallel_hashmap/conanfile.py
+475 −294 parallel_hashmap/phmap.h
+243 −293 parallel_hashmap/phmap_base.h
+7 −7 parallel_hashmap/phmap_bits.h
+33 −43 parallel_hashmap/phmap_config.h
+81 −8 parallel_hashmap/phmap_dump.h
+32 −0 parallel_hashmap/phmap_fwd_decl.h
+54 −25 parallel_hashmap/phmap_utils.h
+17 −6 phmap_lldb.py
+19 −31 tests/btree_test.cc
+4 −4 tests/btree_test.h
+0 −201 tests/compressed_tuple_test.cc
+42 −0 tests/dump_load_test.cc
+10 −0 tests/flat_hash_map_test.cc
+2 −0 tests/parallel_flat_hash_map_mutex_test.cc
+48 −0 tests/parallel_flat_hash_map_test.cc
+29 −0 tests/parallel_hash_map_test.cc
+2 −1 tests/raw_hash_set_allocator_test.cc
+17 −48 tests/raw_hash_set_test.cc
Loading