Support Unbounded Dynamism for torch.export

[lsy323]

🚀 Feature

The unbounded dynamic shape needs to be propagated through ops.

Scope:

• Only for export use case.

Example:

opcode         name           target                   args                     kwargs
-------------  -------------  -----------------------  -----------------------  --------
placeholder    arg0_1         arg0_1                   ()                       {}
placeholder    l_embeddings_  l_embeddings_            ()                       {}
call_function  sym_size_int   aten.sym_size.int        (l_embeddings_, 0)       {}
call_function  mul            <built-in function mul>  (sym_size_int, 2)        {}
call_function  expand         aten.expand.default      (arg0_1, [mul, -1, -1])  {}
output         output         output                   ((expand,),)             {}

In LazyIR, we need to capture the aten.sym_size.int and the subsequent arithmetic operations on the SymInt. So that the semantic can be lowered.

In the lowered HLO graph, we should have something like

main(%arg0: tensor<?x3x224x224xf32>, %arg1: tensor<1x1x768xf32>) -> tensor<?x1x768xf32> {
    %1 = get_dimension_size(%arg0, dim = 0) 
    %2 = expand(%arg1, %1, dim=0)
    return %2 : tensor<?x1x768xf32>
}

Rough Plan

• We need to trace and lower torch ops with SymInt output in LTC. (for aten.sym_size.int, which generates a SymInt)
• The arithmetic on the SymInt needs to be traced. This shouldn't be hard to achieve if the corresponding LazyIR node can be created for the SymInt.
• When the sym_int version of the op is lowered, it needs to retrieve the underlying LazyIR of the SymInt argument

Open questions

1. Would it make sense to handle both bounded and unbounded dynamism under the same workflow/infra? The source of the dynamic dim needs to be traced for unbounded dynamic case, but not for bounded dynamism.

Example

dynamic_dim = input.shape[0]
dynamic_dim = dynamic_dim * 2
expanded = input.expand([dynamic * 2, -1, -1])

Let's say input.shape[0] has a bound of <= 5. In bounded dynamism, only knowing the upper bound in the op is enough. In unbounded dynamism, the arithmetic on SymInt needs to be traced and lowered in LTC.

2. Not sure if there is any API to create an unbounded dynamic tensor, so the graph can be traced with it.

[JackCaoG]

@ezyang We are trying to support unbounded dynamic shape only for export.

The exported FX Graph is dynamic and has sym_size op that extracts the size of a tensor and store it in a SymInt. During LTC tracing, we are using static tensors. (Not sure if there is a way to create a tensor with unbounded dynamic dimension and trace with it) In this process, all tensors are static from PyTorch's perspective so it will try to collapse the symint into concrete integer instead of treating the size as a symbolic things. What's the correct way to make PyTorch always treat symint as unbacked and not collapsing it? We want to lower the symint to an IR in the final HLO.

[ezyang]

Are you dead set on lazy tensor for this, or can you be flexible?

Because if you can be flexible, I would tell you to use torch.export and you'd get a graph with non-collapsed dynamic dimensions and size compute for whatever you needed to be dynamic. And then life is good. To get it to work with LTC... ugh, I don't even wanna think about it lol.

[JackCaoG]

@ezyang we are already using the torch.export. The issue is we need to lower that exported FX to HLO. I am wondering how does inductor handles this, does inductor support unbounded dynamism?

[ezyang]

Inductor does just handle it! For example, when generating kernels for pointwise operations, we support symbolic ranges. We can also deal with symbolic indexing formulas and use sympy to simplify them.

[lsy323]

cc @miladm @vanbasten23

[lsy323]

#6653 Provides support with limitations to this uses case by applying a FX pass to group the sym_size and the following op into a single xla op, the op will be lowered with dynamism semantics in torch_xla. The limitation is if there are some arithmetics on the SymInt, the FX passes is not flexibly enough to capture those operations.

[lsy323]

Code contributions that enabled this unbounded dynamism:

Propagating dynamism information from PyTorch down to StableHLO

• Enable passing down dynamic dimensions from torch to XLA #5790

Support unbounded dynamism for torch ops

• support dynamism on add, mul #6443
• Enable unbounded dynamism on conv, softmax, addmm, slice #6494
• Add dynamism support to conv1d, view, softmax #6777

FX passes to to support dynamism in op argument

• Add fx passes to support unbounded dynamism in torch op arg #6653
• Add dynamism support to aten.embedding and aten.split_with_sizes #6781
• Run shape propagation for inserted fx nodes #6805