Add vectorization in elementwise_util (not working yet)

this works with op_mul, which is vectorized-friendly, but doesn't work
when we roll out to pattern.h because those ops will not work with
Vectorized yet. See TODO in elementwise_util.h

ghstack-source-id: 033b63ce3bee8a0136efdab3e03905cafb79b915
ghstack-comment-id: 2738665976
Pull Request resolved: #9432

Author: swolchok

Diff for: .lintrunner.toml

@@ -264,6 +264,10 @@ exclude_patterns = [
     'examples/**',
     'exir/verification/bindings.cpp',
     'extension/**',
+    # Uses properly-gated (ET_USE_PYTORCH_HEADERS) ATen include.
+    'kernels/portable/cpu/util/elementwise_util.h',
+    'kernels/portable/cpu/util/math_util.h',
+    'kernels/portable/cpu/util/vectorized_math.h',
     'kernels/optimized/**',
     'runtime/core/exec_aten/**',
     # Want to be able to keep c10 in sync with PyTorch core.

Diff for: kernels/portable/cpu/op_atan2.cpp

@@ -60,7 +60,7 @@ Tensor& atan2_out(
         op_name,
         utils::SupportedTensorDtypes::FLOATHBF16>(
         [](const auto val_a, const auto val_b) {
-          return std::atan2(val_a, val_b);
+          return executorch::math::atan2(val_a, val_b);
         },
         ctx,
         a,

Diff for: kernels/portable/cpu/op_elu.cpp

@@ -48,8 +48,7 @@ Tensor& elu_out(
         CTYPE,
         op_name,
         utils::SupportedTensorDtypes::SAME_AS_COMMON>(
-        [negcoef, math_scale, math_input_scale](const auto x) {
-          // TODO: rewrite this to be vectorization-capable.
+        [negcoef, math_scale, math_input_scale](const CTYPE x) {
           return MathT(x) <= MathT(0)
               ? std::expm1(MathT(x) * math_input_scale) * negcoef
               : MathT(x) * math_scale;

Diff for: kernels/portable/cpu/op_fmod.cpp

@@ -61,7 +61,7 @@ Tensor& fmod_Tensor_out(
         utils::SupportedTensorDtypes::REALHBF16>(
         [&div_by_zero_error](
             const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
-          // TODO: rewrite this to be vectorization-capable.
+          // TODO: rewrite this to be vectorization-capable?
           CTYPE_COMPUTE value = 0;
           if (is_integral_type<CTYPE_COMPUTE, /*includeBool=*/true>::value) {
             if (val_b == 0) {
@@ -138,10 +138,8 @@ Tensor& fmod_Scalar_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBF16>(
-        [val_b](const CTYPE_COMPUTE val_a) {
-          // TODO: rewrite this to be vectorization-capable.
-          CTYPE_COMPUTE value = std::fmod(val_a, val_b);
-          return value;
+        [val_b](const auto val_a) {
+          return executorch::math::fmod(val_a, (decltype(val_a))val_b);
         },
         ctx,
         a,

Diff for: kernels/portable/cpu/op_maximum.cpp

@@ -49,7 +49,7 @@ Tensor& maximum_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBBF16>(
-        [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
+        [](const auto val_a, const auto val_b) {
           return utils::max_override(val_a, val_b);
         },
         ctx,

Diff for: kernels/portable/cpu/op_minimum.cpp

@@ -49,8 +49,7 @@ Tensor& minimum_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBBF16>(
-        [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
-          // TODO: rewrite this to be vectorization-capable.
+        [](const auto val_a, const auto val_b) {
           return utils::min_override(val_a, val_b);
         },
         ctx,

Diff for: kernels/portable/cpu/op_mul.cpp

@@ -56,9 +56,7 @@ Tensor& mul_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBBF16>(
-        [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
-          return val_a * val_b;
-        },
+        [](const auto val_a, const auto val_b) { return val_a * val_b; },
         ctx,
         a,
         utils::SupportedTensorDtypes::REALHBBF16,

Diff for: kernels/portable/cpu/op_pow.cpp

@@ -57,9 +57,9 @@ Tensor& pow_Tensor_Tensor_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::REALHBF16>(
-        [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b) {
+        [](const auto val_a, const auto val_b) {
           // TODO: rewrite this to be vectorization-capable.
-          return std::pow(val_a, val_b);
+          return executorch::math::pow(val_a, val_b);
         },
         ctx,
         a,

Diff for: kernels/portable/cpu/op_sigmoid.cpp

@@ -49,8 +49,9 @@ Tensor& sigmoid_out(KernelRuntimeContext& ctx, const Tensor& in, Tensor& out) {
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::FLOATHBF16>(
-        [](const auto val_in) -> CTYPE_COMPUTE {
-          // TODO: rewrite this to be vectorization-capable
+        [](const CTYPE_COMPUTE val_in) {
+          // TODO: rewrite this to be vectorization-capable; need
+          // unary - overload for Vectorized.
           CTYPE_COMPUTE out_val = static_cast<CTYPE_COMPUTE>(1.0) /
               (static_cast<CTYPE_COMPUTE>(1.0) + exp(-val_in));
           return out_val;

Diff for: kernels/portable/cpu/op_where.cpp

@@ -47,7 +47,7 @@ Tensor& where_out(
         CTYPE_COMPUTE,
         op_name,
         utils::SupportedTensorDtypes::SAME_AS_COMMON>(
-        [](const auto val_a, const auto val_b, const auto val_c) {
+        [](const CTYPE_COMPUTE val_a, const CTYPE_COMPUTE val_b, const CTYPE_COMPUTE val_c) {
           return val_c ? val_a : val_b;
         },
         ctx,

Diff for: kernels/portable/cpu/util/elementwise_util.h

@@ -12,9 +12,14 @@
 #include <executorch/kernels/portable/cpu/util/broadcast_indexes_range.h>
 #include <executorch/kernels/portable/cpu/util/broadcast_util.h>
 #include <executorch/kernels/portable/cpu/util/dtype_util.h>
+#include <executorch/kernels/portable/cpu/util/vectorized_math.h> // Make vectorization support easy for clients.
 #include <executorch/runtime/kernel/kernel_runtime_context.h>
 #include <executorch/runtime/kernel/thread_parallel_interface.h>
 
+#ifdef ET_USE_PYTORCH_HEADERS
+#include <ATen/cpu/vec/vec.h>
+#endif // ET_USE_PYTORCH_HEADERS
+
 #include <array>
 #include <utility>
 
@@ -51,6 +56,34 @@ inline int64_t scalar_to<int64_t>(const Scalar& s) {
 }
 
 namespace internal {
+template <typename Ignore, typename T>
+using ignore_first_yield_second = T;
+
+#ifdef ET_USE_PYTORCH_HEADERS
+// Can I call a function of type Op with sizeof...(Args) arguments of type
+// at::vec::Vectorized<CTYPE_COMPUTE>?
+//
+// See [NOTE: Generic lambdas] below for requirements on Op.
+template <typename CTYPE_COMPUTE, typename Op, typename... Args>
+constexpr bool can_use_vectorized() {
+  using Vec = at::vec::Vectorized<CTYPE_COMPUTE>;
+  if constexpr (std::is_invocable_v<
+                    Op,
+                    ignore_first_yield_second<Args, Vec>...>) {
+    // For bool, we will get a false positive if we rely on only the
+    // is_invocable_v check above because at::vec::Vectorized is
+    // implicitly convertible to a pointer, which makes it implicitly
+    // convertible to bool (which was 15 minutes of fun to debug). Also
+    // just seems like good hygiene to make sure we get the Vectorized
+    // we're expecting.
+    return std::is_same_v<
+        std::invoke_result_t<Op, ignore_first_yield_second<Args, Vec>...>,
+        Vec>;
+  }
+  return false;
+}
+#endif // ET_USE_PYTORCH_HEADERS
+
 template <
     typename CTYPE_COMPUTE,
     typename CTYPE_OUT,
@@ -61,8 +94,71 @@ inline void dtype_specialized_elementwise_fn_impl(
     KernelRuntimeContext& ctx,
     const Tensor& out,
     Args... inputs) {
+  static_assert(
+      (std::is_same_v<Args, std::pair<const Tensor*, SupportedTensorDtypes>> &&
+       ...));
   constexpr auto kNumInputs = sizeof...(inputs);
-  ET_DCHECK(((inputs.first->element_size() == sizeof(CTYPE_COMPUTE)) && ...));
+  // All inputs must be of type CTYPE_COMPUTE.
+  ET_DCHECK(
+      ((inputs.first->scalar_type() ==
+        CppTypeToScalarType<CTYPE_COMPUTE>::value) &&
+       ...));
+
+#ifdef ET_USE_PYTORCH_HEADERS
+  if constexpr (can_use_vectorized<CTYPE_COMPUTE, Op, Args...>()) {
+    const bool any_is_broadcasted =
+        !(torch::executor::internal::sizes_match_ignoring_leading_1s(
+              inputs.first->sizes(), out.sizes()) &&
+          ...);
+    if (!any_is_broadcasted) {
+      using Vec = at::vec::Vectorized<CTYPE_COMPUTE>;
+      ::executorch::extension::parallel_for(
+          0,
+          out.numel(),
+          ::executorch::extension::internal::GRAIN_SIZE,
+          [&](const auto begin, const auto end) {
+            std::array<const CTYPE_COMPUTE*, kNumInputs> inputs_data_ptrs = {
+                inputs.first->template const_data_ptr<CTYPE_COMPUTE>()...};
+
+            CTYPE_OUT* const data_out = out.mutable_data_ptr<CTYPE_OUT>();
+
+            const auto vectorized_begin =
+                begin + (Vec::size() - begin % Vec::size()) % Vec::size();
+            const auto vectorized_end = end - (end % Vec::size());
+            // Scalar prologue.
+            for (const auto idx : c10::irange(begin, vectorized_begin)) {
+              std::array<CTYPE_COMPUTE, kNumInputs> loaded_inputs;
+              for (const auto input_idx : c10::irange(kNumInputs)) {
+                loaded_inputs[input_idx] = inputs_data_ptrs[input_idx][idx];
+              }
+              data_out[idx] = std::apply(compute_fun, loaded_inputs);
+            }
+
+            // Main vectorized loop.
+            for (auto idx = vectorized_begin; idx < vectorized_end;
+                 idx += Vec::size()) {
+              std::array<Vec, kNumInputs> loaded_vec_inputs;
+              for (const auto input_idx : c10::irange(kNumInputs)) {
+                loaded_vec_inputs[input_idx] =
+                    Vec::loadu(&inputs_data_ptrs[input_idx][idx]);
+              }
+              auto result_vec = std::apply(compute_fun, loaded_vec_inputs);
+              result_vec.store(&data_out[idx]);
+            }
+
+            // Scalar epilogue.
+            for (const auto idx : c10::irange(vectorized_end, end)) {
+              std::array<CTYPE_COMPUTE, kNumInputs> loaded_inputs;
+              for (const auto input_idx : c10::irange(kNumInputs)) {
+                loaded_inputs[input_idx] = inputs_data_ptrs[input_idx][idx];
+              }
+              data_out[idx] = std::apply(compute_fun, loaded_inputs);
+            }
+          });
+      return;
+    }
+  }
+#endif
 
   ::executorch::extension::parallel_for(
       0,
@@ -240,6 +336,19 @@ inline void apply_unitensor_elementwise_fn(
       compute_fun, ctx, out, out_dtypes, std::make_pair(&a, a_dtypes));
 }
 
+/**
+ * Useful for unary elementwise operators. For each element of the
+ * input, call Op and write to the corresponding element of the
+ * output. Tensor broadcasting is applied wherever it is required.
+ *
+ * [NOTE: Generic lambdas]: If Op is a *generic* lambda (i.e., one with `auto`
+ * parameters; normal lambdas are fine), it must fulfill one of the
+ * following conditions. Either:
+ * 1) It must in fact compile when passed at::vec::Vectorized<CTYPE_COMPUTE>, or
+ * 2) It must be actively SFINAE-friendly, as per the C++17 examples in
+ * https://stackoverflow.com/questions/76525790/detecting-if-a-generic-lambda-with-certain-arguments-is-invocable
+ * .
+ */
 template <
     typename CTYPE_COMPUTE,
     const char* op_name,
@@ -281,6 +390,8 @@ inline void apply_bitensor_elementwise_fn(
  * Useful for bi-tensor elementwise operators. For each element of the inputs,
  * perform a computation and write to the corresponding element of the output.
  * Tensor broadcasting is applied wherever it is required.
+ * See [NOTE: Generic lambdas] if you want to pass a generic lambda for
+ * compute_fun.
  */
 template <
     typename CTYPE_COMPUTE,
@@ -347,6 +458,9 @@ inline void apply_tritensor_elementwise_fn(
  *
  * static constexpr const char op_name[] = "my_op";
  * apply_ternary_elementwise_fn<CTYPE_COMPUTE, op_name>.
+ *
+ * See [NOTE: Generic lambdas] if you want to pass a generic lambda for
+ * compute_fun.
  */
 template <
     typename CTYPE_COMPUTE,

Diff for: kernels/portable/cpu/util/math_util.h

@@ -8,6 +8,10 @@
 
 #pragma once
 
+#ifdef ET_USE_PYTORCH_HEADERS
+#include <ATen/cpu/vec/vec.h>
+#endif
+
 namespace torch {
 namespace executor {
 namespace native {
@@ -138,6 +142,21 @@ T max_override(T a, T b) {
   return b;
 }
 
+#ifdef ET_USE_PYTORCH_HEADERS
+template <typename T>
+at::vec::Vectorized<T> min_override(
+    at::vec::Vectorized<T> a,
+    at::vec::Vectorized<T> b) {
+  return at::vec::minimum(a, b);
+}
+
+template <typename T>
+at::vec::Vectorized<T> max_override(
+    at::vec::Vectorized<T> a,
+    at::vec::Vectorized<T> b) {
+  return at::vec::maximum(a, b);
+}
+#endif
 /**
  * There is a slight difference in how std::fmod works compared to how ATen
  * determines remainders:

Diff for: kernels/portable/cpu/util/targets.bzl

@@ -32,6 +32,7 @@ def define_common_targets():
             "//executorch/kernels/portable/cpu/util:slice_util",
             "//executorch/kernels/portable/cpu/util:elementwise_util",
             "//executorch/kernels/portable/cpu/util:upsample_util",
+            "//executorch/kernels/portable/cpu/util:vectorized_math",
             "//executorch/runtime/kernel:thread_parallel_interface",
         ],
         visibility = ["//executorch/...", "@EXECUTORCH_CLIENTS"],
@@ -110,6 +111,8 @@ def define_common_targets():
             ":broadcast_indexes_range",
             ":broadcast_util",
             ":dtype_util",
+            ":vectorized_math",
+            "//executorch/runtime/core/portable_type/c10/c10:aten_headers_for_executorch",
             "//executorch/runtime/kernel:kernel_runtime_context",
             "//executorch/runtime/kernel:thread_parallel_interface",
         ],
@@ -260,6 +263,9 @@ def define_common_targets():
         srcs = [],
         exported_headers = ["math_util.h"],
         visibility = ["//executorch/kernels/portable/cpu/...", "//executorch/kernels/quantized/..."],
+        exported_deps = [
+            "//executorch/runtime/core/portable_type/c10/c10:aten_headers_for_executorch",
+        ],
     )
 
     runtime.cxx_library(
@@ -307,6 +313,15 @@ def define_common_targets():
         ],
     )
 
+    runtime.cxx_library(
+        name = "vectorized_math",
+        exported_headers = ["vectorized_math.h"],
+        visibility = ["//executorch/..."],
+        exported_deps = [
+            "//executorch/runtime/core/exec_aten/util:scalar_type_util",
+        ],
+    )
+
     # Utility functions that can be used by operators that perform reduction
     for aten_mode in get_aten_mode_options():
         suffix = "_aten" if aten_mode else ""