Validate MaxpoolWithMask kernel rank matches input spatial rank#29253
Open
titaiwangms wants to merge 2 commits into
Open
Validate MaxpoolWithMask kernel rank matches input spatial rank#29253titaiwangms wants to merge 2 commits into
titaiwangms wants to merge 2 commits into
Conversation
Residual hardening for MaxpoolWithMask::Compute: the 1D/2D/3D dispatch is selected by kernel_shape.size() and reads x_shape[2..4] / output_dims[2..4] accordingly, but the only rank guard was x_shape.NumDimensions() >= 3. A kernel_shape rank that exceeds the input spatial rank (e.g. rank-3 X with a 2D kernel_shape) caused out-of-bounds reads of x_shape[3]/x_shape[4] and output_dims[3]/output_dims[4]. ONNX shape inference catches this for shaped ONNX models, but the kernel must validate at Compute time for defense in depth (ORT-format / shapeless models bypass shape inference). Add ORT_RETURN_IF_NOT(kernel rank == input spatial rank) alongside the existing mask/input guards, and a kExpectFailure regression test that bypasses ONNX shape inference via AddShapeToTensorData(false) to exercise the Compute-time guard directly. Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot] Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… check
Follow-up polish on the kernel-rank guard. The 1D/2D/3D dispatch switch has a
default case that already returns INVALID_ARGUMENT, so an oversized kernel rank
was never a silent/uninitialized-output bug, but the default produced a vague
'Unsupported pooling size :' message late in Compute after allocating the output.
Add an explicit early guard requiring the input spatial rank (== pooling kernel
rank) to be in {1, 2, 3}, with a clear message, alongside the existing guards and
before output allocation. Extract the repeated x_shape.NumDimensions() - 2 into a
single input_spatial_rank local (defined after the rank >= 3 guard so it cannot
underflow) and reuse it. Add a MaxPoolWithMask_KernelRankTooLarge regression test
(rank-6 X + 4D kernel_shape) that passes the equality guard but trips the new
supported-rank guard, asserting the specific message substring.
Agent-signed-off: Developer (0b207188) [claude-opus-4.8 via copilot]
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR strengthens input validation for the CPU contrib op MaxpoolWithMask by ensuring the pooling kernel_shape rank matches the input’s spatial rank, and that only 1D/2D/3D pooling ranks are accepted. This prevents out-of-bounds reads in the rank-based dispatch logic and improves error diagnostics for malformed inputs.
Changes:
- Add pre-allocation validation in
MaxpoolWithMask::Computeto requirekernel_shape.size() == (X.rank - 2)and constrain supported spatial ranks to{1, 2, 3}with clear error messages. - Ensure the validation happens before output sizing/allocation to avoid unsafe shape-dependent indexing.
- Add unit tests for kernel-rank mismatch and kernel-rank-too-large scenarios (bypassing ONNX shape inference so the checks are exercised in
Compute()).
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cpu/maxpool_with_mask.h | Adds kernel-rank vs input-spatial-rank validation (and supported rank bounds) before output allocation to prevent invalid dispatch/indexing. |
| onnxruntime/test/contrib_ops/maxpool_mask_test.cc | Adds negative tests that reach Compute() and validate the new error paths for rank mismatch and unsupported rank. |
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
MaxpoolWithMask::Computenow validates that the pooling kernel rank equals the input spatial rank (and that the rank is one of the supported values {1, 2, 3}) before allocating the output, returning a clear error for malformed inputs instead of proceeding with a mismatched configuration.Changes
MaxpoolWithMask::Computethat the kernel rank matches the input spatial rank and is within the supported {1, 2, 3} range, with a descriptive error message.MaxPoolWithMask_KernelRankMismatchandMaxPoolWithMask_KernelRankTooLargeunit tests covering the new validation.Motivation
Improves input validation and error diagnostics for malformed
MaxpoolWithMaskinputs. CPU-only; no behavior change for valid inputs.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com