Fold DyT affine maps into following convolutions (#21953)

[cgreenberg]

Summary:

Companion to FoldDyTAlphaIntoLUTPass. FoldDyTAffineIntoConvPass removes the per-channel gamma Mul and beta Add after the DyT TABLE by folding their exact integer-affine map into the following convolution weights and bias.

The pass evaluates the materialized 256-entry TABLE path and rewrites only when the map is provably exact. It fails closed on nonlinear or saturating maps, unsupported layouts or constants, mismatched quantization metadata, and ambiguous layout provenance. Padded convolutions retain beta when its boundary contribution cannot be represented by one bias.

Coverage includes non-uniform gamma/beta, unpadded depthwise convolutions, output-zero-point equality, unused-constant cleanup, and exact-one-permute layout provenance. This revision contains only the public pass implementation, tests, and package export; pass registration is intentionally outside this public change.

Reviewed By: rascani

Differential Revision: D116573000

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 Unclassified Failure

As of commit a66e691 with merge base 1afd07f ():

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

• trunk / test-models-macos-cpu (efficient_sam, portable) / macos-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  ##[error]The operation was canceled.

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: cgreenberg / name: Charles Greenberg (0cec5b5, bef6d1d)

[meta-codesync]

@cgreenberg has exported this pull request. If you are a Meta employee, you can view the originating Diff in D116573000.

[pytorch-bot]

Workflows were awaiting approval. CI has now been triggered for the ciflow labels on this PR.

[cgreenberg]

@pytorchbot label "release notes: arm"

[cgreenberg]

/easycla

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 6m 0s —— View job

---

PR Review: Fold DyT affine maps into following convolutions

• Read PR diff and changed files
• Analyze FoldDyTAffineIntoConvPass logic
• Review pass manager integration
• Review tests
• Post review feedback

I reviewed the affine-fold commit (bef6d1d) and its companion alpha-fold (0cec5b5) that ships in the same branch. This is careful, well-documented work — the fail-closed philosophy is applied consistently and the tests pin the interesting edge cases (channel-narrowing slice, singleton broadcast, wrong-axis views, non-affine maps, padded beta). A few things worth confirming/tightening below.

Correctness

1. The conv accumulator equivalence assumes conv.input_qparams[0].zp == add_output_rescale.output_zp, but this isn't verified.
In _fold_conv_constants (fold_dyt_affine_into_conv_pass.py:589-599) the conv's input zero-point is overwritten with table_qargs zp. The fold is exact only because _fit_integer_affine centers outputs on add_output_rescale.output_zp (:616-617) — i.e. it implicitly assumes the conv originally subtracted that same zp. In a real DyT lowering that invariant holds (both describe the same add_output edge, set by FoldAndAnnotateQParamsPass), so this is not a live bug. But given the pass otherwise fails closed on every unproven assumption (rank, shape, layout, saturation), it's worth an explicit guard: reject the fold if activation_qargs.get_zp_per_tensor() != match.add_output_rescale.output_zp. Otherwise a graph where they diverge would silently shift the accumulator by sum(weight_centered) * (orig_zp - output_zp). Fix this →

2. Scale is intentionally left as the original conv input scale (:591) even though the input edge is now the TABLE, whose scale differs. This is correct at the code level (the integer accumulator and the output-rescale multiplier in_scale * weight_scale are both preserved, so output codes are byte-identical), but it does leave input_qparams[0] semantically describing the wrong scale for the new edge. Worth a one-line comment noting the scale is deliberately preserved for multiplier-equivalence, not because it describes the TABLE — otherwise a future reader may "fix" it and break exactness.

3. Singleton-x affine (len(points) == 1) yields slope = 0 (:385-387). This zeroes the folded weights for that channel. It's sound (if the TABLE only ever emits one code, the conv input is constant and a zero-slope + offset-in-bias fold reproduces it exactly), but it's a subtle case that no test covers. Consider adding a regression for an all-constant TABLE column to pin the behavior.

Robustness / smaller notes

• _fit_integer_affine uses the first two sorted points to derive the slope, then validates all points (:388-399). This is correct and the full-set verification makes it robust to a coincidentally-collinear leading pair. Good.
• The fit runs a Python double loop over 256 codes × channels (:374-401). Fine for the expected channel counts, but if this ever runs on wide layers it'll be slow — a vectorized fit (sort unique x, check constant dy/dx) would avoid the per-element .item() calls. Not blocking.
• _match is a long method carrying # noqa: C901. It's readable as written, but the operand-classification block (gamma vs table, beta vs activation) is duplicated between the mul and add handling; a small shared helper could reduce the surface. Optional.

Tests

Coverage is strong: exact fold, shared-weight distinct folds, identity removal, padded-beta retention, and every fail-closed path has a dedicated test with an explanatory docstring. The pipeline test correctly documents why MatchArgRanksPass is required (materializing the (1,1,1,C) view the matcher keys on). One gap: no test for the singleton-x / zero-slope path noted above.

CI

The two red jobs (Cadence hifi4, Cadence vision) are the pull_request_target fork-checkout refusal, unrelated to this change. Nothing in this PR touches those workflows.

Overall this looks solid and safe-by-construction. My only substantive ask is guard #1 to keep the pass consistent with its own fail-closed contract.
· branch export-D116573000

[cgreenberg]

@claude review this PR

[rascani]

Review automatically exported from Phabricator review in Meta.