Exact code line for checking if node op is instance of torch.ops.OpOverload

Donghun_Lee: <p>Dear community,</p>
<p>I’ve been digging around torch.comile and torchdynamo but couldn’t find exact line(code) where<br>
<code>make_fx</code> checks if op is instance of torch.ops.OpOverload? (or converting higher level op to OpOverload). It would be really helpful if anyone could point to such code line</p>

bdhirsh: <p>Hey!</p>
<p><code>make_fx()</code> doesn’t directly check if an op is an instance of a <code>torch.ops.OpOverload</code>. Instead, it relies on  <strong>tracing</strong> through all code in torch.* (and the pytorch dispatcher), basically the same as if you were running your code in eager mode.</p>
<p>In eager mode, there are a couple of steps that happen before we execute a kernel. They look roughly like:<br>
(1) Run any python code (e.g. all of the code in <code>torch.nn.MaxPool2d*</code>)<br>
(2) Enter the python ↔ C++ boundary (e.g <code>torch.max_pool2d</code>)<br>
(3) Enter the PythonArgParser, performing overload resolution (this is normally where we desugar from a torch op to an individual “OpOverload”. e.g. `aten.max_pool2d.default)<br>
(4) For that OpOverload, enter the pytorch dispatcher<br>
(5) For that OpOverload, execute it’s cpu/cuda kernel</p>
<p><code>make_fx()</code> tracing involves running all of the same code. However, right before getting to the cpu/cuda kernel, we jump to <code>make_fx</code>’s <code>__torch_dispatch__</code> code, which traces each op into a graph (here: <a href="https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/proxy_tensor.py#L600" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/torch/fx/experimental/proxy_tensor.py#L600</a>)</p>
<hr>
<p>Here are some hopefully-helpful related readings:</p>
<p>A guide stepping slowly through python bindings + the C++ dispatcher for a single operator (torch.mul): <a href="https://github.com/pytorch/pytorch/wiki/PyTorch-dispatcher-walkthrough" class="inline-onebox" rel="noopener nofollow ugc">PyTorch dispatcher walkthrough · pytorch/pytorch Wiki · GitHub</a></p>
<p>Horace’s post about <code>__torch_dispatch__</code>: <a href="https://dev-discuss.pytorch.org/t/what-and-why-is-torch-dispatch/557" class="inline-onebox" rel="noopener nofollow ugc">What (and Why) is __torch_dispatch__? - frontend API - PyTorch Dev Discussions</a></p>
<p>Ed’s blog about the dispatcher: <a href="http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/#:~:text=The%20dispatcher's%20job%20is%20to,pointed%20to%20by%20the%20table" class="inline-onebox" rel="noopener nofollow ugc">Let’s talk about the PyTorch dispatcher : ezyang’s blog</a>.</p>
<p>Also - if you happen to build pytorch from source, and you build with the <code>DEBUG=1</code> build flag, then we have a handy runtime env-var that you can use to print every dispatch key that gets hit:</p>
<pre><code class="lang-auto">def f(x):
    out = torch.mul(x, 2)
    out.sum().backward()
    return out, x.grad

x = torch.ones(2, requires_grad=True)
gm = make_fx(f)(x)

# If you have a debug build, you can try running:
# TORCH_SHOW_DISPATCH_TRACE=1 python script.py
</code></pre>

Donghun_Lee: <p>Thanks so much for detailed explanation!</p>
<p>I have another question regarding PythonArgParser(desugaring part). My use case is that I have a custom function(empty function that returns void) that has been included in the graphmodule via <code>torch._dynamo.allow_in_graph </code> but disappears after running make_fx. I assume that this is natural in the sense that functions added via allow_in_graph are still traced through and since my custom function is empty, it is omitted. Does this exclusion happen in the PythonArgParser or somewhere else?</p>
<p>Thanks in advance.</p>

bdhirsh: <p>If you want your custom function to make it through <code>make_fx</code>, then you will need to register it as a custom operator to the pytorch dispatcher.</p>
<p>(cc <a class="mention" href="/u/richard">@richard</a> - do we have any docs yet for python-only custom ops?)</p>
<p>In the meantime, there are also some test cases that you can look at with examples: <a href="https://github.com/pytorch/pytorch/blob/main/test/test_custom_ops.py#L1552" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/test/test_custom_ops.py#L1552</a>.</p>