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

Method GetNodeProducingMklTensor

tensorflow/core/graph/mkl_layout_pass.cc:2166–2193  ·  view source on GitHub ↗

Get an input node that will feed OneDNN tensor to the new node that we are constructing. An input node could be (1) 'n' if it is OneDNN layer, or (2) a dummy node producing dummy OneDNN tensor if 'n' is not an OneDNN layer.

Source from the content-addressed store, hash-verified

2164// if it is OneDNN layer, or (2) a dummy node producing dummy OneDNN tensor
2165// if 'n' is not an OneDNN layer.
2166void MklLayoutRewritePass::GetNodeProducingMklTensor(
2167 std::unique_ptr<Graph>* g, const Node* orig_node, Node* n,
2168 int n_output_slot, Node** mkl_node, int* mkl_node_output_slot) {
2169 CHECK_NOTNULL(n);
2170 CHECK_NOTNULL(mkl_node);
2171 CHECK_NOTNULL(mkl_node_output_slot);
2172
2173 // If this is an OneDNN op, then it will create extra output for OneDNN layout.
2174 DataType T;
2175 if (TryGetNodeAttr(n->def(), "T", &T) &&
2176 mkl_op_registry::IsMklLayoutDependentOp(n->type_string(), T)) {
2177 // If this is an OneDNN op, then it will generate an edge that will receive
2178 // OneDNN tensor from a node.
2179 // output slot number for OneDNN tensor would be N+slot number of TensorFlow
2180 // tensor, where N is total number of TensorFlow tensors.
2181 *mkl_node = n;
2182 *mkl_node_output_slot =
2183 GetTensorMetaDataIndex(n_output_slot, n->num_outputs());
2184 } else {
2185 // If we have not visited the node and rewritten it, then we need
2186 // to create a dummy node that will feed a dummy OneDNN tensor to this node.
2187 // DummyMklTensor node has no input and generates only 1 output
2188 // (dummy OneDNN tensor) as output slot number 0.
2189 GetDummyMklTensorNode(g, mkl_node, orig_node);
2190 CHECK_NOTNULL(*mkl_node);
2191 *mkl_node_output_slot = 0;
2192 }
2193}
2194
2195int MklLayoutRewritePass::SetUpContiguousInputs(
2196 std::unique_ptr<Graph>* g,

Callers

nothing calls this directly

Calls 3

IsMklLayoutDependentOpFunction · 0.85
GetTensorMetaDataIndexFunction · 0.85
num_outputsMethod · 0.45

Tested by

no test coverage detected