Adding tanh optimizations

dijopaul · Sep 20, 2024 · c0b1005 · c0b1005
1 parent 6ad490a
commit c0b1005
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 0 deletions.
diff --git a/backends/cadence/aot/functions_hifi.yaml b/backends/cadence/aot/functions_hifi.yaml
@@ -92,6 +92,11 @@
     - arg_meta: null
       kernel_name: torch::executor::sub_out
 
+- op: tanh.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::tanh_out
+
 - op: view_copy.out
   kernels:
     - arg_meta: null

diff --git a/backends/cadence/hifi/operators/CMakeLists.txt b/backends/cadence/hifi/operators/CMakeLists.txt
@@ -25,6 +25,7 @@ set(_aten_ops__srcs
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_mul.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_sigmoid.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_sub.cpp"
+    "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_tanh.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_bmm.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_cat.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_clone.cpp"
@@ -37,6 +38,7 @@ set(_aten_ops__srcs
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_to_copy.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_view_copy.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_where.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/pattern/unary_ufunc_realhb_to_floath.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/activation_ops_util.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/broadcast_util.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/util/copy_ops_util.cpp"

diff --git a/backends/cadence/hifi/operators/op_tanh.cpp b/backends/cadence/hifi/operators/op_tanh.cpp
@@ -0,0 +1,40 @@
+/*
+ * 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 <executorch/kernels/portable/cpu/pattern/pattern.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include <cmath>
+#include "kernels.h"
+
+namespace torch {
+namespace executor {
+namespace native {
+
+Tensor& tanh_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
+
+  int fall_back = 0;
+  if((in.scalar_type() != ScalarType::Float) || (out.scalar_type() != ScalarType::Float))
+      fall_back = 1;
+
+  if(!fall_back)
+  {
+    float* data_in = in.mutable_data_ptr<float>();
+    float* data_out = out.mutable_data_ptr<float>();
+    xa_nn_vec_tanh_f32_f32(data_out, data_in, (int)in.numel());
+    return out;
+  }
+  else
+  {
+    return internal::unary_ufunc_realhb_to_floath(std::tanh, ctx, in, out);
+  }
+
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch