/**
 * \file src/opr/impl/dnn/adaptive_pooling.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
 */
#include "megbrain/opr/dnn/adaptive_pooling.h"
#include "../internal/megdnn_opr_wrapper.inl"
#include "megbrain/graph/grad_impl.h"
#include "megbrain/opr/utility.h"

#include "megdnn/opr_param_defs.h"
#include "megdnn/oprs/nn.h"

using namespace mgb;
using namespace opr;

MGB_DYN_TYPE_OBJ_FINAL_IMPL(AdaptivePoolingForward);
AdaptivePoolingForward::AdaptivePoolingForward(VarNode* src, VarNode* out_shape,
                                               const Param& param,
                                               const OperatorNodeConfig& config)
        : Super(OperatorNodeBaseCtorParam{src->owner_graph(),
                                          config,
                                          "adaptive_pooling",
                                          {src, out_shape}}) {
    init_megdnn_opr(*this, param);
    add_input({src, out_shape});
    outshape_by_symvar_enable(1, 1);
}

SymbolVar AdaptivePoolingForward::make(SymbolVar src, SymbolVar out_shape,
                                       const Param& param,
                                       const OperatorNodeConfig& config) {
    return src.insert_single_output_opr<AdaptivePoolingForward>(
            src.node(), out_shape.node(), param, config);
}

void AdaptivePoolingForward::scn_do_execute() {
    megdnn_opr()->exec(input(0)->dev_tensor().as_megdnn(),
                       output(0)->dev_tensor().as_megdnn(),
                       intl::get_megdnn_workspace_from_var(output().back()));
}

void AdaptivePoolingForward::outshape_by_symvar_do_get_output_shape(
        TensorShape& dest, const ShapeInferInfo& shpinfo) {
    TensorShape oshp2d;
    cg::copy_tensor_value_to_shape(oshp2d, *shpinfo.shpval_inp_val.at(0));
    auto src = shpinfo.shape_inp_shp.at(0);
    mgb_assert(src.ndim == 4 && oshp2d.ndim == 2,
               "shape mismatch for AdaptivePooling: src=%s, out2d=%s",
               src.to_string().c_str(), oshp2d.to_string().c_str());

    mgb_assert(param().format == Param::Format::NCHW,
               "AdaptivePooling only support NCHW");
    dest.ndim = 4;
    dest.shape[0] = src.shape[0];
    dest.shape[1] = src.shape[1];
    dest.shape[2] = oshp2d.shape[0];
    dest.shape[3] = oshp2d.shape[1];
}

size_t AdaptivePoolingForward::get_workspace_size_bytes(
        const TensorShapeArray& input_shapes,
        const TensorShapeArray& output_shapes) const {
    return megdnn_opr()->get_workspace_in_bytes(
            {input_shapes[0], this->input(0)->dtype(),
             this->input(0)->format()},
            {output_shapes[0], this->output(0)->dtype(),
             this->output(0)->format()});
}

void AdaptivePoolingForward::init_output_dtype() {
    output(0)->dtype(input(0)->dtype());
}

void AdaptivePoolingForward::add_input_layout_constraint() {
    mixin::megdnn_utils::add_input_layout_constraint_contig(*this);
}

void AdaptivePoolingForward::init_output_static_infer_desc() {
    Super::init_output_static_infer_desc();
    init_output_static_infer_desc_workspace(false);
}

void AdaptivePoolingForward::record_execute_deps(ExecDependencyArray& deps) {
    record_megdnn_opr(deps);
}

#if MGB_ENABLE_GRAD
MGB_IMPL_OPR_GRAD(AdaptivePoolingForward) {
    if (wrt_idx == 0) {
        // wrt src
        SymbolVar grad = AdaptivePoolingBackward::make(
                opr.input(0), opr.input(1), opr.output(0), out_grad[0],
                opr.param());
        return grad.node();
    } else {
        mgb_assert(wrt_idx == 1);
        return InvalidGrad::make(opr, wrt_idx);
    }
}
#endif

MGB_DYN_TYPE_OBJ_FINAL_IMPL(AdaptivePoolingBackward);
AdaptivePoolingBackward::AdaptivePoolingBackward(
        VarNode* src, VarNode* out_shape, VarNode* dst, VarNode* diff,
        const Param& param, const OperatorNodeConfig& config)
        : Super(OperatorNodeBaseCtorParam{src->owner_graph(),
                                          config,
                                          "adaptive_pooling_bwd",
                                          {src}},
                0, true) {
    init_megdnn_opr(*this, param);
    add_input({src, out_shape, dst, diff});
}

SymbolVar AdaptivePoolingBackward::make(SymbolVar src, SymbolVar out_shape,
                                        SymbolVar dst, SymbolVar diff,
                                        const Param& param,
                                        const OperatorNodeConfig& config) {
    return src.insert_single_output_opr<AdaptivePoolingBackward>(
            src.node(), out_shape.node(), dst.node(), diff.node(), param,
            config);
}

void AdaptivePoolingBackward::scn_do_execute() {
    megdnn_opr()->exec(input(0)->dev_tensor().as_megdnn(),
                       input(2)->dev_tensor().as_megdnn(),
                       input(3)->dev_tensor().as_megdnn(),
                       output(0)->dev_tensor().as_megdnn(),
                       intl::get_megdnn_workspace_from_var(output().back()));
}
size_t AdaptivePoolingBackward::get_workspace_size_bytes(
        const TensorShapeArray& input_shapes,
        const TensorShapeArray& output_shapes) const {
    return megdnn_opr()->get_workspace_in_bytes(
            {input_shapes[0], input(0)->dtype(), input(0)->format()},
            {input_shapes[2], input(2)->dtype(), input(2)->format()},
            {input_shapes[3], input(3)->dtype(), input(3)->format()},
            {output_shapes[0], output(0)->dtype(), output(0)->format()});
}

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}