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.
| 2164 | // if it is OneDNN layer, or (2) a dummy node producing dummy OneDNN tensor |
| 2165 | // if 'n' is not an OneDNN layer. |
| 2166 | void 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 | |
| 2195 | int MklLayoutRewritePass::SetUpContiguousInputs( |
| 2196 | std::unique_ptr<Graph>* g, |
nothing calls this directly
no test coverage detected