Validate 'num_scan_inputs' attribute in Scan kernel construction - #31668
Open
titaiwangms wants to merge 3 commits into
Open
Validate 'num_scan_inputs' attribute in Scan kernel construction#31668titaiwangms wants to merge 3 commits into
titaiwangms wants to merge 3 commits into
Conversation
The Scan operator's 'num_scan_inputs' attribute is used to split the node's variadic inputs into loop state variables and scan inputs: num_loop_state_variables = num_variadic_inputs - num_scan_inputs num_scan_outputs = num_outputs - num_loop_state_variables Both subtractions were previously unguarded. A 'num_scan_inputs' value outside [0, num_variadic_inputs] produces a negative num_loop_state_variables, which is later used unguarded as a loop-start index into the node's inputs during Compute(), reading out of bounds. Add ORT_ENFORCE checks in scan::detail::Info::Info (shared by the CPU Scan-8/9 kernels and reused as-is by the CUDA Scan kernel) so invalid attribute values are rejected with a clear error during kernel construction, before any out-of-range arithmetic is used for indexing. Add regression tests for opset 8 and opset 9+ covering an out-of-range 'num_scan_inputs' value. Opset 9+ Scan is caught earlier, during graph resolution's standard ONNX-level shape inference, before the kernel is even constructed, so its test asserts on the shape-inference error instead of the kernel-level message. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 77dcaf1b-748a-4379-94a7-478f7a924d73
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens the CPU (and shared CUDA) Scan kernel setup by validating the num_scan_inputs attribute early, preventing negative/invalid derived counts that can later lead to out-of-range indexing. It also adds regression coverage for invalid num_scan_inputs for opset 8 and opset 9+.
Changes:
- Add
ORT_ENFORCEvalidation inscan::detail::Info::Infoto rejectnum_scan_inputsoutside[0, num_variadic_inputs]and ensure the implied loop-state-variable count is compatible with the node’s outputs. - Add new CPU provider tests that exercise invalid
num_scan_inputsfor opset 8 (kernel-level validation) and opset 9+ (graph-resolution/shape-inference failure path). - Guard the new negative tests with
#if !defined(ORT_NO_EXCEPTIONS).
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| onnxruntime/core/providers/cpu/controlflow/scan_utils.cc | Adds constructor-time validation for num_scan_inputs and derived counts to prevent invalid indexing. |
| onnxruntime/test/providers/cpu/controlflow/scan_test.cc | Adds regression tests for out-of-range num_scan_inputs across opset 8 and opset 9+. |
Comment on lines
+43
to
+45
| ORT_ENFORCE(num_scan_inputs >= 0 && num_scan_inputs <= num_variadic_inputs, | ||
| "Invalid 'num_scan_inputs' of ", num_scan_inputs, ". Value must be between 0 and ", | ||
| num_variadic_inputs, " (the number of variadic inputs) inclusive."); |
Comment on lines
+1072
to
+1079
| Model model("NumScanInputsExceedsVariadicInputs_v9", false, ModelMetaData(), PathString(), | ||
| IOnnxRuntimeOpSchemaRegistryList(), {{"", 11}}, {}, DefaultLoggingManager().DefaultLogger()); | ||
| auto& graph = model.MainGraph(); | ||
| ASSERT_STATUS_OK(CreateSubgraph(graph, options)); | ||
| auto& proto = graph.ToGraphProto(); | ||
|
|
||
| ScanOpTester test{11}; | ||
| test.AddAttribute("body", proto); |
Comment on lines
+1094
to
+1100
| // Unlike opset 8, opset 9+ Scan is type/shape-inferred via the standard ONNX inferencing path, | ||
| // which detects the inconsistent variadic split from the bad 'num_scan_inputs' value and rejects | ||
| // the model during graph resolution -- before the node's kernel (and its own attribute validation) | ||
| // is ever constructed. So the failure surfaces as a graph attribute inferencing error here, rather | ||
| // than the kernel-level message produced by the opset 8 test above. | ||
| test.Run(OpTester::ExpectResult::kExpectFailure, "Graph attribute inferencing failed", | ||
| options.excluded_provider_types); |
- Use narrow<int> instead of static_cast<int> when passing num_scan_inputs to Info's constructor, so an out-of-int-range attribute value fails predictably instead of silently truncating to an in-range value that could bypass Info's own bounds check. - Fix opset mismatch in the Scan9 negative test (was building an opset 11 model) and tighten its expected failure message to a specific ShapeInferenceError substring instead of a generic wrapper message. - Fix clang-format continuation-line indentation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 77dcaf1b-748a-4379-94a7-478f7a924d73
- Move num_scan_inputs validation into Scan<8>/Scan<9>::Init, immediately after the attribute is read, via a new shared ValidateNumScanInputs helper. This closes a gap where the value was previously used to size directions/axes vectors and derive loop-state/output counts before Info's own (later) validation ever ran. - Tighten the lower bound from 0 to 1: the ONNX Scan spec requires one or more scan_input tensors, and a value of 0 left downstream sequence-length state unset, which could result in loop-state outputs never being written. - Improve the second check's error message to not misattribute a too-few-outputs condition solely to num_scan_inputs. - Add regression tests for num_scan_inputs of 0 and negative on opset 8, and update the opset 8 exceeds-inputs test's expected message for the new lower bound. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 77dcaf1b-748a-4379-94a7-478f7a924d73
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
The
Scanoperator'snum_scan_inputsattribute determines how many of the node's variadic inputs are scan inputs (the remainder are loop state variables). Inscan::detail::Info::Info(shared by the CPUScan-8/Scan-9kernels, and reused as-is by the CUDAScankernel), this attribute was used in two unguarded subtractions:If
num_scan_inputsis outside[0, num_variadic_inputs],num_loop_state_variablesbecomes negative. That value is later used, unguarded, as a loop-start index into the node's inputs duringCompute(), resulting in out-of-range indexing.This PR adds validation of
num_scan_inputs(and the derivednum_loop_state_variables) in the constructor, rejecting invalid values up front with a clear error message before any arithmetic derived from them is used for indexing.Testing
Added regression tests for both opset 8 and opset 9+ with an out-of-range
num_scan_inputsvalue:Both tests are guarded by
#if !defined(ORT_NO_EXCEPTIONS)since they rely on exception-based failure reporting.Ran locally (CPU-only Debug build):
onnxruntime_provider_test --gtest_filter='Scan*'— all 34 tests pass, including the 2 new ones.