enable restructure for first stage

paddle/fluid/operators/conv_mkldnn_op.cc

@@ -132,6 +132,8 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     std::shared_ptr<mkldnn::memory> user_src_memory_p;
     std::shared_ptr<mkldnn::memory> dst_memory_p;
     std::vector<primitive> pipeline;
+    std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd;
+    std::shared_ptr<platform::ConvMKLDNNHandler> handler; 
     
     auto prim_key = key + "@conv_p";
     auto dst_key = key + "@dst_mem_p";
@@ -139,6 +141,44 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
     auto user_src_key = key + "@user_src_mem_p";
     auto src_reorder_key = key + "@src_mem_p" + "reorder_p";
     conv_p = std::static_pointer_cast<mkldnn::convolution_forward>(dev_ctx.GetBlob(prim_key));
+    if(conv_p == nullptr){
+      if(is_INT8){
+        CreateINT8Primitive(ctx, is_test, dev_ctx, mkldnn_engine, input, filter,
+        bias, output,
+        strides, paddings,
+        dilations, fuse_relu,
+        fuse_residual_conn,// input_data,
+        filter_data, src_tz,
+        weights_tz, g,
+        dst_tz, key,
+        dst_memory_p,
+        pipeline,
+        key_conv_pd,
+        src_memory_p,
+        user_src_memory_p,
+        conv_p,
+        conv_pd,
+        handler,
+        force_fp32_output);
+      }else{
+        CreateFP32Primitive(ctx, is_test, dev_ctx, mkldnn_engine, input, filter,
+        bias, output,
+        strides, paddings,
+        dilations, fuse_relu,
+        fuse_residual_conn, //input_data,
+        filter_data, src_tz,
+        weights_tz, g,
+        dst_tz, key,
+        dst_memory_p,
+        pipeline,
+        key_conv_pd,
+        src_memory_p,
+        user_src_memory_p,
+        conv_p,
+        conv_pd,
+        handler);
+      }
+    } else {
       auto src_memory_reorder_p = std::static_pointer_cast<mkldnn::memory>(dev_ctx.GetBlob(src_reorder_key));
       src_memory_p = std::static_pointer_cast<mkldnn::memory>(dev_ctx.GetBlob(src_key));
       if(src_memory_reorder_p){
@@ -149,14 +189,11 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
       }
       
       dst_memory_p = std::static_pointer_cast<mkldnn::memory>(dev_ctx.GetBlob(dst_key));
-    
-    std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd;
       conv_pd = std::static_pointer_cast<mkldnn::convolution_forward::primitive_desc>(dev_ctx.GetBlob(key_conv_pd));
-    std::shared_ptr<platform::ConvMKLDNNHandler> handler;
       if(conv_pd){
         handler.reset(new platform::ConvMKLDNNHandler(conv_pd, dev_ctx, mkldnn_engine, key));
       }
-    if (!is_INT8 && dst_memory_p){
+      if (!is_INT8){
         if (fuse_residual_conn) {
           auto residual_param = ctx.Input<Tensor>("ResidualData");
           auto residual_param_data = residual_param->data<T>();
@@ -184,7 +221,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
               output->mutable_data<T>(ctx.GetPlace(), ::paddle::memory::Allocator::kDefault, handler->GetDstMemorySize());
           dst_memory_p->set_data_handle(to_void_cast<T>(output_data)); 
         }
-    } else if(is_INT8 && dst_memory_p){
+      } else if(is_INT8){
         if(fuse_residual_conn) {
           auto residual_param = ctx.Input<Tensor>("ResidualData");
           auto residual_dt = paddle::framework::ToMKLDNNDataType(residual_param->type());
@@ -210,8 +247,48 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
         }
       }
 
-    if(!is_INT8){
-      if(conv_p == nullptr){
+      if(src_memory_reorder_p){
+        pipeline.push_back(*src_memory_reorder_p);
+      }
+      pipeline.push_back(*conv_p);
+    }
+
+    // push primitive to stream and wait until it's executed
+    //pipeline.push_back(*conv_p);
+    stream(stream::kind::eager).submit(pipeline).wait();
+
+    if (need_s8_to_u8) {
+      output->mutable_data<uint8_t>(ctx.GetPlace());
+    }
+
+    output->set_layout(DataLayout::kMKLDNN);
+    output->set_format(GetMKLDNNFormat(*dst_memory_p));
+  };
+
+  private:
+    void CreateFP32Primitive(
+    paddle::framework::ExecutionContext ctx, bool is_test,
+    const paddle::platform::MKLDNNDeviceContext& dev_ctx,
+    const mkldnn::engine&  mkldnn_engine,
+    const paddle::framework::Tensor* input, const paddle::framework::Tensor* filter,
+    const paddle::framework::Tensor* bias, paddle::framework::Tensor* output,
+    std::vector<int> strides, std::vector<int> paddings,
+    std::vector<int> dilations, bool fuse_relu,
+    bool fuse_residual_conn, //const T* input_data,
+    const float* filter_data, std::vector<int> src_tz,
+    std::vector<int> weights_tz, int g,
+    std::vector<int> dst_tz, const std::string key,
+    std::shared_ptr<mkldnn::memory> &dst_memory_p,
+    std::vector<primitive>& pipeline,
+    const std::string &key_conv_pd,
+    std::shared_ptr<mkldnn::memory> src_memory_p,
+    std::shared_ptr<mkldnn::memory> user_src_memory_p,
+    std::shared_ptr<mkldnn::convolution_forward> conv_p,
+    std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd,
+    std::shared_ptr<platform::ConvMKLDNNHandler> handler) const{
+
+      const T* input_data = input->data<T>();
+
       auto user_src_md = platform::MKLDNNMemDesc(
           {src_tz}, platform::MKLDNNGetDataType<T>(), input->format());
       auto user_weights_md = platform::MKLDNNMemDesc(
@@ -322,22 +399,31 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
       }
       // push primitive to stream and wait until it's executed
       pipeline.push_back(*conv_p);
-        stream(stream::kind::eager).submit(pipeline).wait();
-
-        output->set_layout(DataLayout::kMKLDNN);
-        output->set_format(GetMKLDNNFormat(*dst_memory_p));
-      } else {
-        if(src_memory_reorder_p){
-          pipeline.push_back(*src_memory_reorder_p);
-        } 
-        pipeline.push_back(*conv_p);
-        stream(stream::kind::eager).submit(pipeline).wait();
-
-        output->set_layout(DataLayout::kMKLDNN);
-        output->set_format(GetMKLDNNFormat(*dst_memory_p));
-      }
-    } else{
-      if(conv_p == nullptr){
+    };
+
+    void CreateINT8Primitive(
+    const paddle::framework::ExecutionContext& ctx, bool is_test,
+    const paddle::platform::MKLDNNDeviceContext & dev_ctx,
+    const mkldnn::engine & mkldnn_engine,
+    const paddle::framework::Tensor* input, const paddle::framework::Tensor* filter,
+    const paddle::framework::Tensor* bias, paddle::framework::Tensor* output,
+    std::vector<int> strides, std::vector<int> paddings,
+    std::vector<int> dilations, bool fuse_relu,
+    bool fuse_residual_conn,// const T* input_data,
+    const float* filter_data, std::vector<int> src_tz,
+    std::vector<int> weights_tz, int g,
+    std::vector<int> dst_tz, const std::string key,
+    std::shared_ptr<mkldnn::memory>& dst_memory_p,
+    std::vector<primitive>& pipeline,
+    const std::string &key_conv_pd,
+    std::shared_ptr<mkldnn::memory> src_memory_p,
+    std::shared_ptr<mkldnn::memory> user_src_memory_p,
+    std::shared_ptr<mkldnn::convolution_forward> conv_p,
+    std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd,
+    std::shared_ptr<platform::ConvMKLDNNHandler> handler,
+    bool force_fp32_output) const {
+      const T* input_data = input->data<T>();
+      bool is_INT8 = true;
       auto* scale_in = ctx.HasInput("Scale_in") ? ctx.Input<Tensor>("Scale_in") : nullptr;
       auto* scale_in_eltwise = ctx.HasInput("Scale_in_eltwise")? ctx.Input<Tensor>("Scale_in_eltwise") : nullptr;
       auto* scale_weights = ctx.HasInput("Scale_weights")? ctx.Input<Tensor>("Scale_weights") : nullptr;
@@ -381,8 +467,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
         sum_scale[0] = scale_out_data[0] / scale_in_eltwise_data[0];
       }
 
-        std::vector<primitive> pipeline;
-
       auto user_src_md = platform::MKLDNNMemDesc(
               {src_tz}, paddle::framework::ToMKLDNNDataType(input->type()), input->format());
       auto user_weights_md = platform::MKLDNNMemDesc(
@@ -517,32 +601,8 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 
         // push primitive to stream and wait until it's executed
       pipeline.push_back(*conv_p);
-        stream(stream::kind::eager).submit(pipeline).wait();
-
-        if(need_s8_to_u8){
-          output->mutable_data<uint8_t>(ctx.GetPlace());
-        }
-
-        output->set_layout(DataLayout::kMKLDNN);
-        output->set_format(GetMKLDNNFormat(*dst_memory_p));
-      } else {
-        if(src_memory_reorder_p){
-          pipeline.push_back(*src_memory_reorder_p);
-        }
-        pipeline.push_back(*conv_p);
-        stream(stream::kind::eager).submit(pipeline).wait();
-      
-        if (need_s8_to_u8) {
-          output->mutable_data<uint8_t>(ctx.GetPlace());
-        }
+    };
 
-        output->set_layout(DataLayout::kMKLDNN);
-        output->set_format(GetMKLDNNFormat(*dst_memory_p));
-      }
-    }
-  }
-
- private:
     void AppendKey(std::string& key, mkldnn::memory::dims& input_dims,    // NOLINT
                    mkldnn::memory::dims& weights_dims,  // NOLINT
                    std::vector<int>& strides,           // NOLINT