MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / CheckForNodeMerge

Method CheckForNodeMerge

tensorflow/core/graph/mkl_layout_pass.cc:3030–3071  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

3028//////////////////////////////////////////////////////////////////////////
3029
3030Node* MklLayoutRewritePass::CheckForNodeMerge(const Node* a) const {
3031 // TODO(nhasabni) Add check for type of node similar to CheckForNodeRewrite
3032 // once we support BiasAddGrad as OneDNN layer.
3033
3034 // Search for all matching mergeinfo.
3035 // We allow more than one match for extensibility.
3036 std::vector<const MergeInfo*> matching_mi;
3037 for (auto mi = minfo_.cbegin(); mi != minfo_.cend(); ++mi) {
3038 if (a->type_string() == mi->op1 || a->type_string() == mi->op2) {
3039 matching_mi.push_back(&*mi);
3040 }
3041 }
3042
3043 for (const MergeInfo* mi : matching_mi) {
3044 // Get the operand with which 'a' can be merged.
3045 Node* b = nullptr;
3046 if ((b = mi->get_node_to_be_merged(a)) == nullptr) {
3047 continue;
3048 }
3049
3050 // Get the control edges and input of node
3051 const int N_in = a->num_inputs();
3052 gtl::InlinedVector<Node*, 4> a_control_edges;
3053 gtl::InlinedVector<std::pair<Node*, int>, 4> a_in(N_in);
3054 FillInputs(a, &a_control_edges, &a_in);
3055
3056 const int B_in = b->num_inputs();
3057 gtl::InlinedVector<Node*, 4> b_control_edges;
3058 gtl::InlinedVector<std::pair<Node*, int>, 4> b_in(B_in);
3059 FillInputs(b, &b_control_edges, &b_in);
3060
3061 // Shouldn't merge if a and b have different control edges.
3062 if (a_control_edges != b_control_edges) {
3063 continue;
3064 } else {
3065 // We found a match.
3066 return b;
3067 }
3068 }
3069
3070 return nullptr;
3071}
3072
3073Status MklLayoutRewritePass::MergeConv2DWithBiasAdd(std::unique_ptr<Graph>* g,
3074 Node* m, Node* n) {

Callers

nothing calls this directly

Calls 3

FillInputsFunction · 0.70
push_backMethod · 0.45
num_inputsMethod · 0.45

Tested by

no test coverage detected