Poor Performance with CPU Threads (C++ extension)

username: its_sandy
cooked: <p>I’m trying to rewrite the following sparse multiplication operation as a C++ extension.</p>
<pre><code class="lang-auto">def sparse_mul(weights, bias, in_values, active_out_indices):
   # for the ith input in batch, produce output only for the nodes in active_out_indices[i]     
   # active_out_indices.size(1) is expected to be around 10% of weights.size(0)
   active_weights = weights[[active_out_indices]]
   active_bias = bias[[active_out_indices]]
   out_values = torch.bmm(active_weights, in_values.unsqueeze(-1)).squeeze(-1) + active_biases
   return out_values
</code></pre>
<p>I want to do this because the above python implementation performs unnecessary data copies due to the advanced indexing. The following is the C++ code I have as of now. It performs almost 50x slower (on my 4 core CPU) than the python implementation and also the regular dense multiplication. Can someone tell me what I’m doing wrong? Can I not expect a speedup here? I’m new to PyTorch C++ and am trying to follow the solution given in <a href="https://discuss.pytorch.org/t/using-at-parallel-for-in-a-custom-operator/82747/5">https://discuss.pytorch.org/t/using-at-parallel-for-in-a-custom-operator/82747/5</a></p>
<pre><code class="lang-auto">// sparse_mul.cpp
#include &lt;torch/extension.h&gt;
#include &lt;vector&gt;

#define _OPENMP
#include &lt;ATen/ParallelOpenMP.h&gt;

torch::Tensor mul(
        torch::Tensor in_values,
        torch::Tensor active_out_indices,
        torch::Tensor weights,
        torch::Tensor bias) {
    
    auto out_values = torch::empty(active_out_indices.sizes());
    auto active_out_indices_a = active_out_indices.accessor&lt;long,2&gt;();
    int batch_size = active_out_indices_a.size(0);
    int active_out_dim = active_out_indices_a.size(1);

    at::parallel_for(0, batch_size, 0, [&amp;](int64_t start, int64_t end) {
        for (int64_t i = start; i &lt; end; i++){
            torch::set_num_threads(1);
            auto in_values_i = in_values[i];
            for(int j=0; j&lt;active_out_dim; j++){
                out_values[i][j] = torch::dot(in_values_i, weights[active_out_indices_a[i][j]]) + bias[active_out_indices_a[i][j]];
            } 
        }
    });
    return out_values;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("mul", &amp;mul, "sparse multiplicatio");
}
</code></pre>
<pre><code class="lang-auto"># setup.py
from setuptools import setup, Extension
from torch.utils import cpp_extension

cpp_module = cpp_extension.CppExtension('sparse_mul_cpp',
                                        sources=['sparse_mul.cpp'],
                                        extra_compile_args=['-fopenmp'],
                                        extra_link_args=['-lgomp']
                                        )

setup(name='sparse_mul_cpp',
      ext_modules=[cpp_module],
      cmdclass={'build_ext': cpp_extension.BuildExtension})
</code></pre>
<p>I also get an issue “a leaf Variable that requires grad is being used in an in-place operation error” when I try to use larger sizes for the weight matrix. I’m using Pytorch version 1.6.0</p>
<p>Thank you for your time <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=9" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"></p>

username: googlebot
cooked: <p>That code is too high level, so the inner loop sucks.</p>
<ol>
<li>Dispatch by dtype, like built-in kernels do.</li>
<li>Instead of Tensor.operator[] indexing, use either raw pointers (this implies contiguous tensors and doing address arithmetic) or “accessors” (torch API).</li>
<li>Compute dot product manually</li>
</ol>

username: its_sandy
cooked: <p>Thank you for your advice <a class="mention" href="/u/googlebot">@googlebot</a> . I was able to achieve a significant speedup after incorporating your suggestions. Here is my current code in case it is of use to anyone. Is there anything else I could do for improving performance even more? I’m now gonna look into implementing backprop for the same.</p>
<pre><code class="lang-auto">// sparse_mul.cpp
#include &lt;torch/extension.h&gt;
#include &lt;vector&gt;

#define _OPENMP
#include &lt;ATen/ParallelOpenMP.h&gt;

template &lt;typename scalar_t&gt;
void mul2_kernel(
        torch::Tensor&amp; out_values,
        const torch::Tensor&amp; in_values,
        const torch::Tensor&amp; active_out_indices,
        const torch::Tensor&amp; weights,
        const torch::Tensor&amp; bias){

    int64_t batch_size = out_values.size(0);
    int64_t active_out_dim = out_values.size(1);
    int64_t k_size = weights.size(1);

    auto out_values_0 = out_values.accessor&lt;scalar_t, 2&gt;();
    auto in_values_0 = in_values.accessor&lt;scalar_t, 2&gt;();
    auto active_out_indices_0 = active_out_indices.accessor&lt;int64_t,2&gt;();
    auto weights_0 = weights.accessor&lt;scalar_t, 2&gt;();
    auto bias_0 = bias.accessor&lt;scalar_t, 1&gt;();

    at::parallel_for(0, batch_size, 0, [&amp;](int64_t start, int64_t end){
        for (int64_t i = start; i &lt; end; i++){
            torch::set_num_threads(1); //is this required
            auto out_values_1 = out_values_0[i];
            auto in_values_1 = in_values_0[i];
            auto active_out_indices_1 = active_out_indices_0[i];
            
            for(int64_t j=0; j&lt;active_out_dim; j++){
                scalar_t &amp;res = out_values_1[j];
                int64_t out_index = active_out_indices_1[j];
                auto weights_1 = weights_0[out_index];

                res = bias_0[out_index];
                for(int64_t k=0; k&lt;k_size; k++)
                    res += in_values_1[k]*weights_1[k];
            }
        }
        // std::cout &lt;&lt; "in thread " &lt;&lt; omp_get_thread_num() &lt;&lt; std::endl;
    });
}

torch::Tensor mul2(
        const torch::Tensor&amp; in_values,
        const torch::Tensor&amp; active_out_indices,
        const torch::Tensor&amp; weights,
        const torch::Tensor&amp; bias){
    
    // std::cout&lt;&lt; "omp_get_max_threads() " &lt;&lt; omp_get_max_threads() &lt;&lt; std::endl;

    auto out_values = torch::empty(active_out_indices.sizes(), in_values.options());
    AT_DISPATCH_ALL_TYPES(
        in_values.scalar_type(), "mul2", [&amp;] {
            mul2_kernel&lt;scalar_t&gt;(out_values, in_values, active_out_indices, weights, bias);
        });
    return out_values;
}
</code></pre>

username: googlebot
cooked: <p>Looks good. Performance-wise, this won’t compile to AVX instructions, but these are tricky. You may check how dot is calculated in <a href="https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/BlasKernel.cpp" rel="noopener nofollow ugc">BlasKernel.cpp</a> to get an idea.Otherwise, you can use AT_DISPATCH_FLOATING_TYPES and remove set_num_threads line.</p>