Scatter reduce optimization on pytorch2.2

zht-github: <p>Hi, I have a question about the significant performance boost on <em><strong>scatter_reduce</strong></em> ops since the version of pytorch 2.2.</p>
<p>I used to develop on <strong>torch2.0</strong>, and I observe <strong>significant performance decay when unique value count in my index tensor is very low</strong>, that is, for example, 10M elements been scattered and reduced into a small amount of groups (for example 10 groups).</p>
<p>I post a <a>question</a> before. <strong>scatter-and-reducing 10M elements into 10 group will cost me up to 95.84s!</strong></p>
<hr>
<p><strong>Experiemnt</strong> :<br>
I conduct a experiment to compare the performance of scatter reduce in pytorch2.0 and pytorch2.2 respectively.<br>
<strong>device</strong> :Nvidia 3090<br>
<strong>code</strong>: we test the performance difference under different <em><strong>group_number</strong></em>, that is the unique value count in the index tensor.<br>
<strong>test_scale:</strong> we set the length of the tensor as <strong>10M</strong>, and the group number is 10M, 1M ,100K, 10K, 1K…<br>
<strong>test distribution:</strong> I also test whether the distribution of the element in the index will influence the performance.<br>
In each group number, I test the <strong>consecutive distribution</strong>, that is all identical element in the index tensor are placed consecutively like  [0,0,0…,1,1,1…,n,n,n…],versus  <strong>random distribution</strong> that the index tensor and the source tensor are shuffled by randomly generated permute index.</p>
<pre><code class="lang-auto">import torch
import time
scale = int(1e7)
group_num = scale

data = torch.randint(0,10,size=((scale),),device= 'cuda',dtype = torch.int64)
permute_idx = torch.randperm(scale,device = 'cuda')
while group_num &gt;=10:
    # generate idx
    repeat_num = int(scale/group_num)
    # ordered idx like [0,0,0...,1,1,1....,n,n,n....], identical element are placed consecutively 
    ordered_idx = torch.arange(group_num,device = 'cuda').view(-1,1).repeat(1,repeat_num).view(-1)
    # shuffle ordered idx and data according to permute_idx
    rand_idx = ordered_idx[permute_idx]
    rand_data = data[permute_idx]
    # test performance on rand data
    # test for 10 round 
    res_list = []
    for i in range(10):
        res = torch.zeros(group_num,device= 'cuda',dtype=data.dtype)
        start = time.time()
        res.scatter_reduce_(src = rand_data,index = rand_idx, dim = 0,reduce = 'amin')
        torch.cuda.synchronize()
        end = time.time()
        res_list.append(end-start)
    #use the average time of the last 3 round as the result
    val1 = sum(res_list[-3:])/3
    print(val1)
    # print(res)
    #=======================test for consecutive(ordered) version
    res_list = []
    # test performance on rand data
    # test for 10 round 
    for i in range(10):
        res = torch.zeros(group_num,device= 'cuda',dtype=data.dtype)
        start = time.time()
        res.scatter_reduce_(src = data,index = ordered_idx, dim = 0,reduce = 'amin')
        torch.cuda.synchronize()
        end = time.time()
        res_list.append(end-start)
    # use the average time of the last 3 round as the result
    val2 = sum(res_list[-3:])/3
    print(val2)
    # print(res)
    print(f'========{group_num}==========')
    group_num = int(group_num/10)

</code></pre>
<hr>
<p><strong>Test result</strong><br>
<strong>torch 2,0:</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/3/7325ab6ae12a18152117da7bb9ca28c35154ebdd.png" data-download-href="https://discuss.pytorch.org/uploads/default/7325ab6ae12a18152117da7bb9ca28c35154ebdd" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/3/7325ab6ae12a18152117da7bb9ca28c35154ebdd.png" alt="image" data-base62-sha1="gqDCbjtTOJ97GKuMNIB0QuKZJVz" width="167" height="375" data-dominant-color="1D1C1D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">262×586 16.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
<strong>torch2.2</strong>:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/5/c5560936a37aed7e7305f485fbd815d2d823d24e.png" data-download-href="https://discuss.pytorch.org/uploads/default/c5560936a37aed7e7305f485fbd815d2d823d24e" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/5/c5560936a37aed7e7305f485fbd815d2d823d24e.png" alt="image" data-base62-sha1="s9IomrQQw0LgsJIQDcJBrRN86bs" width="180" height="375" data-dominant-color="E3E2E3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">242×501 16 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p><strong>We can observe significant performance difference between the execution of the same code indifferent pytorch version.</strong><br>
<strong>observations in In torch2.0:</strong></p>
<ol>
<li>the performance of <strong>scatter_reduce</strong> is reduced according with the decrease of group number</li>
<li>and the <strong>consecutive version</strong>(lower) suffered from more severe performance drop than  the <strong>random version</strong>(upper).</li>
<li>When the group number equals 10, <strong>random version</strong> in torch2.0 will cost <strong>16s</strong> and <strong>consecutive version</strong> will cost <strong>96s</strong>!.</li>
</ol>
<p>Pytorch use <strong>atomic ops</strong> to implement scatter reduce operations. Therefore I tried to use <strong>memory contention</strong> to explain this observation since we are updating the shared memory under large parallelism scenario, when the group number is low, there might exists high risk that thousands of threads competing the access right of a particular element thereby causing a sharp performance degrade. And the worse performance of consecutive version can support my explaination .</p>
<hr>
<p><strong>observations in In torch2.2:</strong></p>
<p>However the <strong>observation above in torch2.0  disappeared in torch2.2.</strong> They can run perfectly efficient on low group number scenarios.</p>
<hr>
<p>So I am wondering what is the main optimization pytorch did to overcome this problem. And what strategy did the new version applied to get rid of the memory contention problem!</p>

ptrblck: <p>I cannot reproduce a slow execution in <code>torch==2.0.0+cu118</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">2.0.0+cu118
0.004313866297403972
0.00040086110432942707
========10000000==========
0.0027761459350585938
0.00029087066650390625
========1000000==========
0.0005593299865722656
0.0003363291422526042
========100000==========
0.0007781187693277994
0.001141945521036784
========10000==========
0.0008450349171956381
0.004339853922526042
========1000==========
0.006618897120157878
0.006076892217000325
========100==========
0.0068082014719645185
0.006005446116129558
========10==========
</code></pre>

zht-github: <p>Sorry for my late reply <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"> I felt really down in the last week and I forgot to check my email. I am really sorry for that.</p>
<p>I post wrong code, and here is the new code.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import time
scale = int(1e7)
group_num = scale

data = torch.randint(0,10,size=((scale),),device= 'cuda')
permute_idx = torch.randperm(scale,device = 'cuda')
while group_num &gt;=10:
    # generate idx
    repeat_num = int(scale/group_num)
    # ordered idx like [0,0,0...,1,1,1....,n,n,n....], identical element are placed consecutively 
    ordered_idx = torch.arange(group_num,device = 'cuda').view(-1,1).repeat(1,repeat_num).view(-1)
    # shuffle ordered idx and data according to permute_idx
    rand_idx = ordered_idx[permute_idx]
    rand_data = data[permute_idx]
    # test performance on rand data
    # test for 10 round 
    res_list = []
    for i in range(10):
        res = torch.zeros(group_num,device= 'cuda',dtype=data.dtype)
        start = time.time()
        res.scatter_add_(src = rand_data,index = rand_idx, dim = 0)
        torch.cuda.synchronize()
        end = time.time()
        res_list.append(end-start)
    #use the average time of the last 3 round as the result
    val1 = sum(res_list[-3:])/3
    print(val1)
    #=======================test for consecutive(ordered) version
    res_list = []
    for i in range(10):
        res = torch.zeros(group_num,device= 'cuda',dtype=data.dtype)
        start = time.time()
        res.scatter_add_(src = data,index = ordered_idx, dim = 0)
        torch.cuda.synchronize()
        end = time.time()
        res_list.append(end-start)
    #use the average time of the last 3 round as the result
    val2 = sum(res_list[-3:])/3
    print(val2)
    # print(res)
    print(f'========{group_num}==========')
    group_num = int(group_num/10)
</code></pre>
<p>If you are convenient please check this code.  It do behave badly on torch2.0+cu118 during my recent test.<br>
Thanks for your kindness and patience!</p>