What's the time complexity of Tensor.topk?

username: AOZMH
cooked: <p>The following code finds the top-k elements of a tensor.</p>
<pre><code class="lang-auto">import torch
n, k = 100, 5
a = torch.randn(n)
b = a.topk(k)
</code></pre>
<p>I’m wondering what’s the <strong>time complexity</strong> of that operation.<br>
E.g., a simple sorting and [:k] implementation would result in <strong>O(nlogn)</strong>, while a quick-sort-styled partition would be <strong>O(n+k)</strong>, also, when a heap is used the complexity would be <strong>O(n+klogn)</strong>. Clearly, if the first algorithm is used in <em>topk</em> function we may lose some performance and need to manually choose other better implementations when finding top-k elements.<br>
Thanks for any help!</p>

username: Arvid_Frydenlund
cooked: <p>I’m also looking to know the time complexity for the topk operation, especially on the GPU.  Thanks!</p>

username: Eta_C
cooked: <p>I found some clues</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/95f4cd0ba99b04a82d711a47033e16a69f8f1b5e/aten/src/ATen/native/cuda/TensorTopK.cpp#L32-L52">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/95f4cd0ba99b04a82d711a47033e16a69f8f1b5e/aten/src/ATen/native/cuda/TensorTopK.cpp#L32-L52" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/95f4cd0ba99b04a82d711a47033e16a69f8f1b5e/aten/src/ATen/native/cuda/TensorTopK.cpp#L32-L52" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/95f4cd0ba99b04a82d711a47033e16a69f8f1b5e/aten/src/ATen/native/cuda/TensorTopK.cpp#L32-L52</a></h4>



  <pre class="onebox">    <code class="lang-cpp">
      <ol class="start lines" start="32" style="counter-reset: li-counter 31 ;">
          <li>TORCH_IMPL_FUNC(topk_out_cuda)</li>
          <li>  (const Tensor&amp; self,</li>
          <li>   int64_t k, int64_t dim, bool largest, bool sorted,</li>
          <li>   const Tensor&amp; values,</li>
          <li>   const Tensor&amp; indices) {</li>
          <li>  TensorArg topK_arg{values, "topK", 1}, indices_arg{indices, "indices", 2}, input_arg{self, "self", 3};</li>
          <li>  checkAllSameGPU(__func__, {topK_arg, indices_arg, input_arg});</li>
          <li>
          </li>
<li>  dim = at::maybe_wrap_dim(dim, self);</li>
          <li>
          </li>
<li>  if (should_use_sort(self, dim)) {</li>
          <li>    topk_out_with_sort(self, k, dim, largest, values, indices);</li>
          <li>    return;</li>
          <li>  }</li>
          <li>
          </li>
<li>  // If k is 0 the result is an empty tensor, so we don't need to launch a kernel.</li>
          <li>  if (k == 0) {</li>
          <li>    return;</li>
          <li>  }</li>
          <li>
          </li>
<li>  launch_gather_topk_kernel(self, k, dim, largest, sorted, values, indices);</li>
      </ol>
    </code>
  </pre>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/4262c8913c2bddb8d91565888b4871790301faba/aten/src/ATen/native/cuda/TensorTopK.cu#L92-L121">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/4262c8913c2bddb8d91565888b4871790301faba/aten/src/ATen/native/cuda/TensorTopK.cu#L92-L121" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/4262c8913c2bddb8d91565888b4871790301faba/aten/src/ATen/native/cuda/TensorTopK.cu#L92-L121" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/4262c8913c2bddb8d91565888b4871790301faba/aten/src/ATen/native/cuda/TensorTopK.cu#L92-L121</a></h4>



  <pre class="onebox">    <code class="lang-cu">
      <ol class="start lines" start="92" style="counter-reset: li-counter 91 ;">
          <li>for (IndexType i = threadIdx.x; i &lt; numIterations; i += blockDim.x) {</li>
          <li>  bool inRange = (i &lt; inputSliceSize);</li>
          <li>  T v =</li>
          <li>    inRange ? doLdg(&amp;inputSliceStart[i * inputWithinSliceStride]) : static_cast&lt;T&gt;(0);</li>
          <li>  const auto convertedV = at::native::TopKTypeConfig&lt;T&gt;::convert(v);</li>
          <li>  bool hasTopK;</li>
          <li>  if (Order) {</li>
          <li>    hasTopK = inRange &amp;&amp; (convertedV &gt; topKConverted);</li>
          <li>  } else {</li>
          <li>    hasTopK = inRange &amp;&amp; (convertedV &lt; topKConverted);</li>
          <li>  }</li>
          <li>
          </li>
<li>  int index;</li>
          <li>  int carry;</li>
          <li>  at::cuda::exclusiveBinaryPrefixScan&lt;int, true&gt;(</li>
          <li>      smem, hasTopK, &amp;index, &amp;carry, AddOp&lt;int&gt;());</li>
          <li>
          </li>
<li>  if (hasTopK) {</li>
          <li>    int writeIndex = writeIndexStart + index;</li>
          <li>    CUDA_KERNEL_ASSERT(writeIndex &lt; outputSliceSize);</li>
      </ol>
    </code>
  </pre>

  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/4262c8913c2bddb8d91565888b4871790301faba/aten/src/ATen/native/cuda/TensorTopK.cu#L92-L121" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: AOZMH
cooked: <p>Thanks for the efforts! If these are indeed the backends used by the topk function, it points to an O(n+k) implementation.</p>

username: LeoLiu1998
cooked: <p>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)?</p>