Swap reverse-broadcast operands once in handle_broadcast_elementwise (#24809)

### 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.

---------

Co-authored-by: SofianElmotiem <sofianelmotiem@gmail.com>

Author: JakeStevens

kernels/optimized/cpu/binary_ops.h

@@ -223,17 +223,15 @@ std::optional<BroadcastElementwisePlan> plan_broadcast_elementwise(
     const Tensor& b,
     Tensor& out,
     const ElementwiseOptimizedPath selected_optimized_path);
-} // namespace internal
 
 template <typename CTYPE, typename Op>
-Tensor& handle_broadcast_elementwise(
+Tensor& handle_broadcast_elementwise_impl(
     KernelRuntimeContext& ctx,
     const Op& vec_fun,
     const Tensor& a,
     const Tensor& b,
     Tensor& out,
-    const ElementwiseOptimizedPath selected_optimized_path,
-    const std::optional<Scalar>& alpha = {}) {
+    const ElementwiseOptimizedPath selected_optimized_path) {
   if ((selected_optimized_path ==
        ElementwiseOptimizedPath::kBroadcastLastDim) ||
       (selected_optimized_path ==
@@ -257,5 +255,35 @@ Tensor& handle_broadcast_elementwise(
       opt_plan->inner_size);
   return out;
 }
+} // namespace internal
+
+// vec_fun always receives (a, b). The reverse paths hand the broadcasting maps
+// (b, a), so vec_fun's arguments are swapped back on those paths. Pass
+// kCommutative = true when the order doesn't matter to skip compiling the
+// swapped copy.
+template <typename CTYPE, bool kCommutative = false, typename Op>
+Tensor& handle_broadcast_elementwise(
+    KernelRuntimeContext& ctx,
+    const Op& vec_fun,
+    const Tensor& a,
+    const Tensor& b,
+    Tensor& out,
+    const ElementwiseOptimizedPath selected_optimized_path,
+    const std::optional<Scalar>& alpha = {}) {
+  if constexpr (!kCommutative) {
+    if (selected_optimized_path ==
+            ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments ||
+        selected_optimized_path ==
+            ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments ||
+        selected_optimized_path ==
+            ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments) {
+      auto swapped_fun = [&vec_fun](auto x, auto y) { return vec_fun(y, x); };
+      return internal::handle_broadcast_elementwise_impl<CTYPE>(
+          ctx, swapped_fun, a, b, out, selected_optimized_path);
+    }
+  }
+  return internal::handle_broadcast_elementwise_impl<CTYPE>(
+      ctx, vec_fun, a, b, out, selected_optimized_path);
+}
 } // namespace executor
 } // namespace torch

kernels/optimized/cpu/op_add_sub_impl.h

@@ -95,51 +95,17 @@ Tensor& opt_add_sub_out_impl(
       using Vec = at::vec::Vectorized<CTYPE>;
       Vec alpha_val_vec(alpha_val);
       if constexpr (is_sub) {
-        if (selected_optimized_path ==
-                ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments ||
-            selected_optimized_path ==
-                ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments ||
-            selected_optimized_path ==
-                ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments) {
-          auto add_lambda = [&alpha_val_vec](auto x, auto y) {
-            return y - alpha_val_vec * x;
-          };
-          torch::executor::handle_broadcast_elementwise<CTYPE>(
-              ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
-        } else {
-          auto add_lambda = [&alpha_val_vec](auto x, auto y) {
-            return x - alpha_val_vec * y;
-          };
-          torch::executor::handle_broadcast_elementwise<CTYPE>(
-              ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
-        }
+        auto add_lambda = [&alpha_val_vec](auto x, auto y) {
+          return x - alpha_val_vec * y;
+        };
+        torch::executor::handle_broadcast_elementwise<CTYPE>(
+            ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
       } else {
-        if (selected_optimized_path ==
-                ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments ||
-            selected_optimized_path ==
-                ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments ||
-            selected_optimized_path ==
-                ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments) {
-          // Reason we swap out args here is because
-          // handle_broadcast_elementwise handles this selected_optimized_path
-          // option a bit differently. This should really be resolved in
-          // handle_broadcast_elementwise. However, the current blocker is that
-          // handle_broadcast_elementwise tries to be agnostic of op. This
-          // should be fixed, likely by moving lambda creation to
-          // handle_broadcast_elementwise and it be aware of which op is being
-          // executed.
-          auto add_lambda = [&alpha_val_vec](auto x, auto y) {
-            return y + alpha_val_vec * x;
-          };
-          torch::executor::handle_broadcast_elementwise<CTYPE>(
-              ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
-        } else {
-          auto add_lambda = [&alpha_val_vec](auto x, auto y) {
-            return x + alpha_val_vec * y;
-          };
-          torch::executor::handle_broadcast_elementwise<CTYPE>(
-              ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
-        }
+        auto add_lambda = [&alpha_val_vec](auto x, auto y) {
+          return x + alpha_val_vec * y;
+        };
+        torch::executor::handle_broadcast_elementwise<CTYPE>(
+            ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
       }
     });
   } else {

kernels/optimized/cpu/op_div.cpp

@@ -137,20 +137,9 @@ Tensor& opt_div_out(
     // Reason for using alpha is becasuse handle_broadcast_elementwise
     // is used for add and sub as well:
     ET_SWITCH_REALB_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
-      if (selected_optimized_path ==
-              ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments ||
-          selected_optimized_path ==
-              ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments ||
-          selected_optimized_path ==
-              ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments) {
-        auto div_lambda = [](auto x, auto y) { return y / x; };
-        torch::executor::handle_broadcast_elementwise<CTYPE>(
-            ctx, div_lambda, a, b, out, selected_optimized_path);
-      } else {
-        auto div_lambda = [](auto x, auto y) { return x / y; };
-        torch::executor::handle_broadcast_elementwise<CTYPE>(
-            ctx, div_lambda, a, b, out, selected_optimized_path);
-      }
+      auto div_lambda = [](auto x, auto y) { return x / y; };
+      torch::executor::handle_broadcast_elementwise<CTYPE>(
+          ctx, div_lambda, a, b, out, selected_optimized_path);
     });
   } else {
     ScalarType common_type = get_common_type(a.scalar_type(), b.scalar_type());

kernels/optimized/cpu/op_mul.cpp

@@ -101,14 +101,16 @@ Tensor& opt_mul_out(
 
       ET_SWITCH_COMPLEXH_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
         auto mul_lambda = [](auto x, auto y) { return x * y; };
-        torch::executor::handle_broadcast_elementwise<CTYPE>(
-            ctx, mul_lambda, a, b, out, selected_optimized_path);
+        torch::executor::
+            handle_broadcast_elementwise<CTYPE, /*kCommutative=*/true>(
+                ctx, mul_lambda, a, b, out, selected_optimized_path);
       });
     } else {
       ET_SWITCH_REALB_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
         auto mul_lambda = [](auto x, auto y) { return x * y; };
-        torch::executor::handle_broadcast_elementwise<CTYPE>(
-            ctx, mul_lambda, a, b, out, selected_optimized_path);
+        torch::executor::
+            handle_broadcast_elementwise<CTYPE, /*kCommutative=*/true>(
+                ctx, mul_lambda, a, b, out, selected_optimized_path);
       });
     }
   } else {