chromium/third_party/xnnpack/src/src/operators/reduce-nd.c

// Copyright 2023 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "xnnpack.h"
#include "xnnpack/allocator.h"
#include "xnnpack/common.h"
#include "xnnpack/compute.h"
#include "xnnpack/config-types.h"
#include "xnnpack/config.h"
#include "xnnpack/log.h"
#include "xnnpack/microkernel-type.h"
#include "xnnpack/microparams.h"
#include "xnnpack/normalization.h"
#include "xnnpack/operator-type.h"
#include "xnnpack/operator.h"
#include "xnnpack/params.h"
#include "pthreadpool.h"

static enum xnn_status create_mean_nd(
    uint32_t flags,
    uint32_t log2_element_size,
    enum xnn_operator_type operator_type,
    const struct xnn_reduce_config* rdsum_config,
    const struct xnn_reduce_config* rsum_config,
    const struct xnn_unary_elementwise_config* cvt_config,
    const void* params,
    size_t params_size,
    xnn_operator_t* mean_op_out)
{ … }

enum xnn_status xnn_create_mean_nd_f16(
  uint32_t flags,
  xnn_operator_t* mean_op_out)
{ … }

enum xnn_status xnn_create_mean_nd_f32(
  uint32_t flags,
  xnn_operator_t* mean_op_out)
{ … }

static enum xnn_status reshape_mean_nd(
    xnn_operator_t mean_op,
    size_t num_reduction_axes,
    const size_t* reduction_axes,
    size_t num_input_dims,
    const size_t* input_shape,
    size_t* workspace_size,
    size_t* workspace_alignment,
    size_t log2_data_element_size,
    size_t log2_accumulator_element_size,
    enum xnn_operator_type expected_operator_type,
    const void* scale_params,
    size_t scale_params_size,
    const void* cvt_params,
    size_t cvt_params_size,
    void (*update_params)(xnn_operator_t, size_t),
    pthreadpool_t threadpool)
{ … }

static void update_params_mean_f16(
  xnn_operator_t mean_op,
  size_t num_elements)
{ … }

enum xnn_status xnn_reshape_mean_nd_f16(
    xnn_operator_t mean_op,
    size_t num_reduction_axes,
    const size_t* reduction_axes,
    size_t num_input_dims,
    const size_t* input_shape,
    size_t* workspace_size,
    size_t* workspace_alignment,
    pthreadpool_t threadpool)
{ … }

static void update_params_mean_f32(
  xnn_operator_t mean_op,
  size_t num_elements)
{ … }

enum xnn_status xnn_reshape_mean_nd_f32(
    xnn_operator_t mean_op,
    size_t num_reduction_axes,
    const size_t* reduction_axes,
    size_t num_input_dims,
    const size_t* input_shape,
    pthreadpool_t threadpool)
{ … }

static enum xnn_status setup_mean_nd(
    xnn_operator_t mean_op,
    void* workspace,
    const float* input,
    float* output,
    enum xnn_operator_type expected_operator_type)
{ … }

enum xnn_status xnn_setup_mean_nd_f16(
    xnn_operator_t mean_op,
    void* workspace,
    const void* input,
    void* output)
{ … }

enum xnn_status xnn_setup_mean_nd_f32(
    xnn_operator_t mean_op,
    const float* input,
    float* output)
{ … }