Dynamic shapes for torch.export

[peri044]

This RFC aims to enhance our dynamic shape support and close the gaps to fully compile a model with dynamic shapes.

Two backends we support :

torch.export

We already support dynamic shapes via Input API frontend. The left over tasks are in this story. #2274

torch.compile

The proposal is to pass torch_tensorrt.Input API objects via backend options as follows:

inputs = torch_tensorrt.Input(min_shape=(1, 3, 4),
                                  opt_shape=(4, 3, 4),
                                  max_shape=(8, 3, 4),
                                  dtype=torch.float32,
                                  name='x')
compile_spec = {"inputs": [inputs],
                "debug": True,
                "min_block_size": 1}
# Torch compile path
trt_model = torch.compile(model, backend="tensorrt", dynamic=False, options=compile_spec)
trt_model(torch_inputs)

These input objects pass the (min, opt, max) range to TensorRT conversion.

• Note 1: Make sure the number of aotautograd sample inputs and provided TRT inputs match. Also, there should be correspondence or validation between these two.
• Note 2: Check if the aotautograd inputs have a name and if they do, this name should be used for declaring dynamic torch_tensorrt.Inputs.

Case 1: No fallback case

If the graph is end-end compiled by TRT, recompilation shouldn't happen if the input shape is changed during inference. Explore recompilation triggers and when this happens.

Case 2: Fallback

1. If the graph has PyT -> TRT -> PyT .... partitions, ensure recompilation is not triggered due to pytorch segments. By default, dynamic=True should handle these cases hopefully for us.

[peri044]

1. [torch.export] Given an mapping from user inputs and our mapping, how to build a TRT engine ?
   torch_tensorrt.Input(min, max, opt) -> sym_ints are created in pytorch graph.
   In full graph, we dummy run the graph to get shapes - get rid of this.
   Use sym_ints directly to estimate the intermediate tensor shapes. Real proof this works is testing data dependant graphs. In such cases, shape tensors are the only way we can estimate tensor shapes.

2. [torch.export] Provided there is no graph breaks, how do we match the torch tensorrt Inputs and inputs to the graph ?

   • What if there is a collections input (dictionary- key with strings) ?

3. In torch.export, if there are subgraphs partitioning, how do we handle shape propagation b/w TRT and PyT boundaries ?
   SymInt tells us the info around subgraphs. Maybe name - shape mapping will help us ?

4. In torch.compile, no direct access to user inputs. How do we get initial user shapes ? Do some operations on torch_tensorrt.Input to get the shape for subgraph ?
   Check dynamic=None and use mark_dynamic API instead.

5. Adding guards to avoid recompilation ? If it's out of of the range, what shape would be needed to recompile ?

6. cache the entire user settings for each engine along with the each engine

[peri044]

After prototyping 1 and 2, converters validation is required for dynamic shapes.
For converters, we need validator for dynamic shapes.