Skip to content

Commit aed854f

Browse files
Arm backend: Add as_strided_copy.default decomposition and test (pytorch#16702)
Change-Id: Iffedfed51bba0ef63524b0aed7ce63c36881eaea ### Summary Arm backend: Add as_strided_copy.default decomposition and test @zingo @per @freddan80 @oscarandersson8218 @Sebastian-Larsson cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai --------- Signed-off-by: Xingguo Li <xingguo.li@arm.com> Co-authored-by: Oscar Andersson <87121123+oscarandersson8218@users.noreply.github.com>
1 parent 06f10b9 commit aed854f

9 files changed

Lines changed: 443 additions & 6 deletions

File tree

‎backends/arm/_passes/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from .decompose_add_sub_alpha_pass import DecomposeAddSubAlphaPass # noqa
3232
from .decompose_addmm_pass import DecomposeAddmmPass # noqa
3333
from .decompose_any_pass import DecomposeAnyPass # noqa
34+
from .decompose_as_strided_copy_pass import DecomposeAsStridedCopyPass # noqa
3435
from .decompose_asin_and_acos_pass import DecomposeAsinAndAcosPass # noqa
3536
from .decompose_asinh_pass import DecomposeAsinhPass # noqa
3637
from .decompose_atan_pass import DecomposeAtanPass # noqa

‎backends/arm/_passes/arm_pass_manager.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
DecomposeAnyPass,
3939
DecomposeAsinAndAcosPass,
4040
DecomposeAsinhPass,
41+
DecomposeAsStridedCopyPass,
4142
DecomposeAtanhPass,
4243
DecomposeAtanPass,
4344
DecomposeAvgPool2dPass,
@@ -333,6 +334,7 @@ def _tosa_pipeline(
333334
ConvertExpandCopyToRepeatPass(),
334335
UnsqueezeBeforeRepeatPass(),
335336
DecomposeCumsumPass(exported_program),
337+
DecomposeAsStridedCopyPass(),
336338
DecomposeMaxPool2dPass(),
337339
SizeAdjustInputPass(),
338340
DecomposeSelectPass(),
Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from typing import Dict, Optional, Set, Tuple, Type
7+
8+
import torch
9+
10+
from executorch.backends.arm._passes import ArmPass
11+
from executorch.backends.arm.common.as_strided_utils import (
12+
contiguous_strides,
13+
maybe_static_sequence,
14+
to_int,
15+
to_int_tuple,
16+
)
17+
from executorch.exir.dialects._ops import ops as exir_ops
18+
from executorch.exir.pass_base import ExportPass
19+
20+
21+
class DecomposeAsStridedCopyPass(ArmPass):
22+
"""
23+
Replace contiguous `aten.as_strided_copy` with `aten.view_copy`.
24+
25+
The TOSA backend only supports the contiguous-as-strided case where the stride matches
26+
row-major layout and the storage offset is zero. In that scenario the operator is
27+
equivalent to a reshape with copy semantics and can be lowered via `view_copy`.
28+
"""
29+
30+
_passes_required_after: Set[Type[ExportPass]] = set()
31+
32+
_EDGE_OPS = (exir_ops.edge.aten.as_strided_copy.default,)
33+
_ATEN_OPS = (torch.ops.aten.as_strided_copy.default,)
34+
35+
def _extract_args(
36+
self, args: Tuple[object, ...], kwargs: dict
37+
) -> Optional[Tuple[Tuple[int, ...], Tuple[int, ...], int]]:
38+
"""Return (size, stride, storage_offset) when they are statically known."""
39+
if len(args) < 3:
40+
return None
41+
42+
size_arg = args[1]
43+
stride_arg = args[2]
44+
offset_arg = (
45+
kwargs.get("storage_offset") if "storage_offset" in kwargs else None
46+
)
47+
if offset_arg is None and len(args) > 3:
48+
offset_arg = args[3]
49+
50+
size_seq = maybe_static_sequence(size_arg)
51+
stride_seq = maybe_static_sequence(stride_arg)
52+
if size_seq is None or stride_seq is None:
53+
return None
54+
55+
size_tuple = to_int_tuple(size_seq)
56+
stride_tuple = to_int_tuple(stride_seq)
57+
if size_tuple is None or stride_tuple is None:
58+
return None
59+
60+
if len(size_tuple) != len(stride_tuple):
61+
return None
62+
63+
if any(stride < 0 for stride in stride_tuple):
64+
return None
65+
66+
if offset_arg is None:
67+
storage_offset = 0
68+
else:
69+
parsed_offset = to_int(offset_arg)
70+
if parsed_offset is None:
71+
return None
72+
storage_offset = parsed_offset
73+
74+
return size_tuple, stride_tuple, storage_offset
75+
76+
def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
77+
if op not in (*self._EDGE_OPS, *self._ATEN_OPS):
78+
return super().call_operator(op, args, kwargs, meta, updated)
79+
80+
extracted = self._extract_args(args, kwargs)
81+
if extracted is None:
82+
return super().call_operator(op, args, kwargs, meta, updated)
83+
84+
size_tuple, stride_tuple, storage_offset = extracted
85+
if storage_offset != 0:
86+
return super().call_operator(op, args, kwargs, meta, updated)
87+
88+
expected_strides = contiguous_strides(size_tuple)
89+
90+
def _stride_matches(idx: int, dim: int) -> bool:
91+
stride = stride_tuple[idx]
92+
expected = expected_strides[idx]
93+
if idx == len(size_tuple) - 1:
94+
return stride >= expected
95+
if dim == 1 or expected == 0:
96+
return True
97+
return stride == expected
98+
99+
if any(not _stride_matches(i, dim) for i, dim in enumerate(size_tuple)):
100+
return super().call_operator(op, args, kwargs, meta, updated)
101+
102+
view_args = (args[0], tuple(size_tuple))
103+
view_kwargs: Dict[str, object] = {}
104+
105+
view_op = (
106+
exir_ops.edge.aten.view_copy.default
107+
if op in self._EDGE_OPS
108+
else torch.ops.aten.view_copy.default
109+
)
110+
111+
return super().call_operator(
112+
view_op, view_args, view_kwargs, meta, updated=True
113+
)
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
"""Utility helpers shared across as_strided_copy handling."""
6+
7+
from __future__ import annotations
8+
9+
import numbers
10+
11+
from collections.abc import Sequence
12+
from typing import Optional, Tuple, TypeVar
13+
14+
import torch
15+
import torch.fx as fx
16+
17+
T = TypeVar("T", bound=Sequence)
18+
19+
20+
def to_int(value: object) -> Optional[int]:
21+
"""Return an int for supported numeric types, otherwise None."""
22+
if isinstance(value, (numbers.Integral, torch.SymInt)):
23+
return int(value)
24+
return None
25+
26+
27+
def maybe_static_sequence(value: object) -> Optional[Sequence]:
28+
"""
29+
Return a Python sequence for literal or FX-constant values.
30+
31+
FX exporters often wrap constant lists in nodes where the materialised
32+
value is stored in ``node.meta["val"]``. This helper unwraps that so the
33+
rest of the logic can treat them uniformly.
34+
"""
35+
if isinstance(value, (str, bytes)):
36+
return None
37+
if isinstance(value, fx.Node):
38+
const_val = value.meta.get("val")
39+
if isinstance(const_val, Sequence):
40+
return const_val
41+
return None
42+
if isinstance(value, Sequence):
43+
return value
44+
return None
45+
46+
47+
def to_int_tuple(value: object) -> Optional[Tuple[int, ...]]:
48+
"""Best-effort conversion of a sequence of integers/SymInts to a tuple[int, ...]."""
49+
seq = maybe_static_sequence(value)
50+
if seq is None:
51+
return None
52+
53+
result: list[int] = []
54+
for item in seq:
55+
converted = to_int(item)
56+
if converted is None:
57+
return None
58+
result.append(converted)
59+
return tuple(result)
60+
61+
62+
def contiguous_strides(shape: Sequence[int]) -> Tuple[int, ...]:
63+
"""Compute row-major contiguous strides for the provided shape."""
64+
strides = [0] * len(shape)
65+
running = 1
66+
for idx in reversed(range(len(shape))):
67+
dim_val = shape[idx]
68+
strides[idx] = running if dim_val != 0 else 1
69+
running *= max(dim_val, 1)
70+
return tuple(strides)

‎backends/arm/operator_support/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66

77
from . import ( # noqa
8+
as_strided_copy_support,
89
clone_dim_order_support,
910
control_flow_support,
1011
convolution_support,
Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
"""Declare operator support for aten.as_strided_copy in the TOSA backend."""
6+
7+
from collections.abc import Mapping
8+
from typing import Any, Optional
9+
10+
import torch.fx as fx
11+
from executorch.backends.arm.common.as_strided_utils import (
12+
contiguous_strides,
13+
maybe_static_sequence,
14+
to_int,
15+
to_int_tuple,
16+
)
17+
18+
from executorch.backends.arm.operator_support.tosa_supported_operators import (
19+
register_tosa_support_check,
20+
SupportedTOSAOperatorCheck,
21+
)
22+
from executorch.backends.arm.tosa.specification import TosaSpecification
23+
from executorch.exir.dialects._ops import ops as exir_ops
24+
25+
26+
def _arg_from_node(node: fx.Node, position: int, keyword: str) -> object | None:
27+
"""Fetch an argument either by keyword or positional index."""
28+
kwargs: Mapping[str, Any] = node.kwargs
29+
if keyword in kwargs:
30+
return kwargs[keyword]
31+
if len(node.args) > position:
32+
return node.args[position]
33+
return None
34+
35+
36+
def _extract_static_args(
37+
node: fx.Node,
38+
) -> Optional[tuple[tuple[int, ...], tuple[int, ...], int]]:
39+
"""Return static size/stride/offset if they are compatible."""
40+
size_arg = _arg_from_node(node, 1, "size")
41+
stride_arg = _arg_from_node(node, 2, "stride")
42+
offset_arg = _arg_from_node(node, 3, "storage_offset")
43+
44+
if (
45+
maybe_static_sequence(size_arg) is None
46+
or maybe_static_sequence(stride_arg) is None
47+
):
48+
return None
49+
50+
size_tuple = to_int_tuple(size_arg)
51+
stride_tuple = to_int_tuple(stride_arg)
52+
if size_tuple is None or stride_tuple is None:
53+
return None
54+
55+
if len(size_tuple) != len(stride_tuple):
56+
return None
57+
58+
if any(stride < 0 for stride in stride_tuple):
59+
return None
60+
61+
storage_offset = 0
62+
if offset_arg is not None:
63+
parsed_offset = to_int(offset_arg)
64+
if parsed_offset is None:
65+
return None
66+
storage_offset = parsed_offset
67+
68+
return size_tuple, stride_tuple, storage_offset
69+
70+
71+
@register_tosa_support_check
72+
class AsStridedCopySupported(SupportedTOSAOperatorCheck):
73+
"""Support check ensuring as_strided_copy is contiguous with zero offset."""
74+
75+
targets = [exir_ops.edge.aten.as_strided_copy.default]
76+
tosa_specs = [
77+
TosaSpecification.create_from_string("TOSA-1.0+INT"),
78+
TosaSpecification.create_from_string("TOSA-1.0+FP"),
79+
]
80+
81+
def is_node_tosa_supported(
82+
self, node: fx.Node, tosa_spec: TosaSpecification # noqa: D417
83+
) -> bool:
84+
extracted = _extract_static_args(node)
85+
if extracted is None:
86+
self.reporter.report_reject(
87+
node, "Size/stride must be static with non-negative strides."
88+
)
89+
return False
90+
91+
size_tuple, stride_tuple, storage_offset = extracted
92+
93+
if storage_offset != 0:
94+
self.reporter.report_reject(
95+
node, "Non-zero storage offsets are unsupported."
96+
)
97+
return False
98+
99+
expected_strides = contiguous_strides(size_tuple)
100+
101+
def _stride_matches(idx: int, dim: int) -> bool:
102+
stride = stride_tuple[idx]
103+
expected = expected_strides[idx]
104+
if idx == len(size_tuple) - 1:
105+
return stride >= expected
106+
if dim == 1:
107+
return True
108+
return stride == expected
109+
110+
if any(not _stride_matches(i, dim) for i, dim in enumerate(size_tuple)):
111+
self.reporter.report_reject(
112+
node,
113+
f"Stride {stride_tuple} is not contiguous for shape {size_tuple}.",
114+
)
115+
return False
116+
117+
return True

‎backends/arm/quantizer/quantization_annotator.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,7 @@ def _match_pattern(
432432
torch.ops.aten.unfold_copy.default,
433433
torch.ops.aten.index_select.default,
434434
torch.ops.aten.index.Tensor,
435+
torch.ops.aten.as_strided_copy.default,
435436
# Neg operator flips the range, but keps the magnitude the same.
436437
# That is why we force it to use the same qparams and avoid
437438
# dequant -> neg -> requant chain.

‎backends/arm/test/models/test_mobilenet_v2_arm.py‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
# Copyright (c) Meta Platforms, Inc. and affiliates.
22
# All rights reserved.
3-
# Copyright 2024-2025 Arm Limited and/or its affiliates.
3+
# Copyright 2024-2026 Arm Limited and/or its affiliates.
44
#
55
# This source code is licensed under the BSD-style license found in the
66
# LICENSE file in the root directory of this source tree.
@@ -64,11 +64,6 @@ def test_mv2_tosa_FP_channels_last():
6464
exir_op=[],
6565
use_to_edge_transform_and_lower=True,
6666
)
67-
# Changing memory format leads to an unsupported as_strided_copy op being inserted into the graph,
68-
# leading to a graph break.
69-
pipeline.change_args(
70-
"check_count.exir", {"torch.ops.higher_order.executorch_call_delegate": 2}
71-
)
7267
pipeline.run()
7368

7469

0 commit comments

Comments
 (0)