Fix MatMulInteger heap-buffer-overflow with mismatched 1D inputs#27965
Open
bmehta001 wants to merge 2 commits into
Open
Fix MatMulInteger heap-buffer-overflow with mismatched 1D inputs#27965bmehta001 wants to merge 2 commits into
bmehta001 wants to merge 2 commits into
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR addresses a heap out-of-bounds read in MatMulInteger by adding missing shape validation for the 1D×1D (vector dot-product) case, plus additional validation changes across Slice/Pad/Split implementations in native and JS Web backends.
Changes:
- Add a missing K-dimension compatibility check for 1D×1D MatMul in
MatMulComputeHelper. - Add regression tests for MatMulInteger 1D dimension mismatch and a valid 1D dot product.
- Add/adjust validation in WebGPU/CUDA/CPU (and JS Web) implementations for Slice/Pad/Split inputs and computed shapes.
Reviewed changes
Copilot reviewed 12 out of 12 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/providers/cpu/math/matmul_integer_test.cc | Adds regression tests for 1D MatMulInteger mismatch and valid dot product. |
| onnxruntime/core/providers/cpu/math/matmul_helper.h | Adds missing K-dimension check for 1D×1D MatMul. |
| onnxruntime/core/providers/webgpu/tensor/slice.cc | Adds axis range validation and new output-dimension validation. |
| onnxruntime/core/providers/webgpu/tensor/pad.cc | Adds new output-dimension validation. |
| onnxruntime/core/providers/cuda/tensor/split.cc | Adds validation to reject negative values in split input. |
| onnxruntime/core/providers/cuda/tensor/pad.cc | Adds new output-dimension validation. |
| onnxruntime/core/providers/cpu/tensor/split.h | Adds validation to reject negative values in split input. |
| onnxruntime/core/providers/cpu/tensor/pad.cc | Adds new output-dimension validation in reshaped and final output dims. |
| js/web/lib/wasm/jsep/webgpu/ops/slice.ts | Adds JS WebGPU Slice output-dimension validation. |
| js/web/lib/wasm/jsep/util.ts | Fixes normalizeAxis() condition and adds padShape output-dimension validation. |
| js/web/lib/onnxjs/util.ts | Fixes normalizeAxis() condition and adds padShape output-dimension validation. |
| js/web/lib/onnxjs/backends/webgl/ops/slice.ts | Adds WebGL Slice output-dimension validation. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
bmehta001
force-pushed
the
fix/matmulinteger-heap-oob-read
branch
from
e0ffdbb to
f656138
Compare
MatMulComputeHelper::Compute skipped K-dimension validation when both inputs were 1D vectors (the vector x vector dot-product case). This allowed mismatched shapes like A=[5] and B=[1] to reach the MLAS GEMM backend, which assumed B had K elements and read past its allocation. Add the missing ORT_RETURN_IF_NOT check in the num_output_dims==0 branch to reject mismatched K dimensions before dispatch. This closes the heap-buffer-overflow (CVE-class) in MlasGemmQuantCopyPackB triggered via MatMulInteger with malformed 1D quantized inputs. Files changed: - onnxruntime/core/providers/cpu/math/matmul_helper.h - onnxruntime/test/providers/cpu/math/matmul_integer_test.cc Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
bmehta001
force-pushed
the
fix/matmulinteger-heap-oob-read
branch
from
f656138 to
f592854
Compare
bmehta001
changed the title
bmehta001
force-pushed
the
fix/matmulinteger-heap-oob-read
branch
2 times, most recently
from
6ecb6fe to
bb5b207
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
bmehta001
force-pushed
the
fix/matmulinteger-heap-oob-read
branch
2 times, most recently
from
413eba0 to
c376cea
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Address review feedback: tests now use actual 1D inputs to exercise the vector x vector code path this PR fixes. AddShapeToTensorData(false) skips ONNX shape inference which rejects 1D MatMulInteger inputs at graph build time. Also removed unnecessary static_cast in matmul_helper.h to match the style of other K-dimension checks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
bmehta001
force-pushed
the
fix/matmulinteger-heap-oob-read
branch
from
c376cea to
9f5b222
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fix a heap-buffer-overflow (out-of-bounds read) in
MatMulIntegertriggered when both inputs are 1D vectors with mismatched K dimensions (e.g.,A=[5],B=[1]).Root Cause
MatMulComputeHelper::Computeinmatmul_helper.hvalidated K-dimension mismatches for all input rank combinations except the vector×vector case (both 1D,num_output_dims == 0). This allowed mismatched shapes to pass through to the MLAS GEMM backend, which assumed B had K elements and read past its allocation inMlasGemmQuantCopyPackB.Fix
Added the missing
ORT_RETURN_IF_NOT(K_ == right_shape[0], "MatMul dimension mismatch")check in the vector×vector branch ofMatMulComputeHelper::Compute— matching the pattern already used in the other branches.Testing
MatMulInteger_1D_DimensionMismatch: verifies mismatched 1D shapes are rejectedMatMulInteger_1D_Valid: verifies correct 1D dot products still workMotivation and Context
Security fix — prevents out-of-bounds heap reads when processing untrusted ONNX models with malformed
MatMulIntegerinputs on the CPUExecutionProvider.