From 1ee5b88dcad04098b7e73072f743ca95cad297d0 Mon Sep 17 00:00:00 2001 From: kiwigitops Date: Fri, 31 Jul 2026 21:32:00 -0400 Subject: [PATCH] fix: support dynamic slice_scatter fallback shapes --- .../dynamo/conversion/impl/slice_scatter.py | 40 ++++++++++++++++--- .../conversion/test_slice_scatter_aten.py | 19 +++++++++ 2 files changed, 54 insertions(+), 5 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py b/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py index c277a4ca6b..529323d611 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py @@ -22,8 +22,11 @@ from typing import Optional, Tuple import numpy as np +import tensorrt as trt +from tensorrt import ITensor as TRTTensor from torch.fx.node import Target from torch_tensorrt.dynamo._SourceIR import SourceIR +from torch_tensorrt.dynamo.conversion import impl from torch_tensorrt.dynamo.conversion._ConversionContext import ( AliasedOutput, AliasKind, @@ -34,9 +37,8 @@ set_layer_name, ) from torch_tensorrt.dynamo.conversion.impl import select - -import tensorrt as trt -from tensorrt import ITensor as TRTTensor +from torch_tensorrt.dynamo.conversion.impl.shape import shape as get_shape +from torch_tensorrt.dynamo.utils import DYNAMIC_DIM logger = logging.getLogger(__name__) @@ -215,8 +217,36 @@ def slice_scatter( target_shape[dim] = len(indices_np) indices_np = indices_np.reshape(target_shape) src_shape = tuple(src.shape) - indices_np = np.broadcast_to(indices_np, src_shape).astype(np.int64) - indices_tensor = get_trt_tensor(ctx, indices_np, name + "_fallback_indices") + if DYNAMIC_DIM in src_shape: + indices_tensor = get_trt_tensor(ctx, indices_np, name + "_fallback_indices") + runtime_src_shape = [] + for axis, size in enumerate(src_shape): + if axis == dim % rank: + runtime_src_shape.append(len(indices_np)) + elif size == DYNAMIC_DIM: + runtime_src_shape.append( + get_shape( + ctx, + target, + source_ir, + name + f"_fallback_src_shape_{axis}", + src, + axis, + ) + ) + else: + runtime_src_shape.append(size) + indices_tensor = impl.slice.expand( + ctx, + target, + source_ir, + name + "_fallback_indices_expand", + indices_tensor, + tuple(runtime_src_shape), + ) + else: + indices_np = np.broadcast_to(indices_np, src_shape).astype(np.int64) + indices_tensor = get_trt_tensor(ctx, indices_np, name + "_fallback_indices") return select.scatter( ctx, diff --git a/tests/py/dynamo/conversion/test_slice_scatter_aten.py b/tests/py/dynamo/conversion/test_slice_scatter_aten.py index 1034327360..6423353590 100644 --- a/tests/py/dynamo/conversion/test_slice_scatter_aten.py +++ b/tests/py/dynamo/conversion/test_slice_scatter_aten.py @@ -21,6 +21,7 @@ import torch from parameterized import parameterized from torch.testing._internal.common_utils import run_tests +from torch_tensorrt import Input from .harness import DispatchTestCase @@ -76,6 +77,24 @@ def test_fallback_step_two(self): update = torch.randn(2, 4, 8, 8) self.run_test(module, [cache, update]) + def test_fallback_dynamic_shape(self): + module = _SliceScatterNotInputModule(2, 1, 60, step=3) + input_specs = [ + Input( + min_shape=(2, 4, 64, 8), + opt_shape=(3, 4, 64, 12), + max_shape=(4, 4, 64, 16), + dtype=torch.float32, + ), + Input( + min_shape=(2, 4, 20, 8), + opt_shape=(3, 4, 20, 12), + max_shape=(4, 4, 20, 16), + dtype=torch.float32, + ), + ] + self.run_test_with_dynamic_shape(module, input_specs) + def test_full_overwrite_is_identity(self): """When start=0, end=dim_size, step=1, the converter short-circuits and returns ``src`` directly. Wrap the returned tensor in a small op