[MLAS] Fix scratch buffer over-allocation in batched/grouped NCHW Conv

[hariharans29]

Description

#26103 which was shipped with ORT 1.23.2 introduced logic to parallelize on the batch and/or group count parameters for NCHW Conv if either were > 1 to improve performance of batched grouped convolutions - prior to that the parallelization was only on the output spatial map. The per-thread scratch buffer allocated by the B/G parallelization path introduced in #26103 is excessive and is not needed for the algorithm being used. Under the hood, MlasConvOperation() uses a double tiled approach which means that the per-thread scratch buffer needs just enough memory to fit that tile of data and nothing more than that. This change fixes the conservative over-allocation which shows up as a memory bloat in OP's model in issue #28466

Note: the work partitioning across (batch × group) introduced in #26103 is unchanged — only the per-thread scratch buffer sizing is corrected. The batched/grouped Conv speedup from #26103 is retained.

Motivation and Context

Fix #28466

[Copilot]

Pull request overview

This PR fixes a large, unnecessary scratch-buffer allocation in the MLAS NCHW Conv path when parallelizing over (batch × group) for the MlasConvAlgorithmExpandThenGemmSegmented algorithm. The change keeps the existing work-partitioning behavior from #26103, but corrects the per-thread working-buffer sizing and indexing to match the actual tiled (StrideN×StrideK) im2col usage inside MlasConvOperation().

Changes:

• Fix per-thread working-buffer pointer arithmetic in the batched/grouped threaded Conv path to use MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD instead of OutputSize * K.
• Fix working-buffer sizing in MlasConvPrepare() for the batched/grouped segmented algorithm path to allocate TargetThreadCount * MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD.
• Add regression tests covering batched Conv3D working-buffer sizing (MLAS unit test) and a CPU batched Conv3D FusedConv execution case.

Reviewed changes

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

File	Description
onnxruntime/core/mlas/lib/convolve.cpp	Corrects batched/grouped segmented Conv working-buffer sizing and per-thread buffer indexing to match the tile-based im2col usage.
onnxruntime/test/mlas/unittest/test_conv2d.h	Adds a regression test to ensure batched Conv3D does not allocate a full im2col buffer per worker.
onnxruntime/test/contrib_ops/fused_conv_test.cc	Adds a CPU batched Conv3D FusedConv test to exercise the relevant execution path.

---

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

[tianleiwu]

APPROVE

Clean, well-motivated fix for a significant memory over-allocation regression in the MLAS batch/group convolution path (introduced in #26103, shipped with ORT 1.23.2).

Analysis

MlasConvOperation uses a double-tiled approach: it dynamically adjusts StrideN and StrideK (one doubles while the other halves), preserving the invariant StrideN × StrideK = MLAS_SGEMM_STRIDEN × MLAS_SGEMM_STRIDEK = 16,384. The im2col/vol2col functions write exactly CountK × CountN ≤ StrideK × StrideN elements to ColumnBuffer. This confirms each thread needs exactly MLAS_CONV_WORKING_BUFFER_SIZE_PER_THREAD — matching what the existing output-space threading path (MlasConvOperationThreaded) already uses.

The impact for the reported Conv3D case (InputShape=[32,64,64], Kernel=[7,7,7], BatchCount=12): old per-thread allocation was OutputSize × K ≈ 45M floats ≈ 172 MB/thread (over 2 GB total). The fix reduces this to 64 KB/thread (~768 KB total).

Both the allocation sizing and per-thread buffer indexing are updated consistently. Work partitioning across (batch × group) is untouched. Tests cover the fix from two angles: end-to-end crash/ASAN safety and direct buffer-size assertion.

One minor test improvement suggested inline.

[tianleiwu]

No remaining findings on the current head. The buffer-size fix is consistent with MlasConvOperation()'s tiled StrideN x StrideK working-set invariant, and the earlier test concern was addressed by switching the new FusedConv regression to non-zero inputs with numerically meaningful expected values. Approving.

[hariharans29]

Will merge this after verifying a customer model perf on ARM server. Their model heavily relies on this path and I want to make sure its perf is not affected (in theory it should not be, but worth double checking)

[hariharans29]

Will merge this after verifying a customer model perf on ARM server. Their model heavily relies on this path and I want to make sure its perf is not affected (in theory it should not be, but worth double checking)

Double checked the perf numbers with this change. The perf stays the same. Merging.

[albertopasqualetto]

Hey @hariharans29 when will this code be released officially?

[hariharans29]

Hey @hariharans29 when will this code be released officially?

It will be part of 1.27

[albertopasqualetto]

Hey @hariharans29 when will this code be released officially?

It will be part of 1.27

Is there an indicative release date? (The Web page is not updated)

[hariharans29]

Hey @hariharans29 when will this code be released officially?

It will be part of 1.27

Is there an indicative release date? (The Web page is not updated)

I think 1.27 Python wheels should already be available on PyPi. If the nuget packages and/or other release artifacts are not yet available, it will be very soon (folks are working on it as we speak)

[albertopasqualetto]

Hey @hariharans29 when will this code be released officially?

It will be part of 1.27

Is there an indicative release date? (The Web page is not updated)

I think 1.27 Python wheels should already be available on PyPi. If the nuget packages and/or other release artifacts are not yet available, it will be very soon (folks are working on it as we speak)

You are right, it is on pypi. Thanks! I wasn't seeing the release on github.

[hariharans29]

Hey @hariharans29 when will this code be released officially?

It will be part of 1.27

Is there an indicative release date? (The Web page is not updated)

I think 1.27 Python wheels should already be available on PyPi. If the nuget packages and/or other release artifacts are not yet available, it will be very soon (folks are working on it as we speak)

You are right, it is on pypi. Thanks! I wasn't seeing the release on github.

We will officially announce it on GH once all the release artifacts are published and the entire release process is officially complete