Prologue fusion

MichihiroHorie: <p>Hi everyone, I’m new in PyTorch.</p>
<p>I have questions about fusion in Inductor.</p>
<p>I think Inductor currently supports epilogue fusion, but does not support prologue fusion.<br>
For example, Inductor fuses two nodes: mm followed by relu (that is, epilogue), but when the order is relu followed by mm (that is, prologue), no fusion happens. In Inductor IR, it looks that mm always receives StarDep, which does not match MemoryDep that is the output of relu.</p>
<p>If I want to enable prologue fusion, how can I extend inductor code in the most reasonable way?</p>
<p>In scheduler.py, can_fuse_vertical has conditions to partly accept unmet_dependencies. I’m wondering if the conditions can be extended so that StarDep can compare with MemoryDep somehow. For example, can StarDep be cast to MemoryDep under some conditions?</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/scheduler.py#L2010-L2018">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/scheduler.py#L2010-L2018" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/scheduler.py#L2010-L2018" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/_inductor/scheduler.py#L2010-L2018</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="2010" style="counter-reset: li-counter 2009 ;">
          <li>if (</li>
          <li>    rd.name == cd.name</li>
          <li>    and type(rd) == type(cd)</li>
          <li>    and not free_symbol_has(rd.index, "tmp")</li>
          <li>    and not free_symbol_has(cd.index, "tmp")</li>
          <li>    and rd.index == cd.index</li>
          <li>    and len(rd.size) &gt;= len(cd.size)</li>
          <li>    and rd.size[: len(cd.size)] == cd.size</li>
          <li>):</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>Also, I’m wondering if TemplateBuffer can also handle MemoryDep.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/ir.py#L3327">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/ir.py#L3327" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/ir.py#L3327" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/_inductor/ir.py#L3327</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="3317" style="counter-reset: li-counter 3316 ;">
          <li>    name = self.get_name()</li>
          <li>    indexer = self.layout.make_indexer()</li>
          <li></li>
          <li>    def dummy(index, rindex):</li>
          <li>        assert len(rindex) == 0</li>
          <li>        return ops.store(name, indexer(index), "fake")</li>
          <li></li>
          <li>    deps = dependencies.extract_read_writes(</li>
          <li>        dummy, self.get_size(), (), normalize=True</li>
          <li>    )</li>
          <li class="selected">    deps.reads = {dependencies.StarDep(x.get_name()) for x in self.inputs}</li>
          <li>    return deps</li>
          <li></li>
          <li>def get_reduction_size(self):</li>
          <li>    return 1</li>
          <li></li>
          <li>def get_reduction_type(self):</li>
          <li>    return None</li>
          <li></li>
          <li>def is_no_op(self):</li>
          <li>    return False</li>
      </ol>
    </code></pre>



  </article>

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

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

<p>Any comments would be helpful. Thank you.</p>