Guard MlasConvPrepare working-buffer products with SafeInt

[Akshay Sonawane (apsonawane)]

This pull request strengthens the safety of arithmetic operations in the convolution implementation by using SafeInt to prevent integer overflows, especially in cases where tensor shapes may be attacker-controlled. It also removes some compiler-specific warning pragmas that are no longer needed due to these changes.

Enhanced integer safety in convolution calculations:

• Replaced raw size_t multiplications with SafeInt<size_t> for calculating input/output sizes and kernel dimensions, ensuring that any arithmetic overflow is caught and handled appropriately. This is particularly important for preventing security vulnerabilities from attacker-controlled tensor shapes. [1] [2]
• Used SafeInt<size_t> for all working buffer size calculations, including those involving thread counts and batch/group products, to prevent buffer overflows and ensure correct memory allocation. [1] [2]

Code cleanup and maintenance:

• Removed MSVC-specific warning suppression pragmas since SafeInt now guards against arithmetic overflow, making these warnings obsolete. [1] [2]
• Added the core/common/safeint.h include to provide access to the SafeInt functionality.

[Copilot]

Pull request overview

This PR hardens MLAS convolution preparation (MlasConvPrepare) against size_t overflow by moving several shape-dependent products into SafeInt<size_t> computations, and adds unit tests that validate overflow cases throw rather than silently producing wrapped sizes.

Changes:

• Switched InputSize, OutputSize, and K accumulation in MlasConvPrepare to SafeInt<size_t> and materialized them as size_t only after checked multiplication.
• Guarded working-buffer element-count products (OutputSize * K, TargetThreadCount * per_thread) using SafeInt<size_t>.
• Added a new unit test file that constructs overflow-inducing shapes and validates exception behavior.

Reviewed changes

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

File	Description
onnxruntime/core/mlas/lib/convolve.cpp	Uses SafeInt<size_t> to validate shape-derived products and working-buffer size calculations in MlasConvPrepare.
onnxruntime/test/mlas/unittest/test_conv_prepare_safeint.cpp	Adds gtest coverage to ensure overflow cases throw (instead of wrapping) and that small “happy path” inputs still work.

[Tianlei Wu (tianleiwu)]

Summary

Focused, well-scoped defense-in-depth change. The cross-tensor products in MlasConvPrepare (InputSize/OutputSize/K accumulators, the OutputSize * K im2col buffer size, and the BatchCount * GroupCount thread clamp) are now computed in SafeInt's checked domain, so attacker-controlled shapes that would previously wrap size_t and under-size the working buffer now throw instead of producing a heap-buffer-overflow on the downstream im2col write. Overall this looks good — one non-blocking note below.

Strengths

• Standalone-MLAS split is correct — SafeInt.hpp under BUILD_MLAS_NO_ONNXRUNTIME, core/common/safeint.h otherwise. safeint_interface is on the include path for every MLAS target (the onnxruntime_mlas.cmake foreach loop), so both branches compile. This resolves the earlier concern about an unconditional ORT-core include in MLAS.
• SafeInt<size_t> is the right primitive here (vs the ORT-only SafeMul<> helper in core/common/safeint.h), because it also works in the standalone branch where the ORT helper isn't available.
• BatchGroupProduct reuse removes the duplicated BatchCount * GroupCount multiply while preserving behavior; the later static_cast<ptrdiff_t>(BatchGroupProduct) is only reached when the product is <= MaximumThreadCount, so the narrowing stays safe.
• Good test construction — kHalfShift is portable across 32/64-bit pointer widths; the 3D shapes bypass the KleidiAI 2D MlasConvPrepareOverride; and the non-unit stride keeps execution out of the pointwise GemmDirect early return so each guarded product is actually exercised. The !ORT_NO_EXCEPTIONS && !BUILD_MLAS_NO_ONNXRUNTIME gate and matching #endif comment are correct, and the tests are auto-collected by the *.cpp glob in the MLAS unittest target.

Non-blocking note

• The removed #pragma warning(disable : 26451) push/pop spanned a broad region (from before the anonymous namespace, through MlasConvSupportsDenseChannelsLast2DFloatKernel, MlasConvSupportsDepthwiseChannelsLast2DFloatKernel, and MlasConvPrepare), whereas the SafeInt conversion only touched MlasConvPrepare. The remaining arithmetic in that region is all in the size_t domain, so 26451 most likely won't re-fire — but it's worth confirming the MSVC static-analysis leg (onnxruntime_ENABLE_STATIC_ANALYSIS -> /analyze) stays warning-clean before merge, since 26451 is a /analyze-only check that is not covered by the target's global /wd suppressions. If it does re-fire, scoping the pragma to just around MlasConvPrepare (rather than removing it entirely) would be the minimal fix.

The earlier automated review comments (unconditional include, test gating, #endif comment text, hard-coded PR number) all appear to be addressed already in the current head.

[Tianlei Wu (tianleiwu)]

Re-review on the latest head

The one non-blocking note from my previous round is now addressed: the #pragma warning(disable : 26451) push/pop is re-added but scoped tightly around MlasConvPrepare instead of spanning the anonymous namespace and the MlasConvSupports* helpers. I confirmed the now-unsuppressed helpers above (ComputeChannelsLastDilatedKernelSize / ComputeChannelsLastConvOutSize) only contain size_t multiplies under #if defined(USE_KLEIDIAI) && defined(MLAS_TARGET_ARM64), so they aren't compiled on the MSVC /analyze leg and 26451 won't re-fire there. The push/pop remain balanced (pop is still at end-of-function).

Verified

• Overflow guard placement is complete for the generic path: the InputSize/OutputSize/K accumulators, the OutputSize * K im2col buffer size, and the BatchCount * GroupCount thread clamp are all computed in SafeInt's checked domain. The downstream byte allocation in the caller is already guarded (alloc->Alloc(sizeof(float) * SafeInt<size_t>(WorkingBufferSize)) in core/providers/cpu/nn/conv.cc), so element-count and byte-count math are both covered.
• Portability is correct — SafeInt.hpp under BUILD_MLAS_NO_ONNXRUNTIME and core/common/safeint.h otherwise; safeint_interface is on every MLAS target's include path, and BUILD_MLAS_NO_ONNXRUNTIME is WASM-only where the tests are gated off. Using the raw SafeInt<size_t> primitive (rather than the ORT-only SafeMul<> helper) is the right call since it also compiles in the standalone branch.
• Tests exercise the guarded products — kHalfShift overflows on both 32/64-bit widths, 3D shapes bypass the KleidiAI 2D override, and the non-unit stride keeps execution off the pointwise GemmDirect early-return so each SafeInt product is actually reached. Auto-collected by the *.cpp glob in the MLAS unittest target.

All earlier automated-review concerns (unconditional ORT-core include, test gating, #endif comment, hard-coded PR number) also remain addressed. LGTM.