modify scale reuse to decrease map find

paddle/fluid/operators/conv_mkldnn_op.cc

@@ -373,26 +373,32 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
         src_tz, weights_tz, strides, paddings, dilations, groups,
         ctx.op().Output("Output"));
     const std::string key_conv_pd = key + "@conv_pd";
-    static std::unordered_map<std::string, std::vector<float>> scale_map;
+    static std::unordered_map<std::string, std::vector<std::vector<float>>> scale_map;
     //scale_map.insert({key_conv_pd,{1.0f}});
     //scale_map[key_conv_pd]={0.1f};
     bool scale_reuse = true;
-    auto scale_in_key = key + "@scale_in";
-    auto scale_weights_key = key + "@scale_weights";
-    auto scale_out_key = key + "@scale_out";
-    auto output_shift_scale_key = key + "@output_shift_scale";
-    auto sum_scale_key = key + "@sum_scale";
-    auto scale_in_eltwise_key = key + "@scale_in_eltwise";
+    //auto scale_in_key = key + "@scale_in";
+    //auto scale_weights_key = key + "@scale_weights";
+    //auto scale_out_key = key + "@scale_out";
+    //auto output_shift_scale_key = key + "@output_shift_scale";
+    //auto sum_scale_key = key + "@sum_scale";
+    //auto scale_in_eltwise_key = key + "@scale_in_eltwise";
     std::vector<float> scale_in_data;
     std::vector<float> scale_out_data;
     std::vector<float> scale_weights_data;
-    std::vector<float> scale_in_eltwise_data;
+    std::vector<float> scale_in_eltwise_data = {1.0f};
     std::vector<float> output_shift_scale;
     std::vector<float> sum_scale = {1.0f};
-    std::vector<float> none_scale = {0.0f};
+    std::vector<float> scale_bias_data = {1.0f};
+    std::vector<std::vector<float>> none_scale = {{0.0f}};
+    std::vector<std::vector<float>> scale_datas(7,{1.0f});
 
-    if (is_INT8 && GetScaleMap(scale_map, scale_in_key) == none_scale){
+//scale_in_data 0, scale_in_eltwise_data 1, scale_weights_data 2, scale_bias_data 3, scale_out_data 4, output_shift_scale 5, sum_scale 6
+
+    if (is_INT8 && GetScaleMap(scale_map, key) == none_scale){
         scale_reuse = false;
+    } else{
+        scale_datas = GetScaleMap(scale_map, key);
     }
 //std::cout<<"scale_reuse = "<<scale_reuse<<std::endl;
     if(is_INT8){
@@ -417,24 +423,30 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
             if(fuse_residual_conn){
                 scale_in_eltwise_data = {*(scale_in_eltwise->data<float>())};
                 sum_scale[0] = scale_out_data[0] / scale_in_eltwise_data[0];
-                SetScaleMap(scale_map, scale_in_eltwise_key, scale_in_eltwise_data);
+                //SetScaleMap(scale_map, scale_in_eltwise_key, scale_in_eltwise_data);
             }
 
             //scale reuse
-            SetScaleMap(scale_map, scale_in_key, scale_in_data);
-            SetScaleMap(scale_map, scale_weights_key, scale_weights_data);
-            SetScaleMap(scale_map, scale_out_key, scale_out_data);
-            SetScaleMap(scale_map, output_shift_scale_key, output_shift_scale);
-            SetScaleMap(scale_map, sum_scale_key, sum_scale);
+            scale_datas[0] = scale_in_data;
+            scale_datas[1] = scale_in_eltwise_data;
+            scale_datas[2] = scale_weights_data;
+            scale_datas[4] = scale_out_data;
+            scale_datas[5] = output_shift_scale;
+            scale_datas[6] = sum_scale;
+            //SetScaleMap(scale_map, key, scale_datas);
+            //SetScaleMap(scale_map, scale_weights_key, scale_weights_data);
+            //SetScaleMap(scale_map, scale_out_key, scale_out_data);
+            //SetScaleMap(scale_map, output_shift_scale_key, output_shift_scale);
+            //SetScaleMap(scale_map, sum_scale_key, sum_scale);
         } else{
-            scale_in_data = GetScaleMap(scale_map, scale_in_key);
-            scale_out_data = GetScaleMap(scale_map, scale_out_key);
-            scale_weights_data = GetScaleMap(scale_map, scale_weights_key);
+            scale_in_data = scale_datas[0];
+            scale_out_data = scale_datas[3];
+            scale_weights_data = scale_datas[2];
             if(fuse_residual_conn){
-                scale_in_eltwise_data = GetScaleMap(scale_map, scale_in_eltwise_key);
+                scale_in_eltwise_data = scale_datas[1];
             }
-            output_shift_scale = GetScaleMap(scale_map, output_shift_scale_key);
-            sum_scale = GetScaleMap(scale_map, sum_scale_key); 
+            output_shift_scale = scale_datas[5];
+            sum_scale = scale_datas[6]; 
             //printf("pause!!!");
         }
 
@@ -693,8 +705,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 
     // create convolution op primitive
     std::shared_ptr<mkldnn::convolution_forward> conv_p;
-    std::vector<float> scale_bias_data;
-    auto scale_bias_key = key + "@scale_bias";
+    //auto scale_bias_key = key + "@scale_bias";
     auto user_bias_md_key = key + "@user_bias_md";
     if (bias) {
       const float* bias_data = bias->data<float>();
@@ -718,9 +729,10 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
               for(int i=0; i<count; i++){
                   scale_bias_data[i] = scale_in_data[0] * scale_weights_data[i];
               }
-              SetScaleMap(scale_map, scale_bias_key, scale_bias_data);
+              scale_datas[3] = scale_bias_data;
+              //SetScaleMap(scale_map, scale_bias_key, scale_bias_data);
           } else{
-              scale_bias_data = GetScaleMap(scale_map, scale_bias_key);
+              scale_bias_data = scale_datas[3];
           }
           bias_memory_p =
               handler.AcquireBiasMemoryFromPrimitive(user_bias_memory_p, pipeline, is_test, is_INT8, scale_bias_data, mask_reorder);
@@ -735,6 +747,7 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
                                           dst_memory_p);
     }
 
+    SetScaleMap(scale_map, key, scale_datas);
 
     // push primitive to stream and wait until it's executed
     pipeline.push_back(*conv_p);
@@ -750,33 +763,33 @@ class ConvMKLDNNOpKernel : public paddle::framework::OpKernel<T> {
 
  private:
 
-    void SetScaleMap(std::unordered_map<std::string, std::vector<float>> &scale_map,
-                       const std::string& name, std::vector<float> scale_data) const {
+    void SetScaleMap(std::unordered_map<std::string, std::vector<std::vector<float>>> &scale_map,
+                       const std::string& name, std::vector<std::vector<float>> scale_datas) const {
       auto it = scale_map.find(name);
       if (it == scale_map.end()) {
-        scale_map[name] = scale_data;  // create new blob
+        scale_map[name] = scale_datas;  // create new blob
       } else {
-        (*it).second = scale_data;  // set data to existing blob
+        (*it).second = scale_datas;  // set data to existing blob
       }
       return;
     }
 
-    std::vector<float> GetScaleMap(std::unordered_map<std::string, std::vector<float>> scale_map,
+    std::vector<std::vector<float>> GetScaleMap(std::unordered_map<std::string, std::vector<std::vector<float>>> scale_map,
          const std::string& name) const {
       auto it = scale_map.find(name);
       if (it != scale_map.end()) {
         return (*it).second;
       }
-      return {0.0f};
+      return {{0.0f}};
     }
 
     void SetMdMap(std::unordered_map<std::string, std::shared_ptr<mkldnn::memory::desc>> &md_map,
-                       const std::string& name, std::shared_ptr<mkldnn::memory::desc> md) const {
+                       const std::string& name, std::shared_ptr<mkldnn::memory::desc> mds) const {
       auto it = md_map.find(name);
       if (it == md_map.end()) {
-        md_map[name] = md;  // create new blob
+        md_map[name] = mds;  // create new blob
       } else {
-        (*it).second = md;  // set data to existing blob
+        (*it).second = mds;  // set data to existing blob
       }
       return;
     }