nnfw::cker::functor::ReduceOuterDimensions< Device, InputT, AccumT, OutputT, BinaryFunctor > Struct Template Reference

#include <redux_functor.h>

Public Member Functions
	ReduceOuterDimensions ()
template<int num_dims>
void	operator() (const Device &device, const Eigen::DSizes< Eigen::Index, num_dims > &input_dims, const Tensor &input, Tensor *output) const

Detailed Description

template<typename Device, typename InputT, typename AccumT, typename OutputT, typename BinaryFunctor>
struct nnfw::cker::functor::ReduceOuterDimensions< Device, InputT, AccumT, OutputT, BinaryFunctor >

Definition at line 36 of file redux_functor.h.

Constructor & Destructor Documentation

ReduceOuterDimensions()

template<typename Device , typename InputT , typename AccumT , typename OutputT , typename BinaryFunctor >

nnfw::cker::functor::ReduceOuterDimensions< Device, InputT, AccumT, OutputT, BinaryFunctor >::ReduceOuterDimensions ( )

inline

Definition at line 38 of file redux_functor.h.

{}

Member Function Documentation

operator()()

template<typename Device , typename InputT , typename AccumT , typename OutputT , typename BinaryFunctor >

template<int num_dims>

void nnfw::cker::functor::ReduceOuterDimensions< Device, InputT, AccumT, OutputT, BinaryFunctor >::operator() ( const Device &  device, const Eigen::DSizes< Eigen::Index, num_dims > &  input_dims, const Tensor &  input, Tensor *  output  ) const

inline

Definition at line 41 of file redux_functor.h.

  {
    // Compute inner and outer dim after reshaping into 2d tensor.
    const int num_output_dims = output->shape.DimensionsCount();
    auto output_dims = output->template flat<OutputT>().dimensions();
 
    Eigen::Index inner_dim = 1, outer_dim = 1;
    for (int i = 0; i < num_dims - num_output_dims; ++i)
      outer_dim *= input_dims[i];
    for (int i = num_dims - num_output_dims; i < num_dims; ++i)
      inner_dim *= input_dims[i];
 
    if (1 == outer_dim)
    {
      // Nothing to do but passing input to output.
      output->template flat<OutputT>() =
        input.template flat<InputT>().template cast<OutputT>().reshape(output_dims);
      return;
    }
 
    // Get device thread num.
    const Eigen::Index num_threads = device.numThreads();
 
    // If the inner dim parallelism is large enough
    // TODO(ezhulenev): There seems to be no benefits in going this route. Check
    // if this can be improved, or use better heuristic?
    if (inner_dim > num_threads * 32)
    {
      // Do not create more blocks than there are threads in a pool.
      const Eigen::Index num_blocks = num_threads;
 
      // Block size along the outer dimension.
      const Eigen::Index inner_block_size = Eigen::divup(inner_dim, num_blocks);
      const InputT *input_data = input.template flat<InputT>().data();
 
      // Allocate temporary buffer for partial reductions.
      Eigen::Tensor<AccumT, 1, Eigen::RowMajor, Eigen::Index> buffer({inner_dim});
      buffer.setZero();
      AccumT *buffer_data = buffer.data();
 
      using Buffer =
        Eigen::TensorMap<Eigen::Tensor<AccumT, 1, Eigen::RowMajor, Eigen::Index>, Eigen::Unaligned>;
 
      using Input = Eigen::TensorMap<Eigen::Tensor<const InputT, 1, Eigen::RowMajor, Eigen::Index>,
                                     Eigen::Unaligned>;
 
      const auto compute = [inner_dim, outer_dim, inner_block_size, input_data,
                            buffer_data](Eigen::Index start, Eigen::Index limit) -> void {
        Eigen::Index inner_dim_start = start * inner_block_size;
        Eigen::Index inner_dim_limit = limit * inner_block_size;
        inner_dim_limit = std::min(inner_dim, inner_dim_limit);
        Eigen::Index my_job_len = inner_dim_limit - inner_dim_start;
 
        const InputT *my_job_start = input_data + inner_dim_start;
        Buffer buf(buffer_data + inner_dim_start, my_job_len);
 
        for (Eigen::Index i = 0; i < outer_dim; ++i)
        {
          auto in = Input(my_job_start + i * inner_dim, my_job_len);
          auto cast = in.template cast<AccumT>();
          buf =
            Eigen::TensorCwiseBinaryOp<BinaryFunctor, const decltype(buf), const decltype(cast)>(
              buf, cast);
        }
      };
 
      // Compute cost of reducing a single block.
      const Eigen::Index compute_size = outer_dim * inner_block_size;
      const Eigen::Index compute_input_bytes = compute_size * sizeof(InputT);
      const Eigen::TensorOpCost cost(compute_input_bytes,
                                     0, // We'll be mostly writing to L1, assume store cost is 0
                                     compute_size *
                                       Eigen::internal::functor_traits<BinaryFunctor>::Cost);
 
      device.parallelFor(num_blocks, cost, compute);
 
      // Write final result to the output.
      output->template flat<OutputT>() = buffer.template cast<OutputT>().reshape(output_dims);
    }
    else
    {
      // Compute block size along the outer dimension for efficiency.
      const Eigen::Index parallel_cell_size = inner_dim;
      const Eigen::Index total_workload = outer_dim * inner_dim;
      const Eigen::Index max_parallelism = total_workload / parallel_cell_size;
 
      const Eigen::Index min_block_workload = 2000;
      const Eigen::Index min_block_size = Eigen::divup(min_block_workload, parallel_cell_size);
      const Eigen::Index max_num_blocks =
        std::min(max_parallelism, Eigen::divup(total_workload, min_block_size));
 
      // Do not create more blocks than there are threads in a pool.
      const Eigen::Index num_blocks = std::min(max_num_blocks, num_threads);
 
      // Block size along the outer dimension.
      const Eigen::Index outer_block_size = Eigen::divup(outer_dim, num_blocks);
 
      const InputT *input_data = input.template flat<InputT>().data();
 
      // Allocate temporary buffer for partial reductions.
      std::vector<AccumT> buffer(num_blocks * inner_dim);
      AccumT *buffer_data = buffer.data();
 
      using Buffer =
        Eigen::TensorMap<Eigen::Tensor<AccumT, 1, Eigen::RowMajor, Eigen::Index>, Eigen::Unaligned>;
 
      using Input = Eigen::TensorMap<Eigen::Tensor<const InputT, 1, Eigen::RowMajor, Eigen::Index>,
                                     Eigen::Unaligned>;
 
      const auto compute = [inner_dim, outer_block_size, buffer_data, input_data,
                            outer_dim](Eigen::Index start, Eigen::Index limit) -> void {
        Eigen::Index outer_dim_start = start * outer_block_size;
        Eigen::Index outer_dim_limit = limit * outer_block_size;
        outer_dim_limit = std::min(outer_dim, outer_dim_limit);
 
        Buffer buf(buffer_data + start * inner_dim, inner_dim);
        for (Eigen::Index i = outer_dim_start; i < outer_dim_limit; ++i)
        {
          auto in = Input(input_data + i * inner_dim, inner_dim);
          auto cast = in.template cast<AccumT>();
          buf =
            Eigen::TensorCwiseBinaryOp<BinaryFunctor, const decltype(buf), const decltype(cast)>(
              buf, cast);
        }
      };
 
      // Compute cost of reducing a single block.
      const Eigen::Index compute_size = outer_block_size * inner_dim;
      const Eigen::Index compute_input_bytes = compute_size * sizeof(InputT);
      const Eigen::TensorOpCost cost(compute_input_bytes,
                                     0, // We'll be mostly writing to L1, assume store cost is 0
                                     compute_size *
                                       Eigen::internal::functor_traits<BinaryFunctor>::Cost);
 
      device.parallelFor(num_blocks, cost, compute);
 
      // Aggregate partial results from temporary buffer into first block.
      auto buf0 = Buffer(buffer_data, inner_dim);
      // Just sum the buffer up, as inner dimensions is not large in this case.
      for (int i = 1; i < num_blocks; ++i)
      {
        auto buf = Buffer(buffer_data + i * inner_dim, inner_dim);
        buf0 = Eigen::TensorCwiseBinaryOp<BinaryFunctor, const decltype(buf0), const decltype(buf)>(
          buf0, buf);
      }
      // Write final result to the output.
      output->template flat<OutputT>() = buf0.template cast<OutputT>().reshape(output_dims);
    }
  }

---

The documentation for this struct was generated from the following file:

• compute/cker/include/cker/eigen/redux_functor.h