From b4855b6eceb82a1eea7ce1b50564dd49ab56c9e1 Mon Sep 17 00:00:00 2001
From: Matthias Fey <matthias.fey@tu-dortmund.de>
Date: Sat, 11 Jan 2025 21:43:35 +0100
Subject: [PATCH] `CPUHashMap` benchmark (#379)

---
 benchmark/classes/hash_map.py | 81 ++++++++++++++++++++++++++++++++++-
 1 file changed, 80 insertions(+), 1 deletion(-)

diff --git a/benchmark/classes/hash_map.py b/benchmark/classes/hash_map.py
index 02dbc7d7..e1058307 100644
--- a/benchmark/classes/hash_map.py
+++ b/benchmark/classes/hash_map.py
@@ -1,3 +1,82 @@
+import argparse
+import time
+
+import pandas as pd
+import torch
+
 from pyg_lib.classes import HashMap
 
-print(HashMap)
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cuda')
+    parser.add_argument('--num_keys', type=int, default=10_000_000)
+    parser.add_argument('--num_queries', type=int, default=1_000_000)
+    args = parser.parse_args()
+
+    args.num_queries = min(args.num_queries, args.num_keys)
+
+    num_warmups, num_steps = 50, 100
+    if args.device == 'cpu':
+        num_warmups, num_steps = num_warmups // 10, num_steps // 10
+
+    key = torch.randperm(args.num_keys, device=args.device)
+    query = torch.randperm(args.num_queries, device=args.device)
+    query = query[:args.num_queries]
+
+    t_init = t_get = 0
+    for i in range(num_warmups + num_steps):
+        torch.cuda.synchronize()
+        t_start = time.perf_counter()
+        hash_map = HashMap(key)
+        torch.cuda.synchronize()
+        if i >= num_warmups:
+            t_init += time.perf_counter() - t_start
+
+        t_start = time.perf_counter()
+        hash_map.get(query)
+        torch.cuda.synchronize()
+        if i >= num_warmups:
+            t_get += time.perf_counter() - t_start
+
+    print(f'HashMap Init: {t_init / num_steps:.4f}s')
+    print(f'HashMap  Get: {t_get / num_steps:.4f}s')
+    print('=====================')
+
+    t_init = t_get = 0
+    for i in range(num_warmups + num_steps):
+        torch.cuda.synchronize()
+        t_start = time.perf_counter()
+        hash_map = torch.full((args.num_keys, ), fill_value=-1,
+                              dtype=torch.long, device=args.device)
+        hash_map[key] = torch.arange(args.num_keys)
+        torch.cuda.synchronize()
+        if i >= num_warmups:
+            t_init += time.perf_counter() - t_start
+
+        t_start = time.perf_counter()
+        hash_map[query]
+        torch.cuda.synchronize()
+        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('=====================')
+
+    if key.is_cpu:
+        t_init = t_get = 0
+        for i in range(num_warmups + num_steps):
+            t_start = time.perf_counter()
+            hash_map = pd.CategoricalDtype(categories=key.numpy(),
+                                           ordered=True)
+            if i >= num_warmups:
+                t_init += time.perf_counter() - t_start
+
+            t_start = time.perf_counter()
+            ser = pd.Series(query.numpy(), dtype=hash_map)
+            ser.cat.codes.to_numpy()
+            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')