Adding minimum, maximum and clamp operations (#7)

Co-authored-by: dijopaul <[email protected]>
dijopaul · Aug 13, 2024 · ba58c07 · ba58c07
1 parent f4cf6c8
commit ba58c07
Show file tree

Hide file tree

Showing 11 changed files with 2,548 additions and 4 deletions.
diff --git a/backends/cadence/aot/functions_hifi.yaml b/backends/cadence/aot/functions_hifi.yaml
@@ -37,6 +37,11 @@
     - arg_meta: null
       kernel_name: torch::executor::cat_out
 
+- op: clamp.Tensor_out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::clamp_tensor_out
+
 - op: clone.out
   kernels:
     - arg_meta: null
@@ -62,6 +67,16 @@
     - arg_meta: null
       kernel_name: torch::executor::full_out
 
+- op: maximum.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::maximum_out
+
+- op: minimum.out
+  kernels:
+    - arg_meta: null
+      kernel_name: torch::executor::minimum_out
+
 - op: mul.out
   kernels:
     - arg_meta: null

diff --git a/backends/cadence/hifi/kernels/CMakeLists.txt b/backends/cadence/hifi/kernels/CMakeLists.txt
@@ -9,10 +9,13 @@ add_library(
   cadence_kernels
   kernels.cpp
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/matmul_asym8uxasym8u_asym8u.cpp
+  ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_broadcast_32_32.c
+  ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_minimum_maximum_f32.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_add_f32_broadcast.c 
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_mul_f32_broadcast.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_transpose_32.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_where_f32xf32_f32.c
+  ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_clamp_f32.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_div_broadcast_f32.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_floor_div_broadcast_f32.c
 )

diff --git a/backends/cadence/hifi/kernels/kernels.h b/backends/cadence/hifi/kernels/kernels.h
@@ -15,6 +15,58 @@
 /* For NNLIB APIs */
 #include "xa_nnlib_kernels_api.h"
 
+extern "C" WORD32 xa_nn_broadcast_32_32( WORD32* __restrict__ p_out,      /* pointer to write broadcasted output data to */
+        const int *const out_shape,         /* output shape resulting after broadcast */
+
+        WORD32* __restrict__ p_in,    /* pointer to unextended input data */
+        const int * const in_shape,         /* input shape */
+        int num_dims);
+
+extern "C" WORD32 xa_nn_elm_clamp_f32xf32xf32_f32(FLOAT32 * __restrict__ p_out,
+                               const FLOAT32 * __restrict__ p_inp,
+                               const FLOAT32 * __restrict__ p_min,
+                               const FLOAT32 * __restrict__ p_max,
+                               WORD32 num_elm);
+
+extern "C" WORD32 xa_nn_elm_clamp_broadcast_4D_f32Xf32xf32_f32(FLOAT32 * __restrict__ p_out,
+                      const WORD32 *const p_out_shape,
+                      const FLOAT32 * __restrict__ p_inp,
+                      const WORD32 *const p_inp_shape,
+                      const FLOAT32 * __restrict__ p_min,
+                      const WORD32 *const p_min_shape,
+                      const FLOAT32 * __restrict__ p_max,
+                      const WORD32 *const p_max_shape
+                      );
+
+extern "C" WORD32 xa_nn_elm_maximum_f32xf32_f32(FLOAT32 * __restrict__ p_out,
+                               const FLOAT32 * __restrict__ p_inp1,
+                               const FLOAT32 * __restrict__ p_inp2,
+                               WORD32 num_elm);
+
+extern "C" WORD32 xa_nn_elm_minimum_f32xf32_f32(FLOAT32 * __restrict__ p_out,
+                               const FLOAT32 * __restrict__ p_inp1,
+                               const FLOAT32 * __restrict__ p_inp2,
+                               WORD32 num_elm);
+
+extern "C" WORD32 xa_nn_elm_where_f32xf32_f32(FLOAT32 * __restrict__ p_out, 
+                                const FLOAT32 * __restrict__ p_inp1, 
+                                const FLOAT32 * __restrict__ p_inp2, 
+                                const unsigned char *__restrict__ p_condition, 
+                                WORD32 num_elm);
+
+extern "C" WORD32 xa_nn_elm_maximum_broadcast_4D_f32xf32_f32(FLOAT32 * __restrict__ p_out,
+                      const WORD32 *const p_out_shape,
+                      const FLOAT32 * __restrict__ p_inp1,
+                      const WORD32 *const p_inp1_shape,
+                      const FLOAT32 * __restrict__ p_inp2,
+                      const WORD32 *const p_inp2_shape);
+
+extern "C" WORD32 xa_nn_elm_minimum_broadcast_4D_f32xf32_f32(FLOAT32 * __restrict__ p_out,
+                      const WORD32 *const p_out_shape,
+                      const FLOAT32 * __restrict__ p_inp1,
+                      const WORD32 *const p_inp1_shape,
+                      const FLOAT32 * __restrict__ p_inp2,
+                      const WORD32 *const p_inp2_shape);
 
 extern "C" WORD32 xa_nn_elm_floor_div_f32xf32_f32(
                                     FLOAT32 * __restrict__ p_out,
@@ -39,7 +91,7 @@ extern "C" WORD32 xa_nn_elm_div_broadcast_4D_f32xf32_f32(
                                     const WORD32 *const p_inp2_shape);
 
 extern "C" WORD32 xa_nn_elm_where_f32xf32_f32(FLOAT32 * __restrict__ p_out, const FLOAT32 * __restrict__ p_inp1, const FLOAT32 * __restrict__ p_inp2, const unsigned char *__restrict__ p_condition, WORD32 num_elm);
-                          
+
 extern "C" WORD32 xa_nn_elm_where_broadcast_4D_f32xf32_f32(FLOAT32 * __restrict__ p_out, const WORD32 *const p_out_shape,
                       const FLOAT32 * __restrict__ p_inp1,
                       const WORD32 *const p_inp1_shape,

diff --git a/backends/cadence/hifi/operators/CMakeLists.txt b/backends/cadence/hifi/operators/CMakeLists.txt
@@ -20,6 +20,9 @@ endif()
 
 # ATen compliant ops that are needed to run this model.
 set(_aten_ops__srcs
+    "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_clamp.cpp"
+    "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_maximum.cpp"
+    "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_minimum.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_add.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_embedding.cpp"
     "${EXECUTORCH_ROOT}/kernels/portable/cpu/op_full.cpp"