| 482 | } |
| 483 | |
| 484 | void Compute(OpKernelContext* context) override { |
| 485 | try { |
| 486 | // Input tensors |
| 487 | const Tensor& src_tensor = MklGetInput(context, kInputIndex_Src); |
| 488 | const Tensor& filter_tensor = MklGetInput(context, kInputIndex_Filter); |
| 489 | MklDnnShape src_mkl_shape, filter_mkl_shape; |
| 490 | GetMklShape(context, kInputIndex_Src, &src_mkl_shape, eager_mode); |
| 491 | GetMklShape(context, kInputIndex_Filter, &filter_mkl_shape, eager_mode); |
| 492 | |
| 493 | OP_REQUIRES(context, !filter_mkl_shape.IsMklTensor(), |
| 494 | errors::InvalidArgument("Filter should not be in " |
| 495 | "Mkl Layout")); |
| 496 | |
| 497 | MklDnnData<Tinput> src(&cpu_engine_); |
| 498 | MklDnnData<Tfilter> filter(&cpu_engine_); |
| 499 | |
| 500 | memory::dims src_dims, filter_dims, padding_left, padding_right, |
| 501 | dilations, strides; |
| 502 | memory::dims dst_dims_tf_order, dst_dims_mkl_order; |
| 503 | |
| 504 | // For Quantized-Conv2D and Pad fusion, we get padding from the |
| 505 | // `padding_list` attribute. Otherwise, we get it from one of the inputs. |
| 506 | bool quantized_pad_enabled = false; |
| 507 | for (auto const& padding_val : padding_list_) { |
| 508 | if (padding_val) { |
| 509 | quantized_pad_enabled = true; |
| 510 | break; |
| 511 | } |
| 512 | } |
| 513 | |
| 514 | if (fuse_pad_ || quantized_pad_enabled) { |
| 515 | PadWithConvFusion(context, padding_left, padding_right, |
| 516 | quantized_pad_enabled); |
| 517 | } |
| 518 | |
| 519 | // Get shapes of input tensors in OneDNN order |
| 520 | MklDnnConvUtil conv_utl(context, strides_, padding_, data_format_, |
| 521 | dilations_); |
| 522 | auto src_tf_shape = GetTfShape(context, kInputIndex_Src, eager_mode); |
| 523 | auto filter_tf_shape = |
| 524 | GetTfShape(context, kInputIndex_Filter, eager_mode); |
| 525 | conv_utl.GetConvFwdSizesInMklOrder( |
| 526 | src_tf_shape, filter_tf_shape, &src_dims, &filter_dims, &strides, |
| 527 | &dilations, &dst_dims_tf_order, &dst_dims_mkl_order, &padding_left, |
| 528 | &padding_right, (fuse_pad_ || quantized_pad_enabled), is_depthwise); |
| 529 | |
| 530 | if (!context->status().ok()) return; |
| 531 | |
| 532 | // Check for corner case - if there is nothing to compute, return. |
| 533 | TensorShape dst_tf_shape = MklDnnDimsToTFShape(dst_dims_tf_order); |
| 534 | |
| 535 | // Corner cases: output with 0 elements and 0 batch size. |
| 536 | Tensor* dst_tensor = nullptr; |
| 537 | Tensor tmp_tensor; |
| 538 | bool emit_filter_output = (typeid(Tinput) == typeid(Tfilter) && |
| 539 | typeid(Tinput) == typeid(Toutput) && |
| 540 | (typeid(Tinput) == typeid(float) || |
| 541 | typeid(Tinput) == typeid(bfloat16))) && |
nothing calls this directly
no test coverage detected