TODO(bhavanis): Move the below ConvertMklToTf() to mkl_util.h
| 63 | |
| 64 | // TODO(bhavanis): Move the below ConvertMklToTf() to mkl_util.h |
| 65 | static void ConvertMklToTf(OpKernel* op_kernel, OpKernelContext* context, |
| 66 | string data_format_str, DataType op_data_type, |
| 67 | bool has_avx512f, uint input_number) { |
| 68 | try { |
| 69 | // Check that input tensor is in OneDNN format. |
| 70 | const Tensor& input_tensor = MklGetInput(context, input_number); |
| 71 | MklDnnShape input_shape; |
| 72 | GetMklShape(context, input_number, &input_shape); |
| 73 | |
| 74 | // if input is already in Tf format, then copy input tensor to output. |
| 75 | if (!input_shape.IsMklTensor()) { |
| 76 | context->set_output(input_number, input_tensor); |
| 77 | VLOG(1) << "MKLToTFConversion: No conversion needed, " |
| 78 | << "copying input to output"; |
| 79 | return; |
| 80 | } |
| 81 | |
| 82 | // Check that input data type is same as operator data type and that it |
| 83 | // is same as output data type. |
| 84 | DataType input_data_type = op_kernel->input_type(input_number); |
| 85 | DataType output_data_type = op_kernel->output_type(input_number); |
| 86 | CHECK_EQ(op_data_type, input_data_type); |
| 87 | CHECK_EQ(op_data_type, output_data_type); |
| 88 | |
| 89 | auto cpu_engine = engine(ENGINE_CPU, 0); |
| 90 | MklDnnData<T> input(&cpu_engine); |
| 91 | |
| 92 | // Get OneDNN layout of input tensor. |
| 93 | auto input_mkl_md = input_shape.GetMklLayout(); |
| 94 | // Get TensorFlow layout of input tensor. Expected output of conversion |
| 95 | // has same layout as Tensorflow layout of input tensor. |
| 96 | auto output_tf_md = input_shape.GetTfLayout(); |
| 97 | // Set input OneDNN layout as the user layout. |
| 98 | input.SetUsrMem(input_mkl_md, &input_tensor); |
| 99 | |
| 100 | // Allocate output tensor. |
| 101 | TensorShape output_shape = input_shape.GetTfShape(); |
| 102 | Tensor* output_tensor = nullptr; |
| 103 | OP_REQUIRES_OK(context, context->allocate_output( |
| 104 | input_number, output_shape, &output_tensor)); |
| 105 | DCHECK(output_tensor); |
| 106 | |
| 107 | // Check if input needs to be reordered |
| 108 | if (input.IsReorderNeeded(OUTPUT_TF_MD)) { |
| 109 | // Insert reorder between OneDNN layout and TensorFlow layout |
| 110 | OP_REQUIRES( |
| 111 | context, |
| 112 | input.CheckReorderToOpMem(OUTPUT_TF_MD, output_tensor, context), |
| 113 | errors::Internal("MklToTfOp: Failed to create input reorder")); |
| 114 | } else { |
| 115 | // If not, just forward input tensor to output tensor. |
| 116 | OP_REQUIRES(context, |
| 117 | output_tensor->CopyFrom(input_tensor, output_shape), |
| 118 | errors::Internal( |
| 119 | "MklToTfOp: Failed to forward input tensor to output")); |
| 120 | } |
| 121 | } catch (dnnl::error& e) { |
| 122 | OP_REQUIRES_OK( |
nothing calls this directly
no test coverage detected