From f465b03ef93ba27637890636fcae1a6991cdde3b Mon Sep 17 00:00:00 2001
From: Michal Gallus <michal.gallus@intel.com>
Date: Thu, 20 Sep 2018 15:50:42 +0200
Subject: [PATCH] Enable MKLDNN in Analysis Predictor

Also fix MKL-DNN pooling integration for ceil mode
---
 .../fluid/inference/api/analysis_predictor.cc |  2 +-
 paddle/fluid/operators/pool_mkldnn_op.cc      | 42 ++++++++++++++++---
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
index 684e0ce0e29..08c5d240fba 100644
--- a/paddle/fluid/inference/api/analysis_predictor.cc
+++ b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -76,10 +76,10 @@ bool AnalysisPredictor::Init(
   }
 
   OptimizeInferenceProgram();
-  ctx_ = executor_->Prepare(*inference_program_, 0);
   if (config_._use_mkldnn) {
     executor_->EnableMKLDNN(*inference_program_);
   }
+  ctx_ = executor_->Prepare(*inference_program_, 0);
 
   VLOG(5) << "to create variables";
   PADDLE_ENFORCE(scope_.get());
diff --git a/paddle/fluid/operators/pool_mkldnn_op.cc b/paddle/fluid/operators/pool_mkldnn_op.cc
index 5341187d1ce..b0132ceeb21 100644
--- a/paddle/fluid/operators/pool_mkldnn_op.cc
+++ b/paddle/fluid/operators/pool_mkldnn_op.cc
@@ -46,6 +46,26 @@ static std::string gethash(const memory::dims& input_dims,
          dims2str(paddings) + pooling_type + suffix;
 }
 
+static int ComputeCeiledOutput(int input_size, int kernel_size, int padding,
+                               int stride) {
+  return (input_size - kernel_size + 2 * padding) / stride + 1;
+}
+
+static void CorrectOutputSize(const std::vector<int>& src_tz,
+                              const std::vector<int>& dst_tz,
+                              const std::vector<int>& kernel_size,
+                              const std::vector<int>& paddings,
+                              const std::vector<int>& strides,
+                              std::vector<int>& right_bot_padding) {
+  for (int i = 0; i < right_bot_padding.size(); i++) {
+    int desired_size = ComputeCeiledOutput(src_tz[i + 2], kernel_size[i],
+                                           paddings[i], strides[i]);
+    if (desired_size != dst_tz[i + 2]) {
+      right_bot_padding[i] += strides[i];
+    }
+  }
+}
+
 template <typename T>
 class PoolMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
  public:
@@ -103,6 +123,13 @@ class PoolMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     auto pool_p =
         std::static_pointer_cast<pooling_forward>(dev_ctx.GetBlob(key_pool_p));
     if (pool_p == nullptr) {
+      const std::vector<int>& padding_left_top(paddings);
+      std::vector<int> padding_right_bottom(paddings);
+      bool ceil_mode = ctx.Attr<bool>("ceil_mode");
+      if (ceil_mode) {
+        CorrectOutputSize(src_tz, dst_tz, ksize, paddings, strides,
+                          padding_right_bottom);
+      }
       auto src_md = platform::MKLDNNMemDesc(
           src_tz, platform::MKLDNNGetDataType<T>(), input_format);
 
@@ -114,8 +141,9 @@ class PoolMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
                                             mkldnn::memory::format::any);
 
       std::shared_ptr<mkldnn::pooling_forward::primitive_desc> pool_pd =
-          CreatePrimitiveDesc(src_md, dst_md, strides, paddings, ksize,
-                              pooling_type, mkldnn_engine);
+          CreatePrimitiveDesc(src_md, dst_md, strides, padding_left_top,
+                              padding_right_bottom, ksize, pooling_type,
+                              mkldnn_engine, ceil_mode);
 
       // save pool_pd into global device context to be referred in backward path
       dev_ctx.SetBlob(key_pool_pd, pool_pd);
@@ -171,14 +199,16 @@ class PoolMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
  private:
   std::unique_ptr<mkldnn::pooling_forward::primitive_desc> CreatePrimitiveDesc(
       const mkldnn::memory::desc& src, const mkldnn::memory::desc& dst,
-      const std::vector<int>& stride, const std::vector<int>& padding,
-      const std::vector<int>& kernel, const std::string& pooling_type,
-      const mkldnn::engine& engine) const {
+      const std::vector<int>& stride, const std::vector<int>& padding_left_top,
+      const std::vector<int>& padding_right_bot, const std::vector<int>& kernel,
+      const std::string& pooling_type, const mkldnn::engine& engine,
+      bool ceil_mode) const {
     auto pool_desc = mkldnn::pooling_forward::desc(
         mkldnn::prop_kind::forward,
         pooling_type == "max" ? mkldnn::algorithm::pooling_max
                               : mkldnn::algorithm::pooling_avg,
-        src, dst, stride, kernel, padding, padding, mkldnn::padding_kind::zero);
+        src, dst, stride, kernel, padding_left_top, padding_right_bot,
+        mkldnn::padding_kind::zero);
 
     auto p_pool_pd =
         new mkldnn::pooling_forward::primitive_desc(pool_desc, engine);
-- 
GitLab