Add RAFT ANN benchmark for CAGRA (#1552)

This PR adds the RAFT ANN benchmark for CAGRA Authors: - tsuki (https://github.com/enp1s0) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - Tamas Bela Feher (https://github.com/tfeher) URL: #1552
rapidsai · Jun 23, 2023 · 185e933 · 185e933
1 parent 0db8f15
commit 185e933
Show file tree

Hide file tree

Showing 9 changed files with 395 additions and 73 deletions.
diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt
@@ -21,6 +21,7 @@ option(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ "Include faiss' ivf pq algorithm in bench
 option(RAFT_ANN_BENCH_USE_RAFT_BFKNN "Include raft's brute-force knn algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT "Include raft's ivf flat algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ "Include raft's ivf pq algorithm in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_RAFT_CAGRA "Include raft's CAGRA in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" ON)
 
@@ -36,8 +37,9 @@ endif()
 
 set(RAFT_ANN_BENCH_USE_RAFT OFF)
 if(RAFT_ANN_BENCH_USE_RAFT_BFKNN
-   OR RAFT_ANN_BENCH_USE_RAFT_IVFPQ
-   OR RAFT_ANN_BENCH_USE_RAFT_IVFFLAT
+   OR RAFT_ANN_BENCH_USE_RAFT_IVF_PQ
+   OR RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT
+   OR RAFT_ANN_BENCH_USE_RAFT_CAGRA
 )
   set(RAFT_ANN_BENCH_USE_RAFT ON)
 endif()
@@ -136,14 +138,24 @@ endif()
 if(RAFT_ANN_BENCH_USE_RAFT)
   ConfigureAnnBench(
     NAME
-    RAFT_IVF_PQ
+    RAFT
     PATH
     bench/ann/src/raft/raft_benchmark.cu
     $<$<BOOL:${RAFT_ANN_BENCH_USE_RAFT_IVF_PQ}>:bench/ann/src/raft/raft_ivf_pq.cu>
     $<$<BOOL:${RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT}>:bench/ann/src/raft/raft_ivf_flat.cu>
+    $<$<BOOL:${RAFT_ANN_BENCH_USE_RAFT_CAGRA}>:bench/ann/src/raft/raft_cagra.cu>
     LINKS
     raft::compiled
   )
+  add_compile_definitions(
+    RAFT $<$<BOOL:${RAFT_ANN_BENCH_USE_RAFT_IVF_PQ}>:RAFT_ANN_BENCH_USE_RAFT_IVF_PQ>
+  )
+  add_compile_definitions(
+    RAFT $<$<BOOL:${RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT}>:RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT>
+  )
+  add_compile_definitions(
+    RAFT $<$<BOOL:${RAFT_ANN_BENCH_USE_RAFT_CAGRA}>:RAFT_ANN_BENCH_USE_RAFT_CAGRA>
+  )
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS)

diff --git a/cpp/bench/ann/conf/bigann-100M.json b/cpp/bench/ann/conf/bigann-100M.json
@@ -168,7 +168,35 @@
       "search_result_file" : "result/bigann-100M/ivf_flat/nlist100K"
     },
 
+    {
+      "name" : "cagra.dim32",
+      "algo" : "cagra",
+      "build_param": {
+        "index_dim" : 32
+      },
+      "file" : "index/bigann-100M/cagra/dim32",
+      "search_params" : [
+        "itopk": 32,
+        "itopk": 64,
+        "itopk": 128
+      ],
+      "search_result_file" : "result/bigann-100M/cagra/dim32"
+    },
 
 
+    {
+      "name" : "cagra.dim64",
+      "algo" : "cagra",
+      "build_param": {
+        "index_dim" : 64
+      },
+      "file" : "index/bigann-100M/cagra/dim64",
+      "search_params" : [
+        "itopk": 32,
+        "itopk": 64,
+        "itopk": 128
+      ],
+      "search_result_file" : "result/bigann-100M/cagra/dim64"
+    }
   ]
 }
diff --git a/cpp/bench/ann/conf/deep-100M.json b/cpp/bench/ann/conf/deep-100M.json
@@ -218,6 +218,35 @@
       "search_result_file" : "result/deep-100M/ivf_flat/nlist100K"
     },
 
+    {
+      "name" : "cagra.dim32",
+      "algo" : "raft_cagra",
+      "build_param": {
+        "index_dim" : 32
+      },
+      "file" : "index/deep-100M/cagra/dim32",
+      "search_params" : [
+        {"itopk": 32},
+        {"itopk": 64},
+        {"itopk": 128}
+      ],
+      "search_result_file" : "result/deep-100M/cagra/dim32"
+    },
 
+
+    {
+      "name" : "cagra.dim64",
+      "algo" : "raft_cagra",
+      "build_param": {
+        "index_dim" : 64
+      },
+      "file" : "index/deep-100M/cagra/dim64",
+      "search_params" : [
+        {"itopk": 32},
+        {"itopk": 64},
+        {"itopk": 128}
+      ],
+      "search_result_file" : "result/deep-100M/cagra/dim64"
+    }
   ]
 }
diff --git a/cpp/bench/ann/conf/sift-128-euclidean.json b/cpp/bench/ann/conf/sift-128-euclidean.json
@@ -90,8 +90,8 @@
 
 
 
-      
-	  {
+
+    {
       "name": "raft_bfknn",
       "algo": "raft_bfknn",
       "build_param": {},
@@ -1316,6 +1316,36 @@
         }
       ],
       "search_result_file": "result/sift-128-euclidean/raft_ivf_flat/nlist16384"
+    },
+
+    {
+      "name" : "cagra.dim32",
+      "algo" : "raft_cagra",
+      "build_param": {
+        "index_dim" : 32
+      },
+      "file" : "index/sift-128-euclidean/cagra/dim32",
+      "search_params" : [
+        {"itopk": 32},
+        {"itopk": 64},
+        {"itopk": 128}
+      ],
+      "search_result_file" : "result/sift-128-euclidean/cagra/dim32"
+    },
+
+    {
+      "name" : "cagra.dim64",
+      "algo" : "raft_cagra",
+      "build_param": {
+        "index_dim" : 64
+      },
+      "file" : "index/sift-128-euclidean/cagra/dim64",
+      "search_params" : [
+        {"itopk": 32},
+        {"itopk": 64},
+        {"itopk": 128}
+      ],
+      "search_result_file" : "result/sift-128-euclidean/cagra/dim64"
     }
   ]
 }