Update base for Update on "[ExecuTorch] Add broadcasting support to optimized op_div"

Summary:
Similar to broadcast support in op_mul

Test Plan:
Tests added

Reviewers:

Subscribers:

Tasks:

Tags:

cc larryliu0820 manuelcandales

[ghstack-poisoned]

Author: kimishpatel

kernels/optimized/cpu/binary_ops.h

@@ -49,38 +49,8 @@ enum class ElementwiseOptimizedPath {
   kBroadcastLastDimReverseArguments,
 };
 
-enum class BinaryOpType {
-  kAdd,
-  kSub,
-  kMul,
-  kDiv,
-};
-
 namespace internal {
 
-template <BinaryOpType op_type>
-struct BinaryOpTypeName;
-
-template <>
-struct BinaryOpTypeName<BinaryOpType::kAdd> {
-  static constexpr char kName[] = "add.out";
-};
-
-template <>
-struct BinaryOpTypeName<BinaryOpType::kSub> {
-  static constexpr char kName[] = "sub.out";
-};
-
-template <>
-struct BinaryOpTypeName<BinaryOpType::kMul> {
-  static constexpr char kName[] = "mul.out";
-};
-
-template <>
-struct BinaryOpTypeName<BinaryOpType::kDiv> {
-  static constexpr char kName[] = "div.out";
-};
-
 /*
   Given two tensors, this function returns the broadcast dim if it exists.
   Returns 0 if no broadcast dim is found.
@@ -222,15 +192,15 @@ std::array<int32_t, 3> inline get_normalized_tensor_size(
   return normalized_tensor_size;
 }
 
-template <BinaryOpType op_type, typename Op>
+template <const char* op_name, 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,
-    executorch::aten::optional<Scalar>& alpha = {}) {
+    const executorch::aten::optional<Scalar>& alpha = {}) {
   ScalarType out_type = out.scalar_type();
   const Tensor* lhs;
   const Tensor* rhs;
@@ -251,11 +221,11 @@ 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, internal::BinaryOpTypeName<op_type>::kName, CTYPE, [&]() {
+  ET_SWITCH_REALB_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
     using Vec = executorch::vec::Vectorized<CTYPE>;
-    CTYPE alpha_val;
-    Vec alpha_val_vec(alpha_val);
+    Vec alpha_val_vec;
     if (alpha.has_value()) {
+      CTYPE alpha_val;
       ET_KERNEL_CHECK(
           ctx,
           native::utils::extract_scalar(alpha.value(), &alpha_val),
@@ -276,20 +246,20 @@ Tensor& handle_last_dim_broadcast_elementwise(
   return out;
 }
 
-template <BinaryOpType op_type, typename Op>
+template <const char* op_name, 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,
-    executorch::aten::optional<Scalar> alpha = {}) {
+    const executorch::aten::optional<Scalar>& alpha = {}) {
   if ((selected_optimized_path ==
        ElementwiseOptimizedPath::kBroadcastLastDim) ||
       (selected_optimized_path ==
        ElementwiseOptimizedPath::kBroadcastLastDimReverseArguments)) {
-    return handle_last_dim_broadcast_elementwise<op_type>(
+    return handle_last_dim_broadcast_elementwise<op_name>(
         ctx, vec_fun, a, b, out, selected_optimized_path, alpha);
   }
 
@@ -336,11 +306,11 @@ 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, internal::BinaryOpTypeName<op_type>::kName, CTYPE, [&]() {
+  ET_SWITCH_REALB_TYPES(out_type, ctx, op_name, CTYPE, [&]() {
     using Vec = executorch::vec::Vectorized<CTYPE>;
-    CTYPE alpha_val;
     Vec alpha_val_vec;
     if (alpha.has_value()) {
+      CTYPE alpha_val;
       ET_KERNEL_CHECK(
           ctx,
           native::utils::extract_scalar(alpha.value(), &alpha_val),

kernels/optimized/cpu/op_add.cpp

@@ -67,7 +67,8 @@ Tensor& opt_add_out(
     return opt_add_out(ctx, b, a, alpha, out);
   }
 
-  return torch::executor::kernels::impl::opt_add_sub_out_impl(
+  static constexpr const char op_name[] = "add.out";
+  return torch::executor::kernels::impl::opt_add_sub_out_impl<false, op_name>(
       ctx, a, b, alpha, out);
 }