Speedup torch.argsort, torch.sort. runtime and sorting axis. how to?

username: sbelharbi
cooked: <p>hi,<br>
i have a large float cuda tensor v = (32, 6, 59536).<br>
applying <code>torch.argsort(v, dim=1, descending=True)</code> takes 14ms. (<a href="https://pytorch.org/docs/stable/generated/torch.argsort.html#torch.argsort" rel="noopener nofollow ugc">argsort</a>)<br>
it is called twice. so, it weights in term of runtime.<br>
is there any way to speedup this op?</p>
<p>the same questions goes for <code>torch.sort</code>.</p>
<p>here is a weird behavior: runtime of <code>torch.sort</code> depends on the sorting axis!!!</p>
<pre><code class="lang-python">import numpy as np
import torch


def main(axis=-1):
    seed = 0
    torch.manual_seed(seed)
    np.random.seed(seed)

    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    cpu = np.random.rand(32, 6, 59536).astype(np.float32)
    gpu = torch.tensor(cpu).to(device='cuda:0')

    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()

    np.sort(cpu, axis=axis)[1]

    torch.cuda.synchronize()
    end_event.record()
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print('time cpu: {}'.format(elapsed_time_ms))

    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()

    torch.sort(gpu, dim=axis, descending=True).values[1]

    torch.cuda.synchronize()
    end_event.record()
    torch.cuda.synchronize()
    elapsed_time_ms = start_event.elapsed_time(end_event)
    print('time gpu: {}'.format(elapsed_time_ms))


if __name__ == '__main__':

    for axis in [0, 1, 2]:
        print('*************** sorting axis={} *************'.format(axis))
        for i in range(4):
            print('run {}'.format(i))
            main(axis=axis)
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">*************** sorting axis=0 *************
run 0
time cpu: 272.53204345703125
time gpu: 6.794528007507324
run 1
time cpu: 268.1492614746094
time gpu: 5.733503818511963
run 2
time cpu: 268.2332458496094
time gpu: 5.738399982452393
run 3
time cpu: 271.6305236816406
time gpu: 5.7454400062561035
*************** sorting axis=1 *************
run 0
time cpu: 163.71519470214844
time gpu: 15.351776123046875
run 1
time cpu: 163.94012451171875
time gpu: 15.353407859802246
run 2
time cpu: 165.7310791015625
time gpu: 15.347135543823242
run 3
time cpu: 163.39231872558594
time gpu: 15.351840019226074
*************** sorting axis=2 *************
run 0
time cpu: 817.253173828125
time gpu: 6.221983909606934
run 1
time cpu: 820.2147827148438
time gpu: 5.75600004196167
run 2
time cpu: 817.254638671875
time gpu: 5.731488227844238
run 3
time cpu: 814.9271850585938
time gpu: 5.727200031280518

</code></pre>
<p>sorting using extreme axes seems way faster than sorting using inner axes, at least for pytorch. numpy behaves differently.<br>
any explanation?</p>
<p>thanks</p>

username: pritamdamania87
cooked: <aside class="quote no-group" data-username="sbelharbi" data-post="1" data-topic="131441">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/f475e1/48.png" class="avatar"> sbelharbi:</div>
<blockquote>
<p>i have a large float cuda tensor v = (32, 6, 59536).<br>
applying <code>torch.argsort(v, dim=1, descending=True)</code> takes 14ms. (<a href="https://pytorch.org/docs/stable/generated/torch.argsort.html#torch.argsort" rel="noopener nofollow ugc">argsort</a>)<br>
it is called twice. so, it weights in term of runtime.<br>
is there any way to speedup this op?</p>
</blockquote>
</aside>
<p>You could split the tensor along the other dimension and sort chunks in parallel on different GPUs and then concatenate them together later on a single GPU.</p>
<blockquote>
<p>sorting using extreme axes seems way faster than sorting using inner axes, at least for pytorch. numpy behaves differently.<br>
any explanation?</p>
</blockquote>
<p>This might have to do with contiguous data layout in case of some sorting axes. For example if you are sorting on the last axis, all the data is contiguous.</p>