[Benchmark] Benchmark script for KNN Graph (#3080)

* add knn benchmark result

* update

* update knn benchmark to test on multiple algorithms

* update benchmark script

* update

benchmarks/benchmarks/api/bench_knn_graph.py

@@ -5,21 +5,22 @@ import numpy as np
 
 from .. import utils
 
-
 @utils.benchmark('time', timeout=60)
-@utils.parametrize('k', [3, 5, 10])
-@utils.parametrize('size', [50, 200, 1000])
-@utils.parametrize('dim', [16, 128, 512])
-def track_time(size, dim, k):
+@utils.parametrize('k', [8, 64])
+@utils.parametrize('size', [1000, 10000])
+@utils.parametrize('dim', [4, 32, 256])
+@utils.parametrize_cpu('algorithm', ['bruteforce-blas', 'bruteforce', 'kd-tree', 'nn-descent'])
+@utils.parametrize_gpu('algorithm', ['bruteforce-blas', 'bruteforce', 'bruteforce-sharemem', 'nn-descent'])
+def track_time(size, dim, k, algorithm):
     device = utils.get_bench_device()
     features = np.random.RandomState(42).randn(size, dim)
     feat = torch.tensor(features, dtype=torch.float, device=device)
     # dry run
-    for i in range(3):
-        dgl.knn_graph(feat, k)
+    for i in range(1):
+        dgl.knn_graph(feat, k, algorithm=algorithm)
     # timing
     with utils.Timer() as t:
-        for i in range(20):
-            dgl.knn_graph(feat, k)
+        for i in range(5):
+            dgl.knn_graph(feat, k, algorithm=algorithm)
 
-    return t.elapsed_secs / 20
+    return t.elapsed_secs / 5