Torch.einsum is around ~4x faster than broadcasting torch.matmul for my use case

Santosh_Gupta: <p>My use case is to project the hidden state of every hidden state out of a transformer using a linear layer. There are two ways to do this, broadcast using matmaul or use einsum. I found that using einsum was about 4x faster.</p>
<p>In this colab notebook, I set up the code for each, and profile each method.</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1s5ELkIu_jYEZywExt3uF70tnqLzhKZmb?usp=sharing">
  <header class="source">
      <img src="https://colab.research.google.com/img/favicon.ico?vrz=colab-20220609-060046-RC00_453884947" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1s5ELkIu_jYEZywExt3uF70tnqLzhKZmb?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar" width="260" height="260">

<h3><a href="https://colab.research.google.com/drive/1s5ELkIu_jYEZywExt3uF70tnqLzhKZmb?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colaboratory</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>For those who don’t want to open colab, this are the equivalent operations I am comparing</p>
<p><code>torch.matmul(phase2.unsqueeze(2), weight.unsqueeze(2)).squeeze(2).squeeze(2)</code></p>
<p>and</p>
<p><code>torch.einsum('abc, bc -&gt; ab', phase2, weight)</code></p>
<p>At first I thought it may have been the squeeze and unsqueeze operations, so I did a version where the inputs were pre unsqueezed and no squeeze operations after, and the timings were about the same.</p>
<p>What could be causing the 4x speedup in einsum?</p>

KFrank: <p>Hi Santosh!</p>
<aside class="quote no-group" data-username="Santosh_Gupta" data-post="1" data-topic="153866" data-full="true">
<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/santosh_gupta/48/41337_2.png" class="avatar"> Santosh_Gupta:</div>
<blockquote>
<p><code>torch.matmul(phase2.unsqueeze(2), weight.unsqueeze(2)).squeeze(2).squeeze(2)</code></p>
<p>and</p>
<p><code>torch.einsum('abc, bc -&gt; ab', phase2, weight)</code></p>
</blockquote>
</aside>
<p>I don’t know why <code>einsum()</code> is faster than <code>matmul()</code>, but we have seen things<br>
like this before, for example, as in this post:</p>
<aside class="quote quote-modified" data-post="2" data-topic="152835">
  <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">
    <a href="https://discuss.pytorch.org/t/why-are-batched-operations-slower-than-a-for-loop/152835/2">Why are batched operations slower than a for loop?</a> 
  </div>
  <blockquote>
    Hi Howard! 

Good question – I don’t know. 
I can reproduce your result. including on a recent nightly build. 
It’s interesting to note that a non-loop einsum() version is competitive 
with your “non-batched” loop version, and, at least on my system, turns 
out to be somewhat faster. 
(I see slower, but similar relative timings using just the cpu, with einsum() 
again being the fastest version.) 
Running a tweaked version of your test code, with einsum() added as 
“Batch3,” I get: 
1.10.0
10.2
G…
  </blockquote>
</aside>

<p>(We’ve also seen cases where <code>einsum()</code> is unexpectedly and unreasonably<br>
slow.)</p>
<aside class="quote no-group">
<blockquote>
<p>What could be causing the 4x speedup in einsum?</p>
</blockquote>
</aside>
<p>As an aside, I might guess that it would be better to describe this as an<br>
(unexpected) slowdown in <code>matmul()</code>, rather than as a speedup in <code>einsum()</code>.<br>
Have you considered comparing the <code>matmul()</code> timings with a loop version?</p>
<p>Best.</p>
<p>K. Frank</p>

Santosh_Gupta: <p>Thanks for the detailed answer K.Frank. Will look into a loop version.</p>

Santosh_Gupta: <p>I think I may have found another reason: I suspect einsum uses more GPU memory somehow in this case, because the max inference batch size I can run is lower.</p>
<p>Is there a way to profile this to be sure?</p>