🐛 [Bug] Encountered bug when compiling model with nested Input type using `input_signature` argument

[gs-olive]

Bug Description

When compiling a model with input type Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], compilation fails with the following error:

RuntimeError: [Error thrown at core/ir/ir.cpp:46] Expected vals.size() == specs.size() to be true but got false
Expected dimension specifications for all input tensors, but found 1 input tensors and 2 dimension specs

To Reproduce

Steps to reproduce the behavior:

1. Define a Torch model with forward function having the following form:

def forward(self, input : Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]])

2. Define the Torch-TRT Inputs and compilation settings using input_signature, then compile the scripted model:

compile_settings = {
    "input_signature": (torch_tensorrt.Input((5, 5), dtype=torch.float), (torch_tensorrt.Input((5, 5), dtype=torch.float), torch_tensorrt.Input((5, 5), dtype=torch.float)),),
    "enabled_precisions": {torch.float},
    "truncate_long_and_double": True,
}
trt_ts_module = torch_tensorrt.ts.compile(scripted_model, **compile_settings)

Expected behavior

Model should compile with input signature containing nested Tuple collection.

Environment

• Torch-TensorRT Version: 1.4.0.dev0+f43be5b6
• PyTorch Version: 1.14.0.dev20221114+cu116
• CPU Architecture: Intel Xeon CPU
• OS: Ubuntu 20.04
• How you installed PyTorch: pip
• Build command you used: python setup.py develop
• Are you using local sources or building from archives: local
• Python version: 3.8.13
• CUDA version: 11.6

Additional context

The case of a nested Tuple containing a singleton Tensor followed by a Tuple of Tensors is not parsed as expected by the input signature interpreter, whereas regular tuple Inputs of Tensors are supported, as in:

TensorRT/tests/py/api/test_collections.py

Lines 210 to 213 in b2a5da6

	compile_spec = {
	"input_signature": (
	[torchtrt.Input(self.input.shape), torchtrt.Input(self.input.shape)],
	),

Wrapping the input signature in another Tuple does not resolve the issue, and raises a different error message:

RuntimeError: forward() Expected a value of type 'Tuple[Tensor, Tuple[Tensor, Tensor]]' for argument 'input.1' but instead found type 'Tuple[Tensor]'.
Position: 1
Declaration: forward.forward(__torch__.module self_1, (Tensor, (Tensor, Tensor)) input.1) -> ((Tensor, Tensor, Tensor) 7)

[gs-olive]

Update

Error Origin

The origin of this compilation error derives from the graph inputs sourced from the input_signature argument. As per the documentation of input_signature, the signature is expected as a Python Tuple or List of input arguments to the model's forward function. This implies that the desired format of the inputs for the above case would be:

( (Tensor, (Tensor, Tensor)), )

With the above as the input format, however, the innermost Tensor object is considered to have nesting depth 3, which is disallowed here:

TensorRT/core/ir/GraphInputs.cpp

Lines 42 to 52 in 5fa6374

	if (level == 0) { // a single value like A
	collection_inputs.resize(1);
	collection_inputs[0].push_back(cur_input);
	} else if (level == 1) { // like A in [A, A] or [(B, B), A]
	collection_inputs[index].push_back(cur_input);
	} else if (level == 2) { // like A in [(A, A), C]
	collection_inputs[index].push_back(cur_input);
	} else { // only support 2 level
	LOG_ERROR(
	"Input nesting depth exceeds currently supported depth (3), use 1 level: [A, B], or 2 level: [A, (B, C)]");
	}

Since the initial Tuple coalescing the individual arguments in the input signature is considered as a nesting level, this implies that Tuple[Tuple[Tensor]] input constructs are disallowed as they violate the maximum nesting depth.

Resolution

Since this example did not derive from a specific Input, but was instead a general assessment of input_signature, resolving the error for this specific signature is not critical. One suggestion would be to elicit an error and halt the program when the input_signature nesting level is found to be invalid, so the program does not progress after encountering an invalid input.

A more involved fix might include restructuring the input signature format to allow for arbitrary input nesting depth. A suggestion on this approach would be to refactor the flattened_inputs and collection_inputs_ (featured in the snippet below) to be:

• A vector std::vector<torch::jit::IValue>, containing an ordered list of input objects to the model
• A map from torch::jit::IValue to std::vector<torch_tensorrt::core::ir::Input>, where each IValue is a single input to the model, and the vector stores a flattened list of Inputs

TensorRT/core/ir/GraphInputs.cpp

Lines 64 to 66 in 5fa6374

	GraphInputs::GraphInputs(torch::jit::IValue& input_signature_) {
	std::vector<torch_tensorrt::core::ir::Input> flattened_inputs;
	std::vector<std::vector<torch_tensorrt::core::ir::Input>> collection_inputs_;