diff --git a/onnxruntime/core/providers/cpu/math/gemm.cc b/onnxruntime/core/providers/cpu/math/gemm.cc index fcb8e77761..0b43ec6a80 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.cc +++ b/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()), Gemm); + +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(b_shape[1]) : static_cast(b_shape[0]); + const size_t N = trans_b ? static_cast(b_shape[0]) : static_cast(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(), + trans_b ? K : N, + packed_b_data); + return true; +} + +template +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(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(c_data, N).transpose(); + } else if ((*c_shape)[1] == 1) { + // C is (M, 1) + output_mat.colwise() = ConstEigenVectorMap(c_data, M); + } else { + // C is (M, N), no broadcast needed. + output_mat = ConstEigenMatrixMapRowMajor(c_data, M, N); + } + } +} + +template +void Gemm::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(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::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 +Status Gemm::PrePack(const Tensor& /* tensor */, int /* input_idx */, bool& is_packed) { + is_packed = false; + return Status::OK(); +} + +template <> +Status Gemm::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 +void Gemm::ComputeActivation(T* y_data, size_t y_size, concurrency::ThreadPool* thread_pool) const { + if (activation_) { + std::unique_ptr> f(activation_->Copy()); + f->input = y_data; + f->output = y_data; + std::ptrdiff_t total_len = static_cast(y_size); + double cost = f->Cost(); + functors::ElementWiseRangedTransform* c(f.get()); + concurrency::ThreadPool::TryParallelFor( + thread_pool, total_len, + {static_cast(sizeof(T)), static_cast(sizeof(T)), cost}, + [c](std::ptrdiff_t first, std::ptrdiff_t last) { (*c)(first, last); }); + } +} + +template +Status Gemm::Compute(OpKernelContext* context) const { + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); + + const auto* A = context->Input(0); + const auto* B = context->Input(1); + const auto* C = context->Input(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(); + const T* c_data = C != nullptr ? C->Data() : nullptr; + const TensorShape* c_shape = C != nullptr ? &C->Shape() : nullptr; + + ComputeGemm(trans_A_, trans_B_, M, N, K, alpha_, A->Data(), B->Data(), beta_, + c_data, c_shape, y_data, thread_pool); + + ComputeActivation(y_data, M * N, thread_pool); + + return Status::OK(); +} + +template <> +Status Gemm::Compute(OpKernelContext* context) const { + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); + + const auto* A = context->Input(0); + const auto* B = packed_b_ ? nullptr : context->Input(1); + const auto* C = context->Input(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(); + + const float* c_data = C != nullptr ? C->Data() : nullptr; + const TensorShape* c_shape = C != nullptr ? &C->Shape() : nullptr; + + if (B) { + ComputeGemm(trans_A_, trans_B_, M, N, K, alpha_, A->Data(), B->Data(), 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(M), + static_cast(N), + static_cast(K), + alpha_, + A->Data(), + static_cast(trans_A_ != CblasNoTrans ? M : K), + packed_b_.get(), + c_data != nullptr ? beta_ : 0.0f, + y_data, + static_cast(N), + thread_pool); + } + + ComputeActivation(y_data, M * N, thread_pool); + + return Status::OK(); +} + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/gemm.h b/onnxruntime/core/providers/cpu/math/gemm.h index 6f16765e33..cc91d756f0 100644 --- a/onnxruntime/core/providers/cpu/math/gemm.h +++ b/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 class Gemm : public OpKernel { -private: - class CallWrapper{ -public: - CallWrapper(functors::ElementWiseRangedTransform* b1):b(b1){} - void operator()(std::ptrdiff_t first, std::ptrdiff_t last) const { - (*b)(first, last); - } -private: - functors::ElementWiseRangedTransform* b; -}; - public: - Gemm(const OpKernelInfo& info) : OpKernel(info) - { + Gemm(const OpKernelInfo& info) : OpKernel(info) { int64_t temp; ORT_ENFORCE(info.GetAttr("transA", &temp).IsOK()); trans_A_ = temp == 0 ? CblasNoTrans : CblasTrans; @@ -39,6 +25,10 @@ private: ORT_ENFORCE(info.GetAttr("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(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(c_data, N).transpose(); - } else if ((*c_shape)[1] == 1) { - // C is (M, 1) - output_mat.colwise() = ConstEigenVectorMap(c_data, M); - } else { - // C is (M, N), no broadcast needed. - output_mat = ConstEigenMatrixMapRowMajor(c_data, M, N); - } - } - - math::Gemm(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(0); - const auto* W = context->Input(1); - const auto* B = context->Input(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() : nullptr; - const TensorShape* b_shape = B != nullptr ? &B->Shape() : nullptr; - - T* y_data = Y->MutableData(); - - ComputeGemm(trans_A_, trans_B_, M, N, K, alpha_, X->Data(), W->Data(), beta_, - b_data, b_shape, - y_data, - thread_pool); - - if(activation_){ - std::unique_ptr> f(activation_->Copy()); - f->input = y_data; - f->output = y_data; - std::ptrdiff_t total_len = static_cast(M * N); - double cost = f->Cost(); - CallWrapper c(f.get()); - concurrency::ThreadPool::TryParallelFor(thread_pool, total_len, {static_cast(sizeof(T)), static_cast(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> activation_; + + void ComputeActivation(T* y_data, size_t y_size, concurrency::ThreadPool* thread_pool) const; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/math/gemm_matmul_common.h b/onnxruntime/core/providers/cpu/math/gemm_matmul_common.h new file mode 100644 index 0000000000..ffd229e825 --- /dev/null +++ b/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 diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index 508bdee3a6..adb4419d68 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/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::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(b_shape_[1]) - : static_cast(b_shape_[0]); - const size_t N = trans_b ? static_cast(b_shape_[0]) - : static_cast(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(), - static_cast(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(); } diff --git a/onnxruntime/test/providers/cpu/math/gemm_test.cc b/onnxruntime/test/providers/cpu/math/gemm_test.cc index 3ccbc6424f..36f2af804b 100644 --- a/onnxruntime/test/providers/cpu/math/gemm_test.cc +++ b/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("A", {2, 4}, {1.0f, 2.0f, 3.0f, 4.0f, -1.0f, -2.0f, -3.0f, -4.0f}); - test.AddInput("B", {4, 3}, std::vector(12, 1.0f)); + test.AddInput("B", {4, 3}, std::vector(12, 1.0f), b_is_initializer); test.AddInput("C", {3}, std::vector{1.0f, 2.0f, 3.0f}); test.AddOutput("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("B", {3, 4}, std::vector(12, 1.0f)); + test.AddInput("B", {3, 4}, std::vector(12, 1.0f), b_is_initializer); test.AddInput("C", {3}, std::vector(3, 1.0f)); test.AddOutput("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");