From d9f93c42a50b1f1401d9c186eac0ae8dc9093c3b Mon Sep 17 00:00:00 2001
From: Jianwei Xie <xiejw@google.com>
Date: Wed, 24 Jan 2018 10:02:35 -0800
Subject: Merge changes from github.

PiperOrigin-RevId: 183100142
---
 tensorflow/core/kernels/mkl_maxpooling_op.cc | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

(limited to 'tensorflow/core/kernels/mkl_maxpooling_op.cc')

diff --git a/tensorflow/core/kernels/mkl_maxpooling_op.cc b/tensorflow/core/kernels/mkl_maxpooling_op.cc
index de4d7d2e72..82c5229bab 100644
--- a/tensorflow/core/kernels/mkl_maxpooling_op.cc
+++ b/tensorflow/core/kernels/mkl_maxpooling_op.cc
@@ -517,7 +517,7 @@ class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> {
 
       MklDnnData<T> dnn_data_input(&cpu_engine);
       MklDnnData<T> dnn_data_output(&cpu_engine);
-      MklDnnData<T> dnn_data_wksp(&cpu_engine);
+      MklDnnData<uint8> dnn_data_wksp(&cpu_engine);
 
       // initialize variables for the pooling op
       MklPoolParameters pool_params;
@@ -588,16 +588,16 @@ class MklMaxPoolingOp : public MklPoolingForwardOpBase<T> {
 
     void AllocateWorkspaceTensor(OpKernelContext* context,
                 const pooling_forward::primitive_desc& pool_fwd_prim_desc,
-                MklDnnData<T>* dnn_data_wksp) {
+                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_t_elems = this->GetNumTElements(workspace_pd);
+        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_t_elems);
+        workspace_tf_shape.AddDim(workspace_bytes);
         AllocateOutputSetMklShape(context, kOutputTensorIndexWorkspace,
                                 &workspace_tensor,
                                 workspace_tf_shape, workspace_mkl_shape);
@@ -651,7 +651,7 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> {
         if (!context->status().ok()) return;
 
         MklDnnData<T> grad_dnn_data(&cpu_engine);
-        MklDnnData<T> workspace_dnn_data(&cpu_engine);
+        MklDnnData<uint8> workspace_dnn_data(&cpu_engine);
         MklDnnData<T> output_dnn_data(&cpu_engine);
         Tensor* output_tensor = nullptr;
         MklPoolParameters pool_params;
@@ -770,7 +770,7 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> {
 
     void ConfigureWorkspace(const Tensor& workspace_tensor,
                         memory::primitive_desc workspace_pd,
-                        MklDnnData<T> *workspace_dnn_data) {
+                        MklDnnData<uint8> *workspace_dnn_data) {
         CHECK_NOTNULL(workspace_dnn_data);
 
         workspace_dnn_data->SetUsrMem(workspace_pd, &workspace_tensor);
@@ -811,7 +811,7 @@ class MklMaxPoolingGradOp : public MklPoolingBackwardOpBase<T> {
                     errors::InvalidArgument("Gradient must be "
                     "4-dimensional"));
         }
-        if (this->workspace_enabled_){
+        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"
-- 
cgit v1.2.3