/*
 * 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 lazy_gradient_init.cc
 *
 * \brief Lazily instantiate 0-filled or 1-filled tensors.
 * This pass should be used after reverse-mode ad so that gradient tensors
 * are not instantiated until after the forward pass.
 * 
 * This pass delays or removes memory allocation by converting tensors into 
 * GradCell, an algebraic data type defined in gradient.rly.
 * 
 * This will delay or decrease memory usage. All calls to
 * ones, ones_like, zeros, zeros_like will call the One or Zero constructor
 * of GradCell, which will not instantiate in memory until needed. All other cases result
 * in using the Raw constructor which means the tensor is instantiated in memory.
 * 
 * It also overloads + and * operation which can increase performance when doing
 * operations involving tensors with values of only 0 or 1.
 * 
 * Note: this pass can only be used with functions where the input/output types are
 * a combination of TupleTypes and TensorTypes
 * 
 * This pass optimizes 6 ops:
 * - add
 * - multiply
 * - ones
 * - ones_like
 * - zeros
 * - zeros_like
 * 
 * This pass makes use of three visitor. The most important one visits the entire function,
 * one is used for wrap inputs and one to unwrap outputs.
 * 
 * For example:
 * fn: TensorType[(10,10), float32] -> TensorType[(10,10), float32]
 * 
 * After this pass
 * fn: GradCell[TensorType[(10,10), float32]] -> GradCell[TensorType[(10,10), float32]]
 * 
 * Thus, it is necessary to wrap this outer function so that the input/output types remain the same
 */

#include <tvm/node/structural_equal.h>
#include <tvm/relay/analysis.h>
#include <tvm/relay/expr_functor.h>
#include <tvm/ir/type_functor.h>
#include <tvm/relay/transform.h>
#include "let_list.h"

namespace tvm {
namespace relay {

/*!
* \brief Visitor appropriately wraps tensors with Raw constructor
*
* Recursively looks at the type of the expression (TensorType or TupleType are only supported for now)
* and either call the GradCell constructor if TensorType
* or unfold and recursively visit if TupleType
*/
class InputVisitor: public ExprFunctor<Expr(const Expr&, const Type&)> {
 public:
  explicit InputVisitor(IRModule module): module_(module) {}

  Expr VisitExpr_(const VarNode* op, const Type& t) final {
    std::cout << op->type_annotation << std::endl;
    return WrapExpr(GetRef<Var>(op), op->type_annotation);
  }

  Expr VisitExpr_(const TupleGetItemNode* op, const Type& t) final {
    return WrapExpr(GetRef<TupleGetItem>(op), t);
  }
 private:
  IRModule module_;

  Expr WrapExpr(const Expr expr, const Type& type) {
    if (type.as<TensorTypeNode>()) {
      return Call(module_->GetConstructor("GradCell", "Raw"),
                  {expr}, Attrs(), {type});
    } else if (auto* type_anno = type.as<TupleTypeNode>()) {
      tvm::Array<Expr> fields;
      for (size_t i = 0; i < type_anno->fields.size(); i++) {
        const Type& t = type_anno->fields[i];
        fields.push_back(this->VisitExpr(TupleGetItem(expr, i), t));
      }
      Expr tuple = Tuple(fields);
      return tuple;
    }

    return expr;
  }
};

/*!
* \brief Visitor appropriately unwraps expressions with GradCell type into Tensors
*
* Recursively looks at the type of the expression
* and either use the FromGradCell function if TypeCall to GradCell
* or unfold and recursively visit if TupleType
*/
class OutputVisitor: public ExprFunctor<Expr(const Expr&, const Type&)> {
 public:
  explicit OutputVisitor(IRModule module): module_(module) {}

  Expr VisitExpr_(const CallNode* op, const Type& t) final {
    return UnwrapExpr(GetRef<Call>(op), t);
  }

  Expr VisitExpr_(const TupleGetItemNode* op, const Type& t) final {
    return UnwrapExpr(GetRef<TupleGetItem>(op), t);
  }
 private:
  IRModule module_;

  Expr UnwrapExpr(const Expr expr, const Type& type) {
    if (auto* type_call = type.as<TypeCallNode>()) {
      if (type_call->func.same_as(module_->GetGlobalTypeVar("GradCell"))) {
        return Call(module_->GetGlobalVar("FromGradCell"), {expr});
      }
      return expr;
    } else if (auto* type_anno = type.as<TupleTypeNode>()) {
      tvm::Array<Expr> fields;
      for (size_t i = 0; i < type_anno->fields.size(); i++) {
        const Type& t = type_anno->fields[i];
        fields.push_back(this->VisitExpr(TupleGetItem(expr, i), t));
      }
      Expr tuple = Tuple(fields);
      return tuple;
    }

    return expr;
  }
};

class LazyGradientInitializer: public ExprMutator, public TypeMutator {
 public:
  explicit LazyGradientInitializer(IRModule module):
    module_(module) {
      module_->ImportFromStd("gradient.rly");
    }

  /*!
  * \brief apply LazyGradientInit transformation and wrap function
  * so that function type stays the same
  * 
  * input/output types should only be a combination of TupleTypes and TensorTypes
  */
  Expr Transform(const Expr& e) {
    auto* f = (e).as<FunctionNode>();
    auto* transformed = this->Mutate(e).as<FunctionNode>();

    if (e.same_as(GetRef<Function>(transformed))) {
      return GetRef<Function>(transformed);
    }

    // wrap inputs of Tensor type using InputVisitor class
    tvm::Array<Expr> args;
    for (Var var : f->params) {
      Expr wrappedInput = InputVisitor(module_).VisitExpr(var, var->checked_type());
      args.push_back(wrappedInput);
    }
    Expr transformedExpr = Call(GetRef<Function>(transformed), args);

    // unwrap outputs of GradCell type into Tensor type using OutputVisitor class
    Expr tensorOutput = OutputVisitor(module_).VisitExpr(transformedExpr, transformed->ret_type);
    return Function(f->params, tensorOutput, f->ret_type, Array<TypeVar>());
  }

  Expr VisitExpr_(const ConstantNode* op) final {
    return Call(module_->GetConstructor("GradCell", "Raw"),
                {GetRef<Constant>(op)}, Attrs(), {op->checked_type()});
  }

  Expr VisitExpr_(const CallNode* call_node) final {
    if (auto* op = (call_node->op).as<OpNode>()) {
      Expr op_expr = GetRef<Op>(op);

      if (op_expr == Op::Get("add")) {
        return CallGradCellFunction(call_node, module_->GetGlobalVar("AddGradCell"));
      }

      if (op_expr == Op::Get("multiply")) {
        return CallGradCellFunction(call_node, module_->GetGlobalVar("MultiplyGradCell"));
      }

      if (op_expr == Op::Get("ones") || op_expr == Op::Get("zeros")) {
        // fn() -> T, function returns result of the operation
        Expr func = Function({}, {ExprMutator::VisitExpr_(call_node)},
                              {call_node->checked_type()}, {});
        // call appropriate GradCell constructor
        std::string constructor_name = op_expr == Op::Get("ones") ? "One" : "Zero";
        return Call(module_->GetConstructor("GradCell", constructor_name),
                    {func}, Attrs(), {call_node->checked_type()});
      }

      if (op_expr == Op::Get("ones_like") || op_expr == Op::Get("zeros_like")) {
        // ones_like and zeros_like need TensorType input
        Expr result = CallPrimitiveOp(call_node);
        // fn() -> T, function returns result of operation
        Expr func = Function({}, result, {call_node->checked_type()}, Array<TypeVar>());
        // call appropriate GradCell constructor
        std::string constructor_name = op_expr == Op::Get("ones_like") ? "One" : "Zero";
        return Call(module_->GetConstructor("GradCell", "One"),
                    {func}, Attrs(), {call_node->checked_type()});
      }

      // handle all other ops
      Expr result = CallPrimitiveOp(call_node);
      // wrap result with Raw constructor
      return Call(module_->GetConstructor("GradCell", "Raw"), {result},
                  Attrs(), {call_node->checked_type()});
    }
    // not an op
    return ExprMutator::VisitExpr_(call_node);
  }

  Type VisitType(const Type& t) final {
    return TypeMutator::VisitType(t);
  }

  Type VisitType_(const TensorTypeNode* op) {
    GlobalTypeVar gradCell = module_->GetGlobalTypeVar("GradCell");
    tvm::Array<Type> args;
    args.push_back(GetRef<TensorType>(op));
    return TypeCall(gradCell, args);
  }

 private:
  // Module
  IRModule module_;

  /*!
  * \brief Convert call_node to add/multiply op to use overloaded functions for GradCell type
  */
  Expr CallGradCellFunction(const CallNode* call_node, GlobalVar overloaded_op) {
    // can only use overloaded functions if 2 arguments of same type
    if (call_node->args.size() != 2 ||
        !tvm::StructuralEqual()(call_node->args[0]->checked_type(),
                                call_node->args[1]->checked_type())) {
      Expr result = CallPrimitiveOp(call_node);
      return Call(module_->GetConstructor("GradCell", "Raw"), {result},
                  Attrs(), {call_node->checked_type()});
    }

    tvm::Array<Expr> args;
    // create "fallback" function for overloaded function
    Type paramType = call_node->args[0]->checked_type();
    tvm::Array<Var> params = {Var("lhs", paramType),
                              Var("rhs", paramType)};
    // use primitive op in this case
    Expr callOp = Call(call_node->op, {params[0], params[1]});
    Expr func = Function(params, callOp, paramType, Array<TypeVar>());

    // pass "fallback" function and tensors as arguments
    args.push_back(func);
    for (Expr expr : call_node->args) {
      args.push_back(VisitExpr(expr));
    }
    // return new call to overloaded function
    return Call(overloaded_op, args, Attrs(), {paramType});
  }

  /*!
  * \brief Convert calls to other ops by converting args into TensorType
  * \return call expr returning result of op
  */
  Expr CallPrimitiveOp(const CallNode* call_node) {
    const auto fromFunc = module_->GetGlobalVar("FromGradCell");
    tvm::Array<Expr> args;
    // use FromGradCell to convert args to Tensor
    for (Expr expr : call_node->args) {
      args.push_back(Call(fromFunc,
                          {VisitExpr(expr)}, Attrs(), {expr->checked_type()}));
    }
    // result of operation
    return Call(call_node->op, args);
  }
};

Expr LazyGradientInit(const Expr& e, IRModule mod) {
  return LazyGradientInitializer(mod).Transform(e);
}

namespace transform {
Pass LazyGradientInit() {
  runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func =
    [=](Function f, IRModule m, PassContext pc) {
      return Downcast<Function>(LazyGradientInit(f, m));
    };
    return CreateFunctionPass(pass_func, 2, "LazyGradientInit", {});
}

TVM_REGISTER_GLOBAL("relay._transform.LazyGradientInit")
.set_body_typed(LazyGradientInit);

}  // namespace transform

}  // namespace relay
}  // namespace tvm
