| 2461 | } |
| 2462 | |
| 2463 | void MklLayoutRewritePass::AddWorkSpaceEdgeIfNeeded( |
| 2464 | std::unique_ptr<Graph>* g, const Node* orig_node, NodeBuilder* nb, |
| 2465 | std::vector<NodeBuilder::NodeOut>* ws_tensors, bool* are_ws_tensors_added) { |
| 2466 | bool workspace_edge_added = false; // Default initializer |
| 2467 | CHECK_NOTNULL(are_ws_tensors_added); |
| 2468 | *are_ws_tensors_added = false; // Default initializer |
| 2469 | |
| 2470 | DataType T; |
| 2471 | TF_CHECK_OK(GetNodeAttr(orig_node->def(), "T", &T)); |
| 2472 | for (auto ws : wsinfo_) { |
| 2473 | if (orig_node->type_string() == ws.fwd_op && |
| 2474 | mkl_op_registry::IsMklLayoutDependentOp( |
| 2475 | mkl_op_registry::GetMklOpName(orig_node->type_string()), T)) { |
| 2476 | // If this op is a fwd op, then we need to check if there is an |
| 2477 | // edge from this node's fwd_slot to bwdop's bwd_slot. If there is |
| 2478 | // an edge, then we just add an attribute on this node for setting |
| 2479 | // workspace_passed to true. We don't add actual workspace edge |
| 2480 | // in this node. Actual workspace edge gets added in the backward |
| 2481 | // op for this node. |
| 2482 | for (const Edge* e : orig_node->out_edges()) { |
| 2483 | if (e->src_output() == ws.fwd_slot && |
| 2484 | e->dst()->type_string() == ws.bwd_op && |
| 2485 | e->dst_input() == ws.bwd_slot) { |
| 2486 | nb->Attr("workspace_enabled", true); |
| 2487 | VLOG(1) << "MklLayoutRewritePass: workspace_enabled for " |
| 2488 | << orig_node->type_string(); |
| 2489 | workspace_edge_added = true; |
| 2490 | // We found the edge that we were looking for, so break. |
| 2491 | break; |
| 2492 | } |
| 2493 | } |
| 2494 | |
| 2495 | if (!workspace_edge_added) { |
| 2496 | // If we are here, then we did not find backward operator for this |
| 2497 | // node. |
| 2498 | nb->Attr("workspace_enabled", false); |
| 2499 | } |
| 2500 | } else if (orig_node->type_string() == ws.bwd_op && |
| 2501 | mkl_op_registry::IsMklLayoutDependentOp( |
| 2502 | mkl_op_registry::GetMklOpName(orig_node->type_string()), |
| 2503 | T)) { |
| 2504 | // If this op is a bwd op, then we need to add workspace edge and |
| 2505 | // it's OneDNN tensor edge between its corresponding fwd op and this |
| 2506 | // op. Corresponding fwd op is specified in 'fwd_op' field of |
| 2507 | // workspace info. fwd_slot and bwd_slot in workspace info specify |
| 2508 | // an edge between which slots connect forward and backward op. |
| 2509 | // Once all these criteria match, we add a workspace edge between |
| 2510 | // ws_fwd_slot and ws_bwd_slot. Its corresponding OneDNN tensor is |
| 2511 | // determined by interleaved/contiguous ordering. Function |
| 2512 | // DataIndexToMetaDataIndex tells us the location of OneDNN tensor |
| 2513 | // from the location of the Tensorflow tensor. |
| 2514 | for (const Edge* e : orig_node->in_edges()) { |
| 2515 | if (e->src_output() == ws.fwd_slot && |
| 2516 | // We would have rewritten the forward op, so we need to use |
| 2517 | // GetMklOpName call to get its OneDNN name. |
| 2518 | e->src()->type_string() == |
| 2519 | mkl_op_registry::GetMklOpName(ws.fwd_op) && |
| 2520 | e->dst_input() == ws.bwd_slot) { |
nothing calls this directly
no test coverage detected