Skip to content

馃悰 [Bug] slice does not clamp a finite stop on a dynamic axis, baking the literal stop into the engine#4607

Description

@SrivastavaKshitij

Bug Description

slice_op in py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py :: slice_op clamps the
slice stop bound to the axis length only when the axis is static. On a dynamic axis it takes
the else branch, and that branch's inner condition covers only four shapes of input: a
non-int (ITensor) start/stop, a negative start/stop, an omitted stop (stop_dynamic_None),
and stop == sys.maxsize. A plain finite positive int stop matches none of them, so control
falls straight through the entire dynamic-shape block:

if input.shape[dim] != -1 and isinstance(start, int) and isinstance(stop, int):
    start = get_positive_dim(start, input.shape[dim])
    stop = get_positive_dim(stop, input.shape[dim])   # static path: min(stop, dim_size)
    start_slice[dim] = start
else:
    # the start and stop or None is dynamic along dim or or start or stop is an ITensor
    if (
        not (isinstance(start, int))
        or not (isinstance(stop, int))
        or start < 0
        or stop < 0
        or stop_dynamic_None
        or stop == sys.maxsize
    ):
        ...            # builds a runtime-computed extent and returns from inside here
        return layer.get_output(0)

output_shape[dim] = math.ceil((stop - start) / step)   # reached with a FINITE stop

The result is a constant extent baked into the ISliceLayer on an axis whose real extent is
only known at runtime. Torch's own semantics are that a slice stops at the end of the tensor,
so the two agree only while the axis happens to be at least stop long.

The converter is silent about this. It neither logs nor raises, and nothing in the Python
traceback names slice/ops.py. What surfaces is a TensorRT shape-analysis error during engine
build naming the slice layer, and then a bare AssertionError from assert cuda_engine in
_TRTInterpreter.run. Identification therefore rests on the layer identity in the TensorRT
message plus two controls, both included in the reproducer below:

  • the same finite stop on a static axis compiles, because get_positive_dim clamps there;
  • an omitted stop on the same dynamic axis compiles, because stop == sys.maxsize is one of
    the conditions the dynamic branch does test. x[:, 5:] reaches the converter as
    stop = INT64_MAX, which equals sys.maxsize; the reproducer passes sys.maxsize explicitly
    so the control does not depend on that detail.

One further limitation of what TensorRT can catch, stated up front: only a stop above the
optimization profile's kOPT extent is rejected. A stop below the example extent (case B in
the reproducer, stop=48 with an example extent of 64) compiles cleanly and matches eager for
that input, yet the engine still carries the fixed extent 48 and is wrong for any runtime input
shorter than 48. TensorRT cannot detect that, because the engine is valid for its profile.

To Reproduce

Steps to reproduce the behavior:

docker run --rm --gpus all --ipc=host -v "$PWD":/w -w /w \
        nvcr.io/nvidia/pytorch:26.07-py3 python repro.py
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

repro.py

repro.py

import sys
import traceback
from typing import Final

import torch
import torch_tensorrt  # noqa: F401  (registers the "tensorrt" torch.compile backend)

# The example extent of the dynamic axis, which is also the profile's kOPT extent.
HINT_EXTENT: Final[int] = 64


class FiniteStop(torch.nn.Module):
    """x[:, :stop] on a runtime-length axis, with a plain Python int stop."""

    def __init__(self, stop: int) -> None:
        super().__init__()
        self.stop = stop

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.ops.aten.slice.Tensor(x, 1, 0, self.stop) * 2.0


class OmittedStop(torch.nn.Module):
    """Control: x[:, 5:], which aten canonicalizes to stop=INT64_MAX=sys.maxsize.

    `stop == sys.maxsize` is one of the conditions the dynamic branch does test for, so
    this identical-shaped slice builds a runtime-computed extent and compiles.
    """

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.ops.aten.slice.Tensor(x, 1, 5, sys.maxsize) * 2.0


def run_case(label: str, model: torch.nn.Module, x: torch.Tensor, dynamic: bool) -> bool:
    """Compiles `model` for `x`, marking dim 1 dynamic if asked. True if it worked."""
    torch._dynamo.reset()
    x = x.clone()
    if dynamic:
        torch._dynamo.mark_dynamic(x, 1)
    model = model.eval().cuda()
    optimized = torch.compile(
        model,
        backend="tensorrt",
        options={"pass_through_build_failures": True, "min_block_size": 1},
    )
    print(f"\n===== {label} =====", flush=True)
    try:
        out = optimized(x)
        torch.testing.assert_close(out, model(x))
    except Exception:  # pylint: disable=broad-except
        traceback.print_exc()
        print(f"----- {label}: FAILED", flush=True)
        return False
    print(f"----- {label}: OK {tuple(out.shape)}, matches eager", flush=True)
    return True


def main(argv: list[str] | tuple[str, ...] = ()) -> int:
    """Runs a finite stop above and below the example extent, plus two controls."""
    del argv
    print(f"torch          {torch.__version__}")
    print(f"torch_tensorrt {torch_tensorrt.__version__}")

    x = torch.randn((2, HINT_EXTENT)).cuda()

    above = run_case("A: stop 100 on a dynamic axis", FiniteStop(100), x, dynamic=True)
    below = run_case("B: stop 48 on a dynamic axis", FiniteStop(48), x, dynamic=True)
    control_maxsize = run_case(
        "control: omitted stop on a dynamic axis", OmittedStop(), x, dynamic=True
    )
    control_static = run_case(
        "control: stop 100 on a static axis", FiniteStop(100), x, dynamic=False
    )

    reproduced = (not above or not below) and control_maxsize and control_static
    print(f"\nreproduced: {reproduced}")
    return 0 if reproduced else 1


if __name__ == "__main__":
    sys.exit(main(argv=sys.argv))

output

torch          2.13.0a0+9186a08b2c.nv26.07
torch_tensorrt 2.14.0a0

===== A: stop 100 on a dynamic axis =====
ERROR:torch_tensorrt [TensorRT Conversion Context]:IBuilder::buildEngineWithConfig: Error Code 4: API Usage Error (ISliceLayer [SLICE]-[aten_ops.slice.Tensor]-[slice_1]: [SLICE]-[aten_ops.slice.Tensor]-[slice_1]: ISliceLayer has out of bounds access on axis 1 In processCheck at /_src/optimizer/shapeof/graphShapeAnalyzer.cpp:1042 In analyzeShapes at /_src/optimizer/shapeof/graphShapeAnalyzer.cpp:2785)
CRITICAL:torch_tensorrt.dynamo.backend.backends:Halting compilation on build failure since pass_through_build_failures was specified as True. To return the default Torch implementation and avoid halting compilation on engine build failures, specify pass_through_build_failures=False.
Traceback (most recent call last):
  File "/w/repro.py", line 79, in run_case
    out = optimized(x)
          ^^^^^^^^^^^^
[...]
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_conversion.py", line 280, in interpret_module_to_result
    interpreter_result = interpreter.run()
                         ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 653, in run
    assert cuda_engine
           ^^^^^^^^^^^
torch._dynamo.exc.BackendCompilerFailed: backend='tensorrt' raised:
AssertionError:

----- A: stop 100 on a dynamic axis: FAILED

===== B: stop 48 on a dynamic axis =====
WARNING:torch_tensorrt.dynamo.partitioning.common:Dynamic input arg1_1 (shape: torch.Size([2, s27])) has no max bound for dim 1, attempting to use a sane default (max: min(48) * 2^12). Please set an upper bound using torch._dynamo.mark_dynamic or torch.export.Dim
----- B: stop 48 on a dynamic axis: OK (2, 48), matches eager

===== control: omitted stop on a dynamic axis =====
----- control: omitted stop on a dynamic axis: OK (2, 59), matches eager

===== control: stop 100 on a static axis =====
----- control: stop 100 on a static axis: OK (2, 64), matches eager

reproduced: True

Expected behavior

A finite stop on a dynamic axis should be clamped to the runtime extent of that axis, the same
way get_positive_dim clamps it on the static path. Concretely, the dynamic branch should also
be taken for a finite positive int stop, and the emitted stop should be
min(stop, dim_extent) computed from the input's shape tensor. Then x[:, :100] on an axis of
runtime length 64 would produce 64 columns, as it does in Torch, instead of a slice layer that
reads out of bounds.

If clamping on the dynamic path is not something you want to add, the acceptable alternative is
for the converter to DECLINE the node (so the partitioner falls back to Torch for that slice)
rather than emit a layer that is only correct for a subset of the profile.

Separately, and independently useful: the diagnostics here are poor. The converter makes a
silent decision, the TensorRT error is only logged, and the exception that actually reaches
Python is a bare AssertionError with an empty message from assert cuda_engine in
_TRTInterpreter.run. Please consider raising an exception that carries the TensorRT builder
error text, so a user does not have to correlate a log line with an assertion.

Environment

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

  • Pytorch ngc container : 26.07-py3

Additional context

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions