| 859 | } |
| 860 | |
| 861 | virtual void AllocateOutputTensor(OpKernelContext* context, |
| 862 | const ConvFwdPd& conv_prim_desc, |
| 863 | const memory::dims& output_dims_mkl_order, |
| 864 | MKL_TENSOR_FORMAT output_tf_format, |
| 865 | MklDnnShape* output_mkl_shape, |
| 866 | Tensor** output_tensor, |
| 867 | Tensor* tmp_tensor) { |
| 868 | DCHECK(output_tensor); |
| 869 | auto dst_md = conv_prim_desc.dst_desc(); |
| 870 | |
| 871 | if (!std::is_same<Ttemp_output, Toutput>::value) { |
| 872 | dst_md.data.data_type = |
| 873 | static_cast<dnnl_data_type_t>(MklDnnType<Toutput>()); |
| 874 | } |
| 875 | |
| 876 | // Allocate shape of OneDNN tensor |
| 877 | output_mkl_shape->SetMklTensor(true); |
| 878 | output_mkl_shape->SetMklLayout(&DST_MD); |
| 879 | output_mkl_shape->SetElemType(MklDnnType<Toutput>()); |
| 880 | output_mkl_shape->SetTfLayout(output_dims_mkl_order.size(), |
| 881 | output_dims_mkl_order, output_tf_format); |
| 882 | |
| 883 | // Allocate shape of TF tensor |
| 884 | TensorShape output_tf_shape; |
| 885 | output_tf_shape.AddDim((DST_MD.get_size() / sizeof(Toutput))); |
| 886 | if (eager_mode) { |
| 887 | AllocTmpBuffer<Toutput>(context, tmp_tensor, output_tf_shape); |
| 888 | output_tf_shape = output_mkl_shape->GetTfShape(); |
| 889 | } |
| 890 | |
| 891 | if (fuse_add_) { |
| 892 | const Tensor& add_tensor = MklGetInput(context, kInputIndex_Add); |
| 893 | MklDnnShape add_mkl_shape; |
| 894 | GetMklShape(context, kInputIndex_Add, &add_mkl_shape); |
| 895 | |
| 896 | // Check if reorder is needed |
| 897 | if (add_mkl_shape == *output_mkl_shape && |
| 898 | context->forward_input_to_output_with_shape( |
| 899 | kInputIndex_Add, kOutputIndex_Dst, output_tf_shape, |
| 900 | output_tensor)) { |
| 901 | AllocateOutputSetMklShape(context, kOutputIndex_Dst, *output_mkl_shape); |
| 902 | return; |
| 903 | } else { |
| 904 | AllocateOutputSetMklShape(context, kOutputIndex_Dst, output_tensor, |
| 905 | output_tf_shape, *output_mkl_shape, |
| 906 | eager_mode); |
| 907 | auto output_format_tag = MklTensorFormatToMklDnnDataFormat( |
| 908 | output_mkl_shape->GetTfDataFormat()); |
| 909 | OP_REQUIRES(context, output_format_tag != memory::format_tag::undef, |
| 910 | errors::InvalidArgument( |
| 911 | "MklConvOp: AddN fusion: Invalid data format")); |
| 912 | auto add_md = |
| 913 | add_mkl_shape.IsMklTensor() |
| 914 | ? add_mkl_shape.GetMklLayout() |
| 915 | : memory::desc(output_dims_mkl_order, MklDnnType<Toutput>(), |
| 916 | output_format_tag); |
| 917 | void* add_buf = static_cast<void*>( |
| 918 | const_cast<Toutput*>(add_tensor.flat<Toutput>().data())); |
nothing calls this directly
no test coverage detected