From c74cf4f52f4b718f5b92b14a612fb47cd687b7bc Mon Sep 17 00:00:00 2001
From: Matthijs Douze <matthijs@meta.com>
Date: Tue, 24 Jan 2023 11:24:47 -0800
Subject: [PATCH] Support CPU quantizer parameters in GPU autotune

Summary:
IVF GPU indexes now support CPU quantizers.
In order for the the auto-tuning to support accessing the parameters of the quantizers, GpuAutoTune needs to be adapted.

Reviewed By: algoriddle

Differential Revision: D42633348

fbshipit-source-id: ef472855aa882ccde9d878ae09782204045e38c5
---
 faiss/gpu/GpuAutoTune.cpp        | 16 ++++++++++++++++
 faiss/gpu/test/test_gpu_index.py | 18 ++++++++++++++++++
 faiss/python/swigfaiss.swig      |  1 +
 3 files changed, 35 insertions(+)

diff --git a/faiss/gpu/GpuAutoTune.cpp b/faiss/gpu/GpuAutoTune.cpp
index b5f8ebbb2..dc335e2e5 100644
--- a/faiss/gpu/GpuAutoTune.cpp
+++ b/faiss/gpu/GpuAutoTune.cpp
@@ -53,6 +53,14 @@ void GpuParameterSpace::initialize(const Index* index) {
                 break;
             pr.values.push_back(nprobe);
         }
+
+        ParameterSpace ivf_pspace;
+        ivf_pspace.initialize(ix->quantizer);
+
+        for (const ParameterRange& p : ivf_pspace.parameter_ranges) {
+            ParameterRange& pr = add_range("quantizer_" + p.name);
+            pr.values = p.values;
+        }
     }
     // not sure we should call the parent initializer
 }
@@ -83,6 +91,14 @@ void GpuParameterSpace::set_index_parameter(
         }
     }
 
+    if (name.find("quantizer_") == 0) {
+        if (DC(GpuIndexIVF)) {
+            std::string sub_name = name.substr(strlen("quantizer_"));
+            set_index_parameter(ix->quantizer, sub_name, val);
+            return;
+        }
+    }
+
     // maybe normal index parameters apply?
     ParameterSpace::set_index_parameter(index, name, val);
 }
diff --git a/faiss/gpu/test/test_gpu_index.py b/faiss/gpu/test/test_gpu_index.py
index ae75769c6..048ee3841 100755
--- a/faiss/gpu/test/test_gpu_index.py
+++ b/faiss/gpu/test/test_gpu_index.py
@@ -574,3 +574,21 @@ class TestLSQIcmEncoder(unittest.TestCase):
     def test_multiple_gpu(self):
         ngpu = faiss.get_num_gpus()
         self.subtest_gpu_encoding(ngpu)
+
+
+class TestGpuAutoTune(unittest.TestCase):
+
+    def test_params(self):
+        index = faiss.index_factory(32, "IVF65536_HNSW,PQ16")
+        index = faiss.index_cpu_to_gpu(faiss.StandardGpuResources(), 0, index)
+        ps = faiss.GpuParameterSpace()
+        ps.initialize(index)
+        for i in range(ps.parameter_ranges.size()):
+            pr = ps.parameter_ranges.at(i)
+            if pr.name == "quantizer_efSearch":
+                break
+        else:
+            self.fail("should include efSearch")
+        ps.set_index_parameter(index, "quantizer_efSearch", 123)
+        quantizer = faiss.downcast_index(index.quantizer)
+        self.assertEqual(quantizer.hnsw.efSearch, 123)
diff --git a/faiss/python/swigfaiss.swig b/faiss/python/swigfaiss.swig
index 57a23db76..b1785b5a8 100644
--- a/faiss/python/swigfaiss.swig
+++ b/faiss/python/swigfaiss.swig
@@ -246,6 +246,7 @@ namespace std {
 %template(InvertedListsPtrVector) std::vector<faiss::InvertedLists*>;
 %template(RepeatVector) std::vector<faiss::Repeat>;
 %template(ClusteringIterationStatsVector) std::vector<faiss::ClusteringIterationStats>;
+%template(ParameterRangeVector) std::vector<faiss::ParameterRange>;
 
 #ifndef SWIGWIN
 %template(OnDiskOneListVector) std::vector<faiss::OnDiskOneList>;