mkl split gemm for better perf

Author: tensor-tang

CMakeLists.txt

@@ -136,6 +136,12 @@ else()
     set(THIRD_PARTY_BUILD_TYPE Release)
 endif()
 
+if(WITH_MKL)
+  option(MKL_SPLIT_GEMM "PaddlePaddle MKL gemm would split to small ones" OFF)
+  if (MKL_SPLIT_GEMM)
+    add_definitions(-DPADDLE_MKL_SPLIT_GEMM)
+  endif()
+endif()
 set(WITH_MKLML ${WITH_MKL})
 if (NOT DEFINED WITH_MKLDNN)
     if (WITH_MKL AND AVX2_FOUND)

paddle/fluid/operators/math/blas_impl.h

@@ -209,8 +209,23 @@ void Blas<platform::CPUDeviceContext>::GEMM(CBLAS_TRANSPOSE transA,
                        &beta, C, &ldc);
   } else {
 #endif
-    CBlas<T>::GEMM(CblasRowMajor, transA, transB, M, N, K, alpha, A, lda, B,
-                   ldb, beta, C, ldc);
+
+#ifdef PADDLE_MKL_SPLIT_GEMM
+    constexpr int bs = 2;
+    if (M % bs == 0 && transA == CblasNoTrans && transB == CblasNoTrans) {
+      for (int off = 0; off < M; off += bs) {
+        CBlas<T>::GEMM(CblasRowMajor, CblasNoTrans, CblasNoTrans, off, N, K,
+                       alpha, A + off * lda, lda, B, ldb, beta, C + off * ldb,
+                       ldc);
+      }
+    } else {
+#endif
+      CBlas<T>::GEMM(CblasRowMajor, transA, transB, M, N, K, alpha, A, lda, B,
+                     ldb, beta, C, ldc);
+#ifdef PADDLE_MKL_SPLIT_GEMM
+    }
+#endif
+
 #ifdef PADDLE_WITH_LIBXSMM
   }
 #endif