Repository: pytorch
Issue: [TorchScript JIT] `dtype` propagation pass (meta-issue)
Labels: triaged

## 🚀 Feature

This pass will propagate `dtype` properties from input tensor descriptors throughout the graph.

## Motivation

TorchScript supports a generic Tensor type that can be of arbitrary shape, dtype, and layout (i.e., generic tensors).  While graphs operating on generic tensors are more expressive, they lack important shape, dtype, layout properties for further optimizations for both 3rd party optimizers as well as our internal kernel code generator such as NNC. 

We observe that meta-properties of generic tensors can be _inferred_ from those of input tensor by propagating `dtype` properties along the control-flow grow provided that a `dtype` transfer function is defined for each type of TorchScript IR.  In this feature, we will provide `dtype` propagation as a graph-level pass. In subsequent features we will provide utility functions that allow developers to query the `dtype` of an IR given input tensor `dtype`s.

Note that we have a separate proposal to allow TorchScript users describe properties of model inputs via a mechanism called **Input Descriptor** that describe properties (e.g., shape, rank, dtype, layout, values) of model inputs. 

## A toy IR 

We use a toy TorchScript IR to model the propagation algorithm. With only 10 IRs, it models the main complexity of TorchScript IRs, i.e., control-flow, in-place ops, view ops. Note that the modeling of List[Tensor] is a case where tensors are stored to and retrieved from the heap. Therefore we need to model container heaps with some precision to avoid property blackhole (i.e., dtype of a tensor is lost after being added to a list). 

```
// basic ops
prim::Constant

// arithmetic ops
aten::add // funcional binary operator that may require implicit type promotion
aten::add_ // inplace binary operator that may require implicit type promotion
aten::conv2d // functional operator where output dtype equals that of one of its input

// List[Tensor] ops
prim::ListConstruct
aten::append

// view ops
aten::view() // view ops

// control-flow ops
prim::if
prim:loop
prim::Call
```

## Tensor dtype transfer functions

This doc ([here] (https://github.com/nairbv/pytorch/blob/promote_types_strict/docs/source/tensor_attributes.rst)) describes the tensor dtype promotion rules  (last updated 08/2019). The actual implementation is update_result_type_state  in aten/src/ATen/native/TypeProperties.cpp ([here] (https://github.com/pytorch/pytorch/blob/5fb1142702320f0d52a3d87a94ab4c93220013c9/aten/src/ATen/native/TypeProperties.cpp#L79)) and `TensorIteratorbase::compute_types`.

For both internal and custom aten ops, most output dtypes can be expressed as a transfer function of some input dtypes. The typical dtype rules fall into the following categories

- Specify the output dtype as either *a static type* (e.g., float) or as *the dtype of one of its inputs*
- Specify a conservative dtype as *a set-union* of either a static type or dtype of some of its input
-  *Use type rule configuration variables* to describe dtype rules, as shown in the example below:
```
    void TensorIteratorBase::build_binary_float_op(const Tensor& out, const Tensor& a, const Tensor& b) {
      build(TensorIteratorConfig()
         .set_check_mem_overlap(true)
         .add_output(out)
         .add_input(a)
         .add_input(b)
         *.allow_cpu_scalars(true)
         .promote_inputs_to_common_dtype(true)
         .cast_common_dtype_to_outputs(true)
         .enforce_safe_casting_to_output(true)
         .promote_integer_inputs_to_float(true**)*);
    }
```

If an op (either internal or custom) has a really special dtype rule that can not be expressed above (e.g., the output dtype may depend on the value of an input). Then we may need to define a custom dtype transfer function for that (future work)

## Handling Control-flow

Since TorchScript IR is represented in SSA form, there is an explicit join of control-flow to the definition of a variable. We use the dtype lattice to handle such control-flow join points

### `prim::if`

In SSA representation, `prim::if` serves as phi-nodes for control-flow join. Consider the following example
```
def f(a, b: bool):
    x = a
    if b:
        x = a+1
    return x
```
The graph IR generated is as follows. In the SSA representation, there is only a single definition of %x  as the join of %x.1 and %a.1, whereas x is defined in two places in the source code.
```
  Scripted graph: graph(%a.1 : Tensor,
      %b.1 : bool):
  %6 : int = prim::Constant[value=1]() # test.py:8:14
  %x : Tensor = prim::If(%b.1) # test.py:7:4
    block0():
      %x.1 : Tensor = aten::add(%a.1, %6, %6) # test.py:8:12
      -> (%x.1)
    block1():
      -> (%a.1)
  return (%x)
```
The transfer function for prim-if would be a set union. For instance, in the example above, 
```
%x.dtype = {%.a.1.dtype, %x.1.dtype}
```
Note that, given the small number of distinct types, the set-union may be implemented efficiently using integers.

### prim::Loop

Loop construct is also represented in an SSA form. For output tensor types, we can compute a transitive closure of `dtype` if an output tensor is also used as input tensor in the loop form. 

### prim::Call

If we assume no cycle in TorchScript call graphs and that call targets are statically resolved, we can analyze prim::Call at every call-site using the dtype properties of the actual arguments to the call sites. This would yield a context-sensitive dtype proprogation.

## Handling in-place mutation and aliasing

- Since out_variant of aten ops do not change `dtype`, they can be handled trivially.  
- Mutation operations (e.g., due to aliasing) known to overwrite the entire tensor region of the output tensor, then the dtype transfer function is the same as that of the functional variant. Otherwise, the transfer function becomes
```
output.dtype = output.dtype set-union op-transferfunction(...)
```
- For `List[Tensor]`, if there are `aten::append`, then the dtype of the new tensor needs to be added (as set union) to that of the list of tensors. 

## Other extensions

### dtype propagation utility function

Apart from a pass that propagates tensor properties (expressed as input descriptors) to the entire graph, we may also provide a utility function that would compute output tensor properties based on input tensor properties. Such utility method may be invoked on-demand by other optimizations (e.g., when inserting a cast operation in the graph for AMP optimization). The exact interface is not yet defined but the propagation transfer function should be the same as used in the dtype propagation path

### other tensor property propagations

The algorithm is designed such that it can be easily extended to propagate other types of tensor properties (e.g., is_contiguous, memory_layout, device). 

cc @gmagogsfm