From b8ca70503622ea8ed063d1fd8cffcfd3f3f1c33e Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Thu, 27 Aug 2026 11:24:52 -0700 Subject: [PATCH] Revert "Memory planning: Support shared_allocation with offset (#21840)" This reverts commit 51fa9416ce0b0eb810fe211c1e182aa8479aaafe. --- exir/memory_planning.py | 137 +++------- exir/passes/memory_planning_pass.py | 55 +--- .../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, 49 insertions(+), 443 deletions(-) diff --git a/exir/memory_planning.py b/exir/memory_planning.py index 44b53585455..012cf8dd144 100644 --- a/exir/memory_planning.py +++ b/exir/memory_planning.py @@ -1,6 +1,5 @@ # 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. @@ -90,22 +89,6 @@ def mem_obj_id_match( return lhs_spec.mem_obj_id == rhs_spec.mem_obj_id - @classmethod - def _storage_root_chain(cls, spec: TensorSpec) -> set[TensorSpec]: - """Return spec and all TensorSpecs backing its storage.""" - seen: Set[TensorSpec] = set() - root = spec - while True: - internal_assert( - root not in seen, - "Circular storage_base relationship is not supported.", - ) - seen.add(root) - if root.storage_base is None: - break - root = root.storage_base - return seen - @classmethod def has_overlap(cls, lhs_ivl: List[int], rhs_ivl: List[int]) -> bool: r""" @@ -208,16 +191,13 @@ def verify_storage_reuse( if not allow_lifetime_and_storage_overlap and self.lifetime_overlap( lhs_spec, rhs_spec ): - # Some ops, such as in-place ops, intentionally place one - # TensorSpec inside another TensorSpec's storage despite - # overlapping lifetimes. - # This is OK if one spec is storage-backed by the other. - # Specs that merely share a root, such as siblings, should - # still be checked as normal allocations. - lhs_chain = Verifier._storage_root_chain(lhs_spec) - rhs_chain = Verifier._storage_root_chain(rhs_spec) - is_common_base_pair = lhs_spec in rhs_chain or rhs_spec in lhs_chain - if not is_common_base_pair: + # 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: raise InternalError( f"Unexpected storage overlap: {Verifier._debug_message_from_specs(lhs_spec, rhs_spec)}" ) @@ -711,7 +691,6 @@ 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 @@ -779,6 +758,21 @@ 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. @@ -945,48 +939,17 @@ def _contains_xnnpack_delegate(graph_module: torch.fx.GraphModule) -> bool: return False -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], +def _resolve_inplace_specs( + deferred_inplace: List[TensorSpec], spec2obj: Dict[TensorSpec, SharedObject], greedy_result: MemoryAlgoResult, ) -> None: - remaining = list(deferred_storage_base) + remaining = list(deferred_inplace) while remaining: progress = False next_remaining = [] for spec in remaining: - base = spec.storage_base - internal_assert( - base is not None, - "Deferred storage-backed TensorSpec should have a storage_base.", - ) + base = spec.inplace_base if base not in spec2obj: next_remaining.append(spec) continue @@ -997,47 +960,25 @@ def _resolve_storage_base_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 {base_alloc_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 + storage_base_offset, spec) - ) + sobj.allocations.append(AllocationSpec(base_alloc_offset, spec)) spec2obj[spec] = sobj if not progress: unresolved = ", ".join( f"allocated_memory={s.allocated_memory}" for s in next_remaining ) raise InternalError( - "Circular or unresolvable storage_base dependency chain: " + unresolved + f"Circular or unresolvable in-place dependency chain: {unresolved}" ) remaining = next_remaining @@ -1060,11 +1001,9 @@ def _compute_total_sizes( assert isinstance(bufsizes, list) if len(bufsizes) > mem_id: input_total_size = bufsizes[mem_id] - 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] = materialize_buffer( + shared_objects[mem_id], input_total_size + ) total_sizes[mem_id] += extra_padding for sobj in shared_objects[mem_id]: @@ -1117,7 +1056,7 @@ def greedy( sorted_specs.reverse() - deferred_storage_base: List[TensorSpec] = [] + deferred_inplace: List[TensorSpec] = [] for spec in sorted_specs: spec_alloc_result = greedy_result.spec_dict.get(spec, SpecAllocResult(0, 0, 0)) @@ -1128,8 +1067,8 @@ def greedy( greedy_result.spec_dict[spec] = spec_alloc_result spec.realign(alignment) - if spec.storage_base is not None: - deferred_storage_base.append(spec) + if spec.inplace_base is not None: + deferred_inplace.append(spec) continue spec2obj[spec] = pick_shared_obj( @@ -1138,7 +1077,7 @@ def greedy( allow_overlapping_allocations, ) - _resolve_storage_base_specs(deferred_storage_base, spec2obj, greedy_result) + _resolve_inplace_specs(deferred_inplace, spec2obj, greedy_result) total_sizes = _compute_total_sizes( shared_objects, graph_module, extra_padding, greedy_result, len(spec2obj) @@ -1267,10 +1206,10 @@ def _allocate_buf(bufsizes: List[int], mem_id: int, allocated: int) -> int: bufsizes = cast(List[int], bufsizes) for spec in specs: - if spec.storage_base is not None: + if spec.inplace_base is not None: raise InternalError( - "The naive memory planning algorithm does not support storage-backed " - "TensorSpecs. Use the greedy algorithm instead." + "The naive memory planning algorithm does not support in-place " + "element-wise ops (inplace_base). 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 6cb33b12f13..99a5f3dd8ec 100644 --- a/exir/passes/memory_planning_pass.py +++ b/exir/passes/memory_planning_pass.py @@ -1,6 +1,5 @@ # 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. @@ -137,51 +136,6 @@ 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") - shared_alloc_offset = node.meta.get("_shared_alloc_offset") - - if share_idx is None: - if shared_alloc_offset is not None: - raise ValueError( - "_shared_alloc_offset meta was set but not _share_alloc_with_arg_idx." - ) - return - if shared_alloc_offset is None: - shared_alloc_offset = 0 - - if not isinstance(share_idx, int): - raise TypeError("_share_alloc_with_arg_idx must be an int") - if not isinstance(shared_alloc_offset, int): - raise TypeError("_shared_alloc_offset must be an int") - - 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 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" - ) - - 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) @@ -242,8 +196,13 @@ 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"] - - _move_memory_meta_to_spec(node) + 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 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 947952d7692..28fcc97aaf5 100644 --- a/exir/passes/replace_view_copy_with_view_pass.py +++ b/exir/passes/replace_view_copy_with_view_pass.py @@ -1,6 +1,5 @@ # 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. @@ -109,8 +108,6 @@ 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 199c5adfafe..a4e480ffce0 100644 --- a/exir/tensor.py +++ b/exir/tensor.py @@ -212,35 +212,9 @@ def init_mem_planning_fields(self) -> None: self.mem_id = None self.mem_obj_id = None self.mem_offset = 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 + # 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 @property def dtype(self) -> torch.dtype: diff --git a/exir/tests/test_memory_planning.py b/exir/tests/test_memory_planning.py index 37bc088cfcc..31f3b1844c2 100644 --- a/exir/tests/test_memory_planning.py +++ b/exir/tests/test_memory_planning.py @@ -1,6 +1,5 @@ # 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. @@ -9,7 +8,7 @@ import itertools import unittest -from typing import Any, Callable, cast, List, Optional, Tuple, Type +from typing import Any, Callable, List, Optional, Tuple, Type import executorch.exir as exir @@ -30,7 +29,6 @@ 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, @@ -1655,221 +1653,6 @@ 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, @@ -1942,35 +1725,6 @@ 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 19a536d44d5..6435ca98a13 100644 --- a/exir/tests/test_tensor.py +++ b/exir/tests/test_tensor.py @@ -183,23 +183,6 @@ 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))