Arm backend: Support quantized cond and while

[Erik-Lundell]

Arm backend: Support quantized while_loop

- Add annotation logic.
- Extend cond handling in q-dq folding to while
- Extend InsertCondRescale pass to handle while.

---

Arm backend: Add initial while_loop support.

- Refactor CondSupported to also test while, move to own file
  and split into one check for submodule nodes, and one for ops.
- Add node visitor
- Add tests

---

Arm backend: Initial quantization support for conditional

The standard prepare/convert_pt2 does not seem to support
quantization out of the box. Instead, a quantization call
is introduced in the TOSAQuantizer, that does the neccessary
steps to get correct quantization on submodules. A custom
Quantize step is needed in the ArmTester to make this work
in testing.

Additionally, getting correct quantization parameters needs
some delicate handling. The model is calibrated twice,
once for each code path. Because of this, the observers outside
the if/else submodules see different data than the observers
inside the submodules. Rescales need to be inserted to handle
this. To get a correctly traced graph at all times, we
  1. Fold the outmost quant ops in the submodules at the same
     time as the cond is folded. Add qparam meta to folded
     nodes inside submodule.
  2. Use this meta in the InsertCondRescale pass to
     insert a tosa.RESCALE to handle the different qparams.
  3. After this, the submodule's q-dq nodes can be folded normally.

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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, 3 Unrelated Failures

As of commit 332375f with merge base 9952aef ():

NEW FAILURE - The following job has failed:

• Test CUDA Builds / test-models-cuda (linear) / linux-job (gh)
  RuntimeError: Command docker exec -t 7fd8a7eaa262b30ceb02794f2284d8d6c16164add636a9bf89db62ff3379882d /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_nested_lowered_backend_module
• pull / unittest-editable / linux / linux-job (gh) (trunk failure)
  exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_nested_lowered_backend_module
• trunk / unittest-release / linux / linux-job (gh) (trunk failure)
  exir/backend/test/test_lowered_backend_module.py::TestBackendAPI::test_emit_nested_lowered_backend_module

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

[Erik-Lundell]

@metascroy Hey Scott, this PR is rather large so probably best for you to test internally first. Thanks!

[Copilot]

Pull Request Overview

This PR adds support for quantized while_loop and cond operations to the Arm backend. The implementation extends the TOSA backend to handle control flow constructs with proper quantization handling through specialized passes and annotation logic.

Key changes:

• Introduces while_loop operator support with validation, node visitors, and tests
• Adds quantization infrastructure for control flow operations via InsertControlFlowRescalesPass and updated annotation logic
• Refactors control flow support into dedicated module (control_flow_support.py) for better code organization

Reviewed Changes

Copilot reviewed 22 out of 22 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
backends/arm/tosa/partitioner.py	Renames get_control_flow_submodules to get_cond_while_submodules
backends/arm/tosa/backend.py	Adds containing_graph_module parameter and output node normalization
backends/arm/test/tester/test_pipeline.py	Switches to custom ArmQuantize class
backends/arm/test/tester/quantize.py	New custom quantize stage for Arm-specific quantization with submodule support
backends/arm/test/tester/arm_tester.py	Updates quantize stage to use ArmQuantize
backends/arm/test/ops/test_while.py	New comprehensive test suite for while_loop operations
backends/arm/test/ops/test_cond.py	Adds tests for quantized conditionals and new test cases
backends/arm/scripts/parse_test_names.py	Adds while_loop.default to custom edge ops list
backends/arm/quantizer/quantization_annotator.py	Extends annotation logic to handle cond and while_loop operations
backends/arm/quantizer/arm_quantizer.py	Adds quantize_with_submodules method for proper submodule quantization
backends/arm/process_node.py	Handles placeholder processing for submodule inputs
backends/arm/operators/operator_validation_utils.py	Adds reusable control flow extension validation
backends/arm/operators/op_while.py	New visitor for while_loop operator serialization
backends/arm/operators/op_cond_if.py	Refactors to use shared validation utility
backends/arm/operators/init.py	Imports new op_while module
backends/arm/operator_support/tosa_supported_operators.py	Moves control flow support to dedicated module
backends/arm/operator_support/control_flow_support.py	New dedicated module for control flow partitioning logic
backends/arm/operator_support/init.py	Imports new control flow support module
backends/arm/_passes/insert_rescales_pass.py	Adds InsertControlFlowRescalesPass for quantization parameter handling
backends/arm/_passes/fold_qdq_with_annotated_qparams_pass.py	Extends to handle control flow node quantization folding
backends/arm/_passes/arm_pass_manager.py	Integrates new control flow rescale pass
backends/arm/_passes/init.py	Exports new control flow rescale pass

Comments suppressed due to low confidence (1)

backends/arm/process_node.py:249

• Unnecessary 'pass' statement.

        pass

---

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

[Copilot]

Pull Request Overview

Copilot reviewed 22 out of 22 changed files in this pull request and generated 11 comments.

---

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

[metascroy]

@metascroy Hey Scott, this PR is rather large so probably best for you to test internally first. Thanks!

Importing now

[metascroy]

I need to wait for the previous Arm PRs to commit to the internal code base before importing this one due to merge conflicts. They are landing right now (#15828, #15869).

[Erik-Lundell]

I need to wait for the previous Arm PRs to commit to the internal code base before importing this one due to merge conflicts. They are landing right now (#15828, #15869).

Strange, I rebased now and did not get any merge conflicts. Anyways, just pushed a rebased version + fixes from copilot's review.

[metascroy]

I need to wait for the previous Arm PRs to commit to the internal code base before importing this one due to merge conflicts. They are landing right now (#15828, #15869).

Strange, I rebased now and did not get any merge conflicts. Anyways, just pushed a rebased version + fixes from copilot's review.

It's not the OSS code base that had a merge conflict, but the internal codebase (which can be 10-20 commits behind main in the OSS one due to the import process). Anyway, I was able to import this morning and tests are running.

For my own reference: D87349540

[metascroy]

The internal tests pass on this, so I think we're safe to merge

[meta-codesync]

@metascroy has imported this pull request. If you are a Meta employee, you can view this in D87349540.

[zingo]

Fails are unrelated to Arm changes in this PR