build_cinn_pass.cc 24.2 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 53
namespace paddle {
namespace framework {
namespace paddle2cinn {

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

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

56
namespace {
57 58 59 60
// The delim(`;`) that is used to split the FLAGS_allow_cinn_ops
// & FLAGS_deny_cinn_ops.
constexpr char kDelim[] = ";";

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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
const std::unordered_map<std::string, std::unordered_set<std::string>>
    kDenyParamMap = {{"batch_norm", {"ReserveSpace"}},
                     {"batch_norm_grad", {"ReserveSpace"}}};

std::unordered_set<std::string> GetDenyVarNames(const GraphNodeSet& cluster) {
  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) {
    if (kDenyParamMap.count(op->Name())) {
      const auto* desc = op->Op();
      PADDLE_ENFORCE_NE(desc, nullptr,
                        platform::errors::PreconditionNotMet(
                            "The Op %s's OpDesc should not be NULL, which has "
                            "a parameter in kDenyParamMap.",
                            op->Name().c_str()));

      auto deny_param_names = kDenyParamMap.at(op->Name());
      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;
}

118 119 120 121 122 123 124 125 126 127
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;
}

128 129 130 131 132 133 134 135 136 137 138 139 140 141 142
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)) {
      op_roles.insert(BOOST_GET_CONST(int, n->Op()->GetAttr(attr_name)));
    }
  }
  if (op_roles.size() == 1U) {
    return *(op_roles.begin());
  } else {
    return static_cast<int>(OpRole::kNotSpecified);
  }
}

143 144
// Deal with subgraph's feed input var node:
// create a new input var node and it's feed op node
145 146 147
void AddFeedOpAndVar(const GraphNodeSet& feed_vars, const GraphNodeSet& cluster,
                     const GraphNodeMap& old_op2new_op,
                     const GraphNodeMap& old_var2new_var, Graph* graph) {
148 149 150 151 152 153 154
  for (auto* old_var : feed_vars) {
    // create feed op
    OpDesc desc;
    desc.SetType("feed");
    desc.SetOutput("Out", {old_var->Name()});
    auto op = graph->CreateOpNode(&desc);

155 156
    // get new feed var node
    auto* var = old_var2new_var.at(old_var);
157
    VLOG(4) << "Add Feed Op before: " << var->Name();
158 159

    // link feed op and feed var
160
    IR_NODE_LINK_TO(op, var);
161 162 163 164

    // link feed var to cluster op
    for (auto* old_op : old_var->outputs) {
      if (cluster.count(old_op)) {
165
        IR_NODE_LINK_TO(var, old_op2new_op.at(old_op));
166 167
      }
      // Do not need relink old op or old var here, they will be
168
      // fixed in RemoveSubGraphFromGraph, here we just deal with
169 170 171 172 173 174 175 176
      // new subgraph's node.
    }
  }
}

// Deal with subgraph's parameter var node:
// create a new input var node, it's data will get by scope,
// so it don't need feed op
177 178 179
void AddParamVar(const GraphNodeSet& param_vars, const GraphNodeSet& cluster,
                 const GraphNodeMap& old_op2new_op,
                 const GraphNodeMap& old_var2new_var, Graph* graph) {
180
  for (auto* old_var : param_vars) {
181
    auto* var = old_var2new_var.at(old_var);
182
    VLOG(4) << "Add Param Var Node: " << var->Name();
183 184 185

    for (auto* old_op : old_var->outputs) {
      if (cluster.count(old_op)) {
186
        IR_NODE_LINK_TO(var, old_op2new_op.at(old_op));
187 188 189 190 191 192 193
      }
    }
  }
}

// Deal with subgraph's outputs var node:
// create a new output var node and it's fetch op
194 195 196
void AddOutputVar(const GraphNodeSet& output_vars, const GraphNodeSet& cluster,
                  const GraphNodeMap& old_op2new_op,
                  const GraphNodeMap& old_var2new_var, Graph* graph) {
197
  for (auto* old_var : output_vars) {
198 199 200 201 202 203
    // create fetch op
    OpDesc desc;
    desc.SetType("fetch");
    desc.SetInput("X", {old_var->Name()});
    auto op = graph->CreateOpNode(&desc);

204
    auto* var = old_var2new_var.at(old_var);
205
    VLOG(4) << "Add Output Var Node: " << var->Name();
206

207 208 209
    // link fetch op and fetch var
    IR_NODE_LINK_TO(var, op);

210 211
    for (auto* old_op : old_var->inputs) {
      if (cluster.count(old_op)) {
212
        IR_NODE_LINK_TO(old_op2new_op.at(old_op), var);
213 214 215 216 217
      }
    }
  }
}

218 219 220 221 222
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) {
223 224 225 226 227 228
    const auto* op = OpInfoMap::Instance().GetNullable(op_node->Name());
    // If op not registered in Paddle, skip
    if (!op) {
      continue;
    }
    auto& inferer = op->NoNeedBufferVarsInferer();
229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288
    if (!inferer) {
      continue;
    }
    auto* op_desc = op_node->Op();
    PADDLE_ENFORCE_NOT_NULL(
        op_desc, platform::errors::PreconditionNotMet(
                     "The op desc of node in cluster shouldn't be null."));
    auto inferred_params =
        inferer(op_desc->Inputs(), op_desc->Inputs(), op_desc->GetAttrMap());
    std::unordered_set<std::string> inferred_args;
    std::for_each(inferred_params.begin(), inferred_params.end(),
                  [&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 已提交
289 290
// Create new subgraph with and op nodes are cluster nodes, and all
// var node are from internal nodes
291 292
std::unique_ptr<Graph> CreateNewSubGraph(const GraphNodeSet& cluster,
                                         const GraphNodeSet& cluster_internals,
293 294
                                         const GraphNodeSet& cluster_inputs,
                                         const GraphNodeSet& cluster_outputs) {
J
jiangcheng 已提交
295 296
  // Graph's constructor must has one parameter, and in our code,
  // the ProgramDesc is useless, so here we pass a temporary object.
297
  auto subgraph = std::make_unique<Graph>(framework::ProgramDesc());
J
jiangcheng 已提交
298

299
  GraphNodeMap old_op2new_op;
J
jiangcheng 已提交
300
  for (auto* op : cluster) {
301
    auto sub_node = subgraph->CreateOpNode(op->Op());
J
jiangcheng 已提交
302 303 304
    old_op2new_op[op] = sub_node;
  }

305
  GraphNodeMap old_var2new_var;
J
jiangcheng 已提交
306
  for (auto* var : cluster_internals) {
307 308 309 310 311 312 313 314 315 316 317 318 319
    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;
    }
320
    auto* sub_node = subgraph->CreateVarNode(var->Var());
J
jiangcheng 已提交
321 322
    old_var2new_var[var] = sub_node;
  }
323 324 325 326 327 328 329 330 331 332 333 334
  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 已提交
335

336
  GraphNodeSet need_feed_vars;
337
  std::unordered_set<Node *> param_vars, output_vars;
J
jiangcheng 已提交
338 339 340 341 342
  // 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) {
343 344 345 346
      if (!var->Var()) {
        // skip control var
        continue;
      }
347 348 349 350
      // 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));
351
      } else if (cluster_inputs.count(var) && var->Var() != nullptr) {
352 353 354 355 356 357 358 359 360 361 362
        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 已提交
363 364 365
      }
    }
    for (auto* var : op->outputs) {
366 367 368 369
      if (!var->Var()) {
        // skip control var
        continue;
      }
J
jiangcheng 已提交
370
      if (cluster_internals.count(var)) {
371
        IR_NODE_LINK_TO(old_op2new_op.at(op), old_var2new_var.at(var));
372
      } else if (cluster_outputs.count(var) && var->Var() != nullptr) {
373 374 375 376
        // 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 已提交
377 378 379 380
      }
    }
  }

381 382 383 384 385 386
  AddFeedOpAndVar(need_feed_vars, cluster, old_op2new_op, old_var2new_var,
                  subgraph.get());
  AddParamVar(param_vars, cluster, old_op2new_op, old_var2new_var,
              subgraph.get());
  AddOutputVar(output_vars, cluster, old_op2new_op, old_var2new_var,
               subgraph.get());
387 388 389 390 391 392 393
  // Save lists of input variables, internal variables and output variables
  // of the cluster as attributes of the subgraph for convenience.
  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) {
394
      if (!node->Var() || ignore_names.count(node->Name())) {
395 396 397 398 399 400 401 402 403 404 405 406 407
        continue;
      }
      result->emplace_back(node->Name());
    }
    return result;
  };
  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
408 409
  auto no_need_buffer_feeds = std::make_unique<std::unordered_set<std::string>>(
      ExtractNoNeedBufferFeeds(cluster, cluster_inputs));
410 411 412
  subgraph->Set<std::vector<std::string>>(
      kInputVars,
      collect_names_fn(cluster_inputs, *no_need_buffer_feeds).release());
413 414
  subgraph->Set<std::unordered_set<std::string>>(
      kNoNeedBufferFeeds, no_need_buffer_feeds.release());
415 416
  // initialize empty map for kMemOptVarInfoFromMainGraph attribute,
  // it will be filled on the share_mem_opt_info_to_subgraph pass
417
  subgraph->GetOrInit<Name2VarInfoMap>(kMemOptVarInfoFromMainGraph);
418
  return subgraph;
J
jiangcheng 已提交
419 420 421 422
}

// This interface is used to classify all variables involved in a cluster into
// three types: inputs, outputs, and internals.
423 424 425
// 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 已提交
426
// out-graph should not using this node at all.
427 428
// cluster_inputs & cluster_outputs & cluster_internals == NULL
// cluster_outputs | cluster_internals == all graph op's outputs node
429 430 431 432 433
void AnalyseClusterVariables(
    const GraphNodeSet& cluster,
    const std::unordered_set<std::string>& deny_var_set,
    GraphNodeSet* cluster_inputs, GraphNodeSet* cluster_outputs,
    GraphNodeSet* cluster_internals) {
J
jiangcheng 已提交
434 435
  // collecting all input and output of op
  for (auto* op_node : cluster) {
436
    const auto& op_name = op_node->Name();
J
jiangcheng 已提交
437
    for (auto* input_var_node : op_node->inputs) {
438 439 440 441
      if (!deny_var_set.count(input_var_node->Name())) {
        // ignore deny var node
        cluster_inputs->insert(input_var_node);
      }
J
jiangcheng 已提交
442 443
    }
    for (auto* output_var_node : op_node->outputs) {
444 445 446
      if (!deny_var_set.count(output_var_node->Name())) {
        cluster_outputs->insert(output_var_node);
      }
J
jiangcheng 已提交
447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473
    }
  }
  // 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);
  }
}

474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489
void AddLinkToCinnOp(const GraphNodeSet& cluster_inputs,
                     const GraphNodeSet& cluster_outputs, Node* cinn_op_node) {
  // 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,
490 491 492
                      const std::string& compilation_key,
                      const std::unordered_set<std::string>& deny_var_set,
                      Graph* graph) {
493 494 495
  // Add the cinn launch op
  framework::OpDesc cinn_op_desc;
  cinn_op_desc.SetType(kCinnLaunchOp);
496

497 498
  const auto& subgraph =
      CinnCompiler::GetInstance()->FindGraph(compilation_key);
499
  const auto& no_need_buffer_feeds =
500 501
      subgraph.Get<std::unordered_set<std::string>>(kNoNeedBufferFeeds);

502 503 504 505 506 507 508
  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));
509
  cinn_op_desc.SetAttr(operators::kCompilationKey, compilation_key);
510 511 512 513 514 515
  cinn_op_desc.SetAttr(OpProtoAndCheckerMaker::OpRoleAttrName(),
                       ExtractOpRole(cluster));
  cinn_op_desc.Flush();
  auto* cinn_op_node = graph->CreateOpNode(&cinn_op_desc);
  // Add new links from or to the the cinn launch op node
  AddLinkToCinnOp(cluster_inputs, cluster_outputs, cinn_op_node);
516 517

  VLOG(4) << "Add op [" << kCinnLaunchOp << "] into graph.";
J
jiangcheng 已提交
518 519 520 521 522 523
}

// Removing cluster node and internals node from Graph
void RemoveSubGraphFromGraph(const GraphNodeSet& cluster,
                             const GraphNodeSet& cluster_internals,
                             Graph* graph) {
524 525 526 527 528 529
  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 已提交
530 531
}

532
// Replacing Cinn subgraph to a cinn op node, whose op_type is
J
jiangcheng 已提交
533 534
// kCinnLaunchOp, and inputs ares cluster_inputs and outputs are
// cluster_outputs.
535
// Meanwhile, move all links of cluster to the cinn op.
536 537 538 539 540
void ReplaceSubGraphWithCinnOpNode(
    const GraphNodeSet& cluster, const GraphNodeSet& cluster_inputs,
    const GraphNodeSet& cluster_outputs, const GraphNodeSet& cluster_internals,
    const std::string& compilation_key,
    const std::unordered_set<std::string>& deny_var_set, Graph* graph) {
541 542
  // Add the cinn op node whose name is "kCinnLaunchOp" into graph
  AddCinnOpToGraph(cluster, cluster_inputs, cluster_outputs, compilation_key,
543
                   deny_var_set, graph);
544
  // Remove the cinn subgraph from graph
J
jiangcheng 已提交
545 546 547
  RemoveSubGraphFromGraph(cluster, cluster_internals, graph);
}

S
sneaxiy 已提交
548 549 550 551 552 553 554 555 556
static bool 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;
}

J
jiangcheng 已提交
557 558 559 560
// 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.
561
void SearchAllSubgraphs(Graph* graph) {
562 563 564 565 566 567 568 569 570 571 572 573 574 575 576
  auto allow_ops = StringSplit(FLAGS_allow_cinn_ops, kDelim);
  auto deny_ops = StringSplit(FLAGS_deny_cinn_ops, kDelim);
  auto teller = [&allow_ops, &deny_ops](const Node* node) {
    bool registered = ::cinn::frontend::OpMapperRegistry::Global()->Find(
                          node->Name()) != nullptr;
    // if the op type is registered in CINN and allow_ops is not empty, return
    // true only when it is in allow_ops
    if (allow_ops.size()) {
      return registered && allow_ops.count(node->Name());
    }
    // if the op type is registered in CINN and deny_ops is not empty, return
    // true only when it is not in deny_ops
    if (deny_ops.size()) {
      return registered && !deny_ops.count(node->Name());
    }
S
sneaxiy 已提交
577

578 579
    // if the user doesn't set FLAGS_allow_cinn_ops and FLAGS_deny_cinn_ops,
    // return true only when it is registered in CINN
S
sneaxiy 已提交
580
    return registered && (node->IsOp() && !IsInplaceOp(*node->Op()));
J
jiangcheng 已提交
581
  };
582 583
  VLOG(4) << "The allowed Cinn Ops: " << FLAGS_allow_cinn_ops;
  VLOG(4) << "The denied Cinn Ops: " << FLAGS_deny_cinn_ops;
J
jiangcheng 已提交
584 585 586
  std::vector<GraphNodeVec> clusters =
      framework::ir::SubgraphDetector(graph, teller)();

587 588 589 590 591 592 593 594 595 596
  auto cluster_debug_info = [](const GraphNodeSet& cluster) {
    std::string res = "(";
    for (auto* node : cluster) {
      res.append(node->Name());
      res.append(", ");
    }
    res.append(")");
    return res;
  };

597
  auto* cinn_compiler = CinnCompiler::GetInstance();
J
jiangcheng 已提交
598
  for (const auto& node_vec : clusters) {
599
    // Classify var node to inputs, outputs, and internals.
J
jiangcheng 已提交
600 601
    GraphNodeSet cluster_set(node_vec.begin(), node_vec.end());

602 603
    auto deny_var_set = GetDenyVarNames(cluster_set);

J
jiangcheng 已提交
604
    GraphNodeSet cluster_inputs, cluster_outputs, cluster_internals;
605 606
    AnalyseClusterVariables(cluster_set, deny_var_set, &cluster_inputs,
                            &cluster_outputs, &cluster_internals);
607 608 609 610 611 612 613

    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);

614 615
    // Create a new subgraph according to the found cluster and
    // save it in CinnCompiler
616 617
    std::string compilation_key = cinn_compiler->AddGraph(CreateNewSubGraph(
        cluster_set, cluster_internals, cluster_inputs, cluster_outputs));
618 619
    VLOG(4) << "Compilation Key:\n"
            << cinn_compiler->ReadableKey(compilation_key);
620

621 622
    // Replace the found cluster to a new cinn op node
    ReplaceSubGraphWithCinnOpNode(cluster_set, cluster_inputs, cluster_outputs,
623 624
                                  cluster_internals, compilation_key,
                                  deny_var_set, graph);
J
jiangcheng 已提交
625 626
  }
}
627
}  // namespace
J
jiangcheng 已提交
628

629
void BuildCinnPass::ApplyImpl(Graph* graph) const { SearchAllSubgraphs(graph); }
J
jiangcheng 已提交
630 631 632 633 634 635

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

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