einsum running 100x slower than expected when controlling for number of ops

samringer: <p>For two simple einsum cases, Case 2 is running ~30% slower than Case 1, despite needing to perform 101x fewer ops.</p>
<p><strong>Case 1</strong></p>
<pre><code class="lang-auto">a = torch.randn(10000, 100, 101).cuda()
b = torch.randn(10000, 101, 3).cuda()

%%time
for _ in range(10000):
    torch.einsum("bij, bjf-&gt;bif", a, b)
    torch.cuda.synchronize()
</code></pre>
<p><code>CPU times: user 26.7 s, sys: 18.8 s, total: 45.5 s Wall time: 45.5 s</code></p>
<p><strong>Case 2</strong></p>
<pre><code class="lang-auto">c = torch.randn(10000, 100, 1).cuda()
d = torch.randn(10000, 100, 1, 3).cuda()

%%time
for _ in range(1000):
    torch.einsum("bic, bicf-&gt;bif", c, d)
    torch.cuda.synchronize()
</code></pre>
<p><code>CPU times: user 37.7 s, sys: 23.7 s, total: 1min 1s Wall time: 1min 1s</code></p>
<p>I believe the number of ops needed for each case should be approximately:<br>
<strong>Case 1:</strong>  10,000 x 100 x 101 x 3 ( x 2)<br>
<strong>Case 2:</strong>  10,000 x 100 x 1 x 3 (x 2)<br>
(The factor of 2 is to account for the multiply-accumulate operation used in matrix multiplication)</p>
<p>Despite requiring 101x fewer ops and operating over smaller tensors than Case 1, in the example above Case 2 has a wall clock time 34% greater than Case 1.<br>
Are there methods/kernels that can be used which are more optimised towards the later case?</p>

tom: <p>The problem is that einsum reduces to batch matmul and so copies your data around.<br>
The real solution is to implement a more general contraction. I had a patch using TensorIterators instead a few years ago, but somehow I decided that it would not work on CPU and abandoned it instead of measuring it on GPU. And then it never was enough fun to dust it off.</p>
<p>One easy way would be to try the Keops third party library.</p>
<p>Best regards</p>
<p>Thomas</p>

tom: <p>To not have the same discussion twice:</p>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/41467" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/41467" target="_blank" rel="nofollow noopener">einsum running 100x slower than expected when controlling for number of ops</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-07-15" data-time="10:13:46" data-timezone="UTC">10:13AM - 15 Jul 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/samringer" target="_blank" rel="nofollow noopener">
          <img alt="samringer" src="https://avatars1.githubusercontent.com/u/33346950?v=4" class="onebox-avatar-inline" width="20" height="20">
          samringer
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🐛 Bug
For two simple einsum cases, Case 2 is running ~30% slower than Case 1, despite needing to perform 101x fewer...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: cuda</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: operators</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">topic: performance</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


samringer: <p>Thanks for your help with this Thomas!</p>
<p>I’ll play around with Keops and see if I can get something closer to the desired behaviour.</p>
<p>Best wishes,</p>
<p>Sam</p>