Reduce build size of op_addmm (#6018)

Summary:
Pull Request resolved: #6018

200 K -> 30 K
ghstack-source-id: 246985125
exported-using-ghexport

Reviewed By: swolchok

Differential Revision: D63994874

fbshipit-source-id: 1d26b944be8b0cdc4e2343e9efa6d0e28a0a82e3

Author: manuelcandales

kernels/portable/cpu/op_addmm.cpp

@@ -7,7 +7,7 @@
  */
 
 #include <executorch/kernels/portable/cpu/scalar_utils.h>
-#include <executorch/kernels/portable/cpu/util/broadcast_util.h>
+#include <executorch/kernels/portable/cpu/util/elementwise_util.h>
 #include <executorch/kernels/portable/cpu/util/matmul_ops_util.h>
 #include <executorch/kernels/portable/cpu/vec_ops.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
@@ -53,62 +53,54 @@ Tensor& addmm_out(
 
   ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out);
 
-  ScalarType alpha_dtype = utils::get_scalar_dtype(alpha);
-  ScalarType beta_dtype = utils::get_scalar_dtype(beta);
-  ET_SWITCH_REAL_TYPES_AND(
-      Half, in.scalar_type(), ctx, "addmm.out", CTYPE, [&]() {
-        ET_SWITCH_SCALAR_OBJ_TYPES(
-            alpha_dtype, ctx, "addmm.out", ALPHA_T, [&]() {
-              ET_SWITCH_SCALAR_OBJ_TYPES(
-                  beta_dtype, ctx, "addmm.out", BETA_T, [&]() {
-                    size_t m = mat1.size(0);
-                    size_t n = mat1.size(1);
-                    size_t p = mat2.size(1);
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "addmm.out";
 
-                    if (out.sizes() == in.sizes()) {
-                      // vec_addmm assumes that no broadcasting is required.
-                      vec_addmm<CTYPE, CTYPE>(
-                          out.mutable_data_ptr<CTYPE>(),
-                          in.const_data_ptr<CTYPE>(),
-                          mat1.const_data_ptr<CTYPE>(),
-                          mat2.const_data_ptr<CTYPE>(),
-                          m,
-                          n,
-                          p,
-                          convert<CTYPE>(beta.to<BETA_T>()),
-                          convert<CTYPE>(alpha.to<ALPHA_T>()));
-                    } else {
-                      // If broadcasting is required, them compute the matmul
-                      // and addition separately, using
-                      // apply_binary_elementwise_fn to perform the addition
-                      // while applying broadcasting
-                      vec_matmul<CTYPE, CTYPE>(
-                          out.mutable_data_ptr<CTYPE>(),
-                          mat1.const_data_ptr<CTYPE>(),
-                          mat2.const_data_ptr<CTYPE>(),
-                          m,
-                          n,
-                          p);
+  ET_SWITCH_REALHBF16_TYPES(in.scalar_type(), ctx, op_name, CTYPE, [&]() {
+    CTYPE alpha_val = utils::scalar_to<CTYPE>(alpha);
+    CTYPE beta_val = utils::scalar_to<CTYPE>(beta);
+    size_t m = mat1.size(0);
+    size_t n = mat1.size(1);
+    size_t p = mat2.size(1);
 
-                      CTYPE alpha_val = convert<CTYPE>(alpha.to<ALPHA_T>());
-                      CTYPE beta_val = convert<CTYPE>(beta.to<BETA_T>());
-                      apply_binary_elementwise_fn<CTYPE, CTYPE, CTYPE>(
-                          [alpha_val, beta_val](
-                              const CTYPE val_a, const CTYPE val_b) {
-                            CTYPE a_casted = static_cast<CTYPE>(val_a);
-                            CTYPE b_casted = static_cast<CTYPE>(val_b);
-                            CTYPE value =
-                                a_casted * alpha_val + b_casted * beta_val;
+    if (out.sizes() == in.sizes()) {
+      // vec_addmm assumes that no broadcasting is required.
+      vec_addmm<CTYPE, CTYPE>(
+          out.mutable_data_ptr<CTYPE>(),
+          in.const_data_ptr<CTYPE>(),
+          mat1.const_data_ptr<CTYPE>(),
+          mat2.const_data_ptr<CTYPE>(),
+          m,
+          n,
+          p,
+          beta_val,
+          alpha_val);
+    } else {
+      // If broadcasting is required, them compute the matmul
+      // and addition separately, using
+      // apply_binary_elementwise_fn to perform the addition
+      // while applying broadcasting
+      vec_matmul<CTYPE, CTYPE>(
+          out.mutable_data_ptr<CTYPE>(),
+          mat1.const_data_ptr<CTYPE>(),
+          mat2.const_data_ptr<CTYPE>(),
+          m,
+          n,
+          p);
 
-                            return value;
-                          },
-                          out,
-                          in,
-                          out);
-                    }
-                  });
-            });
-      });
+      utils::apply_bitensor_elementwise_fn<CTYPE, op_name>(
+          [alpha_val, beta_val](const CTYPE val_a, const CTYPE val_b) {
+            return val_a * alpha_val + val_b * beta_val;
+          },
+          ctx,
+          out,
+          utils::SupportedTensorDtypes::REALHBF16,
+          in,
+          utils::SupportedTensorDtypes::REALHBF16,
+          out,
+          utils::SupportedTensorDtypes::REALHBF16);
+    }
+  });
 
   return out;
 }

shim/xplat/executorch/kernels/portable/op_registration_util.bzl

@@ -224,6 +224,8 @@ ATEN_OPS = (
         name = "op_addmm",
         deps = [
             "//executorch/kernels/portable/cpu/util:broadcast_util",
+            "//executorch/kernels/portable/cpu/util:dtype_util",
+            "//executorch/kernels/portable/cpu/util:elementwise_util",
             "//executorch/kernels/portable/cpu/util:matmul_ops_util",
             ":scalar_utils",
             ":vec_ops",