im2row HiFi optimizations

backends/cadence/aot/functions_hifi.yaml

@@ -314,6 +314,11 @@
     - arg_meta: null
       kernel_name: cadence::impl::HiFi::quantized_linear_per_tensor_out
 
+- func: cadence::im2row.out(Tensor input, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, Tensor in_zero_point, bool channel_last=False, *, Tensor(a!) out) -> Tensor(a!)
+  kernels:
+    - arg_meta: null
+      kernel_name: cadence::impl::HiFi::im2row_out
+
 - func: cadence::quantized_relu_per_tensor.out(Tensor X, Tensor X_zero_point, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)
   kernels:
     - arg_meta: null

backends/cadence/hifi/kernels/CMakeLists.txt

@@ -18,6 +18,7 @@ add_library(
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_div_mode_f32_broadcast.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_fmod_broadcast_f32.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_greater_lesser_equal_f32.c
+  ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_im2row.c
   ${EXECUTORCH_ROOT}/backends/cadence/hifi/third-party/nnlib/xa_nn_elm_logicalxor_bool_bool.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_mul_f32_broadcast.c

backends/cadence/hifi/kernels/kernels.h

@@ -196,6 +196,28 @@ extern "C" WORD32 xa_nn_elm_where_broadcast_4D_f32xf32_f32(
     const unsigned char* __restrict__ p_condition,
     const WORD32* const p_condition_shape);
 
+extern "C" WORD32 xa_nn_im2row_quantized(
+    const WORD8* __restrict__ data_im,
+    const WORD32 in_zero_point,
+    /* input parameters*/
+    const WORD32 channels,
+    const WORD32 height,
+    const WORD32 width,
+    /* output parameters */
+    const WORD32 out_height,
+    const WORD32 out_width,
+    /* convolution parameters */
+    const WORD32 kernel_h,
+    const WORD32 kernel_w,
+    const WORD32 pad_h,
+    const WORD32 pad_w,
+    const WORD32 stride_h,
+    const WORD32 stride_w,
+    const WORD32 dilation_h,
+    const WORD32 dilation_w,
+    WORD8* __restrict__ data_col,
+    WORD32 channels_last);
+
 extern "C" WORD32 xa_nn_reduce_mean_4D_f32_f32(
     FLOAT32* __restrict__ p_out,
     const WORD32* const p_out_shape,

backends/cadence/hifi/operators/CMakeLists.txt

@@ -16,6 +16,7 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake)
 
 # ATen compliant ops that are needed to run this model.
 set(_aten_ops__srcs
+    "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/im2row_out.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_add.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_atan2.cpp"
     "${EXECUTORCH_ROOT}/backends/cadence/hifi/operators/op_bitwise_and.cpp"