From 188978a1c9be0d297167b1907a7639190b991b95 Mon Sep 17 00:00:00 2001 From: Oscar Andersson Date: Wed, 10 Jun 2026 17:06:21 +0200 Subject: [PATCH] Arm backend: Materialize symbolic shapes Arguments with symbolic shapes and operations on SymInts are materialized by SymbolicToTosaShapesPass and SymbolicMaterializationHelper. Signed-off-by: Oscar Andersson Change-Id: I083ce910694254c8d07e51407b9890af189abbd8 --- backends/arm/SYMINT_ET_LOWERING.md | 354 ++++++++++++ backends/arm/_passes/arm_pass_manager.py | 3 +- .../arm/_passes/insert_dynamic_padding.py | 30 +- .../resolve_view_copy_inferred_dim_pass.py | 5 +- .../symbolic_materialization_helper.py | 153 ++++++ .../_passes/symbolic_to_tosa_shape_pass.py | 225 +++++++- .../test_symbolic_materialization_helper.py | 234 ++++++++ .../test_insert_dynamic_padding_pass.py | 33 +- .../test_symbolic_to_tosa_shape_pass.py | 516 +++++++++++++++++- 9 files changed, 1494 insertions(+), 59 deletions(-) create mode 100644 backends/arm/SYMINT_ET_LOWERING.md create mode 100644 backends/arm/_passes/symbolic_materialization_helper.py create mode 100644 backends/arm/test/misc/test_symbolic_materialization_helper.py diff --git a/backends/arm/SYMINT_ET_LOWERING.md b/backends/arm/SYMINT_ET_LOWERING.md new file mode 100644 index 00000000000..d1cdfb6b906 --- /dev/null +++ b/backends/arm/SYMINT_ET_LOWERING.md @@ -0,0 +1,354 @@ +# Symbolic TOSA Shape Lowering + +This note explains how `resolve_view_copy_inferred_dim_pass.py`, +`symbolic_to_tosa_shape_pass.py`, and `symbolic_materialization_helper.py` +work together. + +## Goal + +PyTorch export represents dynamic shape values as scalar symbolic values such as +`SymInt`s and FX nodes like `aten.sym_size.int`. TOSA shape ops represent shape +values as one-dimensional shape tensors, modeled in fake execution as +`SymInt[]`. + +The pass converts scalar symbolic shape computations and args containg `SymInt[]` +into TOSA shape operands. + +For example, this FX-style shape computation: + +```python +class GraphModule(torch.nn.Module): + def forward(self, x: "f32[s0, 3, 4]"): + sym_size = torch.ops.aten.sym_size.int(x, 0) + add = operator.add(sym_size, 1) + view = torch.ops.aten.view.default(x, [add, 12]) + return (view,) +``` + +is rewritten conceptually into: + +```python +class GraphModule(torch.nn.Module): + def forward(self, x: "f32[s0, 3, 4]"): + dim = executorch_exir_dialects_edge__ops_backend_tosa_DIM_default( + x, + axis=0, + ) + one = executorch_exir_dialects_edge__ops_backend_tosa_CONST_SHAPE_default( + [1] + ) + add = executorch_exir_dialects_edge__ops_backend_tosa_ADD_SHAPE_default( + dim, + one, + ) + twelve = executorch_exir_dialects_edge__ops_backend_tosa_CONST_SHAPE_default( + [12] + ) + shape = executorch_exir_dialects_edge__ops_backend_tosa_CONCAT_SHAPE_default( + [add, twelve] + ) + view = torch.ops.aten.view.default(x, shape) + return (view,) +``` + +The important representation change is: + +```text +scalar symbolic value -> TOSA shape operand +s0 -> [s0] +s0 + 1 -> [s0 + 1] +[scalar pieces] -> CONCAT_SHAPE([...]) +``` + +## `ResolveViewCopyInferredDimPass` + +`ResolveViewCopyInferredDimPass` runs before `SymbolicToTosaShapesPass` and +normalizes view shapes that use one inferred dimension. For example: + +```text +view_copy(x, [sym_size(x, 0), -1]) +``` + +The `-1` is replaced with the inferred output dimension from fake tensor +metadata. If that dimension is dynamic, the pass uses +`Graph.materialize_symints(...)` to lift the raw `torch.SymInt` expression into +ordinary FX symbolic IR rooted at existing producers, such as +`aten.sym_size.int` and `operator.mul`. After this pass, +`SymbolicToTosaShapesPass` sees an explicit shape expression rather than a +view-specific `-1` convention. + +## `SymbolicToTosaShapesPass` + +`SymbolicToTosaShapesPass` is an `ArmPass` that intercepts symbolic shape nodes +while the FX graph is being transformed. + +It handles three cases. + +### 1. `aten.sym_size.int` + +`aten.sym_size.int(x, dim)` reads one dynamic dimension from a tensor. The pass +lowers it to TOSA `DIM`: + +```text +aten.sym_size.int(x, 0) -> tosa.DIM(x, axis=0) +``` + +The original scalar `SymInt` result becomes a one-element TOSA shape operand. + +Example: + +```python +# before +sym_size = torch.ops.aten.sym_size.int(x, 0) + +# after +dim = tosa.DIM.default(x, axis=0) +``` + +`DIM` is special because the axis stays as a keyword argument. The helper does +not materialize the axis as a TOSA shape operand. + +### 2. Symbolic Arithmetic + +The pass lowers supported Python symbolic arithmetic operators to TOSA shape +arithmetic ops: + +```text +operator.add -> tosa.ADD_SHAPE +operator.sub -> tosa.SUB_SHAPE +operator.mul -> tosa.MUL_SHAPE +operator.mod -> tosa.MOD_SHAPE +operator.floordiv -> tosa.DIV_FLOOR_SHAPE +``` + +For example: + +```python +class GraphModule(torch.nn.Module): + def forward(self, x: "f32[s0, s1, 4]"): + s0 = torch.ops.aten.sym_size.int(x, 0) + s1 = torch.ops.aten.sym_size.int(x, 1) + product = operator.mul(s0, s1) + view = torch.ops.aten.view.default(x, [product, 4]) + return (view,) +``` + +becomes conceptually: + +```python +dim_0 = tosa.DIM.default(x, axis=0) +dim_1 = tosa.DIM.default(x, axis=1) +product = tosa.MUL_SHAPE.default(dim_0, dim_1) +four = tosa.CONST_SHAPE.default([4]) +shape = tosa.CONCAT_SHAPE.default([product, four]) +view = torch.ops.aten.view.default(x, shape) +``` + +Only symbolic arithmetic with at least one TOSA shape operand is lowered. +Ordinary scalar or Python arithmetic is forwarded to the base pass. Unsupported +arithmetic with a shape operand raises `NotImplementedError`. + +Nested expressions are lowered one symbolic operation at a time. For example, +this exported FX form: + +```python +s0 = torch.ops.aten.sym_size.int(x, 0) +s1 = torch.ops.aten.sym_size.int(x, 1) +sub = operator.sub(s0, 1) +add = operator.add(sub, s1) +floordiv = operator.floordiv(add, 2) +``` + +is represented as a chain like: + +```text +DIM(x, axis=0) +CONST_SHAPE([1]) +SUB_SHAPE(dim_0, one) +DIM(x, axis=1) +ADD_SHAPE(sub, dim_1) +CONST_SHAPE([2]) +DIV_FLOOR_SHAPE(add, two) +``` + +### 3. Shape Lists Used By Operators + +Some operators receive shape-like Python lists or tuples. For example `view` +can receive a shape argument like: + +```python +sym_size = torch.ops.aten.sym_size.int(x, 0) +view = torch.ops.aten.view.default(x, [sym_size, 12]) +``` + +After `sym_size` has been lowered to a TOSA shape `ProxyValue`, this list is a +mixed Python container: + +```text +[ProxyValue(DIM), 12] +``` + +The pass detects list or tuple arguments containing at least one TOSA shape +`ProxyValue`, including nested shape proxies. It then asks the helper to turn +the whole container into a single TOSA shape operand: + +```text +[DIM(x, axis=0), 12] + -> CONST_SHAPE([12]) + -> CONCAT_SHAPE([DIM, CONST_12]) +``` + +The rewritten operator receives the `CONCAT_SHAPE` result instead of the Python +list. + +## Shape-Marked Nodes + +The pass uses `meta_has_shape_mark(...)` to distinguish TOSA shape values from +ordinary tensor values. + +If an operator result is already marked as a TOSA shape value, the pass forwards +it unchanged: + +```python +if meta_has_shape_mark(meta.data): + return super().call_operator(op, args, kwargs, meta, updated) +``` + +This avoids recursively trying to lower TOSA shape ops that are already in the +right representation. + +## `SymbolMaterializationHelpers` + +`SymbolMaterializationHelpers` owns the TOSA-specific materialization logic and +cache. The pass decides *when* something should become a TOSA shape value; the +helper decides *how* to build or reuse the required TOSA shape nodes. + +The helper accepts valid shape pieces made from: + +- `ProxyValue` objects that already produce TOSA shape values; +- Python `int` constants; +- nested Python `list` or `tuple` containers of those values. + +It is not responsible for lowering arbitrary raw `torch.SymInt` values. Those +should already have FX producers by the time this helper is used. For inferred +`view_copy` dimensions, `ResolveViewCopyInferredDimPass` creates those FX +producers before TOSA shape materialization runs. + +## `materialize_arglist(...)` + +`materialize_arglist(shape_arg, meta)` converts a Python shape container into a +single TOSA shape operand. + +For a single existing shape proxy, it reuses the proxy: + +```text +[DIM(x, axis=0)] -> DIM(x, axis=0) +``` + +For an integer, it creates or reuses `CONST_SHAPE`: + +```text +[7] -> CONST_SHAPE([7]) +``` + +For multiple pieces, it flattens nested lists/tuples, materializes each piece, +and creates `CONCAT_SHAPE`: + +```text +[[DIM(x, axis=0)], [2, 3]] + -> DIM(x, axis=0) + -> CONST_SHAPE([2]) + -> CONST_SHAPE([3]) + -> CONCAT_SHAPE([DIM, CONST_2, CONST_3]) +``` + +If the same integer constant is needed again, the cached `CONST_SHAPE` node is +reused: + +```text +materialize_arglist([5]) -> CONST_SHAPE([5]) +materialize_arglist([dim, 5]) -> CONCAT_SHAPE([dim, cached_CONST_5]) +``` + +## `materialize_shape_op(...)` + +`materialize_shape_op(target, args, kwargs, meta)` creates or reuses a TOSA +shape op result. + +For non-`DIM` shape ops, each argument is first converted into a TOSA shape +operand with `materialize_arglist(...)`: + +```text +ADD_SHAPE(dim, 1) + -> ADD_SHAPE(dim, CONST_SHAPE([1])) +``` + +For `DIM`, the input tensor is passed through directly and the axis remains in +`kwargs`: + +```text +DIM(x, axis=1) +``` + +The helper caches shape-op outputs by `str(meta.data["val"])`. If the same +symbolic output shape is requested again, the existing `ProxyValue` is reused. + +## Cache Behavior + +The helper has one cache: + +```python +self._shape_to_proxyval: dict[str, ProxyValue] +``` + +It stores: + +- integer constants under `str(value)`, for example `"1"`; +- shape-op results under `str(meta.data["val"])`, for example `"[s0 + 1]"`. + +This cache is local to one helper/pass instance. It is used to avoid duplicating +shape producers while lowering one graph. + +## Metadata + +Shape nodes are created through `ArmPass.call_shape_operator(...)`. The helper +passes through the `NodeMetadata` it received, so metadata such as `val`, debug +handles, and TOSA shape markers can be attached by the pass infrastructure. + +The tests assert that metadata is preserved on emitted `CONST_SHAPE` and +`CONCAT_SHAPE` nodes. + +## Canonicalizing Mixed Shape Containers + +Arm/TOSA lowering code may temporarily build mixed Python containers of integer +constants, shape-producing `ProxyValue`s, and nested lists or tuples. That is a +normal intermediate form while constructing shape arguments. + +`InsertDynamicPaddingPass` is the main current producer of this form. For a +dynamic 2D convolution or pool, it rewrites implicit spatial padding into an +explicit `PAD` op and resets the original op padding to zeros. The new `PAD` +argument is intentionally left as a flattened Python list: + +```text +[0, 0, *spatial_padding, 0, 0] +``` + +where `spatial_padding` may contain shape-producing proxies that came from +`Graph.materialize_symints(...)`. Before this container reaches serialization, +`ResolveViewCopyInferredDimPass` and `SymbolicToTosaShapesPass` run after +`InsertDynamicPaddingPass`. The symbolic pass detects the shape-marked values +in the list and asks the helper to flatten and materialize +the pieces. Conceptually, the 2D padding list becomes a single shape operand such as: + +```text +CONCAT_SHAPE([ + CONST_SHAPE([0]), + CONST_SHAPE([0]), + spatial_pad_0, + spatial_pad_1, + spatial_pad_2, + spatial_pad_3, + CONST_SHAPE([0]), + CONST_SHAPE([0]), +]) +``` \ No newline at end of file diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index a0540d4a639..398f3876cca 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -698,8 +698,9 @@ def _tosa_pipeline( FuseEqualPlaceholdersPass(exported_program), NormalizeTransformInputPlaceholdersPass(exported_program), ExirToTosaPass(exported_program), - SymbolicToTosaShapesPass(), InsertDynamicPaddingPass(), + ResolveViewCopyInferredDimPass(), + SymbolicToTosaShapesPass(), FuseConsecutiveConcatShapesPass(), RemoveNoopPass(), # Fuse duplicates exposed by late rewrites before inserting rescales; diff --git a/backends/arm/_passes/insert_dynamic_padding.py b/backends/arm/_passes/insert_dynamic_padding.py index 405c996af2e..344f7cd5b61 100644 --- a/backends/arm/_passes/insert_dynamic_padding.py +++ b/backends/arm/_passes/insert_dynamic_padding.py @@ -6,8 +6,6 @@ from typing import Set, Type from executorch.backends.arm._passes import ArmOpTargetedPass -from executorch.backends.arm.tosa.dialect.shape import is_shape_op_node - from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, ProxyValue @@ -34,13 +32,10 @@ class InsertDynamicPaddingPass(ArmOpTargetedPass): ) def _is_dynamic_padding( - self, padding: ProxyValue | list[int] | tuple[int, ...] + self, padding: list[int | ProxyValue] | tuple[int | ProxyValue, ...] ) -> bool: - return (isinstance(padding, ProxyValue) and is_shape_op_node(padding.node)) or ( - ( - isinstance(padding, (list, tuple)) - and any(isinstance(p, ProxyValue) for p in padding) - ) + return isinstance(padding, (list, tuple)) and any( + isinstance(p, ProxyValue) for p in padding ) def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue: @@ -62,24 +57,7 @@ def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue: zero_padding_pair = [0, 0] spatial_rank = 3 if op == exir_ops.backend.tosa.CONV3D.default else 2 zero_spatial_padding = [0] * (spatial_rank * 2) - N_padding = super().call_shape_operator( - exir_ops.backend.tosa.CONST_SHAPE.default, - (zero_padding_pair,), - {}, - meta, - True, - ) - C_padding = N_padding - - padding_shape_args = [N_padding, padding, C_padding] - - padding_shape = super().call_shape_operator( - exir_ops.backend.tosa.CONCAT_SHAPE.default, - (padding_shape_args,), - {}, - meta, - True, - ) + padding_shape = [*zero_padding_pair, *padding, *zero_padding_pair] pad_res = super().call_operator( exir_ops.backend.tosa.PAD.default, diff --git a/backends/arm/_passes/resolve_view_copy_inferred_dim_pass.py b/backends/arm/_passes/resolve_view_copy_inferred_dim_pass.py index 0d4a4702a17..df63665586b 100644 --- a/backends/arm/_passes/resolve_view_copy_inferred_dim_pass.py +++ b/backends/arm/_passes/resolve_view_copy_inferred_dim_pass.py @@ -7,6 +7,9 @@ from executorch.backends.arm._passes.arm_pass import ArmPass from executorch.backends.arm._passes.symbolic_shape_utils import materialize_symints +from executorch.backends.arm._passes.symbolic_to_tosa_shape_pass import ( + SymbolicToTosaShapesPass, +) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import PassResult @@ -14,7 +17,7 @@ class ResolveViewCopyInferredDimPass(ArmPass): """Materialize inferred view dimensions before TOSA shape lowering.""" - _passes_required_after = set() + _passes_required_after = {SymbolicToTosaShapesPass} target_ops = { torch.ops.aten.view.default, exir_ops.edge.aten.view_copy.default, diff --git a/backends/arm/_passes/symbolic_materialization_helper.py b/backends/arm/_passes/symbolic_materialization_helper.py new file mode 100644 index 00000000000..5dbe9610c68 --- /dev/null +++ b/backends/arm/_passes/symbolic_materialization_helper.py @@ -0,0 +1,153 @@ +# 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. + +import copy +import logging +from typing import Iterable, List, Tuple + +from executorch.backends.arm._passes import ArmPass +from executorch.backends.arm.tosa.dialect.shape import meta_has_shape_mark +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import NodeMetadata, ProxyValue + + +ShapeList = List | Tuple + +logger = logging.getLogger(__name__) + + +class SymbolMaterializationHelpers: + """Build canonical TOSA shape operands for symbolic shape lowering.""" + + def __init__(self, owning_pass: ArmPass): + self._shape_to_proxyval: dict[str, ProxyValue] = {} + self.builder = owning_pass + + def _ensure_value( + self, + value: ProxyValue | int, + meta: NodeMetadata, + ) -> ProxyValue: + if isinstance(value, ProxyValue): + if not meta_has_shape_mark(value.node.meta) and isinstance(value.data, int): + logger.debug( + "Materializing scalar ProxyValue node=%s data=%s", + value.node.name, + value.data, + ) + return self._materialize_int(value.data, meta) + logger.debug( + "Using existing ProxyValue node=%s data=%s", + value.node.name, + value.data, + ) + return value + elif isinstance(value, int): + logger.debug("Materializing integer shape value=%s", value) + return self._materialize_int(value, meta) + else: + logger.debug("Unsupported symbolic materialization value=%r", value) + raise TypeError( + f"Unsupported value type {type(value)} for symbolic materialization" + ) + + def materialize_arglist( + self, shape_arg: ShapeList, meta: NodeMetadata + ) -> ProxyValue: + logger.debug("Materializing shape arglist=%s", shape_arg) + elements = list(self._iter_materialized_shape_elements(shape_arg, meta)) + logger.debug( + "Materialized arglist elements=%s", + [(element.node.name, element.data) for element in elements], + ) + if len(elements) == 1: + logger.debug( + "Arglist has one element; reusing node=%s", elements[0].node.name + ) + return elements[0] + logger.debug("Creating CONCAT_SHAPE for %d elements", len(elements)) + return self.builder.call_shape_operator( + exir_ops.backend.tosa.CONCAT_SHAPE.default, + (elements,), + {}, + meta, + True, + ) + + def _iter_materialized_shape_elements( + self, + shape_arg: ShapeList, + meta: NodeMetadata, + ) -> Iterable[ProxyValue]: + for element in shape_arg: + if isinstance(element, (list, tuple)): + yield from self._iter_materialized_shape_elements(element, meta) + else: + yield self._ensure_value(element, meta) + + def _register_proxyval(self, key: str, proxyval: ProxyValue) -> None: + logger.debug( + "Registering shape proxy key=%s node=%s data=%s", + key, + proxyval.node.name, + proxyval.data, + ) + self._shape_to_proxyval[key] = proxyval + + def _materialize_int(self, value: int, meta: NodeMetadata) -> ProxyValue: + maybe_proxy = self._shape_to_proxyval.get(str(value), None) + if maybe_proxy is not None: + logger.debug("Reusing CONST_SHAPE for integer value=%s", value) + return maybe_proxy + logger.debug("Creating CONST_SHAPE for integer value=%s", value) + proxy_value = self.builder.call_shape_operator( + exir_ops.backend.tosa.CONST_SHAPE.default, + ([value],), + {}, + meta, + True, + ) + self._register_proxyval(str(value), proxy_value) + return proxy_value + + def materialize_shape_op(self, target, args: Tuple, kwargs, meta) -> ProxyValue: + output_shape = meta.data["val"] + logger.debug( + "Materializing shape op target=%s output_shape=%s args=%s kwargs=%s", + target, + output_shape, + args, + kwargs, + ) + maybe_output_proxy = self._shape_to_proxyval.get(str(output_shape), None) + if maybe_output_proxy is not None: + logger.debug( + "Reusing cached shape op target=%s output_shape=%s node=%s", + target, + output_shape, + maybe_output_proxy.node.name, + ) + return maybe_output_proxy + if target == exir_ops.backend.tosa.DIM.default: + args = (args[0],) + else: + args = tuple([self.materialize_arglist([arg], meta) for arg in args]) + + logger.debug( + "Calling shape op target=%s with materialized args=%s", target, args + ) + shape_meta = copy.copy(meta) + shape_meta.data = dict(meta.data) + if not isinstance(output_shape, (list, tuple)): + shape_meta.data["val"] = [output_shape] + proxy = self.builder.call_shape_operator( + target, + args, + kwargs, + shape_meta, + True, + ) + self._register_proxyval(str(output_shape), proxy) + return proxy diff --git a/backends/arm/_passes/symbolic_to_tosa_shape_pass.py b/backends/arm/_passes/symbolic_to_tosa_shape_pass.py index 3efc6cc260f..5c8df29a607 100644 --- a/backends/arm/_passes/symbolic_to_tosa_shape_pass.py +++ b/backends/arm/_passes/symbolic_to_tosa_shape_pass.py @@ -3,25 +3,238 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging +import operator -from typing import Optional +from typing import Any, cast, Optional import torch -from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.backends.arm._passes.arm_pass import ArmPass + +from executorch.backends.arm._passes.symbolic_materialization_helper import ( + SymbolMaterializationHelpers, +) +from executorch.backends.arm._passes.symbolic_shape_utils import materialize_symints +from executorch.backends.arm.tosa.dialect.shape import meta_has_shape_mark from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import NodeMetadata, ProxyValue + +logger = logging.getLogger(__name__) + + +_SYMBOLIC_SHAPE_OPS: dict[Any, Any] = { + operator.add: exir_ops.backend.tosa.ADD_SHAPE.default, + operator.sub: exir_ops.backend.tosa.SUB_SHAPE.default, + operator.mul: exir_ops.backend.tosa.MUL_SHAPE.default, + operator.mod: exir_ops.backend.tosa.MOD_SHAPE.default, + operator.floordiv: exir_ops.backend.tosa.DIV_FLOOR_SHAPE.default, +} + +class SymbolicToTosaShapesPass(ArmPass): + """Lower PyTorch symbolic shape expressions to TOSA shape operands. -class SymbolicToTosaShapesPass(ArmOpTargetedPass): + This pass owns the when of symbolic shape lowering. It maps + `aten.sym_size.int` to `DIM`, maps supported symbolic arithmetic to TOSA + shape arithmetic, and finds operator arguments that contain + shape-producing proxies. When a shape value must be materialized, it + delegates the how to `SymbolMaterializationHelpers`, which builds and + caches `CONST_SHAPE`, `CONCAT_SHAPE`, and other TOSA shape ops. + + """ _passes_required_after = set() - target_ops = {torch.ops.aten.sym_size.int} + + def __init__(self): + super().__init__() + self.materializer = SymbolMaterializationHelpers(self) + + def _is_shape_proxy(self, arg): + return isinstance(arg, ProxyValue) and meta_has_shape_mark( + getattr(arg.node, "meta", {}) + ) + + def _has_shape_node_arg(self, arg): + if isinstance(arg, torch.fx.Node): + return meta_has_shape_mark(arg.meta) + if isinstance(arg, (list, tuple)): + return any(self._has_shape_node_arg(a) for a in arg) + return False + + def _has_shape_proxy_arg(self, arg): + if self._is_shape_proxy(arg): + return True + if isinstance(arg, (list, tuple)): + return any(self._has_shape_proxy_arg(a) for a in arg) + return False + + def _has_raw_symint_arg(self, arg): + if isinstance(arg, torch.SymInt): + return True + if isinstance(arg, (list, tuple)): + return any(self._has_raw_symint_arg(a) for a in arg) + return False + + def _proxy_value_from_node(self, node: torch.fx.Node) -> ProxyValue: + return ProxyValue(node.meta["val"], self.tracer.proxy(node)) + + def _proxy_value_from_arg(self, arg) -> ProxyValue | int: + if isinstance(arg, torch.fx.Node): + return self._materialize_shape_node(arg) + return arg + + def _meta_value_from_arg(self, arg): + if not isinstance(arg, torch.fx.Node): + return arg + self._ensure_shape_node_meta(arg) + value = arg.meta["val"] + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + def _ensure_shape_node_meta(self, node: torch.fx.Node) -> None: + if "val" in node.meta: + return + target = node.target + if not callable(target): + return + node.meta["val"] = target( + *(self._meta_value_from_arg(arg) for arg in node.args), + **{ + key: self._meta_value_from_arg(value) + for key, value in node.kwargs.items() + }, + ) + + def _materialize_shape_node(self, node: torch.fx.Node) -> ProxyValue: + if meta_has_shape_mark(node.meta): + return self._proxy_value_from_node(node) + if node.target == torch.ops.aten.sym_size.int: + tensor_node = cast(torch.fx.Node, node.args[0]) + tensor = self._proxy_value_from_node(tensor_node) + return self.materializer.materialize_shape_op( + exir_ops.backend.tosa.DIM.default, + (tensor,), + {"axis": node.args[1]}, + NodeMetadata(node.meta), + ) + if node.target in _SYMBOLIC_SHAPE_OPS: + self._ensure_shape_node_meta(node) + return self.materializer.materialize_shape_op( + _SYMBOLIC_SHAPE_OPS[node.target], + tuple(self._proxy_value_from_arg(arg) for arg in node.args), + {}, + NodeMetadata(node.meta), + ) + return self._proxy_value_from_node(node) + + def _erase_temporary_shape_expression( + self, node: torch.fx.Node, original_nodes: set[torch.fx.Node] + ) -> None: + input_nodes = list(node.all_input_nodes) + if node not in original_nodes and not node.users: + self.tracer.graph.erase_node(node) + for input_node in reversed(input_nodes): + self._erase_temporary_shape_expression(input_node, original_nodes) + + def _materialize_raw_symints(self, arg): + if isinstance(arg, torch.SymInt): + original_nodes = set(self.tracer.graph.nodes) + materialized = materialize_symints(self.tracer.graph, [arg])[0] + if isinstance(materialized, torch.fx.Node): + proxy_value = self._materialize_shape_node(materialized) + self._erase_temporary_shape_expression(materialized, original_nodes) + return proxy_value + return materialized + if isinstance(arg, list): + return [self._materialize_raw_symints(a) for a in arg] + if isinstance(arg, tuple): + return tuple(self._materialize_raw_symints(a) for a in arg) + return arg + + def should_run_pass(self, graph_module): + visited_graph_modules = set() + + def graph_needs_shape_materialization(module): + if id(module) in visited_graph_modules: + return False + visited_graph_modules.add(id(module)) + + for node in module.graph.nodes: + if node.op != "call_function": + continue + if node.target == torch.ops.aten.sym_size.int: + return True + if meta_has_shape_mark(node.meta): + continue + if any( + self._has_shape_node_arg(arg) or self._has_raw_symint_arg(arg) + for arg in node.args + ): + return True + + return any( + isinstance(child, torch.fx.GraphModule) + and graph_needs_shape_materialization(child) + for child in module.children() + ) + + return graph_needs_shape_materialization(graph_module) def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False): if op == torch.ops.aten.sym_size.int: - return super().call_shape_operator( + logger.debug("Materializing sym_size.int as TOSA DIM axis=%s", args[1]) + return self.materializer.materialize_shape_op( exir_ops.backend.tosa.DIM.default, (args[0],), {"axis": args[1]}, meta, ) - return super().call_operator(op, args, kwargs, meta, updated) + + if meta_has_shape_mark(meta.data): + logger.debug("Forwarding already shape-marked op=%s", op) + return super().call_operator(op, args, kwargs, meta, updated) + new_args: list[Any] = [] + for arg in args: + if isinstance(arg, (list, tuple)) and len(arg) > 0: + if self._has_raw_symint_arg(arg): + logger.debug( + "Materializing raw SymInt entries for op=%s shape arg: %s", + op, + arg, + ) + arg = self._materialize_raw_symints(arg) + if self._has_shape_proxy_arg(arg): + logger.debug( + "Materializing list arg for op=%s as TOSA shape arg: %s", + op, + arg, + ) + shape_op_arg = self.materializer.materialize_arglist(arg, meta) + new_args.append(shape_op_arg) + else: + new_args.append(arg) + else: + new_args.append(arg) + args = tuple(new_args) + logger.debug("Calling rewritten op=%s args=%s", op, args) + + return super().call_operator(op, args, kwargs, meta) + + def call_sym(self, target, args, meta): + has_shape_arg = any(self._has_shape_proxy_arg(arg) for arg in args) + if target in _SYMBOLIC_SHAPE_OPS and has_shape_arg: + logger.debug( + "Materializing symbolic op target=%s as shape op=%s args=%s", + target, + _SYMBOLIC_SHAPE_OPS[target], + args, + ) + return self.materializer.materialize_shape_op( + _SYMBOLIC_SHAPE_OPS[target], args, {}, meta + ) + if has_shape_arg: + raise NotImplementedError( + f"Symbolic op target {target} not supported in symbolic to TOSA shape pass" + ) + return super().call_sym(target, args, meta) diff --git a/backends/arm/test/misc/test_symbolic_materialization_helper.py b/backends/arm/test/misc/test_symbolic_materialization_helper.py new file mode 100644 index 00000000000..1ef38bc656e --- /dev/null +++ b/backends/arm/test/misc/test_symbolic_materialization_helper.py @@ -0,0 +1,234 @@ +# 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. + +import copy +from typing import cast + +import executorch.backends.arm.tosa.dialect # noqa: F401 + +import torch +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.arm._passes.symbolic_materialization_helper import ( + SymbolMaterializationHelpers, +) +from executorch.backends.arm.tosa.mapping import TosaSpecialDtype +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.backends.test.graph_builder import GraphBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import NodeMetadata, ProxyValue +from torch.fx import Node + + +class _ShapeGraphBuilder(GraphBuilder): + def __init__(self) -> None: + super().__init__() + self.calls: list[torch.fx.Node] = [] + + def call_shape_operator( + self, + op, + args: tuple, + kwargs: dict, + meta: NodeMetadata, + updated: bool = True, + ) -> ProxyValue: + shape_meta = copy.copy(meta) + shape_meta.data = dict(meta.data) + shape_meta.data[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.SHAPE + proxy = self.call_operator(op, args, kwargs, shape_meta) + self.calls.append(proxy.node) + return proxy + + +def _shape_proxy(builder: GraphBuilder) -> ProxyValue: + x = builder.placeholder("x", torch.randn(1, 3)) + return builder.call_operator( + exir_ops.backend.tosa.DIM.default, + (x,), + {"axis": 1}, + NodeMetadata( + { + "val": [3], + TosaSpecialDtype.meta_key(): TosaSpecialDtype.SHAPE, + } + ), + ) + + +def _make_helper() -> tuple[SymbolMaterializationHelpers, _ShapeGraphBuilder]: + builder = _ShapeGraphBuilder() + helper = SymbolMaterializationHelpers(cast(ArmPass, builder)) + return helper, builder + + +def test_materialize_int_emits_and_reuses_const_shape() -> None: + helper, builder = _make_helper() + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + first = helper.materialize_arglist([7], NodeMetadata({})) + second = helper.materialize_arglist([7], NodeMetadata({})) + + assert first.node is second.node + assert first.node.target == exir_ops.backend.tosa.CONST_SHAPE.default + assert first.node.args == ([7],) + assert len(builder.calls) == 1 + + +def test_materialize_single_shape_proxy_reuses_existing_node() -> None: + helper, builder = _make_helper() + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + result = helper.materialize_arglist([proxy], NodeMetadata({})) + + assert result.node is proxy.node + assert builder.calls == [] + + +def test_materialize_arglist_emits_concat_shape_for_mixed_shape_values() -> None: + helper, builder = _make_helper() + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + result = helper.materialize_arglist([proxy, 5], NodeMetadata({})) + + assert result.node.target == exir_ops.backend.tosa.CONCAT_SHAPE.default + assert [node.target for node in builder.calls] == [ + exir_ops.backend.tosa.CONST_SHAPE.default, + exir_ops.backend.tosa.CONCAT_SHAPE.default, + ] + concat_args = result.node.args[0] + assert isinstance(concat_args, list) + assert concat_args[0] is proxy.node + assert builder.calls[0].args == ([5],) + + +def test_materialize_arglist_flattens_nested_shape_lists() -> None: + helper, builder = _make_helper() + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + result = helper.materialize_arglist([[proxy], [2]], NodeMetadata({})) + + assert result.node.target == exir_ops.backend.tosa.CONCAT_SHAPE.default + concat_args = result.node.args[0] + assert isinstance(concat_args, list) + assert concat_args[0] is proxy.node + assert builder.calls[0].target == exir_ops.backend.tosa.CONST_SHAPE.default + assert builder.calls[1].target == exir_ops.backend.tosa.CONCAT_SHAPE.default + + +def test_materialize_arglist_reuses_cached_const_inside_concat() -> None: + helper, builder = _make_helper() + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + cached_const = helper.materialize_arglist([5], NodeMetadata({})) + result = helper.materialize_arglist([proxy, 5], NodeMetadata({})) + + assert result.node.target == exir_ops.backend.tosa.CONCAT_SHAPE.default + concat_args = result.node.args[0] + assert isinstance(concat_args, list) + assert concat_args == [proxy.node, cached_const.node] + assert [node.target for node in builder.calls] == [ + exir_ops.backend.tosa.CONST_SHAPE.default, + exir_ops.backend.tosa.CONCAT_SHAPE.default, + ] + + +def test_materialize_arglist_accepts_tuple_shape_values() -> None: + helper, builder = _make_helper() + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + result = helper.materialize_arglist(((proxy,), (2, 3)), NodeMetadata({})) + + assert result.node.target == exir_ops.backend.tosa.CONCAT_SHAPE.default + concat_args = result.node.args[0] + assert isinstance(concat_args, list) + assert concat_args[0] is proxy.node + assert builder.calls[0].args == ([2],) + assert builder.calls[1].args == ([3],) + assert builder.calls[2].target == exir_ops.backend.tosa.CONCAT_SHAPE.default + + +def test_materialize_arglist_propagates_meta_to_emitted_shape_ops() -> None: + helper, _ = _make_helper() + meta = NodeMetadata({"val": [2, 3], "debug_handle": 123}) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + result = helper.materialize_arglist([2, 3], meta) + + concat_args = result.node.args[0] + assert isinstance(concat_args, list) + const_node = concat_args[0] + assert isinstance(const_node, Node) + assert const_node.meta["debug_handle"] == 123 + assert result.node.meta["debug_handle"] == 123 + + +def test_materialize_shape_op_materializes_non_dim_args() -> None: + helper, builder = _make_helper() + meta = NodeMetadata({"val": [4]}) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + result = helper.materialize_shape_op( + exir_ops.backend.tosa.ADD_SHAPE.default, + (proxy, 1), + {}, + meta, + ) + + assert result.node.target == exir_ops.backend.tosa.ADD_SHAPE.default + assert result.node.args[0] is proxy.node + assert builder.calls[0].target == exir_ops.backend.tosa.CONST_SHAPE.default + assert result.node.args[1] is builder.calls[0] + + +def test_materialize_shape_op_reuses_cached_output_shape() -> None: + helper, builder = _make_helper() + meta = NodeMetadata({"val": [4]}) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + proxy = _shape_proxy(builder) + first = helper.materialize_shape_op( + exir_ops.backend.tosa.ADD_SHAPE.default, + (proxy, 1), + {}, + meta, + ) + second = helper.materialize_shape_op( + exir_ops.backend.tosa.ADD_SHAPE.default, + (proxy, 1), + {}, + meta, + ) + + assert first.node is second.node + assert first.node.target == exir_ops.backend.tosa.ADD_SHAPE.default + assert len(builder.calls) == 2 + + +def test_materialize_dim_does_not_materialize_axis_arg() -> None: + helper, builder = _make_helper() + tensor = builder.placeholder("x", torch.randn(1, 3)) + meta = NodeMetadata({"val": [3]}) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.1+FP+shape")): + result = helper.materialize_shape_op( + exir_ops.backend.tosa.DIM.default, + (tensor,), + {"axis": 1}, + meta, + ) + + assert result.node.target == exir_ops.backend.tosa.DIM.default + assert result.node.args == (tensor.node,) + assert result.node.kwargs == {"axis": 1} + assert len(builder.calls) == 1 diff --git a/backends/arm/test/passes/test_insert_dynamic_padding_pass.py b/backends/arm/test/passes/test_insert_dynamic_padding_pass.py index 72d667ca881..580580f1139 100644 --- a/backends/arm/test/passes/test_insert_dynamic_padding_pass.py +++ b/backends/arm/test/passes/test_insert_dynamic_padding_pass.py @@ -8,6 +8,9 @@ InsertDynamicPaddingPass, ) from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass +from executorch.backends.arm._passes.symbolic_to_tosa_shape_pass import ( + SymbolicToTosaShapesPass, +) from executorch.backends.arm.tosa.specification import ( TosaLoweringContext, TosaSpecification, @@ -24,6 +27,7 @@ def _assert_inserted_padding( target_op, zero_spatial_padding: list[int], expected_full_padding_len: int, + expected_spatial_padding: list, ) -> None: nodes = graph_module.graph.nodes conv_node = next(n for n in nodes if n.target == target_op) @@ -47,37 +51,14 @@ def _assert_inserted_padding( padding_shape_node = padding_node.args[1] assert padding_shape_node.target == exir_ops.backend.tosa.CONCAT_SHAPE.default - n_padding, spatial_padding, c_padding = padding_shape_node.args[0] - assert n_padding.meta["val"] == [0, 0] - assert c_padding.meta["val"] == [0, 0] - pad_list = padding_shape_node.meta["val"] pad_list_vals = [ p.meta["val"] if isinstance(p, torch.fx.Node) else p for p in pad_list ] assert len(pad_list_vals) == expected_full_padding_len assert pad_list_vals[:2] == [0, 0] + assert pad_list_vals[2:-2] == expected_spatial_padding assert pad_list_vals[-2:] == [0, 0] - # For static graphs spatial_padding is a CONST_SHAPE node; for dynamic - # graphs (RewriteConvPass materialized) it is an immutable_list of Nodes/ints. - if hasattr(spatial_padding, "target"): - assert spatial_padding.target == exir_ops.backend.tosa.CONST_SHAPE.default - spatial_padding_value = spatial_padding.meta["val"] - if isinstance(spatial_padding_value, (list, tuple)): - spatial_vals = [ - p.meta["val"] if isinstance(p, torch.fx.Node) else p - for p in spatial_padding_value - ] - assert pad_list_vals[2:-2] == spatial_vals - else: - assert pad_list_vals[2:-2] == spatial_padding_value - else: - # Dynamic case: spatial_padding is the original pad list (possibly Nodes) - spatial_vals = [ - p.meta["val"] if isinstance(p, torch.fx.Node) else p - for p in spatial_padding - ] - assert pad_list_vals[2:-2] == spatial_vals class ConvModule(torch.nn.Module): @@ -133,6 +114,7 @@ def test_insert_dynamic_padding(): ] edge_model = edge_model.transform([InsertDynamicPaddingPass()]) + edge_model = edge_model.transform([SymbolicToTosaShapesPass()]) graph_module = edge_model.exported_program().graph_module conv_node = next( @@ -159,6 +141,7 @@ def test_insert_dynamic_padding(): exir_ops.backend.tosa.CONV2D.default, zero_spatial_padding=[0, 0, 0, 0], expected_full_padding_len=8, + expected_spatial_padding=initial_padding_vals, ) @@ -192,6 +175,7 @@ def test_insert_dynamic_padding_conv3d(): ] edge_model = edge_model.transform([InsertDynamicPaddingPass()]) + edge_model = edge_model.transform([SymbolicToTosaShapesPass()]) graph_module = edge_model.exported_program().graph_module conv_node = next( @@ -217,4 +201,5 @@ def test_insert_dynamic_padding_conv3d(): exir_ops.backend.tosa.CONV3D.default, zero_spatial_padding=[0, 0, 0, 0, 0, 0], expected_full_padding_len=10, + expected_spatial_padding=initial_padding_vals, ) diff --git a/backends/arm/test/passes/test_symbolic_to_tosa_shape_pass.py b/backends/arm/test/passes/test_symbolic_to_tosa_shape_pass.py index 0ea2c235c8c..45e31125d71 100644 --- a/backends/arm/test/passes/test_symbolic_to_tosa_shape_pass.py +++ b/backends/arm/test/passes/test_symbolic_to_tosa_shape_pass.py @@ -3,10 +3,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Any +import operator +from typing import Any, Callable, cast import executorch.backends.arm.tosa.dialect # noqa: F401 +import pytest +import sympy # type: ignore[import-untyped] import torch +from executorch.backends.arm._passes.resolve_view_copy_inferred_dim_pass import ( + ResolveViewCopyInferredDimPass, +) from executorch.backends.arm._passes.symbolic_to_tosa_shape_pass import ( SymbolicToTosaShapesPass, ) @@ -16,8 +22,22 @@ TosaSpecification, ) from executorch.backends.test.graph_builder import GraphBuilder +from executorch.exir import to_edge from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import NodeMetadata, PassResult, ProxyValue +from torch._export.utils import _get_shape_env_from_gm +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.export import Dim, export +from torch.fx.experimental.symbolic_shapes import ShapeEnv + + +class _RecordingMaterializer: + def __init__(self) -> None: + self.calls: list[tuple[Any, tuple[Any, ...], dict[str, Any], NodeMetadata]] = [] + + def materialize_shape_op(self, target, args, kwargs, meta): + self.calls.append((target, args, kwargs, meta)) + return target def _run_symbolic_shape_pass(graph_module: torch.fx.GraphModule) -> PassResult: @@ -27,6 +47,10 @@ def _run_symbolic_shape_pass(graph_module: torch.fx.GraphModule) -> PassResult: return result +def _targets(graph_module: torch.fx.GraphModule) -> list[Any]: + return [node.target for node in graph_module.graph.nodes] + + def _single_node_with_target( graph_module: torch.fx.GraphModule, target: Any, @@ -34,6 +58,38 @@ def _single_node_with_target( return next(node for node in graph_module.graph.nodes if node.target == target) +def _nodes_with_target( + graph_module: torch.fx.GraphModule, + target: Any, +) -> list[torch.fx.Node]: + return [node for node in graph_module.graph.nodes if node.target == target] + + +def _symbolic_binary_op( + builder: GraphBuilder, + op: Callable[..., Any], + args: tuple[Any, Any], + value: int, +) -> ProxyValue: + return builder.call_operator(op, args, meta=NodeMetadata({"val": value})) + + +def _build_view_graph( + shape_builder: Callable[[GraphBuilder, ProxyValue], list[Any]], + output_shape: tuple[int, ...], +) -> torch.fx.GraphModule: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 18)) + shape = shape_builder(builder, x) + view = builder.call_operator( + torch.ops.aten.view.default, + (x, shape), + meta=NodeMetadata({"val": torch.empty(output_shape)}), + ) + builder.output([view]) + return builder.get_graph_module() + + def _sym_size( builder: GraphBuilder, x: ProxyValue, @@ -47,6 +103,19 @@ def _sym_size( ) +def _shape_proxy(builder: GraphBuilder, value: int) -> ProxyValue: + return builder.call_operator( + exir_ops.backend.tosa.CONST_SHAPE.default, + ([value],), + meta=NodeMetadata( + { + "val": [value], + TosaSpecialDtype.meta_key(): TosaSpecialDtype.SHAPE, + } + ), + ) + + def test_symbolic_to_tosa_shapes_rewrites_sym_size_to_dim() -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(2, 18)) @@ -98,4 +167,449 @@ def test_symbolic_to_tosa_shapes_leaves_non_sym_size_ops_unchanged() -> None: assert exir_ops.backend.tosa.DIM.default not in { node.target for node in graph_module.graph.nodes } + + +def test_symbolic_to_tosa_shapes_rewrites_sym_size_and_symbolic_list() -> None: + def shape(builder: GraphBuilder, x: ProxyValue) -> list[Any]: + return [_sym_size(builder, x, 0, 2), 18] + + result = _run_symbolic_shape_pass(_build_view_graph(shape, (2, 18))) + graph_module = result.graph_module + targets = _targets(graph_module) + view_node = _single_node_with_target(graph_module, torch.ops.aten.view.default) + + dim_node = _single_node_with_target(graph_module, exir_ops.backend.tosa.DIM.default) + const_node = _single_node_with_target( + graph_module, exir_ops.backend.tosa.CONST_SHAPE.default + ) + + assert dim_node.kwargs == {"axis": 0} + assert const_node.args == ([18],) + assert exir_ops.backend.tosa.CONCAT_SHAPE.default in targets + assert torch.ops.aten.sym_size.int not in targets + assert ( + getattr(view_node.args[1], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_materializes_raw_symint_list_arg() -> None: + shape_env = ShapeEnv() + width = shape_env.create_symintnode(sympy.Symbol("width"), hint=14) + assert isinstance(width, torch.SymInt) + shape_env.constrain_symbol_range(width.node.expr, compiler_min=1, compiler_max=16) + + with FakeTensorMode(shape_env=shape_env, allow_non_fake_inputs=True): + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty(size=(1, width * 16, 1, 1)) + reshape = graph.call_function( + exir_ops.backend.tosa.RESHAPE.default, + args=(x, [1, width * 16, 1, 1]), + ) + reshape.meta["val"] = torch.empty(size=(1, width * 16, 1, 1)) + graph.output(reshape) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+FP+shape"), shape_env + ): + result = SymbolicToTosaShapesPass()(graph_module) + + assert result is not None + graph_module = result.graph_module + reshape_node = _single_node_with_target( + graph_module, exir_ops.backend.tosa.RESHAPE.default + ) + targets = _targets(graph_module) + + assert result.modified + assert torch.ops.aten.sym_size.int not in targets + assert exir_ops.backend.tosa.DIM.default in targets + assert ( + getattr(reshape_node.args[1], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + assert not any(isinstance(arg, torch.SymInt) for arg in reshape_node.args) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_rewrites_symbolic_expression() -> None: + def shape(builder: GraphBuilder, x: ProxyValue) -> list[Any]: + dim_0 = _sym_size(builder, x, 0, 2) + dim_1 = _sym_size(builder, x, 1, 18) + product = _symbolic_binary_op(builder, operator.mul, (dim_0, dim_1), 36) + return [product] + + result = _run_symbolic_shape_pass(_build_view_graph(shape, (36,))) + graph_module = result.graph_module + targets = _targets(graph_module) + view_node = _single_node_with_target(graph_module, torch.ops.aten.view.default) + + assert targets.count(exir_ops.backend.tosa.DIM.default) == 2 + assert exir_ops.backend.tosa.MUL_SHAPE.default in targets + assert torch.ops.aten.sym_size.int not in targets + assert ( + getattr(view_node.args[1], "target", None) + == exir_ops.backend.tosa.MUL_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_resolves_view_copy_inferred_dim() -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 18)) + sym_size = _sym_size(builder, x, 0, 2) + view = builder.call_operator( + exir_ops.edge.aten.view_copy.default, + (x, [sym_size, 3, -1]), + meta=NodeMetadata({"val": torch.empty(2, 3, 6)}), + ) + builder.output([view]) + + resolve_result = ResolveViewCopyInferredDimPass()(builder.get_graph_module()) + assert resolve_result is not None + assert resolve_result.modified + + result = _run_symbolic_shape_pass(resolve_result.graph_module) + graph_module = result.graph_module + view_node = _single_node_with_target( + graph_module, exir_ops.edge.aten.view_copy.default + ) + const_shape_args = [ + node.args[0] + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.CONST_SHAPE.default + ] + concat_node = _single_node_with_target( + graph_module, exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + + assert [-1] not in const_shape_args + assert [3] in const_shape_args + assert [6] in const_shape_args + assert concat_node.meta["val"] == [2, 3, 6] + assert ( + getattr(view_node.args[1], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + graph_module.graph.lint() + + +def test_resolve_view_copy_inferred_dim_materializes_dynamic_dim() -> None: + class ViewModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten.view_copy.default(x, [x.shape[0], -1]) + + edge_model = to_edge( + export( + ViewModule(), + (torch.randn(2, 3, 4),), + dynamic_shapes={ + "x": { + 0: Dim("batch", min=1, max=5), + 1: Dim("height", min=2, max=6), + 2: Dim("width", min=2, max=8), + } + }, + ) + ) + graph_module = edge_model.exported_program().graph_module + shape_env = _get_shape_env_from_gm(graph_module) + + resolve_result = ResolveViewCopyInferredDimPass()(graph_module) + assert resolve_result is not None + assert resolve_result.modified + + view_node = _single_node_with_target( + resolve_result.graph_module, exir_ops.edge.aten.view_copy.default + ) + resolved_shape = cast(list[Any], view_node.args[1]) + assert -1 not in resolved_shape + assert getattr(resolved_shape[1], "target", None) == operator.mul + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+FP+shape"), shape_env + ): + result = SymbolicToTosaShapesPass()(resolve_result.graph_module) + assert result is not None + + graph_module = result.graph_module + targets = _targets(graph_module) + view_node = _single_node_with_target( + graph_module, exir_ops.edge.aten.view_copy.default + ) + + assert torch.ops.aten.sym_size.int not in targets + assert targets.count(exir_ops.backend.tosa.DIM.default) == 3 + assert exir_ops.backend.tosa.MUL_SHAPE.default in targets + assert ( + getattr(view_node.args[1], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_keeps_view_copy_shape_without_inferred_dim() -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(36)) + sym_size = _sym_size(builder, x, 0, 36) + view = builder.call_operator( + exir_ops.edge.aten.view_copy.default, + (x, [sym_size, 1]), + meta=NodeMetadata({"val": torch.empty(36, 1)}), + ) + builder.output([view]) + + result = _run_symbolic_shape_pass(builder.get_graph_module()) + graph_module = result.graph_module + view_node = _single_node_with_target( + graph_module, exir_ops.edge.aten.view_copy.default + ) + const_shape_nodes = _nodes_with_target( + graph_module, exir_ops.backend.tosa.CONST_SHAPE.default + ) + + assert [node.args for node in const_shape_nodes] == [([1],)] + assert ( + getattr(view_node.args[1], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_runs_for_shape_marked_list_without_sym_size() -> None: + graph = torch.fx.Graph() + shape = graph.call_function(exir_ops.backend.tosa.CONST_SHAPE.default, ([2],)) + shape.meta["val"] = [2] + shape.meta[TosaSpecialDtype.meta_key()] = TosaSpecialDtype.SHAPE + empty = graph.call_function( + torch.ops.aten.empty.memory_format, + ([shape],), + {"device": torch.device("cpu"), "pin_memory": False}, + ) + empty.meta["val"] = torch.empty(2) + graph.output(empty) + + result = _run_symbolic_shape_pass(torch.fx.GraphModule(torch.nn.Module(), graph)) + graph_module = result.graph_module + empty_node = _single_node_with_target( + graph_module, + torch.ops.aten.empty.memory_format, + ) + + assert getattr(empty_node.args[0], "target", None) == ( + exir_ops.backend.tosa.CONST_SHAPE.default + ) + graph_module.graph.lint() + + +@pytest.mark.parametrize( + "symbolic_op,tosa_op", + [ + (operator.add, exir_ops.backend.tosa.ADD_SHAPE.default), + (operator.sub, exir_ops.backend.tosa.SUB_SHAPE.default), + (operator.mul, exir_ops.backend.tosa.MUL_SHAPE.default), + (operator.mod, exir_ops.backend.tosa.MOD_SHAPE.default), + (operator.floordiv, exir_ops.backend.tosa.DIV_FLOOR_SHAPE.default), + ], +) +def test_symbolic_to_tosa_shapes_maps_symbolic_arithmetic_ops( + symbolic_op, + tosa_op, +) -> None: + shape_pass = SymbolicToTosaShapesPass() + materializer = _RecordingMaterializer() + shape_pass.materializer = materializer + builder = GraphBuilder() + lhs = _shape_proxy(builder, 4) + rhs = _shape_proxy(builder, 2) + meta = NodeMetadata({"val": [4]}) + + result = shape_pass.call_sym(symbolic_op, (lhs, rhs), meta) + + assert result == tosa_op + assert materializer.calls == [(tosa_op, (lhs, rhs), {}, meta)] + + +def test_symbolic_to_tosa_shapes_rejects_unsupported_symbolic_op() -> None: + shape_pass = SymbolicToTosaShapesPass() + shape_pass.materializer = _RecordingMaterializer() + builder = GraphBuilder() + lhs = _shape_proxy(builder, 4) + rhs = _shape_proxy(builder, 2) + + with pytest.raises(NotImplementedError, match="Symbolic op target"): + shape_pass.call_sym(operator.truediv, (lhs, rhs), NodeMetadata({"val": [4]})) + + +def test_symbolic_to_tosa_shapes_rewrites_add_expression() -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 18)) + dim_0 = _sym_size(builder, x, 0, 2) + dim_1 = _sym_size(builder, x, 1, 18) + add = _symbolic_binary_op(builder, operator.add, (dim_0, dim_1), 20) + empty = builder.call_operator( + torch.ops.aten.empty.memory_format, + ([add],), + {"device": torch.device("cpu"), "pin_memory": False}, + NodeMetadata({"val": torch.empty(20)}), + ) + builder.output([empty]) + + result = _run_symbolic_shape_pass(builder.get_graph_module()) + graph_module = result.graph_module + targets = _targets(graph_module) + empty_node = _single_node_with_target( + graph_module, + torch.ops.aten.empty.memory_format, + ) + + assert targets.count(exir_ops.backend.tosa.DIM.default) == 2 + assert exir_ops.backend.tosa.ADD_SHAPE.default in targets + assert torch.ops.aten.sym_size.int not in targets + assert ( + getattr(empty_node.args[0], "target", None) + == exir_ops.backend.tosa.ADD_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_materializes_raw_symint_mod_expression() -> None: + shape_env = ShapeEnv() + height = shape_env.create_symintnode(sympy.Symbol("height"), hint=14) + assert isinstance(height, torch.SymInt) + shape_env.constrain_symbol_range(height.node.expr, compiler_min=1, compiler_max=16) + + with FakeTensorMode(shape_env=shape_env, allow_non_fake_inputs=True): + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty(size=(1, height, 1, 1)) + dynamic_dim = (((height - ((height - 1) % 2)) - 1) // 2) + 1 + empty = graph.call_function( + torch.ops.aten.empty.memory_format, + args=([1, dynamic_dim, 1, 1],), + kwargs={"device": torch.device("cpu"), "pin_memory": False}, + ) + empty.meta["val"] = torch.empty(size=(1, dynamic_dim, 1, 1)) + graph.output(empty) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+FP+shape"), shape_env + ): + result = SymbolicToTosaShapesPass()(graph_module) + + assert result is not None + graph_module = result.graph_module + empty_node = _single_node_with_target( + graph_module, torch.ops.aten.empty.memory_format + ) + targets = _targets(graph_module) + + assert exir_ops.backend.tosa.MOD_SHAPE.default in targets + assert exir_ops.backend.tosa.DIV_FLOOR_SHAPE.default in targets + assert ( + getattr(empty_node.args[0], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_concats_static_scalar_proxy() -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 18)) + dim_1 = _sym_size(builder, x, 1, 18) + sub = builder.call_operator( + operator.sub, + (2, 1), + meta=NodeMetadata({"val": 1}), + ) + empty = builder.call_operator( + torch.ops.aten.empty.memory_format, + ([sub, dim_1],), + {"device": torch.device("cpu"), "pin_memory": False}, + NodeMetadata({"val": torch.empty(1, 18)}), + ) + builder.output([empty]) + + result = _run_symbolic_shape_pass(builder.get_graph_module()) + graph_module = result.graph_module + concat_node = _single_node_with_target( + graph_module, exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + empty_node = _single_node_with_target( + graph_module, torch.ops.aten.empty.memory_format + ) + + assert concat_node.meta["val"] == [1, 18] + assert ( + getattr(empty_node.args[0], "target", None) + == exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_concats_scalar_symbolic_expression() -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 18)) + dim_0 = _sym_size(builder, x, 0, 2) + dim_1 = _sym_size(builder, x, 1, 18) + sub = _symbolic_binary_op(builder, operator.sub, (dim_0, 1), 1) + empty = builder.call_operator( + torch.ops.aten.empty.memory_format, + ([sub, dim_1],), + {"device": torch.device("cpu"), "pin_memory": False}, + NodeMetadata({"val": torch.empty(1, 18)}), + ) + builder.output([empty]) + + result = _run_symbolic_shape_pass(builder.get_graph_module()) + graph_module = result.graph_module + sub_node = _single_node_with_target( + graph_module, exir_ops.backend.tosa.SUB_SHAPE.default + ) + concat_node = _single_node_with_target( + graph_module, exir_ops.backend.tosa.CONCAT_SHAPE.default + ) + + assert sub_node.meta["val"] == [1] + assert concat_node.meta["val"] == [1, 18] + graph_module.graph.lint() + + +def test_symbolic_to_tosa_shapes_rewrites_nested_symbolic_expression() -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(8, 3, 2)) + dim_0 = _sym_size(builder, x, 0, 8) + dim_1 = _sym_size(builder, x, 1, 3) + sub = _symbolic_binary_op(builder, operator.sub, (dim_0, 1), 7) + add = _symbolic_binary_op(builder, operator.add, (sub, dim_1), 10) + floordiv = _symbolic_binary_op(builder, operator.floordiv, (add, 2), 5) + empty = builder.call_operator( + torch.ops.aten.empty.memory_format, + ([floordiv],), + {"device": torch.device("cpu"), "pin_memory": False}, + NodeMetadata({"val": torch.empty(5)}), + ) + builder.output([empty]) + + result = _run_symbolic_shape_pass(builder.get_graph_module()) + graph_module = result.graph_module + targets = _targets(graph_module) + empty_node = _single_node_with_target( + graph_module, + torch.ops.aten.empty.memory_format, + ) + + assert targets.count(exir_ops.backend.tosa.DIM.default) == 2 + assert exir_ops.backend.tosa.SUB_SHAPE.default in targets + assert exir_ops.backend.tosa.ADD_SHAPE.default in targets + assert exir_ops.backend.tosa.DIV_FLOOR_SHAPE.default in targets + assert torch.ops.aten.sym_size.int not in targets + assert ( + getattr(empty_node.args[0], "target", None) + == exir_ops.backend.tosa.DIV_FLOOR_SHAPE.default + ) graph_module.graph.lint()