From b2994492af9851c0a54842e0a7a0b645fdb06b22 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Wed, 16 Sep 2020 08:36:27 -0700 Subject: [PATCH] MLAS: add sgemm weight prepacking (#5183) Add support to MLAS to prepack weights for the float GEMM. Support for prepacking has been added to MatMul and Attention for this release. --- onnxruntime/contrib_ops/cpu/bert/attention.cc | 151 ++++- onnxruntime/contrib_ops/cpu/bert/attention.h | 21 - .../quantization/attention_quantization.h | 2 +- onnxruntime/core/mlas/inc/mlas.h | 35 + onnxruntime/core/mlas/lib/mlasi.h | 2 + onnxruntime/core/mlas/lib/qgemm.cpp | 16 +- onnxruntime/core/mlas/lib/sgemm.cpp | 607 +++++++++++++----- onnxruntime/core/providers/cpu/math/matmul.cc | 100 +++ .../test/contrib_ops/attention_op_test.cc | 34 +- onnxruntime/test/mlas/unittest.cpp | 99 ++- .../test/providers/cpu/math/matmul_test.cc | 15 +- 11 files changed, 845 insertions(+), 237 deletions(-) delete mode 100644 onnxruntime/contrib_ops/cpu/bert/attention.h diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.cc b/onnxruntime/contrib_ops/cpu/bert/attention.cc index 1da8aa411d..0ec32bc376 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/attention.cc @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -#include "attention.h" +#include "attention_cpu_base.h" #include "attention_helper.h" #include "core/framework/tensorprotoutils.h" #include "core/graph/onnx_protobuf.h" @@ -14,13 +14,33 @@ using onnxruntime::concurrency::ThreadPool; namespace onnxruntime { namespace contrib { -// These ops are internal-only, so register outside of onnx -#define REGISTER_KERNEL_TYPED(T) \ - ONNX_OPERATOR_TYPED_KERNEL_EX(Attention, kMSDomain, 1, T, kCpuExecutionProvider, \ - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ - Attention); -REGISTER_KERNEL_TYPED(float) +template +class Attention : public OpKernel, public AttentionCPUBase { + public: + explicit Attention(const OpKernelInfo& info); + + Status Compute(OpKernelContext* context) const override; +#if !defined(USE_MKLML_FOR_BLAS) + Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; +#endif + + private: + BufferUniquePtr packed_weights_; + size_t packed_weights_size_; + TensorShape weight_shape_; +}; + +// These ops are internal-only, so register outside of onnx +ONNX_OPERATOR_TYPED_KERNEL_EX( + Attention, + kMSDomain, + 1, + float, + kCpuExecutionProvider, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::GetTensorType()), + Attention); AttentionBase::AttentionBase(const OpKernelInfo& info) { int64_t num_heads = 0; @@ -116,8 +136,7 @@ Status AttentionBase::CheckInputs(const TensorShape& input_shape, if ((static_cast(mask_dims[0]) == batch_size || static_cast(mask_dims[0]) == 1) && static_cast(mask_dims[1]) == 1) { // Mask will have same value after propogation, which has same effect as no mask. mask_index = nullptr; - } - else { + } else { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Inputs 'mask_index' with raw attention mask shall have shape batch_size x (past_sequence_length + sequence_length)"); } } @@ -159,15 +178,68 @@ template Attention::Attention(const OpKernelInfo& info) : OpKernel(info), AttentionCPUBase(info) { } +#if !defined(USE_MKLML_FOR_BLAS) + +template +Status Attention::PrePack(const Tensor& weights, int input_idx, bool& is_packed) { + is_packed = false; + + if (1 != input_idx) { + return Status::OK(); + } + + weight_shape_ = weights.Shape(); + const auto& weights_dims = weight_shape_.GetDims(); + if (weights_dims.size() != 2) { + return Status::OK(); + } + + const size_t hidden_size = static_cast(weights_dims[0]); + const size_t hidden_size_x3 = static_cast(weights_dims[1]); + const size_t head_size = hidden_size / num_heads_; + + // Bail out if the weights shape has an expected shape. + if ((hidden_size == 0) || ((hidden_size % num_heads_) != 0) || (hidden_size_x3 != 3 * hidden_size)) { + return Status::OK(); + } + + const auto* weights_data = weights.Data(); + + packed_weights_size_ = MlasGemmPackBSize(head_size, hidden_size); + if (packed_weights_size_ == 0) { + return Status::OK(); + } + + const size_t loop_len = 3 * num_heads_; + auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); + auto* packed_weights_data = static_cast(alloc->Alloc(packed_weights_size_ * loop_len)); + packed_weights_ = BufferUniquePtr(packed_weights_data, BufferDeleter(alloc)); + + for (size_t i = 0; i < loop_len; i++) { + MlasGemmPackB(CblasNoTrans, head_size, hidden_size, weights_data, hidden_size_x3, packed_weights_data); + packed_weights_data += packed_weights_size_; + weights_data += head_size; + } + + is_packed = true; + return Status::OK(); +} + +#endif + template Status Attention::Compute(OpKernelContext* context) const { const Tensor* input = context->Input(0); - const Tensor* weights = context->Input(1); + const Tensor* weights = packed_weights_ ? nullptr : context->Input(1); const Tensor* bias = context->Input(2); const Tensor* mask_index = context->Input(3); const Tensor* past = context->Input(4); - ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), weights->Shape(), bias->Shape(), mask_index, past)); + ORT_RETURN_IF_ERROR(CheckInputs(input->Shape(), + packed_weights_ ? weight_shape_ : weights->Shape(), + bias->Shape(), + mask_index, + past)); const auto& shape = input->Shape().GetDims(); const int batch_size = static_cast(shape[0]); @@ -187,17 +259,17 @@ Status Attention::Compute(OpKernelContext* context) const { // gemm_data(BS, 3NH) = input(BS, NH) x weights(NH, 3NH) + bias(3NH) auto gemm_data = allocator->Alloc(SafeInt(batch_size) * sequence_length * 3 * hidden_size * element_size); BufferUniquePtr gemm_buffer(gemm_data, BufferDeleter(allocator)); + auto Q = reinterpret_cast(gemm_data); auto K = Q + batch_size * sequence_length * hidden_size; auto V = K + batch_size * sequence_length * hidden_size; - T* QKV[3] = {Q, K, V}; { const int loop_len = 3 * batch_size * num_heads_; - const auto input_data = input->template Data(); - const auto weights_data = weights->template Data(); - const auto bias_data = bias->template Data(); + const auto* input_data = input->template Data(); + const auto* weights_data = packed_weights_ ? nullptr : weights->template Data(); + const auto* bias_data = bias->template Data(); const double cost = static_cast(sequence_length) * static_cast(head_size) * static_cast(hidden_size); @@ -224,22 +296,39 @@ Status Attention::Compute(OpKernelContext* context) const { // A: input (BxSxNxH) (B.)S x NH S x NH // B: weights (NxHx3xNxH) NH x (3.N.)H NH x H // C: QKV[qkv_index] (3xBxNxSxH) (3.B.N.)S x H S x H - - math::GemmEx(CblasNoTrans, // TransA = no - CblasNoTrans, // TransB = no - sequence_length, // M = S - head_size, // N = H - hidden_size, // K = NH - 1.0f, // alpha - input_data + input_offset, // A - hidden_size, // lda = NH - weights_data + weights_offset, // B - 3 * hidden_size, // ldb = 3NH - 1.0f, // beta - qkv_dest + qkv_offset, // C - head_size, // ldc - nullptr // use single-thread - ); + if (packed_weights_) { + const auto* packed_weight = + static_cast(packed_weights_.get()) + packed_weights_size_ * (weights_offset / head_size); + MlasGemm( + CblasNoTrans, // TransA = no + sequence_length, // M = S + head_size, // N = H + hidden_size, // K = NH + 1.0f, // alpha + input_data + input_offset, // A + hidden_size, // lda = NH + packed_weight, // B + 1.0f, // beta + qkv_dest + qkv_offset, // C + head_size, // ldc + nullptr); // use single-thread + } else { + math::GemmEx(CblasNoTrans, // TransA = no + CblasNoTrans, // TransB = no + sequence_length, // M = S + head_size, // N = H + hidden_size, // K = NH + 1.0f, // alpha + input_data + input_offset, // A + hidden_size, // lda = NH + weights_data + weights_offset, // B + 3 * hidden_size, // ldb = 3NH + 1.0f, // beta + qkv_dest + qkv_offset, // C + head_size, // ldc + nullptr // use single-thread + ); + } } }); } diff --git a/onnxruntime/contrib_ops/cpu/bert/attention.h b/onnxruntime/contrib_ops/cpu/bert/attention.h deleted file mode 100644 index bd82f07f81..0000000000 --- a/onnxruntime/contrib_ops/cpu/bert/attention.h +++ /dev/null @@ -1,21 +0,0 @@ -// Copyright (c) Microsoft Corporation. All rights reserved. -// Licensed under the MIT License. - -#pragma once - -#include "attention_cpu_base.h" -#include "core/common/common.h" -#include "core/framework/op_kernel.h" - -namespace onnxruntime { -namespace contrib { - -template -class Attention : public OpKernel, public AttentionCPUBase { - public: - explicit Attention(const OpKernelInfo& info); - Status Compute(OpKernelContext* context) const override; -}; - -} // namespace contrib -} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h index 403ec37a75..fc9c35059b 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h +++ b/onnxruntime/contrib_ops/cuda/quantization/attention_quantization.h @@ -6,7 +6,7 @@ #include "core/common/common.h" #include "core/framework/op_kernel.h" #include "core/providers/cuda/cuda_common.h" -#include "contrib_ops/cpu/bert/attention.h" +#include "contrib_ops/cpu/bert/attention_base.h" namespace onnxruntime { namespace contrib { diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 1fb137d135..697f510367 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -128,6 +128,23 @@ MlasGemm( MLAS_THREADPOOL* ThreadPool ); +void +MLASCALL +MlasGemm( + CBLAS_TRANSPOSE TransA, + size_t M, + size_t N, + size_t K, + float alpha, + const float* A, + size_t lda, + const void* PackedB, + float beta, + float* C, + size_t ldc, + MLAS_THREADPOOL* ThreadPool + ); + void MLASCALL MlasGemm( @@ -225,6 +242,24 @@ MlasGemm( // Buffer packing routines. // +size_t +MLASCALL +MlasGemmPackBSize( + size_t N, + size_t K + ); + +void +MLASCALL +MlasGemmPackB( + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const float* B, + size_t ldb, + void* PackedB + ); + size_t MLASCALL MlasGemmPackBSize( diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index c02ef6a335..93c00ad7d3 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -139,6 +139,8 @@ Abstract: #define MLAS_SGEMM_STRIDEN 128 #define MLAS_SGEMM_STRIDEK 128 +#define MLAS_SGEMM_PACKED_STRIDEN 128 +#define MLAS_SGEMM_PACKED_STRIDEK 256 #define MLAS_DGEMM_STRIDEN 64 #define MLAS_DGEMM_STRIDEK 128 diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index bcb4d77c23..48da1b6d0f 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -157,7 +157,7 @@ Return Value: for (size_t k = 0; k < K; k += CountK) { - CountK = (std::min)(K - k, Strides.K); + CountK = std::min(K - k, Strides.K); // // Step through each slice of matrix B along the N dimension. @@ -167,7 +167,7 @@ Return Value: for (size_t n = 0; n < N; n += CountN) { - CountN = (std::min)(N - n, Strides.N); + CountN = std::min(N - n, Strides.N); // // Copy a panel of matrix B to a local packed buffer. @@ -191,7 +191,7 @@ Return Value: for (size_t m = 0; m < M; m += CountM) { - CountM = (std::min)(M - m, Strides.M); + CountM = std::min(M - m, Strides.M); // // Copy a panel of matrix A to a local packed buffer. @@ -311,7 +311,7 @@ Return Value: for (size_t k = 0; k < K; k += CountK) { - CountK = (std::min)(K - k, Strides.K); + CountK = std::min(K - k, Strides.K); const size_t PackedCountK = (CountK + KernelType::PackedK - 1) / KernelType::PackedK; @@ -328,7 +328,7 @@ Return Value: for (size_t n = 0; n < N; n += CountN) { - CountN = (std::min)(N - n, Strides.N); + CountN = std::min(N - n, Strides.N); if (k == 0) { MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, PackedColumnSumBuffer + n, @@ -347,7 +347,7 @@ Return Value: for (size_t m = 0; m < M; m += CountM) { - CountM = (std::min)(M - m, Strides.M); + CountM = std::min(M - m, Strides.M); // // Copy a panel of matrix A to a local packed buffer. @@ -2126,7 +2126,7 @@ Return Value: for (size_t k = 0; k < K; k += CountK) { - CountK = (std::min)(K - k, StrideK); + CountK = std::min(K - k, StrideK); // // Step through each slice of matrix B along the N dimension. @@ -2141,7 +2141,7 @@ Return Value: constexpr size_t BatchedN = 128; MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[BatchedN], 64); - CountN = (std::min)(N - n, BatchedN); + CountN = std::min(N - n, BatchedN); if (GemmU8X8Operation == &MlasGemmU8X8PackedOperation) { MLAS_GEMM_U8S8_KERNEL_AVX2::CopyPackB(pb, B + n, ldb, CountN, CountK, ColumnSumBuffer, BIsSigned); diff --git a/onnxruntime/core/mlas/lib/sgemm.cpp b/onnxruntime/core/mlas/lib/sgemm.cpp index 79dc0810e9..7489d89095 100644 --- a/onnxruntime/core/mlas/lib/sgemm.cpp +++ b/onnxruntime/core/mlas/lib/sgemm.cpp @@ -32,21 +32,22 @@ Abstract: // struct MLAS_SGEMM_WORK_BLOCK { + int32_t ThreadCountM; + int32_t ThreadCountN; CBLAS_TRANSPOSE TransA; CBLAS_TRANSPOSE TransB; + size_t M; + size_t N; size_t K; + const float* A; size_t lda; + const float* B; size_t ldb; + const void* PackedB; + float* C; size_t ldc; float alpha; float beta; - struct SEGMENT { - size_t M; - size_t N; - const float* A; - const float* B; - float* C; - } Segments[MLAS_MAXIMUM_THREAD_COUNT]; }; void @@ -962,8 +963,8 @@ Return Value: // the A panel needs to be used for transposing. // - uint32_t StrideN = MLAS_SGEMM_STRIDEN; - uint32_t StrideK = MLAS_SGEMM_STRIDEK; + size_t StrideN = MLAS_SGEMM_STRIDEN; + size_t StrideK = MLAS_SGEMM_STRIDEK; if (N >= K) { @@ -985,15 +986,10 @@ Return Value: // size_t CountN; - size_t CountK; for (size_t n = 0; n < N; n += CountN) { - CountN = StrideN; - - if (CountN > (N - n)) { - CountN = N - n; - } + CountN = std::min(N - n, StrideN); // // Multiply the output matrix by beta as needed. @@ -1007,15 +1003,12 @@ Return Value: // Step through each slice of matrix B along the K dimension. // + size_t CountK; bool ZeroMode = (beta == 0.0f); for (size_t k = 0; k < K; k += CountK) { - CountK = StrideK; - - if (CountK > (K - k)) { - CountK = K - k; - } + CountK = std::min(K - k, StrideK); // // Copy or transpose a panel of matrix B to a local packed buffer. @@ -1048,11 +1041,7 @@ Return Value: // Transpose elements from matrix A into a local buffer. // - size_t RowsTransposed = RowsRemaining; - - if (RowsTransposed > MLAS_SGEMM_TRANSA_ROWS) { - RowsTransposed = MLAS_SGEMM_TRANSA_ROWS; - } + size_t RowsTransposed = std::min(RowsRemaining, size_t(MLAS_SGEMM_TRANSA_ROWS)); MlasSgemmTransposeA(PanelA, a, lda, RowsTransposed, CountK); @@ -1074,9 +1063,145 @@ Return Value: } void -MlasSgemmOperationThreaded( +MlasSgemmPackedOperation( + CBLAS_TRANSPOSE TransA, + size_t M, + size_t RangeStartN, + size_t RangeCountN, + size_t K, + float alpha, + const float* A, + size_t lda, + const void* PackedB, + size_t AlignedN, + float beta, + float* C, + size_t ldc + ) +/*++ + +Routine Description: + + This routine implements the single precision matrix/matrix multiply + operation (SGEMM). + +Arguments: + + TransA - Supplies the transpose operation for matrix A. + + M - Supplies the number of rows of matrix A and matrix C. + + RangeStartN - Supplies the starting column from packed matrix B. + + RangeCountN - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + alpha - Supplies the scalar alpha multiplier (see SGEMM definition). + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + PackedB - Supplies the address of packed matrix B. + + AlignedN - Supplies the total number of aligned columns for packed matrix B. + + ldb - Supplies the first dimension of matrix B. + + beta - Supplies the scalar beta multiplier (see SGEMM definition). + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + +Return Value: + + None. + +--*/ +{ + float PanelA[MLAS_SGEMM_TRANSA_ROWS * MLAS_SGEMM_PACKED_STRIDEK]; + + // + // Step through each slice of matrix B along the N dimension. + // + + size_t CountN; + + for (size_t n = 0; n < RangeCountN; n += CountN) { + + const size_t SliceStartN = RangeStartN + n; + + CountN = std::min(RangeCountN - n, size_t(MLAS_SGEMM_PACKED_STRIDEN)); + + // + // Multiply the output matrix by beta as needed. + // + + if (beta != 0.0f && beta != 1.0f) { + MlasSgemmMultiplyBeta(C + n, M, CountN, ldc, beta); + } + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + bool ZeroMode = (beta == 0.0f); + + for (size_t k = 0; k < K; k += CountK) { + + CountK = std::min(K - k, size_t(MLAS_SGEMM_PACKED_STRIDEK)); + + // + // Step through each slice of matrix A along the M dimension. + // + + const float* pb = (const float*)PackedB + AlignedN * k + CountK * SliceStartN; + float* c = C + n; + + if (TransA == CblasNoTrans) { + + MlasSgemmKernelLoop(A + k, pb, c, CountK, M, CountN, lda, ldc, alpha, ZeroMode); + + } else { + + const float* a = A + k * lda; + size_t RowsRemaining = M; + + do { + + // + // Transpose elements from matrix A into a local buffer. + // + + size_t RowsTransposed = std::min(RowsRemaining, size_t(MLAS_SGEMM_TRANSA_ROWS)); + + MlasSgemmTransposeA(PanelA, a, lda, RowsTransposed, CountK); + + RowsRemaining -= RowsTransposed; + a += RowsTransposed; + + // + // Step through the rows of the local buffer. + // + + c = MlasSgemmKernelLoop(PanelA, pb, c, CountK, RowsTransposed, CountN, CountK, ldc, alpha, ZeroMode); + + } while (RowsRemaining > 0); + } + + ZeroMode = false; + } + } +} + +void +MlasSgemmThreaded( void* Context, - int32_t Index + int32_t ThreadId ) /*++ @@ -1089,7 +1214,7 @@ Arguments: Context - Supplies the pointer to the context for the threaded operation. - Index - Supplies the current index of the threaded operation. + ThreadId - Supplies the current index of the threaded operation. Return Value: @@ -1097,89 +1222,111 @@ Return Value: --*/ { - MLAS_SGEMM_WORK_BLOCK* WorkBlock = (MLAS_SGEMM_WORK_BLOCK*)Context; + const auto* WorkBlock = (MLAS_SGEMM_WORK_BLOCK*)Context; - MLAS_SGEMM_WORK_BLOCK::SEGMENT* Segment = &WorkBlock->Segments[Index]; + const int32_t ThreadCountM = WorkBlock->ThreadCountM; + const int32_t ThreadCountN = WorkBlock->ThreadCountN; - MlasSgemmOperation(WorkBlock->TransA, WorkBlock->TransB, Segment->M, - Segment->N, WorkBlock->K, WorkBlock->alpha, Segment->A, WorkBlock->lda, - Segment->B, WorkBlock->ldb, WorkBlock->beta, Segment->C, - WorkBlock->ldc); + const int32_t ThreadIdM = ThreadId / ThreadCountN; + const int32_t ThreadIdN = ThreadId % ThreadCountN; + + // + // Partition the operation along the M dimension. + // + + size_t M = WorkBlock->M; + size_t RangeStartM; + size_t RangeCountM; + + MlasPartitionWork(ThreadIdM, ThreadCountM, M, &RangeStartM, &RangeCountM); + + // + // Partition the operation along the N dimension. + // + + size_t N = WorkBlock->N; + size_t RangeStartN; + size_t RangeCountN; + + const size_t BlockedN = (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) / + MLAS_SGEMM_STRIDEN_THREAD_ALIGN; + + MlasPartitionWork(ThreadIdN, ThreadCountN, BlockedN, &RangeStartN, + &RangeCountN); + + RangeStartN *= MLAS_SGEMM_STRIDEN_THREAD_ALIGN; + RangeCountN *= MLAS_SGEMM_STRIDEN_THREAD_ALIGN; + + RangeCountN = std::min(N - RangeStartN, RangeCountN); + + // + // Dispatch the partitioned operation. + // + + CBLAS_TRANSPOSE TransA = WorkBlock->TransA; + + const size_t lda = WorkBlock->lda; + const size_t ldc = WorkBlock->ldc; + + const float* A = WorkBlock->A + RangeStartM * ((TransA == CblasNoTrans) ? lda : 1); + float* C = WorkBlock->C + RangeStartM * ldc + RangeStartN; + + if (WorkBlock->B != nullptr) { + + CBLAS_TRANSPOSE TransB = WorkBlock->TransB; + + const size_t ldb = WorkBlock->ldb; + + const float* B = WorkBlock->B + RangeStartN * ((TransB == CblasNoTrans) ? 1 : ldb); + + MlasSgemmOperation(TransA, TransB, RangeCountM, RangeCountN, WorkBlock->K, + WorkBlock->alpha, A, lda, B, ldb, WorkBlock->beta, C, ldc); + + } else { + + MlasSgemmPackedOperation(TransA, RangeCountM, RangeStartN, RangeCountN, + WorkBlock->K, WorkBlock->alpha, A, lda, WorkBlock->PackedB, + BlockedN * MLAS_SGEMM_STRIDEN_THREAD_ALIGN, WorkBlock->beta, C, ldc); + } } -inline -bool -MlasSgemmTryMultithread( - CBLAS_TRANSPOSE TransA, - CBLAS_TRANSPOSE TransB, - size_t M, - size_t N, - size_t K, - float alpha, - const float* A, - size_t lda, - const float* B, - size_t ldb, - float beta, - float* C, - size_t ldc, +void +MlasSgemmSchedule( + MLAS_SGEMM_WORK_BLOCK* WorkBlock, MLAS_THREADPOOL* ThreadPool ) /*++ Routine Description: - This routine attempts to launch a single precision matrix/matrix multiply - operation (SGEMM) across multiple threads. + This routine schedules the single precision matrix/matrix multiply + operation (SGEMM) across one or more threads. Arguments: - TransA - Supplies the transpose operation for matrix A. - - TransB - Supplies the transpose operation for matrix B. - - M - Supplies the number of rows of matrix A and matrix C. - - N - Supplies the number of columns of matrix B and matrix C. - - K - Supplies the number of columns of matrix A and the number of rows of - matrix B. - - alpha - Supplies the scalar alpha multiplier (see SGEMM definition). - - A - Supplies the address of matrix A. - - lda - Supplies the first dimension of matrix A. - - B - Supplies the address of matrix B. - - ldb - Supplies the first dimension of matrix B. - - beta - Supplies the scalar beta multiplier (see SGEMM definition). - - C - Supplies the address of matrix C. - - ldc - Supplies the first dimension of matrix C. + WorkBlock - Supplies the structure containing the GEMM parameters. ThreadPool - Supplies the thread pool object to use, else nullptr if the base library threading support should be used. Return Value: - Returns true if the operation was completed across multiple threads, else - false if the operation should fall back to a single thread. + None. --*/ { - MLAS_SGEMM_WORK_BLOCK WorkBlock; - int32_t TargetThreadCount; + const size_t M = WorkBlock->M; + const size_t N = WorkBlock->N; + const size_t K = WorkBlock->K; // // Compute the number of target threads given the complexity of the SGEMM // operation. Small requests should run using the single threaded path. // - double Complexity = double(M) * double(N) * double(K); + const double Complexity = double(M) * double(N) * double(K); + + int32_t TargetThreadCount; if (Complexity < double(MLAS_SGEMM_THREAD_COMPLEXITY * MLAS_MAXIMUM_THREAD_COUNT)) { TargetThreadCount = int32_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1; @@ -1193,90 +1340,36 @@ Return Value: TargetThreadCount = MaximumThreadCount; } - if (TargetThreadCount == 1) { - return false; - } - - // - // Initialize the common fields of the work block. - // - - WorkBlock.TransA = TransA; - WorkBlock.TransB = TransB; - WorkBlock.K = K; - WorkBlock.lda = lda; - WorkBlock.ldb = ldb; - WorkBlock.ldc = ldc; - WorkBlock.alpha = alpha; - WorkBlock.beta = beta; - // // Segment the operation across multiple threads. // - - int32_t Index = 0; + // N.B. Currently, the operation is segmented as a 1D partition, which + // works okay for operations involving skinny matrices. + // if (N > M) { - size_t StrideN = N / TargetThreadCount; + const size_t BlockedN = (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) / + MLAS_SGEMM_STRIDEN_THREAD_ALIGN; - if ((StrideN * TargetThreadCount) != N) { - StrideN++; + if (size_t(TargetThreadCount) > BlockedN) { + TargetThreadCount = int32_t(BlockedN); } - StrideN = - (StrideN + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1); - - size_t pldb = (TransB == CblasNoTrans) ? 1 : ldb; - - for (size_t CountN, n = 0; n < N; n += CountN) { - - CountN = StrideN; - - if (CountN > (N - n)) { - CountN = N - n; - } - - WorkBlock.Segments[Index].M = M; - WorkBlock.Segments[Index].N = CountN; - WorkBlock.Segments[Index].A = A; - WorkBlock.Segments[Index].B = B + n * pldb; - WorkBlock.Segments[Index].C = C + n; - - Index++; - } + WorkBlock->ThreadCountM = 1; + WorkBlock->ThreadCountN = TargetThreadCount; } else { - size_t StrideM = M / TargetThreadCount; - - if ((StrideM * TargetThreadCount) != M) { - StrideM++; + if (size_t(TargetThreadCount) > M) { + TargetThreadCount = int32_t(M); } - size_t plda = (TransA == CblasNoTrans) ? lda : 1; - - for (size_t CountM, m = 0; m < M; m += CountM) { - - CountM = StrideM; - - if (CountM > (M - m)) { - CountM = M - m; - } - - WorkBlock.Segments[Index].M = CountM; - WorkBlock.Segments[Index].N = N; - WorkBlock.Segments[Index].A = A + m * plda; - WorkBlock.Segments[Index].B = B; - WorkBlock.Segments[Index].C = C + m * ldc; - - Index++; - } + WorkBlock->ThreadCountM = TargetThreadCount; + WorkBlock->ThreadCountN = 1; } - MlasExecuteThreaded(MlasSgemmOperationThreaded, &WorkBlock, Index, ThreadPool); - - return true; + MlasExecuteThreaded(MlasSgemmThreaded, WorkBlock, TargetThreadCount, ThreadPool); } void @@ -1342,12 +1435,216 @@ Return Value: --*/ { + MLAS_SGEMM_WORK_BLOCK WorkBlock; + // - // Try to run the operation across multiple threads or fall back to a - // single thread based on the GEMM parameters and system configuration. + // Capture the GEMM parameters to the work block. // - if (!MlasSgemmTryMultithread(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc, ThreadPool)) { - MlasSgemmOperation(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc); + memset(&WorkBlock, 0, sizeof(MLAS_SGEMM_WORK_BLOCK)); + + WorkBlock.TransA = TransA; + WorkBlock.TransB = TransB; + WorkBlock.M = M; + WorkBlock.N = N; + WorkBlock.K = K; + WorkBlock.A = A; + WorkBlock.lda = lda; + WorkBlock.B = B; + WorkBlock.ldb = ldb; + WorkBlock.C = C; + WorkBlock.ldc = ldc; + WorkBlock.alpha = alpha; + WorkBlock.beta = beta; + + // + // Schedule the operation across a set of worker threads. + // + + MlasSgemmSchedule(&WorkBlock, ThreadPool); +} + +void +MLASCALL +MlasGemm( + CBLAS_TRANSPOSE TransA, + size_t M, + size_t N, + size_t K, + float alpha, + const float* A, + size_t lda, + const void* PackedB, + float beta, + float* C, + size_t ldc, + MLAS_THREADPOOL* ThreadPool + ) +/*++ + +Routine Description: + + This routine implements the single precision matrix/matrix multiply + operation (SGEMM). + +Arguments: + + TransA - Supplies the transpose operation for matrix A. + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + alpha - Supplies the scalar alpha multiplier (see SGEMM definition). + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + PackedB - Supplies the address of packed matrix B. + + beta - Supplies the scalar beta multiplier (see SGEMM definition). + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + + ThreadPool - Supplies the thread pool object to use, else nullptr if the + base library threading support should be used. + +Return Value: + + None. + +--*/ +{ + MLAS_SGEMM_WORK_BLOCK WorkBlock; + + // + // Capture the GEMM parameters to the work block. + // + + memset(&WorkBlock, 0, sizeof(MLAS_SGEMM_WORK_BLOCK)); + + WorkBlock.TransA = TransA; + WorkBlock.M = M; + WorkBlock.N = N; + WorkBlock.K = K; + WorkBlock.A = A; + WorkBlock.lda = lda; + WorkBlock.PackedB = PackedB; + WorkBlock.C = C; + WorkBlock.ldc = ldc; + WorkBlock.alpha = alpha; + WorkBlock.beta = beta; + + // + // Schedule the operation across a set of worker threads. + // + + MlasSgemmSchedule(&WorkBlock, ThreadPool); +} + +size_t +MLASCALL +MlasGemmPackBSize( + size_t N, + size_t K + ) +/*++ + +Routine Description: + + This routine computes the length in bytes for the packed matrix B buffer. + +Arguments: + + N - Supplies the number of columns of matrix B. + + K - Supplies the number of rows of matrix B. + +Return Value: + + Returns the size in bytes for the packed matrix B buffer. + +--*/ +{ + // + // Compute the number of bytes required to hold the packed buffer. + // + + const size_t AlignedN = + (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1); + + const size_t BytesRequired = AlignedN * K * sizeof(float); + const size_t BufferAlignment = MlasGetPreferredBufferAlignment(); + const size_t AlignedBytesRequired = (BytesRequired + BufferAlignment - 1) & + ~(BufferAlignment - 1); + + return AlignedBytesRequired; +} + +void +MLASCALL +MlasGemmPackB( + CBLAS_TRANSPOSE TransB, + size_t N, + size_t K, + const float* B, + size_t ldb, + void* PackedB + ) +/*++ + +Routine Description: + + This routine packs the contents of matrix B to the destination buffer. The + destination buffer should be sized based on MlasGemmPackBSize(). For best + performance, the destination buffer should be aligned to the value returned + from MlasGetPreferredBufferAlignment(). + +Arguments: + + TransB - Supplies the transpose operation for matrix B. + + N - Supplies the number of columns of matrix B. + + K - Supplies the number of rows of matrix B. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + PackedB - Supplies the address of packed matrix B. + +Return Value: + + None. + +--*/ +{ + const size_t AlignedN = + (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1); + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + + for (size_t k = 0; k < K; k += CountK) { + + CountK = std::min(K - k, size_t(MLAS_SGEMM_PACKED_STRIDEK)); + + if (TransB == CblasNoTrans) { + MlasSgemmCopyPackB((float*)PackedB, B + k * ldb, ldb, N, CountK); + } else { + MlasSgemmTransposePackB((float*)PackedB, B + k, ldb, N, CountK); + } + + PackedB = (float*)PackedB + AlignedN * CountK; } } diff --git a/onnxruntime/core/providers/cpu/math/matmul.cc b/onnxruntime/core/providers/cpu/math/matmul.cc index 66762b5b35..728cf40cc0 100644 --- a/onnxruntime/core/providers/cpu/math/matmul.cc +++ b/onnxruntime/core/providers/cpu/math/matmul.cc @@ -5,6 +5,7 @@ #include "core/providers/cpu/math/matmul_helper.h" #include "core/util/math.h" #include "core/util/math_cpuonly.h" +#include "core/mlas/inc/mlas.h" namespace onnxruntime { @@ -16,6 +17,23 @@ class MatMul final : public OpKernel { Status Compute(OpKernelContext* context) const override; }; +#if !defined(USE_MKLML_FOR_BLAS) + +template <> +class MatMul final : public OpKernel { + public: + MatMul(const OpKernelInfo& info) : OpKernel(info) {} + + Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + Status Compute(OpKernelContext* context) const override; + + private: + TensorShape b_shape_; + BufferUniquePtr packed_b_; +}; + +#endif + ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL( MatMul, 1, 8, @@ -98,4 +116,86 @@ Status MatMul::Compute(OpKernelContext* ctx) const { return Status::OK(); } +#if !defined(USE_MKLML_FOR_BLAS) + +Status MatMul::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { + is_packed = false; + + // 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 size_t K = static_cast(b_shape_[0]); + const size_t N = 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(CblasNoTrans, N, K, tensor.Data(), N, packed_b_data); + is_packed = true; + } + return Status::OK(); +} + +Status MatMul::Compute(OpKernelContext* ctx) const { + concurrency::ThreadPool* thread_pool = ctx->GetOperatorThreadPool(); + + const Tensor* a = ctx->Input(0); + const Tensor* b = packed_b_ ? nullptr : ctx->Input(1); + + MatMulComputeHelper helper; + ORT_RETURN_IF_ERROR(helper.Compute(a->Shape(), b ? b->Shape() : b_shape_)); + Tensor* y = ctx->Output(0, helper.OutputShape()); + + // Bail out early if the output is going to be empty + if (y->Shape().Size() == 0) + return Status::OK(); + + const auto* a_data = a->Data(); + const auto* b_data = b ? b->Data() : nullptr; + auto* y_data = y->MutableData(); + + // TODO: replace it with GemmBatch for performance, it's OK for now as GemmBatch unrolls as well + size_t max_len = helper.OutputOffsets().size(); + for (size_t i = 0; i < max_len; i++) { + if (packed_b_) { + MlasGemm(CblasNoTrans, + static_cast(helper.M()), + static_cast(helper.N()), + static_cast(helper.K()), + 1.0f, + a_data + helper.LeftOffsets()[i], + static_cast(helper.K()), + packed_b_.get(), + 0.0f, + y_data + helper.OutputOffsets()[i], + static_cast(helper.N()), + thread_pool); + } else { + math::MatMul( + static_cast(helper.M()), + static_cast(helper.N()), + static_cast(helper.K()), + a_data + helper.LeftOffsets()[i], + b_data + helper.RightOffsets()[i], + y_data + helper.OutputOffsets()[i], + thread_pool); + } + } + + return Status::OK(); +} + +#endif + } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/attention_op_test.cc b/onnxruntime/test/contrib_ops/attention_op_test.cc index 299e55ade4..187145b1d4 100644 --- a/onnxruntime/test/contrib_ops/attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/attention_op_test.cc @@ -19,6 +19,7 @@ enum MaskIndexType { static void RunAttentionTest( const std::vector& input_data, // input: [batch_size, sequence_length, hidden_size] const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] + bool is_weights_constant, const std::vector& bias_data, // bias: [3 * hidden_size] const std::vector& mask_index_data, // mask_index: [batch_size] or [batch_size, past_sequence_length + sequence_length] or empty const std::vector& output_data, // output: [batch_size, sequence_length, hidden_size] @@ -35,7 +36,7 @@ static void RunAttentionTest( MaskIndexType mask_index_type = kMaskIndexEnd) { int min_cuda_architecture = use_float16 ? 530 : 0; - bool enable_cuda = HasCudaEnvironment(min_cuda_architecture); + bool enable_cuda = HasCudaEnvironment(min_cuda_architecture) && !is_weights_constant; bool enable_cpu = (nullptr != DefaultCpuExecutionProvider().get()) && !use_float16; int head_size = hidden_size / number_of_heads; if (enable_cpu || enable_cuda) { @@ -76,12 +77,12 @@ static void RunAttentionTest( if (use_float16) { tester.AddInput("input", input_dims, ToFloat16(input_data)); - tester.AddInput("weight", weights_dims, ToFloat16(weights_data)); + tester.AddInput("weight", weights_dims, ToFloat16(weights_data), is_weights_constant); tester.AddInput("bias", bias_dims, ToFloat16(bias_data)); tester.AddOutput("output", output_dims, ToFloat16(output_data)); } else { tester.AddInput("input", input_dims, input_data); - tester.AddInput("weight", weights_dims, weights_data); + tester.AddInput("weight", weights_dims, weights_data, is_weights_constant); tester.AddInput("bias", bias_dims, bias_data); tester.AddOutput("output", output_dims, output_data); } @@ -120,6 +121,33 @@ static void RunAttentionTest( } } +static void RunAttentionTest( + const std::vector& input_data, // input: [batch_size, sequence_length, hidden_size] + const std::vector& weights_data, // weights: [hidden_size, 3 * hidden_size] + const std::vector& bias_data, // bias: [3 * hidden_size] + const std::vector& mask_index_data, // mask_index: [batch_size] or [batch_size, past_sequence_length + sequence_length] or empty + const std::vector& output_data, // output: [batch_size, sequence_length, hidden_size] + int batch_size, + int sequence_length, + int hidden_size, + int number_of_heads, + bool use_float16 = false, + bool is_unidirectional = false, + bool use_past_state = false, + int past_sequence_length = 0, + const std::vector* past_data = nullptr, + const std::vector* present_data = nullptr, + MaskIndexType mask_index_type = kMaskIndexEnd) { + RunAttentionTest(input_data, weights_data, false, bias_data, mask_index_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, + use_float16, is_unidirectional, use_past_state, past_sequence_length, + past_data, present_data, mask_index_type); + RunAttentionTest(input_data, weights_data, true, bias_data, mask_index_data, output_data, + batch_size, sequence_length, hidden_size, number_of_heads, + use_float16, is_unidirectional, use_past_state, past_sequence_length, + past_data, present_data, mask_index_type); +} + TEST(AttentionTest, AttentionBatch1) { int batch_size = 1; int sequence_length = 2; diff --git a/onnxruntime/test/mlas/unittest.cpp b/onnxruntime/test/mlas/unittest.cpp index 3a4d8cc25d..f60fe83b00 100644 --- a/onnxruntime/test/mlas/unittest.cpp +++ b/onnxruntime/test/mlas/unittest.cpp @@ -67,7 +67,7 @@ public: ReleaseBuffer(); } - T* GetBuffer(size_t Elements) + T* GetBuffer(size_t Elements, bool ZeroFill = false) { // // Check if the internal buffer needs to be reallocated. @@ -125,20 +125,27 @@ public: T* GuardAddress = _GuardAddress; T* buffer = GuardAddress - Elements; - const int MinimumFillValue = -23; - const int MaximumFillValue = 23; + if (ZeroFill) { - int FillValue = MinimumFillValue; - T* FillAddress = buffer; + std::fill_n(buffer, Elements, T(0)); - while (FillAddress < GuardAddress) { + } else { - *FillAddress++ = (T)FillValue; + const int MinimumFillValue = -23; + const int MaximumFillValue = 23; - FillValue++; + int FillValue = MinimumFillValue; + T* FillAddress = buffer; - if (FillValue > MaximumFillValue) { - FillValue = MinimumFillValue; + while (FillAddress < GuardAddress) { + + *FillAddress++ = (T)FillValue; + + FillValue++; + + if (FillValue > MaximumFillValue) { + FillValue = MinimumFillValue; + } } } @@ -206,8 +213,67 @@ public: } }; +template +class MlasFgemmTestBase; + template -class MlasFgemmTest : public MlasTestBase +class MlasFgemmTestBase : public MlasTestBase +{ +public: + void + TestGemm( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t M, + size_t N, + size_t K, + float alpha, + const T* A, + size_t lda, + const T* B, + size_t ldb, + float beta, + T* C, + size_t ldc + ) + { + MlasGemm(TransA, TransB, M, N, K, T(alpha), A, lda, B, ldb, T(beta), C, ldc, threadpool); + } +}; + +template +class MlasFgemmTestBase : public MlasTestBase +{ +public: + void + TestGemm( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB, + size_t M, + size_t N, + size_t K, + float alpha, + const T* A, + size_t lda, + const T* B, + size_t ldb, + float beta, + T* C, + size_t ldc + ) + { + size_t PackedBSize = MlasGemmPackBSize(N, K); + void* PackedB = BufferBPacked.GetBuffer(PackedBSize, true); + MlasGemmPackB(TransB, N, K, B, ldb, PackedB); + MlasGemm(TransA, M, N, K, T(alpha), A, lda, PackedB, T(beta), C, ldc, threadpool); + } + +private: + MatrixGuardBuffer BufferBPacked; +}; + +template +class MlasFgemmTest : public MlasFgemmTestBase { private: void @@ -251,7 +317,7 @@ private: std::fill_n(C, M * N, -0.5f); std::fill_n(CReference, M * N, -0.5f); - MlasGemm(TransA, TransB, M, N, K, T(alpha), A, lda, B, ldb, T(beta), C, ldc, threadpool); + this->TestGemm(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, C, ldc); ReferenceGemm(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, CReference, ldc); for (size_t f = 0; f < M * N; f++) { @@ -392,6 +458,9 @@ public: for (size_t b = 256; b < 320; b += 32) { Test(b, b, b, 1.0f, 0.0f); } + + Test(128, 3072, 768, 1.0f, 0.0f); + Test(128, 768, 3072, 1.0f, 0.0f); } void @@ -2705,10 +2774,12 @@ RunThreadedTests( ) { printf("SGEMM tests.\n"); - onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); + printf("SGEMM packed tests.\n"); + onnxruntime::make_unique>()->ExecuteShort(); #ifdef MLAS_HAS_DGEMM printf("DGEMM tests.\n"); - onnxruntime::make_unique>()->ExecuteShort(); + onnxruntime::make_unique>()->ExecuteShort(); #endif #ifdef MLAS_HAS_QGEMM_U8X8 diff --git a/onnxruntime/test/providers/cpu/math/matmul_test.cc b/onnxruntime/test/providers/cpu/math/matmul_test.cc index 0f813a63aa..045d703b78 100644 --- a/onnxruntime/test/providers/cpu/math/matmul_test.cc +++ b/onnxruntime/test/providers/cpu/math/matmul_test.cc @@ -94,7 +94,7 @@ std::vector> GenerateTestCases() { } template -void RunMatMulTest(int32_t opset_version = 7) { +void RunMatMulTest(int32_t opset_version, bool is_b_constant = false) { std::vector common_input_vals{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}; for (auto t : GenerateTestCases()) { OpTester test("MatMul", opset_version); @@ -105,17 +105,24 @@ void RunMatMulTest(int32_t opset_version = 7) { int64_t size1 = TensorShape::ReinterpretBaseType(t.input1_dims).SizeHelper(0, t.input1_dims.size()); std::vector input1_vals(common_input_vals.cbegin(), common_input_vals.cbegin() + size1); - test.AddInput("B", t.input1_dims, input1_vals); + test.AddInput("B", t.input1_dims, input1_vals, is_b_constant); test.AddOutput("Y", t.expected_dims, t.expected_vals); // OpenVINO EP: Disabled temporarily matmul broadcasting not fully supported - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // Disable TensorRT because of unsupported data type + // Disable TensorRT because of unsupported data type + std::unordered_set excluded_providers{kTensorrtExecutionProvider, kOpenVINOExecutionProvider}; + if (is_b_constant) { + // NNAPI: currently fails for the "test 2D empty input" case + excluded_providers.insert(kNnapiExecutionProvider); + } + test.Run(OpTester::ExpectResult::kExpectSuccess, "", excluded_providers); } } TEST(MathOpTest, MatMulFloatType) { - RunMatMulTest(7); + RunMatMulTest(7, false); + RunMatMulTest(7, true); } TEST(MathOpTest, MatMulDoubleType) {