| 52 | } |
| 53 | |
| 54 | void Compute(OpKernelContext* context) override { |
| 55 | try { |
| 56 | const Tensor& input_tensor = |
| 57 | MklGetInput(context, this->kInputTensorIndexInput); |
| 58 | MklDnnShape dnn_shape_input; |
| 59 | GetMklShape(context, this->kInputTensorIndexInput, &dnn_shape_input); |
| 60 | this->SanityCheckInput(context, input_tensor, dnn_shape_input); |
| 61 | if (!context->status().ok()) return; |
| 62 | |
| 63 | MklDnnData<T> dnn_data_input(&cpu_engine_); |
| 64 | MklDnnData<T> dnn_data_output(&cpu_engine_); |
| 65 | |
| 66 | // Initialize variables for the pooling op. |
| 67 | MklPoolParameters pool_params; |
| 68 | // Check whether pooling is 2D or 3D. |
| 69 | bool is_pool2d = (this->ksize_.size() == 4); |
| 70 | // Get the input tensor and initialize the pooling parameters |
| 71 | TensorShape input_tensor_shape = input_tensor.shape(); |
| 72 | this->InitMklPoolParameters(context, &pool_params, dnn_shape_input, |
| 73 | input_tensor_shape); |
| 74 | OP_REQUIRES_OK(context, context->status()); |
| 75 | |
| 76 | // Declare output tensor |
| 77 | Tensor* output_tensor = nullptr; |
| 78 | // Declare output workspace tensor |
| 79 | Tensor* output_ws_tensor = nullptr; |
| 80 | memory::dims output_dims_mkl_order; |
| 81 | this->GetOutputDims(pool_params, &output_dims_mkl_order); |
| 82 | |
| 83 | // If input is an empty tensor, allocate an empty output tensor and return |
| 84 | if (input_tensor.NumElements() == 0) { |
| 85 | const int kOutputIndex = 0; |
| 86 | this->AllocateEmptyOutputTensor(context, kOutputIndex, &pool_params, |
| 87 | output_dims_mkl_order, &output_tensor); |
| 88 | bool int8_forward_inference = |
| 89 | std::is_same<T, qint8>::value || std::is_same<T, quint8>::value; |
| 90 | |
| 91 | // Allocate an empty workspace tensor if not Quantized MaxPooling |
| 92 | // Because Quantized MaxPooling does not have backward pass |
| 93 | // Therefore no workspace, which is used to help backward pass in OneDNN |
| 94 | if (!int8_forward_inference) { |
| 95 | const int kOutputWorkspaceIndex = 1; |
| 96 | // output_ws_tensor is not really used, so using output_dims_mkl_order |
| 97 | this->AllocateEmptyOutputTensor(context, kOutputWorkspaceIndex, |
| 98 | &pool_params, output_dims_mkl_order, |
| 99 | &output_ws_tensor); |
| 100 | } |
| 101 | return; |
| 102 | } |
| 103 | |
| 104 | // Get the input memory descriptor |
| 105 | memory::desc input_md = |
| 106 | dnn_shape_input.IsMklTensor() |
| 107 | ? dnn_shape_input.GetMklLayout() |
| 108 | : is_pool2d ? memory::desc( |
| 109 | TFShapeToMklDnnDimsInNCHW( |
| 110 | input_tensor_shape, this->data_format_tf_), |
| 111 | MklDnnType<T>(), this->data_format_dnnl_) |
nothing calls this directly
no test coverage detected