🐛 [Bug] Reshaping Empty Tensors Causes Compilation Crash (original) (raw)

Bug Description

Converters including aten::select which use TensorRT's IShuffleLayer can cause compilation failures when run with empty Tensors (those having a 0 in one or more of their dimensions). These tensors are often used in models which take caches as arguments to their forward functions.

The error displayed is:

DEBUG: [Torch-TensorRT - Debug Build] - Gather tensor shape: [1, 20, 0, 768] ERROR: [Torch-TensorRT TorchScript Conversion Context] - 4: [graphShapeAnalyzer.cpp::analyzeShapes::1872] Error Code 4: Miscellaneous (IShuffleLayer %3 : Tensor = aten::select(%0, %2, %1): reshape of empty tensor to non-empty tensor. Reshaping [1,20,0,768] to [20,0,768].)

To Reproduce

Steps to reproduce the behavior:

1. Define a network with a forward function taking a single tensor, and invoke torch.select(), returning the result.
2. Define an input shape with at least one zero dimension
3. Compile the model via TorchScript

Expected behavior

The model should compile successfully.

Environment

• Torch-TensorRT Version: b2a5da6
• PyTorch Version: 2.0.0.dev20230123+cu117
• 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.7

Additional context

The bug originates from the fact that 0 has a special meaning as a placeholder in the setReshapeDimensions function of the IShuffleLayer. Specifically, unless explicitly specified, zero denotes the size of the previous-most-significant aligned dimension. Thus, improved input checking is needed to specify 0 as not being a placeholder when the input tensor is empty.

Fixing this bug would also address #1616.