Skip to content

mlas/arm64: add NEON conv asm kernels and tune NCHWC kernel selection#27099

Merged
hariharans29 merged 11 commits intomicrosoft:mainfrom
milpuz01:aarch64_convolutions
Feb 13, 2026
Merged

mlas/arm64: add NEON conv asm kernels and tune NCHWC kernel selection#27099
hariharans29 merged 11 commits intomicrosoft:mainfrom
milpuz01:aarch64_convolutions

Conversation

@milpuz01
Copy link
Contributor

Overview

This PR adds ARM64 NEON assembly micro‑kernels for NCHW, depthwise, and pointwise convolution, wires them into the MLAS build, and adds shape‑based selection heuristics for NCHWC depthwise/pointwise to favor the asm kernels in safe cases (stride‑1 pointwise; wider depthwise outputs). The BF16 path is unchanged.

Key changes

  • cmake/onnxruntime_mlas.cmake
    • Add new AArch64 assembly sources for NCHW, depthwise, and pointwise conv to the MLAS build.
  • onnxruntime/core/mlas/lib/aarch64/SconvKernelNeon.S
    • New vectorised NCHW convolution micro‑kernel.
  • onnxruntime/core/mlas/lib/aarch64/SconvDepthwiseKernelNeon.S
    • New vectorised depthwise micro‑kernel (fast path for in‑bounds loads, slow path for padding).
  • onnxruntime/core/mlas/lib/aarch64/SconvPointwiseKernelNeon.S
    • New vectorised pointwise micro‑kernel (multi‑output reuse).
  • onnxruntime/core/mlas/lib/mlasi.h, onnxruntime/core/mlas/lib/platform.cpp
    • Declare/register new asm kernels and prefer them on ARM64.
  • onnxruntime/core/mlas/lib/snchwc.cpp
    • Heuristics: use pointwise asm when StrideHeight == 1 && StrideWidth == 1 and OutputThisIteration >= 4; use depthwise asm when OutputWidth >= 4.
  • onnxruntime/core/mlas/lib/sbconv_kernel_neon.cpp
    • Include fix for the conv kernel flags header.

Performance

Numbers below are expressed as multipliers vs the non‑NCHWC baseline (same model and perf_test settings):

Baseline (no --enable_arm_neon_nchwc)

  • 8 cores: 1.00×
  • 16 cores: 1.00×

With --enable_arm_neon_nchwc (no asm additions/heuristics)

  • 8 cores: 1.18×
  • 16 cores: 1.24×

With this PR (asm kernels + heuristics)

  • 8 cores: 1.77×
  • 16 cores: 2.54×

Testing

  • ./build.sh --config Release --build_shared_lib --parallel --compile_no_warning_as_error --skip_submodule_sync --skip_tests --enable_pybind --build_wheel --enable_arm_neon_nchwc
  • OMP_NUM_THREADS=8 ./build/Linux/Release/onnxruntime_perf_test -I -m times -r 1000 --x 8 ~/mobilenetv2-7.onnx

@aviralagrawal
Copy link

Interesting contribution - thank you!

A few questions -

  1. Pointwise convolution is currently implemented via direct GEMM - which I assume is optimized. How does this kernel beat the performance of GEMM?
  2. Can you share the link to the mobilenet model that you used for performance benchmarking?
  3. How does it perform on single threaded experiments? Afaik, the original nchwc kernels in NEON kernels for NCHWc Convolution and Pooling #25580 suffered in the single threaded setting but outperformed the default in thread counts>8.

@milpuz01
Copy link
Contributor Author

Hi @aviralagrawal, thank you vey much for your prompt feedback.

  1. Pointwise convolution is currently implemented via direct GEMM - which I assume is optimized. How does this kernel beat the performance of GEMM?

Compared to direct GEMM implementation of pointwise convolution asm kernel computes 1x1 conv directly:

  • it explicitly tiles 4 outputs: computes up to 4 output positions in parallel and reuses filter loads across those outputs so with single load we are able to accumulate 4 outputs while direct GEMM doesn't tile multiple outputs together
  • fuses accumulate/bias/ReLU into store path instead of separate passes with direct GEMM
  • unrolls the block size explicitly with 16 invocations to keep accumulators in registers and minimise loop overheads thus reducing dispatch/param overhead and output read-modify-write passes compared to direct GEMM

As usual there are trade-offs so direct GEMM would be faster when output count is small because then asm kernel drops to single-output path which has less ILP and won't be able to reuse filter loads, non-unit stride and non-contigious output regions hence why we have heuristics checking for stride width and height and very large K/M when GEMM blocking can make better use of caches then a fixed 4-output tile.

This is best illustrated if we extract pointwise convolutions from mobilnet that we ran and we can see that on average asm implementation is 1.07x faster, and significant speed ups are when number of channels is high and K/M are small (in the image those are H and W dimensions). In convolution heavy networks the convolutions that are dominant are ones with large number of channels and low height and width so we see visible performance improvements as optimisations from this PR are weighted in that direction.

image
  1. Can you share the link to the mobilenet model that you used for performance benchmarking?

For benchmarking we used the model from: https://github.com/onnx/models/blob/main/validated/vision/classification/mobilenet/model/mobilenetv2-7.onnx

  1. How does it perform on single threaded experiments? Afaik, the original nchwc kernels in NEON kernels for NCHWc Convolution and Pooling #25580 suffered in the single threaded setting but outperformed the default in thread counts>8.

Running OMP_NUM_THREADS=1 ./build/Linux/Release/onnxruntime_perf_test -I -m times -r 1000 --x 1 ~/mobilenetv2-7.onnx on Graviton 4 with binary that was built with --enable_arm_neon_nchwc it is slower by 0.89x then building without that flag, while with this PR it is actually 1.25x faster than the baseline.

@Rohanjames1997
Copy link
Contributor

Thanks @milpuz01 for the detailed description & comment!

A couple questions from my side:

  1. Is there a reason why ConvNchwcFloatKernel was not optimized? Afaik, It is not very different from ConvNchwFloatKernel. The x86 asm implementations of these two kernels differ very slightly too. It is a much heavier kernel than Pointwise and Depthwise, and many larger Conv models stress this kernel. An example for this type of model is in this comment: NEON kernels for NCHWc Convolution and Pooling #25580 (comment).

  2. Can we switch the default path of Fp32 Conv on Arm64 to use these new kernels? (effectively voiding --enable_arm_neon_nchwc like it was before) Asking because this PR improves upon the single-threaded performance as well. I'd love to hear your thoughts, but also would be wise to hear from @hariharans29 before implementing.

@milpuz01
Copy link
Contributor Author

Hi @Rohanjames1997, thank you very much for your comments.

  1. Is there a reason why ConvNchwcFloatKernel was not optimized?

No, particular reason. Mostly because the focus for this PR was on MobileNet model and lack of bandwidth. Thank you for sharing the model where ConvNchwcFloatKernel is invoked. We will take a look at optimising it to, but I would suggest that we add optimisation as a follow up PR so that we do not overload this PR with too many changes to review.

  1. Can we switch the default path of Fp32 Conv on Arm64 to use these new kernels? (effectively voiding --enable_arm_neon_nchwc like it was before) Asking because this PR improves upon the single-threaded performance as well. I'd love to hear your thoughts, but also would be wise to hear from @hariharans29 before implementing.

Yes, I think that is great idea and would be interesting to hear from @hariharans29 too what other testing we should make to try to make these kernels default. As you can see above this change is not going to accelerate all possible pointwise convolutions for example but on average it will show the improvements so if we could agree on a set of performance targets we can use that to drive the decision.

Also thank you for your code review I will address them in a separate commit.

@hariharans29
Copy link
Member

hariharans29 commented Jan 23, 2026

Hi @Rohanjames1997, thank you very much for your comments.

  1. Is there a reason why ConvNchwcFloatKernel was not optimized?

No, particular reason. Mostly because the focus for this PR was on MobileNet model and lack of bandwidth. Thank you for sharing the model where ConvNchwcFloatKernel is invoked. We will take a look at optimising it to, but I would suggest that we add optimisation as a follow up PR so that we do not overload this PR with too many changes to review.

  1. Can we switch the default path of Fp32 Conv on Arm64 to use these new kernels? (effectively voiding --enable_arm_neon_nchwc like it was before) Asking because this PR improves upon the single-threaded performance as well. I'd love to hear your thoughts, but also would be wise to hear from @hariharans29 before implementing.

Yes, I think that is great idea and would be interesting to hear from @hariharans29 too what other testing we should make to try to make these kernels default. As you can see above this change is not going to accelerate all possible pointwise convolutions for example but on average it will show the improvements so if we could agree on a set of performance targets we can use that to drive the decision.

Also thank you for your code review I will address them in a separate commit.

Unfortunately, I don't have a comprehensive list of performance targets to be met to make the feature default. Since, the performance testing may not include all possible Conv shapes, I would like to err on the side of caution and atleast provide one release timeline heads-up to the users before considering making the feature default. I would also encourage you to open a discussion to solicit feedback from other ORT users on ARM if they see speed-up for their models with this feature. It would provide greater confidence and a strong data point to turn it on by default.

Thanks for this contribution, we will review it shortly !

@milpuz01
Copy link
Contributor Author

Unfortunately, I don't have a comprehensive list of performance targets to be met to make the feature default. Since, the performance testing may not include all possible Conv shapes, I would like to err on the side of caution and atleast provide one release timeline heads-up to the users before considering making the feature default. I would also encourage you to open a discussion to solicit feedback from other ORT users on ARM if they see speed-up for their models with this feature. It would provide greater confidence and a strong data point to turn it on by default.

Thanks @hariharans29. I agree with erring on the side of caution. If this PR goes through and it is in the main release is it possible to add a note that we would like to make --enable_arm_neon_nchwc as a default in the future releases so that we can try to get some feedback on that via that route too? Thanks again.

@hariharans29
Copy link
Member

hariharans29 commented Jan 26, 2026

Unfortunately, I don't have a comprehensive list of performance targets to be met to make the feature default. Since, the performance testing may not include all possible Conv shapes, I would like to err on the side of caution and atleast provide one release timeline heads-up to the users before considering making the feature default. I would also encourage you to open a discussion to solicit feedback from other ORT users on ARM if they see speed-up for their models with this feature. It would provide greater confidence and a strong data point to turn it on by default.

Thanks @hariharans29. I agree with erring on the side of caution. If this PR goes through and it is in the main release is it possible to add a note that we would like to make --enable_arm_neon_nchwc as a default in the future releases so that we can try to get some feedback on that via that route too? Thanks again.

Thanks @milpuz01. The PR should go through in main eventually but I don't think it will go in 1.24.0 unfortunately as the release branch is cut and the bar to take in new code at this point is critical bug fixes and urgent customer asks only. I will try to take this in for 1.24.1 when it happens and sure I will add a note about considering making it default in one of the future releases, but ultimately, as discussed in the comment #27099 (comment), I expect the NchwcFloatKernel needs optimizations before considering that.

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds new AArch64 NEON assembly micro-kernels for NCHW, depthwise NCHWc, and pointwise NCHWc convolution, integrates them into the MLAS build, and updates NCHWc kernel-selection heuristics to prefer the asm kernels in selected shapes.

Changes:

  • Add new AArch64 .S convolution micro-kernels (NCHW, depthwise NCHWc, pointwise NCHWc) and wire them into the MLAS build.
  • Update ARM64 platform init and NCHWc execution heuristics to select asm kernels for pointwise (stride-1, larger tiles) and depthwise (wider outputs).
  • Remove the old intrinsics wrapper for the NCHW float kernel in the NCHWc NEON source file.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
cmake/onnxruntime_mlas.cmake Adds new AArch64 asm sources to the ARM NEON NCHWc MLAS build setup.
onnxruntime/core/mlas/lib/snchwc.cpp Adds ARM64 heuristics to prefer asm depthwise/pointwise kernels in “safe” cases.
onnxruntime/core/mlas/lib/sconv_nchwc_kernel_neon.cpp Removes the old NCHW float kernel wrapper implementation from the NCHWc NEON source file.
onnxruntime/core/mlas/lib/platform.cpp Switches ARM64 NCHW conv kernel default to asm; updates commentary around kernel choices.
onnxruntime/core/mlas/lib/mlasi.h Declares new asm kernel entry points for ARM64 NEON NCHWc.
onnxruntime/core/mlas/lib/aarch64/SconvKernelNeon.S Adds new NCHW convolution asm micro-kernel.
onnxruntime/core/mlas/lib/aarch64/SconvDepthwiseKernelNeon.S Adds new depthwise NCHWc asm micro-kernel (fast/slow path for padding).
onnxruntime/core/mlas/lib/aarch64/SconvPointwiseKernelNeon.S Adds new pointwise NCHWc asm micro-kernel (multi-output reuse).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@hariharans29
Copy link
Member

/azp run Linux QNN CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Windows ARM64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline

@azure-pipelines
Copy link

Azure Pipelines successfully started running 4 pipeline(s).

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@hariharans29
Copy link
Member

hariharans29 commented Feb 6, 2026

This change looks good to me. Thanks. Can you please address remaining comments (Copilot + mine) so that it can be merged ?

FYI - I have called out the NCHWc layout support on ARM in the 1.24 release notes, so that the community can give a try and share feedback/issues if any - https://github.com/microsoft/onnxruntime/releases. CC: @Rohanjames1997

@aviralagrawal
Copy link

Nice 🚀
Thanks @hariharans29!
I checked the contributor list in the 1.24.1 release and I think it's wrong. I don't see @Rohanjames1997 and a few other names in the list. And I see a few names that did not contribute to this release. What can explain this? 🙂

@hariharans29
Copy link
Member

hariharans29 commented Feb 6, 2026

Nice 🚀 Thanks @hariharans29! I checked the contributor list in the 1.24.1 release and I think it's wrong. I don't see @Rohanjames1997 and a few other names in the list. And I see a few names that did not contribute to this release. What can explain this? 🙂

Thanks for bringing this to my attention. I am not sure how the contributors list is generated myself. I ll pass along the information for folks to take a look. Meanwhile, I have added Rohan manually to the list and apologies.

EDIT: Filed as issue for tracking: #27274

Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

hariharans29 added a commit that referenced this pull request Feb 10, 2026
### Description

Initially the NCHWc code was built only on Mac CIs to keep the build
path regression-free.

There were some Linux-specific paths introduced in
#26838 and there is more
community interest in contributing to these code paths. See
#27099. Hence, it makes
sense to keep these code paths built and tested on Linux and Windows
too.

### Motivation and Context
Improve CI quality with regards to ARM64 NCHWc builds

CC: @Rohanjames1997
@hariharans29
Copy link
Member

/azp run Linux QNN CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Windows ARM64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline

@hariharans29
Copy link
Member

Can you please rebase with main ?

@azure-pipelines
Copy link

Azure Pipelines successfully started running 4 pipeline(s).

milpuz01 and others added 10 commits February 12, 2026 22:12
Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
)

<!-- Describe your changes. -->

<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
Description
Enables the file mapping of weights as well as the overall context bin.
This feature is currently only enabled for ARM64 WIN devices

Motivation and Context
Currently, when reading the context bin, ORT allocates a large buffer on
the heap. Assuming the same model is used, each ORT session will
allocate a buffer for the context bin. This is incredibly wasteful when
large models are used. Instead, WIN file mapping can be leveraged to map
the context bin, then every time a context needs to be created with the
context bin, the pointer to the context bin can be retrieved and used
instead of some pre-allocated buffer, thus making QNN EP more
memory-efficient. In the case of multiple ORT sessions, the context bin
will only be loaded once for all sessions, increasing memory efficiency
and overall initialization performance. This is very useful regarding
the use of LLMs going forward.

---------

Co-authored-by: quic_calvnguy <quic_calvnguy@quic_inc.com>
…ft#27151)

Previously in `MatMulReadFnSource()` we use duplicated code to read data
from two inputs `a` and `b`. This patch implements another overload of
`MatMulReadFnSource()` to only read data from one input to reduce
duplicated code and get ready for further use.
Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
…ck spill

Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
@milpuz01 milpuz01 force-pushed the aarch64_convolutions branch from 2d05853 to bd38b0e Compare February 12, 2026 22:25
Signed-off-by: Milos Puzovic <milos.puzovic@arm.com>
@milpuz01
Copy link
Contributor Author

Can you please rebase with main ?

Just rebased.

@hariharans29
Copy link
Member

/azp run Linux QNN CI Pipeline,Win_TRT_Minimal_CUDA_Test_CI,Windows ARM64 QNN CI Pipeline,Windows GPU Doc Gen CI Pipeline

@azure-pipelines
Copy link

Azure Pipelines successfully started running 4 pipeline(s).

@hariharans29 hariharans29 enabled auto-merge (squash) February 13, 2026 02:29
@hariharans29 hariharans29 merged commit bd8f781 into microsoft:main Feb 13, 2026
88 checks passed
@milpuz01 milpuz01 deleted the aarch64_convolutions branch February 13, 2026 07:35
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants