Very slow torch.median() compared to CuPy

arnaudvl: <p>To my surprise <code>torch.median()</code> is well over an order of magnitude slower than the equivalent <code>cupy.median()</code> on matrices of dimension 1000x1000 or more. It also gets worse as the matrix size grows. This is even more surprising given that unlike CuPy, PyTorch returns element <code>N // 2 - 1</code> of the sorted array as median for arrays with an even number of entries instead of the average of the 2 middle numbers, which should lead to a speed bump. On the other hand, it seems to be less memory hungry than CuPy (about 25% less for 10,000x10,000 matrices). Any insight what might cause this?</p>
<p>The following code can be used to reproduce the results:</p>
<pre><code class="lang-python">import cupy as cp
import numpy as np
from timeit import default_timer as timer
import torch

dim = [100, 1000, 5000, 10000]
n_runs = 10
n_warmup = 2
n_tot = n_runs + n_warmup

device = torch.device('cuda')

# timers
t_start = torch.cuda.Event(enable_timing=True)
t_end = torch.cuda.Event(enable_timing=True)

# PyTorch
for d in dim:
    t_pt = np.zeros(n_tot)
    for n in range(n_tot):
        np.random.seed(0)
        x = np.random.randn(*(d, d)).astype(np.float32)
        x_pt = torch.from_numpy(x).to(device)
        torch.cuda.synchronize()
        t_start.record()
        torch.median(x_pt)
        t_end.record()
        torch.cuda.synchronize()
        t_pt[n] = 1e-3 * t_start.elapsed_time(t_end)  # 1e-3 b/c prints in ms
    t_mean, t_std = t_pt[n_warmup:].mean(), t_pt[n_warmup].std()
    print(f'PyTorch: {d} x {d}: {t_mean:.5f} +- {t_std:.5f}')

# CuPy
for d in dim:
    t_cp = np.zeros(n_tot)
    for n in range(n_tot):
        np.random.seed(0)
        x = np.random.randn(*(d, d)).astype(np.float32)
        x_cp = cp.asarray(x)
        cp.cuda.Stream.null.synchronize()
        t_start = timer()
        cp.median(x_cp)
        cp.cuda.Stream.null.synchronize()
        t_cp[n] = timer() - t_start
    t_mean, t_std = t_cp[n_warmup:].mean(), t_cp[n_warmup].std()
    print(f'CuPy: {d} x {d}: {t_mean:.5f} +- {t_std:.5f}')
</code></pre>
<p>Which on a RTX2080Ti, CuPy 8.5.0 for CUDA 10.0 and torch 1.8.0 gives as mean (in sec) and standard deviation over 10 runs after the warmup phase:</p>
<pre><code class="lang-auto">PyTorch: 100 x 100: 0.00016 +- 0.00000
PyTorch: 1000 x 1000: 0.00186 +- 0.00000
PyTorch: 5000 x 5000: 0.05195 +- 0.00000
PyTorch: 10000 x 10000: 0.22688 +- 0.00000
CuPy: 100 x 100: 0.00024 +- 0.00000
CuPy: 1000 x 1000: 0.00041 +- 0.00000
CuPy: 5000 x 5000: 0.00470 +- 0.00000
CuPy: 10000 x 10000: 0.01795 +- 0.00000
</code></pre>

crcrpar: <p>Disclaimer: I didn’t thoroughly read the source code, I just git grep-ed and googled some information. So what I say might be wrong.</p>
<p>PyTorchv1.8’s <code>median</code> calls <code>sort</code> internally:</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/37c1f4a7fef115d719104e871d0cf39434aa9d56/aten/src/ATen/native/cuda/Sorting.cu#L393" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/37c1f4a7fef115d719104e871d0cf39434aa9d56/aten/src/ATen/native/cuda/Sorting.cu#L393" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/37c1f4a7fef115d719104e871d0cf39434aa9d56/aten/src/ATen/native/cuda/Sorting.cu#L393</a></h4>
<pre class="onebox"><code class="lang-cu"><ol class="start lines" start="383" style="counter-reset: li-counter 382 ;">
<li>  return std::forward_as_tuple(values, indices);</li>
<li>}</li>
<li>
<li>Tensor median_impl(const Tensor&amp; self, bool ignore_nan) {</li>
<li>  NoNamesGuard guard;</li>
<li>
<li>  int64_t size = self.numel();</li>
<li>  TORCH_CHECK(size &gt; 0, "median() input tensor cannot be empty");</li>
<li>
<li>  // Sort input tensor to efficiently query for median element</li>
<li class="selected">  Tensor sorted = std::get&lt;0&gt;(self.flatten().sort());</li>
<li>
<li>  if (!ignore_nan) {</li>
<li>    // For torch.median return either the middle element or nan (sorted as</li>
<li>    // largest) if there are any</li>
<li>    int64_t k = (size - 1) / 2;</li>
<li>    return at::where(sorted[-1].isnan(), sorted[-1], sorted[k]);</li>
<li>  } else {</li>
<li>    // For torch.nanmedian return the middle element among the non-nan values</li>
<li>    Tensor k = ((size - 1) - sorted.isnan().sum()) / 2;</li>
<li>    return sorted[k.toType(kLong)];</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


<p>On the other hand, CuPyv8.5’s <code>median</code> calls <code>partition</code> internally:</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/cupy/cupy/blob/6d91f956ee3c83d07fe0548268345d46fa319ccd/cupy/core/_routines_statistics.pyx#L410-L414" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/cupy/cupy/blob/6d91f956ee3c83d07fe0548268345d46fa319ccd/cupy/core/_routines_statistics.pyx#L410-L414" target="_blank" rel="noopener nofollow ugc">cupy/cupy/blob/6d91f956ee3c83d07fe0548268345d46fa319ccd/cupy/core/_routines_statistics.pyx#L410-L414</a></h4>
<pre class="onebox"><code class="lang-pyx"><ol class="start lines" start="410" style="counter-reset: li-counter 409 ;">
<li>if axis is None:</li>
<li>    part = part.ravel()</li>
<li>    part.partition(kth)</li>
<li>else:</li>
<li>    part.partition(kth, axis=axis)</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>In NumPy, <code>sort</code>'s order is <code>O(n^2)</code> in the worst case while <code>partition</code>, <code>O(n)</code> according to <a href="https://stackoverflow.com/a/43589598/8335699" rel="noopener nofollow ugc">https://stackoverflow.com/a/43589598/8335699</a>. And if this is true to PyTorch and CuPy, the gap you found can be possible.</p>