Add optimized ELU implementation

This uses PyTorch code sharing, so we'll need a pin bump to pick up
pytorch/pytorch#149673 (and
pytorch/pytorch#149684 and, when it lands,
pytorch/pytorch#149780).

ghstack-source-id: 50f2d19171aa73280a6d68e297202d9ea421df8d
ghstack-comment-id: 2744722101
Pull Request resolved: #9521

Author: swolchok

kernels/optimized/cpu/op_elu.cpp

@@ -0,0 +1,96 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <ATen/native/cpu/Elu.h>
+
+#include <executorch/kernels/portable/cpu/scalar_utils.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <executorch/runtime/kernel/thread_parallel_interface.h>
+#include <executorch/runtime/platform/assert.h>
+
+namespace torch::executor::native {
+
+namespace {
+template <typename CTYPE>
+void elu(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    const Scalar& alpha,
+    const Scalar& scale,
+    const Scalar& input_scale,
+    Tensor& out) {
+  const CTYPE* in_data = input.const_data_ptr<CTYPE>();
+  CTYPE* out_data = out.mutable_data_ptr<CTYPE>();
+  using MathT =
+      std::conditional_t<c10::is_reduced_floating_point_v<CTYPE>, float, CTYPE>;
+  MathT math_alpha = 0;
+  MathT math_scale = 0;
+  MathT math_input_scale = 0;
+  ET_EXTRACT_SCALAR(alpha, math_alpha);
+  ET_EXTRACT_SCALAR(scale, math_scale);
+  ET_EXTRACT_SCALAR(input_scale, math_input_scale);
+  const auto scalar_func =
+      at::native::get_scalar_elu_elementwise_func<CTYPE, MathT>(
+          math_alpha, math_scale, math_input_scale);
+  const auto vec_func = at::native::get_vectorized_elu_elementwise_func<CTYPE>(
+      math_alpha, math_scale, math_input_scale);
+
+  ::executorch::extension::parallel_for(
+      0,
+      out.numel(),
+      ::executorch::extension::internal::GRAIN_SIZE,
+      [&](const auto begin, const auto end) {
+        using Vec = at::vec::Vectorized<CTYPE>;
+        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)) {
+          out_data[idx] = scalar_func(in_data[idx]);
+        }
+
+        // Main vectorized loop.
+        for (auto idx = vectorized_begin; idx < vectorized_end;
+             idx += Vec::size()) {
+          auto result_vec = vec_func(Vec::loadu(&in_data[idx]));
+          result_vec.store(&out_data[idx]);
+        }
+
+        // Scalar epilogue.
+        for (const auto idx : c10::irange(vectorized_end, end)) {
+          out_data[idx] = scalar_func(in_data[idx]);
+        }
+      });
+}
+} // namespace
+
+Tensor& opt_elu_out(
+    KernelRuntimeContext& ctx,
+    const Tensor& in,
+    const Scalar& alpha,
+    const Scalar& scale,
+    const Scalar& input_scale,
+    Tensor& out) {
+  ET_KERNEL_CHECK(ctx, tensors_have_same_dtype(in, out), InvalidArgument, out);
+  ET_KERNEL_CHECK(
+      ctx, resize_tensor(out, in.sizes()) == Error::Ok, InvalidArgument, out);
+
+  ET_KERNEL_CHECK(
+      ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(ctx, tensor_is_floating_type(in), InvalidArgument, out);
+
+  ET_KERNEL_CHECK(ctx, tensors_have_same_dtype(in, out), InvalidArgument, out);
+
+  ET_SWITCH_FLOATHBF16_TYPES(in.scalar_type(), ctx, "elu.out", CTYPE, [&]() {
+    elu<CTYPE>(ctx, in, alpha, scale, input_scale, out);
+  });
+  return out;
+}
+
+} // namespace torch::executor::native

kernels/optimized/cpu/targets.bzl

@@ -25,6 +25,14 @@ _OPTIMIZED_ATEN_OPS = (
             "//executorch/kernels/portable/cpu/util:broadcast_util",
         ],
     ),
+    op_target(
+        name = "op_elu",
+        deps = [
+            "//executorch/extension/threadpool:threadpool",
+            "//executorch/kernels/portable/cpu:scalar_utils",
+            "//executorch/runtime/core/portable_type/c10/c10:aten_headers_for_executorch",
+        ],
+    ),
     op_target(name = "op_exp"),
     op_target(
         name = "op_fft_r2c",

kernels/optimized/optimized.yaml

@@ -37,6 +37,11 @@
     - arg_meta: null
       kernel_name: torch::executor::opt_div_scalar_out
 
+- op: elu.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::opt_elu_out
+
 - op: exp.out
   kernels:
     - arg_meta: null

kernels/test/CMakeLists.txt

@@ -274,6 +274,7 @@ set(_optimized_kernels_test_sources
     "op_add_test.cpp"
     "op_bmm_test.cpp"
     "op_div_test.cpp"
+    "op_elu_test.cpp"
     "op_exp_test.cpp"
     "op_fft_r2c_test.cpp"
     "op_gelu_test.cpp"

kernels/test/targets.bzl

@@ -215,7 +215,7 @@ def define_common_targets():
     _common_op_test("op_detach_copy_test", ["aten", "portable"])
     _common_op_test("op_diagonal_copy_test", ["aten", "portable"])
     _common_op_test("op_div_test", ["aten", "portable", "optimized"])
-    _common_op_test("op_elu_test", ["aten", "portable"])
+    _common_op_test("op_elu_test", ["aten", "portable", "optimized"])
     _common_op_test("op_embedding_test", ["aten", "portable"])
     _common_op_test("op_empty_test", ["aten", "portable"])
     _common_op_test("op_eq_test", ["aten", "portable"])