| 48 | } |
| 49 | |
| 50 | void Compute(OpKernelContext* context) override { |
| 51 | try { |
| 52 | const Tensor& input_tensor = |
| 53 | MklGetInput(context, this->kInputTensorIndexInput); |
| 54 | MklDnnShape dnn_shape_input; |
| 55 | GetMklShape(context, this->kInputTensorIndexInput, &dnn_shape_input); |
| 56 | this->SanityCheckInput(context, input_tensor, dnn_shape_input); |
| 57 | if (!context->status().ok()) return; |
| 58 | |
| 59 | MklDnnData<T> dnn_data_input(&cpu_engine_); |
| 60 | |
| 61 | // Initialize variables for the pooling op. |
| 62 | MklPoolParameters pool_params; |
| 63 | // Check whether pooling is 2D or 3D. |
| 64 | bool is_pool2d = (this->ksize_.size() == 4); |
| 65 | // Get the input tensor and initialize the pooling parameters. |
| 66 | TensorShape input_tensor_shape = input_tensor.shape(); |
| 67 | this->InitMklPoolParameters(context, &pool_params, dnn_shape_input, |
| 68 | input_tensor_shape); |
| 69 | OP_REQUIRES_OK(context, context->status()); |
| 70 | |
| 71 | Tensor* output_tensor = nullptr; |
| 72 | memory::dims output_dims_mkl_order; |
| 73 | this->GetOutputDims(pool_params, &output_dims_mkl_order); |
| 74 | |
| 75 | // If input is an empty tensor, allocate an empty output tensor. |
| 76 | if (input_tensor.NumElements() == 0) { |
| 77 | const int kOutputIndex = 0; |
| 78 | this->AllocateEmptyOutputTensor(context, kOutputIndex, &pool_params, |
| 79 | output_dims_mkl_order, &output_tensor); |
| 80 | return; |
| 81 | } |
| 82 | |
| 83 | memory::dims filter_dims, strides, padding_left, padding_right; |
| 84 | // Get src/filter/stride/padding information. |
| 85 | this->PoolParamsToDims(&pool_params, &filter_dims, &strides, |
| 86 | &padding_left, &padding_right, is_pool2d); |
| 87 | |
| 88 | // Get the input memory descriptor. |
| 89 | memory::dims src_dims = |
| 90 | dnn_shape_input.IsMklTensor() |
| 91 | ? dnn_shape_input.GetSizesAsMklDnnDims() |
| 92 | : is_pool2d ? TFShapeToMklDnnDimsInNCHW(input_tensor.shape(), |
| 93 | this->data_format_tf_) |
| 94 | : TFShapeToMklDnnDimsInNCDHW(input_tensor.shape(), |
| 95 | this->data_format_tf_); |
| 96 | memory::desc input_md = dnn_shape_input.IsMklTensor() |
| 97 | ? dnn_shape_input.GetMklLayout() |
| 98 | : memory::desc(src_dims, MklDnnType<T>(), |
| 99 | this->data_format_dnnl_); |
| 100 | |
| 101 | // Get an average pooling primitive from the op pool. |
| 102 | MklPoolingFwdPrimitive<T>* pooling_fwd = nullptr; |
| 103 | prop_kind pooling_prop_kind; |
| 104 | bool int8_forward_inference = |
| 105 | std::is_same<T, qint8>::value || std::is_same<T, quint8>::value; |
| 106 | if (int8_forward_inference) |
| 107 | pooling_prop_kind = prop_kind::forward_inference; |
nothing calls this directly
no test coverage detected