Add query blocking to IndexFlatPanorama search

Process queries in blocks so each DB level-storage block is streamed from
memory once and reused across the block while resident in a faster cache
level, raising cache-bandwidth efficiency on the bandwidth-bound Panorama dot
kernel. The DB-batch loop stays outside the block (preserving each query's
threshold evolution) and within a batch the loop runs level-outer /
query-inner.

Controlled by the global panorama_query_block_size, following the existing
FAISS tunable idiom (cf. distance_compute_blas_query_bs). It defaults to an
enabled value of 32; 0 or 1 selects the original query-at-a-time path. Pure
loop-order transform: each query keeps its own active set, threshold, and
pruning decisions, so results are bit-identical to the original path (verified
432/432 equality cases + unchanged SIFT1M recall). Top-k only; range search
falls back to the original path since RangeSearchPartialResult requires
per-query-contiguous appends.

Speedup up to ~2.8x (n_levels=4) on SIFT1M 1M-vec single core; holds
multithreaded by relieving shared-cache contention.

Also add --query-block-size (comma-separated sweep) and --nq to
bench_flat_l2_panorama so the speedup can be measured directly: the Panorama
index is built once and only the block-size global is toggled between runs,
printing speedups vs plain Flat and vs the Panorama baseline.

Signed-off-by: Mulugeta Mammo <mulugeta.mammo@intel.com>

Author: mulugetam

benchs/bench_flat_l2_panorama.py

@@ -18,14 +18,28 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--dataset", default="gist1m", choices=["sift1m", "gist1m"])
+parser.add_argument(
+    "--nq",
+    type=int,
+    default=1000,
+    help="number of queries (query blocking needs a real batch to help)",
+)
+parser.add_argument(
+    "--query-block-size",
+    default="0",
+    help="comma-separated block sizes to sweep for the Panorama index; "
+    "0 or 1 selects the original query-at-a-time path",
+)
 args = parser.parse_args()
 
+query_block_sizes = [int(v) for v in args.query_block_size.split(",")]
+
 if args.dataset == "sift1m":
     ds = DatasetSIFT1M()
 else:
     ds = DatasetGIST1M()
 
-nq = 10
+nq = args.nq
 xq = ds.get_queries()[:nq]
 xb = ds.get_database()
 gt = ds.get_groundtruth()[:nq]
@@ -73,32 +87,53 @@ def build_index(name):
 
 plt.figure(figsize=(8, 6), dpi=80)
 
-names = [
-    "Flat",
-    f"PCA{d},FlatL2Panorama{nlevels}_{batch_size}",
-]
+pano_name = f"PCA{d},FlatL2Panorama{nlevels}_{batch_size}"
 
 labels = []
 qps_values = []
 
-for name in names:
-    print(f"======{name}")
-    index = build_index(name)
-    recall, qps = eval_qps(index)
-    labels.append(f"{name}\n(r@{recall:.3f})")
+# Plain Flat baseline.
+print("======Flat")
+flat_index = build_index("Flat")
+flat_recall, flat_qps = eval_qps(flat_index)
+labels.append(f"Flat\n(r@{flat_recall:.3f})")
+qps_values.append(flat_qps)
+
+# Panorama index, swept over query block sizes. The index is built once; only
+# the global query-block toggle changes between runs, so results must match.
+print(f"======{pano_name}")
+pano_index = build_index(pano_name)
+pano_qps_values = []
+for qbs in query_block_sizes:
+    faiss.cvar.panorama_query_block_size = qbs
+    tag = "baseline" if qbs <= 1 else f"block={qbs}"
+    print(f"---{tag}")
+    recall, qps = eval_qps(pano_index)
+    labels.append(f"Pano\n{tag}\n(r@{recall:.3f})")
     qps_values.append(qps)
+    pano_qps_values.append(qps)
+faiss.cvar.panorama_query_block_size = 0
+
+# Report speedups: Panorama-vs-Flat, and blocked-vs-baseline-Panorama.
+pano_baseline_qps = pano_qps_values[0]
+print("\n=== summary ===")
+print(f"Flat: {flat_qps:.1f} QPS")
+for qbs, qps in zip(query_block_sizes, pano_qps_values):
+    tag = "baseline" if qbs <= 1 else f"block={qbs}"
+    vs_flat = qps / flat_qps
+    vs_base = qps / pano_baseline_qps
+    print(
+        f"Pano {tag:>10}: {qps:8.1f} QPS  "
+        f"({vs_flat:.2f}x vs Flat, {vs_base:.2f}x vs Pano baseline)"
+    )
 
-x = np.arange(len(names))
-plt.bar(x, qps_values, color=["#1f77b4", "#ff7f0e"])
-speedup = qps_values[1] / qps_values[0]
+x = np.arange(len(qps_values))
+colors = ["#1f77b4"] + ["#ff7f0e"] * len(pano_qps_values)
+plt.bar(x, qps_values, color=colors)
 ax = plt.gca()
-ax.text(
-    x[1],
-    qps_values[1] * 1.01,
-    f"{speedup:.2f}x",
-    ha="center",
-    va="bottom",
-)
+# Annotate each Panorama bar with its speedup over plain Flat.
+for xi, qps in zip(x[1:], pano_qps_values):
+    ax.text(xi, qps * 1.01, f"{qps / flat_qps:.2f}x", ha="center", va="bottom")
 plt.xticks(x, labels, rotation=0)
 plt.ylabel("QPS")
 dataset_label = args.dataset.upper()

faiss/IndexFlat.cpp

@@ -23,6 +23,8 @@
 
 namespace faiss {
 
+size_t panorama_query_block_size = 32;
+
 IndexFlat::IndexFlat(idx_t d_, MetricType metric)
         : IndexFlatCodes(sizeof(float) * d_, d_, metric) {}
 
@@ -571,6 +573,120 @@ inline auto dispatch_metric_compare(MetricType metric, Fn&& fn) {
     return fn.template operator()<C>();
 }
 
+/// Query-blocked search core. For each block of up to `qbs` queries, iterates
+/// DB batches in order (preserving each query's threshold evolution), and
+/// within each batch iterates level-outer / query-inner so the level storage
+/// block is streamed from memory once and reused across the block while
+/// resident in a faster cache level. Produces results identical to
+/// flat_pano_search_core.
+template <bool use_radius, typename C, typename BlockHandler>
+inline void flat_pano_search_core_blocked(
+        const IndexFlatPanorama& index,
+        BlockHandler& handler,
+        idx_t n,
+        const float* x,
+        float radius,
+        const SearchParameters* params,
+        size_t qbs) {
+    using SingleResultHandler = typename BlockHandler::SingleResultHandler;
+
+    IDSelector* sel = params ? params->sel : nullptr;
+    bool use_sel = sel != nullptr;
+
+    const size_t d = index.d;
+    const size_t bs = index.batch_size;
+    const size_t n_lp1 = index.pano.n_levels + 1;
+    size_t n_batches = (index.ntotal + bs - 1) / bs;
+
+    size_t n_blocks = (size_t(n) + qbs - 1) / qbs;
+    [[maybe_unused]] int nt = std::min(int(n_blocks), omp_get_max_threads());
+
+#pragma omp parallel num_threads(nt)
+    {
+        // One persistent SingleResultHandler per query slot so each heap and
+        // threshold survive across DB batches (begin() re-heapifies, so it must
+        // be called exactly once per query, not once per batch).
+        std::vector<SingleResultHandler> ress;
+        ress.reserve(qbs);
+        for (size_t qb = 0; qb < qbs; qb++) {
+            ress.emplace_back(handler);
+        }
+
+        std::vector<float> query_cum_norms(qbs * n_lp1);
+        std::vector<uint32_t> active_indices(qbs * bs);
+        std::vector<uint8_t> active_byteset(qbs * bs);
+        std::vector<float> exact_distances(qbs * bs);
+        std::vector<float> dot_buffer(bs);
+        std::vector<float> thresholds(qbs);
+        std::vector<size_t> num_active(qbs);
+
+#pragma omp for
+        for (int64_t blk = 0; blk < int64_t(n_blocks); blk++) {
+            size_t q0 = size_t(blk) * qbs;
+            size_t block_size = std::min(qbs, size_t(n) - q0);
+
+            PanoramaStats local_stats;
+            local_stats.reset();
+
+            // Per-query: compute cum sums and open the result heap (once).
+            for (size_t qb = 0; qb < block_size; qb++) {
+                const float* xi = x + (q0 + qb) * d;
+                index.pano.compute_query_cum_sums(
+                        xi, query_cum_norms.data() + qb * n_lp1);
+                ress[qb].begin(q0 + qb);
+            }
+
+            for (size_t batch_no = 0; batch_no < n_batches; batch_no++) {
+                size_t batch_start = batch_no * bs;
+
+                // Snapshot each query's current threshold for this batch.
+                for (size_t qb = 0; qb < block_size; qb++) {
+                    if constexpr (use_radius) {
+                        thresholds[qb] = radius;
+                    } else {
+                        thresholds[qb] = ress[qb].threshold;
+                    }
+                }
+
+                with_metric_type(index.metric_type, [&]<MetricType M>() {
+                    index.pano.progressive_filter_block<C, M>(
+                            index.codes.data(),
+                            index.cum_sums.data(),
+                            x + q0 * d,
+                            query_cum_norms.data(),
+                            block_size,
+                            batch_no,
+                            index.ntotal,
+                            sel,
+                            nullptr,
+                            use_sel,
+                            active_indices.data(),
+                            active_byteset.data(),
+                            exact_distances.data(),
+                            dot_buffer.data(),
+                            thresholds.data(),
+                            num_active.data(),
+                            local_stats);
+                });
+
+                // Push survivors into each query's heap (updates threshold).
+                for (size_t qb = 0; qb < block_size; qb++) {
+                    const uint32_t* ai = active_indices.data() + qb * bs;
+                    const float* ed = exact_distances.data() + qb * bs;
+                    for (size_t j = 0; j < num_active[qb]; j++) {
+                        ress[qb].add_result(ed[ai[j]], batch_start + ai[j]);
+                    }
+                }
+            }
+
+            for (size_t qb = 0; qb < block_size; qb++) {
+                ress[qb].end();
+            }
+            indexPanorama_stats.add(local_stats);
+        }
+    }
+}
+
 template <bool use_radius, typename C, typename BlockHandler>
 inline void flat_pano_search_core(
         const IndexFlatPanorama& index,
@@ -579,6 +695,17 @@ inline void flat_pano_search_core(
         const float* x,
         float radius,
         const SearchParameters* params) {
+    // Query blocking is implemented for top-k search only. Range search uses a
+    // RangeSearchPartialResult that finalizes once per handler and requires
+    // per-query-contiguous appends, which is incompatible with the interleaved
+    // block schedule; it falls through to the original path.
+    size_t qbs = panorama_query_block_size;
+    if (qbs > 1 && !use_radius) {
+        flat_pano_search_core_blocked<use_radius, C>(
+                index, handler, n, x, radius, params, qbs);
+        return;
+    }
+
     using SingleResultHandler = typename BlockHandler::SingleResultHandler;
 
     IDSelector* sel = params ? params->sel : nullptr;

faiss/IndexFlat.h

@@ -17,6 +17,14 @@
 
 namespace faiss {
 
+/// Query-block size for IndexFlatPanorama search (default 32). When > 1,
+/// queries are processed in blocks so each DB level-block is loaded from memory
+/// once and reused across the block (raising cache-bandwidth efficiency). 0 or
+/// 1 selects the original query-at-a-time path. Results are identical for any
+/// value; this only trades off cache behavior. Set before searching; changing
+/// it concurrently with in-flight searches is not thread-safe.
+FAISS_API extern size_t panorama_query_block_size;
+
 /** Index that stores the full vectors and performs exhaustive search */
 struct IndexFlat : IndexFlatCodes {
     explicit IndexFlat(