[OneDNN] Matmul kernel optimize for swin model (#59394)

zhanglirong1999 · web-flow · commit aa8ef889b910 · 2023-11-29T16:29:13.000+08:00
diff --git a/paddle/phi/kernels/fusion/onednn/fused_matmul_kernel.cc b/paddle/phi/kernels/fusion/onednn/fused_matmul_kernel.cc
@@ -129,23 +129,26 @@ class FusedMatmulOneDNNHandler
                                                 scale_y,
                                                 scale_in_eltwise,
                                                 scale_out,
-                                                force_fp32_output);
+                                                force_fp32_output,
+                                                out_ddims);
 
     this->AcquireForwardPrimitiveDescriptor(matmul_attrs, x_md, y_md, out_md);
   }
 
-  dnnl::primitive_attr CreateMatmulAttrs(const OneDNNContext &dev_ctx,
-                                         const DenseTensor *residual_data,
-                                         const float matmul_alpha,
-                                         const std::string &fuse_activation,
-                                         const float fuse_alpha,
-                                         const float fuse_beta,
-                                         const float fused_output_scale,
-                                         const float scale_x,
-                                         const float scale_y,
-                                         const float scale_in_eltwise,
-                                         const float scale_out,
-                                         const bool force_fp32_output) {
+  dnnl::primitive_attr CreateMatmulAttrs(
+      const OneDNNContext &dev_ctx,
+      const DenseTensor *residual_data,
+      const float matmul_alpha,
+      const std::string &fuse_activation,
+      const float fuse_alpha,
+      const float fuse_beta,
+      const float fused_output_scale,
+      const float scale_x,
+      const float scale_y,
+      const float scale_in_eltwise,
+      const float scale_out,
+      const bool force_fp32_output,
+      const std::vector<int64_t> &out_ddims) {
     dnnl::primitive_attr matmul_attrs;
     dnnl::post_ops post_operations;
 
@@ -164,9 +167,20 @@ class FusedMatmulOneDNNHandler
 
     if (residual_data) {
       auto residual_data_tz = vectorize(residual_data->dims());
-      auto residual_data_md = memory::desc(residual_data_tz,
-                                           funcs::OneDNNGetDataType<OT>(),
-                                           dnnl::memory::format_tag::any);
+      auto chosen_memory_format = funcs::OneDNNMemoryFormat::any;
+      dnnl::memory::desc residual_data_md;
+      if (residual_data_tz.size() == 4 && residual_data_tz[0] == 1 &&
+          residual_data_tz[1] > 1 && residual_data_tz[2] > 1 &&
+          residual_data_tz[3] > 1) {
+        chosen_memory_format = funcs::OneDNNMemoryFormat::nchw;
+        residual_data_md = memory::desc(
+            out_ddims, funcs::OneDNNGetDataType<OT>(), chosen_memory_format);
+      } else {
+        residual_data_md = memory::desc(residual_data_tz,
+                                        funcs::OneDNNGetDataType<OT>(),
+                                        chosen_memory_format);
+      }
+
       post_operations.append_binary(dnnl::algorithm::binary_add,
                                     residual_data_md);
       if (scale_in_eltwise != 0.0f) {
@@ -193,6 +207,30 @@ class FusedMatmulOneDNNHandler
         this->fwd_pd_->weights_desc(), funcs::to_void_cast<YT>(input_data));
   }
 
+  std::shared_ptr<dnnl::memory> AcquireSrcMemoryStride(
+      const DenseTensor *input) {
+    const XT *input_data = input->data<XT>();
+    std::shared_ptr<dnnl::memory> src_mem =
+        this->AcquireMemoryFromPrimitive(this->fwd_pd_->dst_desc());
+    auto residual_vec = vectorize(input->dims());
+    int IC = residual_vec[1];
+    int IH = residual_vec[2];
+    int IW = residual_vec[3];
+    size_t size = this->fwd_pd_->dst_desc().get_size() / sizeof(XT);
+    XT *dst = static_cast<XT *>(src_mem->get_data_handle());
+
+#if defined(_OPENMP)
+#pragma omp parallel for
+#endif
+    for (size_t i = 0; i < size; ++i) {
+      auto mod_i =
+          static_cast<int>(i - floor(i / (IC * IH * IW)) * (IC * IH * IW));
+      // Make 1*C*H*W to N*C*H*W to avoid broadcast overhead
+      dst[i] = input_data[mod_i];
+    }
+    return src_mem;
+  }
+
   std::shared_ptr<dnnl::memory> AcquireDstMemory(const OneDNNContext &dev_ctx,
                                                  DenseTensor *output) {
     // We cannot use base AcquireDstMemory as it makes an allocation request
@@ -263,7 +301,15 @@ void ExecuteFusedMatmul(const OneDNNContext &dev_ctx,
       {DNNL_ARG_DST, *dst_memory_p}};
 
   if (residual_data) {
-    const auto residual_data_memory_p = handler.AcquireSrcMemory(residual_data);
+    auto residual_data_vec = vectorize(residual_data->dims());
+    std::shared_ptr<dnnl::memory> residual_data_memory_p;
+    if (residual_data_vec.size() == 4 && residual_data_vec[0] == 1 &&
+        residual_data_vec[1] > 1 && residual_data_vec[2] > 1 &&
+        residual_data_vec[3] > 1) {
+      residual_data_memory_p = handler.AcquireSrcMemoryStride(residual_data);
+    } else {
+      residual_data_memory_p = handler.AcquireSrcMemory(residual_data);
+    }
     matmul_args.insert({DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1,
                         *residual_data_memory_p});
   }
diff --git a/test/ir/inference/test_mkldnn_matmul_elementwise_add_fuse_pass.py b/test/ir/inference/test_mkldnn_matmul_elementwise_add_fuse_pass.py
@@ -84,5 +84,68 @@ def test(self):
         )
 
 
+class TestMatmulElementwiseAddMkldnnFuse1CHWPass(PassAutoScanTest):
+    def sample_program_config(self, draw):
+        axis = draw(st.sampled_from([-1, 0, 1]))
+        matmul_as_x = draw(st.booleans())
+        batch_size = draw(st.integers(min_value=1, max_value=1))
+        channel = draw(st.sampled_from([16, 32, 64]))
+        input_dim = draw(st.sampled_from([16, 32, 64]))
+
+        def generate_input():
+            return np.random.random(
+                [batch_size, channel, input_dim, input_dim]
+            ).astype(np.float32)
+
+        matmul_op = OpConfig(
+            type='matmul',
+            inputs={'X': ['matmul_x'], 'Y': ['matmul_y']},
+            outputs={'Out': ['matmul_output']},
+            attrs={
+                'use_mkldnn': True,
+            },
+        )
+
+        if matmul_as_x:
+            inputs = {'X': ['matmul_output'], 'Y': ['elementwise_addend']}
+        else:
+            inputs = {'X': ['elementwise_addend'], 'Y': ['matmul_output']}
+
+        elt_add_op = OpConfig(
+            type='elementwise_add',
+            inputs=inputs,
+            outputs={'Out': ['elementwise_add_output']},
+            attrs={'axis': axis, 'use_mkldnn': True},
+        )
+
+        model_net = [matmul_op, elt_add_op]
+
+        program_config = ProgramConfig(
+            ops=model_net,
+            weights={},
+            inputs={
+                'matmul_x': TensorConfig(data_gen=partial(generate_input)),
+                'matmul_y': TensorConfig(data_gen=partial(generate_input)),
+                'elementwise_addend': TensorConfig(
+                    data_gen=partial(generate_input)
+                ),
+            },
+            outputs=['elementwise_add_output'],
+        )
+
+        return program_config
+
+    def sample_predictor_configs(self, program_config):
+        config = self.create_inference_config(
+            use_mkldnn=True, passes=['matmul_elementwise_add_mkldnn_fuse_pass']
+        )
+        yield config, ['fused_matmul'], (1e-5, 1e-5)
+
+    def test(self):
+        self.run_and_statis(
+            quant=False, passes=['matmul_elementwise_add_mkldnn_fuse_pass']
+        )
+
+
 if __name__ == '__main__':
     unittest.main()
