[Performance] Upstream MLAS backend optimization for better thread partitioning in multi-group or large batch convolutions

[zoeczy]

Describe the issue

We have implemented a performance optimization in the MLAS backend of ONNX Runtime that improves CPU utilization in convolution workloads with multiple groups or large batch sizes. As a result, in multi-group or large-batch convolution scenarios, we observe near-linear performance scaling with increasing intra_op_num_threads.

We would like to ask:

1. Whether this kind of optimization aligns with the upstream design goals for MLAS?
2. If so, would the ONNX Runtime team be open to reviewing a PR that introduces this optimization?
3. Are there any guidelines or preferred patterns for contributing such low-level performance improvements to MLAS?

To reproduce

This optimization targets CPU inference scenarios with:

• Convolution models using group > 1, or
• Convolutions with large batch size (e.g., batch ≥ 32)

To observe the issue without optimization:

1. Use ort.SessionOptions().intra_op_num_threads = N (e.g., 16 or 32) to explicitly control the number of intra-op threads.
2. Benchmark the execution time and thread utilization of convolution operators with multiple groups or large batch sizes
3. Measure the overall inference latency of such models and compare core utilization across threads.

Urgency

No response

Platform

Linux

OS Version

Ubuntu 20.04 (x86_64)

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

master

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Other / Unknown

Execution Provider Library Version

MLAS(cpu default execution provider)

Model File

No response

Is this a quantized model?

No

[snnn]

I wonder if you have tried mlperf inference benchmark with resnet models.

[snnn]

These models are CNN models, which heavily use conv.
https://github.com/mlcommons/inference/tree/master/vision/classification_and_detection

And, how can I reproduce the results? For example, I have some Intel/AMD desktop CPUs, ARM SBC devices, cloud VMs. Can I see perf gain there?

[zoeczy]

Thanks for the questions! Here's a detailed response:

Q1: Have you tried MLPerf inference benchmark with ResNet models?
Our current optimization targets Conv1D workloads, which are commonly used in models like TTS . ResNet models in MLPerf use MlasNchwcConv from the mlas backend. We've specifically focused on optimizing MlasConv (used by Conv1D), so this improvement doesn't directly apply to ResNet-style convolutions.

Q2: How can I reproduce the results?
We tested our optimization on the FastSpeech model using onnxruntime on an Intel Xeon W3-2435 machine. Use --intra_op_num_threads=4, batch size = 1 setting, inference time improved from 0.172s to 0.143s. The performance gain becomes more significant as batch size or number of groups increases, which is where our thread partitioning strategy becomes more effective.

Our optimization focuses on improving the thread partitioning strategy in the MlasConv backend, particularly for cases where batch size > 1 or group > 1. This change leads to better parallelism and performance in such scenarios. I’d like to contribute my patch — please let me know if there’s anything I should do beforehand.

These models are CNN models, which heavily use conv. https://github.com/mlcommons/inference/tree/master/vision/classification_and_detection

And, how can I reproduce the results? For example, I have some Intel/AMD desktop CPUs, ARM SBC devices, cloud VMs. Can I see perf gain there?

[skottmckay]

Feel free to create a PR. Make sure there are unit tests and the code is well commented.

Have you tested on a wider range of CPUs? The biggest concern would be that an improvement on one type of CPU would be a regression on others, and a Xeon is typically used in server scenarios. i.e. we wouldn't want to negatively affect consumer Intel CPUs or AMD CPUs