Method FixMklMetaDataEdges

tensorflow/core/graph/mkl_layout_pass.cc:4031–4086

}

bool MklLayoutRewritePass::FixMklMetaDataEdges(std::unique_ptr<Graph>* g,
                                               Node* n) {
  bool result = false;

  // If graph node is not OneDNN node, then return.
  DataType T = DT_INVALID;
  if (!TryGetNodeAttr(n->def(), "T", &T) ||
      !mkl_op_registry::IsMklLayoutDependentOp(n->type_string(), T)) {
    return result;
  }

  // If it is OneDNN node, then check if the input edges to this node that carry
  // OneDNN metadata are linked up correctly with the source node.

  // For OneDNN nodes, we generate twice the number of input tensors (n for OneDNN
  // data tensors + n for OneDNN metadata tensors). We need to check for correct
  // connection of n metadata tensors only.
  int num_data_inputs = n->num_inputs() / 2;
  for (int idx = 0; idx < num_data_inputs; idx++) {
    // Get the edge connecting input slot with index (idx).
    const Edge* e = nullptr;
    TF_CHECK_OK(n->input_edge(idx, &e));

    // If e is control edge, then skip.
    if (e->IsControlEdge()) {
      continue;
    }

    // Check that the source node for edge 'e' is OneDNN node. If it is not an OneDNN
    // node, then we don't need to do anything.
    Node* e_src = e->src();
    if (TryGetNodeAttr(e_src->def(), "T", &T) &&
        mkl_op_registry::IsMklLayoutDependentOp(e_src->type_string(), T)) {
      // Source node for edge 'e' is OneDNN node.
      // Destination node and destination input slot of e is node 'n' and 'idx'
      // resp.
      CHECK_EQ(e->dst(), n);
      CHECK_EQ(e->dst_input(), idx);

      // Let's get edge that carries OneDNN metadata corresponding to OneDNN data edge
      // 'e'. For that, let's first get the input slot of 'n' where the meta
      // edge will feed the value.
      int e_meta_in_slot =
          GetTensorMetaDataIndex(e->dst_input(), n->num_inputs());
      const Edge* e_meta = nullptr;
      TF_CHECK_OK(n->input_edge(e_meta_in_slot, &e_meta));

      // Let's check if we need to fix this meta edge.
      if (FixMklMetaDataEdgeIfNeeded(g, e, e_meta)) {
        result = true;
      }
    }
  }

  return result;
}

///////////////////////////////////////////////////////////////////////////////