Validate SkipLayerNorm prepacked lengths - #31676
Validate SkipLayerNorm prepacked lengths#31676Akshay Sonawane (apsonawane) wants to merge 1 commit into
Conversation
Co-authored-by: Copilot <223556219@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
The new size-validation code has a signed/unsigned arithmetic issue that should be fixed, and the added checks for beta/bias lack corresponding unit test coverage.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
This PR tightens runtime validation around prepacked (FP16→FP32 converted) inputs for the contrib CPU SkipLayerNormalization / SkipSimplifiedLayerNormalization kernels, so malformed or mismatched prepacked tensors are rejected early with clearer error messages. It also adds unit tests to ensure the new validation is exercised.
Changes:
- Added member state to track prepacked
gamma/beta/biastensor lengths and initialized them in the kernel constructor. - Added explicit size checks in
SkipLayerNorm::Computeto validate prepackedskip/gamma/beta/biaslengths againsthidden_size, returning descriptive failures. - Added two unit tests verifying failures for incorrect prepacked
gammaandskipsizes.
File summaries
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc | Adds unit tests for prepack rejection on invalid gamma and skip sizes. |
| onnxruntime/contrib_ops/cpu/skip_layer_norm.h | Introduces new members to track prepacked tensor sizes. |
| onnxruntime/contrib_ops/cpu/skip_layer_norm.cc | Implements hidden-size validation and prepacked size checks; records prepacked sizes in PrePack. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 2
- Review effort level: Lite
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
| int hidden_size = static_cast<int>(input_dims[input_dims_size - 1]); | ||
| const size_t hidden_size_size = static_cast<size_t>(hidden_size); | ||
| ORT_RETURN_IF(hidden_size <= 0, "hidden_size must be positive."); | ||
|
|
||
| if (prepacked_skip_fp32_data_) { | ||
| ORT_RETURN_IF(prepacked_skip_fp32_size_ < hidden_size_size || (prepacked_skip_fp32_size_ % hidden_size_size) != 0, | ||
| "Prepacked skip length does not match hidden_size. hidden_size=", hidden_size, | ||
| ", prepacked skip length=", prepacked_skip_fp32_size_, "."); | ||
| } |
| kNnapiExecutionProvider, kQnnExecutionProvider}); | ||
| } | ||
|
|
||
| TEST(SkipLayerNormTest, SkipLayerNormPrePackRejectsShortGamma) { |
This pull request improves the robustness and correctness of the
SkipLayerNormoperator in ONNX Runtime by adding stricter validation of prepacked tensor sizes and enhancing test coverage for these checks. The changes ensure that mismatches in tensor dimensions are detected early, preventing potential runtime errors.Validation and Error Checking Enhancements:
skip,gamma,beta,bias) inSkipLayerNorm::Compute, ensuring their lengths match the expectedhidden_sizeand providing clear error messages when mismatches occur.gamma,beta, andbiastensors, and updated their management in the constructor andPrePackmethod. [1] [2] [3]Testing Improvements:
gammaandskip, confirming that the new validation logic works as intended.