diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 4afb804b41..b0de6f7e2e 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -789,7 +789,7 @@ if(onnxruntime_BUILD_BENCHMARKS) file(GLOB_RECURSE MLAS_BENCH_SOURCE_FILES "${MLAS_BENCH_DIR}/*.cpp" "${MLAS_BENCH_DIR}/*.h") onnxruntime_add_executable(onnxruntime_mlas_benchmark ${MLAS_BENCH_SOURCE_FILES}) target_include_directories(onnxruntime_mlas_benchmark PRIVATE ${ONNXRUNTIME_ROOT}/core/mlas/inc) - target_link_libraries(onnxruntime_mlas_benchmark PRIVATE benchmark::benchmark onnxruntime_mlas onnxruntime_common) + target_link_libraries(onnxruntime_mlas_benchmark PRIVATE benchmark::benchmark onnxruntime_mlas onnxruntime_common onnxruntime_framework onnxruntime_util) if(NOT WIN32) target_link_libraries(onnxruntime_mlas_benchmark PRIVATE nsync_cpp) endif() diff --git a/onnxruntime/test/mlas/bench/bench_qgemm.cpp b/onnxruntime/test/mlas/bench/bench_qgemm.cpp new file mode 100644 index 0000000000..cb2c6b94c6 --- /dev/null +++ b/onnxruntime/test/mlas/bench/bench_qgemm.cpp @@ -0,0 +1,111 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "mlas.h" +#include "bench_util.h" +#include "core/util/thread_utils.h" + +#include +#include +#include +#include + +static const std::vector qgemm_arg_names = {"M", "N", "K", "Batch", "Threads"}; + +void QGEMM(benchmark::State& state, bool pack_b) { + const bool b_is_signed = true; + const uint8_t a_zero_point = 29; + const uint8_t b_zero_point = 179; + + const int64_t M = state.range(0); + const int64_t N = state.range(1); + const int64_t K = state.range(2); + + const int64_t batch = state.range(3); + const int64_t threads = state.range(4); + + if (M <= 0) throw std::invalid_argument("M must greater than 0!"); + if (N <= 0) throw std::invalid_argument("N must greater than 0!"); + if (K <= 0) throw std::invalid_argument("K must greater than 0!"); + if (batch <= 0) throw std::invalid_argument("Batch must greater than 0!"); + if (threads <= 0) throw std::invalid_argument("Threads must greater than 0!"); + + OrtThreadPoolParams tpo; + tpo.thread_pool_size = int(threads); + tpo.auto_set_affinity = true; + std::unique_ptr tp( + onnxruntime::concurrency::CreateThreadPool(&onnxruntime::Env::Default(), + tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP)); + + auto A_holder = RandomVectorUniform(static_cast(M * K * batch), uint8_t(-100), uint8_t(100)); + auto B_holder = RandomVectorUniform(static_cast(N * K * batch), uint8_t(-110), uint8_t(110)); + std::vector C_holder(static_cast(M * N * batch)); + std::vector pack_b_holder; + + MLAS_GEMM_U8X8_PARAMETERS gemm_params; + + gemm_params.M = static_cast(M); + gemm_params.N = static_cast(N); + gemm_params.K = static_cast(K); + gemm_params.lda = gemm_params.K; + gemm_params.ZeroPointA = a_zero_point; + gemm_params.ZeroPointB = &b_zero_point; + gemm_params.BIsSigned = b_is_signed; + gemm_params.C = C_holder.data(); + gemm_params.ldc = gemm_params.N; + + size_t packed_b_size = 0; + if (pack_b) { + packed_b_size = MlasGemmPackBSize(N, K, b_is_signed); + pack_b_holder.resize(packed_b_size * batch); + for (int i = 0; i < batch; i++) { + MlasGemmPackB(N, K, B_holder.data() + N * K * i, N, b_is_signed, (void*)(pack_b_holder.data() + packed_b_size * i)); + } + gemm_params.BIsPacked = true; + } else { + gemm_params.ldb = gemm_params.N; + } + + for (auto _ : state) { + for (int i = 0; i < batch; i++) { + gemm_params.A = A_holder.data() + M * K * i; + gemm_params.C = C_holder.data() + M * N * i; + if (pack_b) { + gemm_params.B = (void*)(pack_b_holder.data() + packed_b_size * i); + } else { + gemm_params.B = B_holder.data() + N * K * i; + } + MlasGemm(&gemm_params, tp.get()); + } + } +} + +static void QGemmSize(benchmark::internal::Benchmark* b) { + b->ArgNames(qgemm_arg_names); + // Args for "M", "N", "K", "Batch", + std::vector> mnk_batch = { + {512, 64, 512, 12}, + {512, 512, 64, 12}, + {512, 768, 768, 1}, + {512, 768, 3072, 1}, + {128, 768, 2304, 1}, + {128, 768, 2304, 6}, + {128, 1024, 4096, 1}, + {128, 1024, 4096, 6}, + {128, 2048, 8192, 1}, + {128, 4096, 16384, 1} + }; + + // Args for Threads + for (int64_t threads : {2, 4, 8, 16}) { + for (auto& shape : mnk_batch) { + std::vector copy(shape); + copy.push_back(threads); + b->Args(copy); + } + } +} + + +BENCHMARK_CAPTURE(QGEMM, PackB, true)->Apply(QGemmSize)->UseRealTime(); +BENCHMARK_CAPTURE(QGEMM, NoPackB, false)->Apply(QGemmSize)->UseRealTime(); diff --git a/onnxruntime/test/mlas/bench/bench_sgemm.cpp b/onnxruntime/test/mlas/bench/bench_sgemm.cpp index 8743e04746..4336311a90 100644 --- a/onnxruntime/test/mlas/bench/bench_sgemm.cpp +++ b/onnxruntime/test/mlas/bench/bench_sgemm.cpp @@ -18,8 +18,8 @@ void SGEMM(benchmark::State& state, bool pack_b, bool trans_a, bool trans_b, flo if (N <= 0) throw std::invalid_argument("N must greater than 0!"); if (K <= 0) throw std::invalid_argument("K must greater than 0!"); - auto A = RandomVectorUniform(static_cast(M * K), -1.0, 1.0); - auto B = RandomVectorUniform(static_cast(N * K), -1.0, 1.0); + auto A = RandomVectorUniform(static_cast(M * K), -1.0f, 1.0f); + auto B = RandomVectorUniform(static_cast(N * K), -1.0f, 1.0f); std::vector C(static_cast(M * N)); if (pack_b) { diff --git a/onnxruntime/test/mlas/bench/bench_util.cpp b/onnxruntime/test/mlas/bench/bench_util.cpp index 875088eb31..925f9394fd 100644 --- a/onnxruntime/test/mlas/bench/bench_util.cpp +++ b/onnxruntime/test/mlas/bench/bench_util.cpp @@ -3,7 +3,6 @@ #include "bench_util.h" #include -#include #include std::vector BenchArgsVector(benchmark::State& state, size_t& start, size_t count) { @@ -16,18 +15,6 @@ std::vector BenchArgsVector(benchmark::State& state, size_t& start, siz return shape; } -std::vector RandomVectorUniform(size_t N, float min_value, float max_value) { - if (min_value >= max_value) { - return std::vector(N, min_value); - } - std::default_random_engine generator(static_cast(N)); - std::uniform_real_distribution distribution(min_value, max_value); - std::vector r(N); - for (size_t i = 0; i < N; i++) { - r[i] = distribution(generator); - } - return r; -} std::vector RandomVectorUniform(std::vector shape, float min_value, float max_value) { int64_t sz = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies()); diff --git a/onnxruntime/test/mlas/bench/bench_util.h b/onnxruntime/test/mlas/bench/bench_util.h index 0e0c1ee5c9..33b0b72a3a 100644 --- a/onnxruntime/test/mlas/bench/bench_util.h +++ b/onnxruntime/test/mlas/bench/bench_util.h @@ -5,10 +5,29 @@ #include +#include + + void ArgsProduct(benchmark::internal::Benchmark* bench, const std::vector>& arglists); -std::vector RandomVectorUniform(size_t N, float min_value, float max_value); +template +std::vector RandomVectorUniform( + size_t N, + ElementType min_value = std::numeric_limits::lowest(), + ElementType max_value = std::numeric_limits::max()) { + if (min_value >= max_value) { + return std::vector(N, min_value); + } + std::default_random_engine generator(static_cast(N)); + std::uniform_real_distribution distribution(static_cast(min_value), static_cast(max_value)); + + std::vector r(N); + for (size_t i = 0; i < N; i++) { + r[i] = static_cast(distribution(generator)); + } + return r; +} std::vector RandomVectorUniform(std::vector shape, float min_value, float max_value);