C++ extension performance issue

Sia_Rezaei: <p>I’m writing a C++ extension to optimize the performance of a custom function. As a starting point I’m only focused on the CPU version.<br>
My optimized version should in principle be faster by ~30%.<br>
However, it runs in about the same time. Below are my code:</p>
<p><strong>Unoptimized version:</strong></p>
<pre><code class="lang-auto">torch::Tensor symsum_forward_unoptimized(torch::Tensor input){
    const int shift = input.size(1) / 2;
    float data[] = {0};
    const torch::Tensor&amp; zero = torch::from_blob(data, {});
    auto relu = torch::maximum(input, zero);
    auto inv_relu = torch::minimum(input, zero);
    auto output = torch::roll(inv_relu, shift, 1) + relu;
    return output;
}
</code></pre>
<p><strong>Optimized version:</strong></p>
<pre><code class="lang-auto">torch::Tensor symsum_forward_ptr(torch::Tensor input){
    // Create and initialize the output tensor.
    const torch::Tensor&amp; output = torch::zeros_like(input);

    // Batch stride is needed for pointer arithmetic.
    const int batch_stride = input.strides()[0]; 
    const int half_batch_stride = batch_stride / 2;

    // Define and initialize `input` and `output` tensor pointer
    float *in_i1 = input.data_ptr&lt;float&gt;();  // input index 1
    float *in_i2 = in_i1 + half_batch_stride;  // input index 2
    float *out_i1 = output.data_ptr&lt;float&gt;();  //output index 1
    float *out_i2 = out_i1 + half_batch_stride;    //output index 2

   // Define initialize end of batch and end of tensor pointers for `input`
    float *batch_end = in_i2 + half_batch_stride;
    const float* end = input.data_ptr&lt;float&gt;() + input.numel();
     
    // This loop goes over samples in the batch
    while(in_i2 &lt; end){ 
        // This loop goes over elements in a sample
        while (in_i2 &lt; batch_end){
            // `std::signbit` is used instead of `(*in_i1) &lt; 0`. It is faster.
            if (std::signbit(*in_i1)) {
                if (std::signbit(*in_i2)){
                    *out_i1 = *in_i2;
                    *out_i2 = *in_i1;                    
                }
                else{
                    *out_i2 = (*in_i2) + (*in_i1);
                }
            }
            else{
                if (std::signbit(*in_i2)){
                    *out_i1 = (*in_i2) + (*in_i1);
                }
                else{
                    *out_i1 = *in_i1;
                    *out_i2 = *in_i2;                    
                }                
            }
            in_i1++;
            in_i2++;
            out_i1++;
            out_i2++;  
        }

        in_i1 += half_batch_stride;
        in_i2 += half_batch_stride;        
        out_i1 += half_batch_stride;
        out_i2 += half_batch_stride; 
        batch_end += batch_stride;
    }    
    return output;
}

}
</code></pre>
<p>The idea behind the optimized version is simple. In the unoptimized version each element of <code>input</code> is compared with zero twice (once in <code>maximum</code> and once in <code>minimum</code>). However, this only needs to happen once. Because if an element is larger or equal to zero, then I know it is not smaller than zero.</p>
<p>The optimized version takes advantage of this insight. It performs the same function, but with only one comparison per element. Furthermore, it allocates one third the memory of the unoptimized one. So it should be faster in principle.<br>
Any insight about why it is not faster is much appreciated.</p>

ptrblck: <p>I’m not sure if your theoretical analysis is correct, as it seems you are comparing a nested loop approach while the native implementation would use the <code>TensorIterator</code> with the <code>at::parallel_for</code> approach.<br>
You could check the source code via <a href="https://github.com/pytorch/pytorch/blob/043098ef7f259d6f389115b47a4466d7efe62f4c/aten/src/ATen/native/TensorCompare.cpp#L444-L451"><code>max_out</code></a>, which should call into the <code>max_stub</code>,  <a href="https://github.com/pytorch/pytorch/blob/ecf9c82f24a9cd93d65242b1c426507c78e11110/aten/src/ATen/native/cpu/TensorCompareKernel.cpp#L126-L157"><code>max_kernel_impl</code></a>, <a href="https://github.com/pytorch/pytorch/blob/ecf9c82f24a9cd93d65242b1c426507c78e11110/aten/src/ATen/native/cpu/TensorCompareKernel.cpp#L22-L62"><code>compare_base_kernel_core</code></a>, and eventually into <a href="https://github.com/pytorch/pytorch/blob/6df7b7518684f2018469471ed8ac0662a4d2c700/aten/src/ATen/TensorIterator.cpp#L736-L747"><code>TensorIteratorBase::for_each</code></a>.</p>

googlebot: <p>even in single thread mode, your code won’t be auto-vectorized, while pytorch utilizes AVX2.</p>

Sia_Rezaei: <p>Thank you for the comments.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> Each iteration of the loop performs the comparison for two elements (that’s why the nested <code>if</code>). So in the end each element is only compared once, and the theoretical analysis is sound I think.</p>
<p>I’m testing the code with <code>torch.set_num_threads(1)</code>. So I’m assuming  both algorithms are playing in level ground in terms of parallelization(?)</p>
<p><a class="mention" href="/u/googlebot">@googlebot</a> Do you anticipate that if I disabled AVX2 (or tried the code on a CPU without AVX instructions), then the optimized code would runs faster?</p>

googlebot: <p>Um, no, I mean that pytorch processes 8 items at once, that’s potential reason why you can’t reach its speed. It may be possible to rewrite your loop with AVX intinsics, but it is not trivial; it may also auto-vectorize if you replace branches with b?x:y operations, and add __restrict to pointers, I’m just not sure if signbit is supported.</p>
<p>PS Oh, I misundserstood - if you’ll use “slow” pytorch for comparision, then yes, I don’t rule out that your routine will be faster. It screws the branch predictor though.</p>

Matej_Kompanek: <p>I tried implementing AVX version. On my PC it runs 3 times faster<br>
Duration symsum_avx: 0.127393 s<br>
Duration symsum_forward_unoptimized: 0.370589 s<br>
EDIT: changed loading and storing to unaligned versions as suggested</p>
<p>here’s the code, if you’re interested</p>
<pre><code class="lang-auto">torch::Tensor symsum_avx(torch::Tensor input)
{
	const torch::Tensor&amp; output = torch::empty_like(input);

	const int batch_stride = input.strides()[0];
	const int half_batch_stride = batch_stride / 2;

	float* in_i1 = input.data_ptr&lt;float&gt;(); 
	float* in_i2 = in_i1 + half_batch_stride;
	float* out_i1 = output.data_ptr&lt;float&gt;();
	float* out_i2 = out_i1 + half_batch_stride;

	float* batch_end = in_i2 + half_batch_stride;
	const float* end = input.data_ptr&lt;float&gt;() + input.numel();

	const __m256 avx_zeros = _mm256_setzero_ps();

	float all_bits_on = 0.0f;
	memset(&amp;all_bits_on, ~0, 4);
	const __m256 avx_all_bits_on = _mm256_set1_ps(all_bits_on);

	while (in_i2 &lt; end)
	{
		while ((batch_end - in_i2) &gt;= 8)
		{
			auto in1 = _mm256_loadu_ps(in_i1);
			auto in2 = _mm256_loadu_ps(in_i2);

			const auto b1 = _mm256_cmp_ps(in1, avx_zeros, _CMP_LT_OQ);
			const auto b2 = _mm256_cmp_ps(in2, avx_zeros, _CMP_LT_OQ);

			const auto c1 = _mm256_and_ps(b1, b2);
			const auto c2 = _mm256_andnot_ps(b1, b2);
			const auto c3 = _mm256_andnot_ps(b2, b1); 
			const auto c4 = _mm256_andnot_ps(b1, _mm256_andnot_ps(b2, avx_all_bits_on));

			const auto sum = _mm256_add_ps(in1, in2);

			auto out1 = _mm256_blendv_ps(avx_zeros, in1, c4);
			out1 = _mm256_blendv_ps(out1, sum, c2);
			out1 = _mm256_blendv_ps(out1, in2, c1);

			auto out2 = _mm256_blendv_ps(avx_zeros, in2, c4);
			out2 = _mm256_blendv_ps(out2, sum, c3);
			out2 = _mm256_blendv_ps(out2, in1, c1);

			_mm256_storeu_ps(out_i1, out1);
			_mm256_storeu_ps(out_i2, out2);
			out_i1 += 8; out_i2 += 8; in_i1 += 8; in_i2 += 8;
		}

		while (in_i2 &lt; batch_end)
		{
			const bool b1 = std::signbit(*in_i1);
			const bool b2 = std::signbit(*in_i2);

			const bool c1 = b1 &amp;&amp; b2;
			const bool c2 = !b1 &amp;&amp; b2;
			const bool c3 = b1 &amp;&amp; !b2;
			const bool c4 = !b1 &amp;&amp; !b2;

			const float sum = (*in_i2) + (*in_i1);

			*out_i1++ = c1 ? *in_i2 : (c2 ? sum : (c4 ? *in_i1 : *out_i1));
			*out_i2++ = c1 ? *in_i1 : (c3 ? sum : (c4 ? *in_i2 : *out_i2));

			++in_i1;
			++in_i2;
		}

		in_i1 += half_batch_stride;
		in_i2 += half_batch_stride;
		out_i1 += half_batch_stride;
		out_i2 += half_batch_stride;
		batch_end += batch_stride;
	}

	return output;
}
</code></pre>
<p>How I tested:</p>
<pre><code class="lang-auto">
void main()
{
	auto input = torch::randn({ 10000, 10000 });

	auto avx_start = std::chrono::steady_clock::now();
	auto output_avx = symsum_avx(input);
	auto avx_end = std::chrono::steady_clock::now();

	auto default_start = std::chrono::steady_clock::now();
	auto output_default = symsum_forward_unoptimized(input);
	auto default_end = std::chrono::steady_clock::now();

	std::cout &lt;&lt; "Duration symsum_avx: " &lt;&lt; (std::chrono::duration_cast&lt;std::chrono::microseconds&gt;(avx_end - avx_start).count() / 1000000.0) &lt;&lt; " s" &lt;&lt; std::endl;
	std::cout &lt;&lt; "Duration symsum_forward_unoptimized: " &lt;&lt; (std::chrono::duration_cast&lt;std::chrono::microseconds&gt;(default_end - default_start).count() / 1000000.0) &lt;&lt; " s" &lt;&lt; std::endl;

	std::cout &lt;&lt; "Max error: " &lt;&lt; (output_avx - output_default).abs().max() &lt;&lt; std::endl;
}
</code></pre>

googlebot: <p>nice</p>
<p>you may need to use _mm256_loadu_ps / _mm256_storeu_ps I believe</p>

Matej_Kompanek: <p>I thought that pytorch’s memory is properly aligned for AVX (by looking <a href="https://github.com/pytorch/pytorch/issues/50043" rel="noopener nofollow ugc">here</a>). Am I wrong ?</p>

tom: <p>Freshly allocated tensors, yes, but of course, if I slice away the first element, not so much, but I guess your code isn’t completely safe w.r.t. memory layout, anyways.</p>

Sia_Rezaei: <p><a class="mention" href="/u/matej_kompanek">@Matej_Kompanek</a> wow, you are an angel! On my PC it runs ~4x faster than my unoptimized version. Now I need to decipher this to write the backwards function. Can anyone suggest a good source to learn AVX?</p>
<p>Do you think a similar gain achievable with optimizing the CUDA version?</p>

Matej_Kompanek: <blockquote>
<p>Do you think a similar gain achievable with optimizing the CUDA version?</p>
</blockquote>
<p>No idea, I didn’t expect such an improvement with avx.<br>
I think there is no point in messing around with avx, you should probably focus on CUDA since you’re going to train on GPU anyway (I presume).</p>

googlebot: <p>on CUDA, you gain more from single pass processing (fusion), so you may not need any special “optimizing” beyond that; if you replace your 3 (max,min,+) kernel calls with 1, speedup will be significant, as most overhead is from memory IO (esp. if data doesn’t fit in cache).</p>
<p>As for AVX, “Intel® 64 and IA-32 Architectures Software Developer’s Manual” is the canonical guide, search for pdf, it is on Intel’s site somewhere. Reference is at <a href="https://www.intel.com/content/www/us/en/docs/intrinsics-guide/index.html" class="inline-onebox" rel="noopener nofollow ugc">Intel® Intrinsics Guide</a></p>
<p>But this is assembler level programming really, as such it is not portable. It is better to write auto-vectorizable code, by satisfying some compiler specific requirements.</p>

Sia_Rezaei: <aside class="quote no-group" data-username="googlebot" data-post="12" data-topic="139572">
<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>
<p>compiler specific requirements</p>
</blockquote>
</aside>
<p>Is there a way to find out what those <em>compiler specific requirements</em> are for auto vectorization?</p>

googlebot: <p>Yes, there is usually <em>some</em> documentation, for example <a href="https://gcc.gnu.org/projects/tree-ssa/vectorization.html" class="inline-onebox" rel="noopener nofollow ugc">Auto-vectorization in GCC - GNU Project</a>.</p>
<p>I tried to vectorize your loop out of curiosity - <a href="https://godbolt.org/z/M99x3PWYs" class="inline-onebox" rel="noopener nofollow ugc">Compiler Explorer</a> . Haven’t tested it, but it should be the same as with intristics.</p>
<p>Note that I:</p>
<ol>
<li>added const __restrict modifiers</li>
<li>added an indexing (induction) variable</li>
<li>removed if-else branches<br>
All this to simplify job of code generator. They fail at auto-vectorizing very easily in all compilers, unfortunately.</li>
</ol>
<p>I’ll paste just in case</p>
<pre><code class="lang-auto">void f(const float * __restrict in_i1, float * __restrict out_i1, int batch_stride, int numel)
{
    const int half_batch_stride = batch_stride / 2;
    const float *in_i2 = in_i1 + half_batch_stride;  // input index 2
    float *out_i2 = out_i1 + half_batch_stride;
    const float *batch_end = in_i2 + half_batch_stride;
    const float* end = in_i1 + numel;

    while(in_i2 &lt; end){ 
        for(int i=0; i&lt;batch_stride; i++) {
            float v1 = in_i1[i];
            float v2 = in_i2[i];
            float sum = v1+v2;
            bool b1 = v1 &lt; 0, b2 = v2 &lt; 0;

            out_i1[i] = b1 ? (b2 ? v2 : sum) : (b2 ? sum : v1);
            out_i2[i] = b1 ? (b2 ? v1 : sum) : (b2 ? sum : v2);
        }
        in_i1 += batch_stride;
        in_i2 += batch_stride;
        out_i1 += batch_stride;
        out_i2 += batch_stride; 

        in_i1 += half_batch_stride;
        in_i2 += half_batch_stride;        
        out_i1 += half_batch_stride;
        out_i2 += half_batch_stride; 
   }
}
</code></pre>

Sia_Rezaei: <aside class="quote no-group" data-username="googlebot" data-post="14" data-topic="139572">
<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>
<p>I’ll paste just in case</p>
</blockquote>
</aside>
<p>I updated the function inspired by your code. It takes 1.5x longer than the unoptimized version to run on my PC. I don’t think it is being vectorized.</p>
<pre><code class="lang-auto">torch::Tensor f_forward_vec(torch::Tensor input){
    const torch::Tensor&amp; output = torch::zeros_like(input);
    const int batch_stride = input.strides()[0]; 
    const int half_batch_stride = batch_stride / 2;
    const int onehalf_batch_stride = batch_stride + half_batch_stride;
    float *__restrict__ in_i1 = input.data_ptr&lt;float&gt;();
    float *__restrict__ in_i2 = in_i1 + half_batch_stride;
    float *__restrict__ out_i1 = output.data_ptr&lt;float&gt;();
    float *__restrict__ out_i2 = out_i1 + half_batch_stride;
    const float* end = input.data_ptr&lt;float&gt;() + input.numel();
      
    // This loop goes over samples in the batch
    while(in_i2 &lt; end){ 
        // This loop goes over elements in a sample
        for (int i=0; i&lt;batch_stride; i++){ 
            float v1 = in_i1[i];
            float v2 = in_i2[i];
            bool b1 = std::signbit(v1), b2 = std::signbit(v2);
            out_i1[i] = b1 ? (b2 ? v2 : v1+v2) : (b2 ? v1+v2 : v1);
            out_i2[i] = b1 ? (b2 ? v1 : v1+v2) : (b2 ? v1+v2 : v2);      
        }
        in_i1 += onehalf_batch_stride;
        in_i2 += onehalf_batch_stride;
        out_i1 += onehalf_batch_stride;
        out_i2 += onehalf_batch_stride;
    }    
    return output;
}
</code></pre>
<p>It’s no big deal though. I care about the Cuda version which turned out to work great.</p>

googlebot: <p>Yes, your changes break vectorization - as I said it is very brittle. If you view godbolt. org, my code has characteristic AVX instructions (suffix ‘ps’ for packed single) and “add     rsi, 32” line that is a step over 8 floats. Big chunk of code below handles 0-7 tail elements.</p>
<p>More practical way is enabling compiler diagnostics, that may say the reason why auto-vectorization failed (the message may be obscure though).</p>

tom: <aside class="quote no-group" data-username="Sia_Rezaei" data-post="10" data-topic="139572">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sia_rezaei/48/43436_2.png" class="avatar"> Sia_Rezaei:</div>
<blockquote>
<p>Can anyone suggest a good source to learn AVX?</p>
</blockquote>
</aside>
<p>PyTorch has internal vec256 headers which abstract avx and might be a good resource.</p>
<p>Best regards</p>
<p>Thomas</p>