[ExecuTorch] Add broadcast support for optimized add op

kernels/optimized/cpu/binary_ops.h

@@ -8,6 +8,8 @@
 
 #pragma once
 
+#include <executorch/kernels/optimized/vec/functional.h>
+#include <executorch/kernels/portable/cpu/scalar_utils.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 
 namespace torch {
@@ -190,5 +192,113 @@ std::array<int32_t, 3> inline get_normalized_tensor_size(
   return normalized_tensor_size;
 }
 
+template <typename CTYPE, typename Op>
+Tensor& handle_last_dim_broadcast_elementwise(
+    KernelRuntimeContext& ctx,
+    const Op& vec_fun,
+    const Tensor& a,
+    const Tensor& b,
+    Tensor& out,
+    const ElementwiseOptimizedPath selected_optimized_path,
+    const executorch::aten::optional<Scalar>& alpha = {}) {
+  ScalarType out_type = out.scalar_type();
+  const Tensor* lhs;
+  const Tensor* rhs;
+  if (selected_optimized_path ==
+      ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments) {
+    lhs = &b;
+    rhs = &a;
+  } else {
+    lhs = &a;
+    rhs = &b;
+  }
+  auto error = resize_tensor(out, lhs->sizes());
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      error == Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor.");
+  const size_t outer_size = getLeadingDims(out, out.dim() - 1);
+  const auto broadcast_size = out.size(out.dim() - 1);
+  executorch::vec::broadcasting_map_broadcast_last_dim<CTYPE, Op>(
+      vec_fun,
+      out.mutable_data_ptr<CTYPE>(),
+      lhs->const_data_ptr<CTYPE>(),
+      rhs->const_data_ptr<CTYPE>(),
+      outer_size,
+      broadcast_size);
+  return out;
+}
+
+template <typename CTYPE, 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 executorch::aten::optional<Scalar>& alpha = {}) {
+  if ((selected_optimized_path ==
+       ElementwiseOptimizedPath::kBroadcastLastDim) ||
+      (selected_optimized_path ==
+       ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments)) {
+    return handle_last_dim_broadcast_elementwise<CTYPE>(
+        ctx, vec_fun, a, b, out, selected_optimized_path);
+  }
+
+  const Tensor* lhs;
+  const Tensor* rhs;
+  if ((selected_optimized_path ==
+       ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments) ||
+      (selected_optimized_path ==
+       ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments)) {
+    lhs = &b;
+    rhs = &a;
+  } else {
+    // Catch failure to update logic when adding new broadcasting possibility.
+    ET_DCHECK(
+        (selected_optimized_path ==
+         ElementwiseOptimizedPath::kBroadcast2dBy1d) ||
+        (selected_optimized_path ==
+         ElementwiseOptimizedPath::kBroadcastNdByNd));
+    lhs = &a;
+    rhs = &b;
+  }
+  auto error = resize_tensor(out, lhs->sizes());
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      error == Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor.");
+  int64_t outer_size = 1;
+  int64_t broadcast_size;
+  int64_t inner_size;
+  if ((selected_optimized_path == ElementwiseOptimizedPath::kBroadcastNdByNd) ||
+      (selected_optimized_path ==
+       ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments)) {
+    int32_t broadcast_dim = internal::get_broadcast_dim(*lhs, *rhs);
+    int32_t broadcast_dim_lhs = lhs->dim() + broadcast_dim;
+    auto normalized_tensor_size_lhs =
+        get_normalized_tensor_size(*lhs, broadcast_dim_lhs);
+    outer_size = normalized_tensor_size_lhs[0];
+    broadcast_size = normalized_tensor_size_lhs[1];
+    inner_size = normalized_tensor_size_lhs[2];
+  } else {
+    broadcast_size = lhs->sizes()[lhs->dim() - 2];
+    inner_size = lhs->sizes()[lhs->dim() - 1];
+  }
+  executorch::vec::broadcasting_map_3d_and_unsqueezed_3d<CTYPE, Op>(
+      vec_fun,
+      out.mutable_data_ptr<CTYPE>(),
+      lhs->const_data_ptr<CTYPE>(),
+      rhs->const_data_ptr<CTYPE>(),
+      outer_size,
+      broadcast_size,
+      inner_size);
+  return out;
+}
 } // namespace executor
 } // namespace torch

kernels/optimized/cpu/op_add.cpp

@@ -140,40 +140,42 @@ Tensor& opt_add_out(
           out.numel());
     });
   } else if (selected_optimized_path != ElementwiseOptimizedPath::kNone) {
-    const Tensor* lhs;
-    const Tensor* rhs;
-    if (selected_optimized_path ==
-        ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments) {
-      lhs = &b;
-      rhs = &a;
-    } else {
-      // Catch failure to update logic when adding new broadcasting possibility.
-      ET_DCHECK(
-          selected_optimized_path ==
-          ElementwiseOptimizedPath::kBroadcast2dBy1d);
-      lhs = &a;
-      rhs = &b;
-    }
-    auto error = resize_tensor(out, lhs->sizes());
-    ET_KERNEL_CHECK_MSG(
-        ctx,
-        error == Error::Ok,
-        InvalidArgument,
-        out,
-        "Failed to resize output tensor.");
+    ScalarType out_type = out.scalar_type();
     ET_SWITCH_REALB_TYPES(out_type, ctx, "add.out", CTYPE, [&]() {
       CTYPE alpha_val;
-      ET_KERNEL_CHECK(
-          ctx, utils::extract_scalar(alpha, &alpha_val), InvalidArgument, );
-
+      ET_KERNEL_CHECK_MSG(
+          ctx,
+          utils::extract_scalar(alpha, &alpha_val),
+          InvalidArgument,
+          out,
+          "Failed to extract scalar alpha.");
       using Vec = executorch::vec::Vectorized<CTYPE>;
-      executorch::vec::broadcasting_map_2d_by_1d<CTYPE>(
-          [alpha_val](Vec x, Vec y) { return x + Vec(alpha_val) * y; },
-          out.mutable_data_ptr<CTYPE>(),
-          lhs->const_data_ptr<CTYPE>(),
-          rhs->const_data_ptr<CTYPE>(),
-          lhs->sizes()[lhs->dim() - 2],
-          lhs->sizes()[lhs->dim() - 1]);
+      Vec alpha_val_vec(alpha_val);
+      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;
+        };
+        return 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;
+        };
+        return torch::executor::handle_broadcast_elementwise<CTYPE>(
+            ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
+      }
     });
   } else {
     ScalarType common_type =

kernels/optimized/cpu/op_mul.cpp

@@ -68,114 +68,6 @@ template <
 struct MulInner<false, CTYPE_A, CTYPE_B, CTYPE_IN, CTYPE_OUT>
     : public ReportCanCastBug {};
 
-Tensor& handle_last_dim_broadcast(
-    KernelRuntimeContext& ctx,
-    const Tensor& a,
-    const Tensor& b,
-    Tensor& out,
-    const ElementwiseOptimizedPath selected_optimized_path) {
-  ScalarType out_type = out.scalar_type();
-  const Tensor* lhs;
-  const Tensor* rhs;
-  if (selected_optimized_path ==
-      ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments) {
-    lhs = &b;
-    rhs = &a;
-  } else {
-    lhs = &a;
-    rhs = &b;
-  }
-  auto error = resize_tensor(out, lhs->sizes());
-  ET_KERNEL_CHECK_MSG(
-      ctx,
-      error == Error::Ok,
-      InvalidArgument,
-      out,
-      "Failed to resize output tensor.");
-  const size_t outer_size = getLeadingDims(out, out.dim() - 1);
-  const auto broadcast_size = out.size(out.dim() - 1);
-  ET_SWITCH_REALB_TYPES(out_type, ctx, "mul.out", CTYPE, [&]() {
-    using Vec = executorch::vec::Vectorized<CTYPE>;
-    executorch::vec::broadcasting_map_broadcast_last_dim<CTYPE>(
-        [](Vec x, Vec y) { return x * y; },
-        out.mutable_data_ptr<CTYPE>(),
-        lhs->const_data_ptr<CTYPE>(),
-        rhs->const_data_ptr<CTYPE>(),
-        outer_size,
-        broadcast_size);
-  });
-  return out;
-}
-
-Tensor& handle_broadcast_mul(
-    KernelRuntimeContext& ctx,
-    const Tensor& a,
-    const Tensor& b,
-    Tensor& out,
-    const ElementwiseOptimizedPath selected_optimized_path) {
-  if ((selected_optimized_path ==
-       ElementwiseOptimizedPath::kBroadcastLastDim) ||
-      (selected_optimized_path ==
-       ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments)) {
-    return handle_last_dim_broadcast(ctx, a, b, out, selected_optimized_path);
-  }
-
-  ScalarType out_type = out.scalar_type();
-  const Tensor* lhs;
-  const Tensor* rhs;
-  if ((selected_optimized_path ==
-       ElementwiseOptimizedPath::kBroadcast2dBy1dReverseArguments) ||
-      (selected_optimized_path ==
-       ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments)) {
-    lhs = &b;
-    rhs = &a;
-  } else {
-    // Catch failure to update logic when adding new broadcasting possibility.
-    ET_DCHECK(
-        (selected_optimized_path ==
-         ElementwiseOptimizedPath::kBroadcast2dBy1d) ||
-        (selected_optimized_path ==
-         ElementwiseOptimizedPath::kBroadcastNdByNd));
-    lhs = &a;
-    rhs = &b;
-  }
-  auto error = resize_tensor(out, lhs->sizes());
-  ET_KERNEL_CHECK_MSG(
-      ctx,
-      error == Error::Ok,
-      InvalidArgument,
-      out,
-      "Failed to resize output tensor.");
-  int64_t outer_size = 1;
-  int64_t broadcast_size;
-  int64_t inner_size;
-  if ((selected_optimized_path == ElementwiseOptimizedPath::kBroadcastNdByNd) ||
-      (selected_optimized_path ==
-       ElementwiseOptimizedPath::kBroadcastNdByNdReverseArguments)) {
-    int32_t broadcast_dim = internal::get_broadcast_dim(*lhs, *rhs);
-    int32_t broadcast_dim_lhs = lhs->dim() + broadcast_dim;
-    auto normalized_tensor_size_lhs =
-        get_normalized_tensor_size(*lhs, broadcast_dim_lhs);
-    outer_size = normalized_tensor_size_lhs[0];
-    broadcast_size = normalized_tensor_size_lhs[1];
-    inner_size = normalized_tensor_size_lhs[2];
-  } else {
-    broadcast_size = lhs->sizes()[lhs->dim() - 2];
-    inner_size = lhs->sizes()[lhs->dim() - 1];
-  }
-  ET_SWITCH_REALB_TYPES(out_type, ctx, "mul.out", CTYPE, [&]() {
-    using Vec = executorch::vec::Vectorized<CTYPE>;
-    executorch::vec::broadcasting_map_3d_and_unsqueezed_3d<CTYPE>(
-        [](Vec x, Vec y) { return x * y; },
-        out.mutable_data_ptr<CTYPE>(),
-        lhs->const_data_ptr<CTYPE>(),
-        rhs->const_data_ptr<CTYPE>(),
-        outer_size,
-        broadcast_size,
-        inner_size);
-  });
-  return out;
-}
 } // namespace
 
 Tensor& opt_mul_out(
@@ -238,7 +130,12 @@ Tensor& opt_mul_out(
           out.numel());
     });
   } else if (selected_optimized_path != ElementwiseOptimizedPath::kNone) {
-    return handle_broadcast_mul(ctx, a, b, out, selected_optimized_path);
+    ScalarType out_type = out.scalar_type();
+    ET_SWITCH_REALB_TYPES(out_type, ctx, "mul.out", CTYPE, [&]() {
+      auto mul_lambda = [](auto x, auto y) { return x * y; };
+      return torch::executor::handle_broadcast_elementwise<CTYPE>(
+          ctx, mul_lambda, a, b, out, selected_optimized_path);
+    });
   } else {
     ScalarType common_type =
         promoteTypes(a_type, b_type, /*half_to_float*/ true);