Reduce build size of op_cumsum (#6021)

Summary:
Pull Request resolved: #6021

98 K -> 9 K
ghstack-source-id: 246985127
exported-using-ghexport

Reviewed By: malfet, swolchok

Differential Revision: D63997235

fbshipit-source-id: c78379335ea51d84891a0ed414944cdaca3695c7

Author: manuelcandales

kernels/portable/cpu/op_cumsum.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/kernels/portable/cpu/util/dtype_util.h>
 #include <executorch/kernels/portable/cpu/util/kernel_ops_util.h>
 #include <executorch/runtime/kernel/kernel_includes.h>
 #include <executorch/runtime/platform/assert.h>
@@ -34,17 +35,22 @@ namespace {
  * the memory level, thereby increasing the speed of memory IO as
  * well as reducing the number of cache misses.
  */
-template <typename CTYPE_IN, typename CTYPE_OUT>
-void cumsum_tensors(const Tensor& self, int64_t dim, Tensor& out) {
+template <typename CTYPE_OUT, typename LoadFn = CTYPE_OUT (*)(const void*)>
+void cumsum_tensors(
+    const Tensor& self,
+    LoadFn load_self,
+    int64_t dim,
+    Tensor& out) {
   if (self.numel() == 0) {
     return;
   }
 
-  const CTYPE_IN* input_data_base = self.const_data_ptr<CTYPE_IN>();
+  const char* const input_data_base =
+      reinterpret_cast<const char*>(self.const_data_ptr());
   CTYPE_OUT* output_data_base = out.mutable_data_ptr<CTYPE_OUT>();
 
   if (self.dim() == 0) {
-    output_data_base[0] = input_data_base[0];
+    output_data_base[0] = load_self(&input_data_base[0]);
     return;
   }
 
@@ -57,15 +63,16 @@ void cumsum_tensors(const Tensor& self, int64_t dim, Tensor& out) {
 
     for (size_t idx = 0; idx < trailing_dims; idx++) {
       output_data_base[start_loc + idx] =
-          static_cast<CTYPE_OUT>(input_data_base[start_loc + idx]);
+          load_self(&input_data_base[(start_loc + idx) * self.element_size()]);
     }
 
     for (size_t j = 1; j < dim_size; j++) {
       size_t cur_round_base = start_loc + j * trailing_dims;
       size_t prev_round_base = start_loc + (j - 1) * trailing_dims;
       for (size_t idx = 0; idx < trailing_dims; idx++) {
         output_data_base[cur_round_base + idx] =
-            static_cast<CTYPE_OUT>(input_data_base[cur_round_base + idx]) +
+            load_self(&input_data_base
+                          [(cur_round_base + idx) * self.element_size()]) +
             output_data_base[prev_round_base + idx];
       }
     }
@@ -101,13 +108,15 @@ Tensor& cumsum_out(
 
   dim = (self.dim() == 0) ? 0 : dim < 0 ? dim + self.dim() : dim;
 
-  ET_SWITCH_REAL_TYPES_AND(
-      Bool, self.scalar_type(), ctx, "cumsum", CTYPE_SELF, [&] {
-        ET_SWITCH_REAL_TYPES_AND(
-            Bool, out.scalar_type(), ctx, "cumsum", CTYPE_OUT, [&] {
-              cumsum_tensors<CTYPE_SELF, CTYPE_OUT>(self, dim, out);
-            });
-      });
+  // @lint-ignore CLANGTIDY facebook-hte-CArray
+  static constexpr const char op_name[] = "cumsum.out";
+
+  ET_SWITCH_REALHBBF16_TYPES(out.scalar_type(), ctx, op_name, CTYPE_OUT, [&] {
+    const auto load_self =
+        utils::internal::get_load_to_common_fn<CTYPE_OUT, op_name>(
+            self, utils::SupportedTensorDtypes::REALHBBF16);
+    cumsum_tensors<CTYPE_OUT>(self, load_self, dim, out);
+  });
 
   return out;
 }

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

@@ -447,6 +447,7 @@ ATEN_OPS = (
     op_target(
         name = "op_cumsum",
         deps = [
+            "//executorch/kernels/portable/cpu/util:dtype_util",
             "//executorch/runtime/core/exec_aten/util:scalar_type_util",
             "//executorch/runtime/core/exec_aten/util:tensor_util",
             "//executorch/kernels/portable/cpu/util:kernel_ops_util",