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)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: 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 proposedalgorithm is motivated by an accurate accelerator performance model that takes intoaccount 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 withan analytical guarantee of recall in expectation and does not require maintainingsophisticated index data structure or tuning, making it suitable for applicationswith 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.
Revised: 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.

Original: ,2019; Omar et al., 2013), recommendation systems (Sarwar et al., 2002;
Revised: ,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 asserving as a component for a downstream tasks (Borgeaud et al., 2021; Guu et al., 2020; Lindgrenet 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 closestdatapoints 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 ofdesigning a good K -NN algorithm is to compute accurate K -NN results while being computationallyefﬁ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 theindustry (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 candeliver hundreds of Tera Floating Point Operations Per Seconds (TFLOPS) vital to the neighbordistance 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 needto 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 performancemodel (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, 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.

Original: The main contributions of this work are:
Revised: The main contributions of this work are:

Original: • We extend the rooﬂine model to address the operation throughput differences of the instruc-
Revised: • We extend the rooﬂine model to address the operation throughput differences of the instruc-

Original: tions, essential to the algorithm analysis in this paper.
Revised: tions, 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 basedon our proposed rooﬂine model.

Original: • We conduct experiments verifying our TPU implementation of the algorithm accurately
Revised: • We conduct experiments verifying our TPU implementation of the algorithm accurately

Original: aligned with the performance model and achieves state-of-the-art speed-recall trade-offs onstandard nearest neighbor search benchmarks.
Revised: aligned with the performance model and achieves state-of-the-art speed-recall trade-offs onstandard 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
Revised: Givena matrix

Original: 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: 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 .

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 of 41 the matrix X in a D dimensional vector space.
Revised: 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 of 41 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 areused 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 ⇤ ⇢
Revised: Given a database X 2 R N ⇥ D and aquery vector q 2 R D , ﬁnd the subset S ⇤ ⇢

Original: X collecting the K -closest data points to q :
Revised: 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 ` 2 ( x , y ) :

Original: = k x  y k 2 or the cosinedistance D cos ( x , y ) :
Revised: = k x  y k 2 or the cosinedistance D cos ( x , y ) :

Original: = 1  h x , y i k x kk y k .
Revised: = 1  h x , y i k x kk y k .

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 ` 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 theliterature (Wang et al., 2014, 2015) is to ﬁnd approximate nearest neighbors (ANN) in exchange ofspeed.

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-compressedproblem 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 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.

Original: Since only a subset of data pointsrequires the exact distance computation, the overall cost is reduced.
Revised: Since only a subset of data pointsrequires 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
Revised: • Fractional search has a poor cache reuse rate because the candidate data points for each query

Original: rarely overlaps.
Revised: 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 acceleratorsin 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
Revised: The key

Original: result of Beyer et al.
Revised: 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 withincreasing 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ﬁciency with zero search space prun-
Revised: 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ﬁciency with zero search space prun-

Original: ing.
Revised: ing.

Original: Moreover, since our method computes all the distances, it is immune to the curse of high
Revised: 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 toaccelerate 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 thatdeliver 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 higheroperation 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 tensorcores 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 algorithmdesign.

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 acceleratorsis still an active research area (Monroe et al., 2011; Shanbhag et al., 2018; Johnson et al., 2021; Zhaoet 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 inaccelerators 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 relatedproblems, and we see that the cache inefﬁciency of the compressed domain methods introduces asigniﬁ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 eachkernel 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 kernelmay be composed of one or several high-level operators: Einsum, ReLU, ArgMax, etc., and eachkernel 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 theoperational 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 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

Original: [Equation]
Revised: [Equation]

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 thesearch 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 domainsearch methods minimize  but sacriﬁce cache efﬁciency.

Original: Our method use an alternative route tooptimize P instead.
Revised: Our method use an alternative route tooptimize P instead.

Original: [Table]
Revised: [Table]

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 platformswith 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 arithmeticintensity 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 dedicatedhardware as well as other coefﬁcient-wise operations.

Original: The coefﬁcient-wise operations are abbreviatedas COPs.
Revised: The coefﬁcient-wise operations are abbreviatedas 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, 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).

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

Original: [Equation] [Equation]
Revised: [Equation] [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 2 .

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 commonlyused 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 followingsections 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: Weassociate 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 BLASdescribes vector operations on non-consecutive memory access, such as computing distances whiletraversing 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 storeddata 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 inbrute-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
Revised: 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.
Revised: X 2 R N ⇥ D for denoting queries and database.

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

Original: The corresponding memory arithmetic intensity is:
Revised: The corresponding memory arithmetic intensity is:

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

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

Original: 7 TFLOP/s ⇡ = 

Revised: 7 TFLOP/s ⇡ =

Original: TFLOP/s ⇡ = 312 TFLOP/s ⇡ =
Revised: TFLOP/s ⇡ = 312 TFLOP/s ⇡ =

Original: 126 TFLOP/s ⇡
Revised: 126 TFLOP/s ⇡

Original: = 274 TFLOP/
Revised: = 274 TFLOP/

Original: s
Revised: s

Original: [Figure]
Revised: [Figure]

Original: B
Revised: B

Original: [Figure]
Revised: [Figure]

Original: The largest term in the denominator of (7) is the 4 MN bytes of the query-database distances.
Revised: The largest term in the denominator of (7) is the 4 MN bytes of the query-database distances.

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: 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.

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 memorybandwidth 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 theresults 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 Cdenotes the number of COPs used per dot-product score in a kernel equipped with COPs and matrixmultiplication 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 kernelis 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 andsubstitute 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 thecoefﬁ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

Original: Our algorithm consists of two kernels:
Revised: Our algorithm consists of two kernels:

Original: PartialReduce kernel computes the distances and partially aggregate the results from M ⇥ N
Revised: PartialReduce kernel computes the distances and partially aggregate the results from M ⇥ N

Original: distances to M ⇥
Revised: distances to M ⇥

Original: L distances with original indices.
Revised: L distances with original indices.

Original: Output values
Revised: Output values

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
Revised: 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 *
Revised: 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.
Revised: end 9 end 2.

Original: 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: Thecomplexity 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 anoutline 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 foreach query, resulting high arithmetic intensities:

Original: [Equation]
Revised: [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 section6.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 intensitiesare 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 reformulateour 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 arethrown 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 does not share the bin with other balls.
Revised: Welet Z denotes the random variable of the number of balls that do not have collisions.

Original: Following the recall deﬁnition (3) we have:
Revised: Following therecall 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 .
Revised: 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 .
Revised: Finding L is simple because (13) is invertible in the natural range 0 < r < 1 .

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: TPU
Revised: TPU

Original: [Figure]
Revised: [Figure]

Original: TPU V
Revised: TPU V

Original: [Figure]
Revised: [Figure]

Original: (Williams et al.
Revised: (Williams et al.

Original: [Figure]
Revised: [Figure]

Original: Instruction bandwidth rooﬂine (Ours)
Revised: 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 mostapplications K ⌧ 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 theclaim 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 is near the hardware limitand leads 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 limitand lead to superior performance over the baselines of similar recalls.

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

Original: In our ﬁrst evaluation,we compares the measured FLOP/s to the theoretical peak governed by the proposed reﬁnement ofthe rooﬂine model (6), proclaiming our implementation is reaching the hardware peak performance.
Revised: 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.

Original: In the second benchmark, we compare the end-to-end performance with competitive baselines withpre-tuned parameters.
Revised: Inthe second benchmark, we compare the end-to-end performance with competitive baselines withpre-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 thestate-of-the-art.

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 characteristicover 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 3 (Pennington et al., 2014) and Sift 4 (Jegou et al., 2010) datasets from the ANNbenchmarks.

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: Seethe code snippets in Appendix A.1 and A.2.

Original: See Figure 2, the colored lines represent machines’ max performances, and the dots representeach benchmark with its measured FLOP/s.
Revised: See Figure 2, the colored lines represent machines’ max performances, and the dots representeach benchmark with its measured FLOP/s.

Original: 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: 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 ⇡ /  .

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 searchdoes 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 performanceregression is caused by hitting the coefﬁcient-wise operation throughput wall. 
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. Therefore we afﬁrms [Figure]

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

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

Original: Therefore we afﬁrms 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: 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.

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.
Revised: 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), NYTimes (Dua and Graff, 2017), and Last.fm (Bertin-Mahieux et al., 2011).

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

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

Original: Alternatively, we selected thefollowing GPUs with parity in peak performance to TPU (Table 1).
Revised: Alternatively, we selected the following GPUs with parity in peakperformance 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-ﬂat, and IVF-PQ.
Revised: Faiss provides threealgorithms: 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 (Jegou et al., 2010; Johnson et al., 2021).
Revised: IVF-Flat and IVF-PQ algorithms corresponds to the inverted ﬁle method with and without the productquantization (Jegou et al., 2010; Johnson et al., 2021).

Original: We use the repository’s suggested inverted ﬁlesize (16384) in the IVF methods.
Revised: We use the repository’s suggested inverted ﬁlesize (16384) in the IVF methods.

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

Original: 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: 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.

Original: Since our method scores all thepair-wise distances, our method is immune from the curse of high dimensionality.{+k 1 . 
 M 1 . 5 M 2 . 
+}
Revised: Since our method scores all thepair-wise distances, our method is {+k 1 . 
 M 1 . 5 M 2 . 
+}

Original: Discussion and future work{+M 2 . 5 M 3 . 
 M 1 M 2 M 3 M 4 M 5 M 6 M 7 M 8 M 9 M 10 M Q P S+}
Revised: {+M 2 . 5 M 3 . 
 M 1 M 2 M 3 M 4 M 5 M 6 M 7 M 8 M 9 M 10 M Q P S+}

Revised: Database size N

Revised: Faiss IVF,Flat A 0.

Revised: Ours TPU V4 1 chip 0.

Revised: Ours TPU V4 2 chips 0.

Revised: Ours TPU V4 4 chips 0.

Revised: Yandex Deep D =

Revised: 6.3 Scalability benchmark

Original: We limit our experiments and discussion to single-chip accelerator K -NN performance  of densevectors.
Revised: In the ﬁnal benchmark, we examine the scalability of the algorithm from three aspects. First, weverify if the measured performance is inverse proportional to the database size. Second, we comparethe scaling characteristics to the fastest GPU implementation. Last but not least, we are interested inknowing if our algorithm can horizontally scale by the number of TPUs.

Original: Accelerators performance on sparse vectors follow a completely different paradigm due torandom memory access.
Revised: We conduct our evaluation on TPU V4 and Nvidia GPU A100, which have similar peak performanceand memory bandwidth.

Original:  Our algorithm can be naturally extended to multi-chip or distributed settingsto handle billion scale datasets through Tensorﬂow’s tf.distribute (Abadi et al., 2015) or Jax’sjax.pmap (Bradbury et al., 2018) programming interfaces.
Revised: 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. Figure 4 veriﬁes allmeasurements align with the ideal scalability model: QPS / #chips /N . Our method remains topperformance on all database sizes and linearly scales with the TPU chips..  Discussion and future work

Original: It is also possible to use our operations in conjunction with other strategies, including dimensionreduction, quantization and tree search (Van Der Maaten et al., 2009; Jegou et al., 2010; Wang et al.,2014), because many compressed domain search methods use brute-force distance computation on itsauxiliary data structures before performing the fractional search.
Revised: In contrast to classical fractional search methods, our method computes all the pair-wise distancesfor hardware efﬁciency. 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.

Original:  We note that heterogeneous architectures with off-HBM storage such as host-RAM or even SSD (Chen et al.,2021; Jayaram Subramanyaet al., 2019; Ren et al., 2020) are great starting points for future research.Finally, while our reﬁnement of the rooﬂine model handles kernel with mixture of instructions thathave different throughput, there are many studies that have extended the rooﬂine model to relatedtopics in recent years: GPU warp instruction rooﬂine (Ding and Williams, 2019), time-based rooﬂine(Wang et al., 2020), rooﬂine for multiple cache tiers (Yang et al., 2021), and energy rooﬂines (Choiet al., 2013; Lopes et al., 2017). Readers may refer to these models for some analysis that are left out,such as the auxiliary work in (5).
Revised: An motivating example is multi-billion nearest neighbor search, where ﬁtting the dataset into devicememory is possible but not economic. 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.

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 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.

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: 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.

Original: Ourreﬁned model foreshadowed non-trivial performance regression caused by the coefﬁcient-wise instructions bandwidth.
Revised: Ourreﬁ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 peakperformance 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. [Figure] Checklist For all authors...
Revised: 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.

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes] See Section 4 for how we model

Original: the hardware limitations and Section 6.1 for real world evaluations.

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] See Section5.1, we formulate the problems in terms of the classic balls into bins.

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See Section 5.1 andAppendix A.4.

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [Yes] See Appendix

Original: A.1 and A.

Original: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [N/A]

Original: (c) Did you report error bars (e

Original: g., with respect to the random seed after running experiments multiple times)?

Original: [N/A]

Original: (d) Did you include the total amount of compute and the type of resources used (e

Original: g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Table 1 and 2.

Original: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...

Original: (a)

Original: If your work uses existing assets, did you cite the creators?

Original: [Yes] See Section 6.1.(b)

Original: Did you mention the license of the assets?

Original: [Yes] See the footnotes in Section 6.1.(c)

Original: Did you include any new assets either in the supplemental material or as a URL?

Original: [No]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: (a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?

Original: [N/A]

Original: you describe any

Original: participant risks, with links to Institutional Review

Original: [Table]

Original: Board

Original: approvals,

Original: applicable

Original: N

Original: [Table]

Original: you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation

Original: [N/A]

Original: [Table]