TensorRT doesn't work with Convolution with dynamic batch in some cases

[pranavm-nvidia]

Below is a minimal reproducer:

import nvtripy as tp


kernel_dims = (3, 3)
conv = tp.Conv(
    128, 128, kernel_dims, (1, 1), padding=[(dim // 2, dim // 2) for dim in kernel_dims], groups=128, bias=False
)
conv.weight = tp.ones((128, 1, *kernel_dims))

tp.logger.verbosity = "ir"


def func(x):
    # x.shape is (B, 128, 7, 7)
    B = x.shape[0]
    x = tp.reshape(x, (B, -1, 7, 7))
    out = conv(x)
    return out


# This works (or any other static batch size):
compiled_func = tp.compile(func, args=[tp.InputInfo([(10, 10, 10), 128, 7, 7], dtype=tp.float32)])
# This fails:
compiled_func = tp.compile(func, args=[tp.InputInfo([(1, 10, 10), 128, 7, 7], dtype=tp.float32)])

The error is:

MTRTException: failed to run pass pipeline
    ITensor::getDimensions: Error Code 4: API Usage Error ([tensorrt.convolution] %t86,%t2;;<out>;;%t87: IConvolutionLayer `num_groups` must divide input's channel count. Input channel count is 1 but `num_groups` is 128.)
    (%t87) error: ranks of MLIR Tensor is 4 while TRT ITensor has rank -1

[pranavm-nvidia]

My best guess is that when the batch dimension is dynamic, the channel dimension becomes "dynamic" as well (because it's being inferred by -1) - though obviously in this case it's not truly dynamic. The new channel dimension will always be:

C = (B * C * 7 * 7) / (B * 7 * 7)

which is clearly always just C. Yet somehow, getDimensions() can't resolve this and assumes the dimension is 1.