From 0f4e2717961c7aa5c8acc7de0607465ebeeb432e Mon Sep 17 00:00:00 2001 From: Erik Lundell Date: Tue, 4 Aug 2026 16:25:02 +0200 Subject: [PATCH] Memory planning: Support shared_allocation with offset Previously, there was a mechanism that let passes annotate nodes with a meta field "_share_alloc_with_arg_idx", to indicate to memory planning algorithms that a node output tensor is shared. Generalize this by adding the meta field "_shared_alloc_offset", to allow setting an offset from the base shared allocation. This meta is picked up by the MemoryPlanningPass, and transfered to the TensorSpec. To do this, a new field "storage_base_offset" is added to the TensorSpec. "inplace_base" is renamed to "storage_base", while keeping a inplace_base as an alias. When inplace_base is used, the storage_base_offset is not allowed to be non-zero. This work exposed an issue where planning of inplace_base/storage_base backing tensors didn't respect the lifetimes of backed tensors. Fix this by extending lifetimes accordingly in update_all_tensors_lifetime. Handling of unexpected values of TensorSpec in _move_memory_meta_to_spec were also changed to errors instead of silent returns. Finally, modify the greedy algorithm to respect the new offset. Signed-off-by: Erik Lundell Change-Id: I8d8348487aaf0d29e87fbcc563183b1d9c1d73b1 --- exir/memory_planning.py | 131 ++++++--- exir/passes/memory_planning_pass.py | 46 +++- .../replace_view_copy_with_view_pass.py | 3 + exir/tensor.py | 32 ++- exir/tests/test_memory_planning.py | 248 +++++++++++++++++- exir/tests/test_tensor.py | 17 ++ 6 files changed, 428 insertions(+), 49 deletions(-) diff --git a/exir/memory_planning.py b/exir/memory_planning.py index 012cf8dd144..6750d18441e 100644 --- a/exir/memory_planning.py +++ b/exir/memory_planning.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -89,6 +90,19 @@ def mem_obj_id_match( return lhs_spec.mem_obj_id == rhs_spec.mem_obj_id + @classmethod + def storage_root(cls, spec: TensorSpec) -> TensorSpec: + seen: Set[TensorSpec] = set() + root = spec + while root.storage_base is not None: + internal_assert( + root not in seen, + "Circular storage_base relationship is not supported.", + ) + seen.add(root) + root = root.storage_base + return root + @classmethod def has_overlap(cls, lhs_ivl: List[int], rhs_ivl: List[int]) -> bool: r""" @@ -191,13 +205,13 @@ def verify_storage_reuse( if not allow_lifetime_and_storage_overlap and self.lifetime_overlap( lhs_spec, rhs_spec ): - # In-place element-wise ops intentionally share storage - # between input and output despite overlapping lifetimes. - is_inplace_pair = ( - lhs_spec.inplace_base is rhs_spec - or rhs_spec.inplace_base is lhs_spec - ) - if not is_inplace_pair: + # Some ops, such as in-place ops, intentionally place one + # TensorSpec inside another TensorSpec's storage despite + # overlapping lifetimes. + is_common_base_pair = Verifier.storage_root( + lhs_spec + ) is Verifier.storage_root(rhs_spec) + if not is_common_base_pair: raise InternalError( f"Unexpected storage overlap: {Verifier._debug_message_from_specs(lhs_spec, rhs_spec)}" ) @@ -691,6 +705,7 @@ def update_all_tensors_lifetime( ): update_tensor_lifetime(node, spec, node_idx, max_node_idx, graph_signature) specs.add(spec) + _extend_storage_base_lifetimes(specs) return specs @@ -758,21 +773,6 @@ class MemoryAlgoResult: bufsizes: List[int] -def materialize_buffer( - shared_objects: List[SharedObject], input_total_size: int = 0 -) -> int: - r""" - Assign concrete location in the buffer for each SharedObject.offset. - - Assuming all the passed in shared objects belong to the same memory buffer. - """ - total_size = input_total_size - for sobj in shared_objects: - sobj.offset = total_size - total_size += sobj.size - return total_size - - def _does_not_overlap(sobj: SharedObject, spec: TensorSpec) -> bool: r""" Check if a shared object and a tensor do not overlap. @@ -939,17 +939,48 @@ def _contains_xnnpack_delegate(graph_module: torch.fx.GraphModule) -> bool: return False -def _resolve_inplace_specs( - deferred_inplace: List[TensorSpec], +def _extend_storage_base_lifetimes(specs: Set[TensorSpec]) -> None: + for spec in specs: + if spec.storage_base is None: + continue + internal_assert( + spec.lifetime[0] is not None and spec.lifetime[1] is not None, + "Storage-backed TensorSpec must have a lifetime.", + ) + start = cast(int, spec.lifetime[0]) + end = cast(int, spec.lifetime[1]) + seen: Set[TensorSpec] = {spec} + base = spec.storage_base + while base is not None: + internal_assert( + base not in seen, + "Circular storage_base relationship is not supported.", + ) + seen.add(base) + internal_assert( + base.lifetime[0] is not None and base.lifetime[1] is not None, + "storage_base TensorSpec must have a lifetime.", + ) + base.lifetime[0] = min(cast(int, base.lifetime[0]), start) + base.lifetime[1] = max(cast(int, base.lifetime[1]), end) + base = base.storage_base + + +def _resolve_storage_base_specs( + deferred_storage_base: List[TensorSpec], spec2obj: Dict[TensorSpec, SharedObject], greedy_result: MemoryAlgoResult, ) -> None: - remaining = list(deferred_inplace) + remaining = list(deferred_storage_base) while remaining: progress = False next_remaining = [] for spec in remaining: - base = spec.inplace_base + base = spec.storage_base + internal_assert( + base is not None, + "Deferred storage-backed TensorSpec should have a storage_base.", + ) if base not in spec2obj: next_remaining.append(spec) continue @@ -960,25 +991,47 @@ def _resolve_inplace_specs( spec_alloc_result = greedy_result.spec_dict[spec] spec_alloc_result.mem_id = base_alloc_result.mem_id + allocated_memory = spec.allocated_memory + storage_base_offset = spec.storage_base_offset + internal_assert( + storage_base_offset >= 0, + "storage_base_offset must be non-negative.", + ) base_alloc_offset = None + base_allocated_memory = None for alloc_entry in sobj.allocations: if alloc_entry.spec is base: base_alloc_offset = alloc_entry.offset + base_allocated_memory = alloc_entry.spec.allocated_memory break assert base_alloc_offset is not None, ( f"Base allocation entry not found in shared object for spec " f"with allocated_memory={spec.allocated_memory}" ) + assert base_allocated_memory is not None, ( + f"Base allocation entry not found in shared object for spec " + f"with allocated_memory={allocated_memory}" + ) + internal_assert( + (base_alloc_offset + storage_base_offset) % spec.alignment == 0, + f"Storage-backed TensorSpec allocation must respect alignment, got offset {storage_base_offset} inside parent with offset {storage_base_offset} for alignment {spec.alignment}.", + ) + internal_assert( + storage_base_offset + allocated_memory <= base_allocated_memory, + "Storage-backed TensorSpec allocation must fit within storage_base.", + ) sobj.first_used_index = min(sobj.first_used_index, spec.lifetime[0]) sobj.last_used_index = max(sobj.last_used_index, spec.lifetime[1]) - sobj.allocations.append(AllocationSpec(base_alloc_offset, spec)) + sobj.allocations.append( + AllocationSpec(base_alloc_offset + storage_base_offset, spec) + ) spec2obj[spec] = sobj if not progress: unresolved = ", ".join( f"allocated_memory={s.allocated_memory}" for s in next_remaining ) raise InternalError( - f"Circular or unresolvable in-place dependency chain: {unresolved}" + "Circular or unresolvable storage_base dependency chain: " + unresolved ) remaining = next_remaining @@ -1001,9 +1054,11 @@ def _compute_total_sizes( assert isinstance(bufsizes, list) if len(bufsizes) > mem_id: input_total_size = bufsizes[mem_id] - total_sizes[mem_id] = materialize_buffer( - shared_objects[mem_id], input_total_size - ) + total_size = input_total_size + for sobj in shared_objects[mem_id]: + sobj.offset = total_size + total_size += sobj.size + total_sizes[mem_id] = total_size total_sizes[mem_id] += extra_padding for sobj in shared_objects[mem_id]: @@ -1056,7 +1111,7 @@ def greedy( sorted_specs.reverse() - deferred_inplace: List[TensorSpec] = [] + deferred_storage_base: List[TensorSpec] = [] for spec in sorted_specs: spec_alloc_result = greedy_result.spec_dict.get(spec, SpecAllocResult(0, 0, 0)) @@ -1067,8 +1122,8 @@ def greedy( greedy_result.spec_dict[spec] = spec_alloc_result spec.realign(alignment) - if spec.inplace_base is not None: - deferred_inplace.append(spec) + if spec.storage_base is not None: + deferred_storage_base.append(spec) continue spec2obj[spec] = pick_shared_obj( @@ -1077,7 +1132,7 @@ def greedy( allow_overlapping_allocations, ) - _resolve_inplace_specs(deferred_inplace, spec2obj, greedy_result) + _resolve_storage_base_specs(deferred_storage_base, spec2obj, greedy_result) total_sizes = _compute_total_sizes( shared_objects, graph_module, extra_padding, greedy_result, len(spec2obj) @@ -1206,10 +1261,10 @@ def _allocate_buf(bufsizes: List[int], mem_id: int, allocated: int) -> int: bufsizes = cast(List[int], bufsizes) for spec in specs: - if spec.inplace_base is not None: + if spec.storage_base is not None: raise InternalError( - "The naive memory planning algorithm does not support in-place " - "element-wise ops (inplace_base). Use the greedy algorithm instead." + "The naive memory planning algorithm does not support storage-backed " + "TensorSpecs. Use the greedy algorithm instead." ) spec_alloc_result = naive_result.spec_dict.get(spec, SpecAllocResult(0, 0, 0)) diff --git a/exir/passes/memory_planning_pass.py b/exir/passes/memory_planning_pass.py index 99a5f3dd8ec..6fe5fe1539d 100644 --- a/exir/passes/memory_planning_pass.py +++ b/exir/passes/memory_planning_pass.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -136,6 +137,42 @@ def _check_default_mem_ids(gm: torch.fx.GraphModule): ) +def _move_memory_meta_to_spec(node: Node) -> None: + """Move storage sharing metadata from node.meta to node.meta["spec"]. + + Only applies if _share_alloc_with_arg_idx is set. + """ + share_idx = node.meta.get("_share_alloc_with_arg_idx") + if share_idx is None: + return + output_spec = node.meta.get("spec") + if not isinstance(output_spec, TensorSpec): + raise TypeError( + "_share_alloc_with_arg_idx requires node.meta['spec'] to be a TensorSpec" + ) + if not isinstance(share_idx, int): + raise TypeError("_share_alloc_with_arg_idx must be an int") + if share_idx < 0 or share_idx >= len(node.args): + raise IndexError("_share_alloc_with_arg_idx must index node.args") + + input_node = node.args[share_idx] + if not isinstance(input_node, Node): + raise TypeError("_share_alloc_with_arg_idx must reference a Node argument") + + base_spec = input_node.meta.get("spec") + if not isinstance(base_spec, TensorSpec): + raise TypeError( + "_share_alloc_with_arg_idx must reference an argument with a TensorSpec" + ) + + shared_alloc_offset = node.meta.get("_shared_alloc_offset", 0) + if not isinstance(shared_alloc_offset, int): + raise TypeError("_shared_alloc_offset must be an int") + + output_spec.storage_base = base_spec + output_spec.storage_base_offset = shared_alloc_offset + + @dataclass class _MemoryPlanningState: mutable_buffers: Dict[str, Set[TensorSpec]] = field(default_factory=dict) @@ -196,13 +233,8 @@ def _set_alloc_node_spec(self, graph_module: torch.fx.GraphModule) -> None: if len(out_arg_names) == 1: out_alloc_node = node.kwargs[out_arg_names[0]] out_alloc_node.meta["spec"] = node.meta["spec"] - share_idx = node.meta.get("_share_alloc_with_arg_idx") - if share_idx is not None and share_idx < len(node.args): - input_node = node.args[share_idx] - if isinstance(input_node, Node): - base_spec = input_node.meta.get("spec") - if isinstance(base_spec, TensorSpec): - node.meta["spec"].inplace_base = base_spec + + _move_memory_meta_to_spec(node) continue specs = get_node_tensor_specs(node) i = 0 diff --git a/exir/passes/replace_view_copy_with_view_pass.py b/exir/passes/replace_view_copy_with_view_pass.py index 28fcc97aaf5..947952d7692 100644 --- a/exir/passes/replace_view_copy_with_view_pass.py +++ b/exir/passes/replace_view_copy_with_view_pass.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -108,6 +109,8 @@ def __init__(self, base: TensorSpec, shape: List[int]) -> None: "mem_id", "mem_obj_id", "mem_offset", + "storage_base", + "storage_base_offset", "dtype", # property "extra_tensor_info", # property "device", diff --git a/exir/tensor.py b/exir/tensor.py index a4e480ffce0..199c5adfafe 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -212,9 +212,35 @@ def init_mem_planning_fields(self) -> None: self.mem_id = None self.mem_obj_id = None self.mem_offset = None - # Set by InPlaceElemWiseLikeOpsPass: the base TensorSpec whose memory - # this spec should share (output allocated in-place over the input). - self.inplace_base: Optional["TensorSpec"] = None + # Optional TensorSpec whose storage this spec is backed by. This is + # metadata for memory planners; mem_offset remains an absolute offset + # after the winning memory plan is written back. + self.storage_base: Optional["TensorSpec"] = None + # Byte offset into storage_base when this spec is storage-backed. + self.storage_base_offset: int = 0 + + @property + def inplace_base(self) -> Optional["TensorSpec"]: + """Zero-offset compatibility alias for storage_base. + + Use storage_base and storage_base_offset directly for aliases with a + non-zero offset. + """ + internal_assert( + self.storage_base is None or self.storage_base_offset == 0, + "inplace_base is only valid for TensorSpecs whose storage_base " + "has offset 0.", + ) + return self.storage_base + + @inplace_base.setter + def inplace_base(self, base: Optional["TensorSpec"]) -> None: + internal_assert( + self.storage_base_offset == 0, + "inplace_base can only be set when storage_base_offset is 0. " + "Use storage_base directly for non-zero-offset aliases.", + ) + self.storage_base = base @property def dtype(self) -> torch.dtype: diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index 31f3b1844c2..37bc088cfcc 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -8,7 +9,7 @@ import itertools import unittest -from typing import Any, Callable, List, Optional, Tuple, Type +from typing import Any, Callable, cast, List, Optional, Tuple, Type import executorch.exir as exir @@ -29,6 +30,7 @@ from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.memory_planning import ( _do_user_inputs_exist, + _extend_storage_base_lifetimes, _is_inplace_node, apply_algo, collect_specs_from_nodes, @@ -1653,6 +1655,221 @@ def test_disabled_falls_back_to_cpu(self) -> None: self.assertNotIn("non_const_buffer_device", gm.meta) +class TestStorageBaseMemoryPlanning(unittest.TestCase): + def _empty_graph_module(self) -> GraphModule: + graph = Graph() + graph.output(()) + return GraphModule({}, graph) + + def _make_storage_backed_specs(self) -> Tuple[TensorSpec, TensorSpec]: + base = TensorSpec.from_tensor(torch.empty(10)) + child = TensorSpec.from_tensor(torch.empty(2)) + + base.lifetime = [0, 1] + child.lifetime = [0, 1] + base.mem_id = 1 + child.mem_id = 1 + child.storage_base = base + child.storage_base_offset = 16 + return base, child + + def test_greedy_places_storage_backed_spec_inside_base_object(self) -> None: + base, child = self._make_storage_backed_specs() + + algo = MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + algo( + 16, + {base, child}, + self._empty_graph_module(), + cast(ExportGraphSignature, None), + 0, + ) + + self.assertEqual(child.mem_id, base.mem_id) + self.assertEqual(child.mem_obj_id, base.mem_obj_id) + base_mem_offset = base.mem_offset + self.assertIsNotNone(base_mem_offset) + assert base_mem_offset is not None + self.assertEqual(child.mem_offset, base_mem_offset + 16) + + def test_greedy_result_contains_storage_backed_full_plan(self) -> None: + base, child = self._make_storage_backed_specs() + base.realign(1) + child.realign(1) + + result = greedy( + 1, + {base, child}, + self._empty_graph_module(), + cast(ExportGraphSignature, None), + 0, + ) + + base_result = result.spec_dict[base] + child_result = result.spec_dict[child] + self.assertEqual(child_result.mem_id, base_result.mem_id) + self.assertEqual(child_result.mem_obj_id, base_result.mem_obj_id) + self.assertEqual(child_result.mem_offset, base_result.mem_offset + 16) + + def test_greedy_resolves_chained_storage_base(self) -> None: + # Build a storage chain where `base` owns the allocation, `child` + # aliases `base`, and `grandchild` aliases `child`. + base = TensorSpec.from_tensor(torch.empty(16, dtype=torch.uint8)) + child = TensorSpec.from_tensor(torch.empty(6, dtype=torch.uint8)) + grandchild = TensorSpec.from_tensor(torch.empty(2, dtype=torch.uint8)) + for spec in (base, child, grandchild): + spec.lifetime = [0, 1] + spec.mem_id = 1 + child.storage_base = base + child.storage_base_offset = 8 + grandchild.storage_base = child + grandchild.storage_base_offset = 4 + + # Greedy should resolve the chain in dependency order and assign all + # three specs to the same memory object. + algo = MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + algo( + 1, + {base, child, grandchild}, + self._empty_graph_module(), + cast(ExportGraphSignature, None), + 0, + ) + + self.assertEqual(child.mem_id, base.mem_id) + self.assertEqual(grandchild.mem_id, base.mem_id) + self.assertEqual(child.mem_obj_id, base.mem_obj_id) + self.assertEqual(grandchild.mem_obj_id, base.mem_obj_id) + base_mem_offset = base.mem_offset + self.assertIsNotNone(base_mem_offset) + assert base_mem_offset is not None + # Offsets are accumulated through the chain: child is +8 from base, + # grandchild is +4 from child, so grandchild is +12 from base. + self.assertEqual(child.mem_offset, base_mem_offset + 8) + self.assertEqual(grandchild.mem_offset, base_mem_offset + 12) + + def test_greedy_reserves_storage_base_lifetime_before_reuse(self) -> None: + base = TensorSpec.from_tensor(torch.empty(16, dtype=torch.uint8)) + child = TensorSpec.from_tensor(torch.empty(8, dtype=torch.uint8)) + other = TensorSpec.from_tensor(torch.empty(12, dtype=torch.uint8)) + for spec in (base, child, other): + spec.mem_id = 1 + base.lifetime = [0, 1] + child.lifetime = [4, 5] + other.lifetime = [4, 5] + child.storage_base = base + child.storage_base_offset = 8 + + _extend_storage_base_lifetimes({base, child, other}) + + algo = MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + algo( + 1, + {base, child, other}, + self._empty_graph_module(), + cast(ExportGraphSignature, None), + 0, + ) + + self.assertEqual(base.lifetime, [0, 5]) + self.assertEqual(child.mem_id, base.mem_id) + self.assertEqual(child.mem_obj_id, base.mem_obj_id) + self.assertNotEqual(other.mem_obj_id, base.mem_obj_id) + + def test_set_alloc_node_spec_uses_shared_alloc_offset(self) -> None: + base = TensorSpec.from_tensor(torch.empty(10)) + child = TensorSpec.from_tensor(torch.empty(2)) + + graph = Graph() + input_node = graph.placeholder("input") + input_node.meta["spec"] = base + other_node = graph.placeholder("other") + other_node.meta["spec"] = base + out_node = graph.placeholder("out") + add_node = graph.call_function( + torch.ops.aten.add.out, + args=(input_node, other_node), + kwargs={"out": out_node}, + ) + add_node.meta["spec"] = child + add_node.meta["_share_alloc_with_arg_idx"] = 0 + add_node.meta["_shared_alloc_offset"] = 16 + graph.output(add_node) + graph_module = GraphModule({}, graph) + + MemoryPlanningPass()._set_alloc_node_spec(graph_module) + + self.assertIs(child.storage_base, base) + self.assertEqual(child.storage_base_offset, 16) + + def test_verifier_allows_storage_base_overlap(self) -> None: + base, child = self._make_storage_backed_specs() + + algo = MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + algo( + 1, + {base, child}, + self._empty_graph_module(), + cast(ExportGraphSignature, None), + 0, + ) + + graph = Graph() + base_node = graph.placeholder("base") + base_node.meta["spec"] = base + child_node = graph.placeholder("child") + child_node.meta["spec"] = child + graph.output((base_node, child_node)) + graph_module = GraphModule({}, graph) + + verifier = Verifier( + graph_module, + alloc_graph_input=True, + alloc_graph_output=True, + alloc_mutable_buffers=True, + ) + verifier.verify_storage_reuse() + + def test_verifier_allows_chained_storage_base_overlap(self) -> None: + outer = TensorSpec.from_tensor(torch.empty(10)) + base = TensorSpec.from_tensor(torch.empty(6)) + child = TensorSpec.from_tensor(torch.empty(2)) + for spec in (outer, base, child): + spec.lifetime = [0, 1] + spec.mem_id = 1 + base.storage_base = outer + base.storage_base_offset = 8 + child.storage_base = base + child.storage_base_offset = 4 + + algo = MemoryPlanningAlgorithmSuite(algo_list=[greedy]) + algo( + 1, + {outer, base, child}, + self._empty_graph_module(), + cast(ExportGraphSignature, None), + 0, + ) + + graph = Graph() + outer_node = graph.placeholder("outer") + outer_node.meta["spec"] = outer + base_node = graph.placeholder("base") + base_node.meta["spec"] = base + child_node = graph.placeholder("child") + child_node.meta["spec"] = child + graph.output((outer_node, base_node, child_node)) + graph_module = GraphModule({}, graph) + + verifier = Verifier( + graph_module, + alloc_graph_input=True, + alloc_graph_output=True, + alloc_mutable_buffers=True, + ) + verifier.verify_storage_reuse() + + class TestInPlaceElemWise(unittest.TestCase): def _run_inplace_pipeline( self, @@ -1725,6 +1942,35 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ) verifier.verify_storage_reuse() + def test_verifier_allows_chained_inplace_overlap(self) -> None: + class Model(torch.nn.Module): + def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + c = a + b + d = c * b + e = d * b + return e + + gm = self._run_inplace_pipeline( + Model(), + (torch.randn(10), torch.randn(10)), + {exir_ops.edge.aten.mul.Tensor}, + ) + + inplace_nodes = [ + node + for node in gm.graph.nodes + if node.op == "call_function" and _is_inplace_node(node) + ] + self.assertEqual(len(inplace_nodes), 2) + + verifier = Verifier( + gm, + alloc_graph_input=True, + alloc_graph_output=True, + alloc_mutable_buffers=True, + ) + verifier.verify_storage_reuse() + def test_multi_user_blocks_inplace(self) -> None: class Model(torch.nn.Module): def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: diff --git a/exir/tests/test_tensor.py b/exir/tests/test_tensor.py index 6435ca98a13..19a536d44d5 100644 --- a/exir/tests/test_tensor.py +++ b/exir/tests/test_tensor.py @@ -183,6 +183,23 @@ def test_allocation_info_fails(self) -> None: with self.assertRaisesRegex(Exception, test_case[1], msg=f"{kwargs}"): make_allocation_info(**kwargs) + def test_inplace_base_aliases_storage_base_at_offset_zero(self) -> None: + base = TensorSpec.from_tensor(torch.empty(4)) + child = TensorSpec.from_tensor(torch.empty(4)) + + child.inplace_base = base + + self.assertIs(child.storage_base, base) + self.assertEqual(child.storage_base_offset, 0) + self.assertIs(child.inplace_base, base) + + child.storage_base_offset = 4 + with self.assertRaisesRegex(Exception, "offset 0"): + child.inplace_base + with self.assertRaisesRegex(Exception, "storage_base_offset is 0"): + child.inplace_base = base + self.assertEqual(child.storage_base_offset, 4) + def test_contiguous_stride_from_shape(self) -> None: shape = (2, 3, 4) stride = contiguous_stride_from_shape(torch.Size(shape))