xref: /aosp_15_r20/external/pytorch/torch/_inductor/fx_passes/README.md (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1# Implicit Invariants for writing FX Graph Passes
2## Fake Tensor metadata on node
3Each FX node has metadata on it, and in particular, stores a faketensor representing the metadata of that node `node.meta['val']`. This FakeTensor has properties like 1. shape, 2. stride, and 3. aliasing information. However, various passes may change the faketensor values, and so we need to maintain consistency.
4
5The current way we do this is through FakeTensorUpdater (in _inductor/fx_utils.py). Read it for more details, and run it if your pass needs accurate faketensor metadata.
6
7## Mutations throughout the stack
8The invariant about mutation we have is:
9
10**After AOTDispatch tracing and before Inductor, we have no mutation in our graph, except for a copy_ epilogue at the end of the graph.**
11
12For example, passes operating on the joint_graph and post_grad graph do not need to worry about mutation at all.
13
14However, we do still have aliasing in the graph. This does not matter most of the time, but it does mean that **our passes are not allowed to cause any additional inputs/outputs to alias if they did not alias in the original graph**.
15
16For example
17```python
18def f(x: Tensor):
19    return x.clone()
20```
21cannot be turned into a no-op, as this would change the semantics of the compiled graph.
22
23In addition, AOTDispatch can introduce a copy_ epilogue into the graph. For example, we may have a graph like
24```python
25def f(x: Tensor):
26    y = x.clone()
27    x.copy_(y)
28    return y
29```
30In this case, we are also not allowed to eliminate `x.clone()`. Luckily, the
31condition for when this can cause problems is the same as with aliasing,
32which is that **our passes are not allowed to cause the input and output to
33alias if they did not alias in the original graph**. To check whether the
34inputs and outputs have any aliasing, it suffices to check whether the
35storages of the input and the storages of the output have any overlap. See
36`remove_noop_ops` for an example of how to do this.
37
38Additionally, we do have one pass that *does* introduce mutation - `reinplace_inplaceable_ops`. This pass must run *just before Inductor lowering*, as otherwise this breaks our invariant.
39