Speed up IndexFlatPanorama search with query blocking

benchs/bench_flat_l2_panorama.py

@@ -18,32 +18,65 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--dataset", default="gist1m", choices=["sift1m", "gist1m"])
+parser.add_argument(
+    "--nq",
+    default="1000",
+    help="comma-separated query counts to sweep "
+    "(query blocking needs a real batch to help)",
+)
+parser.add_argument(
+    "--threads",
+    default="1",
+    help="comma-separated search thread counts to sweep "
+    "(index build always uses all cores)",
+)
+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; the first value is "
+    "the reference for the vs-Pano speedup in the summary",
+)
+parser.add_argument(
+    "--repeat",
+    type=int,
+    default=1,
+    help="timed repetitions per configuration, reporting the fastest; "
+    "raise this for small --nq where a single run is noisy",
+)
 args = parser.parse_args()
 
+query_block_sizes = [int(v) for v in args.query_block_size.split(",")]
+nq_values = [int(v) for v in args.nq.split(",")]
+thread_values = [int(v) for v in args.threads.split(",")]
+
 if args.dataset == "sift1m":
     ds = DatasetSIFT1M()
 else:
     ds = DatasetGIST1M()
 
-nq = 10
-xq = ds.get_queries()[:nq]
+max_nq = max(nq_values)
+xq_all = ds.get_queries()[:max_nq]
 xb = ds.get_database()
-gt = ds.get_groundtruth()[:nq]
+gt_all = ds.get_groundtruth()[:max_nq]
 
 xt = ds.get_train()
 
 nb, d = xb.shape
 nt, d = xt.shape
 
 k = 10
-gt = gt[:, :k]
+gt_all = gt_all[:, :k]
 
 
-def eval_qps(index):
+def eval_qps(index, xq, gt):
+    nq = len(xq)
     faiss.cvar.indexPanorama_stats.reset()
-    t0 = time.time()
-    _, I = index.search(xq, k=k)
-    t = time.time() - t0
+    t = np.inf
+    for _ in range(args.repeat):
+        t0 = time.time()
+        _, I = index.search(xq, k=k)
+        t = min(t, time.time() - t0)
     speed = t * 1000 / nq  # ms/query
     qps = 1000 / speed
 
@@ -63,45 +96,82 @@ def build_index(name):
     faiss.omp_set_num_threads(mp.cpu_count())
     index.train(xt)
     index.add(xb)
-
-    faiss.omp_set_num_threads(1)
     return index
 
 
 nlevels = 16 if args.dataset == "gist1m" else 8
 
-plt.figure(figsize=(8, 6), dpi=80)
-
-names = [
-    "Flat",
-    f"PCA{d},FlatL2Panorama{nlevels}",
-]
-
-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})")
-    qps_values.append(qps)
-
-x = np.arange(len(names))
-plt.bar(x, qps_values, color=["#1f77b4", "#ff7f0e"])
-speedup = qps_values[1] / qps_values[0]
-ax = plt.gca()
-ax.text(
-    x[1],
-    qps_values[1] * 1.01,
-    f"{speedup:.2f}x",
-    ha="center",
-    va="bottom",
-)
-plt.xticks(x, labels, rotation=0)
-plt.ylabel("QPS")
-dataset_label = args.dataset.upper()
-plt.title(f"Flat Indexes on {dataset_label}")
-
-plt.tight_layout()
-plt.savefig(f"bench_flat_l2_panorama_{args.dataset}.png", bbox_inches="tight")
+pano_name = f"PCA{d},FlatL2Panorama{nlevels}_{batch_size}"
+
+# Both indexes are built once; only nq, the thread count, and the global
+# query-block toggle change between runs, so results must match.
+print("======building Flat")
+flat_index = build_index("Flat")
+print(f"======building {pano_name}")
+pano_index = build_index(pano_name)
+
+
+def qbs_tag(qbs):
+    return "baseline" if qbs <= 1 else f"block={qbs}"
+
+
+rows = []  # one entry per (nq, threads): flat QPS + per-block-size results
+for nq in nq_values:
+    xq = xq_all[:nq]
+    gt = gt_all[:nq]
+    for nthr in thread_values:
+        faiss.omp_set_num_threads(nthr)
+        print(f"====== nq={nq} threads={nthr}")
+        print("---Flat")
+        flat_recall, flat_qps = eval_qps(flat_index, xq, gt)
+        per_qbs = []
+        for qbs in query_block_sizes:
+            faiss.cvar.panorama_query_block_size = qbs
+            print(f"---Pano {qbs_tag(qbs)}")
+            recall, qps = eval_qps(pano_index, xq, gt)
+            per_qbs.append((qbs, recall, qps))
+        faiss.cvar.panorama_query_block_size = 0
+        rows.append((nq, nthr, flat_recall, flat_qps, per_qbs))
+
+# Report speedups: Panorama-vs-Flat, and each block size vs the first swept
+# block size at the same nq/threads.
+ref_tag = qbs_tag(query_block_sizes[0])
+print("\n=== summary ===")
+for nq, nthr, _, flat_qps, per_qbs in rows:
+    ref_qps = per_qbs[0][2]
+    print(f"nq={nq} threads={nthr}: Flat {flat_qps:.1f} QPS")
+    for qbs, recall, qps in per_qbs:
+        vs_flat = qps / flat_qps
+        vs_ref = qps / ref_qps
+        print(
+            f"  Pano {qbs_tag(qbs):>10}: {qps:10.1f} QPS  "
+            f"({vs_flat:.2f}x vs Flat, {vs_ref:.2f}x vs Pano {ref_tag})"
+        )
+
+# The bar chart only makes sense for a single nq/threads combination; sweeps
+# rely on the summary table above.
+if len(rows) == 1:
+    nq, nthr, flat_recall, flat_qps, per_qbs = rows[0]
+    labels = [f"Flat\n(r@{flat_recall:.3f})"]
+    qps_values = [flat_qps]
+    for qbs, recall, qps in per_qbs:
+        labels.append(f"Pano\n{qbs_tag(qbs)}\n(r@{recall:.3f})")
+        qps_values.append(qps)
+
+    plt.figure(figsize=(8, 6), dpi=80)
+    x = np.arange(len(qps_values))
+    colors = ["#1f77b4"] + ["#ff7f0e"] * len(per_qbs)
+    plt.bar(x, qps_values, color=colors)
+    ax = plt.gca()
+    # Annotate each Panorama bar with its speedup over plain Flat.
+    for xi, qps in zip(x[1:], qps_values[1:]):
+        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()
+    plt.title(f"Flat Indexes on {dataset_label}")
+
+    plt.tight_layout()
+    plt.savefig(f"bench_flat_l2_panorama_{args.dataset}.png", bbox_inches="tight")

faiss/IndexFlat.cpp

@@ -570,6 +570,130 @@ 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;
+
+    // qbs is a maximum: when queries are scarce, shrink the effective block
+    // size so the thread count matches the unblocked path
+    // (min(n, max_threads)) rather than dropping to one thread per full
+    // block.
+    int max_nt = omp_get_max_threads();
+    size_t qbs_eff = std::min(qbs, (size_t(n) + max_nt - 1) / max_nt);
+    size_t n_blocks = (size_t(n) + qbs_eff - 1) / qbs_eff;
+    // Cap threads at n_blocks: every spawned thread allocates the scratch
+    // buffers below, even if it gets no loop iterations.
+    [[maybe_unused]] int nt = std::min(int(n_blocks), max_nt);
+
+#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_eff);
+        for (size_t qb = 0; qb < qbs_eff; qb++) {
+            ress.emplace_back(handler);
+        }
+
+        std::vector<float> query_cum_norms(qbs_eff * n_lp1);
+        std::vector<uint32_t> active_indices(qbs_eff * bs);
+        std::vector<uint8_t> active_byteset(qbs_eff * bs);
+        std::vector<uint8_t> first_level_full(qbs_eff);
+        std::vector<float> exact_distances(qbs_eff * bs);
+        std::vector<float> dot_buffer(bs);
+        std::vector<float> thresholds(qbs_eff);
+        std::vector<size_t> num_active(qbs_eff);
+
+#pragma omp for
+        for (int64_t blk = 0; blk < int64_t(n_blocks); blk++) {
+            size_t q0 = size_t(blk) * qbs_eff;
+            size_t block_size = std::min(qbs_eff, 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(),
+                            first_level_full.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,
@@ -578,6 +702,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/impl/Panorama.cpp

@@ -15,6 +15,8 @@
 
 namespace faiss {
 
+size_t panorama_query_block_size = 32;
+
 namespace {
 
 /// Helper function to compute cumulative sums by iterating backwards through