Skip to content

Commit c4dbf62

Browse files
mcremon-metameta-codesync[bot]
mcremon-meta
authored and
meta-codesync[bot]
committed
Replace tosa_dim_order with explicit NCHW↔NHWC permutes
Summary: Replace implicit `tosa_dim_order`-based layout handling with explicit `permute_copy` ops around TOSA operators that require NHWC layout. ### Rewrite passes insert explicit NCHW↔NHWC permutes `RewriteConvPass`, `RewriteAvgPool2dPass`, and `RewriteMaxPool2dPass` now insert `aten.permute_copy` nodes (NCHW→NHWC before the TOSA op, NHWC→NCHW after) instead of relying on `ToTosaMemoryFormatPass` for layout conversion. This makes layout transitions visible in the graph. ### Grouped conv decomposition in NHWC `RewriteConvPass` decomposes grouped convolutions (non-depthwise) into per-group `TOSA.CONV2D` ops operating entirely in NHWC, with a single input/output permute pair wrapping the whole group. Supports INT8, INT16 (with and without bias) quantisation paths, including the full INT16+bias chain: CONV2D → RESCALE(INT48→INT32) → ADD(bias) → RESCALE(INT32→INT16). ### `ToTosaMemoryFormatPass` scoped down Now only assigns non-identity dim_order to parameter/buffer placeholders (for weight serialisation) and graph I/O. Inserts `permute_copy` instead of `tosa.TRANSPOSE`. Skips users that already carry a matching permute (inserted by the rewrite passes). ### TOSA dialect op metas expect NHWC All TOSA op meta functions (`CONV2D`, `CONV3D`, `DEPTHWISE_CONV2D`, `AVG_POOL2D`, `MAX_POOL2D`, `TRANSPOSE_CONV2D`) now assume NHWC input layout and produce NHWC output shapes. ### Removed `tosa_dim_order` shape remapping `tosa_shape()` no longer reorders dimensions—just resolves symints. `_get_matching_fake_tensor()` returns `node.meta["val"]` directly. Serialisation mapping always uses identity dim_order. ### Operator serialisation simplified `op_amax`, `op_amin`, `op_any`, `op_cat`, `op_sum`, and `op_permute` no longer remap reduction/concat axes through `dim_order` since tensors are already in the layout expected by TOSA. ### Permute optimisation passes added Six shared passes from `executorch/backends/transforms/` are now run after TOSA lowering to fuse, cancel, and simplify the permutes introduced above: - `RemovePermutesAroundElementwiseOps` (extended for `RESCALE`) - `FuseTransposeOrPermuteOpPairsPass` (extended for `RESCALE`) - `ReplaceNopTransposeOrPermuteWithViewPass` - `PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView` - `FuseCascadedTransposeOrPermuteOps` - `FuseCascadedViewOps` ### Removed passes `DecomposeConvWithInt16ActivationPass` and `DecomposeGroupedConvPass` are removed—their logic is now handled inline by `RewriteConvPass`. `RewriteSlicePass` is repositioned after the permute optimisations. ### Ethos-U55 partitioner simplified The dual NCHW/NHWC permute constraint check is removed since tensors are always in the expected layout at partition time. Differential Revision: D100712787
1 parent 6e540da commit c4dbf62

26 files changed

Lines changed: 965 additions & 308 deletions

‎backends/arm/_passes/TARGETS‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@ runtime.python_library(
2525
"//executorch/backends/arm:common",
2626
"//executorch/backends/arm/tosa:utils",
2727
"//executorch/backends/arm/tosa/dialect:lib",
28+
"//executorch/backends/transforms:fuse_cascaded_transpose_or_permute_ops",
29+
"//executorch/backends/transforms:fuse_cascaded_view_ops",
30+
"//executorch/backends/transforms:fuse_transpose_or_permute_op_pairs_pass",
31+
"//executorch/backends/transforms:remove_permutes_around_elementwise_ops",
32+
"//executorch/backends/transforms:postpone_permute_below_squeeze_view",
33+
"//executorch/backends/transforms:replace_nop_transpose_or_permute_with_view",
2834
"//executorch/exir:lib",
2935
],
3036
)

‎backends/arm/_passes/arm_pass_manager.py‎

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@
4444
DecomposeAtanPass,
4545
DecomposeAvgPool2dPass,
4646
DecomposeBatchNormNoStatsPass,
47-
DecomposeConvWithInt16ActivationPass,
4847
DecomposeCoshPass,
4948
DecomposeCosineSimilarityPass,
5049
DecomposeCumsumPass,
@@ -58,7 +57,6 @@
5857
DecomposeFloorDividePass,
5958
DecomposeGeluPass,
6059
DecomposeGluPass,
61-
DecomposeGroupedConvPass,
6260
DecomposeGroupNormPass,
6361
DecomposeGruPass,
6462
DecomposeIndexCopyPass,
@@ -141,7 +139,6 @@
141139
RewriteUpsamplePass,
142140
ScalarsToAttributePass,
143141
SizeAdjustInputPass,
144-
ToTosaMemoryFormatPass,
145142
UnsqueezeBeforeRepeatPass,
146143
UnsqueezeScalarPlaceholdersPass,
147144
)
@@ -455,7 +452,6 @@ def _tosa_pipeline(
455452
DecomposeFloorDividePass(),
456453
DecomposeGeluPass(),
457454
DecomposeAddSubAlphaPass(),
458-
DecomposeGroupedConvPass(),
459455
DecomposeUnfoldToGatherPass(),
460456
DecomposeEmbeddingPass(),
461457
DecomposeIndexSelectToGatherPass(),
@@ -518,21 +514,19 @@ def _tosa_pipeline(
518514
ConvertPermuteSingletonToViewPass(),
519515
RewriteHighRankSingletonPermutePass(),
520516
FuseViewCopyTransformPass(),
521-
DecomposeConvWithInt16ActivationPass(),
522517
DecomposeSumPass(),
523518
InsertTableOpsPass(exported_program),
524519
]
525520
)
526521

527-
# Aten -> TOSA transformation passes
522+
# Aten -> TOSA transformation passes
528523
self.add_passes(
529524
[
530525
RewriteUpsamplePass(),
531526
RewriteMaxPool2dPass(),
532527
RewriteConvPass(exported_program),
533528
RewriteMatmulPass(),
534529
RewritePadPass(),
535-
RewriteSlicePass(),
536530
InsertConstShapesPass(),
537531
]
538532
)
@@ -542,14 +536,61 @@ def _tosa_pipeline(
542536
[
543537
CastInt64BuffersToInt32Pass(exported_program),
544538
FuseEqualPlaceholdersPass(exported_program),
539+
FuseConstantArgsPass(exported_program),
545540
FuseConsecutiveConcatShapesPass(),
546-
ToTosaMemoryFormatPass(exported_program),
547541
RemoveNoopPass(),
548542
InsertRescalePass(),
549543
InsertDataLayoutCastsPass(),
550544
]
551545
)
552546

547+
# Additional optimization passes for permutes
548+
from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import (
549+
FuseCascadedTransposeOrPermuteOps,
550+
)
551+
from executorch.backends.transforms.fuse_cascaded_view_ops import (
552+
FuseCascadedViewOps,
553+
)
554+
from executorch.backends.transforms.fuse_transpose_or_permute_op_pairs_pass import (
555+
FuseTransposeOrPermuteOpPairsPass,
556+
)
557+
from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
558+
RemovePermutesAroundElementwiseOps,
559+
)
560+
from executorch.backends.transforms.postpone_permute_below_squeeze_view import (
561+
PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
562+
)
563+
from executorch.backends.transforms.replace_nop_transpose_or_permute_with_view import (
564+
ReplaceNopTransposeOrPermuteWithViewPass,
565+
)
566+
from executorch.exir.dialects._ops import ops as exir_ops
567+
568+
# Fuse identity permute pairs across RESCALE ops
569+
fuse_pairs = FuseTransposeOrPermuteOpPairsPass()
570+
fuse_pairs.bypass_ops = fuse_pairs.bypass_ops | {
571+
exir_ops.backend.tosa.RESCALE.default,
572+
}
573+
574+
# Remove permutes around elementwise ops including RESCALE
575+
remove_around = RemovePermutesAroundElementwiseOps()
576+
remove_around.permutable_ops = remove_around.permutable_ops | {
577+
exir_ops.backend.tosa.RESCALE.default,
578+
}
579+
580+
self.add_passes(
581+
[
582+
remove_around,
583+
RewriteSlicePass(),
584+
fuse_pairs,
585+
ReplaceNopTransposeOrPermuteWithViewPass(),
586+
PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView(),
587+
FuseCascadedTransposeOrPermuteOps(),
588+
FuseCascadedViewOps(),
589+
AnnotateOutputDimOrderPass(),
590+
InsertConstShapesPass(),
591+
]
592+
)
593+
553594
# Apply all pass insertions once after all passes are collected
554595
self._apply_pass_insertions()
555596

‎backends/arm/_passes/rewrite_avg_pool2d_pass.py‎

Lines changed: 120 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -7,69 +7,139 @@
77

88
import torch
99
from executorch.backends.arm._passes import ArmPass
10+
from executorch.backends.arm._passes.arm_pass_utils import (
11+
create_node,
12+
get_first_fake_tensor,
13+
)
1014
from executorch.backends.arm.operators.operator_validation_utils import (
1115
adjust_pooling_pad_if_needed,
1216
)
1317
from executorch.exir.dialects._ops import ops as exir_ops
14-
from executorch.exir.pass_base import ExportPass
18+
from executorch.exir.pass_base import ExportPass, PassResult
1519

1620
from .fuse_constant_ops_pass import ComputeConstantOpsAOTPass
1721

22+
_NCHW_TO_NHWC = [0, 2, 3, 1]
23+
_NHWC_TO_NCHW = [0, 3, 1, 2]
24+
1825

1926
class RewriteAvgPool2dPass(ArmPass):
20-
"""Rewrite aten.avg_pool2d calls to TOSA AVG_POOL2D op."""
27+
"""Rewrite aten.avg_pool2d calls to TOSA AVG_POOL2D op with NHWC layout."""
2128

22-
# Target the original avg_pool2d operator
2329
targeted_ops = {exir_ops.edge.aten.avg_pool2d.default}
2430
_passes_required_after: Set[Type[ExportPass]] = {
2531
ComputeConstantOpsAOTPass,
2632
}
2733

28-
def call_operator(self, op, args, kwargs, meta, updated=False):
29-
30-
# Only rewrite avg_pool2d
31-
if op not in self.targeted_ops:
32-
return super().call_operator(op, args, kwargs, meta, updated)
33-
34-
x = args[0]
35-
pad_h, pad_w = args[3]
36-
# Make sure pad corresponds to TOSA
37-
pad = [pad_h, pad_w, pad_h, pad_w]
38-
39-
_, _, h, w = x.data.shape
40-
kernel_h, kernel_w = args[1]
41-
stride_h, stride_w = args[2]
42-
43-
ceil_mode = args[4] if len(args) > 4 else False
44-
45-
# Adjust padding if necessary
46-
pad[1] = adjust_pooling_pad_if_needed(h, kernel_h, stride_h, pad[1], ceil_mode)
47-
pad[3] = adjust_pooling_pad_if_needed(w, kernel_w, stride_w, pad[3], ceil_mode)
48-
49-
# Materialize zero-point constants
50-
in_qparams = meta.data.get("input_qparams", {})
51-
in_zp_val = in_qparams[0].get_zp_per_tensor() if 0 in in_qparams else 0
52-
# Materialize input zero-point as a scalar tensor
53-
input_zp = super().call_scalar(in_zp_val, meta)
54-
55-
out_qparams = meta.data.get("output_qparams", {})
56-
out_zp_val = out_qparams[0].get_zp_per_tensor() if 0 in out_qparams else 0
57-
# Materialize output zero-point as a scalar tensor
58-
output_zp = super().call_scalar(out_zp_val, meta)
59-
60-
# Determine accumulator dtype for AVG_POOL2D: INT32 for integer inputs, FP32 otherwise
61-
if x.data.dtype in (torch.int8, torch.int16):
62-
acc_type = torch.int32
63-
else:
64-
acc_type = torch.float32
65-
66-
tosa_args = (args[0], input_zp, output_zp, *args[1:3], pad, acc_type)
67-
68-
# Emit TOSA AVG_POOL2D with normalized args
69-
return super().call_operator(
70-
exir_ops.backend.tosa.AVG_POOL2D.default,
71-
tosa_args,
72-
{},
73-
meta,
74-
True,
34+
@staticmethod
35+
def _insert_permute(graph_module, anchor_node, input_node, perm, before=True):
36+
ctx = (
37+
graph_module.graph.inserting_before(anchor_node)
38+
if before
39+
else graph_module.graph.inserting_after(anchor_node)
7540
)
41+
with ctx:
42+
return create_node(
43+
graph=graph_module.graph,
44+
op_target=exir_ops.edge.aten.permute_copy.default,
45+
args=(input_node, perm),
46+
from_node=input_node,
47+
)
48+
49+
def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
50+
modified = False
51+
52+
for node in list(graph_module.graph.nodes):
53+
if node.op != "call_function" or node.target not in self.targeted_ops:
54+
continue
55+
56+
modified = True
57+
x = node.args[0]
58+
59+
pad_h, pad_w = node.args[3]
60+
pad = [pad_h, pad_w, pad_h, pad_w]
61+
62+
input_fake = get_first_fake_tensor(x)
63+
_, _, h, w = input_fake.shape
64+
kernel_h, kernel_w = node.args[1]
65+
stride_h, stride_w = node.args[2]
66+
67+
ceil_mode = node.args[4] if len(node.args) > 4 else False
68+
69+
pad[1] = adjust_pooling_pad_if_needed(h, kernel_h, stride_h, pad[1], ceil_mode)
70+
pad[3] = adjust_pooling_pad_if_needed(w, kernel_w, stride_w, pad[3], ceil_mode)
71+
72+
# Determine zero-points and accumulator type
73+
in_qparams = node.meta.get("input_qparams", {})
74+
in_zp_val = in_qparams[0].get_zp_per_tensor() if 0 in in_qparams else 0
75+
76+
out_qparams = node.meta.get("output_qparams", {})
77+
out_zp_val = out_qparams[0].get_zp_per_tensor() if 0 in out_qparams else 0
78+
79+
if input_fake.dtype in (torch.int8, torch.int16):
80+
acc_type = torch.int32
81+
else:
82+
acc_type = torch.float32
83+
84+
# Insert NCHW → NHWC permute on input
85+
x_permuted = self._insert_permute(
86+
graph_module, node, x, _NCHW_TO_NHWC, before=True
87+
)
88+
89+
# Materialize zp scalars as graph constants using aten.full with
90+
# explicit dtype matching the input tensor. This ensures the
91+
# pre-computed buffer placeholders carry the correct type for
92+
# INT-only TOSA profiles (avoids defaulting to float32).
93+
zp_kwargs = {"dtype": input_fake.dtype, "device": input_fake.device}
94+
with graph_module.graph.inserting_before(node):
95+
input_zp_node = create_node(
96+
graph=graph_module.graph,
97+
op_target=exir_ops.edge.aten.full.default,
98+
args=((1,), in_zp_val),
99+
kwargs=zp_kwargs,
100+
from_node=node,
101+
)
102+
output_zp_node = create_node(
103+
graph=graph_module.graph,
104+
op_target=exir_ops.edge.aten.full.default,
105+
args=((1,), out_zp_val),
106+
kwargs=zp_kwargs,
107+
from_node=node,
108+
)
109+
110+
kernel = list(node.args[1])
111+
stride = list(node.args[2])
112+
113+
tosa_args = (x_permuted, input_zp_node, output_zp_node, kernel, stride, pad, acc_type)
114+
115+
# Create TOSA AVG_POOL2D node
116+
with graph_module.graph.inserting_after(node):
117+
tosa_op = create_node(
118+
graph=graph_module.graph,
119+
op_target=exir_ops.backend.tosa.AVG_POOL2D.default,
120+
args=tosa_args,
121+
from_node=node,
122+
inherit_qparams=True,
123+
)
124+
125+
# Compute correct NHWC FakeTensor
126+
input_fake_nhwc = input_fake.permute(_NCHW_TO_NHWC)
127+
input_zp_fake = torch.tensor(in_zp_val, dtype=input_fake.dtype)
128+
output_zp_fake = torch.tensor(out_zp_val, dtype=input_fake.dtype)
129+
tosa_node_fake = exir_ops.backend.tosa.AVG_POOL2D.default(
130+
input_fake_nhwc, input_zp_fake, output_zp_fake, kernel, stride, pad, acc_type
131+
)
132+
tosa_op.meta["val"] = tosa_node_fake
133+
134+
# Insert NHWC → NCHW permute on output
135+
output_permute = self._insert_permute(
136+
graph_module, tosa_op, tosa_op, _NHWC_TO_NCHW, before=False
137+
)
138+
139+
node.replace_all_uses_with(output_permute)
140+
graph_module.graph.erase_node(node)
141+
142+
if modified:
143+
graph_module.recompile()
144+
graph_module = super().call(graph_module).graph_module
145+
return PassResult(graph_module, modified)

0 commit comments

Comments
 (0)