Scatter add much slower when compiled

BeachHut: <p>I wish to add to a Tensor at specified indices. This can be achieved in PyTorch, with the Tensor reshaped to 1D for simplicity, using simple indexing (<code>u[x] += v</code>) or with <code>u.scatter_add_(0, x, v)</code>. The former is slow, but in eager mode the latter provides performance almost equivalent to C. However, if I use <code>torch.compile</code>, <code>scatter_add_</code> becomes substantially slower, especially when the number of additions is small. I am currently running on a CPU, using PyTorch 2.5.1.</p>
<p>Here is a demonstration:</p>
<p><code>scatter_add.c</code>:</p>
<pre><code class="lang-auto">#include &lt;stdint.h&gt;

void scatter_add(float *restrict const u, float const *restrict const v,
                 int64_t const *restrict const x, int64_t n) {
  for (int64_t i = 0; i &lt; n; ++i) u[x[i]] += v[i];
}
</code></pre>
<p>Compile with <code>gcc -O2 -march=native -fPIC -shared scatter_add.c -o scatter_add.so</code></p>
<p>Python code to implement Python approaches and run timing:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import timeit
from ctypes import CDLL, c_void_p, c_int64

dll = CDLL('./scatter_add.so')
dll.scatter_add.restype = None
dll.scatter_add.argtypes = [c_void_p, c_void_p, c_void_p, c_int64]


def scatter_add1(u, v, x):
    u[x] += v


def scatter_add2(u, v, x):
    u.scatter_add_(0, x, v)


number = 1000
repeat = 10

times = []  # 2d array of runtimes: [p, implementation index]

for ns in [10**p for p in range(6)]:  # ns = number of adds
    u0 = torch.randn(100000)  # Tensor to add to
    x = torch.randperm(len(u0))[:ns]  # indexes to add to
    v = torch.randn(ns)  # values to add
    ns_times = []  # timings for this "ns": [implementation index]

    # C implementation
    u = u0.clone()  # Clone u0 so same for all implementations
    ns_times.append(
        min(
            timeit.repeat(lambda: dll.scatter_add(u.data_ptr(), v.data_ptr(),
                                                  x.data_ptr(), ns),
                          number=number,
                          repeat=repeat)))

    # Py1 implementation (scatter_add1)
    u = u0.clone()
    ns_times.append(
        min(
            timeit.repeat(lambda: scatter_add1(u, v, x),
                          number=number,
                          repeat=repeat)))

    # Py2 implementation (scatter_add2)
    u = u0.clone()
    ns_times.append(
        min(
            timeit.repeat(lambda: scatter_add2(u, v, x),
                          number=number,
                          repeat=repeat)))

    # compiled Py1
    scatter_add1c = torch.compile(scatter_add1)
    u = u0.clone()
    ns_times.append(
        min(
            timeit.repeat(lambda: scatter_add1c(u, v, x),
                          number=number,
                          repeat=repeat)))

    # compiled Py2
    scatter_add2c = torch.compile(scatter_add2)
    u = u0.clone()
    ns_times.append(
        min(
            timeit.repeat(lambda: scatter_add2c(u, v, x),
                          number=number,
                          repeat=repeat)))

    times.append(ns_times)
</code></pre>
<p>Plot results:</p>
<pre><code class="lang-auto">import matplotlib.pyplot as plt

timeslog10 = torch.tensor(times).log10().numpy()
plt.plot(timeslog10,
         label=['C', 'Py 1', 'Py 2', 'Py 1 compiled', 'Py 2 compiled'])
plt.legend()
plt.ylabel('Runtime (log10)')
plt.xlabel('Number of adds (log10)')
plt.savefig('runtimes.png')
plt.show()
</code></pre>
<p>Result:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/d/cd9ce479d441f90e0c45c2e04be845e49fc11c5b.png" data-download-href="https://discuss.pytorch.org/uploads/default/cd9ce479d441f90e0c45c2e04be845e49fc11c5b" title="runtimes"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/d/cd9ce479d441f90e0c45c2e04be845e49fc11c5b.png" alt="runtimes" data-base62-sha1="tkW1mIr2Pi1iZVSuRtEDUQP9Uvh" width="640" height="480"><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">runtimes</span><span class="informations">640×480 30 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>Notice that “Py 2” (green line, using <code>scatter_add_</code>) has almost the same performance as C in eager mode, but when compiled (purple line) is substantially slower.</p>
<p>Is there any way to improve the performance? I can guarantee that the indices are all unique, so atomic adds are not needed. I can also guarantee that the indices are all within bounds.</p>

BeachHut: <p>The problem appears to be caused by the implementation of the compiled version of <code>scatter_add_</code> looping (twice) over all elements of the target (see below), rather than only looping over the additions. This is not ideal for cases where there are a large number of elements in the target but only a few additions to it. Unfortunately that is exactly the situation that I have, and my model needs to call <code>scatter_add_</code> in this way many times in the forward pass. Can you suggest anything that might help me?</p>
<p>From the output code when using TORCH_COMPILE_DEBUG=1:</p>
<pre data-code-wrap="python"><code class="lang-python">cpp_fused_scatter_add_0 = async_compile.cpp_pybinding(['const float*', 'float*'], '''
extern "C"  void kernel(const float* in_ptr0,
                       float* out_ptr0)
{
    #pragma omp parallel num_threads(2)
    {
        int tid = omp_get_thread_num();
        {
            #pragma omp for
            for(int64_t x0=static_cast&lt;int64_t&gt;(0L); x0&lt;static_cast&lt;int64_t&gt;(100000L); x0+=static_cast&lt;int64_t&gt;(8L))
            {
                auto tmp0 = at::vec::Vectorized&lt;float&gt;::loadu(in_ptr0 + static_cast&lt;int64_t&gt;(x0), static_cast&lt;int64_t&gt;(8));
                tmp0.store(out_ptr0 + static_cast&lt;int64_t&gt;(x0));
            }
        }
    }
}
''')


cpp_fused_1 = async_compile.cpp_pybinding(['const float*', 'float*'], '''
extern "C"  void kernel(const float* in_ptr0,
                       float* out_ptr0)
{
    #pragma omp parallel num_threads(2)
    {
        int tid = omp_get_thread_num();
        {
            #pragma omp for
            for(int64_t x0=static_cast&lt;int64_t&gt;(0L); x0&lt;static_cast&lt;int64_t&gt;(100000L); x0+=static_cast&lt;int64_t&gt;(8L))
            {
                auto tmp0 = at::vec::Vectorized&lt;float&gt;::loadu(in_ptr0 + static_cast&lt;int64_t&gt;(x0), static_cast&lt;int64_t&gt;(8));
                tmp0.store(out_ptr0 + static_cast&lt;int64_t&gt;(x0));
            }
        }
    }
}
''')

def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    assert_size_stride(arg0_1, (100000, ), (1, ))
    assert_size_stride(arg1_1, (10, ), (1, ))
    assert_size_stride(arg2_1, (10, ), (1, ))
    buf0 = empty_strided_cpu((100000, ), (1, ), torch.float32)
    cpp_fused_scatter_add_0(arg0_1, buf0)
    aten.scatter_reduce_.two(buf0,0,arg1_1,arg2_1, reduce='sum', include_self=True)
    del arg1_1
    del arg2_1
    cpp_fused_1(buf0, arg0_1)
    del arg0_1
    return ()
</code></pre>