build_cinn_pass.cc 25.5 KB
Newer Older
J
jiangcheng 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
/* 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. */

#include "paddle/fluid/framework/paddle2cinn/build_cinn_pass.h"

17 18
#include <algorithm>
#include <iterator>
J
jiangcheng 已提交
19
#include <memory>
20
#include <regex>
J
jiangcheng 已提交
21 22 23
#include <string>
#include <unordered_map>
#include <unordered_set>
24
#include <utility>
J
jiangcheng 已提交
25 26
#include <vector>

27 28
#include "cinn/frontend/op_mapper_registry.h"
#include "cinn/frontend/op_mappers/use_op_mappers.h"
29 30
#include "gflags/gflags.h"
#include "glog/logging.h"
J
jiangcheng 已提交
31
#include "paddle/fluid/framework/ir/graph.h"
32
#include "paddle/fluid/framework/ir/graph_pattern_detector.h"
J
jiangcheng 已提交
33 34
#include "paddle/fluid/framework/ir/node.h"
#include "paddle/fluid/framework/ir/subgraph_detector.h"
35
#include "paddle/fluid/framework/op_info.h"
36
#include "paddle/fluid/framework/op_proto_maker.h"
37
#include "paddle/fluid/framework/paddle2cinn/cinn_compiler.h"
38
#include "paddle/fluid/operators/cinn/cinn_launch_op.h"
39 40
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/platform/errors.h"
J
jiangcheng 已提交
41

42 43 44
DECLARE_string(allow_cinn_ops);
DECLARE_string(deny_cinn_ops);

J
jiangcheng 已提交
45 46 47 48 49 50 51 52
namespace paddle {
namespace framework {
namespace paddle2cinn {

using framework::ir::Graph;
using framework::ir::Node;

using GraphNodeVec = std::vector<Node*>;
53
using GraphNodeMap = std::unordered_map<Node*, Node*>;
J
jiangcheng 已提交
54

55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85
OpTransInfo::OpTransInfo() {
  // judgment condition for the dynamic slice
  dynamic_op_cond_.emplace("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");
    return std::find_if(infer_flags.begin(), infer_flags.end(), [](int v) {
             return v < 0;
           }) != infer_flags.end();
  });

  // judgment condition for the dynamic reshape
  dynamic_op_cond_.emplace("reshape", [](const ir::Node& node) -> bool {
    if (!node.IsOp()) {
      return false;
    }
    auto* op_desc = node.Op();
    bool has_shape_tensor = op_desc->Inputs().count("ShapeTensor") &&
                            op_desc->Inputs().at("ShapeTensor").size();
    bool has_shape = op_desc->Inputs().count("Shape") &&
                     op_desc->Inputs().at("Shape").size();
    return has_shape_tensor || has_shape;
  });

  // judgment condition for the dynamic reshape2
  dynamic_op_cond_.emplace("reshape2", dynamic_op_cond_.at("reshape"));
}

86 87
std::unordered_set<std::string> OpTransInfo::GetDenyVarNames(
    const GraphNodeSet& cluster) const {
88 89 90 91 92 93 94 95 96 97 98 99 100
  std::unordered_set<std::string> deny_var_set;

  auto get_debug_info = [](const std::unordered_set<std::string>& var_names) {
    std::string debug_info = "[";
    for (auto& var : var_names) {
      debug_info.append(var);
      debug_info.append(", ");
    }
    debug_info.append("]");
    return debug_info;
  };

  for (auto* op : cluster) {
101
    if (deny_param_cond_.count(op->Name())) {
102
      const auto* desc = op->Op();
103 104
      PADDLE_ENFORCE_NE(desc,
                        nullptr,
105 106
                        platform::errors::PreconditionNotMet(
                            "The Op %s's OpDesc should not be NULL, which has "
107
                            "a parameter in deny_param_cond_.",
108 109
                            op->Name().c_str()));

110
      auto deny_param_names = deny_param_cond_.at(op->Name());
111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
      VLOG(4) << "We found deny param " << get_debug_info(deny_param_names)
              << " in op [" << op->Name() << "].";

      for (const auto& param_name : deny_param_names) {
        if (desc->Inputs().count(param_name)) {
          const auto& arg_names = desc->Input(param_name);
          for (const auto& arg_name : arg_names) {
            deny_var_set.insert(arg_name);
            VLOG(4) << "deny param [" << param_name << "]'s argument name"
                    << " is [" << arg_name << "].";
          }
        }

        if (desc->HasOutput(param_name)) {
          const auto& arg_names = desc->Output(param_name);
          for (const auto& arg_name : arg_names) {
            deny_var_set.insert(arg_name);
            VLOG(4) << "deny param [" << param_name << "]'s argument name"
                    << " is [" << arg_name << "].";
          }
        }
      }
    }
  }

  VLOG(4) << "All deny var names are " << get_debug_info(deny_var_set);

  return deny_var_set;
}

141 142 143 144 145 146 147 148 149
bool OpTransInfo::IsInplaceOp(const OpDesc& op_desc) {
  auto inputs = op_desc.InputArgumentNames();
  std::unordered_set<std::string> input_set(inputs.begin(), inputs.end());
  for (auto& name : op_desc.OutputArgumentNames()) {
    if (input_set.count(name) > 0) return true;
  }
  return false;
}

150 151 152 153 154
namespace {
// The delim(`;`) that is used to split the FLAGS_allow_cinn_ops
// & FLAGS_deny_cinn_ops.
constexpr char kDelim[] = ";";

155 156 157 158 159 160 161 162 163 164
std::unordered_set<std::string> StringSplit(const std::string& str,
                                            const std::string& delim) {
  std::regex reg(delim);
  std::unordered_set<std::string> elems{
      std::sregex_token_iterator(str.begin(), str.end(), reg, -1),
      std::sregex_token_iterator()};
  elems.erase("");
  return elems;
}

165 166 167 168 169
int ExtractOpRole(const GraphNodeSet& cluster) {
  std::unordered_set<int> op_roles;
  std::string attr_name = OpProtoAndCheckerMaker::OpRoleAttrName();
  for (auto* n : cluster) {
    if (n->Op() && n->Op()->HasAttr(attr_name)) {
R
Ruibiao Chen 已提交
170
      op_roles.insert(PADDLE_GET_CONST(int, n->Op()->GetAttr(attr_name)));
171 172 173 174 175 176 177 178 179
    }
  }
  if (op_roles.size() == 1U) {
    return *(op_roles.begin());
  } else {
    return static_cast<int>(OpRole::kNotSpecified);
  }
}

180
// Deal with input var nodes of the target subgraph:
181
// create a new input var node and it's feed op node
182
void AddFeedOpAndVar(const GraphNodeSet& input_vars,
183
                     const GraphNodeSet& cluster,
184
                     const GraphNodeMap& old_op2new_op,
185 186
                     const GraphNodeMap& old_var2new_var,
                     Graph* graph) {
187
  for (auto* old_var : input_vars) {
188 189 190 191 192 193
    // create feed op
    OpDesc desc;
    desc.SetType("feed");
    desc.SetOutput("Out", {old_var->Name()});
    auto op = graph->CreateOpNode(&desc);

194 195
    // get new feed var node
    auto* var = old_var2new_var.at(old_var);
196
    VLOG(4) << "Add Feed Op before the input var: " << var->Name();
197 198

    // link feed op and feed var
199
    IR_NODE_LINK_TO(op, var);
200 201 202 203

    // link feed var to cluster op
    for (auto* old_op : old_var->outputs) {
      if (cluster.count(old_op)) {
204
        IR_NODE_LINK_TO(var, old_op2new_op.at(old_op));
205 206
      }
      // Do not need relink old op or old var here, they will be
207
      // fixed in RemoveSubGraphFromGraph, here we just deal with
208 209 210 211 212 213 214
      // new subgraph's node.
    }
  }
}

// Deal with subgraph's outputs var node:
// create a new output var node and it's fetch op
215 216
void AddOutputVar(const GraphNodeSet& output_vars,
                  const GraphNodeSet& cluster,
217
                  const GraphNodeMap& old_op2new_op,
218 219
                  const GraphNodeMap& old_var2new_var,
                  Graph* graph) {
220
  for (auto* old_var : output_vars) {
221 222 223 224 225 226
    // create fetch op
    OpDesc desc;
    desc.SetType("fetch");
    desc.SetInput("X", {old_var->Name()});
    auto op = graph->CreateOpNode(&desc);

227
    auto* var = old_var2new_var.at(old_var);
228
    VLOG(4) << "Add Output Var Node: " << var->Name();
229

230 231 232
    // link fetch op and fetch var
    IR_NODE_LINK_TO(var, op);

233 234
    for (auto* old_op : old_var->inputs) {
      if (cluster.count(old_op)) {
235
        IR_NODE_LINK_TO(old_op2new_op.at(old_op), var);
236 237 238 239 240
      }
    }
  }
}

241 242 243 244 245
std::unordered_set<std::string> ExtractNoNeedBufferFeeds(
    const GraphNodeSet& cluster, const GraphNodeSet& cluster_inputs) {
  // 1. Find op with NoNeedBufferVarsInferer defined and collect its input nodes
  std::unordered_map<Node*, GraphNodeSet> op_node2no_need_buffer_nodes;
  for (auto* op_node : cluster) {
246 247 248 249 250 251
    const auto* op = OpInfoMap::Instance().GetNullable(op_node->Name());
    // If op not registered in Paddle, skip
    if (!op) {
      continue;
    }
    auto& inferer = op->NoNeedBufferVarsInferer();
252 253 254 255 256
    if (!inferer) {
      continue;
    }
    auto* op_desc = op_node->Op();
    PADDLE_ENFORCE_NOT_NULL(
257 258 259
        op_desc,
        platform::errors::PreconditionNotMet(
            "The op desc of node in cluster shouldn't be null."));
260 261 262
    auto inferred_params =
        inferer(op_desc->Inputs(), op_desc->Inputs(), op_desc->GetAttrMap());
    std::unordered_set<std::string> inferred_args;
263 264
    std::for_each(inferred_params.begin(),
                  inferred_params.end(),
265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
                  [&op_desc, &inferred_args](const std::string& param) {
                    const auto& args = op_desc->Input(param);
                    inferred_args.insert(args.begin(), args.end());
                  });
    auto& no_need_buffer_nodes = op_node2no_need_buffer_nodes[op_node];
    for (auto* input_node : op_node->inputs) {
      if (input_node->Var() && inferred_args.count(input_node->Name())) {
        VLOG(4) << "Input node(" << input_node->Name() << ") of op("
                << op_node->Name() << ") is no_need_buffer";
        no_need_buffer_nodes.insert(input_node);
      }
    }
  }

  // 2. Extract no_need_buffer nodes from cluster_inputs by checking
  // all of their outputs are op nodes with NoNeedBufferVarsInferer
  // and they used as no_need_buffer inputs.
  auto check_all_used_as_no_need_buffer_fn =
      [&op_node2no_need_buffer_nodes](Node* var_node) -> bool {
    for (auto* output_node : var_node->outputs) {
      auto it = op_node2no_need_buffer_nodes.find(output_node);
      if (it == op_node2no_need_buffer_nodes.end()) {
        VLOG(4) << "Var node(" << var_node->Name() << ")'s output node("
                << output_node->Name()
                << ") doesn't have NoNeedBufferVarsInferer";
        return false;
      }
      if (it->second.count(var_node) == 0) {
        VLOG(4) << "Var node("
                << ") is not used as no_need_buffer inputs";
        return false;
      }
    }
    return true;
  };
  std::unordered_set<std::string> result;
  for (const auto& op2inputs_pair : op_node2no_need_buffer_nodes) {
    for (auto* input_node : op2inputs_pair.second) {
      if (cluster_inputs.count(input_node) &&
          check_all_used_as_no_need_buffer_fn(input_node)) {
        VLOG(4) << "Input node(" << input_node->Name()
                << ") is declared as no_need_buffer cluster_inputs";
        result.insert(input_node->Name());
      }
    }
  }
  return result;
}

J
jiangcheng 已提交
314 315
// Create new subgraph with and op nodes are cluster nodes, and all
// var node are from internal nodes
316 317
std::unique_ptr<Graph> CreateNewSubGraph(const GraphNodeSet& cluster,
                                         const GraphNodeSet& cluster_internals,
318 319
                                         const GraphNodeSet& cluster_inputs,
                                         const GraphNodeSet& cluster_outputs) {
J
jiangcheng 已提交
320 321
  // Graph's constructor must has one parameter, and in our code,
  // the ProgramDesc is useless, so here we pass a temporary object.
322
  auto subgraph = std::make_unique<Graph>(framework::ProgramDesc());
J
jiangcheng 已提交
323

324
  GraphNodeMap old_op2new_op;
J
jiangcheng 已提交
325
  for (auto* op : cluster) {
326
    auto sub_node = subgraph->CreateOpNode(op->Op());
J
jiangcheng 已提交
327 328 329
    old_op2new_op[op] = sub_node;
  }

330
  GraphNodeMap old_var2new_var;
J
jiangcheng 已提交
331
  for (auto* var : cluster_internals) {
332 333 334 335 336 337 338 339 340 341 342 343 344
    if (!var->Var()) {
      // skip control var

      // TODO(jiangcheng05): CINN not support control var now, so here we skip
      // it, but it may incur result incorrect problem. In detail, for two
      // unconnected ops, with control var, an op must run before another op.
      // If we remove the control var, the program wouldn't guarantee the run
      // ordering, in other words, the result may incorrect.
      VLOG(4)
          << "The internal var [" << var->Name() << "]'s vardesc empty,"
          << " it may be a control var, but CINN not support control var now.";
      continue;
    }
345
    auto* sub_node = subgraph->CreateVarNode(var->Var());
J
jiangcheng 已提交
346 347
    old_var2new_var[var] = sub_node;
  }
348 349 350 351 352 353 354 355 356 357 358 359
  for (auto* var : cluster_inputs) {
    if (var->Var()) {
      auto* sub_node = subgraph->CreateVarNode(var->Var());
      old_var2new_var[var] = sub_node;
    }
  }
  for (auto* var : cluster_outputs) {
    if (var->Var()) {
      auto* sub_node = subgraph->CreateVarNode(var->Var());
      old_var2new_var[var] = sub_node;
    }
  }
J
jiangcheng 已提交
360

361
  GraphNodeSet need_feed_vars;
362
  std::unordered_set<Node*> param_vars, output_vars;
J
jiangcheng 已提交
363 364 365 366 367
  // the subgraph is independently, so here we only need link
  // to the node in new subgraph, and discard the link to
  // out-graph.
  for (auto* op : cluster) {
    for (auto* var : op->inputs) {
368 369 370 371
      if (!var->Var()) {
        // skip control var
        continue;
      }
372 373 374 375
      // one output var maybe an input of the cluster
      if (cluster_internals.count(var) ||
          (cluster_outputs.count(var) && old_var2new_var.count(var))) {
        IR_NODE_LINK_TO(old_var2new_var.at(var), old_op2new_op.at(op));
376
      } else if (cluster_inputs.count(var) && var->Var() != nullptr) {
377 378 379 380 381 382 383 384 385 386 387
        if (var->Var()->IsParameter()) {
          // Parameters have been preserved in scope, compared to feed var,
          // param just need add new var and don't need add feed op.
          // The var is used for check whether we need preserve the tensor
          // when transform paddle scope to CINN scope.
          param_vars.insert(var);
        } else {
          // When the var is subgraph input and the var is not parameter,
          // we need add a new feed op to feed the var.
          need_feed_vars.insert(var);
        }
J
jiangcheng 已提交
388 389 390
      }
    }
    for (auto* var : op->outputs) {
391 392 393 394
      if (!var->Var()) {
        // skip control var
        continue;
      }
J
jiangcheng 已提交
395
      if (cluster_internals.count(var)) {
396
        IR_NODE_LINK_TO(old_op2new_op.at(op), old_var2new_var.at(var));
397
      } else if (cluster_outputs.count(var) && var->Var() != nullptr) {
398 399 400 401
        // Create new output var node to guarantee the independency of
        // subgraph. In other words, the subgraph has no connection with
        // other graph, even the input graph.
        output_vars.insert(var);
J
jiangcheng 已提交
402 403 404 405
      }
    }
  }

406 407
  AddFeedOpAndVar(
      need_feed_vars, cluster, old_op2new_op, old_var2new_var, subgraph.get());
408
  AddFeedOpAndVar(
409 410 411
      param_vars, cluster, old_op2new_op, old_var2new_var, subgraph.get());
  AddOutputVar(
      output_vars, cluster, old_op2new_op, old_var2new_var, subgraph.get());
412 413
  // Save lists of input variables, internal variables and output variables
  // of the cluster as attributes of the subgraph for convenience.
414 415 416 417 418 419 420 421 422 423 424 425
  auto collect_names_fn =
      [](const GraphNodeSet& nodes,
         const std::unordered_set<std::string>& ignore_names) {
        auto result = std::make_unique<std::vector<std::string>>();
        for (auto* node : nodes) {
          if (!node->Var() || ignore_names.count(node->Name())) {
            continue;
          }
          result->emplace_back(node->Name());
        }
        return result;
      };
426 427 428 429 430 431 432
  subgraph->Set<std::vector<std::string>>(
      kInternalVars, collect_names_fn(cluster_internals, {}).release());
  subgraph->Set<std::vector<std::string>>(
      kOutputVars, collect_names_fn(cluster_outputs, {}).release());
  // Divide input variables into two parts: one is common and will be used
  // in execution, the other may be empty and it is those variables whose
  // buffer are not needed and only be used in graph symbolization
433 434
  auto no_need_buffer_feeds = std::make_unique<std::unordered_set<std::string>>(
      ExtractNoNeedBufferFeeds(cluster, cluster_inputs));
435 436 437
  subgraph->Set<std::vector<std::string>>(
      kInputVars,
      collect_names_fn(cluster_inputs, *no_need_buffer_feeds).release());
438 439
  subgraph->Set<std::unordered_set<std::string>>(
      kNoNeedBufferFeeds, no_need_buffer_feeds.release());
440 441
  // initialize empty map for kMemOptVarInfoFromMainGraph attribute,
  // it will be filled on the share_mem_opt_info_to_subgraph pass
442
  subgraph->GetOrInit<Name2VarInfoMap>(kMemOptVarInfoFromMainGraph);
443
  return subgraph;
J
jiangcheng 已提交
444 445 446 447
}

// This interface is used to classify all variables involved in a cluster into
// three types: inputs, outputs, and internals.
448 449 450
// The input node is some subgraph op's input but not any subgraph op's output.
// The output node is some subgraph op's output and some out-graph op's input.
// Specially, the internal node is a node that only used by subgraph, and
J
jiangcheng 已提交
451
// out-graph should not using this node at all.
452 453
// cluster_inputs & cluster_outputs & cluster_internals == NULL
// cluster_outputs | cluster_internals == all graph op's outputs node
454 455 456
void AnalyseClusterVariables(
    const GraphNodeSet& cluster,
    const std::unordered_set<std::string>& deny_var_set,
457 458
    GraphNodeSet* cluster_inputs,
    GraphNodeSet* cluster_outputs,
459
    GraphNodeSet* cluster_internals) {
J
jiangcheng 已提交
460 461
  // collecting all input and output of op
  for (auto* op_node : cluster) {
462
    const auto& op_name = op_node->Name();
J
jiangcheng 已提交
463
    for (auto* input_var_node : op_node->inputs) {
464 465 466 467
      if (!deny_var_set.count(input_var_node->Name())) {
        // ignore deny var node
        cluster_inputs->insert(input_var_node);
      }
J
jiangcheng 已提交
468 469
    }
    for (auto* output_var_node : op_node->outputs) {
470 471 472
      if (!deny_var_set.count(output_var_node->Name())) {
        cluster_outputs->insert(output_var_node);
      }
J
jiangcheng 已提交
473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499
    }
  }
  // remove output node from cluster_inputs,
  // and add cluster_internals node
  for (auto* var_node : *cluster_outputs) {
    if (cluster_inputs->count(var_node) > 0) {
      // if a input node also exists in output list, remove
      cluster_inputs->erase(var_node);

      // the internal node is must an output node of sub-graph,
      // but not any input node of out-graph.
      bool is_only_used_internal = true;
      for (auto* next_op_node : var_node->outputs) {
        is_only_used_internal &= (cluster.count(next_op_node) > 0);
      }
      if (is_only_used_internal) {
        cluster_internals->insert(var_node);
      }
    }
  }

  // if a output node also exists in internal list, remove.
  for (auto* var_node : *cluster_internals) {
    cluster_outputs->erase(var_node);
  }
}

500
void AddLinkToCinnOp(const GraphNodeSet& cluster_inputs,
501 502
                     const GraphNodeSet& cluster_outputs,
                     Node* cinn_op_node) {
503 504 505 506 507 508 509 510 511 512 513 514 515 516
  // add new link from cluster_inputs to cinn_op_node
  for (auto* var_node : cluster_inputs) {
    IR_NODE_LINK_TO(var_node, cinn_op_node);
  }

  // add new link from cinn_op_node to cluster_outputs
  for (auto* var_node : cluster_outputs) {
    IR_NODE_LINK_TO(cinn_op_node, var_node);
  }
}

void AddCinnOpToGraph(const GraphNodeSet& cluster,
                      const GraphNodeSet& cluster_inputs,
                      const GraphNodeSet& cluster_outputs,
517
                      int64_t compilation_key,
518 519
                      const std::unordered_set<std::string>& deny_var_set,
                      Graph* graph) {
520 521 522
  // Add the cinn launch op
  framework::OpDesc cinn_op_desc;
  cinn_op_desc.SetType(kCinnLaunchOp);
523

524 525
  const auto& subgraph =
      CinnCompiler::GetInstance()->FindGraph(compilation_key);
526
  const auto& no_need_buffer_feeds =
527 528
      subgraph.Get<std::unordered_set<std::string>>(kNoNeedBufferFeeds);

529 530 531 532 533 534 535
  cinn_op_desc.SetInput(operators::kX,
                        subgraph.Get<std::vector<std::string>>(kInputVars));
  cinn_op_desc.SetInput(operators::kNoNeedBufferX,
                        std::vector<std::string>(no_need_buffer_feeds.begin(),
                                                 no_need_buffer_feeds.end()));
  cinn_op_desc.SetOutput(operators::kOutputs,
                         subgraph.Get<std::vector<std::string>>(kOutputVars));
536
  cinn_op_desc.SetAttr(operators::kCompilationKey, compilation_key);
537 538 539 540
  cinn_op_desc.SetAttr(OpProtoAndCheckerMaker::OpRoleAttrName(),
                       ExtractOpRole(cluster));
  cinn_op_desc.Flush();
  auto* cinn_op_node = graph->CreateOpNode(&cinn_op_desc);
541
  // Add new links from or to the cinn launch op node
542
  AddLinkToCinnOp(cluster_inputs, cluster_outputs, cinn_op_node);
543 544

  VLOG(4) << "Add op [" << kCinnLaunchOp << "] into graph.";
J
jiangcheng 已提交
545 546 547 548 549 550
}

// Removing cluster node and internals node from Graph
void RemoveSubGraphFromGraph(const GraphNodeSet& cluster,
                             const GraphNodeSet& cluster_internals,
                             Graph* graph) {
551 552 553 554 555 556
  const std::unordered_set<const Node*> const_cluster{cluster.cbegin(),
                                                      cluster.cend()};
  const std::unordered_set<const Node*> const_internals{
      cluster_internals.cbegin(), cluster_internals.cend()};
  ir::GraphSafeRemoveNodes(graph, const_cluster);
  ir::GraphSafeRemoveNodes(graph, const_internals);
J
jiangcheng 已提交
557 558
}

559
// Replacing Cinn subgraph to a cinn op node, whose op_type is
J
jiangcheng 已提交
560 561
// kCinnLaunchOp, and inputs ares cluster_inputs and outputs are
// cluster_outputs.
562
// Meanwhile, move all links of cluster to the cinn op.
563
void ReplaceSubGraphWithCinnOpNode(
564 565 566 567
    const GraphNodeSet& cluster,
    const GraphNodeSet& cluster_inputs,
    const GraphNodeSet& cluster_outputs,
    const GraphNodeSet& cluster_internals,
568
    int64_t compilation_key,
569 570
    const std::unordered_set<std::string>& deny_var_set,
    Graph* graph) {
571
  // Add the cinn op node whose name is "kCinnLaunchOp" into graph
572 573 574 575 576 577
  AddCinnOpToGraph(cluster,
                   cluster_inputs,
                   cluster_outputs,
                   compilation_key,
                   deny_var_set,
                   graph);
578
  // Remove the cinn subgraph from graph
J
jiangcheng 已提交
579 580 581 582 583 584 585
  RemoveSubGraphFromGraph(cluster, cluster_internals, graph);
}

// Search all subgraphs which all op node supported by CINN,
// Here we using SubgraphDetector to detecte the subgraph that
// all of op node supported by CINN. We using OpMapperRegistry
// to check whether the op node supported by CINN.
586
void SearchAllSubgraphs(Graph* graph) {
587 588
  auto allow_ops = StringSplit(FLAGS_allow_cinn_ops, kDelim);
  auto deny_ops = StringSplit(FLAGS_deny_cinn_ops, kDelim);
589 590
  OpTransInfo trans_info;
  auto teller = [&allow_ops, &deny_ops, &trans_info](const Node* node) {
591
    const auto& node_name = node->Name();
592
    bool registered = ::cinn::frontend::OpMapperRegistry::Global()->Find(
593
                          node_name) != nullptr;
594 595
    // skip the dynamic ops
    bool is_dynamic = false;
596 597
    if (trans_info.dynamic_op_cond().count(node_name)) {
      is_dynamic = trans_info.dynamic_op_cond().at(node_name)(*node);
598
    }
599 600 601 602

    bool is_support =
        registered && !trans_info.default_deny_ops().count(node_name) &&
        !is_dynamic && (node->IsOp() && !trans_info.IsInplaceOp(*node->Op()));
603 604
    // if the op type is registered in CINN and allow_ops is not empty, return
    // true only when it is in allow_ops
605
    if (!allow_ops.empty()) {
606
      return is_support && allow_ops.count(node_name);
607 608 609
    }
    // if the op type is registered in CINN and deny_ops is not empty, return
    // true only when it is not in deny_ops
610
    if (!deny_ops.empty()) {
611
      return is_support && !deny_ops.count(node_name);
612
    }
S
sneaxiy 已提交
613

614 615
    // if the user doesn't set FLAGS_allow_cinn_ops and FLAGS_deny_cinn_ops,
    // return true only when it is registered in CINN
616
    return is_support;
J
jiangcheng 已提交
617
  };
618 619
  VLOG(4) << "The allowed Cinn Ops: " << FLAGS_allow_cinn_ops;
  VLOG(4) << "The denied Cinn Ops: " << FLAGS_deny_cinn_ops;
J
jiangcheng 已提交
620 621
  std::vector<GraphNodeVec> clusters =
      framework::ir::SubgraphDetector(graph, teller)();
622 623
  LOG(INFO) << "--- [build_cinn_pass] detected " << clusters.size()
            << " cinn supported subgraphs";
J
jiangcheng 已提交
624

625 626 627 628 629 630 631 632 633 634
  auto cluster_debug_info = [](const GraphNodeSet& cluster) {
    std::string res = "(";
    for (auto* node : cluster) {
      res.append(node->Name());
      res.append(", ");
    }
    res.append(")");
    return res;
  };

635
  auto* cinn_compiler = CinnCompiler::GetInstance();
J
jiangcheng 已提交
636
  for (const auto& node_vec : clusters) {
637
    // Classify var node to inputs, outputs, and internals.
J
jiangcheng 已提交
638 639
    GraphNodeSet cluster_set(node_vec.begin(), node_vec.end());

640
    auto deny_var_set = trans_info.GetDenyVarNames(cluster_set);
641

J
jiangcheng 已提交
642
    GraphNodeSet cluster_inputs, cluster_outputs, cluster_internals;
643 644 645 646 647
    AnalyseClusterVariables(cluster_set,
                            deny_var_set,
                            &cluster_inputs,
                            &cluster_outputs,
                            &cluster_internals);
648 649 650 651 652 653 654

    VLOG(4) << "Cluster Ops: " << cluster_debug_info(cluster_set);
    VLOG(4) << "Cluster input vars: " << cluster_debug_info(cluster_inputs);
    VLOG(4) << "Cluster output vars: " << cluster_debug_info(cluster_outputs);
    VLOG(4) << "Cluster internal vars: "
            << cluster_debug_info(cluster_internals);

655 656
    // Create a new subgraph according to the found cluster and
    // save it in CinnCompiler
657
    auto compilation_key = cinn_compiler->AddGraph(CreateNewSubGraph(
658
        cluster_set, cluster_internals, cluster_inputs, cluster_outputs));
659 660
    VLOG(4) << "Compilation Key:\n"
            << cinn_compiler->ReadableKey(compilation_key);
661

662
    // Replace the found cluster to a new cinn op node
663 664 665 666 667 668 669
    ReplaceSubGraphWithCinnOpNode(cluster_set,
                                  cluster_inputs,
                                  cluster_outputs,
                                  cluster_internals,
                                  compilation_key,
                                  deny_var_set,
                                  graph);
J
jiangcheng 已提交
670 671
  }
}
672
}  // namespace
J
jiangcheng 已提交
673

674
void BuildCinnPass::ApplyImpl(Graph* graph) const { SearchAllSubgraphs(graph); }
J
jiangcheng 已提交
675 676 677 678 679 680

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

REGISTER_PASS(build_cinn_pass, paddle::framework::paddle2cinn::BuildCinnPass);