Improve MLAS NCHWc conv thread utilization via cost-weighted work par…

[mirounga]

…titioning

The NCHWc convolution algorithms (pointwise and direct) partition work uniformly by item index over FilterSetCount x OutputHeight items, but the cost of an item is proportional to its FilterSet's FilterCount: full sets process FilterSetSize (4) NCHWc blocks while a ragged last set (when Cout/16 % 4 != 0) processes as few as 1. Threads landing on full sets do up to 4x the FLOPs of threads on the tail set, capping scheduling efficiency at 62-75% for common shapes (e.g. Cout=96: sets of 4+2 blocks -> 0.75 efficiency at 4 threads).

Replace the uniform index split with a split proportional to FLOP cost, measured in block-rows (one NCHWc output block x one output row): MlasPartitionWork distributes cost intervals, and CostToWorkIndex maps cost boundaries back to work-item boundaries with exact coverage and no overlap. Applied to both MLAS_NCHWC_CONV_POINTWISE_ALGORITHM and MLAS_NCHWC_CONV_NCHWC_ALGORITHM; the NCHW first-layer, depthwise, and pooling algorithms keep the uniform split (no imbalance there).

Outputs are bitwise identical: repartitioning moves whole output-row items between threads without changing any element's accumulation order.

Also:

• Add session config "mlas.nchwc_conv_max_input_channel_batch" to override the pointwise algorithm's input-channel batch (default 128, rounded up to a BlockSize multiple; 0/unset keeps the default) via MLAS_BACKEND_KERNEL_SELECTOR_CONFIG, for perf experimentation.
• Fix copy-paste guards in the SconvKernelAvx512F.S ReLU post-process (FilterCount 2 -> 3/4 for the zmm14/18/22 and zmm15/19/23 rows); behavior was accidentally correct but the FilterCount=2/OutputCount=6 variant emitted 6 vmaxps on dead registers.

Measured on AMD Ryzen AI MAX+ 395 (Zen 5, AVX-512), batch=1 fp32, intra_op=4, 512-iteration medians vs unmodified main:

yolox_tiny 8.009 -> 7.149 ms (-10.7%) 3x3-conv heavy, 86% of its
3x3 time in ragged shapes
mv3_large 0.962 -> 0.922 ms (-4.2%)
mv3_small 0.458 -> 0.455 ms (-1% noise)
mobileclip unchanged (+/-1% noise; its conv shapes are 64-aligned)

Hottest yolox shape (96->96 3x3 @52x52): 939 -> 1167 GFLOP/s (91% of machine peak). All 29,910 onnxruntime_mlas_test cases pass.

Description

Motivation and Context

[Copilot]

Pull request overview

This PR improves CPU MLAS NCHWc convolution thread utilization by switching from uniform work-item partitioning to a cost-weighted partitioning scheme, reducing load imbalance for ragged last filter sets. It also adds a new session config knob to tune the pointwise (1x1) NCHWc input-channel batching behavior for performance experimentation.

Changes:

• Implement cost-weighted work partitioning for MLAS NCHWc direct and pointwise convolution algorithms via CostToWorkIndex + PrepareWorkWeighted.
• Add session config key mlas.nchwc_conv_max_input_channel_batch and plumb it into MLAS backend kernel selector config parsing.
• Use the new config value to override (and block-size-round) the pointwise conv max input-channel batch (default remains 128).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/providers/cpu/mlas_backend_kernel_selector_config_utils.h	Parse the new NCHWc pointwise tuning option from session config into MLAS selector config.
onnxruntime/core/mlas/lib/snchwc.cpp	Add cost-weighted partitioning for NCHWc conv work distribution and apply configurable pointwise input-channel batching.
onnxruntime/core/mlas/inc/mlas.h	Extend MLAS_BACKEND_KERNEL_SELECTOR_CONFIG with nchwc_conv_max_input_channel_batch.
include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h	Define and document the new session option key mlas.nchwc_conv_max_input_channel_batch.

Suppressed comments (1)

onnxruntime/core/mlas/lib/snchwc.cpp:723

• PrepareWorkWeighted calls SeekToWork(WorkIndexBegin) even when CostCount == 0, which can happen when TotalCost < tids. In that case WorkIndexBegin may equal TotalWork, and SeekToWork will advance pointers past the end of the input/output/filter buffers (undefined pointer arithmetic). Only seek when there is actual work to execute.

        WorkRemaining = WorkIndexEnd - WorkIndexBegin;

        SeekToWork(WorkIndexBegin);