Adding sigmoid optimizations

Author: dijopaul

backends/cadence/hifi/operators/CMakeLists.txt

@@ -23,14 +23,14 @@ set(_aten_ops__srcs
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_add.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_div.cpp"
     "${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}/kernels/portable/cpu/op_bmm.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_cat.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_clone.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_embedding.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_full.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_permute_copy.cpp"
-    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_sigmoid.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_slice_copy.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_softmax.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_split_with_sizes_copy.cpp"

backends/cadence/hifi/operators/op_sigmoid.cpp

@@ -0,0 +1,72 @@
+/*
+ * 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 <cmath>
+
+#include <executorch/kernels/portable/cpu/util/functional_util.h>
+#include <executorch/runtime/kernel/kernel_includes.h>
+#include "kernels.h"
+
+namespace torch {
+namespace executor {
+namespace native {
+
+using Tensor = exec_aten::Tensor;
+
+Tensor& sigmoid_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
+  (void)ctx;
+
+  ET_KERNEL_CHECK(
+      ctx, in.scalar_type() != ScalarType::Bool, InvalidArgument, out);
+  ET_KERNEL_CHECK(ctx, tensor_is_floating_type(out), InvalidArgument, out);
+
+  // Resize for dynamic shape
+  ET_KERNEL_CHECK_MSG(
+      ctx,
+      resize_tensor(out, in.sizes()) == Error::Ok,
+      InvalidArgument,
+      out,
+      "Failed to resize output tensor.");
+
+  ScalarType in_type = in.scalar_type();
+  ScalarType out_type = out.scalar_type();
+  
+  int fall_back = 0;
+  if((in_type != ScalarType::Float) || (out_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_sigmoid_f32_f32(data_out, data_in, in.numel());
+  }
+  else
+  {
+    ET_SWITCH_REALHB_TYPES(in_type, ctx, "sigmoid.out", CTYPE_IN, [&]() {
+      ET_SWITCH_FLOATH_TYPES(out_type, ctx, "sigmoid.out", CTYPE_OUT, [&]() {
+        apply_unary_map_fn(
+            [](const CTYPE_IN val_in) {
+              // perform math in double to preserve precision
+              double in_casted = static_cast<double>(val_in);
+              double out_val = 1.0 / (1.0 + exp(-in_casted));
+              return static_cast<CTYPE_OUT>(out_val);
+            },
+            in.const_data_ptr<CTYPE_IN>(),
+            out.mutable_data_ptr<CTYPE_OUT>(),
+            in.numel());
+      });
+    }); 
+  }
+
+  return out;
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch