Add QGEMM benchmark (#7268)

* Add QGEMM benchmark

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()

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 <stdexcept>
+#include <memory>
+#include <numeric>
+#include <algorithm>
+
+static const std::vector<std::string> 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<onnxruntime::concurrency::ThreadPool> tp(
+      onnxruntime::concurrency::CreateThreadPool(&onnxruntime::Env::Default(),
+      tpo, onnxruntime::concurrency::ThreadPoolType::INTRA_OP));
+
+  auto A_holder = RandomVectorUniform<uint8_t>(static_cast<size_t>(M * K * batch), uint8_t(-100), uint8_t(100));
+  auto B_holder = RandomVectorUniform<uint8_t>(static_cast<size_t>(N * K * batch), uint8_t(-110), uint8_t(110));
+  std::vector<int32_t> C_holder(static_cast<size_t>(M * N * batch));
+  std::vector<uint8_t> pack_b_holder;
+
+  MLAS_GEMM_U8X8_PARAMETERS gemm_params;
+
+  gemm_params.M = static_cast<size_t>(M);
+  gemm_params.N = static_cast<size_t>(N);
+  gemm_params.K = static_cast<size_t>(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<std::vector<int64_t>> 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<int64_t> 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();

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<size_t>(M * K), -1.0, 1.0);
-  auto B = RandomVectorUniform(static_cast<size_t>(N * K), -1.0, 1.0);
+  auto A = RandomVectorUniform(static_cast<size_t>(M * K), -1.0f, 1.0f);
+  auto B = RandomVectorUniform(static_cast<size_t>(N * K), -1.0f, 1.0f);
   std::vector<float> C(static_cast<size_t>(M * N));
 
   if (pack_b) {

onnxruntime/test/mlas/bench/bench_util.cpp

@@ -3,7 +3,6 @@
 
 #include "bench_util.h"
 #include <numeric>
-#include <random>
 #include <stdexcept>
 
 std::vector<int64_t> BenchArgsVector(benchmark::State& state, size_t& start, size_t count) {
@@ -16,18 +15,6 @@ std::vector<int64_t> BenchArgsVector(benchmark::State& state, size_t& start, siz
   return shape;
 }
 
-std::vector<float> RandomVectorUniform(size_t N, float min_value, float max_value) {
-  if (min_value >= max_value) {
-    return std::vector<float>(N, min_value);
-  }
-  std::default_random_engine generator(static_cast<unsigned>(N));
-  std::uniform_real_distribution<float> distribution(min_value, max_value);
-  std::vector<float> r(N);
-  for (size_t i = 0; i < N; i++) {
-    r[i] = distribution(generator);
-  }
-  return r;
-}
 
 std::vector<float> RandomVectorUniform(std::vector<int64_t> shape, float min_value, float max_value) {
   int64_t sz = std::accumulate(shape.begin(), shape.end(), 1LL, std::multiplies<int64_t>());

onnxruntime/test/mlas/bench/bench_util.h

@@ -5,10 +5,29 @@
 
 #include <benchmark/benchmark.h>
 
+#include <random>
+
+
 void ArgsProduct(benchmark::internal::Benchmark* bench,
                  const std::vector<std::vector<int64_t>>& arglists);
 
-std::vector<float> RandomVectorUniform(size_t N, float min_value, float max_value);
+template<typename ElementType>
+std::vector<ElementType> RandomVectorUniform(
+    size_t N,
+    ElementType min_value = std::numeric_limits<ElementType>::lowest(),
+    ElementType max_value = std::numeric_limits<ElementType>::max()) {
+  if (min_value >= max_value) {
+    return std::vector<ElementType>(N, min_value);
+  }
+  std::default_random_engine generator(static_cast<unsigned>(N));
+  std::uniform_real_distribution<double> distribution(static_cast<double>(min_value), static_cast<double>(max_value));
+  
+  std::vector<ElementType> r(N);
+  for (size_t i = 0; i < N; i++) {
+    r[i] = static_cast<ElementType>(distribution(generator));
+  }
+  return r;
+}
 
 std::vector<float> RandomVectorUniform(std::vector<int64_t> shape, float min_value, float max_value);