Method FixMklMetaDataEdgeIfNeeded

tensorflow/core/graph/mkl_layout_pass.cc:4004–4029

////////////////////////////////////////////////////////////////////////// Post-rewrite OneDNN metadata fixup pass //////////////////////////////////////////////////////////////////////////

//              Post-rewrite OneDNN metadata fixup pass
///////////////////////////////////////////////////////////////////////////////
bool MklLayoutRewritePass::FixMklMetaDataEdgeIfNeeded(std::unique_ptr<Graph>* g,
                                                      const Edge* e_data,
                                                      const Edge* e_metadata) {
  if (g == nullptr || e_data == nullptr || e_metadata == nullptr) {
    return false;
  }

  Node* n_data = e_data->src();
  int n_data_op_slot = e_data->src_output();
  int n_metadata_op_slot =
      GetTensorMetaDataIndex(n_data_op_slot, n_data->num_outputs());

  // If the source of meta edge is a constant node (producing dummy OneDNN metadata
  // tensor), then we will need to fix.
  if (IsConstant(e_metadata->src())) {
    Node* e_metadata_dst = e_metadata->dst();
    int e_metadata_in_slot = e_metadata->dst_input();
    auto new_edge = (*g)->AddEdge(n_data, n_metadata_op_slot, e_metadata_dst,
                                  e_metadata_in_slot);
    DCHECK(new_edge);
    (*g)->RemoveEdge(e_metadata);
    return true;
  }

  return false;
}

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