/*********************************************************
 *
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 *
 ************************************************************/
#include "pooling.h"

#include <cmath>

namespace singa {

PoolingHandle::PoolingHandle(const Tensor &input,
                             const std::vector<int> &kernel_size,
                             const std::vector<int> &stride,
                             const std::vector<int> &padding,
                             const bool is_max) {
  kernel_h = kernel_size[0];
  kernel_w = kernel_size[1];

  pad_h = padding[0];
  pad_w = padding[1];

  stride_h = stride[0];
  stride_w = stride[1];

  batchsize = input.shape(0);
  channels = input.shape(1);
  height = input.shape(2);
  width = input.shape(3);

  pooled_height = 1;

  if (stride_h > 0)
    pooled_height =
        std::floor(((height + 2 * pad_h - kernel_h) / stride_h)) + 1;
  pooled_width = std::floor(((width + 2 * pad_w - kernel_w) / stride_w)) + 1;
  is_max_pooling = is_max;

#ifdef USE_DNNL
  if (input.device()->lang() == kCpp) {
    auto x_dims =
        dnnl::memory::dims(input.shape().begin(), input.shape().end());
    auto y_dims =
        dnnl::memory::dims({batchsize, channels, pooled_height, pooled_width});
    auto s_dims = dnnl::memory::dims(stride.begin(), stride.end());
    auto k_dims = dnnl::memory::dims(kernel_size.begin(), kernel_size.end());

    auto p_dims = dnnl::memory::dims(padding.begin(), padding.end());

    auto dtype_ = dnnl::memory::data_type::f32;
    auto format_tag_ = get_dnnl_format_tag(input);
    x_md = dnnl::memory::desc({x_dims}, dtype_, format_tag_);
    y_md = dnnl::memory::desc({y_dims}, dtype_, format_tag_);

    // allow max or avg (follow cudnn implementation convention)
    auto pooling_algo = dnnl::algorithm::pooling_avg_exclude_padding;
    if (is_max_pooling) pooling_algo = dnnl::algorithm::pooling_max;

    auto pool_fwd_d = dnnl::pooling_forward::desc(
        dnnl::prop_kind::forward_training, pooling_algo, x_md, y_md, s_dims,
        k_dims, p_dims, p_dims);
    auto pool_bwd_d = dnnl::pooling_backward::desc(
        pooling_algo, x_md, y_md, s_dims, k_dims, p_dims, p_dims);

    auto eng = input.device()->context(0)->dnnl_engine;
    pool_fwd_pd = dnnl::pooling_forward::primitive_desc(pool_fwd_d, eng);
    pool_bwd_pd =
        dnnl::pooling_backward::primitive_desc(pool_bwd_d, eng, pool_fwd_pd);

    auto ws_md = pool_fwd_pd.workspace_desc();
    ws_mem = dnnl::memory(ws_md, eng);
  }
#endif  // USE_DNNL
}

PoolingHandle::~PoolingHandle() {}

#ifdef USE_DNNL

Tensor CpuPoolingForward(const PoolingHandle &ph, const Tensor &x) {
  CHECK_EQ(x.device()->lang(), kCpp);
  Tensor y({(unsigned long)ph.batchsize, (unsigned long)ph.channels,
            (unsigned long)ph.pooled_height, (unsigned long)ph.pooled_width},
           x.device(), x.data_type());

  y.device()->Exec(
      [y, x, &ph](Context *ctx) mutable {
        auto eng = ctx->dnnl_engine;
        using namespace dnnl;

        memory x_mem(ph.x_md, eng, x.block()->mutable_data());
        memory y_mem(ph.y_md, eng, y.block()->mutable_data());

        pooling_forward(ph.pool_fwd_pd)
            .execute(ctx->dnnl_stream, {{DNNL_ARG_SRC, x_mem},
                                        {DNNL_ARG_DST, y_mem},
                                        {DNNL_ARG_WORKSPACE, ph.ws_mem}});
        ctx->dnnl_stream.wait();
      },
      {x.block()}, {y.block()});

  return y;
}

Tensor CpuPoolingBackward(const PoolingHandle &ph, const Tensor &grad,
                          const Tensor &x, const Tensor &y) {
  CHECK_EQ(x.device()->lang(), kCpp);
  CHECK_EQ(grad.device()->lang(), kCpp);
  CHECK_EQ(y.device()->lang(), kCpp);
  Tensor in_grad;
  in_grad.ResetLike(x);

  in_grad.device()->Exec(
      [x, y, in_grad, grad, &ph](Context *ctx) mutable {
        auto eng = ctx->dnnl_engine;
        using namespace dnnl;

        memory dx_mem(ph.x_md, eng, in_grad.block()->mutable_data());
        memory dy_mem(ph.y_md, eng, grad.block()->mutable_data());

        pooling_backward(ph.pool_bwd_pd)
            .execute(ctx->dnnl_stream, {{DNNL_ARG_DIFF_DST, dy_mem},
                                        {DNNL_ARG_DIFF_SRC, dx_mem},
                                        {DNNL_ARG_WORKSPACE, ph.ws_mem}});
        ctx->dnnl_stream.wait();
      },
      {x.block(), y.block(), grad.block()}, {in_grad.block()});

  return in_grad;
}
#endif  // USE_DNNL

#ifdef USE_CUDNN

CudnnPoolingHandle::CudnnPoolingHandle(const Tensor &input,
                                       const std::vector<int> &kernel_size,
                                       const std::vector<int> &stride,
                                       const std::vector<int> &padding,
                                       const bool is_max)
    : PoolingHandle(input, kernel_size, stride, padding, is_max) {
  // nan_prop = CUDNN_NOT_PROPAGATE_NAN;

  DataType dtype = input.data_type();

  CUDNN_CHECK(cudnnCreateTensorDescriptor(&x_desc));
  CUDNN_CHECK(cudnnCreateTensorDescriptor(&y_desc));
  CUDNN_CHECK(cudnnCreatePoolingDescriptor(&pool_desc));

  CUDNN_CHECK(cudnnSetTensor4dDescriptor(x_desc, CUDNN_TENSOR_NCHW,
                                         GetCudnnDataType(dtype), batchsize,
                                         channels, height, width));
  // LOG(ERROR) << batchsize << " " << channels << " " << pooled_height << " "
  // << pooled_width;
  CUDNN_CHECK(cudnnSetTensor4dDescriptor(
      y_desc, CUDNN_TENSOR_NCHW, GetCudnnDataType(dtype), batchsize, channels,
      pooled_height, pooled_width));
  auto pool_method = CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING;
  if (is_max) pool_method = CUDNN_POOLING_MAX;

  CUDNN_CHECK(cudnnSetPooling2dDescriptor(pool_desc, pool_method, nan_prop,
                                          kernel_h, kernel_w, pad_h, pad_w,
                                          stride_h, stride_w));
};

CudnnPoolingHandle::~CudnnPoolingHandle() {
  if (pool_desc != nullptr)
    CUDNN_CHECK(cudnnDestroyPoolingDescriptor(pool_desc));
  if (x_desc != nullptr) CUDNN_CHECK(cudnnDestroyTensorDescriptor(x_desc));
  if (y_desc != nullptr) CUDNN_CHECK(cudnnDestroyTensorDescriptor(y_desc));
}

Tensor GpuPoolingForward(const CudnnPoolingHandle &cph, const Tensor &x) {
  CHECK_EQ(x.device()->lang(), kCuda);
  CHECK_EQ(x.nDim(), 4u);

  Tensor output = Tensor(
      Shape({cph.batchsize, cph.channels, cph.pooled_height, cph.pooled_width}),
      x.device(), x.data_type());

  output.device()->Exec(
      [output, x, &cph](Context *ctx) mutable {
        float alpha = 1.0f, beta = 0.0f;
        cudnnPoolingForward(ctx->cudnn_handle, cph.pool_desc, &alpha,
                            cph.x_desc, x.block()->data(), &beta, cph.y_desc,
                            output.block()->mutable_data());
      },
      {x.block()}, {output.block()});

  return output;
}

Tensor GpuPoolingBackward(const CudnnPoolingHandle &cph, const Tensor &dy,
                          const Tensor &x, const Tensor &y) {
  CHECK_EQ(dy.device()->lang(), kCuda);
  CHECK_EQ(dy.nDim(), 4u);

  Tensor dx;
  dx.ResetLike(x);

  dx.device()->Exec(
      [dx, dy, x, y, &cph](Context *ctx) mutable {
        float alpha = 1.0f, beta = 0.0f;
        cudnnPoolingBackward(ctx->cudnn_handle, cph.pool_desc, &alpha,
                             cph.y_desc, y.block()->data(), cph.y_desc,
                             dy.block()->data(), cph.x_desc, x.block()->data(),
                             &beta, cph.x_desc, dx.block()->mutable_data());
      },
      {dy.block(), y.block(), x.block()}, {dx.block()});

  return dx;
};
#endif  // USE_CUDNN

}  // namespace singa