Integrating NNLIB optimization for rsqrt op

Author: Rushi-cad

backends/cadence/aot/functions_hifi.yaml

@@ -101,6 +101,11 @@
   kernels:
     - arg_meta: null
       kernel_name: torch::executor::where_out
+      
+- op: rsqrt.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::rsqrt_out        
 
 # custom ops
 - func: cadence::quantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)

backends/cadence/hifi/operators/CMakeLists.txt

@@ -36,6 +36,7 @@ set(_aten_ops__srcs
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_cat.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_clone.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_div.cpp"
+    "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_rsqrt.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_mul.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_permute_copy.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_sigmoid.cpp"
@@ -44,7 +45,8 @@ set(_aten_ops__srcs
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_split_with_sizes_copy.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_sub.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_to_copy.cpp"
-    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_where.cpp")
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_where.cpp"
+    "${EXECUTORCH_ROOT}/kernels/portable/cpu/pattern/unary_ufunc_realhb_to_floath.cpp")
 add_library(aten_ops_cadence ${_aten_ops__srcs})
 target_link_libraries(aten_ops_cadence PUBLIC executorch)
 target_link_libraries(aten_ops_cadence PRIVATE cadence_kernels)

backends/cadence/hifi/operators/op_rsqrt.cpp

@@ -0,0 +1,50 @@
+/*
+ * 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 "kernels.h"
+
+namespace torch {
+namespace executor {
+namespace native {
+namespace {
+
+double rsqrt(double x) {
+  return 1.0 / std::sqrt(x);
+}
+
+} // namespace
+
+Tensor& rsqrt_out(RuntimeContext& ctx, const Tensor& in, Tensor& out) {
+  
+  if(in.scalar_type() == ScalarType::Float)
+  {
+    WORD32 num_elm = out.numel();
+    
+    FLOAT32 * __restrict__ p_out = (FLOAT32 * __restrict__ )out.mutable_data_ptr<float>();
+    const FLOAT32 * __restrict__ p_inp = (const FLOAT32 * __restrict__)in.const_data_ptr<float>();
+    
+    xa_nn_elm_rsqrt_f32_f32(
+                        p_out,
+                        p_inp,
+                        num_elm); 
+              
+    return out;
+  }
+  else
+  {
+    return internal::unary_ufunc_realhb_to_floath(rsqrt, ctx, in, out);
+  }
+
+}
+
+} // namespace native
+} // namespace executor
+} // namespace torch