Update on "[ExecuTorch] Add broadcast support for optimized add op"

Summary:
This brings add op to feature parity, wrt, broadcasting, to mul op in
optimized kernels lib

Test Plan:
tests added

Reviewers:

Subscribers:

Tasks:

Tags:

cc larryliu0820 manuelcandales

[ghstack-poisoned]

Author: kimishpatel

kernels/optimized/cpu/binary_ops.h

@@ -192,7 +192,7 @@ std::array<int32_t, 3> inline get_normalized_tensor_size(
   return normalized_tensor_size;
 }
 
-template <const char* op_name, typename Op>
+template <typename CTYPE, typename Op>
 Tensor& handle_last_dim_broadcast_elementwise(
     KernelRuntimeContext& ctx,
     const Op& vec_fun,
@@ -221,32 +221,17 @@ Tensor& handle_last_dim_broadcast_elementwise(
       "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, op_name, CTYPE, [&]() {
-    using Vec = executorch::vec::Vectorized<CTYPE>;
-    Vec alpha_val_vec;
-    if (alpha.has_value()) {
-      CTYPE alpha_val;
-      ET_KERNEL_CHECK(
-          ctx,
-          native::utils::extract_scalar(alpha.value(), &alpha_val),
-          InvalidArgument, );
-      alpha_val_vec = Vec(alpha_val);
-    }
-    auto vec_fun_alpha = [vec_fun, alpha_val_vec](const Vec& a, const Vec& b) {
-      return vec_fun(a, b, alpha_val_vec);
-    };
-    executorch::vec::broadcasting_map_broadcast_last_dim<CTYPE>(
-        vec_fun_alpha,
-        out.mutable_data_ptr<CTYPE>(),
-        lhs->const_data_ptr<CTYPE>(),
-        rhs->const_data_ptr<CTYPE>(),
-        outer_size,
-        broadcast_size);
-  });
+  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 <const char* op_name, typename Op>
+template <typename CTYPE, typename Op>
 Tensor& handle_broadcast_elementwise(
     KernelRuntimeContext& ctx,
     const Op& vec_fun,
@@ -259,11 +244,10 @@ Tensor& handle_broadcast_elementwise(
        ElementwiseOptimizedPath::kBroadcastLastDim) ||
       (selected_optimized_path ==
        ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments)) {
-    return handle_last_dim_broadcast_elementwise<op_name>(
-        ctx, vec_fun, a, b, out, selected_optimized_path, alpha);
+    return handle_last_dim_broadcast_elementwise<CTYPE>(
+        ctx, vec_fun, a, b, out, selected_optimized_path);
   }
 
-  ScalarType out_type = out.scalar_type();
   const Tensor* lhs;
   const Tensor* rhs;
   if ((selected_optimized_path ==
@@ -306,30 +290,14 @@ Tensor& handle_broadcast_elementwise(
     broadcast_size = lhs->sizes()[lhs->dim() - 2];
     inner_size = lhs->sizes()[lhs->dim() - 1];
   }
-  ET_SWITCH_REALB_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
-    using Vec = executorch::vec::Vectorized<CTYPE>;
-    Vec alpha_val_vec;
-    if (alpha.has_value()) {
-      CTYPE alpha_val;
-      ET_KERNEL_CHECK(
-          ctx,
-          native::utils::extract_scalar(alpha.value(), &alpha_val),
-          InvalidArgument, );
-      alpha_val_vec = Vec(alpha_val);
-    }
-    auto vec_fun_alpha = [vec_fun, alpha_val_vec](const Vec& a, const Vec& b) {
-      return vec_fun(a, b, alpha_val_vec);
-    };
-    executorch::vec::
-        broadcasting_map_3d_and_unsqueezed_3d<CTYPE, decltype(vec_fun_alpha)>(
-            vec_fun_alpha,
-            out.mutable_data_ptr<CTYPE>(),
-            lhs->const_data_ptr<CTYPE>(),
-            rhs->const_data_ptr<CTYPE>(),
-            outer_size,
-            broadcast_size,
-            inner_size);
-  });
+  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

kernels/optimized/cpu/op_add.cpp

@@ -140,32 +140,43 @@ Tensor& opt_add_out(
           out.numel());
     });
   } else if (selected_optimized_path != ElementwiseOptimizedPath::kNone) {
-    static constexpr const char op_name[] = "add.out";
-    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 = [](auto x, auto y, auto alpha_val) {
-        return y + alpha_val * x;
-      };
-      return torch::executor::handle_broadcast_elementwise<op_name>(
-          ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
-    } else {
-      auto add_lambda = [](auto x, auto y, auto alpha_val) {
-        return x + alpha_val * y;
-      };
-      return torch::executor::handle_broadcast_elementwise<op_name>(
-          ctx, add_lambda, a, b, out, selected_optimized_path, alpha);
-    }
+    ScalarType out_type = out.scalar_type();
+    ET_SWITCH_REALB_TYPES(out_type, ctx, "add.out", CTYPE, [&]() {
+      CTYPE alpha_val;
+      ET_KERNEL_CHECK_MSG(
+          ctx,
+          utils::extract_scalar(alpha, &alpha_val),
+          InvalidArgument,
+          out,
+          "Failed to extract scalar alpha.");
+      using Vec = executorch::vec::Vectorized<CTYPE>;
+      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 =
         promoteTypes(a_type, b_type, /*half_to_float*/ true);

kernels/optimized/cpu/op_mul.cpp

@@ -130,15 +130,12 @@ Tensor& opt_mul_out(
           out.numel());
     });
   } else if (selected_optimized_path != ElementwiseOptimizedPath::kNone) {
-    // Reason for using alpha even when used for mul is becasuse
-    // handle_broadcast_elementwise is used for add and sub as well
-    // and it uses alpha.
-    auto mul_lambda = [](auto x, auto y, [[maybe_unused]] auto alpha) {
-      return x * y;
-    };
-    static constexpr const char op_name[] = "mul.out";
-    return torch::executor::handle_broadcast_elementwise<op_name>(
-        ctx, mul_lambda, 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);