| 3235 | } |
| 3236 | |
| 3237 | Status MklLayoutRewritePass::MergePadWithConv2D(std::unique_ptr<Graph>* g, |
| 3238 | Node* m, Node* n) { |
| 3239 | DCHECK((m->type_string() == csinfo_.pad && |
| 3240 | (n->type_string() == csinfo_.conv2d || |
| 3241 | n->type_string() == csinfo_.fused_conv2d)) || |
| 3242 | (n->type_string() == csinfo_.pad && |
| 3243 | (m->type_string() == csinfo_.conv2d || |
| 3244 | m->type_string() == csinfo_.fused_conv2d))); |
| 3245 | |
| 3246 | bool is_fused_conv2d = n->type_string() == csinfo_.fused_conv2d || |
| 3247 | m->type_string() == csinfo_.fused_conv2d; |
| 3248 | // Conv2D is successor node, and Pad predecessor node. |
| 3249 | Node* pred = m->type_string() == csinfo_.pad ? m : n; |
| 3250 | Node* succ = m->type_string() == csinfo_.pad ? n : m; |
| 3251 | |
| 3252 | // 1. Get all attributes from input nodes. |
| 3253 | DataType T_pred, T_succ; |
| 3254 | string padding; |
| 3255 | std::vector<int32> strides; |
| 3256 | std::vector<int32> dilations; |
| 3257 | string data_format_pred, data_format_succ; |
| 3258 | |
| 3259 | TF_CHECK_OK(GetNodeAttr(pred->def(), "T", &T_pred)); |
| 3260 | TF_CHECK_OK(GetNodeAttr(succ->def(), "T", &T_succ)); |
| 3261 | TF_CHECK_OK(GetNodeAttr(succ->def(), "padding", &padding)); |
| 3262 | TF_CHECK_OK(GetNodeAttr(succ->def(), "strides", &strides)); |
| 3263 | TF_CHECK_OK(GetNodeAttr(succ->def(), "dilations", &dilations)); |
| 3264 | // Check if the devices of both succ and pred are the same. |
| 3265 | // Assert is not used because it can be too strict. |
| 3266 | // Don't need to check for data formats because it is not available in Pad. |
| 3267 | if (T_pred != T_succ || |
| 3268 | pred->assigned_device_name() != succ->assigned_device_name() || |
| 3269 | pred->def().device() != succ->def().device()) { |
| 3270 | return Status(error::Code::INVALID_ARGUMENT, |
| 3271 | "T attribute or devices of Conv2D and " |
| 3272 | "Pad do not match. Will skip node merge optimization"); |
| 3273 | } |
| 3274 | |
| 3275 | const int succ_num = succ->num_inputs(); |
| 3276 | gtl::InlinedVector<Node*, 4> succ_control_edges; |
| 3277 | gtl::InlinedVector<std::pair<Node*, int>, 4> succ_in(succ_num); |
| 3278 | FillInputs(succ, &succ_control_edges, &succ_in); |
| 3279 | |
| 3280 | const int pred_num = pred->num_inputs(); |
| 3281 | gtl::InlinedVector<Node*, 4> pred_control_edges; |
| 3282 | gtl::InlinedVector<std::pair<Node*, int>, 4> pred_in(pred_num); |
| 3283 | FillInputs(pred, &pred_control_edges, &pred_in); |
| 3284 | |
| 3285 | // We need to ensure that Pad only feeds to Conv2D (some other operator is |
| 3286 | // not expecting output of Pad). If this is not the case, then we cannot |
| 3287 | // merge Conv2D with Pad. |
| 3288 | const int kFirstOutputSlot = 0; |
| 3289 | for (const Edge* e : pred->out_edges()) { |
| 3290 | if (e->src_output() == kFirstOutputSlot && e->dst() != succ) { |
| 3291 | return Status(error::Code::INVALID_ARGUMENT, |
| 3292 | "Pad does not feed to Conv2D, or " |
| 3293 | "it feeds Conv2D but has multiple outputs. " |
| 3294 | "Will skip node merge optimization"); |
no test coverage detected