| 2133 | } |
| 2134 | |
| 2135 | void MklLayoutRewritePass::GetNodesProducingMklTensorList( |
| 2136 | std::unique_ptr<Graph>* g, const Node* orig_node, |
| 2137 | const gtl::InlinedVector<std::pair<Node*, int>, 4>& inputs, int* input_idx, |
| 2138 | int list_length, std::vector<NodeBuilder::NodeOut>* output_nodes) { |
| 2139 | CHECK_LT(*input_idx, inputs.size()); |
| 2140 | CHECK_GT(list_length, 0); |
| 2141 | CHECK_NOTNULL(output_nodes); |
| 2142 | output_nodes->reserve(list_length); |
| 2143 | |
| 2144 | while (list_length != 0) { |
| 2145 | CHECK_GT(list_length, 0); |
| 2146 | CHECK_LT(*input_idx, inputs.size()); |
| 2147 | Node* n = inputs[*input_idx].first; |
| 2148 | int slot = inputs[*input_idx].second; |
| 2149 | // If 'n' is producing a single tensor, then create a single OneDNN tensor |
| 2150 | // node. |
| 2151 | Node* mkl_node = nullptr; |
| 2152 | int mkl_node_output_slot = 0; |
| 2153 | GetNodeProducingMklTensor(g, orig_node, n, slot, &mkl_node, |
| 2154 | &mkl_node_output_slot); |
| 2155 | output_nodes->push_back( |
| 2156 | NodeBuilder::NodeOut(mkl_node, mkl_node_output_slot)); |
| 2157 | (*input_idx)++; |
| 2158 | list_length--; |
| 2159 | } |
| 2160 | } |
| 2161 | |
| 2162 | // Get an input node that will feed OneDNN tensor to the new |
| 2163 | // node that we are constructing. An input node could be (1) 'n' |