TPU-KNN: K Nearest Neighbor Search at Peak FLOP/s
Authors: Felix Chern
Blake Hechtman
Andy Davis
Ruiqi Guo
David Majnemer
Sanjiv Kumar
Venue: NeurIPS 2022 Accept (2022)
Keywords: TPU
K-nearest neighbor search
Approximate nearest neighbor search
roofline model
accelerator

Abstract: This paper presents a novel nearest neighbor search algorithm achieving TPU (Google Tensor Processing Unit) peak performance, outperforming state-of-the-art GPU algorithms with similar level of recall. The design of the proposed algorithm is motivated by an accurate accelerator performance model that takes into account both the  memory and instruction bottlenecks. Our algorithm comes with an analytical guarantee of recall in expectation and does not require maintaining sophisticated index data structure or tuning, making it suitable for applications with frequent updates. Our work is available in the open-source package of Jax and Tensorflow on TPU.

TPU-KNN
K Nearest Neighbor Search at Peak FLOP/s
Felix Chern Google Research rchern@goggle.comm
Blake Hechtman
Andy Davis Google Core ML
Google Core ML
blaxelhcchhunnggggglee.cco
andydavis@google.com
Ruiqi Guo Google Research guorq$google.com com
David Majnemer Google Core ML majnemer@google. com
Sanjiv Kumar Google Research janjivk@google. com
Abstract
This paper presents novel nearest neighbor search algorithm achieving TPU (Google Tensor Processing Unit) peak performance, outperforming state-of-the- art GPU algorithms with similar level of recall. The design of the proposed algorithm is motivated by an accurate accelerator performance model that takes into account both the memory and instruction bottlenecks. Our algorithm comes with an analytical guarantee of recall in expectation and does not require maintaining sophisticated index data structure tuning, making it suitable for applications with frequent updates. Our work is available in the open-source package of Jax and Tensorflow on TPU.
Introduction
The K K-nearest neighbor (K-NN) search problem has wide range of applications in machine learning and information retrieval systems, including image search (Jia et al. [021] Babenko and Lempitsky
2016), semantic textual retrieval (Liu et al. 2009 Cer et al. 2018), anomaly detection (Gu al.
2019 Omar et al. 2013). recommendation systems (arwar et al. (2002) Zhao et al., 2019). as well as
serving component for downstream tasks (Borgeaud al.) 2021 Guu al. 2020 Lindgren
ct al. 2021 Shazeer et al. 2017). Given a query, the objective of K-NN is to identify K closest
datapoints from database of finite number of data points a vector space. The main challenge of
designing a good K-NN algorithm is to compute accurate K-NN results while being computationally
efficient.
Solving the K-NN problem on accelerators has emerging interests from both the academia and the industry @Jhnsso al. 2021 Shanbhag et al. 2018 Zhao al. 2020). Many accelerators deliver hundreds of Tera Floating Point Operations Per Seconds (TFLOPS) vital to the neighbor distance computation. However, utilizing accelerators in K-NN problems is straightforward; multiple issues data locality, memory bandwidth, and multiple types of hardware parallelism need be carefully considered to achieve high utilization. In this paper extend the rooffine performance model (Williams et al. 2009) to quantify the hardware characteristies accurately. As a result, we designed K-NN algorithm to reach peak performance by the precise modeling of the accelerators. and our TPU implementation aligned with our predicted performance.
The main contributions of this work are:
Eqqual contributions.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

We extend the roofline model to address the operation throughput differences of the instruc- tions, essential to the algorithm analysis this paper.
We design approximate K-NN algorithm with recall and performance guarantees based on our proposed roofline model.
We conduct experiments verifying our TPU implementation of the algorithm accurately aligned with the performance model and achieves state-of-the-ard speed-recall trade-offs on standard nearest neighbor search benchmarks.
Preliminaries
This section covers the necessary notations to work with the nearest neighbor search problem. Given a matrix A E IR MRM X we let di denote the item at the ith row and jth column of A, and a; denote the ith row-vector of A. We use the matrix X E Nx to abbreviate a set-representation database X = 1. with N data points, where each data point x; R is vector of the matrix X in a D dimensional vector space. The set and matrix representation of database X are used interchangeably in this paper.
a
The K nearest neighbor search problem is stated as follows. Given database X RRND and query vector IR_ find the subset S C X collecting the K-closest data points to qF
S
K-argmin D(q, x);
x
where D(x, y) is distance measure such as Euclidean distance Dea De=(x,)) y) =- y the cosine (x,y) distance Deom y) == x A related problem is the maximum inner product search (MIPS),
where the goal is to find the data points that have the highest inner products with the query:
Sq
K-argmax (q. x).
(2)
xeX
MIPS is equivalent to the cosine similarity search when all data points are n-oomaalized.
3 Related work
Exhaustively searching all pair-wise distances between the query and the entire database is compute- intensive and often infeasible on many platforms. Therefore, a problem extensively discussed in the literature (Wang et al.] [2144] 2015) is to find approximate nearest neighbors (ANN) in exchange of speed. By convention, the quality of ANN is measured by
|S 0 Sg |Sq"
Recall :=
(3)
where Sq C denotes the set of data points retrieved by the search method.
Compressed domain search One class of ANN approaches is to search on lossy-compressed problem domain. These methods are composed in two steps: a) search on compressed representatiorf of the original problem to find set of candidate data points, b) compute the distances between the query and candidate data points to select the top- results. Since only a subset of data points requires the exact distance computation, the overall cost is reduced.
The two steps can be composed in arbitrary ways. Locality sensitive hashing (Andoni al., 2015
Neyshabur and Srebro 2015) applies search followed by scoring: tree-search Muja and Lowe 2014
Dasgupta and Freund 2008) applies the two steps recursively; graph-search MMakkov and Yashunin 2018) iterates between two steps until the stopping condition is met. And the inverted file (IVF)
Here data structures like tree, graph, locality sensitive hash etc.

---

method (Jegou et al. 2010 Babenko and Lempitsky 2014 Baranchuk et al. 2018 Guo et al. 2020 search on subset of data points indexed by the k-means centroids.
We see that there are two major challenges with the compressed domain search:
Fractional search has poor cache reuse rate because the candidate data points for each query rarely overlaps. We show optimizing the cache usage has a huge headroom for accelerators in Section 4w
Tweaking the speed-recall trade-off is data-dependent and non-trivial to tune. The key result Beyer al. (1999) states that the distance contrast of neighbors diminishes with increasing dimensionality (also known the curse of high dimensionality). Furthermore, the key result of Rubinstein (0018) states that sub-linear time nearest neighbor search with high recall is impossible for Euclidean, Manhattan, or Hamming distance; otherwise, contradicts the Strong Exponential Time Hypothesis (Impagliazoo and Paturi 1999).
Our work takes an opposite approach focus on machine efficiency with zero search space prun- ing. Moreover, since our method computes all the distances, is immune to the curse of high dimensionality.
Accelerators In this paper, the phrase accelerators represents a class of specialized hardware to accelerate machine learning workloads. In particular, we are interested in the novel platforms that deliver high FLOP/s for distance computation, namely Google TPU V3, V4, Nvidia GPU V100, and A100 in our analysis and evaluation.
Modern accelerators have special computation units for matrix multiplication, providing higher operation throughput over the regular coefficient-we operations. The corresponding units are tensor cores in Nvidia GPUs (Markidis al. 2018) and systolic arrays in Google TPUs (Jouppi et al.) 2017] Norrie et al. 2021). Addressing these operation throughput differences is essential to our algorithm design.
While accelerators excel in parallelism, developing an efficient K-selection algorithm on accelerators still active research area (Monroe al. 201T Shanbhag al., 2018 Johnson al. 2021 Zhao
et al. 2020). Accelerators with higher FLOP/s introduce a higher opportunity cost of computing the K-selection problem instead of the distance computation. The trend of the increasing FLOP/s in accelerators motivated us to optimize the FLOP/s usage by reducing the time required for computing K-selection.
4 Methodology
This section presents a performance model to identify non-trivial bottlenecks on multiple plat- forms and demonstrates problems, and we see that the cache inefficiency of the compressed domain methods introduces a significant on accelerators.
fundamental limits when designing algorithms for K-NN and related
We model the accelerator's runtime as executing a sequence computation kernels, where each kernel is compiled subroutine the accelerator used by the main program on the CPU. A kernel may be composed one or several high-level operators: Einsum, ReLU, ArgMax, etc., and each kernel can have different performance characteristics.
Given a sequence of kernels k;, we let W; denotes the total amount of work and P. denotes the operational speed. Our goal is to estimate the total time of a program:
Σ
W P
(4)
In the following example, we focus on the MIPS problem. Let Q E MxD and X E R denote the queries and the database, the runtime of a generic approximate- program can be modeled as
AW P
AWD
O(Auxiliary) 2
(5)
3

---

Table 1: Hardware specifications for the generalized roofline model
Name
(TFLOP/s)
3 (GB/s)
7 (TCOP/s)
GPU V100
125
900
15.7
GPU A100
312
1555
19.5
TPU V3
126
858
4.0
TPU V4
274
1144
4.3
where WD denotes the total FLOPss required for searching the entire database, and À denotes the search fraction. note that varies by algorithm and platform. Traditionally, compressed domain search methods minimize À but sacrifice cache efficiency. Our method use an alternative route to optimize P instead.
4.1 Instruction throughput-aware rooffine model
This subsection describes how we model the kernel-dependent performance P multiple platforms with small extension roofline model.
The classic roofline model (Williams et al., 2009) is a function of machine peak performance r measured in FLOP/s, machine peak memory bandwidth measured in bytes/s, and arithmetic intensity MEM expressed as the ratio of floating-pont point operations performed to data movement (FLOP/byte). The model states the performance is bounded by P min( (MMM).
We desire to model kernels that has a mixture of floating point operations accelerated by dedicated hardware well as other coefficient-wi wise operations. The coefficient-we operations are abbreviated as COPs. Almost every non matrix multiplication operations are COPs, including vectorized add, multiply, compare, conditional-move, etc. We the symbol for peak COP/s on platforms, and define the instruction throughput intensity ICOP as the ratio between the number FLOPSs and the number of COPs performed in kernel (FI.OP/COP). The attainable performance of kernel is bounded by:
π
P < min
3 x MEM
(6)
X ICOP.
The statement is self-explanatory because the inadequate resources impede the kernel throughput. Table lists the properties of selected accelerators for our analysis The roofline model is commonly used in accelerator profiling tools but not as frequently discussed in algorithm designs. The following sections show how the model prevents pitfalls due to the hardware constraints.
4.2 The memory bandwidth bound
This subsection demonstrates how to evaluate if a kernel hits the memory bandwidth wall. We associate the distance computation with three levels of BLAS (Dongarra describes vector operations on non-consecutive memory access, such as computing distances while traversing through graph. Level 2 BLAS represents scoring query with consecutively stored data points. Level 3 BLAS expresses batched qutry-database distance computation, often used in brute-force scoring.
1990]. Level BLAS
Compressed domain searches are either level or 2 BLAS due to the cache inefficiency. It has O(1) memory arithmetic intensity because the number FLOPs is proportion the bytes read. Combining 5 and 6 we have the following remark:
Remark 1. Distance computations in compressed domain searches are memory bandwidth bounded. In our model, the runtime lower bounded by: t O (AW+/3).
Readers find these numbers from the accelerators' specification sheets.

---

GPU V100
TPU V3
Skylake
GPU A100
TPU V4
350000
BLAS 1,2, IMEM = 0(1)
312 TFLOP/s
Performance P (GFLOP/S)
BLAS 3, D 128, IMEM 64
300000
274 TFLOP/s
250000
= 1555 GB/s
200000
ApproxTopk
1144.4 GB/s
150000
126 TFLOP/s
900 GB/S
125 TFLOPPS
100000
858.4 GB/s
50000
156 GB/s
T.7.P/
50
100
150
4700
4750
4800
Memory arithmetic intensity /MEM (FLOP/byte)
Figure I: Memory rooflines of accelerators and a dual-sockets Intel skylake machine as a baseline. Each colored line denotes the maximum performance platform could achieve, and each vertical line represents the memory arithmetic intensity of an algorithm. The intersections of the lines show the maximum performance of algorithm could achieve platform. We label three levels of BLAS kernels and our algorithm described in Section S
To estimate the memory arithmetic intensity for level BLAS, we continue use Q RR MxD and X E R for denoting queries and database. In many K-NN applications N and M are much greater than D. The corresponding memory arithmetic intensity is:
2MND
D
MEM 4MN + ((MN)
The largest term in the denominator of is the 4MN bytes of the query-database distances. We omit the insignificant terms and refer readers to (Golub and Van Loan 2013) Section 1.5.4) for comprchensive review memory transfers in block matrix multiplications.
Figure shows that the distance scoring kernels of different BLAS levels can easily hit the memory bandwidth wall. In order to attain high performance, we designed our algorithm to aggregate the results within the kernel to avoid writing the (M bytes into memory.
4.3
The instruction bandwidth bound
The use of COPs (non matrix multiplication instructions) introduce another slowdown. We let C denotes the number of COPs used per dot-product score in a kernel equipped with COPs and matrix multiplication instructions. There are M x N dot-product scores, so the total COPs used in kerne CMN. To prevent hitting the COPs bandwidth wall, we must satisfy:
2AND CMN ICOP
π
(8)
2D X
The number of COPs we can afford in the kernels is scarce. We take D = 128 example and substitute it into D. We can only use coefficient-wi instructions per dot- poodprrduct for TPU V4, and 16 for GPU A10O. We conclude with the following remark:
Remark 2. Exact and generic K-selection algorithm cannot be efficiently implemented with the coefficient-wise operations for the selected platforms (GPU V100, A100, TPU V3 and V4).
Because of Remark 2
develop approximate approach to achieve the peak performances.
5

---

Algorithm
Algorithm 1: PartialReduce for MIPS
Input: Q E RMx Batch queries
Input: X E RNxD Database
Input: 2 Bin size
Output: V MR Top-K K values
Output: A E NMXL Top- K indices
for i 4- to .M do
for j 1 to N do
Ji,j - (q(,,,):
[ + ShiftRight (j. W);
Unrolled does not cost COP
b Ya.j 0..
/* /a COP 1: Vectorized compare */
1y - if b then Be. else 24,1 ; t if b then j else a,1
Vectorized conditional move
/* COP 3: Vectorized conditional move */
end
end
Our algorithm consists of two kernels:
1. PartialReduce kernel computes the distances and partially aggregate the results from M x N distances to M distances with original indices.
2. ExactRescoring kernel is optional kernel that aggregates the final top-K o results. The complexity is (ML log (L)) by bitonic sort followed by a truncation.
The PartialReduce kernel is where most of the time and compute takes place. See Algorithm I for an outline the algorithm. We collect top-1 distances from the L non-overlapping bins size 2 for each query, resulting high arithmetic intensities:
MEM ~= O (min (M, ICOP
(10)
2MND 2D
(11)
CAN C
We show these arithmetic intensities can achieve high performance on real world database in section See Appendix A.3 for the detailed expansion of the algorithm and how the arithmetic intensities are derived.
5.1 Recall estimation
This section shows the PartialReduce kernel can achieve high recall with good speed. We reformulate our problem in terms of balls and bins. We have K balls representing the top- K distances that are thrown into L bins. The location of each ball is chosen independently and uniformly random. We let denotes the random variable of the number balls that do not have collisions. Following the recall definition B we have:
Z Recall >
(12)
K
which is a standard Birthday problem:
Z
L
K-1
Recall] >
(13)
K
L

---

TPU V3
TPU V4
Memory bandwidth roofline (Williams et al.)
Instruction bandwidth roofline (Ours)
300000
300000
250000
MIPS.
250000
MIPS
GFLOP/s
200000
200000
150000
150000
100000
e MIPS.
100000
g
MIPS
50000
50000
0
1000 2000
3000 4000
5000
6000
10
20
30
40
50
60
70
80
1MPM (FLOP/byte)
Icop (FLOP/COP)
Figure 2: Roofline plots for MIPS and e search benchmarks using the PartialReduce kernel. The colored lines denotes the attainable performance derived from Table 1 The figure the left shows none of the benchmark is memory bandwidth limited. The figure on the right shows that our model gives much tighter bound for e on TPU V4. See also Appendix A.5 for detailed deviation of the numbers.
Our goal is to find the minimal L. such that the expected recall is greater equals to the target recall r. Finding is simple because [3) is invertible in the natural range r <
1
K - 1 1-r
E. Recall 2 2 12 1 11(KK1)
~
(14)
The approximation in (14 follows from Appendix A.4
Since L is at the order of K, and in most applications K < N, the cost the ExactRescoring kernel is amortized out. Thus we affirm the claim that our method attains high performance with an analytical recall guarantee.
Evaluation
In this section, we show that proposed algorithm and implementation are near the hardware limit and lead to superior performance over the baselines of similar recalls. We applied our algorithm to two datasets from the public ANN benchmarks (Aumüller al. 2020). In first evaluation, compare the measured FLOP/s to the theoretical peak governed by the proposed refinement of the roofline model (6). proclaiming our implementation is reaching the hardware peak performance. In the second benchmark, we compare the end-to-end performance with competitive baselines with pre-tuned parameters. We plot each algorithm's speed-recall curve and show ours achieves the late-oof-the-ar. Finally, we measure the algorithm's scalability by varying the dataset size and number of TPUs used.
6.1 Comparison with the theoretical peak
This section shows that our refined roofline model
0 captures additional performance characteristic
over the classic roofline model, and demonstrates our kernels are having near optimal performances. We select Glove (Pennington et al. 20144 and Sif [Jegou al. 2010) datasets from the ANN benchmarks. Their corresponding distances are the cosine distance and the Euclidean distance. See the code snippets in Appendix A. and
'Released in Apache license 2.0.
RReleased in CCO public domain.

---

800.0k
700.0k
Glovel.2M
SitiM
600.0%
D D=50
500.0k
600.0k
D=128
500.0%
QPS
400.0k
400.0k
300.0k
300.04
200.04
200.0k
100.0k
100.0k
0.0
0.0
0.7
0.8
0.9
0.6
0.7
0.8
0.9
700.04
1.2M
Glovel 22M
Nytimest 3M
600.04
10-000
1.0M
D=256
500.0k
800.0k
400.0k
QPS
300.0%
200.04
400.0k
100.04
0.0
0.0
300.04
1.8.
Glovei MMM
1.6M
LM
250.0k
D=200
1. 1.1M
200.0k
1.2M
QPS
150.0k
QPS
1.0M
800.0K
100.0k
600.0k
50.0k
200.0k
0.0
0.7
0.8
0.9
0.6
0.7
0.8
0.9
IVF,Flat V100
IVF.Flat A IVE.PQ A100
Our TPU V3
IVE,PQ V100
Our TPU V4
Flat V100
Flal A
Figure 3: Recall-speed trade-off benchmarks. The x-axis is recall for k 10: up and to the right the better of the trade-off. The GPU methods (IVF-Flat, IVF-QP, and Flat) are released by Faiss (Johnson al. 2021 For each IVF* benchmark, the search fractions are (0.24%, 0.61%, 1.22%).
See Figure 2 the colored lines represent machines' performances, and the dots represent each benchmark with its measured FLOP/s. The classic roofline on the left shows that our in- cache aggregation strategy has a large memory arithmetic intensity (~4,700) exceeding the memory bandwidth ridge points T/B. However, is difficult to diagnose why the Euclidean distance search does not perform well on TPU V4 from the classic roofline plot.
Fortunately, when combined with the instruction bandwidth roofline we can tell the performance regression caused by hitting the coefficient- wiie operation throughput wall. Therefore affirms the claim that our MIPS solution reaching the peak FLOP/s, and our Euclidean distance search solution is meeting the compute bound on TPU V4 and attaining the peak FLOP/s TPU V3.
6.2 Recall-speed benchmark
To evaluate the effectiveness of the K-NN algorithm in realistic setting, we adopted the methodology of public ANN benchmarks (Aumüller et al. 2020) to compare the end-to-end performance against other methods on the following datasets: Glove (Pennington et al. 2014). Sift (Jegou et al. 2010). NYTimes (Dua and Graff 2017), and Last.fm (Bertin-Mahieux al. 2011) The typical ANN benchmarks are only performed on a single platform. However, it is ton-trivial to either port our TPU algorithm to GPU or vice versa. Alternatively, we selected the following GPUs with parity in peak performance to TPU (Table D..
2021) implementation as our baseline. Faiss provides three algorithms: Flat, IVF-Flatt and IVF-PQ. The Flat algorithm performs a brute-force search, and the IVF-Flat and IVF-PQ algorithms corresponds to the inverted file method with and without the product quantization (Tegou et al. 2010) Johnson et al. 202T). We use the repository's suggested inverted file size (16384) the IVF methods.
We select the Faiss GPU JJohsson et al.

---

Yandex Deep D 96
3.0M
Faiss IVF.Flat A100
Ours TPU V4 1 chip
2.5M
Ours TPU V4 2 chips
Ous TPU V4 clips
2.0M
QPS
1.5M
0.99
1.0M
500.0k
70397
0.00
1.98.......
8.00
11.55
0.99
00:02
11:00
0.0
1M
6MM
7M
10M
Datahase size N
Figure 4: Scalability benchmark. The labeled numbers are the measured recalls. The multi-TPU implementation is listed in Appendix A.6
Figure shows our performance significantly outperforms competing methods in the high recall regions. We highlight that our method has consistent recall-speed trade-off over different datasets, because our recall only rely on the order statistics instead the information encoded in the com- pression domain search methods, which may vary by the datasets. Since method scores all the pair-wise distances, our method is immune from the curse of high dimensionality.
6.3 Scalability benchmark
In the final benchmark, we examine the scalability of the algorithm from three aspects. First, we verify if the measured performance is inverse proportional to the database size. Second, we compare the scaling characteristics to the fastest GPU implementation. Last but not least, are interested in knowing our algorithm can horizontally scale by the number of TPUs.
We conduct our evaluation on TPU V4 and Nvidia GPU A100, which have similar peak performance and memory bandwidth. We sample the Yandex Deep datase (Bbeenko and Lempitsky/ 2016) into ten different scales and measure the QPS of each approach with similar recall. Figure verifies all measurements align with the ideal scalability model: QPS x #chips/N. Our method remains top performance on all database sizes and linearly scales with the number of TPU chips.
7 Discussion and future work
In Section 5 we benchmark our method against others on platforms with similar performances. Some questions might arise: "Is the performance gain algorithmic optimization due platform efficiency?" "Can we achieve the same performance gain on GPU?" "The existence of efficient Prational-scarch accelerators?" We address these questions in this section.
7.1 Platform discussions
first discuss the modeling perspective performance differences between platforms. Section show that the memory bandwidth and instruction throughput bound applies to both GPU and TPU. For instance, it follows that to attain peak performance on every hardware platform, having number of instructions used for collecting (approximate) top-k elements within 2 D/x per distance computation a necessary condition.
Although our Algorithm iss ientfform-indeeenddeen achieving the hardware peak performance requires many level implementation details at the machine level, including cache management, preventing cross-core memory synchronization, in-register accumulation, and instruction scheduling. Typical high-performange libraries such as MKL, cuBLAS, and Google TPU compiler use platform-specifee assembly to take full control of the stated requirements.
'Released in CC BY 4.0.

---

Nevertheless, we cannot use the high-level interface of these libraries, because Algorithm | onlyy performs well when it is integrated into the inner loop of distance computation Moreover, these libraries are all close-sourced, thus increases the difficulty on the implementation.
Fortunately, have the access to TPU compiler internals, and have integrated Algorithm T into the compiler to generate the desired assembly code to solidify our analysis. Thus we leave implementations of other platforms to future works.
7.2 Algorithm discussions
The roofline complexity of the fractional search is identical to BLASS-2 (matrix-vector multiplication) which is memory bandwidth bound. When the cycles spend on data transfer mutually exclusive to our method, it introduces an enormous opportunity cost. Nevertheless, we see an opportunity in heterogeneous architecture because fractional search the host is not mutually exclusive to applying our method to accelerators.
A motivating example is the multi-billion nearest neighbor search, where fitting the dataset into device memory is possible (through device sharding, which TensorFlow and Jax have native support) but not economical. Since brute-force distance computations structures when performing the fractional search, we may replace the brute-force portion with TPU in conduction with the remaining search off-device. We note that heterogeneous architectures with off-device storage such as host-RAM or even SSI) (Jayaram Subramanya et al., 2019. Ren et al. 2020
often involved in the auxiliary data
Chen al. 2021
great starting points for future research.
8
Conclusion
Accelerator-baser machine learning has become the mainstream in academics and industries. How- ever, the performance characteristies of accelerators are counter-intuitive and difficult to program. In this paper, we propose roofline-based complexity analysis framework to discuss the optimality of the algorithms without low-level optimization details: unrolling factors, batch window sizes, vectorization, and systolic array scheduling, which are Wlttorrm-depender and lengthy to read. We demonstrated several examples of inferring hardware performance limits by simply addressing the kernel's total FLOPs, byte transferred, and the number of coefficient-we wise instructions used. Our refined model foreshadowed non-trivial performance regression caused by the coefficient-wis in- structions bandwidth. We took it into account to design new algorithm for K-NN and achieved peak performance on TPU. Finally, our experiments showed that our method outperformed state-of-the-an baselines platforms with similar performance characteristics, which known to be hard to beat.
Acknowledgments and Disclosure of Funding
We would like to thank the XLA team for the continuous effort on developing the state-of-the-art compiler and the full support on enabling our new op: approx_max __x____a__kkk_ We are also grateful the Google ScaNN team for the joint effort on bridging the impactful K-NN problem into the accelerator ecosystem. Last but not least, we thank to Peter Hawkins, Edward Schwartz, and Mani Varadarajan for code reviews in Jax and Tensorflow, and Erik Lindgren for the proof reading of this paper.
This work was performed and funded by Google.
References
Andoni, A., Indyk. P., Laarhoven, T., Razenshteyn, I., and Schmidt. L. (2015). Practical and optimal Ish for angular distance. Advances in neural information processing systems, 28.
Aumüller, M., Bernhardsson, E., and Faithfull, A. (2020). Ann-benchmarks: A benchmarking tool for approximate nearest neighbor algorithms. Information Systems, 87:101374.
Babenko, A. and Lempitsky, V. (2014). The inverted multi-index. IEEE transactions on pattern
analysis and machine intelligence, 37(6): (166):1247-11 1267-1260.
See Appendix A.3.3 for the performance analysis and comparisons.
10

---

Babenko, A. and Lempitsky. V. (2016). Efficient indexing of billion-scale datasets of deep descriptors. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 2055-2063.
Baranchuk, D., Babenko, A., and Malkov, Y. (2018). Revisiting the inverted indices for billion-scale approximate nearest neighbors. In Proceedings of the European Conference on Computer Vision (ECCV), pages 202-216.
BrrtinMMahieux, T., Ellis, D. P., Whitman, B., and Lamere, P. 1). The million song dataset. In Proceedings of the 12th International Conference on Music Information Retrieval (ISMIR 2011).
Beyer, K., Goldstein, J., Ramakrishnan, R., and Shaft, U. (1999). When is "nearest neighbor" meaningful? In International conference on database theory, pages 217-235. Springer.
Borgeaud, S., Mensch, A., Hoffmann, J., Cai, T., Rutherford, E.,, Millican, K., Driessche, G. d., Lespiau, J.-B., Damoc. B., Clark, A., et al. (2021). Improving language models by retrieving from trillions of tokens. arXiv preprint arXiv:2112.0m 044266.
Cer, D., Yang, Y., Kong, S.-y.. Hua, N., Limtiaco, N., John, R. S., Constant, N., Guajardo-Cespeder M., Yuan, S., Tar, C., al. (2018). Universal sentence encoder. arXiv preprint arXiv: 11155
Chen, Q., Zhao, B., Wang, H., Li, M., Liu, C., Li, Z., Yang, M., and Wang. J. (2021). Spann: Highly- efficient billion-scale approximate nearest neighborhood search. Advances Neural Information Processing Systems, 34:5 199-5212.
Dasgupta, S. and Freund, Y. (2008). Random projection trees and low dimensional manifolds. In Proceedings of the fortieth annual ACM symposium Theory of computing, pages 537-546.
Dongarra, J., Du Croz, J., Hammarling, S., and Duff, S. (1990). A set of level basic linear algebra subprograms. ACM Transactions on Mathematical Software (TOMS), 16(1):1-17.
Dua, D. and Graff, C. (2017). UCI machine learning repository.
Golub, G. H. and Van Loan, C. F. (2013). Matrix computations. JHU press.
Gu, X., Akoglu, L., and Rinaldo, A. (2019). Statistical analysis of nearest neighbor methods for anomaly detection. Advances Neural Information Processing Systems, 32.
Guo, R., Sun, P., Lindgren, E., Geng, Q., Simcha, D., Chern, F., and Kumar, S. (2020). Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, pages 3887-3896. PMLR.
Guu, K., Lee, K., Tung, Z., Pasupat, P., and Chang, M.-W. (2020). Realm: Retrieval-auugmented language model pre-training. arXiv preprint arXiv:2002.08909
Impagliazzo, R. and Paturi, R. (1999). The complexity of k-sat. In Proceedings. Fourteenth Annual IEEE Conference Computational Complexity (Formerly: Structure in Complexity Theory Conference)(Cat. No. 99CB36317), pages 237 237-237. IEEE Computer Society.
Jayaram Subramanya, S., Devvrit, F., Simhadri, H. V., Krishnawamy, R., and Kadekodi, R. (2019). Diskann: Fast accurate billion-point nearest neighbor search on single node. Advances in Neural Information Processing Systemx, 32.
Jegou, H., Douze, M., and Schmid, C. (2010). Product quantization for nearest neighbor search. IEEE transactions on pattern analysis and machine intelligence, 33(1):117-128.
Jia, C., Yang. Y., Xia, Y., Chen, Y.-T., Parekh, Z., Pham, H., Le, Q., Sung, Y.-.., Li, Z., and Duerig, T. (2021). Scaling up visual and languae representation learning with noisy text supervision. In International Conference on Machine Learning, pages 4904 PMLR.
Johnson, J., Douze, M., and Jégou, H. (2021). Billion-scale similarity search with gpus. IEEE Transactions Big Data, 733:5535-547.
11

---

Jouppi, N. P., Young, C., Patil, N., Patterson, D., Agrawal, G., Bajwa, R., Bates, S., Bhatia, S., Boden, N., Borchers, A., al. (2017). In-datacenter performance analysis of a tensor processing unit. In Proceediigss of the 44th annual international symposium on computer architecture, pages 1-12.
Liddgenn, E.., Reddi, S., Guo, R., and Kumar, S. (2021). Efficient training of retrieval models using negative cache. Advances in Neural Information Processing Systems, 34.
Liu, T.-Y. et al. (2009), Learning to rank for information retrieval. Foundations and Trends@ in Information Retrieval, 3(3):225-331.
Malkov, Y. and Yashunin, A. (2018). Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions pattern analysis and machine intelligence, 4244))824-836.
Markidis, S., Der Chien, S. W., Laure, E.,, Peng. I. B., and Vetter, S. (2018). Nvidia tensor core programmability, performance & precision. In 2018 IEEE international parallel and distributed processing symposium workshops (IPDPSW), pages 522-531, IEEE.
Monroe, I.., Wendelherger, J., and Michalak, S. (2011). Randomized selection on the gpu. In Proceedings of the ACM SIGGRAPH Symposium on High Performance Graphics, pages 89-98.
Muja, M. and Lowe, D. G. (2014). Scalable nearest neighbor algorithms for high dimensional data. IEEE transactions pattern analysis and machine intelligence, 36( (22111)2227-24444
Neyshabur, B. and Srebro, N. (2015), On symmetric and asymmetric Ishs for inner product search. In International Conference on Machine Learning pages 1926-1934. PMLR.
Norrie, T., Patil, N., Yoon, D. H., Kurian, G., Li, S., Laudon, J., Young, C., Jouppi, N., and Patterson, D. (2021). The design process for google's training chips: Tpuv2 and tpuv3. IEEE Micro, 41(2):56-63.
Omar, S., Ngadi, A., and Jebur, H. H. (2013). Machine learning techniques for anomaly detection: overview. International Journal of Computer Applications, 79(2).
Pennington, J., Socher, R., and Manning, C. D. (2014). Glove: Global vectors for word representation. In Empirical Methods in Natural Language Processing (EMNLP), pages 1532-1543.
Ren, J., Zhang, M., and Li, D. (2020). Hm-ann: Efficient billion-point nearest neighbor search on heterogeneous memory. Advances in Neural Information Processing Systems, 33:10672-10684,
Rubinstein, A. (2018). Hardness of approximate nearest neighbor search. In Proceedings of the 50th annual ACM SIGACT symposium on theory of computing, pages 1260-1268.
Sarwar, B. M., Karypis, G., Konstan, J., and Riedl, J. (2002). Recommender systems for large- scale e-commerce: Scalable neighborhood formation using clustering. Proceedings of the fifth international conference on computer and information technology, volume 1, pages 291-324. Citeseer.
Shanbhag, A., Pirk, H., and Madden, S. (2018). Efficient top-k query processing on massively parallel hardware. In Proceedings of the 2018 International Conference on Management of Data, pages 1557-1570.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. (2017). Outrageously large neural networks: The sparsely- gated mixture-of-exxe -f-eeeertss layer. arXiv preprint arXiv:1701.06538.
Wang, J., Liu, W., Kumar, S., and Chang, S.-F. (2015). Learning hash for indexing big data survey. Proceedings of the IEEE, 104(1)334-7..
Wang, J., Shen, H. T., Song, J., and Ji, J. (2014). Hashing for similarity search: A survey. arXiv preprint 4r0iv:1408.25 2927..
Williams, S., Waterman, A., and Patterson, D. (2009). Roofline: an insightful visual performance model for multicore architectures. Communications of the ACM, 52(4):65-76.
12

---

Zhao, W., Tan. S., and Li, P. (2020). Song: Approximate nearest neighbor search on gpu. In 2020 IEEE 36th International Conference on Data Engineering (ICDE), pages 1033-1044.
Zhao, Z., Hong, L., Wei, L., Chen, J., Nath, A., Andrews, S., Kumthekar, A., Sathiamoorthy, M., Yi, X., and Chi, E. (2019). Recommending what video to watch next: a multitask ranking system. In Proceedings of the 13th ACM Conference on Recommender Systems, pages 43-51.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section for how we model
the hardware limitations and Section 6 for real world evaluations.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Section we formulate the problems in terms of the classic balls into bins.
(b) Did you include complete proofs of all theoretical results? [Yes] See Section 5. and
Appendix A.4
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] See Appendix
and
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Table Madd
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] See Section 6.1
(b) Did you mention the license of the assets? [Yes] See the footnotes in Section 6.1
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
13


Reviews:
[{"review_id": "wAGfSpbpEN8", "reviewer": "Reviewer_Vrzf", "summary": "This paper presents a new algorithm implementation of K nearest neighbor search that is able to increase the arithmetic intensity. The distances between all queries and entries in the database are calculated in L3 BLAS fashion. Then only the top-1 entry is kept in each bin. The bin size can be adjusted based on the requirement of the recall. The roofline analysis shows a better peak throughput while the end-to-end evaluation demonstrates a better throughput-recall trade-off than the previous solutions.", "questions": "1. Is the proposed algorithm specialized for TPU? If it is, what special architectural support from TPU does it utilize?\n\n2. How do you define **ONE** COP? For example, are vectorized comparisons between two 256 arrays and the comparison between two 16 arrays each counted as one COP? If not, how should I interpret the data presented in Table 1?\n\n3. Could you provide a direct comparison of the proposed algorithm on the same hardware platform?", "limitations": "The authors discussed the limitation of the current implementation.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n\n1. This paper presents a detailed analytical model of the kernel implementation of the BLAS-based operation. The authors conclude the impact of COPs and introduce the algorithm accordingly.\n\n2. The proposed solution is simple yet effective. The adjustment of the bin size provides a simple method to balance recall and throughput.\n\n## Weaknesses\n\n1. **Unclear hardware motivation.** The paper is entitled 'TPU-KNN'. However, throughout this paper, I'm not able to find any information indicating that the proposed algorithm requires any special architectural support from TPU rather than other accelerators. It seems that the same analysis could also apply to GPU or other accelerators. \n\n2. **Vague definition of COPs.** I cannot understand the usage of the concept of coefficient-wise operations (COPs). In Table 1, comparing with the datasheet of A100/V100, I can understand COPs as a FLOP in the vector cores (or SMs in NV's terminology) while the FLOP means a half-precision FLOP in the tensor core. However, in the later description like Algorithm 1, one COP seems to be a general non-matrix operation without considering the problem size. In this case, one COP, for example, vectorized comparison, could include multiple FLOPs and seems to be unmatched by what is presented in Table 1.\n\n3. **Incomplete Evaluation.** Based on what I've mentioned in 1, the evaluation part lacks a fair baseline. I can understand the advantage of recall of the proposed algorithm. However, it is hard to understand whether the performance gain over GPU comes from the algorithm itself or the higher efficiency of TPU. To isolate this factor, the authors should either provide the result of the proposed algorithm implemented on GPUs or the baseline algorithms implemented on TPU.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "oJH0rqMMocg", "reviewer": "Reviewer_sFRW", "summary": "The paper presents a new NN-search algorithm, that reaches the peak performance on TPUs. The paper is based on observations of hardware architectural properties and the so called roofline performance model. The algorithm is implemented for TPUs. The evaluation is done on two TPU versions using two KNN datasets, and compared to several GPU algorithms. \n", "questions": "* The algorithm descriptions (Algorithm 1 and Algorithm 2 (suppl. mtrl)) looks relatively clear and straight-forward to implement on other platforms. Can you elaborate a bit on why it would be such a substantial effort to do?\n", "limitations": "I think the authors have adequately addressed the limitations of their work.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n+ Nice connection between theoretical observations and practical results\n+ Good performance\n+ Can have significant practical impact\n\nWeaknesses\n- Limited evaluation, only evaluated on two versions on TPUs, and two datasets\n- Would have been interesting to see how general the algorithm is, i.e., would it reach the same performance limits in other hardware platforms also?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "P952r1Cx9iJ", "reviewer": "Reviewer_S1Tw", "summary": "This paper presents an ANN algorithm on TPU and analyzes the memory and instruction bandwidth of ANN algorithms.\n", "questions": "1. Could we have other methods on GPUs/TPUs, e.g., hashing and graph-based methods besides the FAISS baseline, on more measures, e.g. cosine?\n3. Could we show that the proposed algorithm can achieve high recalls for most ANN datasets?", "limitations": "I suggest the authors to discuss the limitation of the algorithm, e.g., how to adapt the problem to other common similarity measures.", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strong Points\n----\n1. The problem is well-motivated. \n2. The related work is comprehensively studied and discussed. And the entire story is easy to follow for readers. \n3. The methodology to study the algorithm from the hardware bottlenecks is interesting.\n\nWeak Points\n----\n1. I suggest the authors to include a brief discussion of TPU, e.g., what it can be done and what it can not be done (efficiently).\n2. The experimental evaluation is a bit problematic. How about other methods on GPUs/TPUs, e.g., hashing and graph-based methods besides the FAISS baseline?\n3. In Figure 3, the highest recall shown for Glove is 0.9. Is there a recall limitation for the proposed algorithm? \n4. The technical contribution of the bi-stage partial reduction and scoring is limited.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "D2lV8NiA7LF", "reviewer": "Reviewer_xsza", "summary": "This paper studies using TPU for approximate nearest neighbor search. With careful analysis, it finds that search performance is bound by memory access and coefficient-wise operation instead of distance computation. Thus, it proposes to conduct partial reduce to reduce memory access such that search can reach peak FLOPs. Experiment results show that the proposed solution outperforms existing ones in both QPS and recall.", "questions": "See weakness.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n\n1.\tThe extended roofline analysis is interesting and shows the bottleneck of nearest neighbor search with TPU.\n\n2.\tThe proposed partial reduce scheme reduces the amount of memory access and comes with theoretical analysis.\n\n3.\tThe experiment results show that the proposed solution has good performance.\n\n\nWeakness\n\n1.\tIt would improve the paper if the authors can give a brief introduction about the characteristics of GPU and TPU. Are TPUs widely available as a commodity hardware?\n\n2.\tThe proposed solution may also work for GPUs as GPUs are also limited by memory access from Table 1. Moreover, as shown in Figure 2, the proposed solution is still bound by coefficient-wise operations for Euclidean distance search. Any thoughts on reducing the number of coefficient-wise operations?\n\n3.\tThe experiment study is far from the NeurIPS standard. Does the method still outperform the partial search methods (e.g., IVF) when the dataset is large (e.g., SIFT10M or 100M)? How does the proposed method perform when changing the window size W?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept