///////////////////////////////////////////////////////////////////// Helper functions for node fusion /////////////////////////////////////////////////////////////////////
| 3863 | // Helper functions for node fusion |
| 3864 | ////////////////////////////////////////////////////////////////////////// |
| 3865 | Status MklLayoutRewritePass::FuseTransposeMklOpTranspose( |
| 3866 | std::unique_ptr<Graph>* g, std::vector<Node*>& nodes, |
| 3867 | std::function<void(const Node*, NodeBuilder* nb, bool)> copy_attrs, |
| 3868 | string data_format) { |
| 3869 | Node* transpose_to_nhwc = nodes[0]; |
| 3870 | Node* mklop = nodes[1]; |
| 3871 | Node* transpose_to_nchw = nodes[2]; |
| 3872 | |
| 3873 | const int transpose_nhwc_num_inputs = transpose_to_nhwc->num_inputs(); |
| 3874 | gtl::InlinedVector<Node*, 4> transpose_nhwc_control_edges; |
| 3875 | gtl::InlinedVector<std::pair<Node*, int>, 4> transpose_nhwc_in( |
| 3876 | transpose_nhwc_num_inputs); |
| 3877 | FillInputs(transpose_to_nhwc, &transpose_nhwc_control_edges, |
| 3878 | &transpose_nhwc_in); |
| 3879 | |
| 3880 | const int mklop_num_inputs = mklop->num_inputs(); |
| 3881 | gtl::InlinedVector<Node*, 4> mklop_control_edges; |
| 3882 | gtl::InlinedVector<std::pair<Node*, int>, 4> mklop_in(mklop_num_inputs); |
| 3883 | FillInputs(mklop, &mklop_control_edges, &mklop_in); |
| 3884 | |
| 3885 | const int transpose_nchw_num_inputs = transpose_to_nchw->num_inputs(); |
| 3886 | gtl::InlinedVector<Node*, 4> transpose_nchw_control_edges; |
| 3887 | gtl::InlinedVector<std::pair<Node*, int>, 4> transpose_nchw_in( |
| 3888 | transpose_nchw_num_inputs); |
| 3889 | FillInputs(transpose_to_nchw, &transpose_nchw_control_edges, |
| 3890 | &transpose_nchw_in); |
| 3891 | |
| 3892 | // We use same name as original node, but change the op |
| 3893 | // type. |
| 3894 | NodeBuilder nb(mklop->name(), mklop->type_string()); |
| 3895 | |
| 3896 | // Storing the output slots of the input nodes. |
| 3897 | for (int i = 0; i < mklop_num_inputs; i++) { |
| 3898 | if (mklop_in[i].first == transpose_to_nhwc) { |
| 3899 | // Fill "x": |
| 3900 | nb.Input(transpose_nhwc_in[0].first, transpose_nhwc_in[0].second); |
| 3901 | } else { |
| 3902 | // Fill inputs other than "x": |
| 3903 | nb.Input(mklop_in[i].first, mklop_in[i].second); |
| 3904 | } |
| 3905 | } |
| 3906 | |
| 3907 | copy_attrs(const_cast<const Node*>(mklop), &nb, true); |
| 3908 | nb.Attr("data_format", data_format); |
| 3909 | |
| 3910 | // Copy the device assigned to old node to new node. |
| 3911 | nb.Device(mklop->def().device()); |
| 3912 | |
| 3913 | // Create node. |
| 3914 | Node* new_node; |
| 3915 | TF_CHECK_OK(nb.Finalize(&**g, &new_node)); |
| 3916 | // No need to check if new_node is null because it will be null only when |
| 3917 | // Finalize fails. |
| 3918 | |
| 3919 | // Fill outputs. |
| 3920 | for (const Edge* e : transpose_to_nchw->out_edges()) { |
| 3921 | if (!e->IsControlEdge()) { |
| 3922 | const int kTransposeWithMklOpOutputSlot = 0; |
nothing calls this directly
no test coverage detected