# How to replace a callable with an opaque python callable in an Inductor pass?

**URL:** <https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082>\
**Category:** compiler\
**Created:** [May 13, 2024, 11:19pm UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082 "2024-05-13T23:19:05Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![sanchit](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/sanchit/32/1784_2.png) [@sanchit](https://dev-discuss.pytorch.org/u/sanchit)\
**Post date:** [May 13, 2024, 11:19pm UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/1 "2024-05-13T23:19:05Z")

</div>

Hi,

[HigherOrderOperators](https://docs.google.com/document/d/1i1mee7rw8PorTwQJfjGSqX6GuHMOoLt4aIAGHo-ghes/edit) seem to be helpful for tracing models with _opaque_ python code during a torch.compile invocation.

But can some HigherOrderOperator that uses some arbitrary python be used to replace a Node’s callable (or create a new node with a custom callable that can replace an existing node) in a torch.fx graph during an Inductor pass, or is there some other existing mechanism for doing so?

Basically, I’m creating a python class whose objects (override the \__call_\_() method, and are hence callables) are created during an Inductor pass, and are subsequently used to replace some torch.fx Nodes’ callables.

Currently, I’m making such an implementation work locally with some code in [Inductor’s graph.py, ir.py & lowering.py](https://github.com/pytorch/pytorch/compare/main...sanchitintel:pytorch:onednn_graph_int8?expand=1#diff-e7f291bf7d847fea77a46bd4c0ce6ee0eacf851714a397b4271e54f36fedf369), which adds some conditional checks for the custom python class I created, similar to some conditional checks that already exist for HigherOrderOps. Basically, [make\_fallback()](https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/_inductor/lowering.py#L1668) is called for the callables I create, but simply calling `make_fallback()` is not enough, as the aforementioned conditional checks also need to be added.

Please advise if adding such conditional code would be okay, or if there’s currently some other way to do this.  
In case there’s no existing mechanism to do this, could we create a generic one that could also be (re)used by any other FallBack ops in the future (because it seems `make_fallback()` alone may be insufficient)?

Thanks!

cc @ezyang @jansel

---

<div class="post-metadata">

**Author:** ![jansel](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jansel/32/74_2.png) [@jansel](https://dev-discuss.pytorch.org/u/jansel)\
**Post date:** [May 14, 2024, 11:33pm UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/2 "2024-05-14T23:33:22Z")

</div>

If you want to embed an arbitrary callable in Inductor’s output, then `make_fallback()` should work. HigherOrderOperators do not exist past inductor’s lowering step, so while they might be useful at the aten-FX-graph level you need to map them to something in Inductor’s IR, and a fallback would be the easiest thing to map to since it requires very little code.

What are the conditional checks you want to add? Could you provide an example?

If they are compile-time checks, then you could either do them as an [fx pass](https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/_inductor/fx_passes/post_grad.py#L113-L114), or by defining a custom [lowering function](https://github.com/pytorch/pytorch/blob/1e47c7b11b312b47a621efd547f5c90081f0d9cb/torch/_inductor/lowering.py#L314-L319).

If they are run-time checks, you could put the check inside your custom Python function and just use make\_fallback to put that function in the generated code. You could also define your own type of IR node to codegen the conditional in the output code, then write a lowering to that.

---

<div class="post-metadata">

**Author:** ![sanchit](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/sanchit/32/1784_2.png) [@sanchit](https://dev-discuss.pytorch.org/u/sanchit)\
**Post date:** [May 15, 2024, 3:00am UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/3 "2024-05-15T03:00:25Z")

</div>

Thanks a lot for your response, @jansel!

I meant that despite using `make_fallback()`, in order to make Inductor compilation actually work, I still had to modify some asserts statements to add an exception for the class of the python callables I’m adding.

Here’s another one from ir.py -

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/9/91c5b85ee6c0f9c6c3217f63a93448deddf2224e.png)

Thanks!

---

<div class="post-metadata">

**Author:** ![jansel](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jansel/32/74_2.png) [@jansel](https://dev-discuss.pytorch.org/u/jansel)\
**Post date:** [May 15, 2024, 3:16am UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/4 "2024-05-15T03:16:04Z")

</div>

Ah, I think it is fine to relax that assert. I’d suggest adding a `strict=False` flag to `make_fallback` that just skips the assert entirely. For implicit fallbacks (created on an unknown IR node) I think we want the assert, but if someone calls make\_fallback directly it is not needed.

Feel free to submit a PR for that.

---

<div class="post-metadata">

**Author:** ![sanchit](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/sanchit/32/1784_2.png) [@sanchit](https://dev-discuss.pytorch.org/u/sanchit)\
**Post date:** [May 15, 2024, 5:17am UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/5 "2024-05-15T05:17:36Z")

</div>

Thanks again for the clarification, @jansel! 🙂

Sorry, I should’ve mentioned that my implementation is calling `make_fallback()` in the implicit fallback path.

While adding a `strict=False` flag in `make_fallback()` would help prevent modifying an assert in `make_fallback()` for this use-case, it would not resolve the problem that the \__init()\__ of `FallBackKernel` [requires a](https://github.com/pytorch/pytorch/blob/b950217f199ec06e083b53abd86866239600572b/torch/_inductor/ir.py#L5167) `kernel` to be either an instance of `HigherOrderOp`, or an `OpOverload`.

Is it possible to allow `FallBackKernel` to be created for an instance of another class besides these two? Thanks!

---

<div class="post-metadata">

**Author:** ![jansel](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jansel/32/74_2.png) [@jansel](https://dev-discuss.pytorch.org/u/jansel)\
**Post date:** [May 15, 2024, 4:03pm UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/6 "2024-05-15T16:03:12Z")

</div>

Another option is you could create a PR that adds a class like:

```auto
class InductorImplicitFallback:
     def __call__ (self, *args, **kwargs):
         raise NotImplementedError

```

Then add that class to the assert, and make your class a subclass of that one.

That way people can add this type of thing out of tree without needing to change PyTorch.

---

<div class="post-metadata">

**Author:** ![sanchit](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/sanchit/32/1784_2.png) [@sanchit](https://dev-discuss.pytorch.org/u/sanchit)\
**Post date:** [May 15, 2024, 8:49pm UTC](https://dev-discuss.pytorch.org/t/how-to-replace-a-callable-with-an-opaque-python-callable-in-an-inductor-pass/2082/7 "2024-05-15T20:49:49Z")

</div>

Thanks again, @jansel! 🙂

I’ll work on this solution.