# A new strategy for automatic custom operators functionalization

**URL:** <https://dev-discuss.pytorch.org/t/a-new-strategy-for-automatic-custom-operators-functionalization/2733>\
**Category:** compiler\
**Created:** [January 7, 2025, 5:07am UTC](https://dev-discuss.pytorch.org/t/a-new-strategy-for-automatic-custom-operators-functionalization/2733 "2025-01-07T05:07:46Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![laithsakka](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/laithsakka/32/2343_2.png) [@laithsakka](https://dev-discuss.pytorch.org/u/laithsakka)\
**Post date:** [January 7, 2025, 5:07am UTC](https://dev-discuss.pytorch.org/t/a-new-strategy-for-automatic-custom-operators-functionalization/2733/1 "2025-01-07T05:07:46Z")

</div>

# A new strategy for automatic custom operators functionalization that enables re-inplacing on view tensors.

with @zou3519

**TLDR** We shipped a new auto functionalization strategy for custom operators that automatically and efficiently handles mutatble custom operators (including ones that mutate views). This unblocked the vLLM integration with torch.compile by fixing OOMs with custom KV-cache operators used in LLMs, which we started seeing a lot of in the past year. We enabled this new **auto\_functionalized\_v2** strategy for torch.compile but not yet for torch.export.

For the most part, this is a technical note that explains the new custom op functionalization strategy.

## **Summary**

The torch.compile stack requires functionalization of all operators (built-in and custom) so that it can create a functional IR. This functional IR makes it easy for us to run optimization passes, like identifying fusion opportunities and reordering operations.

Today, we automatically handle the functionalization of mutable custom operators by creating a functional variant of it. The functional variant of the operator runs the mutable variant, but with additional copies. After we have run most optimization passes, a de-functionalizing (re-inplacing) pass reverts the effect of functionalization and removes the introduced copies.

Previously, the auto functionalization strategy (auto\_functionalized\_v1) did not support re-inplacing custom operators that mutate inputs that are views. We introduced auto\_functionalized\_v2 to do that; read on for the technical details.

### The Details

This rest of the note dives into the details of auto functionalization. The starts by explaining how auto\_functionalize (v1) works, its limitations with respect to views, and finally how auto\_funcitonalize(v2) mitigates it. This is [note](https://l.workplace.com/l.php?u=https%3A%2F%2Fdev-discuss.pytorch.org%2Ft%2Ffunctionalization-in-pytorch-everything-you-wanted-to-know%2F965&h=AT24aKn2ZYByi-xVU0oHdqhdxWvqmWGfffQUPbgm01_Qxf3uDM5HUf5vjMKPSOwsKgDDpKNYd_tRiLMd0e3wlF4bJk3pAu4PtzBqbawcNP9f5X9jX0LRjK0dlhqvTAjCZWRyI6_ZoHLukrylCARIQCGImzbL77NKmuF2b0mK_wAWDRdwmg& __tn__ =-UK-R&c%5B0%5D=AT2BlXutJgnOqfhBd--p6FTAOPfOyPdtbbU5ZpbURs0EZxfB8I_3vSto0un8cOBxU1vvAbJowY1VksugjIgMtlHvZ4976FFNwg0VtlbTNTfdPImLmq1HzcSc1rTmn3AitmSQzUopn4mpxyGibdqtN28EdRFcI2M9u4vSwRY50b7L6pyB9hLI0H-9tZdALZZ9-modljGeip9q_jjTtaIFQZbx)by [**Brian Hirsh**](https://fb.workplace.com/profile.php?id=100053541379589) is a good intro to functionalization that may provide more context.

### Auto\_functionalize (v1)

auto\_functionalized is a higher-order operator ([_HOP_](https://fb.workplace.com/notes/1326456081409687)) that is used to generate the “functional” version of a mutable custom operator. In the program below, _foo_ is a custom op that mutates its input.

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/0/01b26d5c6d1e0919e1b08f1d04b84d0e4e0fa1f7.png)

After running functionalization we get the following functional program:

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/d/d8833a96d793bdf1434a073ef7680967e3030efc.jpeg)

The call to foo is wrapped inside auto\_functionalized which is a HOP that wraps custom operators and converts it into a functional construct. The semantics of auto\_functionalized are the following:

- Copies all the mutable inputs (node _full_ in the graph above)

- Calls the custom op **foo** passing the copies of the mutable inputs.

- Returns the output of **foo** , followed by each of the copies of the mutable inputs.

We also have to replace all usages of the mutable inputs with auto\_functionalized output ( **return f** in this case is replaced with **return getitem1** ).

To visualize that, we can print the post-grad graph while manually disabling re-inplacing and notice the following; (1) the call to clone the mutable input, (2) the return of the function being the copied input after it is mutated. (do not get confused by as\_strided nodes, they are just aliases here).

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/2/2dd5cade0b28bd70aa65f1cacedae1d960f1f114.jpeg)

After all post-grad passes run on the functional graph, a re-inplace pass runs to reverse the functionalization. The algorithm for re-replacing is simple, for each mutable input to the auto\_functionalized custom operator:

_ **if the input in the functional graph is not read after auto\_functionalized then it can be re-inplaced** _ by doing the following:

1. Do not generate the copy of input (_full_ in the graph above) and pass the input itself to foo during auto\_functionalized decomposition.

2. Replace every usage of the copy of the inputs, which is the output of auto\_functionalize (_getItem1_ in the example above) back to the input itself (_full_ in the graph above).

Running the re-inplacing pass in our example yields the following final outputs. foo is called with the original input directly, and there are no unnecessary clones.

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/4/420f5d3d19fabd0762271185692b005152b2d5a7.jpeg)

### Auto\_functionalize (v1) and views

Consider the following small alteration to the previous program, where we pass f[0] instead of f to foo:

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

After functionalization we get the following forward graph.

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/0/0d89741ac9e1d0e63bce4b54d10b61f68e1237b8.jpeg)

Note the **select\_scatter**. This is an operation that creates a copy of _full_ and writes the portion of _select\_scatter_[0] from _get\_item_. Since the original program returns f and the portion f[0] was mutated by foo, we need to return a copy of f, that observes the effect of _foo_ on f[0], and the mutation is observed from _get\_item_.

The select\_scatter interferes with re-inplacing! The problem is in the re-inplacing part, the criteria for re-inplacing is :"_ **if the input in the functional graph is not read after auto\_functionalized then it can be re-inplaced"** _

It’s tricky to say yes to this question with the existence of select\_scatter in the graph, and with the fact that select (the input to foo) is a view. The select scatters takes full as argument? but full share storage with select, does that mean that we read select?

We know the answer for this case but in general this is a hard question to answer and requires an alias analysis. The view operations before auto\_functionalize and the scatter operations after can get really messy and complicated !

Due to that complexity, auto\_functionalized does not try to re-inplace view operations, giving us the following graph with additional clones:

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/8/886c024933534620be498748438945ea7ee8e229.jpeg)

### Auto\_functionalize (v2)

We experimented with several designs until we settled on this one.

The main question is the following: _is it possible to do the functionalization in a way that avoids dealing with views and scatters so we do not need to worry about the aliasing complexity._

The answer is yes, and the key to the new approach is reformulating inplace-on-view as an in-place on the base tensor followed by the reconstruction of said views. Under this scheme, auto\_functionalized\_v2’s strategy is that it clones the bases, and mutates those clones, and then the rest of the program reconstructs the views from these mutated clones.

What this does is that it removes the two complexities (view generations before auto\_functionalize is called, and select scatter after it is called) so we do not need to worry about any alias analysis.

Here’s an example to demonstrate what we mean. Below is the functional graph that we get when **auto\_functionalized (v2)** is enabled\*\*.\*\*

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/0/0c229530429aeb56b9f89df39397bf25d0dbb581.jpeg)

auto\_functionalized\_v2 takes the following arguments:

- **all\_bases** , a list of unique bases of all mutable inputs of the custom operator. In this case, the base of `full[0]` is the `full` node.

- For each mutable input **\_x\_bases\_index** , which identify the index base of x in all\_bases

- **ViewInfo** Some metadata that tells how to regenerate the view from the base for each arg if its a view, in the case bellow **\_x\_size** , **\_x\_stride** and **\_x\_storage\_offset** indicates that we can generate x by calling `all_bases[0].as_strided((), (), 0)`.

auto\_functionalized\_v2 have the following semantics:

- Copy all the tensors in all\_bases!

- Regenerate all the views (if any from all\_bases) and pass the new copies (or views) to the custom operator (foo in this case).

- Returns the output of custom-op followed by the copies of all\_bases.

- We have to replace every usage of any base, after foo is called with the new copy. (_This will not generate select\_scatters since every variable is either the base itself or a view on top of the new output base_).

If we were to decompose auto\_functionalized\_v2 (without re-inplacing), here’s what it would look like.

 ![image](https://canada1.discourse-cdn.com/flex036/uploads/pytorch1/original/2X/b/b58442ee8243f80cdc24c948ed5b43b142d711da.jpeg)

so many as\_strided nodes in the code above, but do not get confused they are all aliases except for _as\_strided\_default\_2_ nod_e_!!

as\_strided\_default\_2 is generated when auto\_functionalized\_v2 is decomposed using the information serialized in the auto\_functionalize\_v2 call, it generates a view on top of clone\_default[0].

Cool Cool , so now let’s check re-inplacing.

We use the same criteria as auto\_functionalized(\_v1), but apply it to all\_bases instead “_ **if the base in the functional graph is not read after auto\_functionalized\_v2 then it can be re-inplaced** _\*”\* . Applying it to _full_ on the forward graph above, the answer is clear: full is not read past auto\_functionalized\_v2 and we can re-inplace it!!

Inductor is able to re-inplace auto\_functionalized\_v2 to create a copy/clone-free graph.

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

## Other Interesting stuff :

**1. Tracking base in inference mode:**  
In inference mode, tensors are not aware of their base, so we did some hacking to track the base of tensors by analyzing the shared storages during functionalization.

**2. Do we always use as\_strided to regenerate the views?**  
as\_strided is very generic and can be used to generate any view; however, [**Elias Ellison**](https://fb.workplace.com/profile.php?id=100021691660923) noted that it can potentially cause performance issues in the inductor and is harder to analyze by passes. (it’s harder to reason about what it access). To mitigate that we attempt to infer 4 types of views from the striding information (slice, alias, subscript(not yet implemented) and fall back to calling as\_strided if we cannot do the inference.

## Future Work

**Export mode:** auto\_functionalized\_v2 is not enabled in torch.export; we didn’t find motivating use cases for it. Note that some exported models already serialize an auto\_functionalized\_v1 HOP, so even if we introduced auto\_functionalized\_v2, we cannot immediately delete auto\_functionalized\_v1.

**Triton kernel re-inplacing does not go through v2:** This may be important; we’re working on logging how impactful this is.

---

<div class="post-metadata">

**Author:** ![guoyejun](https://avatars.discourse-cdn.com/v4/letter/g/3e96dc/32.png) [@guoyejun](https://dev-discuss.pytorch.org/u/guoyejun)\
**Post date:** [February 8, 2025, 2:16pm UTC](https://dev-discuss.pytorch.org/t/a-new-strategy-for-automatic-custom-operators-functionalization/2733/2 "2025-02-08T14:16:57Z")

</div>

what does “post-grad graph” mean, thanks.

---

<div class="post-metadata">

**Author:** ![jeromeku](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/jeromeku/32/925_2.png) [@jeromeku](https://dev-discuss.pytorch.org/u/jeromeku)\
**Post date:** [September 9, 2025, 4:30pm UTC](https://dev-discuss.pytorch.org/t/a-new-strategy-for-automatic-custom-operators-functionalization/2733/3 "2025-09-09T16:30:42Z")

</div>

@laithsakka @bdhirsh

On pytorch nightly `2.9.0.dev20250907+cu128`, running the above example gives the following result:

```python
import torch
import torch._inductor.config as inductor_config
from torch._logging import set_logs
import logging

@torch.library.custom_op("mylib::my_func", mutates_args={"x"})
def my_func(x: torch.Tensor) -> None:
    x[0] = 1

@torch.compile()
def func():
    f = torch.ones(10)

    my_func(f)

    return f

if __name__ == " __main__":
    set_logs(inductor=logging.DEBUG)
    with inductor_config.patch({"enable_auto_functionalized_v2": False}):
        func()

```

Toggling `enable_auto_functionalized_v2` has no effect.

Both cases result in the following logs:

```auto
I0909 16:26:19.457000 3205004 torch/_inductor/fx_passes/reinplace.py:560] For node auto_functionalized_v2, attempted to reinplace range(0, 1). We were unable to reinplace []; [] (if non-empty) are possible missed reinplacing opportunities that may be bad for memory usage and performance. Total size of missed opportunities with static shapes is : 0 bytes.
=== BEFORE ===
graph():
    %x_1 : [num_users=2] = placeholder[target=x_1]
    %auto_functionalized_v2 : [num_users=2] = call_function[target=torch.ops.higher_order.auto_functionalized_v2](args = (mylib.inc_.default,), kwargs = {_x_base_index: 0, _all_bases: [%x_1]})
    %getitem : [num_users=0] = call_function[target=operator.getitem](args = (%auto_functionalized_v2, 0), kwargs = {})
    %getitem_1 : [num_users=1] = call_function[target=operator.getitem](args = (%auto_functionalized_v2, 1), kwargs = {})
    %copy_ : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%x_1, %getitem_1), kwargs = {})
    return copy_
=== AFTER ===
graph():
    %x_1 : [num_users=2] = placeholder[target=x_1]
    %inc__default : [num_users=0] = call_function[target=torch.ops.mylib.inc_.default](args = (%x_1,), kwargs = {})
    %copy_ : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%x_1, %x_1), kwargs = {})
    return copy_

```

A few questions:

- Is the `copy` in the post-grad graph a no-op since it’s copying `x` into `x`?
- Why is there no difference whether autofunc v2 is enabled?
- How should I interpret the reinplace warning: `For node auto_functionalized_v2, attempted to reinplace range(0, 1). We were unable to reinplace...Total size of missed opportunities with static shapes is : 0 bytes.`

---

<div class="post-metadata">

**Author:** ![lifuhuang](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/lifuhuang/32/2970_2.png) [@lifuhuang](https://dev-discuss.pytorch.org/u/lifuhuang)\
**Post date:** [September 28, 2025, 1:47am UTC](https://dev-discuss.pytorch.org/t/a-new-strategy-for-automatic-custom-operators-functionalization/2733/4 "2025-09-28T01:47:53Z")

</div>

Thank you for the great work! Out of curiosity, in the case when re-inplacement is not possible, would v2 result in additional overheads (as the whole base tensor, not just the view, is copied)? Thanks!