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