Dynamic shape support in dynamo

[peri044]

Dynamic shape support in dynamo

TL;DR

In line with Torchscript front end, we should support dynamic shapes in dynamo workflow by using torch_tensorrt.Input(min_shape=<>, opt_shape=<>, max_shape=<>) and passing it to torch_tensorrt.compile.

Goals

The current workflow for TS is really convenient for users to input dynamic shapes to the model. Using the same workflow will be helpful for users to port their code to dynamo easily.

input = torch_tensorrt.Input(min_shape=(1, 3, 224, 224), opt_shape=(1, 3, 224, 224), max_shape=(1, 3, 224, 224))
trt_model = torch_tensorrt.compile(model, ir="dynamo", [input])

Usecase that we can't support directly is from torch.compile.

One workflow for dynamo is

inputs = torch.randn(shape) # Inputs in pytorch tensors
trt_model = torch.compile(model, backend="torch_tensorrt")
trt_model(inputs)

This can't be used since there is no clean way to provide dynamic shape info for inputs.

Proposed APIs / UX / Internal implementation

We would create input_to_dynamic_range mapping and pass it as an attribute to CompilationSettings object. input_to_dynamic_range can have keys named input0, input1 etc.
This mapping should be used in conjunction to the inputs list here https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/dynamo/backend/conversion.py#L29 to create torch_tensorrt.Input objects which have the dynamic shape information.

Details specific for TorchScript Support

Implementation exists.

Details specific for FX support

No changes here. This support already exists in fx_ts_compat

Implementation Phases

• Implement the above feature
• Add additional test cases in our current models which have dynamic shape inputs.

Prototype - Medium

MVP 2.0 - Small / Medium

Future work / Extensions

• Explore dynamic flag in torch.compile to see if we can make use of the symbolic shapes it returns.

[gs-olive]

This plan overall looks good to me - I have a few comments/clarification requests. Number 1 and 3 are similar.

1. For intermediate subgraphs which we turn into TRT engines, how would we determine the ranges of the input tensors from the user-provided shapes, in both the compile and export paths?

• Would we use Torch's SymInt/FakeTensor right away, or begin with something similar to what we have in TorchScript, where we run a shape analysis pass on each min, opt, and max shape?

2. Does the fact that users can specify options to the backend for torch.compile change the usecase we can't support? See PR #2005 and the example below.

torch.compile(model, backend="torch_tensorrt", options={"min_block_size": 4, ...})

3. With the amount of graph modification that Dynamo + AOT perform, what is a reliable way to map subgraph inputs to user-provided inputs?

• Could we use a _get_name() function, or some Tensor attribute, or set a name like input_0, input_1, etc. as you had suggested, which is preserved through the passes?

[narendasan]

We should be adding support for disjoint optimization ranges which is something we missed in TorchScript