[MLAS][KleidiAI]Catlaw01/sgemm epilogue neon opt

[Laan33]

Description

This change updates the KleidiAI SGEMM post-processing path in onnxruntime/core/mlas/lib/kleidiai/sgemm_kleidiai.cpp with two parts:

• Correctness fix: in the alpha == 0 || K == 0 fast path, beta handling is now applied for every batch entry (not just batch 0), so batched SGEMM behaviour is correct.
• NEON SGEMM epilogue optimisation: adds a vectorised alpha/beta post-processing path for contiguous outputs, with guarded fallback to scalar for non-contiguous or small cases. The 2D epilogue path also routes contiguous tiles through the contiguous 1D epilogue path to enable vectorisation.

Motivation and Context

This change addresses correctness and performance in the SGEMM post-processing stage:

• The batched alpha == 0 || K == 0 path previously used only Data[0], which could produce incorrect results for BatchSize > 1.
• The post-processing loop (C = alpha * (A*B) + beta * C) is a known latency contributor when memcpy fast paths are not applicable. The NEON epilogue changes are intended to reduce this cost on supported ARM platforms while preserving existing fallback behaviour.

[Laan33]

@microsoft-github-policy-service agree company="Arm"

[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

Updates the KleidiAI-backed SGEMM post-processing (alpha/beta “epilogue”) in MLAS to fix batched correctness in an early-exit path and to improve ARM NEON performance for contiguous outputs.

Changes:

• Fix batched alpha == 0 || K == 0 fast path to apply beta reduction for every batch entry.
• Add a contiguous-only NEON-vectorized alpha/beta epilogue path with scalar fallback for small/non-contiguous cases.
• Route contiguous 2D tiles through the 1D contiguous path to reuse vectorization.

---

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

[hariharans29]

Please includ ethis - #27618 when it goes through eventually

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

Can you please fix the CI issues ?

[Laan33]

Pushed new commits there to fix issues caused by when I merged main into the branch

[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).