🐛 [Bug] slice_op not able to handle start value with type trt.ITensor

[shahe-nv]

Bug Description

In py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py, function slice_op, Line 112, it only considers the case when start is int, but not the case when start is trt.ITensor.

I tried to add support for start being trt.ITensor, however, I'm not able to run successfully because TRT requires the start being a shape tensor to be able to do shape inference but it seems start is not a shape tensor. Therefore, currently I'm only able to do a workaround by hard coding all the parameters of the slice layer.

To Reproduce

Steps to reproduce the behavior:

Expected behavior

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Torch-TensorRT Version (e.g. 1.0.0):
• PyTorch Version (e.g. 1.0):
• CPU Architecture:
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, libtorch, source):
• Build command you used (if compiling from source):
• Are you using local sources or building from archives:
• Python version:
• CUDA version:
• GPU models and configuration:
• Any other relevant information:

Additional context

[shahe-nv]

This is the workaround code I added at Line 112:

            if isinstance(start, int) and start < 0:
                start_slice = input.shape
                start_slice[dim] = -1 * start
            elif isinstance(start, trt.ITensor):
                if stop_dynamic_None or stop == sys.maxsize:
                    # This is a workaround assuming start = neg = -L and stop = None or maxsize
                    dim_size_1d = get_shape(
                        ctx, target, source_ir, name + "_start_dim_size", input, dim
                    )
                    # The workaround hard codes this:
                    # START for every dim = 0 (including the sliced dim).
                    # SIZE for sliced dim = L = dim_size_1d (take all elements).
                    start_parts: list = []
                    size_parts: list = []
                    for i in range(len(input.shape)):
                        start_parts.append(0)
                        if i == dim:
                            size_parts.append(dim_size_1d)
                        elif input.shape[i] == DYNAMIC_DIM:
                            size_parts.append(
                                get_shape(
                                    ctx, target, source_ir,
                                    name + f"_start_size_dim{i}", input, i,
                                )
                            )
                        else:
                            size_parts.append(input.shape[i])
                    start_slice_tensor = cat(
                        ctx, target, source_ir,
                        name + "_start_slice_concat",
                        tuple(start_parts), 0, cast_dtype=trt.int32,
                    )
                    size_tensor = cat(
                        ctx, target, source_ir,
                        name + "_size_slice_concat",
                        tuple(size_parts), 0, cast_dtype=trt.int32,
                    )
                    stride_slice_tensor = cat(
                        ctx, target, source_ir,
                        name + "_stride_slice_concat",
                        tuple(stride_slice), 0, cast_dtype=trt.int32,
                    )
                    layer = ctx.net.add_slice(
                        input, start=trt.Dims(), shape=trt.Dims(), stride=trt.Dims()
                    )
                    layer.set_input(1, start_slice_tensor)
                    layer.set_input(2, size_tensor)
                    layer.set_input(3, stride_slice_tensor)
                    return layer.get_output(0)

[narendasan]

@apbose please investigate or reassign to code owner