Accept AOTI's compact clones of view buffers in the CUDA weight collector (pytorch#23355)

## What is wrong today

When the CUDA backend compiles a model, AOTInductor hands back every
constant as a pair: the tensor to save, and a `TensorProperties` record
that describes it (storage size, offset, shape, strides). The CUDA
weight collector writes the tensor's storage to the `.ptd` file and
records that metadata so the runtime can rebuild the tensor.

AOTInductor clones every buffer into compact storage first. The clone
covers only the bytes the buffer needs and starts at offset 0. But
`TensorProperties` is built from the original buffer.

That breaks when a buffer is a view into a larger tensor:

```python
big = torch.arange(32.0, device="cuda").reshape(8, 4)
self.register_buffer("w", big[2:5])  # 48 bytes, at offset 8 of a 128-byte storage
```

The clone is 48 bytes at offset 0, while `TensorProperties` says 128
bytes at offset 8. The collector compares the two, and lowering with
`CudaPartitioner` fails:

```
RuntimeError: AOTI cloned storage is smaller than its TensorProperties (48 < 128 bytes)
```

A strided view fails in the bounds check instead, because its
`TensorProperties` has no storage size. A column slice `big[:, 1:3]` of
an 8x8 tensor fails with:

```
RuntimeError: AOTI view 'w' requires 236 bytes from a 232-byte cloned storage
```

In both cases the clone holds every byte the buffer needs. Only the
metadata describes a different storage.

## What this change does

When the tensor is AOTInductor's compact clone of the view (a different
storage, with the same shape and strides as `TensorProperties`), the
collector skips the comparison with the original storage size and takes
the offset from the clone it writes. Any other tensor goes through the
existing path unchanged.

The check that the view fits inside the bytes actually written still
runs for every weight, so a truncated clone is still rejected. A compact
clone is also checked to span its storage exactly, so a clone with bytes
past its view is rejected too.

## What was tested

- New unit tests in `test_cuda_partitioner.py`: a clone of `big[2:5]`
paired with the view's `TensorProperties` is written as 48 bytes at
offset 0; a clone of the strided `big[:, 1:3]` is written as its
232-byte storage, byte for byte; a clone with too little storage or with
bytes past its view is rejected; a parameter that slices a fused tensor
and keeps its storage still records its own offset and storage size; a
value with the same shape but other strides is not taken for a clone;
and a tensor with another layout keeps the offset its `TensorProperties`
records. The two clone tests fail before this change with the errors
above.
- On a CUDA GPU (Linux aarch64), the module above lowered with
`CudaPartitioner`: before, the error above; after, it exports, and
running the `.pte` with the ExecuTorch runtime gives exactly the eager
output.
- `test_cuda_partitioner.py` and `test_cuda_weight_metadata.py` pass
with this change.

Author: shoumikhin

backends/cuda/cuda_weight_collector.py

@@ -399,15 +399,29 @@ def materialize(
         storages: Dict[str, FileBackedData] = {}
 
         for fqn, (tensor, properties) in weights.items():
+            is_offgraph_kv = _is_offgraph_kv_fqn(fqn)
             storage = tensor.untyped_storage()
             storage_nbytes = storage.nbytes()
+            # AOTI clones buffers into compact storage that starts at the view,
+            # with the same shape and strides, so the original storage size and
+            # offset do not describe it.
+            is_compact_clone = (
+                not is_offgraph_kv
+                and getattr(properties, "storage_ptr", None)
+                not in (None, storage.data_ptr())
+                and tuple(tensor.shape) == tuple(getattr(properties, "shape", ()))
+                and tuple(tensor.stride()) == tuple(getattr(properties, "stride", ()))
+            )
             del storage
             device_type = device_type_for_weight(tensor)
-            is_offgraph_kv = _is_offgraph_kv_fqn(fqn)
             expected_storage_nbytes = int(
                 getattr(properties, "storage_size", None) or 0
             )
-            if not is_offgraph_kv and storage_nbytes < expected_storage_nbytes:
+            if (
+                not is_offgraph_kv
+                and not is_compact_clone
+                and storage_nbytes < expected_storage_nbytes
+            ):
                 raise RuntimeError(
                     "AOTI cloned storage is smaller than its TensorProperties "
                     f"({storage_nbytes} < {expected_storage_nbytes} bytes)"
@@ -427,7 +441,11 @@ def materialize(
             strides = tuple(
                 int(stride) for stride in getattr(properties, "stride", tensor.stride())
             )
-            storage_offset = int(getattr(properties, "offset", tensor.storage_offset()))
+            storage_offset = (
+                tensor.storage_offset()
+                if is_compact_clone
+                else int(getattr(properties, "offset", tensor.storage_offset()))
+            )
             required_nbytes = _required_view_nbytes(
                 fqn, sizes, strides, storage_offset, tensor.element_size()
             )
@@ -436,6 +454,11 @@ def materialize(
                     f"AOTI view {fqn!r} requires {required_nbytes} bytes from a "
                     f"{storage_nbytes}-byte cloned storage"
                 )
+            if is_compact_clone and required_nbytes != storage_nbytes:
+                raise RuntimeError(
+                    f"AOTI compact clone of {fqn!r} should span exactly its "
+                    f"{storage_nbytes}-byte storage, but its view needs {required_nbytes}"
+                )
             if is_offgraph_kv:
                 # Preserve the AOTI view contract; the runtime supplies storage.
                 storage_nbytes = max(

backends/cuda/tests/test_cuda_partitioner.py

@@ -39,6 +39,7 @@
 from executorch.exir.backend.partitioner import PartitionResult
 from executorch.exir.delegate import executorch_call_delegate
 from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
+from torch._inductor.compile_fx import clone_preserve_strides
 from torch.export import export
 from torch.export.pt2_archive._package_weights import TensorProperties, Weights
 from torch.fx.passes.utils.fuser_utils import validate_partition
@@ -232,8 +233,8 @@ def test_different_fqn_views_have_distinct_logical_storage(self) -> None:
         weights = Weights(
             {
                 "base": (base, TensorProperties(base)),
-                # AOTI may return a cloned value tensor; TensorProperties is
-                # the source of truth for reconstructing the original view.
+                # The value shares the view's storage, so the offset from
+                # TensorProperties applies.
                 "view": (base, TensorProperties(view)),
             }
         )
@@ -251,6 +252,101 @@ def test_different_fqn_views_have_distinct_logical_storage(self) -> None:
             for storage in artifact.storages.values():
                 storage.close()
 
+    def test_compact_clone_of_a_view_uses_the_written_storage(self) -> None:
+        # AOTI clones a buffer that views a larger tensor into compact storage
+        # but reports TensorProperties of the original view.
+        view = torch.arange(32, dtype=torch.float32).reshape(8, 4)[2:5]
+        clone = view.clone()
+        weights = Weights({"w": (clone, TensorProperties(view))})
+
+        with tempfile.TemporaryDirectory() as directory:
+            artifact = self._materialize(weights, directory)
+            entry = artifact.entries[0]
+            self.assertEqual(0, entry.storage_offset)
+            self.assertEqual(48, entry.storage_nbytes)
+            self.assertEqual((3, 4), entry.sizes)
+            data = artifact.storages[entry.storage_key].to_bytes()
+            self.assertEqual(bytes(clone.untyped_storage()), data)
+            for storage in artifact.storages.values():
+                storage.close()
+
+    def test_compact_clone_with_too_little_storage_is_rejected(self) -> None:
+        view = torch.arange(32, dtype=torch.float32).reshape(8, 4)[2:5]
+        truncated = view.clone()
+        truncated.untyped_storage().resize_(8 * truncated.element_size())
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertRaisesRegex(RuntimeError, "requires 48 bytes"):
+                self._materialize(
+                    Weights({"w": (truncated, TensorProperties(view))}), directory
+                )
+
+    def test_compact_clone_with_bytes_past_its_view_is_rejected(self) -> None:
+        view = torch.arange(32, dtype=torch.float32).reshape(8, 4)[2:5]
+        padded = view.clone()
+        padded.untyped_storage().resize_(16 * padded.element_size())
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertRaisesRegex(RuntimeError, "span exactly its 64-byte"):
+                self._materialize(
+                    Weights({"w": (padded, TensorProperties(view))}), directory
+                )
+
+    def test_value_with_the_same_shape_but_other_strides_is_not_a_clone(self) -> None:
+        view = torch.arange(32, dtype=torch.float32).reshape(8, 4)[2:5]
+        transposed = torch.arange(12, dtype=torch.float32).reshape(4, 3).t()
+        with tempfile.TemporaryDirectory() as directory:
+            with self.assertRaisesRegex(
+                RuntimeError, "smaller than its TensorProperties"
+            ):
+                self._materialize(
+                    Weights({"w": (transposed, TensorProperties(view))}), directory
+                )
+
+    def test_compact_clone_of_a_strided_view_spans_its_storage(self) -> None:
+        # A strided view's clone keeps the gaps between its rows, so it spans
+        # more bytes than its elements take.
+        view = torch.arange(64, dtype=torch.float32).reshape(8, 8)[:, 1:3]
+        clone = clone_preserve_strides(view)
+        weights = Weights({"w": (clone, TensorProperties(view))})
+
+        with tempfile.TemporaryDirectory() as directory:
+            artifact = self._materialize(weights, directory)
+            entry = artifact.entries[0]
+            self.assertEqual(0, entry.storage_offset)
+            self.assertEqual(232, entry.storage_nbytes)
+            self.assertEqual((8, 1), entry.strides)
+            data = artifact.storages[entry.storage_key].to_bytes()
+            self.assertEqual(bytes(clone.untyped_storage()), data)
+            for storage in artifact.storages.values():
+                storage.close()
+
+    def test_view_sharing_its_storage_keeps_the_view_offset(self) -> None:
+        # AOTI does not clone parameters, so a parameter that slices a fused
+        # tensor arrives in the fused tensor's storage.
+        fused = torch.arange(36, dtype=torch.float32).reshape(9, 4)
+        view = fused[3:6]
+        weights = Weights({"w": (view, TensorProperties(view))})
+
+        with tempfile.TemporaryDirectory() as directory:
+            artifact = self._materialize(weights, directory)
+            entry = artifact.entries[0]
+            self.assertEqual(12, entry.storage_offset)
+            self.assertEqual(144, entry.storage_nbytes)
+            for storage in artifact.storages.values():
+                storage.close()
+
+    def test_value_with_other_layout_keeps_the_view_offset(self) -> None:
+        # A value in a different storage that is not a compact clone of the view
+        # (other shape) keeps the offset its TensorProperties records.
+        base = torch.arange(32, dtype=torch.float32).reshape(8, 4)
+        view = base[:, 1:]
+        weights = Weights({"w": (base.clone(), TensorProperties(view))})
+
+        with tempfile.TemporaryDirectory() as directory:
+            artifact = self._materialize(weights, directory)
+            self.assertEqual(1, artifact.entries[0].storage_offset)
+            for storage in artifact.storages.values():
+                storage.close()
+
     def test_identical_values_keep_distinct_fqn_keys(self) -> None:
         first = torch.zeros(4)
         second = torch.zeros(4)