Method Optimize

tensorflow/core/grappler/optimizers/dice_fusion.cc:270–362

}  // namespace dicefusion

Status DiceFusion::Optimize(Cluster* cluster, const GrapplerItem& item,
                            GraphDef* output) {
  *output = item.graph;
#if defined(__GNUC__) && (__GNUC__ > 6) && (__AVX512F__)
  Status status;
  utils::MutableGraphView graph_view(output, &status);
  TF_RETURN_IF_ERROR(status);
  TF_RETURN_IF_ERROR(graph_view.SortTopologically(/*ignore_cycles=*/false, {}));
  const int num_nodes = item.graph.node_size();
  // invalidated_nodes - nodes that have been changed into a fused op
  // nodes_to_delete -  nodes that were fused into a fused op and are not needed
  // anymore
  std::vector<bool> invalidated_nodes(num_nodes);
  std::vector<bool> nodes_to_delete(num_nodes);
  const GraphDef* graph = graph_view.graph();

  VLOG(3) << "Before Dice fusion rewrites: " << graph->DebugString();

  for (int i = 0; i < num_nodes; ++i) {
    if (invalidated_nodes[i] || nodes_to_delete[i]) continue;

    dicefusion::DicePattern base;
    if (dicefusion::FindDicePattern(graph_view, i, &base)) {
      const auto& fused_node = graph->node(i);
      VLOG(2) << "Optimizing fused dice node " << SummarizeNodeDef(fused_node);

      std::vector<int> dead_nodes = {
          base.sub_1_id, base.sub_2_id, base.sigmoid_id, base.mul_1_id,
          base.mul_2_id, base.mul_3_id, base.mul_4_id,   base.sub_2_input_id};
      for (auto& id : dead_nodes) nodes_to_delete[id] = true;
      invalidated_nodes[base.add_id] = true;

      // Adding fused dice op
      const NodeDef& add = graph->node(base.add_id);

      NodeDef dice_op;
      dice_op.set_name(add.name());
      dice_op.set_op("Dice");
      dice_op.set_device(add.device());

      // Get input
      std::vector<int> nodes_with_inputs = {base.mul_4_id, base.sub_1_id,
                                            base.mul_1_id, base.mul_3_id};
      std::set<std::string> names;
      for (int i = 0; i < dead_nodes.size(); ++i) {
        const auto* node_view = graph_view.GetNode(dead_nodes.at(i));
        const auto* node_def = node_view->node();
        names.insert(node_def->name());
      }
      std::set<std::string> added_input;
      for (int i = 0; i < nodes_with_inputs.size(); ++i) {
        const int id = nodes_with_inputs.at(i);
        const auto* node_view = graph_view.GetNode(id);
        const auto* node_def = node_view->node();
        const std::size_t num_inputs = node_view->NumRegularFanins();
        for (size_t j = 0; j < num_inputs; ++j) {
          std::string node_name =
              dicefusion::get_node_by_tensor(node_def->input(j));