🐛 [Bug] slice does not clamp a finite stop on a dynamic axis, baking the literal stop into the engine
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3k
- Forks
- 410
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 78
Description
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_dimclamps there; - an omitted stop on the same dynamic axis compiles, because
stop == sys.maxsizeis one of
the conditions the dynamic branch does test.x[:, 5:]reaches the converter as
stop = INT64_MAX, which equalssys.maxsize; the reproducer passessys.maxsizeexplicitly
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
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
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py at slice_op, then run the supplied repro.py to compare finite-stop dynamic cases with the static and omitted-stop controls. Done means a finite stop on a dynamic axis follows Torch clamping semantics or the converter declines the node, with the repro succeeding; consider the separate diagnostic concern involving _TRTInterpreter.run only if scope permits.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100