Validate SkipLayerNorm prepacked input shapes

[Akshay Sonawane (apsonawane)]

This pull request makes SkipLayerNorm prepacking preserve the same input-shape validation semantics as the non-prepacked path.

Prepacked input validation:

• Records the full shape of each successfully prepacked MLFloat16 input instead of retaining only its element count.
• Reuses the existing skip, gamma, beta, and bias shape validators, including their INVALID_ARGUMENT status codes.
• Rejects prepacked skip tensors whose last two dimensions do not match the input, preventing incorrect modulo-based wrapping.
• Validates that the hidden dimension fits in int before narrowing.
• Tracks packed state independently from allocation pointers so zero-element prepacked skips remain valid no-op inputs.

Unit test coverage:

• Covers short prepacked gamma, beta, and bias inputs, including simplified-layer-normalization bias input 3.
• Covers short and shape-mismatched prepacked skip tensors.
• Covers a valid zero-element prepacked skip.

Packed shape metadata is populated only when conversion actually succeeds, keeping data and metadata state synchronized.

[Copilot]

🟡 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/bias tensor lengths and initialized them in the kernel constructor.
• Added explicit size checks in SkipLayerNorm::Compute to validate prepacked skip/gamma/beta/bias lengths against hidden_size, returning descriptive failures.
• Added two unit tests verifying failures for incorrect prepacked gamma and skip sizes.

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.

[Tianlei Wu (tianleiwu)]

Review: prepacked length validation for SkipLayerNorm

Verdict: comment / non-blocking suggestions. The change closes a real out-of-bounds read window and I'd be happy to see it merged with a couple of adjustments.

Why the fix is correct: when PrePack converts a constant MLFloat16 initializer, Compute() nulls out the corresponding p_ctx->Input<Tensor>(i), so CheckPotentiallyPrepackedInputs skips CheckSkip/CheckGamma for that input entirely. Because hidden_size comes from the runtime input shape, a model with a dynamic last dim could reach ComputeJob with a prepacked buffer shorter than hidden_size. The new checks restore that validation, and the skip condition (size >= hidden_size && size % hidden_size == 0) is exactly the invariant ComputeJob needs for skip_data + (offset % skip_size) to stay in bounds. Placing the hidden_size > 0 guard before the modulo is also the right ordering.

Main points (details inline):

1. The prepacked skip check restores memory safety but not shape semantics — it is strictly weaker than CheckSkip, so some shapes that are rejected on the non-prepacked path now silently wrap.
2. ORT_RETURN_IF returns common::FAIL, while every sibling shape check in skip_layer_norm_helper.h returns INVALID_ARGUMENT. The same user error now yields a different status code depending on whether the tensor happened to be prepacked.

Cross-cutting suggestion (not anchorable to the diff): CheckPotentiallyPrepackedInputs in onnxruntime/contrib_ops/cpu/skip_layer_norm_helper.h already takes prepacked_skip / prepacked_gamma booleans precisely so prepacked-input validation lives in one place. Threading the recorded lengths (or shapes) into that helper instead of validating in Compute() would keep all input validation together and let both paths share the same message text and error code.

Test coverage looks good — all four prepackable tensors plus the simplified bias-at-input-3 variant, and the tests correctly use MLFloat16 + is_initializer=true, which is the only combination that actually reaches the prepack path.

Nit: the PR description says "Added two new unit tests"; the branch has five. Worth refreshing before merge.

Performance/ABI: no hot-path allocation, no C API surface touched, no stream/sync concerns. The added cost is a few integer compares plus one modulo per Compute(), negligible next to the per-row normalization.

[Tianlei Wu (tianleiwu)]

Reviewed the current head after the shape-preservation follow-up. The prepacked path now uses the same skip/gamma/beta/bias shape predicates and INVALID_ARGUMENT behavior as runtime tensors, keeps prepack metadata synchronized with successful conversion, handles the empty-skip case, and validates before output allocation. I found no remaining issues.

Validation: rebuilt onnxruntime_provider_test from this head and ran the seven new SkipLayerNormTest prepack regression tests; all 7 passed.