Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The paper presents a new NN-search algorithm, that reaches the peak performance on TPUs. The paper is based on observations of hardware architectural properties and the so called roofline performance model. The algorithm is implemented for TPUs. The evaluation is done on two TPU versions using two KNN datasets, and compared to several GPU algorithms. 

strengths_and_weaknesses: Strengths
+ Nice connection between theoretical observations and practical results
+ Good performance
+ Can have significant practical impact

Weaknesses
- Limited evaluation, only evaluated on two versions on TPUs, and two datasets
- Would have been interesting to see how general the algorithm is, i.e., would it reach the same performance limits in other hardware platforms also?

questions: * The algorithm descriptions (Algorithm 1 and Algorithm 2 (suppl. mtrl)) looks relatively clear and straight-forward to implement on other platforms. Can you elaborate a bit on why it would be such a substantial effort to do?

limitations: I think the authors have adequately addressed the limitations of their work.
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: This paper presents an ANN algorithm on TPU and analyzes the memory and instruction bandwidth of ANN algorithms.

strengths_and_weaknesses: Strong Points
----
1. The problem is well-motivated. 
2. The related work is comprehensively studied and discussed. And the entire story is easy to follow for readers. 
3. The methodology to study the algorithm from the hardware bottlenecks is interesting.

Weak Points
----
1. I suggest the authors to include a brief discussion of TPU, e.g., what it can be done and what it can not be done (efficiently).
2. The experimental evaluation is a bit problematic. How about other methods on GPUs/TPUs, e.g., hashing and graph-based methods besides the FAISS baseline?
3. In Figure 3, the highest recall shown for Glove is 0.9. Is there a recall limitation for the proposed algorithm? 
4. The technical contribution of the bi-stage partial reduction and scoring is limited.

questions: 1. Could we have other methods on GPUs/TPUs, e.g., hashing and graph-based methods besides the FAISS baseline, on more measures, e.g. cosine?
3. Could we show that the proposed algorithm can achieve high recalls for most ANN datasets?
limitations: I suggest the authors to discuss the limitation of the algorithm, e.g., how to adapt the problem to other common similarity measures.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper presents a new algorithm implementation of K nearest neighbor search that is able to increase the arithmetic intensity. The distances between all queries and entries in the database are calculated in L3 BLAS fashion. Then only the top-1 entry is kept in each bin. The bin size can be adjusted based on the requirement of the recall. The roofline analysis shows a better peak throughput while the end-to-end evaluation demonstrates a better throughput-recall trade-off than the previous solutions.
strengths_and_weaknesses: ## Strengths

1. This paper presents a detailed analytical model of the kernel implementation of the BLAS-based operation. The authors conclude the impact of COPs and introduce the algorithm accordingly.

2. The proposed solution is simple yet effective. The adjustment of the bin size provides a simple method to balance recall and throughput.

## Weaknesses

1. **Unclear hardware motivation.** The paper is entitled 'TPU-KNN'. However, throughout this paper, I'm not able to find any information indicating that the proposed algorithm requires any special architectural support from TPU rather than other accelerators. It seems that the same analysis could also apply to GPU or other accelerators. 

2. **Vague definition of COPs.** I cannot understand the usage of the concept of coefficient-wise operations (COPs). In Table 1, comparing with the datasheet of A100/V100, I can understand COPs as a FLOP in the vector cores (or SMs in NV's terminology) while the FLOP means a half-precision FLOP in the tensor core. However, in the later description like Algorithm 1, one COP seems to be a general non-matrix operation without considering the problem size. In this case, one COP, for example, vectorized comparison, could include multiple FLOPs and seems to be unmatched by what is presented in Table 1.

3. **Incomplete Evaluation.** Based on what I've mentioned in 1, the evaluation part lacks a fair baseline. I can understand the advantage of recall of the proposed algorithm. However, it is hard to understand whether the performance gain over GPU comes from the algorithm itself or the higher efficiency of TPU. To isolate this factor, the authors should either provide the result of the proposed algorithm implemented on GPUs or the baseline algorithms implemented on TPU.
questions: 1. Is the proposed algorithm specialized for TPU? If it is, what special architectural support from TPU does it utilize?

2. How do you define **ONE** COP? For example, are vectorized comparisons between two 256 arrays and the comparison between two 16 arrays each counted as one COP? If not, how should I interpret the data presented in Table 1?

3. Could you provide a direct comparison of the proposed algorithm on the same hardware platform?
limitations: The authors discussed the limitation of the current implementation.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This paper studies using TPU for approximate nearest neighbor search. With careful analysis, it finds that search performance is bound by memory access and coefficient-wise operation instead of distance computation. Thus, it proposes to conduct partial reduce to reduce memory access such that search can reach peak FLOPs. Experiment results show that the proposed solution outperforms existing ones in both QPS and recall.
strengths_and_weaknesses: Strength

1.	The extended roofline analysis is interesting and shows the bottleneck of nearest neighbor search with TPU.

2.	The proposed partial reduce scheme reduces the amount of memory access and comes with theoretical analysis.

3.	The experiment results show that the proposed solution has good performance.


Weakness

1.	It would improve the paper if the authors can give a brief introduction about the characteristics of GPU and TPU. Are TPUs widely available as a commodity hardware?

2.	The proposed solution may also work for GPUs as GPUs are also limited by memory access from Table 1. Moreover, as shown in Figure 2, the proposed solution is still bound by coefficient-wise operations for Euclidean distance search. Any thoughts on reducing the number of coefficient-wise operations?

3.	The experiment study is far from the NeurIPS standard. Does the method still outperform the partial search methods (e.g., IVF) when the dataset is large (e.g., SIFT10M or 100M)? How does the proposed method perform when changing the window size W?

questions: See weakness.
limitations: Yes
review_title: null

Meta-review: metareview: The authors design an efficient implementation of nearest neighbor search on a TPU accelerator unit. The implementation is motivated by a refined roofline performance model that takes into account the memory and instruction bottlenecks that are found to be significant and that are not typically optimized for. Empirical results demonstrate that the proposed TPU solver outperforms state-of-the art GPU solvers. The package is available on Tensorflow.

The reviewers agree that the paper is well written, well structured and the proposed method can have significant practical impact. 
Some concerns regarding the evaluation came up in the reviews but the additional experiments provided could address most of these concerns. What remains is a question on whether the performance gain over GPU comes from the algorithmic optimization itself or the higher efficiency of the TPU, and whether the algorithmic optimization would be similarly effective on other accelerators. The reviewers agree that performing such an analysis is outside the scope of this work. However, I want to encourage the authors to incorporate additional discussion to help the reader understand what parts of the work are specific to TPUs.

Overall this paper represents a well executed piece of work at the intersection between algorithm design and systems with an open source package that is available to the community. I recommend acceptance.

confidence: Certain
Decision: Accept