Skip to content

Fix test_mutation_not_double_allocated to not use hardcoded indices #16667

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mergennachin merged 1 commit into from
Feb 5, 2026
Merged

Fix test_mutation_not_double_allocated to not use hardcoded indices #16667

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 44 additions & 15 deletions exir/tests/test_memory_planning.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,22 +665,51 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:

et = to_edge(export(model, inputs, strict=True)).to_executorch()

# 0 and 11 should refer to the same tensor. 0 is the input, 11 is the output of copy_
self.assertEqual(
et.executorch_program.execution_plan[0]
.values[0]
.val.allocation_info.memory_offset_low,
et.executorch_program.execution_plan[0]
.values[11]
.val.allocation_info.memory_offset_low,
)
# The mutable buffer (5x5 float32 = 100 bytes) should not be double allocated.
# The input and output of copy_ should share the same memory location.
values = et.executorch_program.execution_plan[0].values
expected_buffer_size = 5 * 5 * 4 # 5x5 float32

# Collect all tensor allocations by their (memory_id, offset) and track sizes
# Size is computed from tensor's sizes and scalar_type, not from allocation_info
# (memory_offset_low/high are low/high 32-bit parts of a 64-bit offset, not bounds)
scalar_type_sizes = {
0: 1, # BYTE
1: 1, # CHAR
2: 2, # SHORT
3: 4, # INT
4: 8, # LONG
5: 2, # HALF
6: 4, # FLOAT
7: 8, # DOUBLE
}
Comment on lines +676 to +685
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The scalar_type_sizes dictionary is incomplete and missing many scalar types defined in ScalarType enum (exir/scalar_type.py). Missing types include: BOOL (11), QINT8 (12), QUINT8 (13), QINT32 (14), BFLOAT16 (15), QUINT4x2 (16), QUINT2x4 (17), BITS16 (22), FLOAT8E5M2 (23), FLOAT8E4M3FN (24), FLOAT8E5M2FNUZ (25), FLOAT8E4M3FNUZ (26), UINT16 (27), UINT32 (28), UINT64 (29), and COMPLEX32 (8), COMPLEX64 (9), COMPLEX128 (10).

While the test currently only expects float32 tensors, if other tensor types appear in the graph, the code will fall back to the default size of 4 bytes, which could produce incorrect size calculations and cause the test to incorrectly pass or fail. Consider either:

  1. Adding all scalar types to the mapping (see devtools/inspector/_inspector_utils.py:172-184 for a similar comprehensive mapping), or
  2. Raising an error for unmapped types instead of silently using a default value

Copilot uses AI. Check for mistakes.
offset_to_indices = {}
for i, val in enumerate(values):
tensor = val.val
if hasattr(tensor, "allocation_info") and tensor.allocation_info:
alloc = tensor.allocation_info
# Compute tensor size from sizes and scalar_type
num_elements = 1
for dim in tensor.sizes:
num_elements *= dim
element_size = scalar_type_sizes.get(int(tensor.scalar_type), 4)
size = num_elements * element_size
key = (alloc.memory_id, alloc.memory_offset)
if key not in offset_to_indices:
offset_to_indices[key] = {"indices": [], "size": size}
Copy link

Copilot AI Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When multiple values share the same (memory_id, memory_offset), the code only stores the size from the first value encountered and doesn't verify that subsequent values at the same location have the same computed size. If values sharing memory have different sizes, this could indicate a memory planning bug that would go undetected. Consider adding a check to verify that all values at the same memory location have the same computed size, or at least store and compare them.

Suggested change
offset_to_indices[key] = {"indices": [], "size": size}
offset_to_indices[key] = {"indices": [], "size": size}
else:
existing_size = offset_to_indices[key]["size"]
self.assertEqual(
existing_size,
size,
f"Inconsistent tensor sizes for shared memory location "
f"{key}: previously {existing_size}, now {size}",
)

Copilot uses AI. Check for mistakes.
offset_to_indices[key]["indices"].append(i)

# Find shared allocations matching the mutable buffer size (before/after copy_)
mutable_buffer_shares = [
info
for info in offset_to_indices.values()
if len(info["indices"]) == 2 and info["size"] == expected_buffer_size
]
self.assertEqual(
et.executorch_program.execution_plan[0]
.values[0]
.val.allocation_info.memory_offset_high,
et.executorch_program.execution_plan[0]
.values[11]
.val.allocation_info.memory_offset_high,
len(mutable_buffer_shares),
1,
f"Expected exactly one shared allocation of size {expected_buffer_size} "
f"with 2 values (copy_ input/output), found: {mutable_buffer_shares}",
)

def test_mutable_buffers_infinite_lifespan(self) -> None:
Expand Down
Loading