[JakeStevens]

Summary

On the *ReverseArguments paths, handle_broadcast_elementwise hands the broadcasting maps (b, a), so every non-commutative caller had to pass a hand-swapped lambda on those paths. op_div and op_add_sub_impl did, op_le did not (#24485). This swaps vec_fun's arguments once inside handle_broadcast_elementwise, so the lambda always receives (a, b), and removes the per-op swaps, resolving the TODO in op_add_sub_impl.h.

op_mul is commutative, so it passes kCommutative = true to skip compiling the swapped copy, which would otherwise add about 16 KB to op_mul. With it, op_mul is unchanged and op_div/op_add/op_sub/op_le shrink slightly.

Stacked on #24486, whose tests cover le on all three reverse paths; land after it.

Authored with AI assistance (Claude Code).

Test plan

optimized_kernels_test passes (1839 tests, unchanged). With the central swap disabled, 9 add/sub/div/le tests fail, so the existing tests cover the swap for every non-commutative caller.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit aa54c1f with merge base 3f08f3f ():
💚 Looks good so far! There are no failures yet. 💚

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[JakeStevens]

@claude review effort:max

[claude]

Claude finished @JakeStevens's task in 2m 32s —— View job

---

I found no correctness bugs: on every broadcast path the lambda now gets (a, b), and the other points below are minor (a quiet change to a public header, an unused alpha parameter and a stale comment).

Detailed review

PR Review: #24809

Summary

This PR moves the reverse-broadcast operand swap into handle_broadcast_elementwise and removes the hand-swapped lambdas from op_div and op_add_sub_impl. On the three *ReverseArguments paths, the new check in the wrapper matches exactly the paths where handle_last_dim_broadcast_elementwise and plan_broadcast_elementwise set lhs=b and rhs=a. Every in-tree caller (op_add_sub_impl.h, op_div.cpp, op_le.cpp, op_mul.cpp) now behaves the same as before, or is fixed in the case of le. What's left is a minor compatibility point and some cleanup. This comes from reading the code only; I didn't build or run the tests.

Backward Compatibility

• Out-of-tree callers would silently get swapped twice (kernels/optimized/cpu/binary_ops.h:264-287).
  • Why it matters: binary_ops.h is exported with visibility = ["PUBLIC"] in kernels/optimized/cpu/targets.bzl, and the name and signature of handle_broadcast_elementwise haven't changed. Code outside the repo may have copied the old op_div pattern, passing y / x on the reverse paths.
  • What happens: that code still compiles, but now returns b op a on the reverse paths, with no error or warning.
  • A second, louder break: callers that write the template arguments out in full, handle_broadcast_elementwise<CTYPE, Op>(...), will fail to compile, because Op now lands in the bool kCommutative slot.
  • Options: rename the function so stale callers fail to build, or at least call out the change in meaning in the PR, for example with a release notes: label. Whether any internal callers exist could only be checked against fbcode.

API Design

• Two public helpers now order arguments differently (kernels/optimized/cpu/binary_ops.h:175-209).
  • handle_last_dim_broadcast_elementwise is still public in torch::executor. On kBroadcastLastDimReverseArguments it still calls vec_fun(b, a).
  • handle_broadcast_elementwise in the same header now always passes (a, b).
  • Its only caller is internal::handle_broadcast_elementwise_impl, so moving it into internal too would keep the header to a single contract.

Code Quality

• alpha is accepted but never used (kernels/optimized/cpu/binary_ops.h:272).
  • This predates the PR, but the PR rewrites this exact signature and already dropped alpha from _impl.
  • Removing it from the wrapper means op_add_sub_impl.h would also stop passing it.
• Stale comment (kernels/optimized/cpu/op_div.cpp:137-138). The comment "Reason for using alpha is becasuse handle_broadcast_elementwise is used for add and sub as well" sits right above the lines this PR rewrites. op_div doesn't pass alpha, so the comment describes nothing and should be deleted.

Testing

• Reverse paths with alpha ≠ 1 are only partly tested.
  • Calls on the reverse paths with alpha ≠ 1 appear only in the 3D/4D broadcast tests (test_broadcast_3D/test_broadcast_4D in op_sub_test.cpp; op_add_test.cpp:174).
  • The 2d-by-1d and last-dim reverse tests all use alpha = 1, so they can't tell a - alpha*b from alpha*a - b.
  • The swap no longer depends on the path, so this is low risk. One sub case on a 2d-by-1d reverse path with alpha = 2 would cover the gap.

Recommendation

Approve: no blockers. Consider dropping the unused alpha parameter and the stale op_div comment, and decide whether the change in meaning for out-of-tree callers needs a rename or a note.

[shoumikhin]

heads up, unittest is red on main on every row (linux, macos, windows, plus the editable and release variants) since this landed. two tests fail:

OpLeTensorOutTest.BroadcastSameDtypeOut and OpLeTensorOutTest.BroadcastBoolInputs

op_le_test.cpp:75: Failure
Expected: is equal to ETensor(sizes={2, 4}, dtype=Int, data={0, 1, 1, 1, 1, 1, 0, 0})
  Actual: ETensor(sizes={2, 4}, dtype=Int, data={1, 1, 0, 0, 0, 0, 1, 1})

looks like a merge race with #24486, which landed an hour before this one. it added its own swap for the reverse broadcast paths inside op_le.cpp. with the central swap in handle_broadcast_elementwise from this pr, le now swaps twice. the parent commit was green.

dropping the reverse path branch from op_le.cpp (back to a single x.le(y) lambda) should fix it. cc @SofianElmotiem