diff options
| author |  Guozhong Zhuang <guozhong.zhuang@intel.com> | 2018-08-20 15:06:18 -0700 |
|---|---|---|
| committer | Guozhong Zhuang <guozhong.zhuang@intel.com> | 2018-08-20 15:06:18 -0700 |
| commit | 86d9ce130c5691cdba16024f7cc7987082acd294 (patch) | |
| tree | ff3bb2ad111fdeb1afcdd739b4f4b41c64ee6966 /tensorflow/core | |
| parent | 2459fd5f01b2a135335b588803fd8946ea761387 (diff) | |
do not control primitive caching for small batch size
Diffstat (limited to 'tensorflow/core')
| -rw-r--r-- | tensorflow/core/kernels/mkl_conv_ops.cc | 1 | ||||
| -rw-r--r-- | tensorflow/core/util/mkl_util.h | 2 |
2 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/core/kernels/mkl_conv_ops.cc b/tensorflow/core/kernels/mkl_conv_ops.cc index a5763e4b74..7351e9a526 100644 --- a/tensorflow/core/kernels/mkl_conv_ops.cc +++ b/tensorflow/core/kernels/mkl_conv_ops.cc @@ -906,6 +906,7 @@ class MklConvOp : public OpKernel { // 1. Legacy CPU without AVX512/AVX2, or // 2. 1x1 convolution with stride != 1 not_cache_ = MklPrimitiveFactory<T>::IsPrimitiveMemOptEnabled() && + (src_dims[MklDnnDims::Dim_N] > kSmallBatchSize) && (MklPrimitiveFactory<T>::IsLegacyPlatform() || IsConv1x1StrideNot1(filter_dims, strides)); diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h index ee02debddf..5d28b8a3e5 100644 --- a/tensorflow/core/util/mkl_util.h +++ b/tensorflow/core/util/mkl_util.h @@ -105,6 +105,8 @@ typedef enum { Dim3d_I = 1 } MklDnnDims3D; +static const int kSmallBatchSize = 32; + #ifdef INTEL_MKL_ML_ONLY class MklShape { public: |