| 3071 | } |
| 3072 | |
| 3073 | Status MklLayoutRewritePass::MergeConv2DWithBiasAdd(std::unique_ptr<Graph>* g, |
| 3074 | Node* m, Node* n) { |
| 3075 | CHECK_EQ(((m->type_string() == csinfo_.bias_add && |
| 3076 | n->type_string() == csinfo_.conv2d)) || |
| 3077 | ((n->type_string() == csinfo_.bias_add && |
| 3078 | m->type_string() == csinfo_.conv2d)), |
| 3079 | true); |
| 3080 | |
| 3081 | // If 'm' is BiasAdd, then 'n' is Conv2D. Since Conv2D feeds BiasAdd, |
| 3082 | // BiasAdd is successor node, and Conv2D predecessor node. |
| 3083 | Node* pred = m->type_string() == csinfo_.bias_add ? n : m; |
| 3084 | Node* succ = m->type_string() == csinfo_.bias_add ? m : n; |
| 3085 | |
| 3086 | // 1. Get all attributes from input nodes. |
| 3087 | DataType T_pred, T_succ; |
| 3088 | string padding; |
| 3089 | std::vector<int32> strides; |
| 3090 | std::vector<int32> dilations; |
| 3091 | string data_format_pred, data_format_succ; |
| 3092 | bool use_cudnn_on_gpu; |
| 3093 | TF_CHECK_OK(GetNodeAttr(pred->def(), "T", &T_pred)); |
| 3094 | TF_CHECK_OK(GetNodeAttr(succ->def(), "T", &T_succ)); |
| 3095 | TF_CHECK_OK(GetNodeAttr(pred->def(), "padding", &padding)); |
| 3096 | TF_CHECK_OK(GetNodeAttr(pred->def(), "strides", &strides)); |
| 3097 | TF_CHECK_OK(GetNodeAttr(pred->def(), "dilations", &dilations)); |
| 3098 | TF_CHECK_OK(GetNodeAttr(pred->def(), "data_format", &data_format_pred)); |
| 3099 | TF_CHECK_OK(GetNodeAttr(succ->def(), "data_format", &data_format_succ)); |
| 3100 | TF_CHECK_OK(GetNodeAttr(pred->def(), "use_cudnn_on_gpu", &use_cudnn_on_gpu)); |
| 3101 | // We check to ensure that data formats of both succ and pred are same. |
| 3102 | // We expect them to be same, so we can enforce this as assert. |
| 3103 | // But assert can be too strict, so we enforce this as a check. |
| 3104 | // If the check fails, then we do not merge two nodes. |
| 3105 | // We also do same check for devices. |
| 3106 | if (data_format_pred != data_format_succ || T_pred != T_succ || |
| 3107 | pred->assigned_device_name() != succ->assigned_device_name() || |
| 3108 | pred->def().device() != succ->def().device()) { |
| 3109 | return Status(error::Code::INVALID_ARGUMENT, |
| 3110 | "data_format or T attribute or devices of Conv2D and " |
| 3111 | "BiasAdd do not match. Will skip node merge optimization"); |
| 3112 | } |
| 3113 | |
| 3114 | const int succ_num = succ->num_inputs(); |
| 3115 | gtl::InlinedVector<Node*, 4> succ_control_edges; |
| 3116 | gtl::InlinedVector<std::pair<Node*, int>, 4> succ_in(succ_num); |
| 3117 | FillInputs(succ, &succ_control_edges, &succ_in); |
| 3118 | |
| 3119 | const int pred_num = pred->num_inputs(); |
| 3120 | gtl::InlinedVector<Node*, 4> pred_control_edges; |
| 3121 | gtl::InlinedVector<std::pair<Node*, int>, 4> pred_in(pred_num); |
| 3122 | FillInputs(pred, &pred_control_edges, &pred_in); |
| 3123 | |
| 3124 | // We need to ensure that Conv2D only feeds to BiasAdd (some other operator is |
| 3125 | // not expecting output of Conv2D). If this is not the case, then we cannot |
| 3126 | // merge Conv2D with BiasAdd. |
| 3127 | const int kFirstOutputSlot = 0; |
| 3128 | for (const Edge* e : pred->out_edges()) { |
| 3129 | if (e->src_output() == kFirstOutputSlot && e->dst() != succ) { |
| 3130 | return Status(error::Code::INVALID_ARGUMENT, |
no test coverage detected