Repository: microsoft/onnxruntime
File: onnxruntime/test/mlas/bench/bench_sconv.cpp (C++)

Code before:

  // warm up first round.
  MlasConv(&Parameters,
           X.data(),
           F.data(),
           nullptr,
           working_buffer.data(),
           Y.data(),
           nullptr);

  for (auto _ : state) {
    MlasConv(&Parameters,
             X.data(),
             F.data(),
             nullptr,
             working_buffer.data(),
             Y.data(),
             nullptr);
  }
}

MLAS_THREADPOOL* GetMlasThreadPool(void) {
  static auto threadpool = std::make_unique<onnxruntime::concurrency::ThreadPool>(
      &onnxruntime::Env::Default(), onnxruntime::ThreadOptions(), nullptr, 4, true);
  return threadpool.get();
}

void SCONV_NCHW_THREADED(benchmark::State& state, const char* /*dummy*/) {
  MLAS_THREADPOOL* tp = GetMlasThreadPool();

  const int64_t rank = state.range(0);                       // Rank
  const int64_t batch_size = state.range(1);                 // N
  const int64_t groups = state.range(2);                     // G
  const int64_t input_channels_per_group = state.range(3);   // Cpg
  const int64_t output_channels_per_group = state.range(4);  // Fpg

  if (rank <= 0) throw std::invalid_argument("Kernel rank must greater than 0!");
  if (batch_size <= 0) throw std::invalid_argument("Batch size must greater than 0!");
  if (groups <= 0) throw std::invalid_argument("Group count must greater than 0!");
  if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must greater than 0!");
  if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must greater than 0!");

  size_t arg_position = 5;
  const auto input_shape = BenchArgsVector(state, arg_position, rank);
  const auto kernel_shape = BenchArgsVector(state, arg_position, rank);
  const auto paddings = BenchArgsVector(state, arg_position, rank * 2);
  const auto strides = BenchArgsVector(state, arg_position, rank);
  const auto dilations = BenchArgsVector(state, arg_position, rank);

  // do not check the size of each vector as they are forced from args.
  if (std::any_of(input_shape.begin(), input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) {

Review comment (refactor):
`GetMlasThreadPool(void)` is already defined here:

https://github.com/microsoft/onnxruntime/blob/cd4ac4943465fbbadbbc622dcfd11f453a27a2c1/onnxruntime/test/mlas/unittest/test_main.cpp#L10-L24

instead of providing another definition, can we maybe refactor the existing one to allow customization?

e.g., maybe have a helper `std::unique_ptr<MLAS_THREADPOOL> CreateMlasThreadPool(int degree_of_parallelism)` and update `GetMlasThreadPool()` to call that and also call that here directly.

Code after:

  // warm up first round.
  MlasConv(&Parameters,
           X.data(),
           F.data(),
           nullptr,
           working_buffer.data(),
           Y.data(),
           nullptr);

  for (auto _ : state) {
    MlasConv(&Parameters,
             X.data(),
             F.data(),
             nullptr,
             working_buffer.data(),
             Y.data(),
             nullptr);
  }
}

static MLAS_THREADPOOL* GetMlasThreadPoolForConvBenchmark(void) {
  static auto threadpool = std::make_unique<onnxruntime::concurrency::ThreadPool>(
      &onnxruntime::Env::Default(), onnxruntime::ThreadOptions(), nullptr, 4, true);
  return threadpool.get();
}

void SCONV_NCHW_THREADED(benchmark::State& state, const char* /*dummy*/) {
  MLAS_THREADPOOL* tp = GetMlasThreadPoolForConvBenchmark();

  const int64_t rank = state.range(0);                       // Rank
  const int64_t batch_size = state.range(1);                 // N
  const int64_t groups = state.range(2);                     // G
  const int64_t input_channels_per_group = state.range(3);   // Cpg
  const int64_t output_channels_per_group = state.range(4);  // Fpg

  if (rank <= 0) throw std::invalid_argument("Kernel rank must greater than 0!");
  if (batch_size <= 0) throw std::invalid_argument("Batch size must greater than 0!");
  if (groups <= 0) throw std::invalid_argument("Group count must greater than 0!");
  if (input_channels_per_group <= 0) throw std::invalid_argument("input_channels_per_group must greater than 0!");
  if (output_channels_per_group <= 0) throw std::invalid_argument("output_channels_per_group must greater than 0!");

  size_t arg_position = 5;
  const auto input_shape = BenchArgsVector(state, arg_position, rank);
  const auto kernel_shape = BenchArgsVector(state, arg_position, rank);
  const auto paddings = BenchArgsVector(state, arg_position, rank * 2);
  const auto strides = BenchArgsVector(state, arg_position, rank);
  const auto dilations = BenchArgsVector(state, arg_position, rank);

  // do not check the size of each vector as they are forced from args.
  if (std::any_of(input_shape.begin(), input_shape.end(), [](const int64_t& dim) { return dim <= 0; })) {