Speed up the implementation without for

username: jetcai1900
cooked: <p>I have a tensor A, a tensor B and a index tensor C<br>
for example:<br>
A = [[1, 2, 3], [4, 5, 6]]<br>
B = [[10, 20, 30], [40, 50, 60]]<br>
C = [0, 0]</p>
<p>I would like to use B to update A,<br>
A[C[0]] += B[0]<br>
A[C[1]] += B[1]</p>
<p>So far, I am using for loop to achieve this, which is very slow. The for loop snippet is<br>
for i, c_i in enumerate (C ):<br>
A[c_i] += B[i]</p>
<p>Any suggestions on how to get rid of for loop or speed up the implementation. Thanks a lot. The above is just a simplified version and in practice, the A and B tensors are very large.</p>

username: KFrank
cooked: <p>Hello Jet!</p>
<aside class="quote no-group quote-modified" data-username="jetcai1900" data-post="1" data-topic="79851" data-full="true">
<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/j/a6a055/48.png" class="avatar"> jetcai1900:</div>
<blockquote>
<p>I would like to use B to update A,<br>
A[C[0]] += B[0]<br>
A[C[1]] += B[1]</p>
<p>So far, I am using for loop to achieve this, which is very slow.</p>
<p>Any suggestions on how to get rid of for loop or speed up the implementation.</p>
</blockquote>
</aside>
<p>I am not aware of any way to do this in pytorch without a loop.</p>
<p>Furthermore, I doubt that there is such a way.  My reasoning runs<br>
as follows:</p>
<p>The purpose of using “tensor” operations (whatever those are) is<br>
to be able to do things “in parallel” (whatever that may mean).  But,<br>
as you describe your desired calculation, it is inherently sequential.<br>
Also, I don’t think it is possible to describe a useful semantics for<br>
this type of operation that isn’t inherently sequential.</p>
<p>If all of the indices in your index tensor <code>C</code> were guaranteed to be<br>
distinct, then one could parallelize the operations.  But your specific<br>
example, <code>C = [0, 0]</code>, provides a good counter-example.</p>
<p>Noting that <code>C[0] = C[1] = 0</code>, we can rewrite:</p>
<pre data-code-wrap="python"><code class="lang-python">A[C[0]] += B[0]
A[C[1]] += B[1]
</code></pre>
<p>as:</p>
<pre data-code-wrap="python"><code class="lang-python">A0tmpa = A[0]
A0tmpa += B[0]
A[0] = A0tmpa
A0tmpb = A[0]   # because C[1] = C[0] = 0
A0tmpb += B[1]
A[0] = A0tmpb
</code></pre>
<p>But suppose – in order to parallelize it – we wrote is as:</p>
<pre data-code-wrap="python"><code class="lang-python">A0tmpa = A[0]
A0tmpb = A[0]   # because C[1] = C[0] = 0
A0tmpa += B[0]
A0tmpb += B[1]
A[0] = A0tmpa
A[0] = A0tmpb
</code></pre>
<p>In this case, you would get a different result, and probably not the<br>
result you would expect or want.</p>
<p>Or just switching two lines would screw things up:</p>
<pre data-code-wrap="python"><code class="lang-python">A0tmpa = A[0]
A0tmpa += B[0]
# A[0] = A0tmpa
# A0tmpb = A[0]   # because C[1] = C[0] = 0
A0tmpb = A[0]   # because C[1] = C[0] = 0
A[0] = A0tmpa
A0tmpb += B[1]
A[0] = A0tmpb
</code></pre>
<p>So the <em>sequence</em> of operations (which, if enforced, will break<br>
parallelization) matters.</p>
<p>Here’s a script that shows – in a somewhat contrived way by<br>
using floating-point round-off – that even if you perform each of<br>
<code>A[C[0]] += B[0]</code> and <code>A[C[1]] += B[1]</code> “atomically,” the order<br>
of the two atomic operations matters:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

delt = 16777216.0  # to get some round-off error

A = torch.FloatTensor ([[1, 2, 3], [4, 5, 6]])
D = A.clone()
B = torch.FloatTensor ([[delt, delt, delt], [-delt, -delt, -delt]])
C = torch.LongTensor ([0, 0])

A
A[C[0]] += B[0]
A
A[C[1]] += B[1]
A

D
D[C[1]] += B[1]
D
D[C[0]] += B[0]
D
</code></pre>
<p>And here is the output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'0.3.0b0+591e73e'
&gt;&gt;&gt;
&gt;&gt;&gt; delt = 16777216.0  # to get some round-off error
&gt;&gt;&gt;
&gt;&gt;&gt; A = torch.FloatTensor ([[1, 2, 3], [4, 5, 6]])
&gt;&gt;&gt; D = A.clone()
&gt;&gt;&gt; B = torch.FloatTensor ([[delt, delt, delt], [-delt, -delt, -delt]])
&gt;&gt;&gt; C = torch.LongTensor ([0, 0])
&gt;&gt;&gt;
&gt;&gt;&gt; A

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]

&gt;&gt;&gt; A[C[0]] += B[0]
&gt;&gt;&gt; A

 1.6777e+07  1.6777e+07  1.6777e+07
 4.0000e+00  5.0000e+00  6.0000e+00
[torch.FloatTensor of size 2x3]

&gt;&gt;&gt; A[C[1]] += B[1]
&gt;&gt;&gt; A

 0  2  4
 4  5  6
[torch.FloatTensor of size 2x3]

&gt;&gt;&gt;
&gt;&gt;&gt; D

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]

&gt;&gt;&gt; D[C[1]] += B[1]
&gt;&gt;&gt; D

-1.6777e+07 -1.6777e+07 -1.6777e+07
 4.0000e+00  5.0000e+00  6.0000e+00
[torch.FloatTensor of size 2x3]

&gt;&gt;&gt; D[C[0]] += B[0]
&gt;&gt;&gt; D

 1  2  3
 4  5  6
[torch.FloatTensor of size 2x3]
</code></pre>
<p>And indeed, the results of iterating over the elements of <code>C</code> in<br>
different orders differs.</p>
<p>One could imagine adding a function to pytorch that does what<br>
you want so that you wouldn’t have to write an explicit loop.  But<br>
because of the relatively strict sequencing required to give this<br>
operation useful semantics, such a pytorch function would pretty<br>
much have to be using a loop (or loop equivalent) under the hood,<br>
so it’s unlikely doing so would buy you any significant speed-up.</p>
<p>Best.</p>
<p>K. Frank</p>

username: googlebot
cooked: <aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="79851">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>One could imagine adding a function to pytorch that does what<br>
you want so that you wouldn’t have to write an explicit loop. But<br>
because of the relatively strict sequencing required to give this<br>
operation useful semantics, such a pytorch function would pretty<br>
much have to be using a loop (or loop equivalent) under the hood,<br>
so it’s unlikely doing so would buy you any significant speed-up.</p>
</blockquote>
</aside>
<p>… And this operation is called scatter_add, it uses atomicAdd on cuda, on cpu it is a loop, but c++ loop != python loop.</p>
<pre><code class="lang-auto">A.scatter_add(0, C.view(2,-1).expand(-1,3), B)
</code></pre>
<p>PS about your rounding example, many cuda ops are non-deterministic, so we live with it, trying not to sum billions with one in float32.</p>

username: KFrank
cooked: <p>Hello Alex!</p>
<aside class="quote no-group quote-modified" data-username="googlebot" data-post="3" data-topic="79851" data-full="true">
<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/g/ba8739/48.png" class="avatar"> googlebot:</div>
<blockquote>
<aside class="quote no-group" data-username="KFrank" data-post="2" data-topic="79851">
<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/k/ecb155/48.png" class="avatar"> KFrank:</div>
<blockquote>
<p>One could imagine adding a function to pytorch that does what<br>
you want</p>
</blockquote>
</aside>
<p>… And this operation is called scatter_add, it uses atomicAdd on cuda, on cpu it is a loop, but c++ loop != python loop.</p>
</blockquote>
</aside>
<p>Thanks for that – you’re absolutely right.  And I agree that the use<br>
of “atomicAdd” (assuming it means what it sounds like) eliminates<br>
the need for explicit sequencing (such as with a <code>for</code> loop).  Multiple<br>
cuda operations can attempt to update a given element of <code>A</code><br>
“simultaneously,” and they should be able to do so on a first-come,<br>
first-served basis, without any conflicts.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">A.scatter_add(0, C.view(2,-1).expand(-1,3), B)
</code></pre>
</blockquote>
</aside>
<p>Sadly I can’t try this, as I’m stuck several layers down in the pytorch<br>
archeological strata with version 0.3.0, which lacks this function.</p>
<aside class="quote no-group">
<blockquote>
<p>PS about your rounding example, many cuda ops are non-deterministic, so we live with it, trying not to sum billions with one in float32.</p>
</blockquote>
</aside>
<p>Yes, that’s fair.  The purpose of my admittedly-contrived example was<br>
to show that when <code>C</code> contains duplicate indices there is an ordering<br>
issue, at least in principle.</p>
<p>As an aside, for <code>scatter()</code>, where the ordering issue is fully real (not<br>
just differing in floating-point round-off), the documentation warns of<br>
this:</p>
<blockquote>
<p>Moreover, as for <code>gather()</code>, the values of <code>index</code> must be between <code>0</code> and <code>self.size(dim) - 1</code> inclusive, and all values in a row along the specified dimension <code>dim</code> must be unique.</p>
</blockquote>
<p>(Although – at least in version 0.3.0 – unique indices are not enforced<br>
nor checked for.)</p>
<p>Best.</p>
<p>K. Frank</p>