adjust structure for PR, peel int8 from fp32

paddle/fluid/operators/conv_mkldnn_op.cc

@@ -296,6 +296,203 @@ template <typename T>
 class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
  public:
   void Compute(const paddle::framework::ExecutionContext& ctx) const override {
+
+    bool is_INT8 = ctx.HasInput("Scale_in")? true : false;
+
+    if(!is_INT8){
+      PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
+                    "It must use CPUPlace.");
+
+      const bool is_test = ctx.Attr<bool>("is_test");
+
+      auto& dev_ctx =
+          ctx.template device_context<paddle::platform::MKLDNNDeviceContext>();
+      const auto& mkldnn_engine = dev_ctx.GetEngine();
+
+      auto* input = ctx.Input<Tensor>("Input");
+      auto* filter = ctx.Input<Tensor>("Filter");
+      auto* bias = ctx.HasInput("Bias") ? ctx.Input<Tensor>("Bias") : nullptr;
+      auto* output = ctx.Output<Tensor>("Output");
+
+      PADDLE_ENFORCE(input->layout() == DataLayout::kMKLDNN &&
+                         input->format() != memory::format::format_undef,
+                     "Wrong layout/format set for Input tensor");
+      PADDLE_ENFORCE(filter->layout() == DataLayout::kMKLDNN &&
+                         filter->format() != memory::format::format_undef,
+                     "Wrong layout/format set for Filter tensor");
+      PADDLE_ENFORCE(input->dims().size() == 4,
+                     "Input must be with 4 dimensions, i.e. NCHW");
+      PADDLE_ENFORCE(filter->dims().size() == 4,
+                     "Filter must be with 4 dimensions, i.e. OIHW");
+      if (bias) {
+        PADDLE_ENFORCE(bias->layout() == DataLayout::kMKLDNN &&
+                           bias->format() != memory::format::format_undef,
+                       "Wrong layout/format set for Bias tensor");
+        PADDLE_ENFORCE(bias->dims().size() == 1,
+                       "Bias must only have 1 dimension, i.e. X");
+      }
+
+      std::vector<int> strides = ctx.Attr<std::vector<int>>("strides");
+      std::vector<int> paddings = ctx.Attr<std::vector<int>>("paddings");
+      std::vector<int> dilations = ctx.Attr<std::vector<int>>("dilations");
+      bool fuse_relu = ctx.Attr<bool>("fuse_relu");
+      bool fuse_residual_conn = ctx.Attr<bool>("fuse_residual_connection");
+      int groups = ctx.Attr<int>("groups");
+
+      // TODO(tpatejko): add support for dilation
+      PADDLE_ENFORCE(
+          dilations.size() == 2 && dilations[0] == 1 && dilations[1] == 1,
+          "dilation in convolution is not implemented yet");
+
+      const T* input_data = input->data<T>();
+      const T* filter_data = filter->data<T>();
+
+      std::vector<int> src_tz = paddle::framework::vectorize2int(input->dims());
+      std::vector<int> weights_tz =
+          paddle::framework::vectorize2int(filter->dims());
+      int g = std::max(groups, 1);
+      if (g > 1) {
+        int o = weights_tz[0];
+        int i = weights_tz[1];
+        int h = weights_tz[2];
+        int w = weights_tz[3];
+        weights_tz.resize(5);
+        weights_tz[0] = g;
+        weights_tz[1] = o / g;
+        weights_tz[2] = i;
+        weights_tz[3] = h;
+        weights_tz[4] = w;
+      }
+      std::vector<int> dst_tz = paddle::framework::vectorize2int(output->dims());
+
+      // Get unique name for storing MKLDNN primitives
+      const std::string key = ConvMKLDNNHandler::GetHash(
+          src_tz, weights_tz, strides, paddings, dilations, groups,
+          ctx.op().Output("Output"));
+      const std::string key_conv_pd = key + "@conv_pd";
+
+      std::vector<primitive> pipeline;
+
+      auto user_src_md = platform::MKLDNNMemDesc(
+          {src_tz}, platform::MKLDNNGetDataType<T>(), input->format());
+      auto user_weights_md = platform::MKLDNNMemDesc(
+          {weights_tz}, platform::MKLDNNGetDataType<T>(),
+          (g == 1) ? filter->format() : mkldnn::memory::format::goihw);
+
+      /* create memory descriptor for convolution without specified format
+       * ('any') which lets a primitive (convolution in this case) choose
+       * the memory format preferred for best performance
+       */
+      std::string data_format = ctx.Attr<std::string>("data_format");
+      auto chosen_memory_format =
+          platform::data_format_to_memory_format(data_format);
+
+      auto src_md = platform::MKLDNNMemDesc(
+          src_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
+      auto weights_md = platform::MKLDNNMemDesc(
+          weights_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
+      std::vector<int> bias_tz;  // TODO(mgallus): avoid empty vector creation.
+                                 // Currently used whenever bias is != nullptr.
+      auto dst_md = platform::MKLDNNMemDesc(
+          dst_tz, platform::MKLDNNGetDataType<T>(), chosen_memory_format);
+
+      // create a conv primitive descriptor and save it for usage in backward
+      std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd;
+      if (bias) {
+        bias_tz = paddle::framework::vectorize2int(bias->dims());
+        auto bias_md = platform::MKLDNNMemDesc(
+            bias_tz, platform::MKLDNNGetDataType<T>(), memory::format::x);
+        conv_pd = ConvFwdPrimitiveDesc(src_md, weights_md, bias_md, dst_md,
+                                       strides, paddings, mkldnn_engine,
+                                       fuse_relu, fuse_residual_conn);
+      } else {
+        conv_pd =
+            ConvFwdPrimitiveDesc(src_md, weights_md, dst_md, strides, paddings,
+                                 mkldnn_engine, fuse_relu, fuse_residual_conn);
+      }
+      // Save conv_pd/src_memory/weights_memory for backward pass
+      dev_ctx.SetBlob(key_conv_pd, conv_pd);
+
+      ConvMKLDNNHandler handler(conv_pd, dev_ctx, mkldnn_engine, key);
+
+      // create mkldnn memory from input tensors (data/weights)
+      auto user_src_memory_p =
+          handler.AcquireSrcMemory(user_src_md, to_void_cast<T>(input_data));
+      auto user_weights_memory_p = handler.AcquireWeightsMemory(
+          user_weights_md, to_void_cast<T>(filter_data));
+
+      // create reorder primitive if the input format is not the preferred one
+      auto src_memory_p =
+          handler.AcquireSrcMemoryFromPrimitive(user_src_memory_p, pipeline);
+      auto weights_memory_p = handler.AcquireWeightsMemoryFromPrimitive(
+          user_weights_memory_p, pipeline, is_test);
+
+      std::shared_ptr<mkldnn::memory> dst_memory_p;
+
+      if (fuse_residual_conn) {
+        auto residual_param = ctx.Input<Tensor>("ResidualData");
+        auto residual_param_data = residual_param->data<T>();
+
+        PADDLE_ENFORCE(
+            residual_param_data != nullptr,
+            "Provide data if you want MKLDNN conv+elementwise_add fusion");
+        PADDLE_ENFORCE_EQ(output->dims(), residual_param->dims(),
+                          "Output and elementwise parameter need to have the "
+                          "same dimension sizes");
+
+        if (residual_param->format() != handler.GetDstFormat()) {
+          auto output_data =
+              output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
+          auto residual_data_tz =
+              paddle::framework::vectorize2int(residual_param->dims());
+          auto residual_data_type =
+              paddle::framework::ToMKLDNNDataType(residual_param->type());
+
+          auto user_residual_md = platform::MKLDNNMemDesc(
+              residual_data_tz, residual_data_type, residual_param->format());
+          auto user_residual_memory_p = handler.AcquireResidualDataMemory(
+              user_residual_md, to_void_cast<T>(residual_param_data));
+
+          dst_memory_p = handler.AcquireDstMemoryFromResidualDataMemory(
+              user_residual_memory_p, to_void_cast<T>(output_data), pipeline);
+        } else {
+          output->ShareDataWith(*residual_param);
+          auto output_data = output->mutable_data<T>(ctx.GetPlace());
+          dst_memory_p =
+              handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
+        }
+      } else {
+        auto output_data =
+            output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
+        dst_memory_p =
+            handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
+      }
+
+      // create convolution op primitive
+      std::shared_ptr<mkldnn::convolution_forward> conv_p;
+      if (bias) {
+        const T* bias_data = bias->data<T>();
+        auto user_bias_md = platform::MKLDNNMemDesc(
+            {bias_tz}, platform::MKLDNNGetDataType<T>(), memory::format::x);
+        auto user_bias_memory_p =
+            handler.AcquireBiasMemory(user_bias_md, to_void_cast<T>(bias_data));
+
+        auto bias_memory_p =
+            handler.AcquireBiasMemoryFromPrimitive(user_bias_memory_p, pipeline);
+        conv_p = handler.AcquireConvolution(src_memory_p, weights_memory_p,
+                                            bias_memory_p, dst_memory_p);
+      } else {
+        conv_p = handler.AcquireConvolution(src_memory_p, weights_memory_p,
+                                            dst_memory_p);
+      }
+
+      // 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{
       PADDLE_ENFORCE(paddle::platform::is_cpu_place(ctx.GetPlace()),
                      "It must use CPUPlace.");
       const bool is_test = ctx.Attr<bool>("is_test");
@@ -314,8 +511,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
       auto* scale_weights = ctx.HasInput("Scale_weights")? ctx.Input<Tensor>("Scale_weights") : nullptr;
       auto* scale_out = ctx.HasInput("Scale_out")? ctx.Input<Tensor>("Scale_out") : nullptr;
 
-    bool is_INT8 = ctx.HasInput("Scale_in")? true : false;
-    bool is_multi_channel = (is_INT8 && scale_weights->memory_size() > 1) ? true : false;
+      bool is_multi_channel = (scale_weights->memory_size() > 1) ? true : false;
 
       PADDLE_ENFORCE(input->layout() == DataLayout::kMKLDNN &&
                          input->format() != memory::format::format_undef,
@@ -374,8 +570,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
           ctx.op().Output("Output"));
       const std::string key_conv_pd = key + "@conv_pd";
       static std::unordered_map<std::string, std::vector<float>> scale_map;
-    //scale_map.insert({key_conv_pd,{1.0f}});
-    //scale_map[key_conv_pd]={0.1f};
+
       bool scale_reuse = false;
       auto scale_in_key = key + "@scale_in";
       auto scale_weights_key = key + "@scale_weights";
@@ -391,11 +586,10 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
       std::vector<float> sum_scale = {1.0f};
       std::vector<float> none_scale = {0};
 
-    if (is_INT8 && GetScaleMap(scale_map, scale_in_key) == none_scale){
+      if (GetScaleMap(scale_map, scale_in_key) == none_scale){
           scale_reuse = true;
       }
 //std::cout<<"scale_reuse = "<<scale_reuse<<std::endl;
-    if(is_INT8){
       if(scale_reuse){
 //std::cout<<"load scale!!!!!!!!"<<std::endl;
         int count = is_multi_channel? (g>1? weights_tz[1]*weights_tz[0] : weights_tz[0]) : 1; 
@@ -438,9 +632,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
           //printf("pause!!!");
       }
 
-    }
-
-
       std::vector<primitive> pipeline;
       auto user_src_md = platform::MKLDNNMemDesc(
               {src_tz}, paddle::framework::ToMKLDNNDataType(input->type()), input->format());
@@ -458,7 +649,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 
       std::shared_ptr<mkldnn::convolution_forward::primitive_desc> conv_pd;
       auto bias_tz = paddle::framework::vectorize2int(bias->dims());
-    if(is_INT8){
+
       auto src_md = platform::MKLDNNMemDesc(
           src_tz, memory::data_type::u8, chosen_memory_format);
       auto weights_md = platform::MKLDNNMemDesc(
@@ -487,27 +678,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
                                    mkldnn_engine, fuse_relu, fuse_residual_conn,
                                    output_shift_scale, sum_scale[0], is_test);
       }
-    } else{
-        auto src_md = platform::MKLDNNMemDesc(
-            src_tz, platform::MKLDNNGetDataType<float>(), chosen_memory_format);
-        auto weights_md = platform::MKLDNNMemDesc(
-            weights_tz, platform::MKLDNNGetDataType<float>(),
-            (g == 1) ? chosen_memory_format : mkldnn::memory::format::goihw);
-        auto dst_md = platform::MKLDNNMemDesc(
-            dst_tz, platform::MKLDNNGetDataType<float>(), chosen_memory_format);
-        // create a conv primitive descriptor and save it for usage in backward
-        if (bias) {
-            auto bias_md = platform::MKLDNNMemDesc(
-                bias_tz, platform::MKLDNNGetDataType<float>(), memory::format::x);
-                conv_pd = ConvFwdPrimitiveDesc(src_md, weights_md, bias_md, dst_md,
-                                               strides, paddings, mkldnn_engine,
-                                               fuse_relu, fuse_residual_conn, is_test);
-        } else {
-                conv_pd =
-                    ConvFwdPrimitiveDesc(src_md, weights_md, dst_md, strides, paddings,
-                                         mkldnn_engine, fuse_relu, fuse_residual_conn, is_test);
-        }
-    }
       // Save conv_pd/src_memory/weights_memory for backward pass
       dev_ctx.SetBlob(key_conv_pd, conv_pd);
 
@@ -524,14 +694,9 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
           handler.AcquireSrcMemoryFromPrimitive(user_src_memory_p, pipeline);
           
       std::shared_ptr<mkldnn::memory> weights_memory_p;
-    if(is_INT8){
       int mask_reorder = is_multi_channel? ((g!= 1) ? (1<<1)+(1<<0) : 1<<0) : 0;
       weights_memory_p = handler.AcquireWeightsMemoryFromPrimitive(
           user_weights_memory_p, pipeline, is_test, is_INT8, scale_weights_data, mask_reorder);
-    } else{
-        weights_memory_p = handler.AcquireWeightsMemoryFromPrimitive(
-            user_weights_memory_p, pipeline, is_test);
-    }
 
       std::shared_ptr<mkldnn::memory> dst_memory_p;
       bool need_s8_to_u8 = false;
@@ -548,7 +713,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
               paddle::framework::ToMKLDNNDataType(residual_param->type());
           auto user_residual_md = platform::MKLDNNMemDesc(
             residual_data_tz, residual_data_type, residual_param->format());
-            if(is_INT8){
             if(residual_dt == mkldnn::memory::data_type::u8){
               auto residual_param_data = residual_param->data<uint8_t>();
               auto user_residual_memory_p = handler.AcquireResidualDataMemory(
@@ -574,23 +738,9 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
               if(fuse_relu)
                 need_s8_to_u8 = true;
             }
-            } else{
-                auto residual_param_data = residual_param->data<T>();
-                auto user_residual_memory_p = handler.AcquireResidualDataMemory(
-                    user_residual_md, to_void_cast<T>(residual_param_data));
-                PADDLE_ENFORCE(
-                      residual_param_data != nullptr,
-                      "Provide data if you want MKLDNN conv+elementwise_add fusion");
-                 auto output_data =
-                     output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
-                 dst_memory_p = handler.AcquireDstMemoryFromResidualDataMemory(
-                      user_residual_memory_p, to_void_cast<T>(output_data), pipeline);
-            }
         } else {
           output->ShareDataWith(*residual_param);
-             if(is_INT8){
           if(residual_dt == mkldnn::memory::data_type::u8){
-
             uint8_t* output_data = output->mutable_data<uint8_t>(ctx.GetPlace());
             dst_memory_p =
                 handler.AcquireDstMemoryFromPrimitive(to_void_cast<uint8_t>(output_data));
@@ -601,14 +751,8 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
             if(fuse_relu)
               need_s8_to_u8 = true;
           }
-             } else{
-                  auto output_data = output->mutable_data<T>(ctx.GetPlace());
-                  dst_memory_p =
-                      handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));               
-             }
         }
       } else {
-        if(is_INT8){
         if(fuse_relu){
           uint8_t* output_data = output->mutable_data<uint8_t>(ctx.GetPlace(), handler.GetDstMemorySize());
           dst_memory_p =
@@ -618,12 +762,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
           dst_memory_p =
               handler.AcquireDstMemoryFromPrimitive(to_void_cast<int8_t>(output_data));
         }
-        } else{
-        auto output_data =
-            output->mutable_data<T>(ctx.GetPlace(), handler.GetDstMemorySize());
-        dst_memory_p =
-            handler.AcquireDstMemoryFromPrimitive(to_void_cast<T>(output_data));
-        }
       }
 
       // create convolution op primitive
@@ -637,7 +775,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
         auto user_bias_memory_p =
             handler.AcquireBiasMemory(user_bias_md, to_void_cast<float>(bias_data));
         std::shared_ptr<mkldnn::memory>  bias_memory_p;
-      if(is_INT8){
         int mask_reorder = is_multi_channel? 1<<0 : 1;
         if(scale_reuse){
             int count = is_multi_channel? (g>1? weights_tz[1]*weights_tz[0] : weights_tz[0]) : 1;
@@ -652,10 +789,6 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
         }
         bias_memory_p =
             handler.AcquireBiasMemoryFromPrimitive(user_bias_memory_p, pipeline, is_test, is_INT8, scale_bias_data, mask_reorder);
-      } else{
-          bias_memory_p =
-              handler.AcquireBiasMemoryFromPrimitive(user_bias_memory_p, pipeline);
-      } 
         conv_p = handler.AcquireConvolution(src_memory_p, weights_memory_p,
                                             bias_memory_p, dst_memory_p);
       } else {
@@ -675,6 +808,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
       output->set_layout(DataLayout::kMKLDNN);
       output->set_format(GetMKLDNNFormat(*dst_memory_p));
     }
+  }
 
  private:
 
@@ -780,7 +914,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
                          const memory::desc& dst, const std::vector<int>& strides,
                          const std::vector<int>& paddings,
                          const mkldnn::engine& engine, const bool fuse_relu,
-                         const bool fuse_residual_conn, bool is_test) const{
+                         const bool fuse_residual_conn, bool is_test=false) const{
       memory::dims stride_dims = {strides[0], strides[1]};
       memory::dims padding_dims = {paddings[0], paddings[1]};
  
@@ -834,7 +968,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
                          const std::vector<int>& strides,
                          const std::vector<int>& paddings,
                          const mkldnn::engine& engine, const bool fuse_relu,
-                         const bool fuse_residual_conn, bool is_test) const{
+                         const bool fuse_residual_conn, bool is_test=false) const{
       memory::dims stride_dims = {strides[0], strides[1]};
       memory::dims padding_dims = {paddings[0], paddings[1]};