/*
 * 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.
 */

/*!
 * \file grad.cc
 * \brief Operators to implement operaor gradients.
 */

#include "grad.h"

#include <tvm/ffi/extra/visit_error_context.h>
#include <tvm/ffi/reflection/registry.h>

#include <utility>

namespace tvm {
namespace relax {

/* relax.grad.no_grad */
Expr no_grad(Expr input) {
  static const Op& op = Op::Get("relax.grad.no_grad");
  return Call(Type::Missing(), op, {std::move(input)}, {}, {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.no_grad", no_grad);
}

Type InferTypeNoGrad(const Call& call, const BlockBuilder& ctx) { return GetType(call->args[0]); }

TVM_REGISTER_OP("relax.grad.no_grad")
    .set_num_inputs(1)
    .add_argument("x", "Expr", "The corresponding input tensor.")
    .set_attr<FInferType>("FInferType", InferTypeNoGrad)
    .set_attr<bool>("FPurity", true);

/* relax.grad.start_checkpoint */
Expr start_checkpoint(Expr input) {
  static const Op& op = Op::Get("relax.grad.start_checkpoint");
  return Call(Type::Missing(), op, {std::move(input)}, {}, {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.start_checkpoint", start_checkpoint);
}

Type InferTypeStartCheckpoint(const Call& call, const BlockBuilder& ctx) {
  if (!call->args[0].as<VarNode>()) {
    TVM_FFI_VISIT_THROW(TypeError, call)
        << "The argument of relax.op.grad.start_checkpoint should be a Var.";
  }
  return GetType(call->args[0]);
}

TVM_REGISTER_OP("relax.grad.start_checkpoint")
    .set_num_inputs(1)
    .add_argument("x", "Expr", "The tensor marking the input of the checkpoint stage.")
    .set_attr<FInferType>("FInferType", InferTypeStartCheckpoint)
    .set_attr<bool>("FPurity", true);

/* relax.grad.end_checkpoint */
Expr end_checkpoint(Expr input) {
  static const Op& op = Op::Get("relax.grad.end_checkpoint");
  return Call(Type::Missing(), op, {std::move(input)}, {}, {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.end_checkpoint", end_checkpoint);
}

Type InferTypeEndCheckpoint(const Call& call, const BlockBuilder& ctx) {
  if (!call->args[0].as<VarNode>()) {
    TVM_FFI_VISIT_THROW(TypeError, call)
        << "The argument of relax.op.grad.end_checkpoint should be a Var.";
  }
  return GetType(call->args[0]);
}

TVM_REGISTER_OP("relax.grad.end_checkpoint")
    .set_num_inputs(1)
    .add_argument("x", "Expr", "The output of the checkpoint stage.")
    .set_attr<FInferType>("FInferType", InferTypeEndCheckpoint)
    .set_attr<bool>("FPurity", true);

/* relax.grad.nll_loss_backward */
Expr nll_loss_backward(Expr output_grad, Expr predictions, Expr targets,
                       ffi::Optional<Expr> weights, ffi::String reduction, int ignore_index) {
  ffi::ObjectPtr<NLLLossAttrs> attrs = ffi::make_object<NLLLossAttrs>();

  attrs->reduction = reduction;
  attrs->ignore_index = ignore_index;

  static const Op& op = Op::Get("relax.grad.nll_loss_backward");
  if (weights.has_value()) {
    return Call(
        Type::Missing(), op,
        {std::move(output_grad), std::move(predictions), std::move(targets), weights.value()},
        Attrs{attrs}, {});
  } else {
    return Call(Type::Missing(), op,
                {std::move(output_grad), std::move(predictions), std::move(targets)}, Attrs{attrs},
                {});
  }
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.nll_loss_backward", nll_loss_backward);
}

Type InferTypeNLLLossBackward(const Call& call, const BlockBuilder& ctx) {
  return GetType(call->args[1]);
}

TVM_REGISTER_OP("relax.grad.nll_loss_backward")
    .set_attrs_type<NLLLossAttrs>()
    .set_num_inputs(4)
    .add_argument("output_grad", "Tensor", "The output gradient.")
    .add_argument("predictions", "Tensor", "The prediction tensor.")
    .add_argument("targets", "Tensor", "The target tensor.")
    .add_argument("weights", "ffi::Optional<Tensor>", "The weight of each target values.")
    .set_attr<FInferType>("FInferType", InferTypeNLLLossBackward)
    .set_attr<bool>("FPurity", true);

/* relax.grad.max_pool2d_backward */
Expr max_pool2d_backward(Expr output_grad, Expr data, ffi::Array<int64_t> pool_size,
                         ffi::Array<int64_t> strides, ffi::Array<int64_t> padding,
                         ffi::Array<int64_t> dilation, bool ceil_mode, bool count_include_pad,
                         ffi::String layout, ffi::Optional<ffi::String> out_layout) {
  auto attrs = ffi::make_object<Pool2DAttrs>();
  attrs->pool_size = std::move(pool_size);
  attrs->strides = std::move(strides);
  attrs->padding = std::move(padding);
  attrs->dilation = std::move(dilation);
  attrs->ceil_mode = ceil_mode;
  attrs->count_include_pad = count_include_pad;
  attrs->layout = layout;
  attrs->out_layout = out_layout.value_or(layout);
  static const Op& op = Op::Get("relax.grad.max_pool2d_backward");
  return Call(Type::Missing(), op, {std::move(output_grad), std::move(data)}, Attrs(attrs), {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.max_pool2d_backward", max_pool2d_backward);
}

Type InferTypeMaxPool2DBackward(const Call& call, const BlockBuilder& ctx) {
  return GetType(call->args[1]);
}

TVM_REGISTER_OP("relax.grad.max_pool2d_backward")
    .set_num_inputs(2)
    .add_argument("output_grad", "Tensor", "The output gradient.")
    .add_argument("data", "Tensor", "The input tensor")
    .set_attrs_type<Pool2DAttrs>()
    .set_attr<FInferType>("FInferType", InferTypeMaxPool2DBackward)
    .set_attr<bool>("FPurity", true);

/* relax.grad.avg_pool2d_backward */
Expr avg_pool2d_backward(Expr output_grad, Expr data, ffi::Array<int64_t> pool_size,
                         ffi::Array<int64_t> strides, ffi::Array<int64_t> padding,
                         ffi::Array<int64_t> dilation, bool ceil_mode, bool count_include_pad,
                         ffi::String layout, ffi::Optional<ffi::String> out_layout) {
  auto attrs = ffi::make_object<Pool2DAttrs>();
  attrs->pool_size = std::move(pool_size);
  attrs->strides = std::move(strides);
  attrs->padding = std::move(padding);
  attrs->dilation = std::move(dilation);
  attrs->ceil_mode = ceil_mode;
  attrs->count_include_pad = count_include_pad;
  attrs->layout = layout;
  attrs->out_layout = out_layout.value_or(layout);
  static const Op& op = Op::Get("relax.grad.avg_pool2d_backward");
  return Call(Type::Missing(), op, {std::move(output_grad), std::move(data)}, Attrs(attrs), {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.avg_pool2d_backward", avg_pool2d_backward);
}

Type InferTypeAvgPool2DBackward(const Call& call, const BlockBuilder& ctx) {
  return GetType(call->args[1]);
}

TVM_REGISTER_OP("relax.grad.avg_pool2d_backward")
    .set_num_inputs(2)
    .add_argument("output_grad", "Tensor", "The output gradient.")
    .add_argument("data", "Tensor", "The input tensor")
    .set_attrs_type<Pool2DAttrs>()
    .set_attr<FInferType>("FInferType", InferTypeAvgPool2DBackward)
    .set_attr<bool>("FPurity", true);

/* relax.grad.take_backward */

Expr take_backward(Expr output_grad, Expr x, Expr indices, ffi::Optional<int64_t> axis) {
  ffi::ObjectPtr<TakeAttrs> attrs = ffi::make_object<TakeAttrs>();
  attrs->axis = std::move(axis);

  static const Op& op = Op::Get("relax.grad.take_backward");
  return Call(Type::Missing(), op, {std::move(output_grad), std::move(x), std::move(indices)},
              Attrs(attrs), {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("relax.op.grad.take_backward", take_backward);
}

Type InferTypeTakeBackward(const Call& call, const BlockBuilder& ctx) {
  return GetType(call->args[1]);
}

TVM_REGISTER_OP("relax.grad.take_backward")
    .set_attrs_type<TakeAttrs>()
    .set_num_inputs(3)
    .add_argument("output_grad", "Tensor", "The output gradient.")
    .add_argument("x", "Tensor", "The source tensor.")
    .add_argument("indices", "Tensor", "The indices of the values to extract.")
    .set_attr<FInferType>("FInferType", InferTypeTakeBackward)
    .set_attr<bool>("FPurity", true);

}  // namespace relax
}  // namespace tvm
