Prepacking in Gemm with merged logic for Matmul and Gemm on PackingB. (#5693)

Prepacking in Gemm with merged logic for Matmul and Gemm on PackingB.

onnxruntime/core/providers/cpu/math/gemm.cc

@@ -2,6 +2,10 @@
 // Licensed under the MIT License.
 
 #include "core/providers/cpu/math/gemm.h"
+#include "core/providers/cpu/math/gemm_matmul_common.h"
+#include "core/util/math_cpuonly.h"
+#include "gemm_helper.h"
+#include "core/mlas/inc/mlas.h"
 
 namespace onnxruntime {
 
@@ -34,4 +38,223 @@ ONNX_CPU_OPERATOR_KERNEL(
     13,
     KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
     Gemm<float>);
+
+bool GemmPackBFp32(const OpKernelInfo& info,
+                   const Tensor& tensor_b,
+                   bool trans_b,
+                   BufferUniquePtr& packed_b,
+                   TensorShape& b_shape) {
+  // Only handle the common case of a 2D weight matrix. Additional matrices
+  // could be handled by stacking the packed buffers.
+  if (tensor_b.Shape().NumDimensions() != 2) {
+    return false;
+  }
+  b_shape = tensor_b.Shape();
+
+  const size_t K = trans_b ? static_cast<size_t>(b_shape[1]) : static_cast<size_t>(b_shape[0]);
+  const size_t N = trans_b ? static_cast<size_t>(b_shape[0]) : static_cast<size_t>(b_shape[1]);
+
+  const size_t packed_b_size = MlasGemmPackBSize(N, K);
+  if (packed_b_size == 0) {
+    return false;
+  }
+
+  auto alloc = info.GetAllocator(0, OrtMemTypeDefault);
+  auto* packed_b_data = alloc->Alloc(packed_b_size);
+  packed_b = BufferUniquePtr(packed_b_data, BufferDeleter(alloc));
+  MlasGemmPackB(trans_b ? CblasTrans : CblasNoTrans,
+                N,
+                K,
+                tensor_b.Data<float>(),
+                trans_b ? K : N,
+                packed_b_data);
+  return true;
+}
+
+template <typename T>
+static void GemmBroadcastBias(int64_t M, int64_t N, float beta,
+                              const T* c_data, const TensorShape* c_shape,
+                              T* y_data) {
+  // Broadcast the bias as needed if bias is given
+  if (beta != 0 && c_data != nullptr) {
+    ORT_ENFORCE(c_shape != nullptr, "c_shape is required if c_data is provided");
+    auto output_mat = EigenMatrixMapRowMajor<T>(y_data, M, N);
+    if (c_shape->Size() == 1) {
+      // C is (), (1,) or (1, 1), set the scalar
+      output_mat.setConstant(*c_data);
+    } else if (c_shape->NumDimensions() == 1 || (*c_shape)[0] == 1) {
+      // C is (N,) or (1, N)
+      output_mat.rowwise() = ConstEigenVectorMap<T>(c_data, N).transpose();
+    } else if ((*c_shape)[1] == 1) {
+      // C is (M, 1)
+      output_mat.colwise() = ConstEigenVectorMap<T>(c_data, M);
+    } else {
+      // C is (M, N), no broadcast needed.
+      output_mat = ConstEigenMatrixMapRowMajor<T>(c_data, M, N);
+    }
+  }
+}
+
+template <typename T>
+void Gemm<T>::ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b,
+                          int64_t M, int64_t N, int64_t K,
+                          float alpha,
+                          const T* a_data, const T* b_data,
+                          float beta,
+                          const T* c_data, const TensorShape* c_shape,
+                          T* y_data,
+                          concurrency::ThreadPool* thread_pool) {
+  // if input is empty tensor, return directly as nothing need to be calculated.
+  if (M == 0 || N == 0)
+    return;
+
+  // Broadcast the bias as needed if bias is given
+  GemmBroadcastBias(M, N, beta, c_data, c_shape, y_data);
+
+  math::Gemm<T>(trans_a, trans_b,
+                M, N, K,
+                alpha,
+                a_data,
+                b_data,
+                // ideally we need to set the output buffer contents to 0 if bias is missing,
+                // but passing 0 for beta is cheaper and it will ignore any junk in the output buffer
+                c_data != nullptr ? beta : 0,
+                y_data,
+                thread_pool);
+}
+
+template void Gemm<float>::ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b,
+                                       int64_t M, int64_t N, int64_t K,
+                                       float alpha,
+                                       const float* a_data, const float* b_data,
+                                       float beta,
+                                       const float* c_data, const TensorShape* c_shape,
+                                       float* y_data,
+                                       concurrency::ThreadPool* thread_pool);
+
+template <typename T>
+Status Gemm<T>::PrePack(const Tensor& /* tensor */, int /* input_idx */, bool& is_packed) {
+  is_packed = false;
+  return Status::OK();
+}
+
+template <>
+Status Gemm<float>::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) {
+  is_packed = false;
+
+  // only pack Matrix B
+  if (input_idx == 1) {
+    is_packed = GemmPackBFp32(Info(), tensor, trans_B_ != CblasNoTrans, packed_b_, b_shape_);
+  }
+  return Status::OK();
+}
+
+template <typename T>
+void Gemm<T>::ComputeActivation(T* y_data, size_t y_size, concurrency::ThreadPool* thread_pool) const {
+  if (activation_) {
+    std::unique_ptr<functors::ElementWiseRangedTransform<T>> f(activation_->Copy());
+    f->input = y_data;
+    f->output = y_data;
+    std::ptrdiff_t total_len = static_cast<std::ptrdiff_t>(y_size);
+    double cost = f->Cost();
+    functors::ElementWiseRangedTransform<T>* c(f.get());
+    concurrency::ThreadPool::TryParallelFor(
+        thread_pool, total_len,
+        {static_cast<float>(sizeof(T)), static_cast<float>(sizeof(T)), cost},
+        [c](std::ptrdiff_t first, std::ptrdiff_t last) { (*c)(first, last); });
+  }
+}
+
+template <typename T>
+Status Gemm<T>::Compute(OpKernelContext* context) const {
+  concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
+
+  const auto* A = context->Input<Tensor>(0);
+  const auto* B = context->Input<Tensor>(1);
+  const auto* C = context->Input<Tensor>(2);
+
+  // Bias could be missing. Treat as scalar 0 if that is the case.
+  GemmHelper helper(A->Shape(), trans_A_ != CblasNoTrans, B->Shape(), trans_B_ != CblasNoTrans,
+                    C != nullptr ? C->Shape() : TensorShape({}));
+
+  if (!helper.State().IsOK())
+    return helper.State();
+
+  int64_t M = helper.M();
+  int64_t N = helper.N();
+  int64_t K = helper.K();
+
+  auto Y = context->Output(0, {M, N});
+
+  // if input is empty tensor, return as nothing need to be calculated and we've set the shape for the output
+  if (M == 0 || N == 0)
+    return Status::OK();
+
+  T* y_data = Y->MutableData<T>();
+  const T* c_data = C != nullptr ? C->Data<T>() : nullptr;
+  const TensorShape* c_shape = C != nullptr ? &C->Shape() : nullptr;
+
+  ComputeGemm(trans_A_, trans_B_, M, N, K, alpha_, A->Data<T>(), B->Data<T>(), beta_,
+              c_data, c_shape, y_data, thread_pool);
+
+  ComputeActivation(y_data, M * N, thread_pool);
+
+  return Status::OK();
+}
+
+template <>
+Status Gemm<float>::Compute(OpKernelContext* context) const {
+  concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
+
+  const auto* A = context->Input<Tensor>(0);
+  const auto* B = packed_b_ ? nullptr : context->Input<Tensor>(1);
+  const auto* C = context->Input<Tensor>(2);
+
+  // Bias could be missing. Treat as scalar 0 if that is the case.
+  GemmHelper helper(A->Shape(), trans_A_ != CblasNoTrans, B ? B->Shape() : b_shape_, trans_B_ != CblasNoTrans,
+                    C != nullptr ? C->Shape() : TensorShape({}));
+
+  if (!helper.State().IsOK())
+    return helper.State();
+
+  int64_t M = helper.M();
+  int64_t N = helper.N();
+  int64_t K = helper.K();
+
+  auto Y = context->Output(0, {M, N});
+
+  // if input is empty tensor, return as nothing need to be calculated and we've set the shape for the output
+  if (M == 0 || N == 0)
+    return Status::OK();
+
+  float* y_data = Y->MutableData<float>();
+
+  const float* c_data = C != nullptr ? C->Data<float>() : nullptr;
+  const TensorShape* c_shape = C != nullptr ? &C->Shape() : nullptr;
+
+  if (B) {
+    ComputeGemm(trans_A_, trans_B_, M, N, K, alpha_, A->Data<float>(), B->Data<float>(), beta_,
+                c_data, c_shape, y_data, thread_pool);
+  } else {
+    GemmBroadcastBias(M, N, beta_, c_data, c_shape, y_data);
+    MlasGemm(
+        trans_A_,
+        static_cast<size_t>(M),
+        static_cast<size_t>(N),
+        static_cast<size_t>(K),
+        alpha_,
+        A->Data<float>(),
+        static_cast<size_t>(trans_A_ != CblasNoTrans ? M : K),
+        packed_b_.get(),
+        c_data != nullptr ? beta_ : 0.0f,
+        y_data,
+        static_cast<size_t>(N),
+        thread_pool);
+  }
+
+  ComputeActivation(y_data, M * N, thread_pool);
+
+  return Status::OK();
+}
+
 }  // namespace onnxruntime

onnxruntime/core/providers/cpu/math/gemm.h

@@ -3,31 +3,17 @@
 
 #pragma once
 
-#include "core/common/common.h"
 #include "core/framework/op_kernel.h"
+#include "core/common/common.h"
 #include "core/util/math.h"
-#include "core/util/math_cpuonly.h"
-#include "gemm_helper.h"
 #include "core/providers/cpu/activation/activations.h"
 
 namespace onnxruntime {
 
 template <typename T>
 class Gemm : public OpKernel {
-private:
-    class CallWrapper{
-public:
-    CallWrapper(functors::ElementWiseRangedTransform<T>* b1):b(b1){}
-    void operator()(std::ptrdiff_t first, std::ptrdiff_t last) const {
-        (*b)(first, last);
-    }
-private:
-    functors::ElementWiseRangedTransform<T>* b;
-};
-
  public:
-  Gemm(const OpKernelInfo& info) : OpKernel(info)
-  {
+  Gemm(const OpKernelInfo& info) : OpKernel(info) {
     int64_t temp;
     ORT_ENFORCE(info.GetAttr<int64_t>("transA", &temp).IsOK());
     trans_A_ = temp == 0 ? CblasNoTrans : CblasTrans;
@@ -39,6 +25,10 @@ private:
     ORT_ENFORCE(info.GetAttr<float>("beta", &beta_).IsOK());
   }
 
+  Status Compute(OpKernelContext* context) const override;
+
+  Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override;
+
   static void ComputeGemm(CBLAS_TRANSPOSE trans_a, CBLAS_TRANSPOSE trans_b,
                           int64_t M, int64_t N, int64_t K,
                           float alpha,
@@ -46,86 +36,7 @@ private:
                           float beta,
                           const T* c_data, const TensorShape* c_shape,
                           T* y_data,
-                          concurrency::ThreadPool* thread_pool) {
-    // if input is empty tensor, return directly as nothing need to be calculated.
-    if (M == 0 || N == 0)
-      return;
-
-    // Broadcast the bias as needed if bias is given
-    if (beta != 0 && c_data != nullptr) {
-      ORT_ENFORCE(c_shape != nullptr, "c_shape is required if c_data is provided");
-      auto output_mat = EigenMatrixMapRowMajor<T>(y_data, M, N);
-      if (c_shape->Size() == 1) {
-        // C is (), (1,) or (1, 1), set the scalar
-        output_mat.setConstant(*c_data);
-      } else if (c_shape->NumDimensions() == 1 || (*c_shape)[0] == 1) {
-        // C is (N,) or (1, N)
-        output_mat.rowwise() = ConstEigenVectorMap<T>(c_data, N).transpose();
-      } else if ((*c_shape)[1] == 1) {
-        // C is (M, 1)
-        output_mat.colwise() = ConstEigenVectorMap<T>(c_data, M);
-      } else {
-        // C is (M, N), no broadcast needed.
-        output_mat = ConstEigenMatrixMapRowMajor<T>(c_data, M, N);
-      }
-    }
-
-    math::Gemm<T>(trans_a, trans_b,
-                  M, N, K,
-                  alpha,
-                  a_data,
-                  b_data,
-                  // ideally we need to set the output buffer contents to 0 if bias is missing,
-                  // but passing 0 for beta is cheaper and it will ignore any junk in the output buffer
-                  c_data != nullptr ? beta : 0,
-                  y_data,
-                  thread_pool);
-  }
-
-  Status Compute(OpKernelContext* context) const override {
-    concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool();
-
-    const auto* X = context->Input<Tensor>(0);
-    const auto* W = context->Input<Tensor>(1);
-    const auto* B = context->Input<Tensor>(2);
-    // Bias could be missing. Treat as scalar 0 if that is the case.
-    GemmHelper helper(X->Shape(), trans_A_ != CblasNoTrans, W->Shape(), trans_B_ != CblasNoTrans,
-                      B != nullptr ? B->Shape() : TensorShape({}));
-
-    if (!helper.State().IsOK())
-      return helper.State();
-
-    int64_t M = helper.M();
-    int64_t N = helper.N();
-    int64_t K = helper.K();
-
-    auto Y = context->Output(0, {M, N});
-
-    // if input is empty tensor, return as nothing need to be calculated and we've set the shape for the output
-    if (M == 0 || N == 0)
-      return Status::OK();
-
-    const T* b_data = B != nullptr ? B->Data<T>() : nullptr;
-    const TensorShape* b_shape = B != nullptr ? &B->Shape() : nullptr;
-
-    T* y_data = Y->MutableData<T>();
-
-    ComputeGemm(trans_A_, trans_B_, M, N, K, alpha_, X->Data<T>(), W->Data<T>(), beta_,
-                b_data, b_shape,
-                y_data,
-                thread_pool);
-
-    if(activation_){
-      std::unique_ptr<functors::ElementWiseRangedTransform<T>> f(activation_->Copy());
-      f->input = y_data;
-      f->output = y_data;
-      std::ptrdiff_t total_len = static_cast<std::ptrdiff_t>(M * N);
-      double cost = f->Cost();
-      CallWrapper c(f.get());
-      concurrency::ThreadPool::TryParallelFor(thread_pool, total_len, {static_cast<float>(sizeof(T)), static_cast<float>(sizeof(T)), cost}, c);
-    }
-    return Status::OK();
-  }
+                          concurrency::ThreadPool* thread_pool);
 
  private:
   CBLAS_TRANSPOSE trans_A_;
@@ -134,8 +45,13 @@ private:
   float beta_;
 
  protected:
-  // For fused gemm + activation  
+  TensorShape b_shape_;
+  BufferUniquePtr packed_b_;
+
+  // For fused gemm + activation
   std::unique_ptr<functors::ElementWiseRangedTransform<T>> activation_;
+
+  void ComputeActivation(T* y_data, size_t y_size, concurrency::ThreadPool* thread_pool) const;
 };
 
 }  // namespace onnxruntime

onnxruntime/core/providers/cpu/math/gemm_matmul_common.h

@@ -0,0 +1,16 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/framework/op_kernel.h"
+
+namespace onnxruntime {
+
+bool GemmPackBFp32(const OpKernelInfo& info,
+                   const Tensor& tensor_b,
+                   bool trans_b,
+                   BufferUniquePtr& packed_b,
+                   TensorShape& b_shape);
+
+};  // namespace onnxruntime

onnxruntime/core/providers/cpu/math/matmul.cc

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "core/providers/cpu/math/matmul.h"
+#include "core/providers/cpu/math/gemm_matmul_common.h"
 #include "core/providers/cpu/math/matmul_helper.h"
 #include "core/util/math.h"
 #include "core/util/math_cpuonly.h"
@@ -130,34 +131,7 @@ Status MatMul<float>::PrePack(const Tensor& tensor, int input_idx, bool& is_pack
 
   // only pack Matrix B
   if (input_idx == 1) {
-    // Only handle the common case of a 2D weight matrix. Additional matrices
-    // could be handled by stacking the packed buffers.
-    b_shape_ = tensor.Shape();
-    if (b_shape_.NumDimensions() != 2) {
-      return Status::OK();
-    }
-
-    const bool trans_b = trans_b_attr_ && b_shape_.NumDimensions() != 1;
-    const size_t K = trans_b ? static_cast<size_t>(b_shape_[1])
-                             : static_cast<size_t>(b_shape_[0]);
-    const size_t N = trans_b ? static_cast<size_t>(b_shape_[0])
-                             : static_cast<size_t>(b_shape_[1]);
-
-    const size_t packed_b_size = MlasGemmPackBSize(N, K);
-    if (packed_b_size == 0) {
-      return Status::OK();
-    }
-
-    auto alloc = Info().GetAllocator(0, OrtMemTypeDefault);
-    auto* packed_b_data = alloc->Alloc(packed_b_size);
-    packed_b_ = BufferUniquePtr(packed_b_data, BufferDeleter(alloc));
-    MlasGemmPackB(trans_b ? CblasTrans : CblasNoTrans,
-                  N,
-                  K,
-                  tensor.Data<float>(),
-                  static_cast<int>(trans_b ? K : N),
-                  packed_b_data);
-    is_packed = true;
+    is_packed = GemmPackBFp32(Info(), tensor, trans_b_attr_, packed_b_, b_shape_);
   }
   return Status::OK();
 }

onnxruntime/test/providers/cpu/math/gemm_test.cc

@@ -75,7 +75,7 @@ TEST(GemmOpTest, GemmNoTrans_f16) {
 }
 #endif
 
-TEST(GemmOpTest, GemmBroadcast) {
+static void TestGemmBroadcast(bool b_is_initializer) {
   OpTester test("Gemm");
 
   test.AddAttribute("transA", (int64_t)0);
@@ -86,7 +86,7 @@ TEST(GemmOpTest, GemmBroadcast) {
   test.AddInput<float>("A", {2, 4},
                        {1.0f, 2.0f, 3.0f, 4.0f,
                         -1.0f, -2.0f, -3.0f, -4.0f});
-  test.AddInput<float>("B", {4, 3}, std::vector<float>(12, 1.0f));
+  test.AddInput<float>("B", {4, 3}, std::vector<float>(12, 1.0f), b_is_initializer);
   test.AddInput<float>("C", {3}, std::vector<float>{1.0f, 2.0f, 3.0f});
   test.AddOutput<float>("Y", {2, 3},
                         {11.0f, 12.0f, 13.0f,
@@ -98,7 +98,15 @@ TEST(GemmOpTest, GemmBroadcast) {
 #endif
 }
 
-TEST(GemmOpTest, GemmTrans) {
+TEST(GemmOpTest, GemmBroadcast) {
+  TestGemmBroadcast(false);
+}
+
+TEST(GemmOpTest, GemmBroadcastBIsInitializer) {
+  TestGemmBroadcast(true);
+}
+
+static void TestGemmTrans(bool b_is_initializer) {
   OpTester test("Gemm");
 
   test.AddAttribute("transA", (int64_t)1);
@@ -111,18 +119,26 @@ TEST(GemmOpTest, GemmTrans) {
                         2.0f, -2.0f,
                         3.0f, -3.0f,
                         4.0f, -4.0f});
-  test.AddInput<float>("B", {3, 4}, std::vector<float>(12, 1.0f));
+  test.AddInput<float>("B", {3, 4}, std::vector<float>(12, 1.0f), b_is_initializer);
   test.AddInput<float>("C", {3}, std::vector<float>(3, 1.0f));
   test.AddOutput<float>("Y", {2, 3},
                         {11.0f, 11.0f, 11.0f,
                          -9.0f, -9.0f, -9.0f});
-#if defined(OPENVINO_CONFIG_GPU_FP16) || defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_MYRIAD) 
+#if defined(OPENVINO_CONFIG_GPU_FP16) || defined(OPENVINO_CONFIG_GPU_FP32) || defined(OPENVINO_CONFIG_MYRIAD)
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider});  // OpenVINO: Temporarily disabled due to accuracy issues
 #else
   test.Run();
 #endif
 }
 
+TEST(GemmOpTest, GemmTrans) {
+  TestGemmTrans(false);
+}
+
+TEST(GemmOpTest, GemmTransBIsInitializer) {
+  TestGemmTrans(true);
+}
+
 // NNAPI EP's GEMM only works as A*B', add case only B is transposed
 TEST(GemmOpTest, GemmTransB) {
   OpTester test("Gemm");