/*
 * 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 relax/ir/transform.cc
 * \brief Relax specific transformation passes.
 */
#include <tvm/ffi/cast.h>
#include <tvm/ffi/extra/dataclass.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/rvalue_ref.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type_functor.h>
#include <tvm/runtime/logging.h>

namespace tvm {
namespace relax {
namespace transform {

TVM_REGISTER_PASS_CONFIG_OPTION("relax.fallback_device_type", IntImm);

class FunctionPass;

/*!
 * \brief Function-level passes are used to implement various global
 * optimizations for a given Relax IRModule. It fetches one function at a time
 * from the function list in the IRModule for optimization.
 *
 * Note that the scope of passes at this level is a Relax function. Therefore,
 * we cannot add or delete a function through these passes as they are not aware
 * of the global information.
 */
class FunctionPassNode : public tvm::transform::PassNode {
 public:
  /* \brief The pass meta data.*/
  PassInfo pass_info;

  /*! \brief The packed pass function sketches the real optimization. For
   * instance, we can implement a pass that works on a Relax function as a
   * `pass_func` and let it run on a given IRModule. The same `pass_func` will
   * then be applied on each function in the IRModule.
   */
  std::function<Function(Function, IRModule, PassContext)> pass_func;

  FunctionPassNode() = default;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FunctionPassNode>().def_ro("pass_info", &FunctionPassNode::pass_info);
  }

  /*!
   * \brief Run a function pass on given pass context.
   *
   * \param mod The IRModule that an optimization pass is applied on.
   * \param pass_ctx The context that an optimization pass executes on.
   *
   * \return Return the updated IRModule.
   */
  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;

  /*!
   * \brief Get the pass information/meta data.
   */
  PassInfo Info() const override { return pass_info; }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.FunctionPass", FunctionPassNode, PassNode);

 private:
};

class FunctionPass : public Pass {
 public:
  /*!
   * \brief The constructor
   * \param pass_func The packed function which implements a pass.
   * \param pass_info The pass info.
   */
  TVM_DLL FunctionPass(std::function<Function(Function, IRModule, PassContext)> pass_func,
                       PassInfo pass_info);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(FunctionPass, Pass, FunctionPassNode);
};

FunctionPass::FunctionPass(std::function<Function(Function, IRModule, PassContext)> pass_func,
                           PassInfo pass_info) {
  auto n = ffi::make_object<FunctionPassNode>();
  n->pass_func = std::move(pass_func);
  n->pass_info = std::move(pass_info);
  data_ = std::move(n);
}

// Perform IRModule -> IRModule optimizations at the Function level.
IRModule FunctionPassNode::operator()(IRModule mod, const PassContext& pass_ctx) const {
  const PassInfo& pass_info = Info();

  TVM_FFI_ICHECK(mod.defined());

  VLOG_CONTEXT << pass_info->name;
  VLOG(0) << "Executing function pass with opt level: " << pass_info->opt_level;
  VLOG(1) << "Input module:" << std::endl << mod;

  IRModule updated_mod = mod->ShallowCopy();

  std::vector<std::pair<GlobalVar, Function> > updates;
  for (const auto& it : updated_mod->functions) {
    // only picks up relax::Function
    if (auto* n = it.second.as<FunctionNode>()) {
      Function func = ffi::GetRef<Function>(n);
      // Enrich at this leaf executor, rendering the location local to the
      // currently-processed function (the access path is function-rooted).
      Function updated_func;
      try {
        updated_func = pass_func(func, updated_mod, pass_ctx);
      } catch (ffi::Error& err) {
        throw tvm::transform::EnrichPassErrorWithContext(err, updated_mod, pass_info->name,
                                                         it.first);
      }
      updates.push_back({it.first, updated_func});
    }
  }

  for (const auto& pair : updates) {
    updated_mod->Add(pair.first, pair.second, true);
  }

  VLOG(1) << "Output module:" << std::endl << updated_mod;

  return updated_mod;
}

Pass CreateFunctionPass(std::function<Function(Function, IRModule, PassContext)> pass_func,
                        int opt_level, ffi::String name, tvm::ffi::Array<ffi::String> required,
                        bool traceable) {
  PassInfo pass_info = PassInfo(opt_level, name, required, traceable);
  return FunctionPass(std::move(pass_func), pass_info);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def(
      "relax.transform.MakeFunctionPass",
      [](ffi::TypedFunction<Function(ffi::RValueRef<Function>, IRModule, PassContext)> pass_func,
         PassInfo pass_info) {
        auto wrapped_pass_func = [pass_func](Function func, IRModule mod, PassContext ctx) {
          return pass_func(ffi::RValueRef<Function>(std::move(func)), mod, ctx);
        };
        return FunctionPass(wrapped_pass_func, pass_info);
      });
}

// Pattern A (RM): auto-default repr from reflection for FunctionPassNode.

class DataflowBlockPass;

/*!
 * \brief DataflowBlock-level passes are used to implement various dataflow block
 * optimizations for a given Relax IRModule. It fetches one dataflow block at a time
 * from the functions in an IRModule, and yields a rewritten DataflowBlock.
 *
 * Note that the scope of passes at this level is a Relax DataflowBlock. Therefore,
 * we cannot modify the global scope Vars and symbolic shape Vars defined inside the dataflow block.
 */
class DataflowBlockPassNode : public tvm::transform::PassNode {
 public:
  /* \brief The pass meta data.*/
  PassInfo pass_info;

  /*! \brief The packed pass function sketches the real optimization. For
   * instance, we can implement a pass that works on a Relax DataflowBlock as a
   * `pass_func` and let it run on a given IRModule. The same `pass_func` will
   * then be applied on each DataflowBlock in the IRModule.
   */
  std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func;

  DataflowBlockPassNode() = default;

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DataflowBlockPassNode>().def_ro("pass_info", &DataflowBlockPassNode::pass_info);
  }

  IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;

  PassInfo Info() const override { return pass_info; }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.DataflowBlockPass", DataflowBlockPassNode, PassNode);
};

/*! \brief Helper to apply the passed function to dataflow blocks.*/
class DataflowBlockMutator : public ExprMutator {
 public:
  DataflowBlockMutator(std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func,
                       IRModule mod, PassContext pass_ctx)
      : pass_func_(pass_func), mod_(mod), pass_ctx_(pass_ctx) {}

  /*!
   * \brief Rewrite the DataflowBlockNode with pass_func_
   *
   * This function will check that there are no rewrites of the global scope Vars
   * and symbolic shape Vars defined inside the dataflow block.
   */
  BindingBlock VisitBindingBlock_(const DataflowBlockNode* n) final {
    // collect Global Scope Vars and Symbolic Vars inside the DataflowBlock
    ffi::Map<ffi::String, Var> global_scope_vars;
    ffi::Map<ffi::String, tirx::Var> symbolic_vars;
    for (const Binding& binding : n->bindings) {
      Var var = binding->var;
      if (const auto* match_cast = binding.as<MatchCastNode>()) {
        auto collected_vars = SymbolicVarCollector::Collect(match_cast->ty);
        for (const tirx::VarNode* var : collected_vars) {
          symbolic_vars.Set(var->name_hint, ffi::GetRef<tirx::Var>(var));
        }
      }
      if (!var.as<DataflowVarNode>()) {
        global_scope_vars.Set(var->name_hint, var);
      }
    }

    // apply pass_func_ to the DataflowBlock
    DataflowBlock block = ffi::GetRef<DataflowBlock>(n);
    DataflowBlock updated_block = pass_func_(block, mod_, pass_ctx_);

    // raise error if there are updates of recorded Global Scope Vars and Symbolic Vars
    for (const Binding& binding : updated_block->bindings) {
      Var var = binding->var;
      if (const auto* match_cast = binding.as<MatchCastNode>()) {
        auto collected_vars = SymbolicVarCollector::Collect(match_cast->ty);
        for (const tirx::VarNode* var : collected_vars) {
          if (symbolic_vars.count(var->name_hint) > 0) {
            tirx::Var old_var = symbolic_vars[var->name_hint];
            TVM_FFI_ICHECK(var == old_var.get())
                << "Error: DataflowBlock Pass should not rewrite any Symbolic Var.";
            symbolic_vars.erase(var->name_hint);
          }
        }
      }
      if (!var.as<DataflowVarNode>() && global_scope_vars.count(var->name_hint) > 0) {
        TVM_FFI_ICHECK(var.same_as(global_scope_vars[var->name_hint]))
            << "Error: DataflowBlock Pass should not rewrite any GlobalScope Var.";
        global_scope_vars.erase(var->name_hint);
      }
    }
    TVM_FFI_ICHECK(global_scope_vars.empty() && symbolic_vars.empty())
        << "Error: DataflowBlock Pass should not delete any GlobalScope/Symbolic Var.";

    return updated_block;
  }

 private:
  class SymbolicVarCollector : public TypeVisitor {
   public:
    static std::unordered_set<const tirx::VarNode*> Collect(const Type& info) {
      SymbolicVarCollector collector;
      collector.VisitType(info);
      return std::move(collector.symbolic_vars_);
    }

   private:
    void VisitTypeExprField(const PrimExpr& expr) final {
      if (const tirx::VarNode* sym_var = expr.as<tirx::VarNode>()) {
        symbolic_vars_.insert(sym_var);
      }
    }

   private:
    std::unordered_set<const tirx::VarNode*> symbolic_vars_;
  };

  std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func_;
  IRModule mod_;
  PassContext pass_ctx_;
};

class DataflowBlockPass : public Pass {
 public:
  /*!
   * \brief The constructor
   * \param pass_func The packed function which implements a pass.
   * \param pass_info The pass info.
   */
  TVM_DLL DataflowBlockPass(
      std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func,
      PassInfo pass_info);

  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataflowBlockPass, Pass, DataflowBlockPassNode);
};

DataflowBlockPass::DataflowBlockPass(
    std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func,
    PassInfo pass_info) {
  auto n = ffi::make_object<DataflowBlockPassNode>();
  n->pass_func = std::move(pass_func);
  n->pass_info = std::move(pass_info);
  data_ = std::move(n);
}

// Perform IRModule -> IRModule transformations at the DataflowBlock level.
IRModule DataflowBlockPassNode::operator()(IRModule mod, const PassContext& pass_ctx) const {
  const PassInfo& pass_info = Info();

  TVM_FFI_ICHECK(mod.defined());

  VLOG_CONTEXT << pass_info->name;
  VLOG(0) << "Executing DataflowBlock pass with opt level: " << pass_info->opt_level;
  VLOG(1) << "Input module:" << std::endl << mod;

  IRModule updated_mod = mod->ShallowCopy();

  DataflowBlockMutator dataflow_block_mutator(pass_func, updated_mod, pass_ctx);
  std::vector<std::pair<GlobalVar, Function> > updates;
  for (const auto& it : updated_mod->functions) {
    // only picks up relax::Function
    if (auto* n = it.second.as<FunctionNode>()) {
      Function func = ffi::GetRef<Function>(n);
      // Enrich at this leaf executor, rendering the location local to the
      // currently-processed function.
      Function updated_func;
      try {
        updated_func = dataflow_block_mutator.VisitExpr(func).as_or_throw<Function>();
      } catch (ffi::Error& err) {
        throw tvm::transform::EnrichPassErrorWithContext(err, updated_mod, pass_info->name,
                                                         it.first);
      }
      updates.push_back({it.first, updated_func});
    }
  }

  for (const auto& pair : updates) {
    updated_mod->Add(pair.first, pair.second, true);
  }

  VLOG(1) << "Output module:" << std::endl << updated_mod;

  return updated_mod;
}

Pass CreateDataflowBlockPass(
    std::function<DataflowBlock(DataflowBlock, IRModule, PassContext)> pass_func, int opt_level,
    ffi::String name, tvm::ffi::Array<ffi::String> required, bool traceable) {
  PassInfo pass_info = PassInfo(opt_level, name, required, traceable);
  return DataflowBlockPass(std::move(pass_func), pass_info);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def(
      "relax.transform.MakeDataflowBlockPass",
      [](ffi::TypedFunction<DataflowBlock(ffi::RValueRef<DataflowBlock>, IRModule, PassContext)>
             pass_func,
         PassInfo pass_info) {
        auto wrapped_pass_func = [pass_func](DataflowBlock func, IRModule mod, PassContext ctx) {
          return pass_func(ffi::RValueRef<DataflowBlock>(std::move(func)), mod, ctx);
        };
        return DataflowBlockPass(wrapped_pass_func, pass_info);
      });
}

// Pattern A (RM): auto-default repr from reflection for DataflowBlockPassNode.

TVM_FFI_STATIC_INIT_BLOCK() {
  FunctionPassNode::RegisterReflection();
  DataflowBlockPassNode::RegisterReflection();
}

}  // namespace transform
}  // namespace relax
}  // namespace tvm
