#include "caffe2/core/common.h"

#ifdef CAFFE2_USE_MKL
#include <mkl.h>
#endif

#include "caffe2/core/context.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/operator.h"
#include "caffe2/operators/conv_pool_op_base.h"
#include "caffe2/operators/leaky_relu_op.h"
#include "caffe2/utils/cpuid.h"
#include "caffe2/utils/math.h"
#include "nnpack.h"

C10_DEFINE_int(
    caffe2_nnpack_num_threads,
    1,
    "The number of nnpack pthreadpool threads.");
C10_DEFINE_bool(
    caffe2_nnpack_use_mkl_num_threads,
    true,
    "If MKL is built, this sets nnpack to use the same number of threads as "
    "MKL does. This overrides caffe2_nnpack_num_threads if set.");

namespace caffe2 {
////////////////////////////////////////////////////////////////////////////////
// Helper Functions
////////////////////////////////////////////////////////////////////////////////

namespace {

bool has_nnpack() {
  // nnp_initialize is a noop after the first call so it's safe to invoke it
  // repeatedly
  auto nnpack_status = nnp_initialize();
  return nnpack_status == nnp_status_success;
}

nnp_convolution_algorithm get_nnp_convolution_algorithm(
    const std::string& algo) {
  if (algo == "AUTO") {
    return nnp_convolution_algorithm_auto;
  }
  if (algo == "WINOGRAD") {
    return nnp_convolution_algorithm_wt8x8;
  }
  if (algo == "FT16") {
    return nnp_convolution_algorithm_ft16x16;
  }
  if (algo == "FT8") {
    return nnp_convolution_algorithm_ft8x8;
  }
  return nnp_convolution_algorithm_auto;
}

nnp_convolution_transform_strategy get_nnp_convolution_transform_strategy(
    const std::string& kts) {
  if (kts == "BLOCK") {
    return nnp_convolution_transform_strategy_block_based;
  }
  if (kts == "TUPLE") {
    return nnp_convolution_transform_strategy_tuple_based;
  }
  return nnp_convolution_transform_strategy_block_based;
}

////////////////////////////////////////////////////////////////////////////////
// Thread Pool
////////////////////////////////////////////////////////////////////////////////

static pthreadpool_t nnpack_threadpool_ = nullptr;

pthreadpool_t nnpack_threadpool() {
  if (nnpack_threadpool_ == nullptr) {
    enum nnp_status nnpack_status = nnp_initialize();
    CAFFE_ENFORCE(
        nnpack_status == nnp_status_success, "NNPack is not supported here!");
    int num_threads = FLAGS_caffe2_nnpack_num_threads;
    if (FLAGS_caffe2_nnpack_use_mkl_num_threads) {
#ifdef CAFFE2_USE_MKL
      num_threads = mkl_get_max_threads();
#else
      VLOG(1) << "I am asked to use MKL num of threads for NNPACK but this "
                 "Caffe2 is not built with MKL. Skipping.";
#endif
    }
    nnpack_threadpool_ = pthreadpool_create(num_threads);
  }
  return nnpack_threadpool_;
}
}

////////////////////////////////////////////////////////////////////////////////
// NNPACK Ops
////////////////////////////////////////////////////////////////////////////////

class NNPACKConvOp final : public ConvPoolOpBase<CPUContext> {
 public:
  NNPACKConvOp(const OperatorDef& operator_def, Workspace* ws)
      : ConvPoolOpBase<CPUContext>(operator_def, ws),
        algo_(get_nnp_convolution_algorithm(
            OperatorBase::GetSingleArgument<std::string>("algo", "AUTO"))),
        kts_(get_nnp_convolution_transform_strategy(
            OperatorBase::GetSingleArgument<std::string>("kts", "TUPLE"))) {
    OPERATOR_NEEDS_FEATURE(
        this->order_ == StorageOrder::NCHW,
        "NNPack only supports NCHW order. Please consider adding "
        "TransposeOp with axes=[0, 3, 1, 2] before NNPack Conv.");
    OPERATOR_NEEDS_FEATURE(
        dilation_h() == 1 && dilation_w() == 1,
        "The NNPack convolution does not support dilation yet.");
    // NNPACK can be built with avx2 support only and might not be able to run
    // on a given machine.
    OPERATOR_NEEDS_FEATURE(has_nnpack(), "NNPack can't run here. No AVX2?");
  }

  bool RunOnDeviceWithOrderNCHW() override {
    auto& X = Input(0);
    auto& filter = Input(1);
    auto& bias = Input(2);
    auto* Y = Output(0);

    const int N = X.dim32(0), C = X.dim32(1), H = X.dim32(2), W = X.dim32(3);
    const int M = filter.dim32(0);

    CAFFE_ENFORCE(X.dim() == 4, "Input dim should be 4");
    CAFFE_ENFORCE(filter.dim(), 4);
    CAFFE_ENFORCE(C % this->group_ == 0, "");
    CAFFE_ENFORCE(M % this->group_ == 0, "");
    CAFFE_ENFORCE(filter.dim32(1) == C / this->group_, "");
    CAFFE_ENFORCE(filter.dim32(2) == this->kernel_h(), "");
    CAFFE_ENFORCE(filter.dim32(3) == this->kernel_w(), "");
    CAFFE_ENFORCE(bias.numel() == M, "");

    ConvPoolOpBase<CPUContext>::SetOutputSize(X, Y, filter.dim32(0));
    const int oH = Y->dim32(2), oW = Y->dim32(3);

    if (N > 1) {
      CAFFE_ENFORCE_EQ(
          this->stride_h(),
          1,
          "NNPack only supports stride = 1 when doing batch feedforward");
      CAFFE_ENFORCE_EQ(
          this->stride_w(),
          1,
          "NNPack only supports stride = 1 when doing batch feedforward");
    }
    std::vector<int> pads(
        {this->pad_t(), this->pad_b(), this->pad_l(), this->pad_r()});
    std::vector<int> stride({this->stride_h(), this->stride_w()});

    const size_t input_channels = X.dim32(1);
    const size_t output_channels = Y->dim32(1);
    const nnp_size input_size = {.width = static_cast<size_t>(X.dim32(3)),
                                 .height = static_cast<size_t>(X.dim32(2))};
    // filter is MCHW
    const nnp_size kernel_size = {
        .width = static_cast<size_t>(filter.dim32(3)),
        .height = static_cast<size_t>(filter.dim32(2))};
    // pad is tblr
    const nnp_padding padding = {.top = static_cast<size_t>(pads[0]),
                                 .right = static_cast<size_t>(pads[3]),
                                 .bottom = static_cast<size_t>(pads[1]),
                                 .left = static_cast<size_t>(pads[2])};

    const nnp_size output_subsample = {
        .width = static_cast<size_t>(stride[1]),
        .height = static_cast<size_t>(stride[0])};
    if (N == 1) {
      VLOG(1) << "Running inference mode";
      for (auto g = 0; g < group_; ++g) {
        const auto status = nnp_convolution_inference(
            algo_,
            kts_,
            C / group_,
            M / group_,
            input_size,
            padding,
            kernel_size,
            output_subsample,
            X.template data<float>() + g * H * W * (C / group_),
            filter.template data<float>() + filter.numel() / group_ * g,
            bias.template data<float>() + bias.numel() / group_ * g,
            Y->template mutable_data<float>() + g * oH * oW * (M / group_),
            nnpack_threadpool(),
            nullptr);
        CAFFE_ENFORCE(nnp_status_success == status, "");
      }
    } else {
      VLOG(1) << "Running batched mode";
      for (auto g = 0; g < group_; ++g) {
        const auto status = nnp_convolution_output(
            algo_,
            N,
            C / group_,
            M / group_,
            input_size,
            padding,
            kernel_size,
            X.template data<float>() + g * H * W * (C / group_),
            filter.template data<float>() + filter.numel() / group_ * g,
            bias.template data<float>() + bias.numel() / group_ * g,
            Y->template mutable_data<float>() + g * oH * oW * (M / group_),
            nnpack_threadpool(),
            nullptr);
        CAFFE_ENFORCE(nnp_status_success == status, "");
      }
    }
    return true;
  }

 private:
  const nnp_convolution_algorithm algo_;
  const nnp_convolution_transform_strategy kts_;
};

class NNPACKMaxPoolOp final : public ConvPoolOpBase<CPUContext> {
 public:
  NNPACKMaxPoolOp(const OperatorDef& operator_def, Workspace* ws)
      : ConvPoolOpBase<CPUContext>(operator_def, ws) {
    OPERATOR_NEEDS_FEATURE(
        this->order_ == StorageOrder::NCHW,
        "NNPack only supports NCHW order. Please consider add "
        "TransposeOp with axes=[0, 3, 1, 2] before NNPack Conv.");
    OPERATOR_NEEDS_FEATURE(
        this->kernel_h() == 2, "NNPack only supports MaxPool kernel size 2*2!");
    OPERATOR_NEEDS_FEATURE(
        this->kernel_w() == 2, "NNPack only supports MaxPool kernel size 2*2!");
    OPERATOR_NEEDS_FEATURE(
        this->stride_h() == 2, "NNPack only supports MaxPool stride size 2*2!");
    OPERATOR_NEEDS_FEATURE(
        this->stride_w() == 2, "NNPack only supports MaxPool stride size 2*2!");
    OPERATOR_NEEDS_FEATURE(
        this->pad_t() == 0,
        "NNPack Pooling differs from Caffe2 Pooling when pad > 0!");
    OPERATOR_NEEDS_FEATURE(
        this->pad_l() == 0,
        "NNPack Pooling differs from Caffe2 Pooling when pad > 0!");
    OPERATOR_NEEDS_FEATURE(
        this->pad_r() == 0,
        "NNPack Pooling differs from Caffe2 Pooling when pad > 0!");
    OPERATOR_NEEDS_FEATURE(
        this->pad_b() == 0,
        "NNPack Pooling differs from Caffe2 Pooling when pad > 0!");
    // NNPACK can be built with avx2 support only and might not be able to run
    // on a given machine.
    OPERATOR_NEEDS_FEATURE(has_nnpack(), "NNPack can't run here. No AVX2?");
  }

  bool RunOnDeviceWithOrderNCHW() override {
    auto& X = Input(0);
    auto* Y = Output(0);
    CAFFE_ENFORCE(X.dim() == 4, "");
    const int H = X.dim32(2), W = X.dim32(3);
    ConvPoolOpBase<CPUContext>::SetOutputSize(X, Y, X.dim32(1));
    std::vector<int> pads(
        {this->pad_t(), this->pad_b(), this->pad_l(), this->pad_r()});
    std::vector<int> stride({this->stride_h(), this->stride_w()});
    std::vector<int> pooling({this->kernel_h(), this->kernel_w()});

    // Input X is in NCHW order
    const size_t batch_size = X.dim32(0);
    const size_t input_channels = X.dim32(1);
    const nnp_size input_size = {.width = static_cast<size_t>(X.dim32(3)),
                                 .height = static_cast<size_t>(X.dim32(2))};
    // pooling kernel
    const nnp_size pooling_size = {.width = static_cast<size_t>(pooling[1]),
                                   .height = static_cast<size_t>(pooling[0])};
    // pad is tblr
    const nnp_padding padding = {.top = static_cast<size_t>(pads[0]),
                                 .right = static_cast<size_t>(pads[3]),
                                 .bottom = static_cast<size_t>(pads[1]),
                                 .left = static_cast<size_t>(pads[2])};

    const nnp_size pooling_stride = {.width = static_cast<size_t>(stride[1]),
                                     .height = static_cast<size_t>(stride[0])};
    const auto status = nnp_max_pooling_output(
        batch_size,
        input_channels,
        input_size,
        padding,
        pooling_size,
        pooling_stride,
        X.template data<float>(),
        Y->template mutable_data<float>(),
        nnpack_threadpool());
    CAFFE_ENFORCE(nnp_status_success == status, "");
    return true;
  }

 private:
};

class NNPACKReluOp final : public Operator<CPUContext> {
 public:
  NNPACKReluOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator<CPUContext>(operator_def, ws) {
    // NNPACK can be built with avx2 support only and might not be able to run
    // on a given machine.
    OPERATOR_NEEDS_FEATURE(has_nnpack(), "NNPack can't run here. No AVX2?");
  }

  bool RunOnDevice() override {
    auto& X = Input(0);
    auto* Y = Output(0);
    const auto status = nnp_relu_output(
        1,
        X.numel(),
        X.template data<float>(),
        Y->template mutable_data<float>(),
        0.0,
        nnpack_threadpool());
    CAFFE_ENFORCE(nnp_status_success == status, "");
    return true;
  }

 private:
};

class NNPACKLeakyReluOp final : public LeakyReluOp<float, CPUContext> {
 public:
  NNPACKLeakyReluOp(const OperatorDef& operator_def, Workspace* ws)
      : LeakyReluOp<float, CPUContext>(operator_def, ws) {
    // NNPACK can be built with avx2 support only and might not be able to run
    // on a given machine.
    OPERATOR_NEEDS_FEATURE(has_nnpack(), "NNPack can't run here. No AVX2?");
  }

  bool RunOnDevice() override {
    auto& X = Input(0);
    auto* Y = Output(0);
    const auto status = nnp_relu_output(
        1,
        X.numel(),
        X.template data<float>(),
        Y->template mutable_data<float>(),
        alpha_,
        nnpack_threadpool());
    CAFFE_ENFORCE(nnp_status_success == status, "");
    return true;
  }

 private:
};

REGISTER_CPU_OPERATOR_WITH_ENGINE(Conv, NNPACK, NNPACKConvOp);
REGISTER_CPU_OPERATOR_WITH_ENGINE(MaxPool, NNPACK, NNPACKMaxPoolOp);
REGISTER_CPU_OPERATOR_WITH_ENGINE(Relu, NNPACK, NNPACKReluOp);
REGISTER_CPU_OPERATOR_WITH_ENGINE(LeakyRelu, NNPACK, NNPACKLeakyReluOp);

} // namespace caffe2