[MLAS] Simplify & optimize Arm64 NCHWc Convolution kernels

[Rohanjames1997]

Description

This PR makes the following changes:

1. Reroutes Pointwise Convolution to use GEMM, as it is essentially the same. This unlocks the performance benefits of GEMM too, as seen in the performance section below.
2. Eliminate some branches in the code. This is achieved by using in-built MLAS functions like MlasBlendFloat32x4
3. Expand the unit test coverage of NCHWc Conv kernels to catch edge cases.

Performance

This speeds up any Conv model that uses the pointwise kernel.
For example, Mobilenet inference speeds up from 500 inf/sec to 550 inf/sec.

Testing

• Build passed: ./build.sh --config=Release --build_shared_lib --parallel --cmake_extra_defines onnxruntime_USE_ARM_NEON_NCHWC=ON
• Unit tests passed: ./build/Linux/Release/onnxruntime_mlas_test --gtest_filter=Conv2dNchwc_*
• Perf: ./build/Linux/Release/onnxruntime_perf_test -x 32 -I -m times -r 2000 ~/scripts/mobilenet.onnx

Happy to run additional perf tests as required.

[Rohanjames1997]

@hariharans29 this may be of interest to you 🙂

TIA!

[hariharans29]

/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]

Azure Pipelines successfully started running 4 pipeline(s).

[hariharans29]

Nice, Thank you! Is there performance uplift from removing the branches in the kernel or is the main perf benefit coming from the pointwise kernel switch to using the Gemm kernel ?

[Rohanjames1997]

The majority of the perf gain is from using GEMM. Making it branchless results in no noticeable gain but I realized that it's better SIMD practice, and MLAS has great support for it with these built-in functions

[Rohanjames1997]

I ran the failing CI tests - onnxruntime_global_thread_pools_test and onnxruntime_test_all - on the main branch & they fail even there 🤔

Looks like these are Required CI's too. Any idea what can be done?

[hariharans29]

/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]

Azure Pipelines successfully started running 4 pipeline(s).

[hariharans29]

I ran the failing CI tests - onnxruntime_global_thread_pools_test and onnxruntime_test_all - on the main branch & they fail even there 🤔

Looks like these are Required CI's too. Any idea what can be done?

Hmm - I don't see it on other PRs. Let's see what happens on this run.

[Rohanjames1997]

Weird. Can you help me understand why this job is failing?

The three failing UTs look unrelated to the changes in this PR, but I wanted to be sure.

[hariharans29]

I think it is hitting the same test segfault that @damdoo01-arm was hitting on Mac M4 (with test ConvNoBiasAddFusion) when NCHWC was enabled (and KleidiAI is enabled too). M4 is an SME compliant machine. I don't think we have SME based machines in our CI pool, so possibly the contents of this PR is triggering the segfault on non-SME machines now. Can you take a look ?

(I am not able to locate the PR where @damdoo01-arm originally raised it and tagged you. Do you recall ?)

[hariharans29]

There are some other NCHWC tests failing on the latest run too. All failing tests and the segfault seem related to the PR (or the NCHWC feature in general). Can you please take a look ?

[Rohanjames1997]

Yup, the discussion was on this PR: #26171

Interesting, thanks @hariharans29. I am going to raise another dummy PR with incremental commits, and I might need you to run CI on it so that I can pinpoint the failure. I want to verify if the kernel is at fault or if there is spillover in the test suite itself (since I can repro the failure on the dev branch as well).

[hariharans29]

/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]

Azure Pipelines successfully started running 4 pipeline(s).

[Copilot]

Pull request overview

This PR optimizes Arm64 NCHWc convolution kernels through two key improvements: (1) routing pointwise (1x1) convolution through optimized GEMM operations instead of manual loops, and (2) making all convolution kernels branchless using SIMD mask operations to eliminate branch mispredictions in hot loops.

Key Changes:

• Refactored pointwise convolution to use batched GEMM, leveraging existing GEMM optimizations
• Converted conditional operations (accumulate output, bias addition, ReLU activation) to branchless SIMD operations using masks
• Added helper function LoadInputVectorWithBounds for bounds-checked vector loading in depthwise convolution

Reviewed changes

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

File	Description
onnxruntime/core/mlas/lib/sconv.h	Added LoadInputVectorWithBounds helper for depthwise kernel bounds checking
onnxruntime/core/mlas/lib/sconv_kernel_neon.cpp	Refactored all three convolution kernels (standard, depthwise, pointwise) with branchless operations and GEMM routing for pointwise
onnxruntime/test/mlas/unittest/test_conv2d_fixture.h	Added test cases for pointwise convolution with 32 and 144 input channels

---

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

[hariharans29]

Hi @Rohanjames1997 - Thank a lot for fixing the tests! Could you take a look at Copilot's comments ?

[Rohanjames1997]

Done! Thank you for helping me triage & running the CI!

[hariharans29]

/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]

Azure Pipelines successfully started running 4 pipeline(s).

[Rohanjames1997]

Thanks @hariharans29 ! Looks like it's all ✅