Registering custom operations and CF for FX Graphs

Romain_Moyard: <p>Hi all, I am working for Catalyst <a href="https://github.com/PennyLaneAI/catalyst" class="inline-onebox" rel="noopener nofollow ugc">GitHub - PennyLaneAI/catalyst: A JIT compiler for hybrid quantum programs in PennyLane</a> a Jit compiler for hybrid quantum programs. We have a pipeline that supports Jax, where we define quantum operations as Jax primitives, this allows us to lower to JaxPr and then to stable hlo and our own MLIR quantum dialect. We are now considering to add support for PyTorch and I have the following questions:</p>
<ol>
<li>Similarly to Jax primitives is there a way to extend PyTorch with custom operations and lower them to FX graph operations?</li>
<li>I know that Dynamo is supporting control flow by splitting into multiple subgraphs. Are there native Torch cf operations or is there a way to add control flow operations to Fx graphs?</li>
</ol>
<p>Thanks!</p>

marksaroufim: <p>I think the closest terminology in inductor is lowerings which you can check out here <a href="https://github.com/pytorch/pytorch/blob/main/torch/_inductor/lowering.py" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/torch/_inductor/lowering.py</a> - I’m not sure this has a registration API quite yet but it’s a reasonable ask</p>
<p>Regarding your second question around control flow there’s been some recent work adding a <code>torch.cond</code> operation, you can click around files changed here to get some sense as to how they work <a href="https://github.com/pytorch/pytorch/pull/98909/files" class="inline-onebox" rel="noopener nofollow ugc">Improve torch.cond useability: Return UserError with actionable error messages by guangy10 · Pull Request #98909 · pytorch/pytorch · GitHub</a></p>
<p>That said considering most of your stack includes stable HLO already, if you wait a couple of weeks we’ll probably soon see something like <code>torch.export(m, backend="stable_hlo")</code> be released, the public API for <code>export()</code> is getting merged soon <a href="https://github.com/pytorch/pytorch/pull/106242" class="inline-onebox" rel="noopener nofollow ugc">Expose torch.export() and utilities by gmagogsfm · Pull Request #106242 · pytorch/pytorch · GitHub</a></p>

Romain_Moyard: <p><a class="mention" href="/u/marksaroufim">@marksaroufim</a> Thanks a lot for your answer! Is there not a way to add an operation to torch and compile the FX graph only by using the FX module and not inductor?</p>
<p>Ideally I would imagine:</p>
<p><span class="hashtag-raw">#register</span> my custom operation</p>
<pre><code class="lang-auto">def func(x, y):
    x = torch.sin(x)
    y = torch.ops.my_operation(y)
    return torch.sum(x, y)

fx_graph = make_fx(f)(torch.tensor(3), torch.tensor(3))
</code></pre>

marksaroufim: <p>Yeah that should work just fine, you can develop a custom backend for torch.compile here <a href="https://pytorch.org/docs/stable/dynamo/custom-backends.html" class="inline-onebox">Custom Backends — PyTorch 2.0 documentation</a> cc <a class="mention" href="/u/sherlocknomad">@SherlockNoMad</a> who has been building out most of the infra here</p>

Romain_Moyard: <p>Thanks for your answer. Correct me if I am wrong but the backend comes into play after building the FX graph?</p>