Repository: pytorch
Issue: [RFC] A PyTorch Tensor Shape DSL For Symbolic Shape Inference
Labels: oncall: jit

# A Tensor Shape DSL For PyTorch

### What do we want ?

* Propagate arbitrary information about tensors shapes on JIT IR, such as input `C, H, W = 3, 32, 32` but that the batch sizes is dynamic
* Reason about operators even when input shapes aren’t known (like with `conv2d`) to eliminate control flow on JIT IR
* Prove Symbolic Shape relationships between tensors to generate fast code for dynamic shapes and do memory planning
* Have separable shape functions that can be invoked at runtime and which do error checking with comprehensible error messages, and which only compute information that is not known statically. We want to be able to compile or run the separable shape function using existing tools.
* Statically check for shape mismatches through symbolic shapes
* PyTorch developers without compiler experience to easily be able to add shape information to a new operator, and have the same easy experience extend to custom operators
* Take inputs and generate what weight sizes should be, and take weights and generate what the input sizes should be
* A recognizable representation of shape functions for users of our IR that will cover the next set of requirements we don’t know today 
* A way to debug shape functions without too much hassle

### How are we going to get there ? 

A Tensor Shape DSL in TorchScript, and a lot of partial evaluation & symbolic analysis!

Here are a couple examples of C++ Shape Functions and their translation to TorchScript. Conv and Broadcasting are two typically challenging shape functions because Conv involves non-Tensor input arithmetic and shared functions that compute over arbitrary dimensions, and Broadcasting can be difficult for languages to express with unknown dimensionality inputs and conditional expressions. 

**Convolution:**

```
# `check_shape_forward` elided for brevity
static inline std::vector<int64_t> conv_output_size(
    IntArrayRef input_size, IntArrayRef weight_size,
    IntArrayRef padding, IntArrayRef stride, IntArrayRef dilation = IntArrayRef()
) {
  // ASSERT(input_size.size() > 2)
  // ASSERT(input_size.size() == weight_size.size())
  bool has_dilation = dilation.size() > 0;
  auto dim = input_size.size();
  std::vector<int64_t> output_size(dim);
  output_size[0] = input_size[input_batch_size_dim];
  output_size[1] = weight_size[weight_output_channels_dim];
  for (size_t d = 2; d < dim; ++d) {
    auto dilation_ = has_dilation ? dilation[d - 2] : 1;
    auto kernel = dilation_ * (weight_size[d] - 1) + 1;
    output_size[d] = (input_size[d] + (2 * padding[d - 2]) - kernel) / stride[d - 2] + 1;
  }
  return output_size;
}
```

```
@torch.jit.script
def conv_output_shape(input: List[int], weight: List[int], stride: List[int], padding: List[int], dilation: List[int], groups: int):
    check_shape_forward(input, weight, bias, stride, padding, dilation, groups)
    
    has_dilation = len(dilation) > 0
    dim = len(input)
    output_size: List[int] = []
    input_batch_size_dim = 0
    weight_output_channels_dim = 0
    output_size.append(input[input_batch_size_dim])
    output_size.append(weight[weight_output_channels_dim])
    
    for d in range(2, dim):
        dilation_ = dilation[d - 2] if has_dilation else 1
        kernel = dilation_ * (weight[d] - 1) + 1
        output_size.append(input_size[d] + (2 * padding[d - 2]) - kernel) // stride[d - 2] + 1)
    
    return output_size
    
@torch.jit.script
def conv2d(input: List[int], weight: List[int], stride: List[int], padding: List[int], dilation: List[int], groups: int):
    assert len(weight) == 4
    return conv_output_shape(input, weight, bias, stride, padding, dilation, groups)
```


**Broadcasting:**

```
std::vector<int64_t> infer_size(IntArrayRef a, IntArrayRef b) {
  size_t dimsA = a.size();
  size_t dimsB = b.size();
  size_t ndim = dimsA > dimsB ? dimsA : dimsB;
  std::vector<int64_t> expandedSizes(ndim);

  // Use ptrdiff_t to ensure signed comparison.
  for (ptrdiff_t i = (ptrdiff_t)ndim - 1; i >= 0; --i) {
    ptrdiff_t offset = ndim - 1 - i;
    ptrdiff_t dimA = dimsA - 1 - offset;
    ptrdiff_t dimB = dimsB - 1 - offset;
    int64_t sizeA = (dimA >= 0) ? a[dimA] : 1;
    int64_t sizeB = (dimB >= 0) ? b[dimB] : 1;

    TORCH_CHECK(
        sizeA == sizeB || sizeA == 1 || sizeB == 1,
        "The size of tensor a (", sizeA,
        ") must match the size of tensor b (", sizeB,
        ") at non-singleton dimension ", i);

      // 1s map to the other size (even 0).
      expandedSizes[i] = sizeA == 1 ? sizeB : sizeA;
  }

  return expandedSizes;
}
```



```
@torch.jit.script
def infer_size(a: List[int], b: List[int]):
    dimsA = len(a)
    dimsB = len(b)
    ndim = max(dimsA, dimsB)
    expandedSizes : List[int] = []

    for i in range(ndim):
        offset = ndim - 1 - i
        dimA = dimsA - 1 - offset
        dimB = dimsB - 1 - offset
        sizeA =  a[dimA] if (dimA >= 0) else 1
        sizeB = b[dimB] if (dimB >= 0) else 1

        if sizeA != sizeB and sizeA != 1 and sizeB != 1:
            raise Exception("The size of tensor a {} must match the size of tensor b ("
                            "{}) at non-singleton dimension {}".format(sizeA, sizeB, i))
        
        expandedSizes.append(sizeB if sizeA == 1 else sizeA)

    return expandedSizes
```


As you can see, the current implementation of shape functions are very easily translated 1-1 to TorchScript. If you are familiar with writing shape functions in C++ and you are also familiar with python, this should not be difficult. It’s not an initial goal but it also leaves open the possibility of writing your shape functions first in TorchScript, and then generating optimized C++ for use in the C++ kernel. 


### How does it work ?

It uses all the usual suspects that clean up IR (Constant Propagation, etc) but the two main passes that give it its juice are a pass that substitutes input Tensor properties into the IR and a pass that rewrites mutation on lists into a functional form that can be optimized. Let’s say we are trying to propagate the Tensor Shape of `aten::mul` (a broadcasting op) between an input with shape `(1, 1, 5, 4)` and an input with shape `(*, 6, 5)`. First,  we are able to substitute `len(a)` with 4 and `len(b)` with 3. Then, we are able to fully unroll the main loop. After unrolling the loop, we have expressions like `b[2]` which the input Tensor property substitution pass can replace with `5`. The calculation has a series of appends to `expandedSizes` and we can rewrite each pair of `li = [values]; li.append(3)` to `li = [values, 3]`. Finally, there is a pass to extract the output tensor shape which will recognize `[1, b[0], 5, 4]` as 3 concrete dimension values and the symbolic shape from the first dimension of `b`. You can click the gif below to look at the code dump of the iterative optimization. 

https://imgur.com/a/Va0N0bS


### How do symbolic shapes work ?

Symbolic Shapes are represented as negative dimension values. The only property that they have is two symbolic shapes of the same value must be equal. You can not compare them to any other symbolic shape or use them in arithmetic. 

Consider the broadcasting logic of two of the same symbolic shapes

```
def broadcast(symbolicA: int, symbolicB: int):
   if symbolicA != symbolicB and symbolicA != 1 and symbolicB != 1:
        raise Exception("The size of tensor a {} must match the size of tensor b ("
                            "{}) at non-singleton dimension {}".format(symbolicA, symbolicB, 0))
        
    return symbolicA if symbolicB == 1 else symbolicA
```


The pass that substitutes input properties will see that symbolicA and symbolicB resolve to the same symbolic shape, and replace all uses of a particular symbolic shape with one Value *. In the example above, we replace `symbolicB` with `symbolicA`. The IR becomes

```
def broadcast(symbolicA: int, symbolicB: int):
   if symbolicA != symbolicA and symbolicA != 1 and symbolicA != 1:
        raise Exception("The size of tensor a {} must match the size of tensor b ("
                            "{}) at non-singleton dimension {}".format(symbolicA, symbolicA, 0))
                                    
    return symbolicA if symbolicA == 1 else symbolicA
```

From here, it becomes very easy to reason that the output will be `symbolicA` because `symbolicA != symbolicA` and `symbolicA if symbolicA == 1 else symbolicA` can be optimized away. 


### Symbolic Shapes Across Equivalent Operators ?

If you have two adjacent `Conv` nodes with the same weights and inputs, you might want to prove that the output of both nodes will have the same shape properties. This works for the batch dimension and channels out since they are directly taken from the input/weight, but not for Hout and Wout just by forward propagation on each node. However, the symbolic shape relationship between the different Tensor outputs is recoverable by joining adjacent `aten` nodes, merging their partially evaluated shape functions into one shape compute graph, and running optimization passes like CSE on it. This can also work through caching of output shape properties for nodes with the same inputs, which also helps as an optimization to speed up shape inference. 

### Why do you need a runtime executable shape function ?

Let’s say you have two input tensors with symbolic shapes, `x` and `y`. If you are calculating the following fusion group:

```
out1 = torch.sigmoid(torch.tanh(x))
out2 = torch.relu(y)
return torch.hardswish(torch.concat((out1, out2))
```

Even if you ahead of time know that every expression in `torch.sigmoid(torch.tanh(x))`, `torch.relu(y)` and `torch.hardswish(...)` have the same symbolic shape, you still need to compute the concat’d shape of `x` and `y`. For this you can merge the shape functions of each individual op into one shape compute graph that is then executed at runtime. When we can prove that there are no exceptions, we can even inline the partially evaluated shape compute graph directly into the NNC kernel, since it’s just TorchScript IR. NVFuser folks said that on simple fusions of dynamic shapes, the slow expression evaluation engine they are using can be just as expensive as the kernel itself. We can also directly inline the shape functions into the Graph to deploy compiled kernels with dynamic shapes on mobile.

You could imagine using the partially evaluated large shape compute graph for other uses as well, such as at the beginning of every invocation of Static Runtime quickly calculating the complete shapes of the graph, and fully allocating one slab of memory. 


### What about other use cases, like bidirectional inference ?

We are capturing the operator shape properties in their full generality, which leaves open other use cases, even if they are not using TorchScript IR to do the analysis. To that end, I’ve a prototyped a TorchScript → Z3 transpiler that works for the scalar math and fixed length lists that occurs in shape analysis. Because the IR becomes very simple after the optimization passes, it took less than a day to write. You can take a look at it [here](https://gist.github.com/eellison/d7f538e3455b2604fbb832a4882efe9a). In example I’ve taken the partially evaluated shape compute graph from specifying weights and other non-input values to the `Conv` shape function, and used it to generate a new set of Conv inputs. Given `nn.Conv2d(16, 33, (3, 5), stride=(2, 1), padding=(4, 2), dilation=(3, 1))` it correctly generates an input of `[1, 16, 8, 12]`. 


### Why not use z3 directly ?

* Not an acceptably easy experience for a developer (PyTorch core or custom operator) to add a new shape function
* Does not support exceptions
* Runtime execution of symbolic shape logic may be slow and build size would preclude dynamic shapes on mobile
* Does not leave open possibility of consolidating PyTorch C++ shape logic with shape formulas
* It’s very difficult to write for non-fixed length inputs, which makes it hard to reuse logic across ops like `Conv1`/`Conv2d`/`Conv3d`.
* I think it would be easier to write and maintain TorchScript shape functions for operators and a TorchScript→Z3 transpiler than to write them directly in Z3, especially for user registered operators. 


### Next Steps 
 
* I have a very barebones [initial PR out](https://github.com/pytorch/pytorch/pull/54809) and I and other JIT developers are working on landing the full version soon
* The concrete use cases which I plan to use it for are enabling dynamic shapes in the JIT with NNC, shape inference to cleanup the graph, calculating memory reuse / high water memory mark, making a single allocation of memory in inference, and trying to do AOT compilation with annotated shape information of inputs from the user
* The other uses cases I have outlined I am not personally planning on working on but I am interested in helping and talking to anyone interested in working on any of these things

Related issues: https://github.com/pytorch/pytorch/issues/52753, https://github.com/pytorch/pytorch/issues/46948, https://github.com/pytorch/pytorch/issues/43501, https://github.com/pytorch/pytorch/issues/40373, https://github.com/pytorch/pytorch/issues/26889






cc @gmagogsfm