Fix slice_scatter open end

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1224,7 +1224,8 @@ def _index_copy_kv_eligible(
     if len(node.args) < 4:
         return False
     if input_node is None:
-        input_node = node.args[0]  # type: ignore[assignment]
+        cache_arg = node.args[0]
+        input_node = cache_arg if isinstance(cache_arg, Node) else None
     dim, _index_node, src_node = node.args[1:4]
 
     if not isinstance(input_node, Node) or input_node.op != "placeholder":
@@ -1309,8 +1310,47 @@ def aten_ops_index_copy_fallback(
     )
 
 
+def slice_scatter_validator(
+    node: Node, settings: Optional[CompilationSettings] = None
+) -> bool:
+    """Reject a write whose bounds need the size of the dim it writes -- a negative
+    index, or the open end torch.export writes for ``x[..., start:]`` -- when that dim
+    is dynamic. The converter cannot resolve those (see ``resolve_slice_scatter_write``)
+    and raises, so they run in PyTorch until it gains dynamic bounds.
+
+    Missing metadata is passed, not rejected: the KV-cache classifier in
+    ``lowering/_buffer_lifting.py`` reads the same metadata, and vetoing a write it
+    classified as engine-aliased fails ``assert_predicted_kv_aliased``.
+    """
+    input_meta = getattr(node.args[0], "meta", {})
+    input_val = input_meta.get("val", input_meta.get("tensor_meta"))
+    if input_val is None:
+        _LOGGER.debug(
+            f"slice_scatter node {node.name} has no shape metadata; leaving its bounds "
+            "for the converter to resolve against the TensorRT shape."
+        )
+        return True
+
+    _start, _end, _step, status = impl.slice_scatter.resolve_slice_scatter_write(
+        tuple(input_val.shape),
+        args_bounds_check(node.args, 2, 0),
+        args_bounds_check(node.args, 3),
+        args_bounds_check(node.args, 4),
+        args_bounds_check(node.args, 5),
+    )
+    if status is impl.slice_scatter.KVWriteStatus.DYNAMIC_DIM_SIZE:
+        _LOGGER.debug(
+            f"slice_scatter node {node.name} needs the size of a dynamic dim to "
+            "resolve its bounds; falling back to PyTorch operation."
+        )
+        return False
+    return True
+
+
 @dynamo_tensorrt_converter(
-    torch.ops.aten.slice_scatter.default, supports_dynamic_shapes=True
+    torch.ops.aten.slice_scatter.default,
+    capability_validator=slice_scatter_validator,
+    supports_dynamic_shapes=True,
 )
 @enforce_tensor_types({0: (TRTTensor,), 1: (TRTTensor,)})
 def aten_ops_slice_scatter(

py/torch_tensorrt/dynamo/conversion/impl/slice_scatter.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 import logging
+import sys
 from enum import Enum, auto
 from typing import Any, Optional, Tuple
 
@@ -79,9 +80,25 @@ class KVWriteStatus(Enum):
     OK = auto()
     FULL_OVERWRITE = auto()
     DYNAMIC_BOUNDS = auto()
+    DYNAMIC_DIM_SIZE = auto()
     BAD_DIM = auto()
 
 
+# torch.export uses INT64_MAX as "open end". `sys.maxsize` is the same value and is
+# what the aten.slice converter matches on, so both are accepted.
+_OPEN_END = (sys.maxsize, 2**63 - 1)
+
+
+def _needs_dim_size(bound: Any) -> bool:
+    """Whether ``bound`` has to be resolved against the dim being written: ``None`` and
+    an open end run to it, a negative index counts back from it. A non-int bound is
+    reported as ``DYNAMIC_BOUNDS`` instead, so it passes here.
+    """
+    if bound is None:
+        return True
+    return isinstance(bound, int) and (bound < 0 or bound in _OPEN_END)
+
+
 def resolve_slice_scatter_write(
     cache_shape: Tuple[Any, ...],
     dim: Any,
@@ -92,9 +109,10 @@ def resolve_slice_scatter_write(
     """Resolve ``slice_scatter``'s slice bounds into the form ``_kv_eligible`` takes.
 
     Returns ``(start, end, step, status)``. Under ``OK`` and ``FULL_OVERWRITE`` the
-    three bounds are Python ``int``s, with the op's defaults filled in and negative
-    indices counted from the end; under ``DYNAMIC_BOUNDS`` and ``BAD_DIM`` all three
-    are ``None``, since nothing resolved. ``status`` is one of:
+    three bounds are Python ``int``s, with the op's defaults filled in, negative
+    indices counted from the end and the whole slice clamped to the dim; under the
+    other statuses all three are ``None``, since nothing resolved. ``status`` is one
+    of:
 
     * ``OK`` — the bounds are concrete, and the caller goes on to
       ``_kv_eligible(cache_shape, dim, start, end - start)``.
@@ -105,6 +123,10 @@ def resolve_slice_scatter_write(
       under this status.
     * ``DYNAMIC_BOUNDS`` — a bound is not a Python ``int``, so the converter
       raises ``NotImplementedError``.
+    * ``DYNAMIC_DIM_SIZE`` — a bound needs the size of the dim it writes (see
+      :func:`_needs_dim_size`) and that dim is dynamic. ``aten_ops_slice_scatter``'s
+      validator rejects these so the partitioner runs them in PyTorch; the converter
+      raises if one reaches it anyway.
     * ``BAD_DIM`` — ``dim`` is either not a Python ``int`` or does not index
       ``cache_shape``; the converter raises ``IndexError`` for both. A
       ``numpy.int64`` is rejected on the type check even when its value is in range.
@@ -123,25 +145,40 @@ def resolve_slice_scatter_write(
     if not isinstance(dim, int) or not -len(cache_shape) <= dim < len(cache_shape):
         return None, None, None, KVWriteStatus.BAD_DIM
     dim_size = cache_shape[dim]
+    # A dynamic dim reaches the converter as DYNAMIC_DIM (-1) and the predictor as a
+    # SymInt; neither is a size, and folding the -1 into an index is how an open-ended
+    # write on a dynamic dim used to resolve to an empty index range.
+    dim_is_static = isinstance(dim_size, int) and dim_size >= 0
 
     if start is None:
         start = 0
-    if isinstance(start, int) and start < 0 and isinstance(dim_size, int):
-        start = dim_size + start
-    if end is None:
-        end = dim_size
-    if isinstance(end, int) and end < 0 and isinstance(dim_size, int):
-        end = dim_size + end
     if step is None:
         step = 1
 
+    if not dim_is_static and (_needs_dim_size(start) or _needs_dim_size(end)):
+        return None, None, None, KVWriteStatus.DYNAMIC_DIM_SIZE
+
+    if end is None:
+        end = dim_size
+    if dim_is_static:
+        if isinstance(start, int) and start < 0:
+            start = dim_size + start
+        if isinstance(end, int) and end < 0:
+            end = dim_size + end
+        # Aten clamps a slice to its dim, and so must this: unclamped, the open end
+        # reaches the fallback's np.arange as an INT64_MAX-long index range.
+        if isinstance(start, int):
+            start = min(max(start, 0), dim_size)
+        if isinstance(end, int):
+            end = min(max(end, 0), dim_size)
+
     # A slice covering the whole dim is a plain copy of the source whatever `step` is
     # made of, so it is settled before the bounds are required to be concrete: `step`
     # only has to compare equal to 1, which a symbolic step can do.
     if (
         isinstance(start, int)
         and isinstance(end, int)
-        and isinstance(dim_size, int)
+        and dim_is_static
         and start == 0
         and end == dim_size
         and step == 1
@@ -285,6 +322,14 @@ def slice_scatter(
         raise NotImplementedError(
             "slice_scatter with dynamic start/end/step is not yet supported"
         )
+
+    if status is KVWriteStatus.DYNAMIC_DIM_SIZE:
+        # The validator keeps these out of TensorRT, so this is only reachable for a
+        # node it could not read a shape from.
+        raise NotImplementedError(
+            f"slice_scatter: dim {dim} of the input is dynamic, so this write's bounds "
+            "cannot be resolved without its size"
+        )
     # OK is the only status left, and it resolves all three bounds to Python ints.
     assert start is not None and end is not None and step is not None
 

py/torch_tensorrt/dynamo/lowering/_buffer_lifting.py

@@ -150,14 +150,16 @@ def _kv_write_will_alias(
     direct network input, so this reuses the converters' own eligibility
     predicates -- and, for ``slice_scatter``, the converter's own derivation of
     the arguments those predicates take (:func:`resolve_slice_scatter_write`), since
-    a divergence there mis-predicts just as effectively as a divergence in the
+    a divergence there mispredicts just as effectively as a divergence in the
     predicate. Returning ``False`` routes the write to copy-back, which is right
     whatever the converter goes on to do with it: it lowers most ineligible writes
     to a non-aliasing scatter, returns the source outright for a full overwrite, and
     raises for the rest (a ``slice_scatter`` with dynamic bounds or a bad ``dim``, an
     ``index_copy`` its fallback cannot express). The first two both have a write-back
     to preserve -- for a full overwrite the source *is* the buffer's new contents --
-    and the ones that raise never get as far as needing one.
+    and the ones that raise never get as far as needing one. One more never reaches a
+    converter: a ``slice_scatter`` whose bounds need the size of a dynamic dim, which
+    its validator runs in PyTorch, and which needs the copy-back like any other.
     Imports are local to avoid a lowering<->conversion import cycle.
     """
     if not (isinstance(value_node, torch.fx.Node) and value_node.op == "call_function"):
@@ -201,8 +203,9 @@ def _kv_write_will_alias(
         # Returning False routes the write to copy-back -- the copy_ is erased
         # either way, and the new value is re-attached as a graph output instead. A
         # full overwrite needs that: it returns the source, emits no KV layer, and
-        # its write still has to be recorded. The other two statuses raise out of the
-        # converter, so the compile aborts and the output is never reached.
+        # its write still has to be recorded. Two of the rest raise out of the
+        # converter, so the compile aborts and the output is never reached; the
+        # dynamic-dim one is validated away to PyTorch and keeps its copy-back.
         dim = args[2] if len(args) > 2 else 0
         start, end, _step, status = resolve_slice_scatter_write(
             tuple(cache_shape),
@@ -484,7 +487,7 @@ def lift_mutated_buffers(
     copyback: List[Tuple[torch.fx.Node, str, str]] = []
     # Input-binding names of writes predicted to alias (KV). compile() asserts each
     # actually appears in the engine's aliased_io, and that no copy-back binding
-    # does, turning either mis-prediction into a loud error instead of a silently
+    # does, turning either misprediction into a loud error instead of a silently
     # dropped write-back or a module that raises on every call. The binding name
     # (buf_*) is the stable key: it survives the buffer renaming inline does later,
     # and it is exactly what aliased_io records on the input side.