Simplify bench/ann scripts to Python based module

conda/environments/bench_ann_cuda-118_arch-x86_64.yaml

@@ -30,6 +30,7 @@ dependencies:
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
 - libfaiss>=1.7.1
+- mathplotlib
 - nccl>=2.9.9
 - ninja
 - nlohmann_json>=3.11.2

cpp/bench/ann/algos.yaml

@@ -0,0 +1,30 @@
+faise_gpu_ivf_flat:
+  executable: FAISS_IVF_FLAT_ANN_BENCH
+  disabled: false
+faiss_gpu_flat:
+  executable: FAISS_IVF_FLAT_ANN_BENCH
+  disabled: false
+faiss_gpu_ivf_pq:
+  executable: FAISS_IVF_PQ_ANN_BENCH
+  disabled: false
+faiss_gpu_ivf_sq:
+  executable: FAISS_IVF_PQ_ANN_BENCH
+  disabled: false
+faiss_gpu_bfknn:
+  executable: FAISS_BFKNN_ANN_BENCH
+  disabled: false
+raft_ivf_flat:
+  executable: RAFT_IVF_FLAT_ANN_BENCH
+  disabled: false
+raft_ivf_pq:
+  executable: RAFT_IVF_PQ_ANN_BENCH
+  disabled: false
+raft_cagra:
+  executable: RAFT_CAGRA_ANN_BENCH
+  disabled: false
+ggnn:
+  executable: GGNN_ANN_BENCH
+  disabled: false
+hnswlib:
+  executable: HNSWLIB_ANN_BENCH
+  disabled: false

cpp/bench/ann/conf/glove-100-inner.json

@@ -789,9 +789,5 @@
 
       ],
       "search_result_file" : "result/glove-100-inner/ggnn/kbuild96-segment64-refine2-k10"
-    },
-
-
-  ]
-
+    }]
 }

cpp/bench/ann/scripts/data_export.py

@@ -0,0 +1,57 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import subprocess
+
+
+def export_results(output_filepath, recompute, groundtruth_filepath,
+                   result_filepaths):
+    print(f"Writing output file to: {os.path.join(os.getcwd(), output_filepath)}")
+    if recompute:
+        p = subprocess.Popen(["scripts/eval.pl", "-f", "-o", output_filepath,
+                              groundtruth_filepath] + result_filepaths)
+    else:
+        p = subprocess.Popen(["scripts/eval.pl", "-o", output_filepath,
+                              groundtruth_filepath] + result_filepaths)
+    p.wait()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--output", help="Path to the CSV output file",
+                        required=True)
+    parser.add_argument("--recompute", action="store_true",
+                        help="Recompute metrics")
+    parser.add_argument("--groundtruth",
+                        help="Dataset whose groundtruth is used",
+                        required=True)
+    args, result_filepaths = parser.parse_known_args()
+
+    # assume "result/<groundtruth_dataset>" folder to be default
+    # if nothing is provided
+    if len(result_filepaths) == 0:
+        result_filepaths = [f"result/{args.groundtruth}"]
+
+    groundtruth_filepath = os.path.join("data", args.groundtruth,
+                                        "groundtruth.neighbors.ibin")
+    export_results(args.output, args.recompute, groundtruth_filepath,
+                   result_filepaths)
+
+
+if __name__ == "__main__":
+    main()

cpp/bench/ann/scripts/get_dataset.py

@@ -0,0 +1,87 @@
+#
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import os
+import subprocess
+from urllib.request import urlretrieve
+
+
+def get_dataset_path(name):
+    if not os.path.exists("data"):
+        os.mkdir("data")
+    return os.path.join("data", f"{name}.hdf5")
+
+
+def download_dataset(url, path):
+    if not os.path.exists(path):
+        print(f"downloading {url} -> {path}...")
+        urlretrieve(url, path)
+
+
+def convert_hdf5_to_fbin(path, normalize):
+    if normalize and "angular" in path:
+        p = subprocess.Popen(["python", "scripts/hdf5_to_fbin.py", "-n",
+                              "%s" % path])
+    else:
+        p = subprocess.Popen(["python", "scripts/hdf5_to_fbin.py",
+                              "%s" % path])
+    p.wait()
+
+
+def move(name, path):
+    if "angular" in name:
+        new_name = name.replace("angular", "inner")
+    else:
+        new_name = name
+    new_path = os.path.join("data", new_name)
+    if not os.path.exists(new_path):
+        os.mkdir(new_path)
+    for bin_name in ["base.fbin", "query.fbin", "groundtruth.neighbors.ibin",
+                     "groundtruth.distances.fbin"]:
+        os.rename(f"data/{name}.{bin_name}",
+                  f"{new_path}/{bin_name}")
+
+
+def download(name, normalize):
+    path = get_dataset_path(name)
+    try:
+        url = f"http://ann-benchmarks.com/{name}.hdf5"
+        download_dataset(url, path)
+
+        convert_hdf5_to_fbin(path, normalize)
+
+        move(name, path)
+    except Exception:
+        print(f"Cannot download {url}")
+        raise
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--name", help="dataset to download",
+                        default="glove-100-angular")
+    parser.add_argument("--normalize",
+                        help="normalize cosine distance to inner product",
+                        action="store_true")
+
+    args = parser.parse_args()
+
+    download(args.name, args.normalize)
+
+
+if __name__ == "__main__":
+    main()