/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.

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

#pragma once

#include <functional>
#include <memory>
#include <string>
#include <unordered_map>
#include <unordered_set>

#include "paddle/fluid/framework/ir/pass.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/errors.h"

namespace paddle {
namespace framework {
namespace ir {
class MemOptVarInfo;
class Node;
}  // namespace ir

namespace paddle2cinn {

constexpr char kCinnLaunchOp[] = "cinn_launch";
constexpr char kInputVars[] = "InputVars";
constexpr char kNoNeedBufferFeeds[] = "NoNeedBufferFeeds";
constexpr char kInternalVars[] = "InternalVars";
constexpr char kOutputVars[] = "OutputVars";
constexpr char kMemOptVarInfoFromMainGraph[] =
    "mem_opt_var_info_from_main_graph";

using Name2VarInfoMap =
    std::unordered_map<std::string,
                       std::shared_ptr<framework::ir::MemOptVarInfo>>;
using GraphNodeSet = std::unordered_set<ir::Node*>;

struct OpTransInfo {
  const std::unordered_set<std::string> default_deny_ops{"feed", "fetch"};

  const std::unordered_map<std::string, std::function<bool(const ir::Node*)>>
      dynamic_op_cond{
          {"slice", [](const ir::Node* node) -> bool {
             if (!node->IsOp()) {
               return false;
             }
             auto* op_desc = node->Op();
             auto infer_flags =
                 op_desc->GetAttrIfExists<std::vector<int>>("infer_flags");
             if (std::find_if(
                     infer_flags.begin(), infer_flags.end(), [](int v) {
                       return v < 0;
                     }) != infer_flags.end()) {
               return true;
             }
             return false;
           }}};

  const std::unordered_map<std::string, std::unordered_set<std::string>>
      deny_param_cond{{"batch_norm", {"ReserveSpace"}},
                      {"batch_norm_grad", {"ReserveSpace"}}};

  std::unordered_set<std::string> GetDenyVarNames(
      const GraphNodeSet& cluster) const;
};

// A pass named BuildCinnPass, the function of this pass is:
//
// a) Detect the subgraphs that can be compiled by the CINN compiler. We call a
// detected subgraph a cluster, which is consisted of several op nodes.
//
// b) Call the CINN compiler to compile each original cluster and get the
// compiled cluster, which is consisted of several kCinnLaunchOp.
//
// c) Replace the original cluster with corresponding compiled cluster on the
// original graph.
//
// In this pass, some questions are handled with cautions:
//
// a) How to determine whether two op nodes can be divided into a cluster?
// Firstly, both op nodes should be compile supported.
// Secondly, there should be a direct path between the two op nodes through a
// var node.
// Thirdly, there should be no extra path between the two op nodes through
// unsupported op nodes.
// Lastly, if op nodes a and b can be divied into a cluster, op nodes b and c
// can be divided into a cluster, a and c can also be divided into a cluster.
// The implementation of cluster detection is encapsulated in the
// SubGraphDetector
// class.
//
// b) How to deal with the links between the var nodes in global graph and the
// op nodes in a cluster?
// We first add links between the var nodes in global graph and the op nodes in
// the compiled cluster, and then remove useless links between the var nodes in
// global graph and the op nodes in the original cluster.
class BuildCinnPass : public framework::ir::Pass {
 protected:
  void ApplyImpl(framework::ir::Graph* graph) const override;
};

}  // namespace paddle2cinn
}  // namespace framework
}  // namespace paddle