JITed GRU too slow

username: Vlado_Boza
cooked: <p>I want to create some custom GRU variants (mainly Layer normalized).<br>
I was following some blog posts and benchmarks (like <a href="https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py" rel="nofollow noopener">https://github.com/pytorch/pytorch/blob/master/benchmarks/fastrnns/custom_lstms.py</a>) and I created following small benchmark <a href="https://gist.github.com/usamec/af21be7b83e6b1a3f38c26136af811f3" rel="nofollow noopener">https://gist.github.com/usamec/af21be7b83e6b1a3f38c26136af811f3</a>, where it seems, that JITed GRU is 10 times slower than cuDNN implemention (but 3times faster than nonJITed). (Using GeForce RTX 2080 Ti).</p>
<p>It is advertised, that at least JIT forward pass runs similarly than cuDNN, so what am I doing wrong?<br>
(I tried running it multiple times, so results are not affected by cold start).</p>

username: G.M
cooked: <p>As far as I know, cuDNN is the fastest in all circumstances. And JIT optimizes the code as it runs, so later batches run faster than the first few batches.</p>

username: Vlado_Boza
cooked: <p>I tried several runs with same inputs. So cold start is not an issue.</p>

username: wanchaol
cooked: <p>can you post your benchmark script? 10 times slower is very weird as we don’t expect that much.</p>

username: Vlado_Boza
cooked: <p>There is a linked script in original post: <a href="https://gist.github.com/usamec/af21be7b83e6b1a3f38c26136af811f3" rel="nofollow noopener">https://gist.github.com/usamec/af21be7b83e6b1a3f38c26136af811f3</a></p>

username: Vlado_Boza
cooked: <p>Also when I increase size of GRU to 1024 features, then the relative difference is much smaller (32 ms JIT, 25 ms cuDNN).</p>

username: Daniel_Gonzalez
cooked: <p>I have the same problem, though my variant is closer to RNN than GRU, so I would think it would be even easier to optimize. Increasing the size has the same effect of shrinking the performance gap, but I need the feature size to be relatively small.</p>
<p>Would be great to figure out how to improve performance of jit RNNs. It’s a really simple model (and inference in C++ is lightning fast), so I’m not sure why it’s so slow.</p>

username: villedepommes
cooked: <p>Could you please open an issue on github and paste in the no. here? I will be taking a look at this very shortly.</p>

username: Vlado_Boza
cooked: <aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/35998" 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/35998" target="_blank" rel="nofollow noopener">JITed GRU too slow</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-04-03" data-time="22:56:33" data-timezone="UTC">10:56PM - 03 Apr 20 UTC</span>
      </div>


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

<div class="github-row">
  <p class="github-content">🐛 Bug
It is advertised, that forward pass of JITed RNNs (e.g. GRU) is as fast as cuDNN implementation.
But it is not...</p>
</div>

<div class="labels">
</div>

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