Dynamic shapes + fallback

[peri044]

Dynamic Shapes with fallback

Goal(s)

Currently our partitioning system supports fallback to torch only for ops with static shapes. The goal is to enable automatic fallback for dynamic shape models.

Usecases

This is required in scenarios where customer has dynamic shapes in the model and

• There are some unsupported ops or supported ops (only with static shapes) which break Torch-TRT compilation today
• fallback certain ops to pytorch for performance/implementation reasons.

Proposed APIs / UX

Example Workflow

Today here is the current workflow in a high level

For every block

1. Generate random inputs based on user provided input shapes.

2. Segment the graph

3. Resolve Non Tensor inputs for TRT blocks

4. Register input and output for segmented blocks

5. Run shape analysis

The only stages that will be affected are 1) and 5).

Limitations

This will not address cases where the output shape of the subgraph is dependent on the input tensor data (a.k.a Data-Dependant shapes).

Internal Implementation

Design

1. Input data generation

std::unordered_map<const torch::jit::Value*, torch::jit::IValue> generateRandomInputs(
    std::unordered_map<const torch::jit::Value*, std::vector<ir::Input>>& inputs,
    std::unordered_map<const torch::jit::Value*, std::vector<c10::optional<at::ScalarType>>>& types)

The return value is a map which maps input tensor pointer to data tensor (with specified shapes).

Proposal :

Return type can be std::unordered_map<const torch::jit::Value*, std::set< torch::jit::IValue>

The set of IValues will contain min, opt and max shaped data tensors.

2. Run shape analysis

   In the getSegmentsOutputByRunning, we hold input data for a subgraph via std::vector<torch::jit::IValue> jit_inputs_ivalues;

   Simple extension is wrapping this and running it for (min, opt, max) shapes for a set of inputs. So the pseudo code would look like

   cur_mod = <defining module out of the graph>
   
   # Handle both static and dynamic case by iterating over (min, opt, max)
   for (int mode=0; mode < 3; mode++){
   	std::vector<torch::jit::IValue> jit_inputs_ivalues
   	for (auto& input : seg_block.raw_inputs()){
   		jit_inputs_ivalues.push_back(ivalues_maps[input][mode].toTensor());
   	}
       auto outputs = cur_mod.forward(jit_inputs_ivalues)
       size_t idx = 0;
       for (auto& output : seg_block.raw_outputs()) {
       	ivalues_maps[output][mode] = jit_results[idx++];
     	}
   }
   
   # Registering input shapes for the segmented blocks takes place now
   seg_block.register_inshapes(input_shapes);
   seg_block.register_intypes(input_types);
   
   # input_shapes will now be std::vector<std::set<ir::Input>>

Extensions Required to Core API implementations

The changes required would be the ones mentioned above

Data Structures

< Figure out if > Since each input ir::Input usage is now replaced with a set of three variants std::set< ir::Input >, we could typedef MinOptMaxInput = std::set< ir::Input >

Implementation Phases

Prototype

There's an existing prototype based on the above design. #1111

Implement the above design with the new centralized partititoning changes.

MVP (1.3)

The feature can be shipped if it passes the following

1. Dynamic shape models of existing torchvision models. Add these tests
2. Force fallback some ops in the above models and implement tests.
3. Test user/internal bugs with this feature

Extension Phase 1 (<TARGET RELEASE VERSION>)

Figure out how aten::size plays with this dynamic shapes.

Extension Phase 2 (<TARGET RELEASE VERSION>)

Explore how FX currently handles dynamic shapes + fallback. Is it easier to switch to FX for dynamic shapes + fallback ?

[narendasan]

MVP should be dynamic batch

[narendasan]

Don't use a set, there should already be support for the 3 fields. Also I think the run through shape analysis then clear results and run again is simpler.

[ncomly-nvidia]

@dheerajperi I'm assuming FX (extension 2) has been absorbed in the aten::size RFC, correct?