Repository: pytorch
Issue: [DRAFT] Auto-casting in JIT - Automatic mixed precision
Labels: oncall: jit,triaged

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->
We want to support automatic precision casting as a JIT optimization pass. This is part of a grand plan to improve PyTorch performance by leveraging hardware acceleration on reduced precision arithmetics.
For more background and motivation, on automatic mixed precision, #25081 provides far more details.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->
Whitelist-blacklist auto-casting has been proven to work on many models through Apex AMP. Dispatcher auto-casting, as proposed in #25081, would work in similar fashion. But auto-casting done through graph mode provides a global view of the computation graph that eager mode does not see, which gives us the opportunity to further improve the performance of auto-casting regarding both computational efficiency and numerical reliability.

Because of the static graph, the decision of cast insertion could be optimized globally. Consider the two following scenarios:
1. Cascading operations: 

```
%2 : Tensor = at::BlackOp0(%1)
%3 : Tensor = at::ClearOp0(%2)
%4 : Tensor = at::WhiteOp0(%3)
```

Eager mode would cast data to fp32 at `%2 : Tensor = at::BlackOp0(%1)` ; Since `%3 : Tensor = at::ClearOp0(%2)` doesn't mind precision, it will transparently handle the data passed in as fp32; Tensor would be casted to fp16 when entering `%4 : Tensor = at::WhiteOp0(%3)`.

Graph mode, as we have visibility of the full sequence, we could make the decision that it is better to cast the output of `%2 : Tensor = at::BlackOp0(%1)` to fp16. This is numerically safe as `at::ClearOp0` would not be affected by reduced precision. More importantly here, It is computationally beneficial as we are reducing memory traffic at `at::ClearOp0` with reduced precision.

2. Branches:
```
%2 : Tensor = at::BlackOp0(%1)
%4 : Tensor = at::WhiteOp0(%3)
%5 : Tensor = at::ClearOp0(%2, %4)
%6 : Tensor = at::ClearOp1(%5)
%7 : Tensor = at::EndOp0(%6)
```
Depends on whether `%7 : Tensor = at::EndOp0(%6)` is numerically safe with reduced precision (white-listed or black-listed), it implores different casting policy for `%5 : Tensor = at::ClearOp0(%2, %4)`.

Because we don't want to use reduced precision when it's not numerically safe or performance-wise beneficial. If there is no whitelist operation on the path between two blacklist operations, it is hard to justify the performance benefit outweighing the numerical risk of reduced precision.

## Pitch

<!-- A clear and concise description of what you want to happen. -->
We want a graph optimization pass mutates the graph by inserting casting ops at appropriate places. The mutated graph would have identical numerical behavior and improved computation efficiency due to reduced precision.

### Implementation:

The optimization pass would be implemented as:

1. Graph Partition: mark nodes in the graph as black/white
2. Insert Cast Ops:
    1. Run `CustomFuseGraph`: Try to fuse white nodes and put them under a custom fusion group;
    2. Insert cast ops at the boundary outside of custom fusion group;
    3. Cancel out inserted cast ops across blocks;
    4. Inlining custom fusion group back into the graph;

A quick note just to avoid very likely confusion: 
- _[Phase 1 Graph Partition]_ do *NOT* mutate the graph. We can be very aggressive in painting white nodes, as long as it’s numerically safe;
- _[Phase 2 Insert cast ops]_ *ONLY* inserts cast ops in the graph. So it should *NOT* undermine future optimization passes;
- _[Phase 2 Insert Cast Ops]_ `CustomFuseGraph` is used here only to handle data dependency, it has nothing to do with kernel fusion (which uses `FuseGraph` in a similar manner);


### Rationals the proposed implementation:

1. Graph Partition

We use a graph painting approach: The technique starts with painting whitelist-blacklist nodes in the graph; It then traces data flow to propagate the color information to not-yet-painted nodes. Doing so it partition the nodes according to whether it’s painted white or not.

One tip here is, we maintain the partition in a separate list without mutating the graph. This allows us to simplify partition logic without worrying about undermining the semantics of the graph.

We can start with the simplified technique:

    1. Mark all blacklist nodes in the graph as black
    2. Mark all whitelist nodes in the graph as white
    3. Mark all nodes in a segment, which starts and ends with black nodes without any white nodes in the path as black
    4. Recursively grow white nodes along both input/output edges to neighboring un-colored nodes until reaches a black node.

2. Mutability-safe optimization passes

Auto-casting will change the semantics of the graph. To safely insert a cast op in the graph, we need to ensure that our casting will not leak through data dependency, implicitly introduced by aliasing in PyTorch JIT IR. For example, if we have a graph like:
```
    %2 : Tensor = prim::opA(%1)
    %3 : Tensor = prim::opB_(%2)
    %4 : Tensor = prim::opC(%2)
```
We should be aware that inserting a cast before node `%3 : Tensor = prim::opB_(%2)` would cause node `%4 : Tensor = prim::opC(%2)` to be in the same precision. This would not be safe, if `%3 : Tensor = prim::opB_(%2)` is painted white, while `%4 : Tensor = prim::opC(%2)` is painted black. In this scenario reduced precision has breached in the unsafe region.

Unfortunately there is no API that allows efficiently data dependency query in our infrastructure yet. Otherwise, we would be able to follow the implicit dependency introduced by aliasing to properly partition graph to avoid such violation.

In order to accommodate this, we need to make sure that we are not inserting cast ops immediately before an in-place operator, which means that graph partition should not cross the edge where down-stream operation either mutates input or returns alias(es) of input(s).

Because our graph partition follows only direct edged in the graph and does not consider aliasing, we can only fix this after the partitioning. The baseline approach is, that we iteratively check all boundary edges and update the painting if we detect violations. Note any update to the painting would change the boundary, so we have to check the new boundary edges as well.

Alternatively, to avoid handling tricky dependency analysis, we use existing `CustomFuseGraph` to fuse all white nodes into a custom FusionGroup. *THIS HAS NOTHING TO DO WITH THE FUSION GROUP WHICH WE USE TO GENERATE FUSED KERNEL.*

`CustomFuseGraph` uses `AliasDb` to track data dependency across the `FusionGroup` when merging nodes, which ensures that we can safely cast inputs to reduced precision without leaking. We then immediately inlining custom FusionGroup back into the original graph. This psuedo-fusion would be transparent and would not interfere following optimization passes.

Comparing to the baselin approach, using `CustomFuseGraph` has the following advantages:
- Simplifies the code;
- Has the potential to benefit from future dependency analysis improvement (maybe we will be extending support to in-place operators).

3. Hierarchical blocks

One limitation of existing pieces: `CustomFuseGraph` and `AliasDb` does not work across hierarchical blocks. This means white nodes fusion would also be limited within a block. This could be addressed by having an additional pass that propagates casts across blocks and fuses (cancels out) consecutive casts.

A complication in JIT IR: hierarchical blocks inside `prim::If` and `prim::Loop` requires some attention when we trace the edges connected to them.:
`prim::If` has two nested blocks that does not have input, as we trace the outputs to `prim::If` nodes, we should redirect that to corresponding `prim::Return` on both blocks;
`prim::Loop` has a single block with input and output, special care should be given at: a. Offset of each tensor as input and output for the node and block is different; b. Loop-carried dependency/dataflow should be considered.

## Future topics

<!-- Add any other context or screenshots about the feature request here. -->
1. JIT pass API

- How to exposure this optimization pass properly? Currently we will just expose this as a manual jit pass under `torch._C._jit_pass_auto_casting`

- How to handle Backward? If we apply our optimization pass before the gradient computation, cast op would be automatically handled, and operations not in `differentiableGraph` in the backward path would also benefit from reduced precision.

2. More complicated graph partition

- NVIDIAN folks have developed more complicated painting logic in TensorFlow-AMP, where they incorporated graylist nodes, which they claim to cause numerical issues for downstream blacklist nodes. We could work on our partition policy if we ever observe similar issues.

3. Compatibility with eager mode scripting

- #25081 proposed to use `With enable_autocasting():` to expose explicitly control to the user. How can we expose this control through graph mode? With statement is not yet fully supported inside script.


cc @suo