What's the time complexity of Tensor.topk?

AOZMH · April 12, 2021, 3:20am

The following code finds the top-k elements of a tensor.

import torch
n, k = 100, 5
a = torch.randn(n)
b = a.topk(k)

I’m wondering what’s the time complexity of that operation.
E.g., a simple sorting and [:k] implementation would result in O(nlogn), while a quick-sort-styled partition would be O(n+k), also, when a heap is used the complexity would be O(n+klogn). Clearly, if the first algorithm is used in topk function we may lose some performance and need to manually choose other better implementations when finding top-k elements.
Thanks for any help!

Arvid_Frydenlund · November 21, 2021, 8:53pm

I’m also looking to know the time complexity for the topk operation, especially on the GPU. Thanks!

Eta_C · November 22, 2021, 3:12am

I found some clues

github.com

pytorch/pytorch/blob/95f4cd0ba99b04a82d711a47033e16a69f8f1b5e/aten/src/ATen/native/cuda/TensorTopK.cpp#L32-L52

    
      
          TORCH_IMPL_FUNC(topk_out_cuda)
            (const Tensor& self,
             int64_t k, int64_t dim, bool largest, bool sorted,
             const Tensor& values,
             const Tensor& indices) {
            TensorArg topK_arg{values, "topK", 1}, indices_arg{indices, "indices", 2}, input_arg{self, "self", 3};
            checkAllSameGPU(__func__, {topK_arg, indices_arg, input_arg});
          
          
  dim = at::maybe_wrap_dim(dim, self);
          
          
  if (should_use_sort(self, dim)) {
              topk_out_with_sort(self, k, dim, largest, values, indices);
              return;
            }
          
          
  // If k is 0 the result is an empty tensor, so we don't need to launch a kernel.
            if (k == 0) {
              return;
            }
          
          
  launch_gather_topk_kernel(self, k, dim, largest, sorted, values, indices);

github.com

pytorch/pytorch/blob/4262c8913c2bddb8d91565888b4871790301faba/aten/src/ATen/native/cuda/TensorTopK.cu#L92-L121

    
      
          for (IndexType i = threadIdx.x; i < numIterations; i += blockDim.x) {
            bool inRange = (i < inputSliceSize);
            T v =
              inRange ? doLdg(&inputSliceStart[i * inputWithinSliceStride]) : static_cast<T>(0);
            const auto convertedV = at::native::TopKTypeConfig<T>::convert(v);
            bool hasTopK;
            if (Order) {
              hasTopK = inRange && (convertedV > topKConverted);
            } else {
              hasTopK = inRange && (convertedV < topKConverted);
            }
          
          
  int index;
            int carry;
            at::cuda::exclusiveBinaryPrefixScan<int, true>(
                smem, hasTopK, &index, &carry, AddOp<int>());
          
          
  if (hasTopK) {
              int writeIndex = writeIndexStart + index;
              CUDA_KERNEL_ASSERT(writeIndex < outputSliceSize);

This file has been truncated. show original

AOZMH · November 22, 2021, 3:47am

Thanks for the efforts! If these are indeed the backends used by the topk function, it points to an O(n+k) implementation.

LeoLiu1998 · February 28, 2022, 9:46pm

Thanks for the time complexity. Given the time complexity being O(n + k), do you know how I could estimate the FLOPs of topk? (do I just plug in n and k)?