text: // Copyright 2020 The MACE Authors. All Rights Reserved.
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //     http://www.apache.org/licenses/LICENSE-2.0
  //
  // Unless required by applicable law or agreed to in writing, software
  // distributed under the License is distributed on an "AS IS" BASIS,
  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  // See the License for the specific language governing permissions and
  // limitations under the License.
  
  #ifndef MICRO_OPS_REDUCE_H_
  #define MICRO_OPS_REDUCE_H_
  
  #include "micro/base/logging.h"
  #include "micro/base/types.h"
  #include "micro/base/utils.h"
  #include "micro/framework/operator.h"
  #include "micro/framework/scratch_buffer.h"
  
  namespace micro {
  namespace ops {
  
  class ReduceOpBase : public framework::Operator {
   public:
    MaceStatus OnInit();
  
   public:
    enum ReduceType {
      MEAN = 0,
      MIN = 1,
      MAX = 2,
      PROD = 3,
      SUM = 4,
    };
  
   protected:
    void Validate();
  
   protected:
    ReduceType reduce_type_;
    const int32_t *axis_;
    uint32_t axis_size_;
    bool keep_dims_;
  
    MACE_OP_INPUT_TAGS(INPUT);
    MACE_OP_OUTPUT_TAGS(OUTPUT);
  };
  
  template<typename T>
  class ReduceOp : public ReduceOpBase {
   public:
    MaceStatus OnInit() {
      input_ = GetInputData<T>(INPUT);
      input_dims_ = GetInputShapeDims(INPUT);
      input_dim_size_ = GetInputShapeDimSize(INPUT);
  
      output_ = GetOutputData<T>(OUTPUT);
      return ReduceOpBase::OnInit();
    }
  
    MaceStatus Run() {
      Validate();
      ScratchBuffer scratch_buffer(engine_config_);
      bool *bitmap = scratch_buffer.GetBuffer<bool>(input_dim_size_);
      int32_t *data_dims = scratch_buffer.GetBuffer<int32_t>(input_dim_size_);
      uint32_t data_dim_size = 0;
      int32_t *output_dims = scratch_buffer.GetBuffer<int32_t>(input_dim_size_);
      uint32_t output_dim_size = 0;
      Simplify(output_dims, &output_dim_size, bitmap,
               input_dim_size_, data_dims, &data_dim_size);
      MACE_RETURN_IF_ERROR(
          ResizeOutputShape(OUTPUT, output_dim_size, output_dims));
      const int32_t output_size =
          base::GetShapeSize(output_dim_size, output_dims);
      Compute(data_dims, data_dim_size, static_cast<uint32_t >(output_size));
  
      return MACE_SUCCESS;
    }
  
   private:
    void Simplify(int32_t *output_dims, uint32_t *output_dim_size,
                  bool *bitmap, int32_t bitmap_size,
                  int32_t *data_dims, uint32_t *data_dim_size) {
      base::memset(bitmap, false, bitmap_size);
      if (axis_size_ == 0) {
        for (uint32_t i = 0; i < input_dim_size_; ++i) {
          bitmap[i] = true;
        }
      } else {
        for (uint32_t i = 0; i < axis_size_; ++i) {
          int32_t index = axis_[i] >= 0 ? axis_[i] : axis_[i] + input_dim_size_;
          DataFormat data_format = static_cast<DataFormat>(GetArgByName(
              "data_format", static_cast<int32_t >(NHWC)));
          if (data_format == NCHW &&
              DataTypeToEnum<T>::value != DT_UINT8 && input_dim_size_ == 4) {
            if (index == 1 || index == 2) {
              index = index + 1;
            } else if (index == 3) {
              index = 1;
            }
          }
          bitmap[index] = true;
        }
      }
      uint32_t out_dim_idx = 0;
      for (uint32_t i = 0; i < input_dim_size_; ++i) {
        if (!bitmap[i]) {
          output_dims[out_dim_idx++] = input_dims_[i];
        } else if (keep_dims_) {
          output_dims[out_dim_idx++] = 1;
        }
      }
      *output_dim_size = out_dim_idx;
  
      int32_t data_dims_idx = 0;
      uint32_t dim_index = 0;
      for (; dim_index < input_dim_size_; ++dim_index) {
        if (input_dims_[dim_index] != 1) break;
      }
      if (dim_index >= input_dim_size_) {
        reduce_first_axis_ = true;
      } else {
        reduce_first_axis_ = bitmap[dim_index];
        data_dims[data_dims_idx++] = input_dims_[dim_index];
        ++dim_index;
        for (; dim_index < input_dim_size_; ++dim_index) {
          const int32_t n = input_dims_[dim_index];
          if (n == 1) {
            bitmap[dim_index] = bitmap[dim_index - 1];
          }
          if (bitmap[dim_index - 1] != bitmap[dim_index]) {
            data_dims[data_dims_idx++] = n;
          } else {
            data_dims[data_dims_idx - 1] *= n;
          }
        }
      }
      *data_dim_size = data_dims_idx;
    }
  
    void Reduce1Dims(ReduceType type, int32_t *data_reshape) {
      if (reduce_first_axis_) {
        if (type == MEAN) {
          T tmp = 0;
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            tmp = tmp + input_[i];
          }
          output_[0] = tmp / data_reshape[0];
        } else if (type == MIN) {
          T tmp = input_[0];
          for (int32_t i = 1; i < data_reshape[0]; ++i) {
            tmp = base::min<T>(tmp, input_[i]);
          }
          output_[0] = tmp;
        } else if (type == MAX) {
          T tmp = input_[0];
          for (int32_t i = 1; i < data_reshape[0]; ++i) {
            tmp = base::max<T>(tmp, input_[i]);
          }
          output_[0] = tmp;
        } else if (type == PROD) {
          T tmp = input_[0];
          for (int32_t i = 1; i < data_reshape[0]; ++i) {
            tmp = tmp * input_[i];
          }
          output_[0] = tmp;
        } else if (type == SUM) {
          T tmp = 0;
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            tmp = tmp + input_[i];
          }
          output_[0] = tmp;
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      } else {
        base::memcpy(output_, input_, data_reshape[0] * sizeof(T));
      }
    }
  
    void Reduce2Dims(ReduceType type, int32_t *data_reshape) {
      if (reduce_first_axis_) {
        if (type == MEAN) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = 0;
            for (int32_t j = 0; j < data_reshape[0]; ++j) {
              tmp += input_[j * data_reshape[1] + i];
            }
            output_[i] = tmp / data_reshape[0];
          }
        } else if (type == MIN) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = input_[i];
            for (int32_t j = 1; j < data_reshape[0]; ++j) {
              tmp = base::min(tmp, input_[j * data_reshape[1] + i]);
            }
            output_[i] = tmp;
          }
        } else if (type == MAX) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = input_[i];
            for (int32_t j = 1; j < data_reshape[0]; ++j) {
              tmp = base::max(tmp, input_[j * data_reshape[1] + i]);
            }
            output_[i] = tmp;
          }
        } else if (type == PROD) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = input_[i];
            for (int32_t j = 1; j < data_reshape[0]; ++j) {
              tmp = tmp * input_[j * data_reshape[1] + i];
            }
            output_[i] = tmp;
          }
        } else if (type == SUM) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = 0;
            for (int32_t j = 0; j < data_reshape[0]; ++j) {
              tmp += input_[j * data_reshape[1] + i];
            }
            output_[i] = tmp;
          }
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      } else {
        if (type == MEAN) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            T tmp = 0;
            for (int32_t j = 0; j < data_reshape[1]; ++j) {
              tmp += input_[i * data_reshape[1] + j];
            }
            output_[i] = tmp / data_reshape[1];
          }
        } else if (type == MIN) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            T tmp = input_[i * data_reshape[1]];
            for (int32_t j = 1; j < data_reshape[1]; ++j) {
              tmp = base::min(tmp, input_[i * data_reshape[1] + j]);
            }
            output_[i] = tmp;
          }
        } else if (type == MAX) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            T tmp = input_[i * data_reshape[1]];
            for (int32_t j = 1; j < data_reshape[1]; ++j) {
              tmp = base::max(tmp, input_[i * data_reshape[1] + j]);
            }
            output_[i] = tmp;
          }
        } else if (type == PROD) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            T tmp = input_[i * data_reshape[1]];
            for (int32_t j = 1; j < data_reshape[1]; ++j) {
              tmp = tmp * input_[i * data_reshape[1] + j];
            }
            output_[i] = tmp;
          }
        } else if (type == SUM) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            T tmp = 0;
            for (int32_t j = 0; j < data_reshape[1]; ++j) {
              tmp += input_[i * data_reshape[1] + j];
            }
            output_[i] = tmp;
          }
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      }
    }
  
    void Reduce3Dims(ReduceType type, int32_t *data_reshape) {
      if (reduce_first_axis_) {
        if (type == MEAN) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[0]; ++k) {
                output_[i] +=
                    input_[(k * data_reshape[1] + i) * data_reshape[2]
                        + j];
              }
            }
            output_[i] /= (data_reshape[0] * data_reshape[2]);
          }
        } else if (type == MIN) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = input_[i * data_reshape[2]];
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[0]; ++k) {
                tmp = base::min(
                    tmp, input_[(k * data_reshape[1] + i) * data_reshape[2] + j]);
              }
            }
            output_[i] = tmp;
          }
        } else if (type == MAX) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = input_[i * data_reshape[2]];
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[0]; ++k) {
                tmp = base::max(
                    tmp, input_[(k * data_reshape[1] + i) * data_reshape[2] + j]);
              }
            }
            output_[i] = tmp;
          }
        } else if (type == PROD) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            T tmp = 1;
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[0]; ++k) {
                tmp *= input_[(k * data_reshape[1] + i) * data_reshape[2] + j];
              }
            }
            output_[i] = tmp;
          }
        } else if (type == SUM) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[0]; ++k) {
                output_[i] +=
                    input_[(k * data_reshape[1] + i) * data_reshape[2] + j];
              }
            }
          }
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      } else {
        if (type == MEAN) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                output_[i * data_reshape[2] + j] +=
                    input_[(i * data_reshape[1] + k) * data_reshape[2] + j];
              }
              output_[i * data_reshape[2] + j] /= data_reshape[1];
            }
          }
        } else if (type == MIN) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              T tmp = input_[i * data_reshape[1] * data_reshape[2] + j];
              for (int32_t k = 1; k < data_reshape[1]; ++k) {
                tmp = base::min(
                    tmp, input_[(i * data_reshape[1] + k) * data_reshape[2] + j]);
              }
              output_[i * data_reshape[2] + j] = tmp;
            }
          }
        } else if (type == MAX) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              T tmp = input_[i * data_reshape[1] * data_reshape[2] + j];
              for (int32_t k = 1; k < data_reshape[1]; ++k) {
                tmp = base::max(
                    tmp, input_[(i * data_reshape[1] + k) * data_reshape[2] + j]);
              }
              output_[i * data_reshape[2] + j] = tmp;
            }
          }
        } else if (type == PROD) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              T tmp = input_[i * data_reshape[1] * data_reshape[2] + j];
              for (int32_t k = 1; k < data_reshape[1]; ++k) {
                tmp *= input_[(i * data_reshape[1] + k) * data_reshape[2] + j];
              }
              output_[i * data_reshape[2] + j] = tmp;
            }
          }
        } else if (type == SUM) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                output_[i * data_reshape[2] + j] +=
                    input_[(i * data_reshape[1] + k) * data_reshape[2] + j];
              }
            }
          }
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      }
    }
  
    void Reduce4Dims(ReduceType type, int32_t *data_reshape) {
      if (reduce_first_axis_) {
        if (type == MEAN) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[3]; ++j) {
              for (int32_t k = 0; k < data_reshape[2]; ++k) {
                for (int32_t t = 0; t < data_reshape[0]; ++t) {
                  output_[i * data_reshape[3] + j] +=
                      input_[((t * data_reshape[1] + i) *
                          data_reshape[2] + k) * data_reshape[3] + j];
                }
              }
              output_[i * data_reshape[3] + j] /=
                  (data_reshape[0] * data_reshape[2]);
            }
          }
        } else if (type == MIN) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[3]; ++j) {
              T tmp = input_[i * data_reshape[2] * data_reshape[3] + j];
              for (int32_t k = 0; k < data_reshape[2]; ++k) {
                for (int32_t t = 0; t < data_reshape[0]; ++t) {
                  tmp = base::min(tmp,
                                  input_[((t * data_reshape[1] + i) *
                                      data_reshape[2] + k) * data_reshape[3]
                                      + j]);
                }
              }
              output_[i * data_reshape[3] + j] = tmp;
            }
          }
        } else if (type == MAX) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[3]; ++j) {
              T tmp = input_[i * data_reshape[2] * data_reshape[3] + j];
              for (int32_t k = 0; k < data_reshape[2]; ++k) {
                for (int32_t t = 0; t < data_reshape[0]; ++t) {
                  tmp = base::max(tmp,  // NOLINT
                                  input_[((t * data_reshape[1] + i) *
                                      data_reshape[2] + k) * data_reshape[3]
                                      + j]);
                }
              }
              output_[i * data_reshape[3] + j] = tmp;
            }
          }
        } else if (type == PROD) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[3]; ++j) {
              T tmp = 1;
              for (int32_t k = 0; k < data_reshape[2]; ++k) {
                for (int32_t t = 0; t < data_reshape[0]; ++t) {
                  tmp = tmp * input_[((t * data_reshape[1] + i) *
                      data_reshape[2] + k) * data_reshape[3] + j];
                }
              }
              output_[i * data_reshape[3] + j] = tmp;
            }
          }
        } else if (type == SUM) {
          for (int32_t i = 0; i < data_reshape[1]; ++i) {
            for (int32_t j = 0; j < data_reshape[3]; ++j) {
              for (int32_t k = 0; k < data_reshape[2]; ++k) {
                for (int32_t t = 0; t < data_reshape[0]; ++t) {
                  output_[i * data_reshape[3] + j] +=
                      input_[((t * data_reshape[1] + i) *
                          data_reshape[2] + k) * data_reshape[3] + j];
                }
              }
            }
          }
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      } else {
        if (type == MEAN) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                for (int32_t t = 0; t < data_reshape[3]; ++t) {
                  output_[i * data_reshape[2] + j] +=
                      input_[((i * data_reshape[1] + k) *
                          data_reshape[2] + j) * data_reshape[3] + t];
                }
              }
              output_[i * data_reshape[2] + j] /=
                  (data_reshape[1] * data_reshape[3]);
            }
          }
        } else if (type == MIN) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              T tmp = input_[(i * data_reshape[1] *
                  data_reshape[2] + j) * data_reshape[3]];
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                for (int32_t t = 0; t < data_reshape[3]; ++t) {
                  tmp = base::min(
                      tmp, input_[((i * data_reshape[1] + k) *
                          data_reshape[2] + j) * data_reshape[3] + t]);
                }
              }
              output_[i * data_reshape[2] + j] = tmp;
            }
          }
        } else if (type == MAX) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              T tmp = input_[(i * data_reshape[1] *
                  data_reshape[2] + j) * data_reshape[3]];
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                for (int32_t t = 0; t < data_reshape[3]; ++t) {
                  tmp = base::max(
                      tmp, input_[((i * data_reshape[1] + k) *
                          data_reshape[2] + j) * data_reshape[3] + t]);
                }
              }
              output_[i * data_reshape[2] + j] = tmp;
            }
          }
        } else if (type == PROD) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              T tmp = 1;
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                for (int32_t t = 0; t < data_reshape[3]; ++t) {
                  tmp = tmp * input_[((i * data_reshape[1] + k) *
                      data_reshape[2] + j) * data_reshape[3] + t];
                }
              }
              output_[i * data_reshape[2] + j] = tmp;
            }
          }
        } else if (type == SUM) {
          for (int32_t i = 0; i < data_reshape[0]; ++i) {
            for (int32_t j = 0; j < data_reshape[2]; ++j) {
              for (int32_t k = 0; k < data_reshape[1]; ++k) {
                for (int32_t t = 0; t < data_reshape[3]; ++t) {
                  output_[i * data_reshape[2] + j] +=
                      input_[((i * data_reshape[1] + k) *
                          data_reshape[2] + j) * data_reshape[3] + t];
                }
              }
            }
          }
        } else {
          MACE_NOT_IMPLEMENTED;
        }
      }
    }
  
    void Compute(int32_t *data_reshape,
                 uint32_t data_reshape_size, uint32_t output_size) {
      base::memset(output_, static_cast<T>(0), output_size);
      switch (data_reshape_size) {
        case 1:Reduce1Dims(reduce_type_, data_reshape);
          break;
        case 2:Reduce2Dims(reduce_type_, data_reshape);
          break;
        case 3:Reduce3Dims(reduce_type_, data_reshape);
          break;
        case 4:Reduce4Dims(reduce_type_, data_reshape);
          break;
        default:LOG(FATAL) << "not implemented in mace"
                           << "data reshape size" << data_reshape_size
                           << "reduce first axis:" << reduce_first_axis_;
          break;
      }
    }
  
   private:
    const T *input_;
    const int32_t *input_dims_;
    uint32_t input_dim_size_;
  
    T *output_;
  
    bool reduce_first_axis_;
  };
  
  }  // namespace ops
  }  // namespace micro
  
  
  #endif  // MICRO_OPS_REDUCE_H_
  