From ce94df4ea8ee1c1d410a09fdd1ce85def2e90bdf Mon Sep 17 00:00:00 2001
From: Alexandr Guzhva <aguzhva@meta.com>
Date: Mon, 17 Oct 2022 15:22:46 -0700
Subject: [PATCH] Speedup IndexRowwiseMinMax::sa_decode() (#2536)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2536

Allocated temporary buffers are of the correct sizes.

Reviewed By: mdouze

Differential Revision: D40439826

fbshipit-source-id: 97087953ce9c1c98b4ab38cab2223c7191ea7025
---
 faiss/IndexRowwiseMinMax.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/faiss/IndexRowwiseMinMax.cpp b/faiss/IndexRowwiseMinMax.cpp
index da42d8e73..de6925c39 100644
--- a/faiss/IndexRowwiseMinMax.cpp
+++ b/faiss/IndexRowwiseMinMax.cpp
@@ -154,8 +154,10 @@ void sa_decode_impl(
     const size_t new_code_size = index->sa_code_size();
 
     // allocate tmp buffers
-    std::vector<uint8_t> tmp(chunk_size * old_code_size);
-    std::vector<StorageMinMaxFP16> minmax(chunk_size);
+    std::vector<uint8_t> tmp(
+            (chunk_size < n_input ? chunk_size : n_input) * old_code_size);
+    std::vector<StorageMinMaxFP16> minmax(
+            (chunk_size < n_input ? chunk_size : n_input));
 
     // all the elements to process
     size_t n_left = n_input;