[Performance] Add fast path for TensorDict copy_at_

Author: vmoens

benchmarks/common/collector_write_benchmarks_test.py

@@ -146,7 +146,7 @@ def _stack_out(rollout, steps, device):
 
 def _copy_at(rollout, steps, device):
     for i, step in enumerate(steps):
-        rollout.copy_at_(step, idx=(slice(None), i))
+        rollout.copy_at_(step, idx=(slice(None), i), fast=True)
     _maybe_synchronize(device)
 
 

tensordict/_td.py

@@ -39,6 +39,8 @@
     _default_is_leaf,
     _device_recorder,
     _expand_to_match_shape,
+    _foreach_copy_,
+    _foreach_copy_compiled,
     _is_leaf_nontensor,
     _is_tensor_collection,
     _load_metadata,
@@ -2714,6 +2716,62 @@ def _set_at_tuple(self, key, value, idx, *, validated, non_blocking: bool):
         )
         return self
 
+    def _update_at_fast(
+        self,
+        input_dict_or_td: dict[str, CompatibleType] | T,
+        idx: IndexType,
+        clone: bool,
+        *,
+        non_blocking: bool,
+        keys_to_update: Sequence[NestedKey] | None,
+    ) -> Any:
+        if (
+            clone
+            or keys_to_update is not None
+            or type(self) is not TensorDict
+            or type(input_dict_or_td) is not TensorDict
+        ):
+            return NotImplemented
+        source_keys, source_values = input_dict_or_td._items_list(True, True)
+        if not source_keys:
+            if input_dict_or_td.is_empty():
+                return self
+            return NotImplemented
+        dest_values = []
+        for key in source_keys:
+            dest = self
+            for subkey in _unravel_key_to_tuple(key):
+                if type(dest) is not TensorDict or subkey not in dest._tensordict:
+                    return NotImplemented
+                dest = dest._tensordict[subkey]
+            dest_values.append(dest)
+        if not isinstance(idx, tuple):
+            idx = (idx,)
+        idx = convert_ellipsis_to_idx(idx, self.batch_size)
+        indexed_dest_values = []
+        for dest, source in zip(dest_values, source_values):
+            if not isinstance(dest, Tensor) or not isinstance(source, Tensor):
+                return NotImplemented
+            try:
+                dest_indexed = dest[idx]
+            except (IndexError, RuntimeError, TypeError):
+                return NotImplemented
+            if dest_indexed.shape != source.shape:
+                return NotImplemented
+            if (
+                dest_indexed is not dest
+                and getattr(dest_indexed, "_base", None) is None
+            ):
+                return NotImplemented
+            indexed_dest_values.append(dest_indexed)
+        if _foreach_copy_ is not None:
+            copy_fn = _foreach_copy_compiled if is_compiling() else _foreach_copy_
+            copy_fn(indexed_dest_values, source_values, non_blocking=non_blocking)
+        else:
+            for dest, source in zip(indexed_dest_values, source_values):
+                dest.copy_(source, non_blocking=non_blocking)
+        return self
+
     @lock_blocked
     def del_(self, key: NestedKey) -> Self:
         key = _unravel_key_to_tuple(key)

tensordict/_tensorcollection.pyi

@@ -985,7 +985,12 @@ class TensorCollection:
     def create_nested(self, key): ...
     def copy_(self, tensordict: T, non_blocking: bool = False) -> Self: ...
     def copy_at_(
-        self, tensordict: T, idx: IndexType, non_blocking: bool = False
+        self,
+        tensordict: T,
+        idx: IndexType,
+        non_blocking: bool = False,
+        *,
+        fast: bool | None = None,
     ) -> Self: ...
     def is_empty(self) -> bool: ...
     def setdefault(

tensordict/base.py

@@ -8553,6 +8553,9 @@ def update_at_(
         """Updates the TensorDict in-place at the specified index with values from either a dictionary or another TensorDict.
 
         Unlike  TensorDict.update, this function will throw an error if the key is unknown to the TensorDict.
+        This method keeps the general ``set_at_`` semantics and supports tensor
+        and non-tensor leaves. For optimized tensor-only copies in hot paths,
+        use :meth:`~tensordict.TensorDictBase.copy_at_` with ``fast=True``.
 
         Args:
             input_dict_or_td (TensorDictBase or dict): input data to be written
@@ -8620,6 +8623,17 @@ def update_at_(
             self.set_at_((firstkey, *nextkeys), value, idx, non_blocking=non_blocking)
         return self
 
+    def _update_at_fast(
+        self,
+        input_dict_or_td: dict[str, CompatibleType] | T,
+        idx: IndexType,
+        clone: bool,
+        *,
+        non_blocking: bool,
+        keys_to_update: Sequence[NestedKey] | None,
+    ) -> Any:
+        return NotImplemented
+
     def replace(self, *args, **kwargs):
         """Creates a shallow copy of the tensordict where entries have been replaced.
 
@@ -8729,9 +8743,70 @@ def copy_(self, tensordict: T, non_blocking: bool = False) -> Self:
         return self.update_(tensordict, non_blocking=non_blocking)
 
     def copy_at_(
-        self, tensordict: T, idx: IndexType, non_blocking: bool = False
+        self,
+        tensordict: T,
+        idx: IndexType,
+        non_blocking: bool = False,
+        *,
+        fast: bool | None = None,
     ) -> Self:
-        """See :obj:`TensorDictBase.update_at_`."""
+        """Copies values from ``tensordict`` into ``self`` at the specified index.
+
+        ``copy_at_`` is an explicit copy-oriented variant of
+        :meth:`~tensordict.TensorDictBase.update_at_`. Unlike ``update_at_``,
+        it may use optimized tensor-only copy paths and is intended for hot
+        paths where the source and destination structures are known to match.
+
+        Args:
+            tensordict (TensorDictBase): input data to be copied in ``self``.
+            idx (int, torch.Tensor, iterable, slice): index of the tensordict
+                where the copy should occur.
+            non_blocking (bool, optional): if ``True`` and this copy is between
+                different devices, the copy may occur asynchronously with respect
+                to the host.
+
+        Keyword Args:
+            fast (bool or None, optional): controls whether ``copy_at_`` may
+                fall back to :meth:`~tensordict.TensorDictBase.update_at_`.
+                If ``True``, only the optimized tensor-only path is used and a
+                ``RuntimeError`` is raised when the fast path is not available.
+                If ``False``, this method delegates directly to ``update_at_``.
+                If ``None``, the current default, ``copy_at_`` warns and falls
+                back to ``update_at_`` when the fast path is unavailable. The
+                default will become ``True`` in v0.14.
+
+        Returns:
+            self
+        """
+        if fast is None:
+            warnings.warn(
+                "copy_at_(..., fast=None) currently falls back to update_at_ "
+                "when the optimized tensor-only copy path is unavailable. "
+                "This default will change to fast=True in v0.14, making "
+                "copy_at_ fast-only by default. Pass fast=False to keep the "
+                "current fallback behavior, or fast=True to require the fast "
+                "path.",
+                FutureWarning,
+                stacklevel=2,
+            )
+        elif fast is False:
+            return self.update_at_(tensordict, idx, non_blocking=non_blocking)
+        result = self._update_at_fast(
+            input_dict_or_td=tensordict,
+            idx=idx,
+            clone=False,
+            non_blocking=non_blocking,
+            keys_to_update=None,
+        )
+        if result is not NotImplemented:
+            return result
+        if fast:
+            raise RuntimeError(
+                "copy_at_(..., fast=True) requires the optimized tensor-only "
+                "copy path, but the source, destination, or index is not "
+                "compatible. Use fast=False or update_at_ for the general "
+                "update semantics."
+            )
         return self.update_at_(tensordict, idx, non_blocking=non_blocking)
 
     def is_empty(self) -> bool:

tensordict/tensorclass.pyi

@@ -1030,7 +1030,12 @@ class TensorClass:
     def create_nested(self, key: NestedKey) -> Self: ...
     def copy_(self, tensordict: T, non_blocking: bool = False) -> Self: ...
     def copy_at_(
-        self, tensordict: T, idx: IndexType, non_blocking: bool = False
+        self,
+        tensordict: T,
+        idx: IndexType,
+        non_blocking: bool = False,
+        *,
+        fast: bool | None = None,
     ) -> Self: ...
     def is_empty(self) -> bool: ...
     def setdefault(

test/test_tensordict.py

@@ -9232,6 +9232,89 @@ def test_update_at_(self, td_name, device):
         td.update_at_(td0, 0)
         assert (td[0] == 0).all()
 
+    def test_update_at_nested_time_slice(self, td_name, device):
+        td = TensorDict(
+            {
+                "a": torch.zeros(4, 3, 5, device=device),
+                "b": TensorDict(
+                    {"c": torch.zeros(4, 3, 2, device=device)},
+                    batch_size=[4, 3],
+                    device=device,
+                ),
+            },
+            batch_size=[4, 3],
+            device=device,
+        )
+        td0 = TensorDict(
+            {
+                "a": torch.ones(4, 5, device=device),
+                "b": TensorDict(
+                    {"c": torch.ones(4, 2, device=device)},
+                    batch_size=[4],
+                    device=device,
+                ),
+            },
+            batch_size=[4],
+            device=device,
+        )
+        td.update_at_(td0, (slice(None), 1))
+        assert (td[:, 1] == td0).all()
+        assert (td[:, 0] == 0).all()
+        assert (td[:, 2] == 0).all()
+
+    def test_update_at_nested_time_slice_locked(self, td_name, device):
+        td = TensorDict(
+            {
+                "a": torch.zeros(4, 3, 5, device=device),
+                "b": TensorDict(
+                    {"c": torch.zeros(4, 3, 2, device=device)},
+                    batch_size=[4, 3],
+                    device=device,
+                ),
+            },
+            batch_size=[4, 3],
+            device=device,
+        ).lock_()
+        td0 = TensorDict(
+            {
+                "a": torch.ones(4, 5, device=device),
+                "b": TensorDict(
+                    {"c": torch.ones(4, 2, device=device)},
+                    batch_size=[4],
+                    device=device,
+                ),
+            },
+            batch_size=[4],
+            device=device,
+        )
+        td.update_at_(td0, (slice(None), 2))
+        assert (td[:, 2] == td0).all()
+
+    @pytest.mark.parametrize("method", ["copy_at_", "update_at_"])
+    def test_update_at_nontensor_data(self, td_name, device, method):
+        td = TensorDict({"val": NonTensorData(data=0, batch_size=[10])}, [10])
+        newdata = TensorDict({"val": NonTensorData(data=1, batch_size=[5])}, [5])
+
+        if method == "copy_at_":
+            getattr(td, method)(newdata, slice(1, None, 2), fast=False)
+        else:
+            getattr(td, method)(newdata, slice(1, None, 2))
+
+        assert td.get("val").tolist() == [0, 1] * 5
+
+    def test_copy_at_fast_transition(self, td_name, device):
+        td = TensorDict({"val": NonTensorData(data=0, batch_size=[10])}, [10])
+        newdata = TensorDict({"val": NonTensorData(data=1, batch_size=[5])}, [5])
+
+        with pytest.warns(FutureWarning, match="fast=True in v0.14"):
+            td.copy_at_(newdata, slice(1, None, 2))
+        assert td.get("val").tolist() == [0, 1] * 5
+
+        td = TensorDict({"val": NonTensorData(data=0, batch_size=[10])}, [10])
+        with pytest.raises(RuntimeError, match="fast=True"):
+            td.copy_at_(newdata, slice(1, None, 2), fast=True)
+        assert td.get("val").tolist() == [0] * 10
+
     # This is needed because update in lazy permute/view etc does not behave correctly when
     # legacy is False. When these classes will be deprecated, we can just remove the decorator
     @set_lazy_legacy(True)