| /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| // See docs in ../ops/nn_ops.cc. |
| #ifdef INTEL_MKL |
| #define EIGEN_USE_THREADS |
| |
| #include <algorithm> |
| |
| #include "mkldnn.hpp" |
| #include "tensorflow/core/framework/op_kernel.h" |
| #include "tensorflow/core/framework/register_types.h" |
| #include "tensorflow/core/kernels/mkl_pooling_ops_common.h" |
| #include "tensorflow/core/lib/core/errors.h" |
| #include "tensorflow/core/util/mkl_util.h" |
| #include "tensorflow/core/util/padding.h" |
| |
| using mkldnn::algorithm; |
| using mkldnn::engine; |
| using mkldnn::error; |
| using mkldnn::memory; |
| using mkldnn::padding_kind; |
| using mkldnn::pooling_backward; |
| using mkldnn::pooling_forward; |
| using mkldnn::prop_kind; |
| |
| namespace tensorflow { |
| |
| typedef Eigen::ThreadPoolDevice CPUDevice; |
| |
| // An implementation of MaxPooling (forward). |
| template <typename Device, typename T> |
| class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> { |
| public: |
| explicit MklMaxPoolingOp(OpKernelConstruction* context) |
| : MklPoolingForwardOpBase<T>(context) { |
| // In Max Pooling, MKLDNN does not allow passing workspace as NULL. |
| // So we set workspace_enabled_ to true. |
| this->workspace_enabled_ = true; |
| } |
| |
| void Compute(OpKernelContext* context) override { |
| try { |
| const Tensor& input_tensor = |
| MklGetInput(context, this->kInputTensorIndexInput); |
| MklDnnShape dnn_shape_input; |
| GetMklShape(context, this->kInputTensorIndexInput, &dnn_shape_input); |
| this->SanityCheckInput(context, input_tensor, dnn_shape_input); |
| if (!context->status().ok()) return; |
| |
| MklDnnData<T> dnn_data_input(&cpu_engine); |
| MklDnnData<T> dnn_data_output(&cpu_engine); |
| |
| // initialize variables for the pooling op |
| MklPoolParameters pool_params; |
| // check whether pooling is 2D or 3D |
| bool is_pool2d = (this->ksize_.size() == 4); |
| // Get the input tensor and initialize the pooling parameters |
| TensorShape input_tensor_shape = input_tensor.shape(); |
| this->InitMklPoolParameters(context, &pool_params, dnn_shape_input, |
| input_tensor_shape); |
| OP_REQUIRES_OK(context, context->status()); |
| |
| // Declare output tensor |
| Tensor* output_tensor = nullptr; |
| // Declare output workspace tensor |
| Tensor* output_ws_tensor = nullptr; |
| memory::dims output_dims_mkl_order; |
| this->GetOutputDims(pool_params, &output_dims_mkl_order); |
| |
| // If input is an empty tensor, allocate an empty output tensor and return |
| if (input_tensor.NumElements() == 0) { |
| const int kOutputIndex = 0; |
| this->AllocateEmptyOutputTensor(context, kOutputIndex, &pool_params, |
| output_dims_mkl_order, &output_tensor); |
| bool int8_forward_inference = |
| std::is_same<T, qint8>::value || std::is_same<T, quint8>::value; |
| |
| // Allocate an empty workspace tensor if not Quantized MaxPooling |
| // Because Quantized MaxPooling does not have backward pass |
| // Therefore no workspace, which is used to help backward pass in MKL |
| if (!int8_forward_inference) { |
| const int kOutputWorkspaceIndex = 1; |
| // output_ws_tensor is not really used, so using output_dims_mkl_order |
| this->AllocateEmptyOutputTensor(context, kOutputWorkspaceIndex, |
| &pool_params, output_dims_mkl_order, |
| &output_ws_tensor); |
| } |
| return; |
| } |
| |
| // Get the input memory descriptor |
| memory::desc input_md = |
| dnn_shape_input.IsMklTensor() |
| ? dnn_shape_input.GetMklLayout() |
| : is_pool2d ? memory::desc( |
| TFShapeToMklDnnDimsInNCHW( |
| input_tensor_shape, this->data_format_tf_), |
| MklDnnType<T>(), this->data_format_mkldnn_) |
| : memory::desc( |
| TFShapeToMklDnnDimsInNCDHW( |
| input_tensor_shape, this->data_format_tf_), |
| MklDnnType<T>(), this->data_format_mkldnn_); |
| |
| // Get src/filter/stride/padding information |
| memory::dims src_dims = |
| dnn_shape_input.IsMklTensor() |
| ? dnn_shape_input.GetSizesAsMklDnnDims() |
| : is_pool2d ? TFShapeToMklDnnDimsInNCHW(input_tensor.shape(), |
| this->data_format_tf_) |
| : TFShapeToMklDnnDimsInNCDHW(input_tensor.shape(), |
| this->data_format_tf_); |
| memory::dims filter_dims, strides, padding_left, padding_right; |
| this->PoolParamsToDims(&pool_params, &filter_dims, &strides, |
| &padding_left, &padding_right, is_pool2d); |
| |
| // Get a pooling op from the cached pool |
| MklPoolingFwdPrimitive<T>* pooling_fwd = nullptr; |
| prop_kind pooling_prop_kind; |
| bool int8_forward_inference = |
| std::is_same<T, qint8>::value || std::is_same<T, quint8>::value; |
| if (int8_forward_inference) |
| pooling_prop_kind = prop_kind::forward_inference; |
| else |
| pooling_prop_kind = prop_kind::forward_training; |
| MklPoolingParams fwdParams(src_dims, output_dims_mkl_order, filter_dims, |
| strides, padding_left, padding_right, |
| algorithm::pooling_max, pooling_prop_kind); |
| pooling_fwd = MklPoolingFwdPrimitiveFactory<T>::Get(fwdParams); |
| |
| // allocate output tensor |
| this->AllocateOutputTensor(context, *(pooling_fwd->GetPoolingFwdPd()), |
| output_dims_mkl_order, |
| this->data_format_mkldnn_, &output_tensor); |
| OP_REQUIRES_OK(context, context->status()); |
| dnn_data_output.SetUsrMem(output_dims_mkl_order, |
| pooling_fwd->GetDstMemoryFormat(), |
| output_tensor); |
| |
| // check wehther we need to reorder src |
| const T* src_data = input_tensor.flat<T>().data(); |
| if (input_md.data.format != pooling_fwd->GetSrcMemoryFormat()) { |
| dnn_data_input.SetUsrMem(input_md, &input_tensor); |
| auto src_target_primitive_desc = memory::primitive_desc( |
| {{src_dims}, MklDnnType<T>(), pooling_fwd->GetSrcMemoryFormat()}, |
| cpu_engine); |
| dnn_data_input.CheckReorderToOpMem(src_target_primitive_desc); |
| src_data = const_cast<T*>( |
| reinterpret_cast<T*>(dnn_data_input.GetOpMem().get_data_handle())); |
| } |
| |
| T* dst_data = output_tensor->flat<T>().data(); |
| |
| if (int8_forward_inference) { |
| // Execute pooling op |
| pooling_fwd->Execute(src_data, dst_data); |
| |
| // pass min, max from input to output |
| const Tensor& min_input_t = MklGetInput(context, 1); |
| const Tensor& max_input_t = MklGetInput(context, 2); |
| const float min_input = min_input_t.flat<float>()(0); |
| const float max_input = max_input_t.flat<float>()(0); |
| |
| Tensor* output_min = nullptr; |
| Tensor* output_max = nullptr; |
| MklDnnShape output_min_mkl_shape, output_max_mkl_shape; |
| output_min_mkl_shape.SetMklTensor(false); |
| output_max_mkl_shape.SetMklTensor(false); |
| AllocateOutputSetMklShape(context, 1, &output_min, {}, |
| output_min_mkl_shape); |
| AllocateOutputSetMklShape(context, 2, &output_max, {}, |
| output_max_mkl_shape); |
| output_min->flat<float>()(0) = min_input; |
| output_max->flat<float>()(0) = max_input; |
| } else { |
| MklDnnData<uint8> dnn_data_wksp(&cpu_engine); |
| AllocateWorkspaceTensor(context, *(pooling_fwd->GetPoolingFwdPd()), |
| &dnn_data_wksp); |
| OP_REQUIRES_OK(context, context->status()); |
| T* ws_data = |
| static_cast<T*>(dnn_data_wksp.GetOpMem().get_data_handle()); |
| |
| // execute pooling op |
| pooling_fwd->Execute(src_data, dst_data, ws_data); |
| } |
| } catch (mkldnn::error& e) { |
| string error_msg = "Status: " + std::to_string(e.status) + |
| ", message: " + string(e.message) + ", in file " + |
| string(__FILE__) + ":" + std::to_string(__LINE__); |
| OP_REQUIRES_OK(context, errors::Aborted("Compute received an exception:", |
| error_msg)); |
| } |
| } |
| |
| private: |
| const int kOutputTensorIndexWorkspace = 1; |
| engine cpu_engine = engine(engine::cpu, 0); |
| |
| void AllocateWorkspaceTensor( |
| OpKernelContext* context, |
| const pooling_forward::primitive_desc& pool_fwd_prim_desc, |
| MklDnnData<uint8>* dnn_data_wksp) { |
| CHECK_NOTNULL(dnn_data_wksp); |
| Tensor* workspace_tensor = nullptr; |
| memory::primitive_desc workspace_pd = |
| pool_fwd_prim_desc.workspace_primitive_desc(); |
| size_t workspace_bytes = workspace_pd.get_size(); |
| MklDnnShape workspace_mkl_shape; |
| workspace_mkl_shape.SetMklTensor(false); |
| TensorShape workspace_tf_shape; |
| workspace_tf_shape.AddDim(workspace_bytes); |
| AllocateOutputSetMklShape(context, kOutputTensorIndexWorkspace, |
| &workspace_tensor, workspace_tf_shape, |
| workspace_mkl_shape); |
| CHECK_NOTNULL(workspace_tensor); |
| dnn_data_wksp->SetUsrMem(workspace_pd, workspace_tensor); |
| } |
| }; |
| |
| // The operation to compute MaxPool gradients. |
| // It takes three inputs: |
| // - The original input tensor |
| // - The original output tensor |
| // - Backprop tensor for output |
| // It produces one output: backprop tensor for input. |
| template <class Device, class T> |
| class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> { |
| public: |
| explicit MklMaxPoolingGradOp(OpKernelConstruction* context) |
| : MklPoolingBackwardOpBase<T>(context) {} |
| void Compute(OpKernelContext* context) override { |
| try { |
| auto cpu_engine = engine(engine::cpu, 0); |
| const Tensor& orig_input_tensor = |
| MklGetInput(context, kInputTensorIndexOrigInput); |
| const Tensor& grad_tensor = |
| MklGetInput(context, kInputTensorIndexGradient); |
| const Tensor& workspace_tensor = |
| MklGetInput(context, kInputTensorIndexWorkspace); |
| MklDnnShape orig_input_mkl_shape, grad_mkl_shape; |
| GetMklShape(context, kInputTensorIndexOrigInput, &orig_input_mkl_shape); |
| GetMklShape(context, kInputTensorIndexGradient, &grad_mkl_shape); |
| if (!context->status().ok()) return; |
| |
| MklDnnData<T> grad_dnn_data(&cpu_engine); |
| MklDnnData<uint8> workspace_dnn_data(&cpu_engine); |
| |
| MklPoolParameters pool_params; |
| TensorShape orig_input_shape = orig_input_tensor.shape(); |
| |
| bool is_pool2d = (this->ksize_.size() == 4); |
| this->InitMklPoolParameters(context, &pool_params, orig_input_mkl_shape, |
| orig_input_shape); |
| |
| memory::dims filter_dims, strides, padding_left, padding_right; |
| this->PoolParamsToDims(&pool_params, &filter_dims, &strides, |
| &padding_left, &padding_right, is_pool2d); |
| |
| memory::dims orig_input_dims_mkl_order = |
| orig_input_mkl_shape.IsMklTensor() |
| ? orig_input_mkl_shape.GetSizesAsMklDnnDims() |
| : is_pool2d ? TFShapeToMklDnnDimsInNCHW(orig_input_shape, |
| this->data_format_tf_) |
| : TFShapeToMklDnnDimsInNCDHW(orig_input_shape, |
| this->data_format_tf_); |
| |
| memory::dims diff_dst_dims = |
| grad_mkl_shape.IsMklTensor() |
| ? grad_mkl_shape.GetSizesAsMklDnnDims() |
| : is_pool2d ? TFShapeToMklDnnDimsInNCHW(grad_tensor.shape(), |
| this->data_format_tf_) |
| : TFShapeToMklDnnDimsInNCDHW(grad_tensor.shape(), |
| this->data_format_tf_); |
| |
| memory::dims output_dims_mkl_order; |
| this->GetOutputDims(pool_params, &output_dims_mkl_order); |
| |
| MklPoolingParams bwdParams( |
| orig_input_dims_mkl_order, output_dims_mkl_order, filter_dims, |
| strides, padding_left, padding_right, algorithm::pooling_max, |
| prop_kind::forward_training); |
| MklPoolingBwdPrimitive<T>* pooling_bwd = |
| MklPoolingBwdPrimitiveFactory<T>::Get(bwdParams); |
| |
| // allocate output tensor and memory primitive |
| Tensor* output_tensor = nullptr; |
| this->AllocateOutputTensor(context, *(pooling_bwd->GetPoolingBwdPd()), |
| orig_input_dims_mkl_order, |
| this->data_format_mkldnn_, &output_tensor); |
| // get diff_dst mem desc |
| memory::desc diff_dst_md = |
| grad_mkl_shape.IsMklTensor() |
| ? grad_mkl_shape.GetMklLayout() |
| : memory::desc(diff_dst_dims, MklDnnType<T>(), |
| this->data_format_mkldnn_); |
| // check if diff_dst needs to be reordered |
| const T* diff_dst_data = grad_tensor.flat<T>().data(); |
| if (diff_dst_md.data.format != pooling_bwd->GetDiffDstFormat()) { |
| auto target_diff_dst = memory::primitive_desc( |
| {{diff_dst_dims}, MklDnnType<T>(), pooling_bwd->GetDiffDstFormat()}, |
| cpu_engine); |
| grad_dnn_data.SetUsrMem(diff_dst_md, &grad_tensor); |
| grad_dnn_data.CheckReorderToOpMem(target_diff_dst); |
| diff_dst_data = const_cast<T*>( |
| reinterpret_cast<T*>(grad_dnn_data.GetOpMem().get_data_handle())); |
| } |
| |
| void* ws_data = static_cast<void*>( |
| const_cast<uint8*>(workspace_tensor.flat<uint8>().data())); |
| |
| auto ws_md = |
| pooling_bwd->GetPoolingFwdPd()->workspace_primitive_desc().desc(); |
| if (ws_md.data.format != pooling_bwd->GetWorkspaceFormat()) { |
| memory::dims ws_dims; |
| ws_dims.assign(ws_md.data.dims, ws_md.data.dims + ws_md.data.ndims); |
| auto target_ws = |
| memory::primitive_desc({{ws_dims}, |
| pooling_bwd->GetWorkspaceDataType(), |
| pooling_bwd->GetWorkspaceFormat()}, |
| cpu_engine); |
| workspace_dnn_data.SetUsrMem(ws_md, &workspace_tensor); |
| workspace_dnn_data.CheckReorderToOpMem(target_ws); |
| ws_data = workspace_dnn_data.GetOpMem().get_data_handle(); |
| } |
| |
| T* diff_src_data = output_tensor->flat<T>().data(); |
| |
| // execute pooling |
| pooling_bwd->Execute(diff_dst_data, diff_src_data, ws_data); |
| } catch (mkldnn::error& e) { |
| string error_msg = "Status:" + std::to_string(e.status) + |
| ", message: " + string(e.message) + ". in file " + |
| string(__FILE__) + ":" + std::to_string(__LINE__); |
| OP_REQUIRES_OK(context, errors::Aborted("Compute received an exception:", |
| error_msg)); |
| } |
| } |
| |
| private: |
| // .Input("orig_input: T") |
| // .Input("orig_output: T") |
| // .Input("grad: T") |
| // .Input("workspace: T") |
| const int kInputTensorIndexOrigInput = 0; |
| const int kInputTensorIndexOrigOutput = 1; |
| const int kInputTensorIndexGradient = 2; |
| const int kInputTensorIndexWorkspace = 3; |
| |
| void ConfigureWorkspace(const Tensor& workspace_tensor, |
| memory::primitive_desc workspace_pd, |
| MklDnnData<uint8>* workspace_dnn_data) { |
| CHECK_NOTNULL(workspace_dnn_data); |
| |
| workspace_dnn_data->SetUsrMem(workspace_pd, &workspace_tensor); |
| } |
| |
| void SanityCheckInputs(OpKernelContext* context, |
| const Tensor& orig_input_tensor, |
| const Tensor& orig_output_tensor, |
| const Tensor& grad_tensor, |
| const Tensor& workspace_tensor, |
| const MklDnnShape& orig_input_mkl_shape, |
| const MklDnnShape& orig_output_mkl_shape, |
| const MklDnnShape& grad_mkl_shape, |
| const MklDnnShape& workspace_mkl_shape) { |
| if (!orig_input_mkl_shape.IsMklTensor()) { |
| OP_REQUIRES(context, orig_input_tensor.dims() == 4, |
| errors::InvalidArgument( |
| "Original input shape must be 4-dimensional")); |
| } else { |
| OP_REQUIRES(context, orig_input_mkl_shape.GetDimension() == 4, |
| errors::InvalidArgument( |
| "Original input shape must be 4-dimensional")); |
| } |
| if (!orig_output_mkl_shape.IsMklTensor()) { |
| OP_REQUIRES( |
| context, orig_output_tensor.dims() == 4, |
| errors::InvalidArgument("Original output must be 4-dimensional")); |
| } else { |
| OP_REQUIRES( |
| context, orig_output_mkl_shape.GetDimension() == 4, |
| errors::InvalidArgument("Original output must be 4-dimensional")); |
| } |
| if (!grad_mkl_shape.IsMklTensor()) { |
| OP_REQUIRES(context, grad_tensor.dims() == 4, |
| errors::InvalidArgument("Gradient must be 4-dimensional")); |
| } else { |
| OP_REQUIRES(context, grad_mkl_shape.GetDimension() == 4, |
| errors::InvalidArgument("Gradient must be 4-dimensional")); |
| } |
| if (this->workspace_enabled_) { |
| // The workspace should not be an MKL tensor |
| OP_REQUIRES(context, workspace_mkl_shape.IsMklTensor() == false, |
| errors::InvalidArgument( |
| "Workspace tensor should not be an MKL Tensor.")); |
| // It should only have one dimension |
| OP_REQUIRES( |
| context, workspace_tensor.dims() == 1, |
| errors::InvalidArgument("Workspace tensor must be 1-dimensional")); |
| } else { |
| OP_REQUIRES( |
| context, this->workspace_enabled_, |
| errors::Unimplemented("MKL-DNN Max Pooling does not " |
| "yet support the use case " |
| "where MaxPoolGrad is called without first" |
| " calling MaxPool.")); |
| } |
| } |
| }; // MklMaxPoolingGradOp |
| |
| #define REGISTER_MKL_MAXPOOL3D_KERNELS(T) \ |
| REGISTER_KERNEL_BUILDER( \ |
| Name("_MklMaxPool3D") \ |
| .Device(DEVICE_CPU) \ |
| .TypeConstraint<T>("T") \ |
| .Label(mkl_op_registry::kMklLayoutDependentOpLabel), \ |
| MklMaxPoolingOp<CPUDevice, T>); \ |
| REGISTER_KERNEL_BUILDER( \ |
| Name("_MklMaxPool3DGrad") \ |
| .Device(DEVICE_CPU) \ |
| .TypeConstraint<T>("T") \ |
| .Label(mkl_op_registry::kMklLayoutDependentOpLabel), \ |
| MklMaxPoolingGradOp<CPUDevice, T>); |
| |
| TF_CALL_float(REGISTER_MKL_MAXPOOL3D_KERNELS); |
| TF_CALL_bfloat16(REGISTER_MKL_MAXPOOL3D_KERNELS); |
| |
| #define REGISTER_MKL_MAXPOOL_KERNELS(T) \ |
| REGISTER_KERNEL_BUILDER( \ |
| Name("_MklMaxPool") \ |
| .Device(DEVICE_CPU) \ |
| .TypeConstraint<T>("T") \ |
| .Label(mkl_op_registry::kMklLayoutDependentOpLabel), \ |
| MklMaxPoolingOp<CPUDevice, T>); \ |
| REGISTER_KERNEL_BUILDER( \ |
| Name("_MklMaxPoolGrad") \ |
| .Device(DEVICE_CPU) \ |
| .TypeConstraint<T>("T") \ |
| .Label(mkl_op_registry::kMklLayoutDependentOpLabel), \ |
| MklMaxPoolingGradOp<CPUDevice, T>); |
| |
| TF_CALL_float(REGISTER_MKL_MAXPOOL_KERNELS); |
| TF_CALL_bfloat16(REGISTER_MKL_MAXPOOL_KERNELS); |
| |
| REGISTER_KERNEL_BUILDER(Name("_MklQuantizedMaxPool") |
| .Device(DEVICE_CPU) |
| .TypeConstraint<quint8>("T") |
| .Label(mkl_op_registry::kMklQuantizedOpLabel), |
| MklMaxPoolingOp<CPUDevice, quint8>); |
| |
| REGISTER_KERNEL_BUILDER(Name("_MklQuantizedMaxPool") |
| .Device(DEVICE_CPU) |
| .TypeConstraint<qint8>("T") |
| .Label(mkl_op_registry::kMklQuantizedOpLabel), |
| MklMaxPoolingOp<CPUDevice, qint8>); |
| |
| } // namespace tensorflow |
| #endif // INTEL_MKL |
