Why is `torch.pdist` so slow?

aRI0U: <p>While <code>torch.pdist</code> being a native C++ function from pytorch, seems that it’s slower to use it than to just use <code>torch.cdist</code> and mask half of the matrix.</p>
<p>Here is the snippet code:</p>
<pre><code class="lang-auto">import torch

print(torch.__version__)
z1 = torch.randn(512, 512)

for device in ["cpu", "cuda"]:
    print(device)
    z1 = z1.to(torch.device(device))

    mask = torch.triu(torch.ones(z1.size(0), z1.size(0), dtype=bool, device=z1.device), diagonal=1)
    torch.testing.assert_close(torch.cdist(z1, z1)[mask], torch.pdist(z1))

    %timeit torch.pdist(z1)
    %timeit torch.cdist(z1, z1)[torch.triu(torch.ones(z1.size(0), z1.size(0), dtype=bool, device=z1.device), diagonal=1)]
    %timeit torch.cdist(z1, z1)[mask]
</code></pre>
<p>and here are the output when running in a colab notebook:</p>
<pre><code class="lang-auto">2.0.0+cu118
cpu
12.5 ms ± 1.29 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)
5.12 ms ± 111 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
4.86 ms ± 231 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
cuda
3.68 ms ± 21.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
320 µs ± 81.7 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
245 µs ± 5.67 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
</code></pre>
<p>Then what is the point of <code>torch.pdist</code>?</p>

ptrblck: <p>Your profiling is invalid for the GPU as you are not synchronizing the code and are profiling the dispatching and kernel launches at best or random syncs in the worst case.<br>
Besides that you might also want to profile different shapes including warmup iterations.</p>

aRI0U: <p>Hi,<br>
Thanks a lot for your answer. I’m not so familiar with profiling and honestly I did not fully understand all you said, but still it seems that:</p>
<ul>
<li>on CPU at least my profiling looks fair and the speed of <code>pdist</code> is twice slower than <code>cdist</code> whereas there are technically twice less operations in <code>pdist</code> than <code>cdist</code> and no CUDA parallelization black magic should be involved</li>
<li>on GPU my measures might not be super precise but still the difference of speed is an order of magnitude so it’s definitely enough to observe that <code>pdist</code> is slower than <code>cdist</code>
</li>
<li>I reported the results only for one shape for simplicity and readability but observed similar behavior with different shapes. I chose (512*512) because it’s a reasonable size for a real-world usage, moreover pdist should be better than cdist particularly when the size of the tensor increases</li>
</ul>
<p>Overall, I’m not interested in precisely benchmarking <code>pdist</code> and <code>cdist</code>, I just wanted to point out that the implementation of <code>torch.pdist</code> looks extremely suboptimal, hence my minimal example. Do you know how <code>pdist</code> is implemented and why the gap between pdist and cdist is so huge?</p>

ptrblck: <aside class="quote no-group" data-username="aRI0U" data-post="3" data-topic="178947">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ari0u/48/24330_2.png" class="avatar"> aRI0U:</div>
<blockquote>
<p>Overall, I’m not interested in precisely benchmarking <code>pdist</code> and <code>cdist</code>,</p>
</blockquote>
</aside>
<p>I understand that you are not interested in <em>precisely</em> profiling workloads, but as already pointed out your profiling is invalid.</p>
<aside class="quote no-group" data-username="aRI0U" data-post="3" data-topic="178947">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ari0u/48/24330_2.png" class="avatar"> aRI0U:</div>
<blockquote>
<p>I just wanted to point out that the implementation of <code>torch.pdist</code> looks extremely suboptimal, hence my minimal example. Do you know how <code>pdist</code> is implemented and why the gap between pdist and cdist is so huge?</p>
</blockquote>
</aside>
<p>You could see the implementation of <code>pdist</code> <a href="https://github.com/pytorch/pytorch/blob/96487d0d1f638ae9c665dcd3d4b79e7e2a976f59/aten/src/ATen/native/cuda/DistanceKernel.cu#L96">here</a> and <code>cdist</code> <a href="https://github.com/pytorch/pytorch/blob/96487d0d1f638ae9c665dcd3d4b79e7e2a976f59/aten/src/ATen/native/cuda/DistanceKernel.cu#L188">here</a>.</p>