Original: TPU-KNN K Nearest Neighbor Search at Peak FLOP/s
Revised: TPU-KNN K Nearest Neighbor Search at Peak FLOP/s

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: This paper presents a novel nearest neighbor search algorithm achieving TPU(Google Tensor Processing Unit) peak performance, outperforming state-of-theart GPU algorithms with similar level of recall.
Revised: This paper presents a novel nearest neighbor search algorithm achieving TPU (Google Tensor Processing Unit) peak performance, outperforming state-of-theart GPU algorithms with similar level of recall.

Original: The design of the proposedalgorithm is motivated by an accurate accelerator performance model that takes intoaccount both the memory and instruction bottlenecks.
Revised: The design of the proposed algorithm is motivated by an accurate accelerator performance model that takes into account both the memory and instruction bottlenecks.

Original: Our algorithm comes withan analytical guarantee of recall in expectation and does not require maintainingsophisticated index data structure or tuning, making it suitable for applicationswith frequent updates.
Revised: 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.

Original: Our work is available in the open-source package of Jax and
Revised: Our work is available in the open-source package of Jax and

Original: Tensorﬂow on TPU.
Revised: Tensorﬂow on TPU

Original:  Introduction
Revised:  Introduction

Original: The K -nearest neighbor ( K -NN) search problem has a wide range of applications in machine learningand information retrieval systems, including image search (Jia et al., 2021; Babenko and Lempitsky,2016), semantic textual retrieval (Liu et al., 2009; Cer et al., 2018), anomaly detection (Gu et al. ,2019; Omar et al., 2013), recommendation systems (Sarwar et al., 2002;
Revised: The K -nearest neighbor ( K -NN) search problem has a wide range of applications in machine learning and information retrieval systems, including image search (Jia et al., 2021; Babenko and Lempitsky, 2016), semantic textual retrieval (Liu et al., 2009; Cer et al., 2018), anomaly detection (Gu et al., 2019; Omar et al., 2013), recommendation systems (Sarwar et al., 2002;

Original: Zhao et al., 2019), as well asserving as a component for a downstream tasks (Borgeaud et al., 2021; Guu et al., 2020; Lindgrenet al., 2021; Shazeer et al., 2017).
Revised: Zhao et al., 2019), as well as serving as a component for a downstream tasks (Borgeaud et al., 2021; Guu et al., 2020; Lindgren et al., 2021; Shazeer et al., 2017).

Original: Given a query, the objective of K -NN is to identify K closestdatapoints from a database of ﬁnite number of data points in a vector space.
Revised: Given a query, the objective of K -NN is to identify K closest datapoints from a database of ﬁnite number of data points in a vector space.

Original: The main challenge ofdesigning a good K -NN algorithm is to compute accurate K -NN results while being computationallyefﬁcient.
Revised: The main challenge of designing a good K -NN algorithm is to compute accurate K -NN results while being computationally efﬁcient.

Original: Solving the K -NN problem on accelerators has emerging interests from both the academia and theindustry (Johnson et al., 2021;
Revised: Solving the K -NN problem on accelerators has emerging interests from both the academia and the industry (Johnson et al., 2021;

Original: Shanbhag et al., 2018; Zhao et al., 2020).
Revised: Shanbhag et al., 2018; Zhao et al., 2020).

Original: Many accelerators candeliver hundreds of Tera Floating Point Operations Per Seconds (TFLOPS) vital to the neighbordistance computation.
Revised: Many accelerators can deliver hundreds of Tera Floating Point Operations Per Seconds (TFLOPS) vital to the neighbor distance computation.

Original: However, utilizing accelerators in K -NN problems is not straightforward;multiple issues in data locality, memory bandwidth, and multiple types of hardware parallelism needto be carefully considered to achieve high utilization.
Revised: However, utilizing accelerators in K -NN problems is not straightforward; multiple issues in data locality, memory bandwidth, and multiple types of hardware parallelism need to be carefully considered to achieve high utilization.

Original: In this paper we extend the rooﬂine performancemodel (Williams et al., 2009) to quantify the hardware characteristics accurately.
Revised: In this paper we extend the rooﬂine performance model (Williams et al., 2009) to quantify the hardware characteristics accurately.

Original: As a result, wedesigned a K -NN algorithm to reach peak performance by the precise modeling of the accelerators,and our TPU implementation aligned with our predicted performance.
Revised: As a result, we designed a K -NN algorithm to reach peak performance by the precise modeling of the accelerators, and our TPU implementation aligned with our predicted performance.

Original: The main contributions of this work are: 
Revised: The main contributions of this work are: ∗ Equal contributions.

Original: • We extend the rooﬂine model to address the operation throughput differences of the instruc- tions, essential to the algorithm analysis in this paper.
Revised: • We extend the rooﬂine model to address the operation throughput differences of the instructions, essential to the algorithm analysis in this paper.

Original: • We design an approximate K -NN algorithm with recall and performance guarantees basedon our proposed rooﬂine model.
Revised: • We design an approximate K -NN algorithm with recall and performance guarantees based on our proposed rooﬂine model.

Original: • We conduct experiments verifying our TPU implementation of the algorithm accurately aligned with the performance model and achieves state-of-the-art speed-recall trade-offs onstandard nearest neighbor search benchmarks.
Revised: • We conduct experiments verifying our TPU implementation of the algorithm accurately aligned with the performance model and achieves state-of-the-art speed-recall trade-offs on standard nearest neighbor search benchmarks.

Original:  Preliminaries
Revised:  Preliminaries

Original: This section covers the necessary notations to work with the nearest neighbor search problem.
Revised: This section covers the necessary notations to work with the nearest neighbor search problem.

Original: Givena matrix A 2 R M ⇥ N , we let a i,j denote the item at the i th row and j th column of A , and a i denotethe i th row-vector of A .
Revised: Given a matrix A ∈ R M × N , we let a i,j denote the item at the i th row and j th column of A , and a i denote the i th row-vector of A .

Original: We use the matrix X 2 R N ⇥ D to abbreviate a set-representation of adatabase X = { x i } i =1 , 2 ,...,N with N data points, where each data point x i 2 R D is a row vector of41 the matrix X in a D dimensional vector space.
Revised: We use the matrix X ∈ R N × D to abbreviate a set-representation of a database X = { x i } i =1 , 2 ,...,N with N data points, where each data point x i ∈ R D is a row vector of the matrix X in a D dimensional vector space.

Original: The set and matrix representation of database X areused interchangeably in this paper.
Revised: The set and matrix representation of database X are used interchangeably in this paper.

Original: The K nearest neighbor search problem is stated as follows.
Revised: The K nearest neighbor search problem is stated as follows.

Original: Given a database X 2 R N ⇥ D and aquery vector q 2 R D , ﬁnd the subset S ⇤ ⇢ X collecting the K -closest data points to q :
Revised: Given a database X ∈ R N × D and a query vector q ∈ R D , ﬁnd the subset S ∗ ⊂ X collecting the K -closest data points to q :

Original: [Equation]
Revised: [Equation]

Original: where D ( x , y ) is a distance measure such as Euclidean distance D ` 2 ( x , y ) :
Revised: where D ( x , y ) is a distance measure such as Euclidean distance D (cid:96) 2 ( x , y ) :

Original: = k x  y k 2 or the cosinedistance D cos ( x , y ) :
Revised: = (cid:107) x − y (cid:107) 2 or the cosine distance D cos ( x , y ) :

Original: = 1  h x , y i k x kk y k .
Revised: = 1 − (cid:104) x , y (cid:105) (cid:107) x (cid:107)(cid:107) y (cid:107) .

Original: A related problem is the maximum inner product search (MIPS),where the goal is to ﬁnd the data points that have the highest inner products with the query:
Revised: A related problem is the maximum inner product search (MIPS), where the goal is to ﬁnd the data points that have the highest inner products with the query:

Original: [Equation]
Revised: [Equation]

Original: MIPS is equivalent to the cosine similarity search when all data points are ` 2 -normalized.
Revised: MIPS is equivalent to the cosine similarity search when all data points are (cid:96) 2 -normalized.

Original:  Related work
Revised:  Related work

Original: Exhaustively searching all pair-wise distances between the query and the entire database is computeintensive and often infeasible on many platforms.
Revised: Exhaustively searching all pair-wise distances between the query and the entire database is computeintensive and often infeasible on many platforms.

Original: Therefore, a problem extensively discussed in theliterature (Wang et al., 2014, 2015) is to ﬁnd approximate nearest neighbors (ANN) in exchange ofspeed.
Revised: Therefore, a problem extensively discussed in the literature (Wang et al., 2014, 2015) is to ﬁnd approximate nearest neighbors (ANN) in exchange of speed.

Original: By convention, the quality of ANN is measured by
Revised: By convention, the quality of ANN is measured by

Original: [Equation]
Revised: [Equation]

Original: where S q ⇢ X denotes the set of data points retrieved by the search method.
Revised: where S q ⊂ X denotes the set of data points retrieved by the search method.

Original: Compressed domain search One class of ANN approaches is to search on a lossy-compressedproblem domain.
Revised: Compressed domain search One class of ANN approaches is to search on a lossy-compressed problem domain.

Original: These methods are composed in two steps: a) search on compressed representation 1of the original problem to ﬁnd a set of candidate data points, b) compute the distances between thequery and the candidate data points to select the top- K results.
Revised: These methods are composed in two steps: a) search on compressed representationof the original problem to ﬁnd a set of candidate data points, b) compute the distances between the query and the candidate data points to select the top- K results.

Original: Since only a subset of data pointsrequires the exact distance computation, the overall cost is reduced.
Revised: Since only a subset of data points requires the exact distance computation, the overall cost is reduced.

Original: The two steps can be composed in arbitrary ways.
Revised: The two steps can be composed in arbitrary ways.

Original: Locality sensitive hashing (Andoni et al., 2015;
Revised: Locality sensitive hashing (Andoni et al., 2015;

Original: Neyshabur and Srebro, 2015) applies search followed by scoring; tree-search (Muja and Lowe, 2014;
Revised: Neyshabur and Srebro, 2015) applies search followed by scoring; tree-search (Muja and Lowe, 2014;

Original: Dasgupta and Freund, 2008) applies the two steps recursively; graph-search (Malkov and Yashunin,2018) iterates between two steps until the stopping condition is met.
Revised: Dasgupta and Freund, 2008) applies the two steps recursively; graph-search (Malkov and Yashunin, 2018) iterates between two steps until the stopping condition is met.

Original: And the inverted ﬁle (IVF)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.
Revised: And the inverted ﬁle (IVF) 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.

Original: We see that there are two major challenges with the compressed domain search:
Revised: We see that there are two major challenges with the compressed domain search:

Original: • Fractional search has a poor cache reuse rate because the candidate data points for each query rarely overlaps.
Revised: • Fractional search has a poor cache reuse rate because the candidate data points for each query rarely overlaps.

Original: We show optimizing the cache usage has a huge headroom for acceleratorsin Section 4.2.
Revised: We show optimizing the cache usage has a huge headroom for accelerators in Section 4.2.

Original: • Tweaking the speed-recall trade-off is data-dependent and non-trivial to tune.
Revised: Tweaking the speed-recall trade-off is data-dependent and non-trivial to tune.

Original: The key result of Beyer et al.
Revised: The key result of Beyer et al.

Original: (1999) states that the distance contrast of neighbors diminishes withincreasing dimensionality (also known as the curse of high dimensionality).
Revised: (1999) states that the distance contrast of neighbors diminishes with increasing dimensionality (also known as the curse of high dimensionality).

Original: Furthermore,the key result of Rubinstein (2018) states that sub-linear time nearest neighbor search withhigh recall is impossible for Euclidean, Manhattan, or Hamming distance; otherwise, itcontradicts the Strong Exponential Time Hypothesis (Impagliazzo and Paturi, 1999).Our work takes an opposite approach to focus on machine efﬁciencywith zero search space prun-
Revised: Furthermore, the key result of Rubinstein (2018) states that sub-linear time nearest neighbor search with high recall is impossible for Euclidean, Manhattan, or Hamming distance; otherwise, it contradicts the Strong Exponential Time Hypothesis (Impagliazzo and Paturi, 1999). Our work takes an opposite approach to focus on machine efﬁciency

Original: ing.
Revised: with

Original:  Moreover, since our method computes all the distances, it is immune to the curse of high
Revised: zero search space pruning Moreover, since our method computes all the distances , it is immune to the curse of high

Original: dimensionality
Revised: dimensionality.

Original: Accelerators In this paper, the phrase accelerators represents a class of specialized hardware toaccelerate machine learning workloads.
Revised: Accelerators In this paper, the phrase accelerators represents a class of specialized hardware to accelerate machine learning workloads.

Original: In particular, we are interested in the novel platforms thatdeliver high FLOP/s for distance computation, namely Google TPU V3, V4, Nvidia GPU V100, and
Revised: 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

Original: A100 in our analysis and evaluation.
Revised: A100 in our analysis and evaluation.

Original: Modern accelerators have special computation units for matrix multiplication, providing a higheroperation throughput over the regular coefﬁcient-wise operations.
Revised: Modern accelerators have special computation units for matrix multiplication, providing a higher operation throughput over the regular coefﬁcient-wise operations.

Original: The corresponding units are tensorcores in Nvidia GPUs (Markidis et al., 2018) and systolic arrays in Google TPUs (Jouppi et al., 2017;
Revised: The corresponding units are tensor cores in Nvidia GPUs (Markidis et al., 2018) and systolic arrays in Google TPUs (Jouppi et al., 2017;

Original: Norrie et al., 2021).
Revised: Norrie et al., 2021).

Original: Addressing these operation throughput differences is essential to our algorithmdesign.
Revised: Addressing these operation throughput differences is essential to our algorithm design.

Original: While accelerators excel in parallelism, developing an efﬁcient K -selection algorithm on acceleratorsis still an active research area (Monroe et al., 2011; Shanbhag et al., 2018; Johnson et al., 2021; Zhaoet al., 2020).
Revised: While accelerators excel in parallelism, developing an efﬁcient K -selection algorithm on accelerators is still an active research area (Monroe et al., 2011; Shanbhag et al., 2018; Johnson et al., 2021; Zhao et al., 2020).

Original: Accelerators with higher FLOP/s introduce a higher opportunity cost of computing the
Revised: Accelerators with higher FLOP/s introduce a higher opportunity cost of computing the

Original: K -selection problem instead of the distance computation.
Revised: K -selection problem instead of the distance computation.

Original: The trend of the increasing FLOP/s inaccelerators motivated us to optimize the FLOP/s usage by reducing the time required for computing
Revised: The trend of the increasing FLOP/s in accelerators motivated us to optimize the FLOP/s usage by reducing the time required for computing

Original: K -selection.
Revised: K -selection.

Original:  Methodology
Revised:  Methodology

Original: This section presents a performance model to identify non-trivial bottlenecks on multiple platforms and demonstrates some fundamental limits when designing algorithms for K -NN and relatedproblems, and we see that the cache inefﬁciency of the compressed domain methods introduces asigniﬁcant cost on accelerators.
Revised: This section presents a performance model to identify non-trivial bottlenecks on multiple platforms and demonstrates some fundamental limits when designing algorithms for K -NN and related problems, and we see that the cache inefﬁciency of the compressed domain methods introduces a signiﬁcant cost on accelerators.

Original: We model the accelerator’s runtime as executing a sequence of computation kernels , where eachkernel is a compiled subroutine on the accelerator used by the main program on the CPU.
Revised: We model the accelerator’s runtime as executing a sequence of computation kernels , where each kernel is a compiled subroutine on the accelerator used by the main program on the CPU.

Original: A kernelmay be composed of one or several high-level operators: Einsum, ReLU, ArgMax, etc., and eachkernel can have different performance characteristics.
Revised: A kernel may be composed of one or several high-level operators: Einsum, ReLU, ArgMax, etc., and each kernel can have different performance characteristics.

Original: Given a sequence of kernels k i , we let W i denotes the total amount of work and P i denotes theoperational speed.
Revised: Given a sequence of kernels k i , we let W i denotes the total amount of work and P i denotes the operational speed.

Original: Our goal is to estimate the total time of a program:
Revised: Our goal is to estimate the total time of a program:

Original: [Equation]
Revised: [Equation]

Original: In the following example, we focus on the MIPS problem.
Revised: In the following example, we focus on the MIPS problem.

Original: Let Q 2 R M ⇥ D and X 2 R N ⇥ D denotethe queries and the database, the runtime of a generic approximate-MIPS program can be modeled as
Revised: Let Q ∈ R M × D and X ∈ R N × D denote the queries and the database, the runtime of a generic approximate-MIPS program can be modeled as

Original: [Equation] 
Revised: [Equation] [Table]

Original: where W D denotes the total FLOPs required for searching the entire database, and  denotes thesearch fraction .
Revised: where W D denotes the total FLOPs required for searching the entire database, and λ denotes the search fraction .

Original: We note that P varies by algorithm and platform.
Revised: We note that P varies by algorithm and platform.

Original: Traditionally, compressed domainsearch methods minimize  but sacriﬁce cache efﬁciency.
Revised: Traditionally, compressed domain search methods minimize λ but sacriﬁce cache efﬁciency.

Original: Our method use an alternative route tooptimize P instead.[Table]
Revised: Our method use an alternative route to optimize P instead.

Original: 4.1 Instruction throughput-aware rooﬂine model
Revised: 4.1 Instruction throughput-aware rooﬂine model

Original: This subsection describes how we model the kernel-dependent performance P on multiple platformswith a small extension of the rooﬂine model.
Revised: This subsection describes how we model the kernel-dependent performance P on multiple platforms with a small extension of the rooﬂine model.

Original: The classic rooﬂine model (Williams et al., 2009) is a function of machine peak performance ⇡measured in FLOP/s, machine peak memory bandwidth  measured in bytes/s, and arithmeticintensity I MEM expressed as the ratio of ﬂoating-point operations performed to data movement(FLOP/byte).
Revised: The classic rooﬂine model (Williams et al., 2009) is a function of machine peak performance π measured in FLOP/s, machine peak memory bandwidth β measured in bytes/s, and arithmetic intensity I MEM expressed as the ratio of ﬂoating-point operations performed to data movement (FLOP/byte).

Original: The model states the performance is bounded by P  min( ⇡ ,  ⇥ I MEM ) .
Revised: The model states the performance is bounded by P ≤ min( π, β × I MEM ) .

Original: We desire to model kernels that has a mixture of ﬂoating point operations accelerated by dedicatedhardware as well as other coefﬁcient-wise operations.
Revised: We desire to model kernels that has a mixture of ﬂoating point operations accelerated by dedicated hardware as well as other coefﬁcient-wise operations.

Original: The coefﬁcient-wise operations are abbreviatedas COPs.
Revised: The coefﬁcient-wise operations are abbreviated as COPs.

Original: Almost every non matrix multiplication operations are COPs, including vectorized add,multiply, compare, conditional-move, etc.
Revised: Almost every non matrix multiplication operations are COPs, including vectorized add, multiply, compare, conditional-move, etc.

Original: We use the symbol  for peak COP/s on platforms, anddeﬁne the instruction throughput intensity I COP as the ratio between the number FLOPs and thenumber of COPs performed in a kernel (FLOP/COP).
Revised: We use the symbol γ for peak COP/s on platforms, and deﬁne the instruction throughput intensity I COP as the ratio between the number FLOPs and the number of COPs performed in a kernel (FLOP/COP).

Original: The attainable performance of a kernel isbounded by:
Revised: The attainable performance of a kernel is bounded by:

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: The statement is self-explanatory because the inadequate resources impede the kernel throughput.
Revised: The statement is self-explanatory because the inadequate resources impede the kernel throughput.

Original: Table 1 lists the properties of selected accelerators for our analysis 2 .
Revised: Table 1 lists the properties of selected accelerators for our analysis 3 .

Original: The rooﬂine model is commonlyused in accelerator proﬁling tools but not as frequently discussed in algorithm designs.
Revised: The rooﬂine model is commonly used in accelerator proﬁling tools but not as frequently discussed in algorithm designs.

Original: The followingsections show how the model prevents pitfalls due to the hardware constraints.
Revised: The following sections show how the model prevents pitfalls due to the hardware constraints.

Original: 4.2 The memory bandwidth bound
Revised: 4.2 The memory bandwidth bound

Original: This subsection demonstrates how to evaluate if a kernel hits the memory bandwidth wall.
Revised: This subsection demonstrates how to evaluate if a kernel hits the memory bandwidth wall.

Original: Weassociate the distance computation with three levels of BLAS (Dongarra et al., 1990).
Revised: We associate the distance computation with three levels of BLAS (Dongarra et al., 1990).

Original: Level 1 BLASdescribes vector operations on non-consecutive memory access, such as computing distances whiletraversing through a graph.
Revised: Level 1 BLAS describes vector operations on non-consecutive memory access, such as computing distances while traversing through a graph.

Original: Level 2 BLAS represents scoring a query with consecutively storeddata points.
Revised: Level 2 BLAS represents scoring a query with consecutively stored data points.

Original: Level 3 BLAS expresses batched query-database distance computation, often used inbrute-force scoring.
Revised: Level 3 BLAS expresses batched query-database distance computation, often used in brute-force scoring.

Original: Compressed domain searches are either level 1 or 2 BLAS due to the cache inefﬁciency.
Revised: Compressed domain searches are either level 1 or 2 BLAS due to the cache inefﬁciency.

Original: It has
Revised: It has

Original: O (1) memory arithmetic intensity because the number of FLOPs is proportion to the bytes read.
Revised: O (1) memory arithmetic intensity because the number of FLOPs is proportion to the bytes read.

Original: Combining (5) and (6) we have the following remark:
Revised: Combining (5) and (6) we have the following remark:

Original: Remark 1.
Revised: Remark 1.

Original: Distance computations in compressed domain searches are memory bandwidth bounded.
Revised: Distance computations in compressed domain searches are memory bandwidth bounded.

Original: In our model, the runtime is lower bounded by: t  O (  W D /  ) .
Revised: In our model, the runtime is lower bounded by: t ≥ O ( λW D /β ) .

Original: To estimate the memory arithmetic intensity for level 3 BLAS, we continue to use Q 2 R M ⇥ D and

Original: X 2 R N ⇥ D for denoting queries and database.

Original: In many K -NN applications N and M are muchgreater than D .

Original: The corresponding memory arithmetic intensity is:

Original: [Equation]

Original: [Figure]
Revised: [Figure]

Original: ⇡ = 6 .
Revised: π = 6 .

Original: 7 TFLOP/s ⇡ = [-
-]
Revised: 7 TFLOP/s π = [-
-]125 TFLOP/

Original: TFLOP/s ⇡ = 312 TFLOP/s ⇡ = 126 TFLOP/s ⇡{+
+}
Revised: s π = 312 TFLOP/ s π = {+
+}

Original:  = 274 TFLOP/ s
Revised: TFLOP/s π = 274 TFLOP/s

Original: [Figure]
Revised: [Figure]

Original: {+β = 9 0 
+}
Revised: {+β = 9 0 
+}

Original: {+G B / s β = 1 5 
+}
Revised: {+G B / s β = 1 5 
+}

Original:  B 
Revised: G B / s β = 8 5 8 . 4 G B / s β = 1 1 4 4 . 4 G B / s

Original: [Figure]
Revised: [Figure]

Revised: L A S 1 ,

Revised: [Figure]

Revised: I M E M =

Revised: O ( 1 )

Revised: B L A S 3 ,

Original: The largest term in the denominator of (7) is the 4 MN bytes of the query-database distances.{+D = 1 2 8 , I M E M = 
 ApproxTopK P e rf o r m a n ce+}
Revised: {+D = 1 2 8 , I M E M = 
 ApproxTopK P e rf o r m a n ce+}

Revised: P ( G

Revised: F L O

Revised: P / s )

Revised: [Figure]

Revised: To estimate the memory arithmetic intensity for level 3 BLAS, we continue to use Q ∈ R M × D and

Revised: X ∈ R N × D for denoting queries and database.

Revised: In many K -NN applications N and M are much greater than D .

Original: Weomit the insigniﬁcant terms and refer readers to (Golub and Van Loan, 2013, Section 1.5.4) for acomprehensive review on memory transfers in block matrix multiplications.
Revised: The corresponding memory arithmetic intensity is: [Equation] The largest term in the denominator of (7) is the 4 MN bytes of the query-database distances. We omit the insigniﬁcant terms and refer readers to (Golub and Van Loan, 2013, Section 1.5.4) for a comprehensive review on memory transfers in block matrix multiplications.

Original: Figure 1 shows that the distance scoring kernels of different BLAS levels can easily hit the memorybandwidth wall.
Revised: Figure 1 shows that the distance scoring kernels of different BLAS levels can easily hit the memory bandwidth wall.

Original: In order to attain high performance, we designed our algorithm to aggregate theresults within the kernel to avoid writing the O ( MN ) bytes into memory.
Revised: In order to attain high performance, we designed our algorithm to aggregate the results within the kernel to avoid writing the O ( MN ) bytes into memory.

Original: 4.3 The instruction bandwidth bound
Revised: 4.3 The instruction bandwidth bound

Original: The use of COPs (non matrix multiplication instructions) introduce another slowdown.
Revised: The use of COPs (non matrix multiplication instructions) introduce another slowdown.

Original: We let Cdenotes the number of COPs used per dot-product score in a kernel equipped with COPs and matrixmultiplication instructions.
Revised: We let C denotes the number of COPs used per dot-product score in a kernel equipped with COPs and matrix multiplication instructions.

Original: There are M ⇥ N dot-product scores, so the total COPs used in a kernelis CMN .
Revised: There are M × N dot-product scores, so the total COPs used in a kernel is CMN .

Original: To prevent hitting the COPs bandwidth wall, we must satisfy:
Revised: To prevent hitting the COPs bandwidth wall, we must satisfy:

Original: [Equation]
Revised: [Equation]

Original: The number of COPs we can afford in the kernels is scarce.
Revised: The number of COPs we can afford in the kernels is scarce.

Original: We take D = 128 as an example andsubstitute it into (9).
Revised: We take D = 128 as an example and substitute it into (9).

Original: We can only use 4 coefﬁcient-wise instructions per dot-product for TPU V4,and 16 for GPU A100.
Revised: We can only use 4 coefﬁcient-wise instructions per dot-product for TPU V4, and 16 for GPU A100.

Original: We conclude with the following remark:
Revised: We conclude with the following remark:

Original: Remark 2.
Revised: Remark 2.

Original: Exact and generic K -selection algorithm cannot be efﬁciently implemented with thecoefﬁcient-wise operations for the selected platforms (GPU V100, A100, TPU V3 and V4).
Revised: Exact and generic K -selection algorithm cannot be efﬁciently implemented with the coefﬁcient-wise operations for the selected platforms (GPU V100, A100, TPU V3 and V4).

Original: Because of Remark 2, we develop an approximate approach to achieve the peak performances.
Revised: Because of Remark 2, we develop an approximate approach to achieve the peak performances.

Original: Algorithm
Revised: Algorithm

Revised: 1 for i ←

Revised: M do for j ← do y i

Revised: j ← q i x j

Revised: ; l ← ShiftRight j

Original:  Our algorithm consists of two kernels:
Revised: b ← v i l ← end 9 end Our algorithm consists of two kernels:

Original: PartialReduce kernel computes the distances and partially aggregate the results from M ⇥ N distances to M ⇥ L distances with original indices.Output values
Revised: PartialReduce kernel computes the distances and partially aggregate the results from M × N distances to M × L distances with original indices.

Original: Output indices 1 for i 1 to M do for j 1 to N do y i,j h q i , x j i ;l

Original: ShiftRight( j, W ) ; /* Unrolled cost COP */ b y i,j > v i,l ; /* COP 1: compare */ v i,l if b then y i,j else v i,l ; /* COP 2: Vectorized conditional move */ a i,l if b then j else a i,l ; move *

Original: end 9 end 2. ExactRescoring kernel is an optional kernel that aggregates the ﬁnal top- K results.
Revised: ExactRescoring kernel is an optional kernel that aggregates the ﬁnal top- K results.

Original: Thecomplexity is O
Revised: The complexity is O

Original: ( ML log 2 ( L )) by a bitonic sort followed by a truncation.
Revised: ( ML log 2 ( L )) by a bitonic sort followed by a truncation.

Original: The PartialReduce kernel is where most of the time and compute takes place.
Revised: The PartialReduce kernel is where most of the time and compute takes place.

Original: See Algorithm 1 for anoutline of the algorithm.
Revised: See Algorithm 1 for an outline of the algorithm.

Original: We collect top-1 distances from the L non-overlapping bins of size 2 W foreach query, resulting high arithmetic intensities:
Revised: We collect top-1 distances from the L non-overlapping bins of size 2 W for each query, resulting high arithmetic intensities:

Original: [Equation] 
Revised: [Equation] [Equation]

Original: We show these arithmetic intensities can achieve high performance on real world database in section6.1.
Revised: We show these arithmetic intensities can achieve high performance on real world database in section 6.1.

Original: See Appendix A.3 for the detailed expansion of the algorithm and how the arithmetic intensitiesare derived.
Revised: See Appendix A.3 for the detailed expansion of the algorithm and how the arithmetic intensities are derived.

Original: 5.1 Recall estimation
Revised: 5.1 Recall estimation

Original: This section shows the PartialReduce kernel can achieve high recall with good speed.
Revised: This section shows the PartialReduce kernel can achieve high recall with good speed.

Original: We reformulateour problem in terms of balls and bins.
Revised: We reformulate our problem in terms of balls and bins.

Original: We have K balls representing the top- K distances that arethrown into L bins.
Revised: We have K balls representing the top- K distances that are thrown into L bins.

Original: The location of each ball is chosen independently and uniformly at random.
Revised: The location of each ball is chosen independently and uniformly at random.

Original: Welet Z denotes the random variable of the number of balls that do not have collisions.
Revised: We let Z denotes the random variable of the number of balls that do not have collisions.

Original: Following therecall deﬁnition (3) we have:
Revised: Following the recall deﬁnition (3) we have:

Original: [Equation]
Revised: [Equation]

Original: which is a standard Birthday problem:
Revised: which is a standard Birthday problem:

Original: [Equation]
Revised: [Equation]

Original: Our goal is to ﬁnd the minimal L such that the expected recall is greater equals to the target recall r .

Original: Finding L is simple because (13) is invertible in the natural range 0 < r < 1 .

Original: [Equation]

Original: [Figure]
Revised: [Figure]

Original: TPU
Revised: Williams et al

Original: [Figure]
Revised: [Figure]

Original: TPU V
Revised: Our goal is to ﬁnd the minimal L such that the expected recall is greater equals to the target recall r .

Original: [Figure] (Williams et al.
Revised: Finding L is simple because (13) is invertible in the natural range 0 < r < 1 .

Original: [Figure]
Revised: [Equation]

Original: Instruction bandwidth rooﬂine (Ours)

Original: The approximation in (14) follows from Appendix A.4.
Revised: The approximation in (14) follows from Appendix A.4.

Original: Since L is at the order of K , and in mostapplications K ⌧ N , the cost of the ExactRescoring kernel is amortized out.
Revised: Since L is at the order of K , and in most applications K (cid:28) N , the cost of the ExactRescoring kernel is amortized out.

Original: Thus we afﬁrm theclaim that our method attains high performance with an analytical recall guarantee.
Revised: Thus we afﬁrm the claim that our method attains high performance with an analytical recall guarantee.

Original:  Evaluation
Revised:  Evaluation

Original: In this section, we show that our proposed algorithm and implementation are near the hardware limitand lead to superior performance over the baselines of similar recalls.
Revised: In this section, we show that our proposed algorithm and implementation are near the hardware limit and lead to superior performance over the baselines of similar recalls.

Original: We applied our algorithm totwo datasets from the public ANN benchmarks (Aumüller et al., 2020).
Revised: We applied our algorithm to two datasets from the public ANN benchmarks (Aumüller et al., 2020).

Original: In our ﬁrst evaluation, wecompare the measured FLOP/s to the theoretical peak governed by the proposed reﬁnement of therooﬂine model (6), proclaiming our implementation is reaching the hardware peak performance.
Revised: In our ﬁrst evaluation, we compare the measured FLOP/s to the theoretical peak governed by the proposed reﬁnement of the rooﬂine model (6), proclaiming our implementation is reaching the hardware peak performance.

Original: Inthe second benchmark, we compare the end-to-end performance with competitive baselines withpre-tuned parameters.
Revised: In the second benchmark, we compare the end-to-end performance with competitive baselines with pre-tuned parameters.

Original: We plot each algorithm’s speed-recall curve and show ours achieves thestate-of-the-art.
Revised: We plot each algorithm’s speed-recall curve and show ours achieves the state-of-the-art.

Original: Finally, we measure the algorithm’s scalability by varying the dataset size and number of TPUs used.
Revised: Finally, we measure the algorithm’s scalability by varying the dataset size and number of TPUs used.

Original: 6.1 Comparison with the theoretical peak
Revised: 6.1 Comparison with the theoretical peak

Original: This section shows that our reﬁned rooﬂine model (6) captures additional performance characteristicover the classic rooﬂine model, and demonstrates our kernels are having near optimal performances.
Revised: This section shows that our reﬁned rooﬂine model (6) captures additional performance characteristic over the classic rooﬂine model, and demonstrates our kernels are having near optimal performances.

Original: We select the Glove 3 (Pennington et al., 2014) and Sift 4 (Jegou et al., 2010) datasets from the ANNbenchmarks.
Revised: We select the Glove 4 (Pennington et al., 2014) and Sift 5 (Jegou et al., 2010) datasets from the ANN benchmarks.

Original: Their corresponding distances are the cosine distance and the Euclidean distance.
Revised: Their corresponding distances are the cosine distance and the Euclidean distance.

Original: Seethe code snippets in Appendix A.1 and A.2.
Revised: See the code snippets in Appendix A.1 and A.2.

Revised: 4 Released in Apache license 2.0.Released in CC0 public domain.

Revised: [Figure]

Original:  See Figure 2, the colored lines represent machines’ max performances, and the dots representeach benchmark with its measured FLOP/s. The classic rooﬂine on the left shows that our incache aggregation strategy has a large memory arithmetic intensity ( ⇠ 4,700) exceeding the memorybandwidth ridge points ⇡ /  .
Revised: Flat) are released by Faiss (Johnson et al., 2021). For each IVF ∗ benchmark, the search fractions are λ = { 0 . See Figure 2, the colored lines represent machines’ max performances, and the dots represent each benchmark with its measured FLOP/s. The classic rooﬂine on the left shows that our incache aggregation strategy has a large memory arithmetic intensity ( ∼ 4,700) exceeding the memory bandwidth ridge points π/β .

Original: However, it is difﬁcult to diagnose why the Euclidean distance searchdoes not perform well on TPU V4 from the classic rooﬂine plot.
Revised: However, it is difﬁcult to diagnose why the Euclidean distance search does not perform well on TPU V4 from the classic rooﬂine plot.

Original: Fortunately, when combined with the instruction bandwidth rooﬂine we can tell the performance regression is caused by hitting the coefﬁcient-wise operation throughput wall.Therefore we afﬁrms [Figure]
Revised: Fortunately, when combined with the instruction bandwidth rooﬂine we can tell the performance regression is caused by hitting the coefﬁcient-wise operation throughput wall.

Original: Flat) are released by Faiss (Johnson et al., 2021).

Original: For each IVF ⇤ benchmark, the search fractions are  = { 0 .

Original:  the claim that our MIPS solution is reaching the peak FLOP/s, and our Euclidean distance searchsolution is meeting the compute bound on TPU V4 and attaining the peak FLOP/s on TPU V3.
Revised: Therefore we afﬁrms the claim that our MIPS solution is 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 on TPU V3.

Original: 6.2 Recall-speed benchmark
Revised: 6.2 Recall-speed benchmark

Original: To evaluate the effectiveness of the K -NN algorithm in a realistic setting, we adopted the methodologyof public ANN benchmarks (Aumüller et al., 2020) to compare the end-to-end performance againstother methods on the following datasets: Glove (Pennington et al., 2014), Sift (Jegou et al., 2010),
Revised: To evaluate the effectiveness of the K -NN algorithm in a 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),

Original: NYTimes (Dua and Graff, 2017), and Last.fm (Bertin-Mahieux et al., 2011).
Revised: NYTimes (Dua and Graff, 2017), and Last.fm (Bertin-Mahieux et al., 2011).

Original: The typical ANNbenchmarks are only performed on a single platform.
Revised: The typical ANN benchmarks are only performed on a single platform.

Original: However, it is non-trivial to either port our TPUalgorithm to GPU or vice versa.
Revised: However, it is non-trivial to either port our TPU algorithm to GPU or vice versa.

Original: Alternatively, we selected the following GPUs with parity in peakperformance to TPU (Table 1).
Revised: Alternatively, we selected the following GPUs with parity in peak performance to TPU (Table 1).

Original: We select the Faiss GPU (Johnson et al., 2021) implementation as our baseline.
Revised: We select the Faiss GPU (Johnson et al., 2021) implementation as our baseline.

Original: Faiss provides threealgorithms: Flat, IVF-Flat, and IVF-PQ.
Revised: Faiss provides three algorithms: Flat, IVF-Flat, and IVF-PQ.

Original: The Flat algorithm performs a brute-force search, and the
Revised: The Flat algorithm performs a brute-force search, and the

Original: IVF-Flat and IVF-PQ algorithms corresponds to the inverted ﬁle method with and without the productquantization (Jegouet al., 2010; Johnson et al., 2021). We use the repository’s suggested inverted ﬁlesize (16384) in the IVF methods. Figure 3 shows our performance signiﬁcantly outperforms competing methods in the high recallregions. We highlight that our method has a consistent recall-speed trade-off over different datasets,because our recall only rely on the order statistics instead of the information encoded in the compression domain search methods, which may vary by the datasets.
Revised: IVF-Flat and IVF-PQ algorithms corresponds to the inverted ﬁle method with and without the product quantization (Jegou

Original: Since our method scores all thepair-wise distances, our method is

Original: k 1 . 

Revised: k 1 .

Original: M 1 .
Revised: M 1 .

Original: 5 M 2 . 

Revised: 5 M 2 .

Original: M 2 .
Revised: M 2 .

Original: 5 M 3 . 

Revised: 5 M 3 .

Original: M 1 M 2 M 3 M 4 M 5 M 6 M 7 M 8 M 9 M 10 M
Revised: M 1 M 2 M 3 M 4 M 5 M 6 M 7 M 8 M 9 M 10 M

Original: Q P S
Revised: Q P S

Original: Database size N
Revised: Database size N

Original: Faiss IVF,Flat A 0.
Revised: Faiss IVF,Flat A 0.

Original: Ours TPU V4 1 chip 0.
Revised: Ours TPU V4 1 chip 0.

Original: Ours TPU V4 2 chips 0. Ours TPU V4 4 chips 0.
Revised: Ours TPU V4 2 chips 0. Ours TPU V4 4 chips 0.

Original: Yandex Deep D =
Revised: Yandex Deep D =

Revised: Figure 3 shows our performance signiﬁcantly outperforms competing methods in the high recall regions.

Revised: We highlight that our method has a consistent recall-speed trade-off over different datasets, because our recall only rely on the order statistics instead of the information encoded in the compression domain search methods, which may vary by the datasets.

Revised: Since our method scores all the pair-wise distances, our method is immune from the curse of high dimensionality.

Original: 6.3 Scalability benchmark
Revised: 6.3 Scalability benchmark

Original: In the ﬁnal benchmark, we examine the scalability of the algorithm from three aspects.
Revised: In the ﬁnal benchmark, we examine the scalability of the algorithm from three aspects.

Original: First, weverify if the measured performance is inverse proportional to the database size.
Revised: First, we verify if the measured performance is inverse proportional to the database size.

Original: Second, we comparethe scaling characteristics to the fastest GPU implementation.
Revised: Second, we compare the scaling characteristics to the fastest GPU implementation.

Original: Last but not least, we are interested inknowing if our algorithm can horizontally scale by the number of TPUs.
Revised: Last but not least, we are interested in knowing if our algorithm can horizontally scale by the number of TPUs.

Original: We conduct our evaluation on TPU V4 and Nvidia GPU A100, which have similar peak performanceand memory bandwidth.
Revised: We conduct our evaluation on TPU V4 and Nvidia GPU A100, which have similar peak performance and memory bandwidth.

Original: We sample the Yandex Deep dataset 5 (Babenko and Lempitsky, 2016) intoten different scales and measure the QPS of each approach with a similar recall.
Revised: We sample the Yandex Deep dataset 6 (Babenko and Lempitsky, 2016) into ten different scales and measure the QPS of each approach with a similar recall.

Original: Figure 4 veriﬁes allmeasurements align with the ideal scalability model: QPS / #chips /N .
Revised: Figure 4 veriﬁes all measurements align with the ideal scalability model: QPS ∝ #chips /N .

Original: Our method remains topperformance on all database sizes and linearly scales with the TPU chips..
Revised: Our method remains top performance on all database sizes and linearly scales with the TPU chips..

Original:  Discussion and future work
Revised:  Discussion and future work

Original: In contrast to classical fractional search methods, our method computes all the pair-wise distancesfor hardware efﬁciency.
Revised: In Section 6, we benchmark our method against others on platforms with similar performances. Some questions may arise, including "is the performance gain an algorithmic optimization or due to platform efﬁciency?" "can we achieve the same performance gain on GPU?" or "the existence of efﬁcient fractional-search on accelerators? " We address these questions in this section.

Revised: 7.1 Platform discussions

Revised: We ﬁrst discuss the modeling perspective of performance differences between platforms.

Revised: In Section 4, we show that the memory bandwidth and instruction throughput bound applies to both GPU and

Revised: TPU.

Revised: For instance, it follows that to attain peak performance on every hardware platform, having the number of instructions used for collecting (approximate) top-k elements within 2 γ · D/π per distance computation is a necessary condition .

Revised: We address the hardware constraints with Algorithm 1, a platform-independent solution that achieves peak performance, and we present a TPU implementation to support our claim.

Revised: Implementations for other platforms (such as GPU, CPU, or FPGA) are also feasible; however, this requires cache management, preventing cross-core memory synchronization, in-register accumulation, and instruction scheduling.

Revised: Typical high-performance libraries such as MKL, cuBLAS, and Google TPU compiler 6 Released in CC BY 4.0.

Original: However, using our solution in conjunction with fractional search methodsis also possible since brute-force distance computations are often involved in the auxiliary datastructures when performing the compressed domain search.
Revised: use platform-speciﬁc assembly to take full control of stated requirements. Moreover, these libraries are all close-sourced. Fortunately, we have access to the TPU compiler to generate the assembly for Algorithm 1, which is sufﬁcient to solidify our analysis.

Revised: Thus we leave implementations of other platforms to future works.

Revised: 7.2 Algorithm discussions

Revised: The rooﬂine complexity of the fractional search is identical to BLAS-2 (matrix-vector multiplication), which is memory bandwidth bound.

Revised: When the cycles spend on data transfer are mutually exclusive to our method, it introduces an enormous opportunity cost.

Revised: Nevertheless, we see an opportunity in a heterogeneous architecture because a fractional search on the host is not mutually exclusive to applying our method to accelerators.

Original: An motivating example is  multi-billion nearest neighbor search, where ﬁtting the dataset into devicememory is possible  but not economic.
Revised: A motivating example is the multi-billion nearest neighbor search, where ﬁtting the dataset into device memory is possible (through device sharding, which TensorFlow and Jax have native support) but not economical.

Revised: Since brute-force distance computations are often involved in the auxiliary data structures when performing the fractional search, we may replace the brute-force portion with TPU in conduction with the remaining search off-device.

Original: We note that heterogeneous architectures with off-devicestorage such as host-RAM or even SSD (Jayaram Subramanya et al., 2019; Ren et al., 2020; Chenet al., 2021) are great starting points for future research.
Revised: We note that heterogeneous architectures with off-device storage such as host-RAM or even SSD (Jayaram Subramanya et al., 2019; Ren et al., 2020; Chen et al., 2021) are great starting points for future research.

Original:  Conclusion
Revised:  Conclusion

Original: Accelerator-based machine learning has become the mainstream in academics and industries.
Revised: Accelerator-based machine learning has become the mainstream in academics and industries.

Original: However, the performance characteristics of accelerators are counter-intuitive and difﬁcult to program.
Revised: However, the performance characteristics of accelerators are counter-intuitive and difﬁcult to program.

Original: In this paper, we propose a rooﬂine-based complexity analysis framework to discuss the optimalityof the algorithms without low-level optimization details: unrolling factors, batch window sizes,vectorization, and systolic array scheduling, which are platform-dependent and lengthy to read.
Revised: In this paper, we propose a rooﬂine-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 platform-dependent and lengthy to read.

Original: Wedemonstrated several examples of inferring the hardware performance limits by simply addressingthe kernel’s total FLOPs, byte transferred, and the number of coefﬁcient-wise instructions used.
Revised: We demonstrated several examples of inferring the hardware performance limits by simply addressing the kernel’s total FLOPs, byte transferred, and the number of coefﬁcient-wise instructions used.

Original: Ourreﬁned model foreshadowed non-trivial performance regression caused by the coefﬁcient-wise instructions bandwidth.
Revised: Our reﬁned model foreshadowed non-trivial performance regression caused by the coefﬁcient-wise instructions bandwidth.

Original: We took it into account to design a new algorithm for K -NN and achieved peakperformance on TPU.
Revised: We took it into account to design a new algorithm for K -NN and achieved peak performance on TPU.

Original: Finally, our experiments showed that our method outperformed state-of-the-artbaselines on platforms with similar performance characteristics, which are known to be hard to beat.5 Released in CC BY 4.0.
Revised: Finally, our experiments showed that our method outperformed state-of-the-art baselines on platforms with similar performance characteristics, which are 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_k . We are also grateful to the

Revised: Google ScaNN team for the joint effort on bridging the impactful K -NN problem into the accelerator ecosystem.

Revised: Last but not least, we thank to Peter Hawkins, Edward Schwartz and Mani Varadarajan for code reviews in Jax and Tensorﬂow.

Revised: This work was performed and funded by Google.

Revised: A Appendix A.1 MIPS implementation import

Revised: jax @jax.jit defMIPS(query, database): scores =

Revised: jax.numpy.einsum('ik,jk->ij', query, database) return jax.lax.approx_max_k(scores, k=10, recall_target=0.95)

Revised: Listing 1: Jax code for maximum inner product search (MIPS)

Revised: Listing 1 demonstrates a maximum inner product search (MIPS) kernel implemented with Jax.

Revised: Tensorﬂow users can use the tf.math.approx_max_k interface; the underlying XLA compiler delivers the same kernel.

Revised: There are several options to control the behavior of approx_max_k , listed as below:

Revised: reduction_dimension speciﬁes the dimension in which to search.

Revised: Default -1 (the last dimension.)

Revised: recall_target derives the number of bins L of the PartialReduce kernel output.

Revised: Default 0.95. 3.

Revised: reduction_input_size_override .

Revised: When set to a positive value, it overrides the size determined by input for evaluating the recall and bin numbers L .

Revised: Users could use this option to control the kernel output size in the distributed environment.

Revised: aggregate_to_topk .

Revised: When set to True emits the ExactRescoring kernel.

Revised: Default: True .

Revised: We also provid a separated approx_min_k interface for ﬁnding minimum distances, which is used in the Euclidean distance search.

Revised: A.2 Euclidean distance search implementation

Revised: @jax.jit defl2nns(qy, db, db_half_norm): dots =

Revised: jax.numpy.einsum('ik,jk->ij', qy, db) dists = db_half_norm - dots return jax.lax.approx_min_k(dists, k=10, recall_target=0.95)

Revised: Listing 2: Jax code for nearest neighbor search in the Euclidean space.

Revised: Listing 2 is the Jax implementation of Euclidean space nearest neighbor search.

Revised: We made a few adjustments to speed up the computation.

Revised: First, for every query vector q , the following search produces the same result:

Revised: [Equation]

Revised: The last equation holds because omitting the query norm does not affect the rank of each result.

Revised: Nevertheless, (18) still uses 2 COPs for the distance computation (one subtract and one multiplication).

Revised: We can further reduce it to 1 COP by pre-computing the halved norm:

Revised: [Equation]

Revised: A.3 MIPS PartialReduce kernel internals

Revised: The MIPS PartialReduce kernel follows the standard numerical computation best practices to utilize the cache usage with the temporal and spatial locality .

Revised: See Algorithm 2 that uncovers the omitted details in Algorithm 1.

Revised: Algorithm 2:

Revised: Input: Q ∈ R

Revised: D Batch queries

Revised: Input: X ∈ R N × D Database

Revised: Input: 2 W Bin size

Revised: Output: V ∈ R M

Revised: ×

Revised: L

Revised: Top- K values

Revised: Output: A ∈ N M × L Top- K indices /* Block iteration over rows */ 1 for ii ← 1 to M step ib do /* Block iteration over columns */ for jj ←

Revised: N step jb do /* l often unrolled or even vectorized */ for i ← ii to ii

Revised: + ib − 1 do /*

Revised: Starts the inner loop of the systolic arrays */ y i ← 0 ; for k ← 1 to D do m ← q i,k ; /* Vectorized FMA fused-multiply-add)

Revised: */ for j ← jj to jj + jb − 1 do y i,j ← y i,j + m · x j,k ; end end /

Revised: Ends loop the systolic arrays */ for j ← jj to jj + jb − 1 do /

Revised: The exact j → l mapping is determined by the compiler backend */ l ← RegisterAlignedShiftRight

Revised: ; b ←

Revised: The temporal locality refers to reusing previously accessed items.

Revised: In line 1, we iterate by blocks of queries.

Revised: The block of queries is reused in the inner loops, achieving the temporal locality.

Revised: The spatial locality refers to accessing items nearby previously accessed.

Revised: The block iteration loads a chunk of data points (line 2) to achieve this optimization.

Revised: The same block iteration structure may apply recursively for multiple cache hierarchies till the register level.

Revised: The inner loops (indexed by i , j , and k in line 3) should be unrolled or even vectorized so that every cycle can produce multiple results via the SIMD (Single Instruction Multiple Data) instructions or systolic arrays.

Revised: # qy shape: f32[1024,128], db shape: f32[1048576,128] # output shapes: f32[1024, 128], i32[1024, 128] @jax.jit defmips_baseline(qy, db): dists =

Revised: jax.numpy.einsum('ik,jk->ij', qy, db) reshaped =

Revised: jax.lax.reshape(dists, [1024, 128, 8192]) return jax.lax.argmax(reshaped, 2, jnp.int32)

Revised: Listing 3: Baseline implementation without the approx_max_k operator

Revised: The algorithm principle is the same on every platform, except that the block factor and vectorization sizes are platform-dependent.

Revised: We refer readers to (Golub and Van Loan, 2013) for more details.

Revised: A.3.1 Estimate memory transfers

Revised: In Algorithm 2, memory transfer for each portion of the data is listed below:

Revised: • Query is only transferred once.

Revised: Takes 4 MD bytes.

Revised: • Database is transferred Mib times.

Revised: Takes 4 ND Mib bytes.

Revised: • Outputs are transferred once.

Revised: Takes 2 × 4 ML bytes.

Revised: The precise formulation for memory arithmetic intensity is

Revised: [Equation]

Revised: which would approach O (min( M, N )) as long as L (cid:28) min( M, N ) and the compiler chooses a large enough ib to minimize the database transfer.

Revised: A.3.2 Estimate COPs used

Revised: The PartialReduce kernel listed in Algorithm 1 and 2 only use C = 3 per dot-product.

Revised: However, there are two cases that would increase the number of COPS on TPU due to the implementation constraints:

Revised: When the dimension D is not multiple of 128, C increases by 1. 2.

Revised: When the database size N is not power-of-2, C increases by 1.

Revised: See Appendix A.5 on how it affects the real world benchmarks.

Revised: A.3.3 Limitation of naive implementation(s)

Revised: A naive implementation of Algorithm 1 and 2 can be composed of Reshape and ArgMax.

Revised: However, the performance is not comparable to the dedicated approx_max_k operator.

Revised: Our experiment setup is as follows: let query be Q ∈ R 1024 × 128 and database be X ∈ R 1048576 × 128 ; we choose the reduction output size as L = 128 , so the algorithm can be written as Listing 3.

Revised: We benchmark the implementations on a single-core TPU V4 instance by 100 times and collect the averaged execution time.

Revised: Listing 3 took 24.9ms to compute; in comparison, our proposed new operator used in Listing 1 only took 2.6ms, which is 9.6x faster.

Revised: A.4 Lower bound approximation of the number of bins

Revised: We care about the number of bins L in the high recall region.

Revised: Let the target recall r = 1 − (cid:15) , we have

Revised: L ≥

Revised: [Equation]

Revised: [Equation]

Revised: The approximation in (23) follows from (1 − (cid:15) )

Revised: a = (1 − (cid:15) )

Revised: [Equation]

Revised: ( − (cid:15)a ) → e − (cid:15)a as (cid:15) → 0 , and (24) follows from the Taylor expansion.

Revised: A.5 Benchmark details

Revised: [Table]

Revised: Table 2 summarizes the dimensions and kernel properties for the two benchmarks.

Revised: The memory arithmetic intensity I MEM is reported by the TPU proﬁler

Revised: [Table]

Revised: and the instruction throughput intensity

Revised: I COP is manually derived.

Revised: The following show how we derive C (COPs per dot-product) for each dataset.

Revised: Glove The Glove dataset uses the cosine distance, which yields same search results as MIPS.

Revised: As described in Appendix A.3, when the database size is not power-of-2, we pay an extra C in the inner loop.

Revised: Therefore the total C used for the Glove benchmark are

Revised: • 3 C by PartialReduce, and • 1 C by non power-of-2 database masking.

Revised: We pre-process the Glove dataset by padding the dimension from 100 to 128 to avoid one C .

Revised: We are not bottleneck on memory bandwidth so the padding is a good trade-off for better performance.

Revised: Sift The Sift dataset uses the Euclidean distance, which requires more coefﬁcient-wise operations.

Revised: In Appendix A.2 we showed that we only need to use one extra C for distance computation.

Revised: However, there are some other inevitable operations used in the benchmark:

Revised: • 3 C by PartialReduce, • 1 C by the relaxed Euclidean distance computation, • 1 C by non power-of-2 database masking, and • 1 C by broadcasting (cid:107) x (cid:107) 2 2 .

Revised: Therefore the total number of C = 6 , resulting a performance regression on TPU V4 as seen in

Revised: Figure 2.

Revised: from functools import partial import jax import jax.numpy as jnp # Computes approx_max_k on each TPU device in parallel.

Revised: @partial( jax.pmap, # static args: db_size, k, recall_target static_broadcasted_argnums=[3, 4, 5], out_axes=(1, 1))

Revised: def pmap_mips(qy, db, db_offset, db_size, k, recall_target): dists =

Revised: jnp.einsum('ik,jk->ij', qy, db) dists, neighbors =

Revised: jax.lax.approx_max_k( dists, k=k, recall_target=

Revised: recall_target, reduction_input_size_override=db_size) return (dists, neighbors + db_offset) # Collects partial results from each device, and # aggregates to the final top-k. @partial( jax.jit, static_argnames=["k"]) def mips_reorder(dists, neighbors, k): flat_dists =

Revised: jax.lax.collapse(dists, 1, 3) flat_nn = jax.lax.collapse(neighbors, 1, 3) _, sorted_nn = jax.lax.sort([-flat_dists, flat_nn]) return jax.lax.slice_in_dim(sorted_nn, 0, k, axis=1)

Revised: # Entry function for distributed MIPS.

Revised: def mips(qy, db, db_offset, db_size, k, recall_target): dists, neighbor = pmap_mips(qy, db, db_offset, db_size, k, recall_target) return mips_reorder(dists, neighbor, k)

Revised: Listing 4: Baseline implementation without the approx_max_k operator

Revised: A.

Revised: Multi-TPU implementation

Revised: Listing 4 three implements a parallel MIPS using Jax’s pmap API.

Revised: The program breaks down into two parts: a) parallel execution of dot-product followed by approx max K , and b) top-

Revised: K aggregation of results collected from each TPU device.

Revised: Let U denotes number of TPU devices, the number of itmes to reorder is 2 KU , where the complexity follows bitonic sort: O ( KU log 2 ( KU )) .

Revised: In the distributed computing use case the dataset size N is often much larger than KU , thus the parallel portion dominates the computation and scales linearly.