Arm backend: Add transpose propagation pass by AdrianLundell · Pull Request #20625 · pytorch/executorch

AdrianLundell · 2026-06-30T15:01:36Z

Adds abstract pass PropagateViewCopyPermutePass and two implementations PropagateViewCopyPermuteUpPass and PropagateViewCopyPermutePass

Adds a helper function to ViewMap to update the view-shape after a view-op swap + additional validity checks.

Adds new InsertDataLayoutCastsPass under the propagation passes to not move permutes/views outside of their intended dtype cast

The move exposed an issue where tosa ops
require the same dtype on both inputs which is not guaranteed at that stage in the pipeline. This is fixed by adding CastInt64BuffersToInt32Pass before TOSA dialect rewrite.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

Adds abstract pass PropagateViewCopyPermutePass and two implementations PropagateViewCopyPermuteUpPass and PropagateViewCopyPermutePass Adds a helper function to ViewMap to update the view-shape after a view-op swap + additional validity checks. Adds new InsertDataLayoutCastsPass under the propagation passes to not move permutes/views outside of their intended dtype cast The move exposed an issue where tosa ops require the same dtype on both inputs which is not guaranteed at that stage in the pipeline. This is fixed by adding CastInt64BuffersToInt32Pass before TOSA dialect rewrite. Signed-off-by: Adrian Lundell <adrian.lundell@arm.com> Change-Id: I31bc29540cebce9eb0ee3d5a3b2c016c19b15809

pytorch-bot · 2026-06-30T15:01:40Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20625

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Unrelated Failure

As of commit daabdf5 with merge base 0d8bb3d ():

NEW FAILURE - The following job has failed:

pull / test-mcu-cortex-m-backend / linux-job (gh)
RuntimeError: Command docker exec -t 45e2c55ec751f73929deba9b636f6cc6879234d3666f1912459b9ce2b0b3c2ac /exec failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Copilot

Pull request overview

This PR extends the Arm backend TOSA lowering pipeline with a new permute/view propagation pass framework to reduce unnecessary transposes, improves view-dimension mapping utilities to support the propagation logic, and adjusts pipeline ordering to avoid dtype-related issues (including adding an earlier int64→int32 cast before TOSA dialect rewrite). It also adds/updates tests to validate the new behavior and updates expected transpose counts for several model cases.

Changes:

Add PropagateViewCopyPermutePass plus Up/Down implementations to propagate and fuse permute_copy / view(_copy) around transparent and arg-updating ops.
Enhance ViewMap dimension mapping and add helpers to remap view target shapes and unit slices after swaps.
Update Arm pass pipeline ordering (including inserting CastInt64BuffersToInt32Pass earlier) and expand InsertDataLayoutCastsPass operator coverage; update tests and transpose-count expectations accordingly.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
backends/arm/_passes/propagate_view_copy_permute_pass.py	New propagation pass framework + Up/Down implementations; integrates canonicalization and permute/view fusion.
backends/arm/_passes/dim_maps.py	Updates `ViewMap.map_dim` logic and adds shape/slice remapping helpers needed by propagation.
backends/arm/_passes/insert_data_layout_casts_pass.py	Extends cast insertion targeting to include additional TOSA data-layout ops (e.g. TRANSPOSE/SLICE/CONCAT, etc.).
backends/arm/_passes/arm_pass_manager.py	Replaces older permute canonicalization/fusion path with new propagation passes; adjusts pass ordering and adds `CastInt64BuffersToInt32Pass` earlier.
backends/arm/_passes/init.py	Exposes new propagation passes from the Arm passes package.
backends/arm/test/passes/test_propagate_permutes_views_pass.py	New comprehensive unit tests for propagation behavior (FX-graph and pipeline-level).
backends/arm/test/passes/test_insert_data_layout_casts_pass.py	Adds slice coverage for InsertDataLayoutCastsPass under bf16 extension profile.
backends/arm/test/passes/test_dim_maps.py	Updates/extends ViewMap mapping expectations and adds unit-slice remap tests.
backends/arm/test/misc/test_transpose_counts.py	Updates expected transpose counts after new propagation behavior reduces transposes.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+def test_propagate_permute_down_through_transparent_ops_tosa_FP() -> None:
+    def predicate(targets: list[object]) -> None:
+        assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(NEG)
+
+    pipeline = PassPipeline[input_t](
+        DownwardPermute(),
+        DownwardPermute.data,
+        quantize=False,
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1,
+        },
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1,
+        },
+        pass_list=[PropagateViewCopyPermuteUpPass],
+        pass_functions=[_assert_call_targets(predicate)],
+    )
+    pipeline.run()
+


+def test_propagate_view_down_through_transparent_ops_tosa_FP() -> None:
+    def predicate(targets: list[object]) -> None:
+        assert targets.index(VIEW) < targets.index(RELU) < targets.index(NEG)
+
+    pipeline = PassPipeline[input_t](
+        DownwardView(),
+        DownwardView.data,
+        quantize=False,
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
+        },
+        ops_after_pass={
+            "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
+        },
+        pass_list=[PropagateViewCopyPermuteUpPass],
+        pass_functions=[_assert_call_targets(predicate)],
+    )
+    pipeline.run()
+


+class StopAtSharedTransformInput(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = x.permute(0, 2, 3, 1)
+        return (y * y.sigmoid()).permute(0, 3, 1, 2)
+
+    data = (torch.randn(1, 2, 3, 4),)
+


+            if len(next_nodes) == 0:
+                assert node.op in (
+                    "placeholder",
+                    "output",
+                ), f"{self.__class__.__name__} reached an endpoint node which is not a placeholder or output: {frontier}"
+                break


+        assert previous_frontier is not None
+        self._move_node(node, frontier, previous_frontier)
+        refresh_permute_view_meta(node)
+        return True


@AdrianLundell

Reverts #20625 FYI @AdrianLundell - This broke some internal tests, I'll own getting it re-merged. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell

Copilot AI review requested due to automatic review settings June 30, 2026 15:01

AdrianLundell requested a review from digantdesai as a code owner June 30, 2026 15:01

AdrianLundell added partner: arm For backend delegation, kernels, demo, etc. from the 3rd-party partner, Arm ciflow/trunk release notes: none Do not include this in the release notes labels Jun 30, 2026

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 30, 2026

Copilot started reviewing on behalf of AdrianLundell June 30, 2026 15:02 View session

github-actions Bot added the module: arm Issues related to arm backend label Jun 30, 2026

Copilot AI reviewed Jun 30, 2026

View reviewed changes

Merge branch 'main' into change-1265162

daabdf5

zingo approved these changes Jul 1, 2026

View reviewed changes

AdrianLundell merged commit 2f02f75 into pytorch:main Jul 2, 2026
486 of 489 checks passed

rascani mentioned this pull request Jul 6, 2026

Revert "Arm backend: Add transpose propagation pass" #20742

Merged

rascani mentioned this pull request Jul 6, 2026

Reland "Arm backend: Add transpose propagation pass" #20748

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Arm backend: Add transpose propagation pass#20625

Arm backend: Add transpose propagation pass#20625
AdrianLundell merged 2 commits into
pytorch:mainfrom
AdrianLundell:change-1265162

AdrianLundell commented Jun 30, 2026 •

edited by pytorch-bot Bot

Loading

Uh oh!

pytorch-bot Bot commented Jun 30, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

AdrianLundell commented Jun 30, 2026 • edited by pytorch-bot Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20625

❌ 1 New Failure, 1 Unrelated Failure

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

AdrianLundell commented Jun 30, 2026 •

edited by pytorch-bot Bot

Loading

pytorch-bot Bot commented Jun 30, 2026 •

edited

Loading