diff --git a/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md b/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md new file mode 100644 index 0000000000000..d39d800fb887c --- /dev/null +++ b/.agents/skills/contrib-op-shape-inference-memory-safety/SKILL.md @@ -0,0 +1,237 @@ +--- +name: contrib-op-shape-inference-memory-safety +description: "Audit and fix out-of-range output writes in ONNX Runtime operator shape-inference functions. Use when reviewing or fixing a contrib (or standard) op TypeAndShapeInference where a getNumOutputs() guard precedes a write to a higher output index - optional trailing outputs make a smaller output count schema-valid, so getOutputType(index) can run one past the declared outputs at Graph::Resolve." +--- + +# Contrib-Op Shape-Inference Output-Index Safety + +Reusable method for finding and fixing the bug class where an operator's +`TypeAndShapeInference` function guards an output write with `getNumOutputs() > N` but then +writes an output index **greater than** `N`. For a node that declares fewer outputs, the +written index is past the end of the inference context's output vector. + +> **Scope**: schema-level shape inference in `onnxruntime/core/graph/contrib_ops/*.cc` and +> `shape_inference_functions.cc`. This runs once during `Graph::Resolve` (model-load time), +> **EP-agnostic** - there is no per-EP (CPU/CUDA/ROCm) kernel duplicate of this code to +> chase. Op *kernels* allocate outputs via the bounds-safe `OpKernelContext::Output(index)` +> and are a separate concern. + +## 1. The pattern + +```cpp +// onnxruntime/core/graph/contrib_ops/bert_defs.cc (before) +propagateElemTypeFromInputToOutput(ctx, 0, 0); +if (ctx.getNumOutputs() > 1) { // guard says "> 1" + propagateElemTypeFromInputToOutput(ctx, 0, 1); + propagateElemTypeFromInputToOutput(ctx, 0, 2); // but writes index 2 +} +``` + +The guard `getNumOutputs() > 1` admits a node with **exactly 2 outputs** (indices 0, 1), yet +the body writes index **2**. The implication "`> 1` ⇒ index 2 exists" is false: `> 1` only +guarantees indices 0 and 1. + +### Why a smaller output count is valid + +Trailing outputs declared `OpSchema::Optional` **lower `min_output`**. ONNX derives +`min_output` = number of required outputs, `max_output` = total declared. The model checker +(`checker::check_node`) only enforces `min_output <= N <= max_output`. + +| Op | Output decls | min / max | A 2-output node? | +|---|---|---|---| +| `DecoderAttention` | out (req), new_key_cache (Opt), new_value_cache (Opt) | 1 / 3 | passes checker | +| `MultiHeadAttention` | out (req), present_key (Opt), present_value (Opt), qk (Opt) | 1 / 4 | passes checker | +| `DecoderMaskedMultiHeadAttention` | out (req) + 3 Optional | 1 / 4 | passes checker | + +So a node with `output=['out','present_key']` is schema-valid, passes the checker, and then +reaches the index-2 write. **A passing checker is not a guarantee the index is in range.** + +## 2. The sink (why the write is not caught) + +```cpp +// onnxruntime/core/graph/graph.cc - InferenceContextImpl +const TypeProto* getInputType(size_t index) const override { + return node_.InputDefs().at(index)->TypeAsProto(); // .at() -> bounds-checked +} +TypeProto* getOutputType(size_t index) override { + return &node_output_types_[index]; // operator[] -> NOT bounds-checked +} +``` + +- `node_output_types_` is sized to `node.OutputDefs().size()` in the `InferenceContextImpl` + ctor, so for a 2-output node it has 2 elements; `getOutputType(2)` returns one past the end. +- `getInputType` uses `.at()` (would throw on a bad index); `getOutputType` uses raw + `operator[]` (no check) - the asymmetry is the root cause. +- The call runs at `Graph::Resolve` → `InferAndVerifyTypeMatch` → `RunInferencing`. The + surrounding `ORT_TRY/ORT_CATCH(const std::exception&)` only catches *thrown* + `fail_shape_inference`; a raw out-of-range `operator[]` does not throw, so the catch does + not help. +- Because this is schema-level inference, it is **EP-independent** - no CUDA/ROCm copy. + +## 3. Audit technique — always sweep siblings + +Do not stop at the reported function. Grep **every** shape-inference guard and compare its +threshold against the **highest output index written before the next guard**. + +```bash +git grep -n 'getNumOutputs' -- \ + onnxruntime/core/graph/contrib_ops/*.cc \ + onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc +``` + +For each `if (ctx.getNumOutputs() > N)` block, find the largest `index` passed to +`propagateElemTypeFromInputToOutput(ctx, _, index)` / `updateOutputShape(ctx, index, _)` / +`getOutputType(index)` inside it. **Rule: the guard must require strictly more outputs than +the highest index written** (write index `k` ⇒ guard must ensure `getNumOutputs() > k`). + +Correct exemplars already in the tree to copy: + +| Exemplar | Pattern | Why it is safe | +|---|---|---| +| `BaseGroupQueryAttention...` | `if (getNumOutputs() >= 3)` then writes idx 2 | guard covers highest index | +| `PagedAttention...` | nested `> 1` + inner `if (getNumOutputs() != 3) fail_shape_inference` | fails before any write | +| `EmbedLayerNormalizationShapeInference` | `> 2` then writes idx 2 | fixed by PR #28176 (precedent) | +| `SkipLayerNormalizationShapeInference` | each idx `k` guarded by `> k` | per-index guard | + +> **Gotcha — conditional writes can hide a vacuous audit.** A write may sit behind an inner +> condition (e.g. `hasInputShape(past_key_index)` before writing index 2). The site is still +> a bug, but you can only *observe* it when that inner condition is also satisfied. Keep this +> in mind both for the audit and for tests (§5). + +## 4. Fix patterns + +**Point fix (required): raise the guard to cover the highest index written.** + +```cpp +// before +if (ctx.getNumOutputs() > 1) { ... writes idx 2 ... } +// after +if (ctx.getNumOutputs() > 2) { // both present_key (idx 1) AND present_value (idx 2) + ... +} +``` + +Justify the threshold with the op's output semantics. For these attention ops the two trailing +outputs - `present_key` (idx 1) and `present_value` (idx 2) for `MultiHeadAttention`, +`new_key_cache` / `new_value_cache` for `DecoderAttention` (see the §1 table for each op's +exact output names) - are a **both-or-neither pair**: there is no valid configuration that +emits one without the other, so requiring all three outputs before populating indices 1 and 2 +is behavior-preserving. (`PagedAttention` encodes the same invariant via its nested `!= 3` +check.) + +**Defense-in-depth (recommended): bound the sink** so a future author cannot reintroduce the +class. + +```cpp +// onnxruntime/core/graph/graph.cc - InferenceContextImpl::getOutputType +TypeProto* getOutputType(size_t index) override { + if (index >= node_output_types_.size()) { + fail_type_inference("output index ", index, " is out of range; node has ", + node_output_types_.size(), " outputs"); + } + return &node_output_types_[index]; +} +``` + +This mirrors `getInputType`'s `.at()` and the existing bounds checks in the sibling +`DataPropagationContextImpl`. Placing it at the base layer transitively protects the NHWC and +quantization wrapper contexts. After the point fix this branch is unreachable through a normal +model (the guard already prevents the out-of-range index), so it is pure defense-in-depth. Its +failure mode is build-dependent: with exceptions enabled, `fail_type_inference` raises +`InferenceError` (a `std::exception`), caught by the existing `ORT_CATCH(const std::exception&)` +around `RunInferencing` and surfaced as a clean load-time error; under `ORT_NO_EXCEPTIONS` it is +**not** compiled out - ONNX's no-exceptions path prints the message to `std::cerr` and calls +`abort()`, a deterministic fail-fast (consistent with `getInputType`'s `.at()`, which likewise +terminates under no-exceptions). Either way the result is a controlled failure rather than an +out-of-range write. + +## 5. Test recipe + +Tests live in `onnxruntime/test/contrib_ops/*.cc` and are **auto-globbed** into the +`onnxruntime_provider_test` target by `cmake/onnxruntime_unittests.cmake` +(`test/contrib_ops/*.cc` pattern) - **no cmake edit needed** for a new file. See the +`ort-test` skill for the executable taxonomy (`onnxruntime_provider_test` vs +`onnxruntime_test_all`). + +Rules that make the regression test actually guard the fix: + +1. **Drive through `Model` + `Graph::Resolve`**, not ONNX's standalone `TestShapeInference`. + Only the full resolve path constructs the real `InferenceContextImpl` and hits the + `getOutputType` sink described in §2. A standalone ONNX shape-inference helper uses a + different context and **bypasses** the sink, so it cannot reproduce the bug. +2. **Negative tests must be NON-VACUOUS** - they must actually enter the write branch on + pre-fix source. If a write is gated by an inner condition (§3 gotcha), satisfy it: e.g. for + `MultiHeadAttention`/`DecoderMaskedMultiHeadAttention`, supply a **shaped `past_key`** + (and `past_sequence_length` / `past_present_share_buffer` as the op requires) so the + index-2 block runs. A negative test that only supplies `query` skips the block and passes + even on pre-fix source - regression-proof in name only. +3. **Add positive (all-outputs) cases**: a node with every output present must still infer the + trailing output types - proves the tightened guard did not over-restrict. +4. **Keep tests throw-free post-fix** so they are valid under `ORT_NO_EXCEPTIONS`. Any case + that is *expected* to `fail_shape_inference` (throws) must be excluded with + `#ifndef ORT_NO_EXCEPTIONS`. The "2 outputs must not go out of range" case is throw-free + after the point fix and is safe in all builds. + +**Verify the negative test is non-vacuous (sanitizer A/B)** - the most reliable way to prove a +negative test enters the previously-out-of-range branch: build the test at the **pre-fix** +commit with **AddressSanitizer** and confirm it flags the out-of-range output access; then +confirm it is clean after the fix. + +```bash +# Functional run (any Debug build): +cmake --build build/Linux/Debug --target onnxruntime_provider_test -j"$(nproc)" +./build/Linux/Debug/onnxruntime_provider_test \ + --gtest_filter='AttentionOptionalOutputsShapeInferenceTest.*' + +# A/B proof (isolated worktree at the pre-fix commit, CPU-only Debug + sanitizer): +git worktree add --detach ../ort-prefix-check ~1 +# copy the new test file in, then: +python3 tools/ci_build/build.py --build_dir build/asan --config Debug --parallel \ + --skip_tests --enable_address_sanitizer --skip_submodule_sync \ + --cmake_generator Ninja --target onnxruntime_provider_test +# Pre-fix: the negative tests fail (the sanitizer flags the out-of-range output access). +# Post-fix (cherry-pick the guard fix): all tests pass, no sanitizer report. +``` + +## 6. Process / wording conventions + +- Run **`lintrunner -a`** before pushing so the `CLANGFORMAT` / Python-format gate passes. See + the `ort-lint` skill. +- Use **correctness/robustness framing** in code, comments, commit messages, and the PR body + - describe the change as fixing an optional-output guard, not as a security fix. This + matches repo convention (compare `python-kwargs-setattr-security`) and keeps the PR neutral. + +## 7. Audit checklist (per-operator review) + +When reviewing or hardening any operator implementation or its shape inference: + +- [ ] Read the op's spec - ONNX standard op page, or for a contrib op its `OpSchema` + registration (`.Input/.Output/.Attr`, and `Optional`/`Variadic` markers). A local ONNX + checkout has the standard-op spec pages; contrib ops are defined only in ORT. +- [ ] Enumerate **all** inputs, attributes, and outputs, noting which are optional and the + resulting `min/max` input and output counts. +- [ ] Validate every input/attribute before indexing into it, to avoid out-of-range reads + (which can cascade into worse failures). Match each output-index write to a guard that + guarantees the index is in range (§3 rule). +- [ ] Prefer `ORT_RETURN_IF` / `ORT_RETURN_IF_NOT` for validation; use `ORT_ENFORCE` in + constructors. In shape inference use `fail_shape_inference` / `fail_type_inference`. +- [ ] Use `SafeInt<>` / `narrow<>()` for index and size arithmetic and casts to avoid overflow + or truncation that yields a wrong index. See `core/common/safeint.h` and + `docs/Coding_Conventions_and_Standards.md`. +- [ ] Ensure tests build and pass under **no-exceptions** builds; `#ifndef ORT_NO_EXCEPTIONS` + around any case expected to throw. +- [ ] Exclude EPs known not to support the op, with a comment explaining why. +- [ ] Check whether **other EPs (notably CUDA/ROCm)** implement the same op and whether the + same issue exists there. (For *shape inference* specifically, the logic is EP-agnostic + and single-source - confirm there is no kernel-side analogue.) + +## References + +- **PR #28176** - "Fix ... in EmbedLayerNormalizationShapeInference": the precedent that fixed + the identical `> 1` → `> 2` primitive in one site; the sibling attention sites were missed, + motivating the sweep in §3. +- **PR #29243** - this fix: guards corrected in `DecoderAttention` / `MultiHeadAttention` / + `DecoderMaskedMultiHeadAttention` shape inference, plus the `getOutputType` bounds check and + non-vacuous regression tests. +- Sibling skill: **`ort-test`** (test executables, `--gtest_filter`, contrib-op test layout); + **`ort-lint`** (`lintrunner -a`); **`ort-build`** (build flags, ASan). diff --git a/docs/Coding_Conventions_and_Standards.md b/docs/Coding_Conventions_and_Standards.md index 02af7ddaa49be..6dcf7e6dd43bd 100644 --- a/docs/Coding_Conventions_and_Standards.md +++ b/docs/Coding_Conventions_and_Standards.md @@ -109,6 +109,7 @@ void foo(gsl::span names) { * Use [SafeInt](https://github.com/dcleblanc/SafeInt) when calculating the size of memory to allocate to protect against overflow errors * `#include "core/common/safeint.h"` * search for `SafeInt` in the code for examples +* In operator shape inference, validate every output index against `getNumOutputs()` before writing it. Optional trailing outputs lower an op's `min_output`, so a node may legally declare fewer outputs than the schema's maximum; guard each optional output by the exact index it populates so a `getNumOutputs() > N` guard never writes an index greater than `N`. * The following C++ warnings should never be disabled in onnxruntime VC++ projects(Required by [Binskim](https://github.com/microsoft/binskim/blob/d9afb65c89a621411efded74c27999281d87867e/src/BinSkim.Rules/PERules/BA2007.EnableCriticalCompilerWarnings.cs)). 1. [4018](https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-3-c4018) 'token' : signed/unsigned mismatch 2. [4146](https://docs.microsoft.com/en-us/cpp/error-messages/compiler-warnings/compiler-warning-level-2-c4146?view=msvc-160) unary minus operator applied to unsigned type, result still unsigned diff --git a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h index b0cf1fc976de0..64b5439b87c7b 100644 --- a/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h @@ -24,8 +24,13 @@ inline Status CheckInputs(const TOpKernelContext* context, bool quantizedVersion const Tensor* beta = context->template Input(6); const Tensor* mask = context->template Input(7); // optional. nullptr if not provided + // Track whether explicit position ids are supplied. When they are not, positions default to + // [0, sequence_length) and must be covered by the position_embedding rows (checked below). + bool has_position_ids = false; + if (!quantizedVersion) { const Tensor* position_ids = context->template Input(8); // optional. nullptr if not provided + has_position_ids = (nullptr != position_ids); if (nullptr != position_ids) { if (input_ids->Shape()[1] != position_ids->Shape()[1]) { @@ -69,6 +74,17 @@ inline Status CheckInputs(const TOpKernelContext* context, bool quantizedVersion "position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size()); } + // When position_ids are not supplied, positions run over [0, sequence_length); the + // position_embedding table must have at least that many rows to be addressable. + if (!has_position_ids) { + int64_t sequence_length = input_ids->Shape()[1]; + if (position_embedding_dims[0] < sequence_length) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "position_embedding must have at least sequence_length (", sequence_length, + ") rows when position_ids is not provided, got ", position_embedding_dims[0]); + } + } + if (nullptr != segment_embedding) { const auto& segment_embedding_dims = segment_embedding->Shape().GetDims(); if (segment_embedding_dims.size() != 2) { diff --git a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h index 7dfb2770a6979..4c7f2268310ff 100644 --- a/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h +++ b/onnxruntime/contrib_ops/cpu/maxpool_with_mask.h @@ -216,6 +216,16 @@ class MaxpoolWithMask : public OpKernel, public PoolBase { "Mask and input spatial dimensions mismatch at dimension ", i, ": mask=", m_shape[i], " input=", x_shape[i]); } + // x_shape.NumDimensions() >= 3 is guaranteed above, so this subtraction cannot underflow. + const size_t input_spatial_rank = x_shape.NumDimensions() - 2; + // The pooling kernel rank drives the 1D/2D/3D dispatch below, which reads x_shape[2..4] and + // output_dims[2..4]. Require it to match the input spatial rank so those reads stay in bounds. + ORT_RETURN_IF_NOT(pool_attrs_.kernel_shape.size() == input_spatial_rank, + "Pooling kernel rank must equal input spatial rank. Got kernel rank: ", + pool_attrs_.kernel_shape.size(), " input spatial rank: ", input_spatial_rank); + // Only 1D/2D/3D pooling is implemented by the dispatch below; a larger rank would match no case. + ORT_RETURN_IF_NOT(input_spatial_rank >= 1 && input_spatial_rank <= 3, + "Only 1D, 2D, and 3D pooling are supported. Got input spatial rank: ", input_spatial_rank); TensorShapeVector pads = pool_attrs_.pads; TensorShapeVector kernel_shape = pool_attrs_.kernel_shape; diff --git a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc index 2094af78f40b7..ba3c8a2ada30b 100644 --- a/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc +++ b/onnxruntime/contrib_ops/cpu/quantization/dynamic_quantize_lstm.cc @@ -178,14 +178,14 @@ Status DynamicQuantizeLSTM::Compute(OpKernelContext* context) const { const Tensor* r_zp = context->Input(11); const TensorShape& W_zp_shape = w_zp->Shape(); - const TensorShape& R_zp_shape = w_zp->Shape(); + const TensorShape& R_zp_shape = r_zp->Shape(); const TensorShape& W_scale_shape = w_scale->Shape(); const TensorShape& R_scale_shape = r_scale->Shape(); WeightCheck(W_zp_shape, W_zero_point); WeightCheck(R_zp_shape, R_zero_point); WeightCheck(W_scale_shape, W_scale); - WeightCheck(W_scale_shape, R_scale); + WeightCheck(R_scale_shape, R_scale); const bool is_W_signed = (W != nullptr) ? W->IsDataType() : is_W_signed_; const bool is_R_signed = (R != nullptr) ? R->IsDataType() : is_R_signed_; diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc index 864e2d1623923..81ba48f800dee 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc @@ -61,9 +61,18 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { int sequence_length = static_cast(input_dims[1]); size_t element_size = sizeof(T); + int word_embedding_length = static_cast(word_embedding->Shape()[0]); + int position_embedding_length = static_cast(position_embedding->Shape()[0]); + int segment_embedding_length = + (nullptr == segment_embedding) ? 0 : static_cast(segment_embedding->Shape()[0]); + const bool broadcast_position_ids = (nullptr != position_ids && position_ids->Shape()[0] == 1); - return LaunchEmbedLayerNormKernel( + // Device flag raised by the kernel when an input id is outside its embedding table. + auto error_flag = GetScratchBuffer(1, context->GetComputeStream()); + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(error_flag.get(), 0, sizeof(int), Stream(context))); + + ORT_RETURN_IF_ERROR(LaunchEmbedLayerNormKernel( Stream(context), output->MutableData(), nullptr == mask_index ? nullptr : mask_index->MutableData(), @@ -82,7 +91,30 @@ Status EmbedLayerNorm::ComputeInternal(OpKernelContext* context) const { element_size, embedding_sum == nullptr ? nullptr : embedding_sum->MutableData(), position_ids == nullptr ? nullptr : position_ids->Data(), - broadcast_position_ids); + broadcast_position_ids, + word_embedding_length, + position_embedding_length, + segment_embedding_length, + error_flag.get())); + + // The kernel always validates ids device-side and skips the embedding reads for any + // out-of-range id, so input safety does not depend on the readback below; the readback only + // upgrades a skipped (silently zeroed) row into a clean error status. cudaStreamSynchronize and + // device-to-host copies are not permitted on a stream that is capturing a CUDA graph, so skip the + // status readback while the stream is capturing. Ids are still kept in range device-side, and + // error surfacing resumes on normal (non-capturing) runs. + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(Stream(context), &capture_status)); + if (capture_status == cudaStreamCaptureStatusNone) { + auto host_error_flag = AllocateBufferOnCPUPinned(1); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_error_flag.get(), error_flag.get(), sizeof(int), + cudaMemcpyDeviceToHost, Stream(context))); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(Stream(context))); + ORT_RETURN_IF(*host_error_flag.get() != 0, + "input id is out of range of the corresponding embedding table."); + } + + return Status::OK(); } } // namespace cuda diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu index a6b97286d1733..53c98e74a2842 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu @@ -123,7 +123,8 @@ template __global__ void EmbedLayerNormKernel( int hidden_size, const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, - float epsilon, T* output, T* embedding_sum, const int* position_ids, const bool broadcast_position_ids) { + float epsilon, T* output, T* embedding_sum, const int* position_ids, const bool broadcast_position_ids, + int word_embedding_length, int position_embedding_length, int segment_embedding_length, int* error_flag) { KeyValuePairSum pair_sum; // 1. lookup word and segment of the block // blockIdx.x = position in the sequence @@ -133,6 +134,7 @@ __global__ void EmbedLayerNormKernel( __shared__ int word_id; __shared__ int segment_id; __shared__ int position_id; + __shared__ bool is_valid_block; const float rld = 1.f / hidden_size; const int sequence_position = blockIdx.y * gridDim.x + blockIdx.x; @@ -150,17 +152,30 @@ __global__ void EmbedLayerNormKernel( } else { position_id = position_ids[sequence_position]; } + + // Flag any id that is out of range of its embedding table so it is rejected instead of indexed. + is_valid_block = (word_id >= 0 && word_id < word_embedding_length) && + (position_id >= 0 && position_id < position_embedding_length) && + (nullptr == segment_embedding || + (segment_id >= 0 && segment_id < segment_embedding_length)); + if (!is_valid_block) { + *error_flag = 1; + } } __syncthreads(); + if (!is_valid_block) { + return; + } + // 2. load pos/segment/word embeddings and add them together // offset into embeddings is given by word_id * hidden_size - const int position_offset = position_id * hidden_size; + const int64_t position_offset = static_cast(position_id) * hidden_size; - const int word_offset = word_id * hidden_size; - const int segment_offset = segment_id * hidden_size; + const int64_t word_offset = static_cast(word_id) * hidden_size; + const int64_t segment_offset = static_cast(segment_id) * hidden_size; // the output offset is given by b * (sequence_length * hidden_size) + s * hidden_size - const int output_offset = sequence_position * hidden_size; + const int64_t output_offset = static_cast(sequence_position) * hidden_size; cub::KeyValuePair thread_data(0.f, 0.f); @@ -182,7 +197,7 @@ __global__ void EmbedLayerNormKernel( thread_data = pair_sum(thread_data, cub::KeyValuePair(rldval, rldval * val_f)); } - // 3. layer norm on the sum + // 3. layer norm on the sum (64-bit output offset to support large tensors) LayerNorm(thread_data, hidden_size, output_offset, beta, gamma, epsilon, output); } @@ -192,7 +207,8 @@ Status EmbedSkipLayerNorm( const int* input_ids, const int* segment_ids, const T* beta, const T* gamma, const T* word_embedding, const T* position_embedding, const T* segment_embedding, float epsilon, T* output, T* embedding_sum, const int* position_ids, - const bool broadcast_position_ids) { + const bool broadcast_position_ids, int word_embedding_length, int position_embedding_length, + int segment_embedding_length, int* error_flag) { constexpr int tpb = 256; const dim3 grid(sequence_length, batch_size, 1); const dim3 block(tpb, 1, 1); @@ -200,7 +216,9 @@ Status EmbedSkipLayerNorm( EmbedLayerNormKernel <<>>(hidden_size, input_ids, segment_ids, beta, gamma, word_embedding, position_embedding, segment_embedding, - epsilon, output, embedding_sum, position_ids, broadcast_position_ids); + epsilon, output, embedding_sum, position_ids, broadcast_position_ids, + word_embedding_length, position_embedding_length, + segment_embedding_length, error_flag); return CUDA_CALL(cudaGetLastError()); } @@ -224,7 +242,11 @@ Status LaunchEmbedLayerNormKernel( const size_t element_size, void* embedding_sum, const int* position_ids, - const bool broadcast_position_ids) { + const bool broadcast_position_ids, + int word_embedding_length, + int position_embedding_length, + int segment_embedding_length, + int* error_flag) { if (mask_index != nullptr) { if (nullptr == input_mask) { CUDA_RETURN_IF_ERROR(cudaMemsetAsync(mask_index, 0, sizeof(int) * batch_size, stream)); @@ -241,7 +263,8 @@ Status LaunchEmbedLayerNormKernel( reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), reinterpret_cast(segment_embedding), epsilon, reinterpret_cast(output), reinterpret_cast(embedding_sum), position_ids, - broadcast_position_ids); + broadcast_position_ids, word_embedding_length, position_embedding_length, + segment_embedding_length, error_flag); } else { return EmbedSkipLayerNorm( stream, hidden_size, batch_size, sequence_length, input_ids, segment_ids, @@ -249,7 +272,8 @@ Status LaunchEmbedLayerNormKernel( reinterpret_cast(word_embedding), reinterpret_cast(position_embedding), reinterpret_cast(segment_embedding), epsilon, reinterpret_cast(output), reinterpret_cast(embedding_sum), position_ids, - broadcast_position_ids); + broadcast_position_ids, word_embedding_length, position_embedding_length, + segment_embedding_length, error_flag); } } diff --git a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h index aaf1d891dab3a..bac7702161792 100644 --- a/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.h @@ -26,7 +26,11 @@ Status LaunchEmbedLayerNormKernel(cudaStream_t stream, const size_t element_size, // size of output element: 2 for half, 4 for float. void* embedding_sum, // Optional output of sum of embeddings const int* position_ids, // Optional input of position ids - const bool broadcast_position_ids); // Whether to broadcast position ids + const bool broadcast_position_ids, // Whether to broadcast position ids + int word_embedding_length, // number of rows in word_embedding + int position_embedding_length, // number of rows in position_embedding + int segment_embedding_length, // rows in segment_embedding (0 if none) + int* error_flag); // device flag set when an id is out of range } // namespace cuda } // namespace contrib diff --git a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh index 0ed2c125405e3..a7bab0b10def4 100644 --- a/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh @@ -78,7 +78,7 @@ struct KeyValuePairSum { template __device__ inline void LayerNorm( - const cub::KeyValuePair& thread_data, const int ld, const int offset, const T* beta, + const cub::KeyValuePair& thread_data, const int ld, const int64_t offset, const T* beta, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation for mean/variance to avoid overflow in fp16/bf16. @@ -98,7 +98,9 @@ __device__ inline void LayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + // 64-bit global element offset: rows*ld (= batch*seq*hidden) can exceed 2^31 for large tensors. + // gamma/beta indices (i) stay int -- bounded by ld (hidden_size). + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); const float b = (nullptr == beta) ? 0.f : static_cast(beta[i]); @@ -108,7 +110,7 @@ __device__ inline void LayerNorm( template __device__ inline void SimplifiedLayerNorm( - const float& thread_data, const int ld, const int offset, const T* gamma, const float epsilon, T* output) { + const float& thread_data, const int ld, const int64_t offset, const T* gamma, const float epsilon, T* output) { // Assuming thread_data is already divided by ld // Uses fp32 accumulation to avoid overflow in fp16/bf16. @@ -124,7 +126,8 @@ __device__ inline void SimplifiedLayerNorm( __syncthreads(); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + // 64-bit element offset (see LayerNorm); large tensors can exceed 2^31. gamma index (i) stays int. + const int64_t idx = offset + i; const float val = static_cast(output[idx]); const float g = static_cast(gamma[i]); output[idx] = static_cast(g * val * rsigma); @@ -133,8 +136,9 @@ __device__ inline void SimplifiedLayerNorm( template __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair& thread_data, - const int ld, const int idx, const T* beta, const T* gamma, + const int ld, const int64_t idx, const T* beta, const T* gamma, const float epsilon, T* output) { + // idx is a 64-bit global element offset (caller passes row*ld = batch*seq*hidden, can exceed 2^31). // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register @@ -182,8 +186,9 @@ __device__ inline void LayerNormSmall(const T* input_v, const cub::KeyValuePair< } template -__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, const int idx, - const T* gamma, const float epsilon, T* output) { +__device__ inline void SimplifiedLayerNormSmall(const T* input_v, const float& thread_data, const int ld, + const int64_t idx, const T* gamma, const float epsilon, T* output) { + // idx is a 64-bit global element offset (caller passes row*ld = batch*seq*hidden, can exceed 2^31). // Assuming thread_data is already divided by ld // Small settings: the block covers the leading dimension TPB >= ld. The input // value is available in a register diff --git a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu index 7f5169639ff8d..e1bab6b38da9b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu @@ -75,7 +75,8 @@ __global__ void SkipLayerNormKernel( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, const int ld, int skip_size) { const float reverse_ld = 1.f / ld; - const int offset = blockIdx.x * ld; + // 64-bit global element offset: blockIdx.x*ld (= row*hidden) can exceed 2^31 for large tensors. + const int64_t offset = static_cast(blockIdx.x) * ld; const bool has_bias = (bias != nullptr); // Reduce sum of x and x^2, and the results are divided by ld. @@ -84,7 +85,7 @@ __global__ void SkipLayerNormKernel( cub::KeyValuePair thread_data(0.f, 0.f); for (int i = threadIdx.x; i < ld; i += TPB) { - const int idx = offset + i; + const int64_t idx = offset + i; T val = input[idx]; if (has_bias) { @@ -116,7 +117,8 @@ __global__ void SkipLayerNormKernelSmall( T* output, T* sum_output, const T* input, const T* skip, const T* bias, const T* gamma, const T* beta, float epsilon, int ld, int skip_size) { const float rld = 1.f / ld; - const int idx = blockIdx.x * ld + threadIdx.x * ILP; + // 64-bit element offset (see SkipLayerNormKernel); large tensors can exceed 2^31. + const int64_t idx = static_cast(blockIdx.x) * ld + threadIdx.x * ILP; using VecT = aligned_vector; T sum_v[ILP]; diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc index 9b91215eba91d..65759c8ca0f27 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cc @@ -47,6 +47,15 @@ template GatherBlockQuantized::GatherBlockQuantized(const OpKernelInfo& info) : CudaKernel(info) { ORT_ENFORCE(info.GetAttr("bits", &bits_).IsOK()); + // bits_ is used as the divisor in 8 / bits_ to derive the uint8 packing factor, so it must + // be one of the supported widths. uint8 data packs 4 or 8 bit values; int4/uint4 are 4 bit. + if constexpr (std::is_same_v) { + ORT_ENFORCE(bits_ == 4 || bits_ == 8, + "GatherBlockQuantized with uint8 data only supports bits == 4 or 8."); + } else { + ORT_ENFORCE(bits_ == 4, "GatherBlockQuantized with int4/uint4 data only supports bits == 4."); + } + block_size_ = info.GetAttrOrDefault("block_size", 0); gather_axis_ = info.GetAttrOrDefault("gather_axis", 0); quantize_axis_ = info.GetAttrOrDefault("quantize_axis", 0); @@ -71,6 +80,46 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) ORT_ENFORCE(quantize_axis_ == static_cast(data_rank) - 1); + // block_size is used as a divisor when mapping an element to its quantization block, + // so it must be strictly positive (the constructor permits the unset default of 0). + ORT_RETURN_IF_NOT(block_size_ > 0, "'block_size' must be greater than 0."); + + // When sub-byte data is packed into uint8_t, the last (quantize) dimension stores + // multiple values per stored element. components is the packing factor. + uint32_t components = 1; + if constexpr (std::is_same_v) { + components = 8 / static_cast(bits_); + } + + // Validate scales shape against data shape, matching the CPU kernel: every dimension + // must be equal except the quantize axis, which is divided into block_size blocks. + auto scales_shape = scales->Shape().GetDims(); + ORT_RETURN_IF_NOT(data_rank == static_cast(scales->Shape().NumDimensions()), + "data and scales must have the same rank."); + for (size_t i = 0; i < data_shape.size(); ++i) { + ORT_RETURN_IF_NOT(static_cast(i) == quantize_axis_ + ? (data_shape[i] * components + block_size_ - 1) / block_size_ == scales_shape[i] + : data_shape[i] == scales_shape[i], + "data and scales do not match shapes."); + } + + // Validate zero_points shape against scales shape when provided. + if (zero_points != nullptr) { + auto zero_points_shape = zero_points->Shape().GetDims(); + ORT_RETURN_IF_NOT(scales->Shape().NumDimensions() == zero_points->Shape().NumDimensions(), + "scales and zero_points must have the same rank."); + for (size_t i = 0; i < scales_shape.size(); ++i) { + if (components > 1 && static_cast(i) == quantize_axis_) { + // For uint8_t with bits < 8, zero_points are packed components per stored element. + ORT_RETURN_IF_NOT((scales_shape[i] + components - 1) / components == zero_points_shape[i], + "scales and zero_points shape does not match."); + } else { + ORT_RETURN_IF_NOT(scales_shape[i] == zero_points_shape[i], + "scales and zero_points must have the same shape."); + } + } + } + TensorShapeVector output_shape; output_shape.reserve(data_rank - 1 + indices_rank); @@ -98,11 +147,8 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) } // Special int4‐in‐uint8 packing tweak: expand the last dim by components - if constexpr (std::is_same_v) { - uint32_t components = 8 / static_cast(bits_); - if (components > 1) { - output_shape.back() *= components; - } + if (components > 1) { + output_shape.back() *= components; } Tensor* output = ctx->Output(0, TensorShape(output_shape)); @@ -123,11 +169,8 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) // after_gather_dim has to be adjusted to match // the unpacked output dims for correct kernel indexing int64_t after_gather_dim_unpacked = after_gather_dim; - if constexpr (std::is_same_v) { - uint32_t components = 8 / static_cast(bits_); - if (components > 1) { - after_gather_dim_unpacked *= components; - } + if (components > 1) { + after_gather_dim_unpacked *= components; } GatherBlockQuantizedParam param; @@ -140,6 +183,12 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) param.gather_axis = gather_axis_; param.N = N; + // Device flag the kernel sets when a gathered index is outside the valid range. + // It is initialized to 0 and read back after the launch to surface a clean status. + auto index_out_of_bounds = GetScratchBuffer(1, ctx->GetComputeStream()); + CUDA_RETURN_IF_ERROR(cudaMemsetAsync(index_out_of_bounds.get(), 0, sizeof(int), param.stream)); + param.index_out_of_bounds = index_out_of_bounds.get(); + const auto dequantized_type = scales->GetElementType(); if (dequantized_type == ONNX_NAMESPACE::TensorProto_DataType_FLOAT) { const auto* scales_ptr = static_cast(scales->DataRaw()); @@ -155,6 +204,23 @@ Status GatherBlockQuantized::ComputeInternal(OpKernelContext* ctx) LaunchGatherBlockQuantizedKernel(data_ptr, indices_ptr, scales_ptr, zero_points_ptr, output_ptr, param); } + auto host_index_out_of_bounds = AllocateBufferOnCPUPinned(1); + // The device-side bounds check and kernel early-return above always run, so memory + // safety does not depend on the host readback below. The readback only upgrades a + // silently incorrect result into a clean error, and it requires a device-to-host copy + // plus a stream synchronize. Both are illegal while the stream is being captured for a + // CUDA graph, so skip the readback during capture and let the device-side guard stand. + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(param.stream, &capture_status)); + if (capture_status == cudaStreamCaptureStatusNone) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(host_index_out_of_bounds.get(), index_out_of_bounds.get(), sizeof(int), + cudaMemcpyDeviceToHost, param.stream)); + CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(param.stream)); + ORT_RETURN_IF(*host_index_out_of_bounds.get() != 0, + "indices element out of data bounds. Each index must be within the inclusive range [", + -param.gather_axis_dim, ", ", param.gather_axis_dim - 1, "]."); + } + return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu index 39286c63e9a08..8539b4e5d2869 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu @@ -44,7 +44,8 @@ __global__ void GatherBlockQuantizedKernel( int64_t block_size, int64_t gather_axis, int64_t N, - bool sign) { + bool sign, + int* index_out_of_bounds) { int64_t out_idx = blockDim.x * blockIdx.x + threadIdx.x; if (out_idx >= N) return; @@ -53,6 +54,18 @@ __global__ void GatherBlockQuantizedKernel( int64_t idx_after = out_idx % after_gather_dim; int64_t idx = (out_idx % (after_gather_dim * ind_dim)) / after_gather_dim; int64_t idx_at_g = indices[idx]; + + // Validate the gathered index and normalize negatives, mirroring the CPU kernel. + // A negative index counts from the end of the gather axis; anything outside + // [-gather_axis_dim, gather_axis_dim) is reported via the device error flag. + if (idx_at_g < -gather_axis_dim || idx_at_g >= gather_axis_dim) { + *index_out_of_bounds = 1; + return; + } + if (idx_at_g < 0) { + idx_at_g += gather_axis_dim; + } + int64_t in_idx = idx_before * gather_axis_dim * after_gather_dim + idx_at_g * after_gather_dim + idx_after; int64_t block_id = in_idx / block_size; @@ -81,8 +94,10 @@ void LaunchGatherBlockQuantizedKernel(const T1* data, int blocksPerGrid = (int)(ceil(static_cast(param.N) / GridDim::maxThreadsPerBlock)); bool sign = std::is_same::value; - GatherBlockQuantizedKernel<<>>(data, indices, scales, zero_points, output, - param.after_gather_dim, param.gather_axis_dim, param.ind_dim, param.bits, param.block_size, param.gather_axis, param.N, sign); + GatherBlockQuantizedKernel<<>>( + data, indices, scales, zero_points, output, + param.after_gather_dim, param.gather_axis_dim, param.ind_dim, param.bits, param.block_size, + param.gather_axis, param.N, sign, param.index_out_of_bounds); } template void LaunchGatherBlockQuantizedKernel(const uint8_t*, const int32_t*, const float*, const uint8_t*, float*, GatherBlockQuantizedParam); diff --git a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh index f5dea3b1f2d9d..36266121ac210 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cuh @@ -19,6 +19,10 @@ struct GatherBlockQuantizedParam { int64_t block_size; int64_t gather_axis; int64_t N; + // Device buffer of a single int. The kernel sets it to 1 if any gathered index + // is outside the valid range [-gather_axis_dim, gather_axis_dim). The host reads + // it back after the launch and returns an error status when it is set. + int* index_out_of_bounds; }; template diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 896774fb5c8d8..9fdcebf57081d 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -29,7 +29,7 @@ namespace contrib { void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx) { // Type inference ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 0); - if (ctx.getNumOutputs() > 1) { + if (ctx.getNumOutputs() > 2) { // has new_key_cache and new_value_cache outputs; a pair, so present only when > 2 ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 1); ONNX_NAMESPACE::propagateElemTypeFromInputToOutput(ctx, 0, 2); } @@ -38,7 +38,7 @@ void DecoderAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx auto& query_shape = getInputShape(ctx, 0); updateOutputShape(ctx, 0, query_shape); } - if (ctx.getNumOutputs() > 1) { + if (ctx.getNumOutputs() > 2) { // has new_key_cache and new_value_cache outputs; a pair, so present only when > 2 if (hasInputShape(ctx, 6) && hasInputShape(ctx, 7)) { auto& cache_shape = getInputShape(ctx, 6); auto& cache_dims = cache_shape.dim(); @@ -199,7 +199,7 @@ void MultiHeadAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& c } } - if (ctx.getNumOutputs() > 1) { // has present output + if (ctx.getNumOutputs() > 2) { // has present_key and present_value outputs; a pair, so present only when > 2 if (hasInputShape(ctx, past_key_index)) { auto& past_shape = getInputShape(ctx, past_key_index); auto& past_dims = past_shape.dim(); diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index f3f2f521ecab2..b2a6b70dbb455 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -4072,8 +4072,10 @@ GatherBlockQuantized is a Gather with data quantized. It is similar to Gather (h if (quantize_axis < -r || quantize_axis >= r) { fail_shape_inference("quantize_axis must be in [-r, r-1]"); } - if (block_size < 0) { - fail_shape_inference("block_size must be non-negative"); + if (block_size <= 0) { + // block_size is used as a divisor below when relating data and scales shapes, + // so it must be strictly positive. + fail_shape_inference("block_size must be greater than 0"); } gather_axis = (gather_axis + r) % r; diff --git a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc index c2d853af86723..2b2a565649175 100644 --- a/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/range_schema_defs.cc @@ -7,6 +7,7 @@ #include "core/graph/constants.h" #include "core/graph/op.h" #include +#include #include namespace onnxruntime { @@ -60,7 +61,13 @@ static T GetFirstElement(const TensorProto* shapeInitializer) { if (utils::HasRawData(*shapeInitializer)) { const std::string& bytes = shapeInitializer->raw_data(); - return *reinterpret_cast(bytes.c_str()); + if (bytes.size() < sizeof(T)) { + fail_shape_inference("Range: raw_data size is smaller than the element size for the given data type."); + } + // std::string data is only char-aligned, so copy the bytes into a properly aligned value. + T value; + std::memcpy(&value, bytes.data(), sizeof(T)); + return value; } return get_data(shapeInitializer); } @@ -75,7 +82,23 @@ static int64_t CalcRangeDim(const TensorProto* startShapeInitializer, if (delta == 0) { fail_shape_inference("delta in Range operator can not be zero!"); } - return static_cast(ceil((1.0 * (limit - start)) / delta)); + // Mirror the CPU kernel (core/providers/cpu/generator/range.cc) which clamps empty or + // backward ranges to 0 so shape inference does not emit a negative dimension value. + double count = ceil((1.0 * (limit - start)) / delta); + if (!std::isfinite(count)) { + fail_shape_inference("Range: the computed number of elements is not a finite value."); + } + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion below. + if (count <= 0) { + return 0; + } + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (count >= 9223372036854775808.0) { + fail_shape_inference("Range: the computed number of elements exceeds the supported range."); + } + return static_cast(count); } static int64_t CalcResultDim(const TensorProto* startShapeInitializer, diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index fe2df6a87d124..ff60f6cb5f741 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -2739,6 +2739,10 @@ class InferenceContextImpl : public ONNX_NAMESPACE::InferenceContext { } TypeProto* getOutputType(size_t index) override { + if (index >= node_output_types_.size()) { + fail_type_inference("output index ", index, " is out of range; node has ", + node_output_types_.size(), " outputs"); + } return &node_output_types_[index]; } diff --git a/onnxruntime/core/providers/cpu/generator/range.cc b/onnxruntime/core/providers/cpu/generator/range.cc index fcd4abcf72f87..752ac2b5405ab 100644 --- a/onnxruntime/core/providers/cpu/generator/range.cc +++ b/onnxruntime/core/providers/cpu/generator/range.cc @@ -80,9 +80,23 @@ static Status ComputeRange( if (delta == T{0}) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "delta in Range operator can not be zero!"); } - int64_t n = static_cast(ceil((1.0 * (limit - start)) / delta)); - if (n <= 0) - n = 0; + double count = ceil((1.0 * (limit - start)) / delta); + if (!std::isfinite(count)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements is not a finite value."); + } + // Empty or backward ranges clamp to 0; handle the non-positive case before the cast so a + // large-magnitude negative count can never reach (and overflow) the int64 conversion. + int64_t n = 0; + if (count > 0) { + // static_cast(INT64_MAX) rounds up to 2^63 (9223372036854775808.0), which is not + // representable as int64_t, so reject any count at or above that boundary before the cast. + if (count >= 9223372036854775808.0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "Range: the computed number of elements exceeds the supported range."); + } + n = static_cast(count); + } TensorShape shape = {n}; T* y = ctx->Output(0, shape)->MutableData(); for (int64_t i = 0; i < n; ++i) { diff --git a/onnxruntime/core/providers/cpu/tensor/grid_sample.cc b/onnxruntime/core/providers/cpu/tensor/grid_sample.cc index 34135ca62e4a8..0299aca0057a3 100644 --- a/onnxruntime/core/providers/cpu/tensor/grid_sample.cc +++ b/onnxruntime/core/providers/cpu/tensor/grid_sample.cc @@ -373,6 +373,14 @@ Status GridSample::Compute(OpKernelContext* context) const { ORT_ENFORCE(input_dims.NumDimensions() == 4 || input_dims.NumDimensions() == 5, "Only 4-D or 5-D tensor is supported"); + // Spatial dimensions must be non-empty for sampling: the output is sized by the grid, so a zero-size + // input spatial dimension would otherwise lead to invalid index computations during interpolation. + for (size_t i = 2; i < input_dims.NumDimensions(); ++i) { + ORT_RETURN_IF_NOT(input_dims[i] > 0, + "Input spatial dimensions must be non-empty for sampling. Dimension ", i, + " has size ", input_dims[i]); + } + auto N = input_dims[0]; auto C = input_dims[1]; ORT_ENFORCE(grid_dims[0] == N, "Grid batch size ", grid_dims[0], " does not match input batch size ", N); diff --git a/onnxruntime/core/providers/cuda/tensor/grid_sample.cc b/onnxruntime/core/providers/cuda/tensor/grid_sample.cc index d97d5fcbb0b5b..e23b5bad573f7 100755 --- a/onnxruntime/core/providers/cuda/tensor/grid_sample.cc +++ b/onnxruntime/core/providers/cuda/tensor/grid_sample.cc @@ -124,6 +124,17 @@ Status GridSample::ComputeInternal(OpKernelContext* context) const { return Status(common::ONNXRUNTIME, common::NOT_IMPLEMENTED, "Cubic mode is only supported in 4-D cases."); } + // Spatial dimensions must be non-empty for sampling: the output is sized by the grid, so a zero-size + // input spatial dimension would otherwise lead to invalid index computations during interpolation. + // Spatial dims are [1, rank-1) for NHWC (N,...,C) and [2, rank) for NCHW (N,C,...). + const size_t spatial_begin = IsNHWC ? 1 : 2; + const size_t spatial_end = IsNHWC ? dims_input.size() - 1 : dims_input.size(); + for (size_t i = spatial_begin; i < spatial_end; ++i) { + ORT_RETURN_IF_NOT(dims_input[i] > 0, + "Input spatial dimensions must be non-empty for sampling. Dimension ", i, + " has size ", dims_input[i]); + } + using Ch = Channels; TensorShapeVector dims_output(dims_input.size()); diff --git a/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc new file mode 100644 index 0000000000000..d2d3a421a6560 --- /dev/null +++ b/onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc @@ -0,0 +1,230 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Shape-inference correctness coverage for attention contrib ops with optional "present" outputs. +// DecoderAttention, MultiHeadAttention and DecoderMaskedMultiHeadAttention each expose an optional +// present_key (output 1) and present_value (output 2). These outputs are produced as a pair, so a +// node may declare either one output, or all three; declaring exactly two (present_key kept, +// present_value omitted) is also valid per the schemas. +// +// The "...Omitted" tests build each op with exactly two outputs and assert that Graph::Resolve() +// completes cleanly without referencing the absent third output. The "...AllPresentOutputs" tests +// build each op with all three outputs and assert that the present_key / present_value branch still +// runs and propagates their element types. Together they pin the guard to exactly "> 2": fewer +// outputs must not touch the missing one, and three outputs must still be inferred. +// +// These tests exercise only graph-load shape inference, which is execution-provider independent, so +// they run on the default CPU build with no provider-specific handling. The resolve path for these +// models is throw-free, so the tests are valid in builds compiled without exceptions +// (ORT_NO_EXCEPTIONS) and need no exception-specific guarding. + +#include "gtest/gtest.h" + +#include "core/graph/constants.h" +#include "core/graph/model.h" +#include "test/test_environment.h" +#include "test/unittest_util/graph_transform_test_builder.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { + +namespace { + +constexpr int kOnnxOpsetVersion = 17; + +// Builds a single-node model via add_node, resolves it (running the node's type/shape inference), +// asserts success, and runs the optional verifier against the resolved graph. +void BuildResolveAndVerify(const std::function& add_node, + const std::function& verify = nullptr) { + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = kOnnxOpsetVersion; + domain_to_version[kMSDomain] = 1; + + Model model("attention_optional_outputs", /*is_onnx_domain_only=*/false, ModelMetaData(), + PathString(), IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, + DefaultLoggingManager().DefaultLogger()); + + ModelTestBuilder builder(model.MainGraph()); + add_node(builder); + builder.SetGraphOutputs(); + + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + if (verify) { + verify(model.MainGraph()); + } +} + +// Asserts that the given output NodeArg received a tensor element type from shape inference. +void ExpectInferredElemType(const NodeArg* output) { + ASSERT_NE(output, nullptr); + const ONNX_NAMESPACE::TypeProto* type = output->TypeAsProto(); + ASSERT_NE(type, nullptr); + EXPECT_TRUE(type->has_tensor_type()); + EXPECT_TRUE(type->tensor_type().has_elem_type()); +} + +} // namespace + +// MultiHeadAttention with present_key kept and present_value omitted (exactly two outputs). +// past_key (input 6), past_value (input 7) and past_sequence_length (input 8) are supplied with +// shapes so the present-output branch is active; with only two outputs declared, inference must not +// reference the absent present_value (output 2). Inputs 1-5 are optional and left empty. +TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionPresentValueOmitted) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* present_key = builder.MakeOutput(std::nullopt); + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + past_key, past_value, past_sequence_length}; + Node& node = builder.AddNode("MultiHeadAttention", inputs, {output, present_key}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +// DecoderMaskedMultiHeadAttention with present_key kept and present_value omitted. +// past_key (input 5) and past_value (input 6) are supplied with shapes and past buffer sharing is +// enabled so the present-output branch is active; with only two outputs declared, inference must not +// reference the absent present_value (output 2). Inputs 1-4 are optional and left empty. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionPresentValueOmitted) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* present_key = builder.MakeOutput(std::nullopt); + std::vector inputs = {query, &empty, &empty, &empty, &empty, past_key, past_value}; + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", inputs, {output, present_key}, + kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + node.AddAttribute("past_present_share_buffer", static_cast(1)); + }); +} + +// DecoderAttention with new_key_cache kept and new_value_cache omitted (exactly two outputs). +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionNewValueCacheOmitted) { + BuildResolveAndVerify([](ModelTestBuilder& builder) { + // DecoderAttention requires inputs 0-4 and 8-11; inputs 5-7 are optional and left empty here. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* key = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* q_weight = builder.MakeInput(std::vector{4, 4}); + NodeArg* kv_weight = builder.MakeInput(std::vector{4, 8}); + NodeArg* bias = builder.MakeInput(std::vector{12}); + NodeArg* static_kv = builder.MakeInput(std::vector{1}); + NodeArg* use_past = builder.MakeInput(std::vector{1}); + NodeArg* has_layer_state = builder.MakeInput(std::vector{1}); + NodeArg* has_key_padding_mask = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + NodeArg* new_key_cache = builder.MakeOutput(std::nullopt); + + std::vector inputs = {query, key, q_weight, kv_weight, bias, &empty, &empty, &empty, + static_kv, use_past, has_layer_state, has_key_padding_mask}; + Node& node = builder.AddNode("DecoderAttention", inputs, {output, new_key_cache}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }); +} + +// MultiHeadAttention with all three outputs: the present_key / present_value branch must still run. +TEST(AttentionOptionalOutputsShapeInferenceTest, MultiHeadAttentionAllPresentOutputs) { + NodeArg* present_key = nullptr; + NodeArg* present_value = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + // present_key/present_value types are propagated from past_key (input 6) and past_value + // (input 7) when past buffer sharing is active, which the op detects from shaped past_key + // (input 6) and past_sequence_length (input 8). Leave inputs 1-5 empty to reach them. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_sequence_length = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + present_key = builder.MakeOutput(); + present_value = builder.MakeOutput(); + + std::vector inputs = {query, &empty, &empty, &empty, &empty, &empty, + past_key, past_value, past_sequence_length}; + Node& node = builder.AddNode("MultiHeadAttention", inputs, + {output, present_key, present_value}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }, + [&](const Graph&) { + ExpectInferredElemType(present_key); + ExpectInferredElemType(present_value); + }); +} + +// DecoderMaskedMultiHeadAttention with all three outputs: shape inference must still populate them. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderMaskedMultiHeadAttentionAllPresentOutputs) { + NodeArg* present_key = nullptr; + NodeArg* present_value = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + // For this op past_key/past_value are inputs 5 and 6; leave inputs 1-4 empty to reach them. + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* past_key = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* past_value = builder.MakeInput(std::vector{2, 2, 3, 2}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + present_key = builder.MakeOutput(); + present_value = builder.MakeOutput(); + + std::vector inputs = {query, &empty, &empty, &empty, &empty, past_key, past_value}; + Node& node = builder.AddNode("DecoderMaskedMultiHeadAttention", inputs, + {output, present_key, present_value}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + node.AddAttribute("past_present_share_buffer", static_cast(1)); + }, + [&](const Graph&) { + ExpectInferredElemType(present_key); + ExpectInferredElemType(present_value); + }); +} + +// DecoderAttention with all three outputs: new_key_cache / new_value_cache types must be inferred. +TEST(AttentionOptionalOutputsShapeInferenceTest, DecoderAttentionAllCacheOutputs) { + NodeArg* new_key_cache = nullptr; + NodeArg* new_value_cache = nullptr; + BuildResolveAndVerify( + [&](ModelTestBuilder& builder) { + NodeArg& empty = builder.graph_.GetOrCreateNodeArg("", nullptr); + NodeArg* query = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* key = builder.MakeInput(std::vector{2, 1, 4}); + NodeArg* q_weight = builder.MakeInput(std::vector{4, 4}); + NodeArg* kv_weight = builder.MakeInput(std::vector{4, 8}); + NodeArg* bias = builder.MakeInput(std::vector{12}); + NodeArg* key_cache = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* value_cache = builder.MakeInput(std::vector{2, 2, 3, 2}); + NodeArg* static_kv = builder.MakeInput(std::vector{1}); + NodeArg* use_past = builder.MakeInput(std::vector{1}); + NodeArg* has_layer_state = builder.MakeInput(std::vector{1}); + NodeArg* has_key_padding_mask = builder.MakeInput(std::vector{1}); + + NodeArg* output = builder.MakeOutput(std::nullopt); + new_key_cache = builder.MakeOutput(); + new_value_cache = builder.MakeOutput(); + + std::vector inputs = {query, key, q_weight, kv_weight, bias, &empty, key_cache, + value_cache, static_kv, use_past, has_layer_state, + has_key_padding_mask}; + Node& node = builder.AddNode("DecoderAttention", inputs, + {output, new_key_cache, new_value_cache}, kMSDomain); + node.AddAttribute("num_heads", static_cast(2)); + }, + [&](const Graph&) { + ExpectInferredElemType(new_key_cache); + ExpectInferredElemType(new_value_cache); + }); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc index d32c7fc224fd8..d50e344674c8f 100644 --- a/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc +++ b/onnxruntime/test/contrib_ops/embed_layer_norm_op_test.cc @@ -227,7 +227,7 @@ TEST(EmbedLayerNormTest, EmbedLayerNormBatch_Distill) { RunTest(embedlayernorm::EmbedLayerNormBatch_Distill()); } -// Regression test: negative position_ids must be rejected to not cause OOB read. +// Input validation test: a negative position id must be rejected rather than used to index the table. TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) { int batch_size = 1; int sequence_size = 2; @@ -282,8 +282,103 @@ TEST(EmbedLayerNormTest, EmbedLayerNormNegativePositionIds) { tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); tester.AddOutput("mask_index", mask_index_dims, {0}); - // Run CPU only - expect failure due to negative position_ids - tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kCudaExecutionProvider, kCudaNHWCExecutionProvider, kDmlExecutionProvider, kOpenVINOExecutionProvider}); + // Both CPU and CUDA reject the out-of-range position id via input validation. The CUDA NHWC EP + // shares the same validated kernel, so it is intentionally left enabled (skipped automatically if + // the kernel is not registered for that EP). Only DML and OpenVINO are excluded here. + tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); +} + +// An input id that points outside the word_embedding table must be rejected rather than indexed. +// CPU validates per-index already; CUDA validates device-side and surfaces the same failure, so +// this case runs on both EPs. +TEST(EmbedLayerNormTest, EmbedLayerNormWordIdOutOfRange) { + int batch_size = 1; + int sequence_size = 2; + int hidden_size = 4; + + std::vector input_ids_dims = {batch_size, sequence_size}; + std::vector word_embedding_dims = {6, hidden_size}; + std::vector position_embedding_dims = {3, hidden_size}; + std::vector segment_embedding_dims = {2, hidden_size}; + std::vector gamma_dims = {hidden_size}; + std::vector beta_dims = {hidden_size}; + std::vector output_dims = {batch_size, sequence_size, hidden_size}; + std::vector mask_index_dims = {batch_size}; + + OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain); + // Second id (99) is outside the 6-row word_embedding table. + tester.AddInput("input_ids", input_ids_dims, {1, 99}); + tester.AddInput("segment_ids", input_ids_dims, {0, 1}); + tester.AddInput("word_embedding", word_embedding_dims, + {0.2f, 0.1f, 0.4f, -0.6f, + 0.3f, 0.2f, 0.5f, 0.6f, + 0.6f, 0.7f, 0.0f, -0.1f, + 0.8f, 0.6f, 0.9f, 1.2f, + 0.1f, 0.3f, 0.5f, 0.9f, + 1.0f, -2.0f, 1.1f, 0.8f}, + /*is_initializer=*/true); + tester.AddInput("position_embedding", position_embedding_dims, + {0.1f, 0.1f, 0.4f, 0.6f, + 0.6f, 0.0f, 0.8f, 0.6f, + 0.3f, 0.9f, -2.0f, 0.8f}, + /*is_initializer=*/true); + tester.AddInput("segment_embedding", segment_embedding_dims, + {0.3f, 0.4f, 0.9f, 0.1f, + 0.7f, 0.3f, 0.5f, 0.2f}, + /*is_initializer=*/true); + tester.AddInput("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, /*is_initializer=*/true); + tester.AddInput("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f}, /*is_initializer=*/true); + tester.AddAttribute("epsilon", embedlayernorm::kEpsilon); + + tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); + tester.AddOutput("mask_index", mask_index_dims, {0}); + + tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); +} + +// Without position_ids, positions default to [0, sequence_length); a position_embedding table with +// fewer rows than sequence_length is rejected by the shared input validation on both CPU and CUDA. +TEST(EmbedLayerNormTest, EmbedLayerNormPositionEmbeddingTooFewRows) { + int batch_size = 1; + int sequence_size = 2; + int hidden_size = 4; + + std::vector input_ids_dims = {batch_size, sequence_size}; + std::vector word_embedding_dims = {6, hidden_size}; + // Only 1 row, but sequence_size is 2 and no position_ids are supplied. + std::vector position_embedding_dims = {1, hidden_size}; + std::vector segment_embedding_dims = {2, hidden_size}; + std::vector gamma_dims = {hidden_size}; + std::vector beta_dims = {hidden_size}; + std::vector output_dims = {batch_size, sequence_size, hidden_size}; + std::vector mask_index_dims = {batch_size}; + + OpTester tester("EmbedLayerNormalization", 1, onnxruntime::kMSDomain); + tester.AddInput("input_ids", input_ids_dims, {1, 3}); + tester.AddInput("segment_ids", input_ids_dims, {0, 1}); + tester.AddInput("word_embedding", word_embedding_dims, + {0.2f, 0.1f, 0.4f, -0.6f, + 0.3f, 0.2f, 0.5f, 0.6f, + 0.6f, 0.7f, 0.0f, -0.1f, + 0.8f, 0.6f, 0.9f, 1.2f, + 0.1f, 0.3f, 0.5f, 0.9f, + 1.0f, -2.0f, 1.1f, 0.8f}, + /*is_initializer=*/true); + tester.AddInput("position_embedding", position_embedding_dims, + {0.1f, 0.1f, 0.4f, 0.6f}, + /*is_initializer=*/true); + tester.AddInput("segment_embedding", segment_embedding_dims, + {0.3f, 0.4f, 0.9f, 0.1f, + 0.7f, 0.3f, 0.5f, 0.2f}, + /*is_initializer=*/true); + tester.AddInput("gamma", gamma_dims, {0.25f, 0.15f, 0.45f, -0.66f}, /*is_initializer=*/true); + tester.AddInput("beta", beta_dims, {0.6f, 0.2f, 0.5f, -0.6f}, /*is_initializer=*/true); + tester.AddAttribute("epsilon", embedlayernorm::kEpsilon); + + tester.AddOutput("output", output_dims, std::vector(batch_size * sequence_size * hidden_size, 0.0f)); + tester.AddOutput("mask_index", mask_index_dims, {0}); + + tester.Run(OpTester::ExpectResult::kExpectFailure, "", {kDmlExecutionProvider, kOpenVINOExecutionProvider}); } } // namespace test diff --git a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc index 7fbe71de5251a..b92d5640745b0 100644 --- a/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc +++ b/onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc @@ -375,6 +375,12 @@ TEST(GatherBlockQuantizedOpTest, UnsupportedTypes) { } #endif +// Runs GatherBlockQuantized with inputs that are expected to be rejected (no zero_points). +// Parameters: +// gather_axis - value for the "gather_axis" attribute +// quantize_axis - value for the "quantize_axis" attribute +// block_size - value for the "block_size" attribute +// bits - value for the "bits" attribute (default 4) template void Test_Fail_WithoutZeroPoints(int64_t gather_axis, int64_t quantize_axis, @@ -507,6 +513,44 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) { } #endif +#ifdef USE_CUDA +// The CUDA EP validates quantization-parameter shapes on the host and gathered-index +// bounds on the device. These cases mirror the CPU coverage above but run on the CUDA +// EP, which previously skipped these checks. They fall back to the CPU EP (which rejects +// the same inputs) when no CUDA device is present. +TEST(GatherBlockQuantizedOpTest, CudaShapeMismatch) { + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); + Test_ShapeMismatch_WithZeroPoints(); +} + +TEST(GatherBlockQuantizedOpTest, CudaInvalidIndices) { + Test_InvalidIndices_WithZeroPoints(); + Test_InvalidIndices_WithZeroPoints(); + Test_InvalidIndices_WithZeroPoints(); +} + +// block_size == 0 is an invalid divisor for relating data and scales shapes, so shape +// inference rejects it with fail_shape_inference during Graph::Resolve. +TEST(GatherBlockQuantizedOpTest, CudaInvalidBlockSizeZero) { + Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); + Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); + Test_Fail_WithoutZeroPoints(/*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/0); +} + +// uint8 data supports bits in {4, 8} and int4/uint4 supports bits == 4 on CUDA; other +// values are rejected by the constructor before the 8 / bits division that derives the +// uint8 packing factor. +TEST(GatherBlockQuantizedOpTest, CudaUnsupportedBits) { + Test_Fail_WithoutZeroPoints( + /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); + Test_Fail_WithoutZeroPoints( + /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/8); + Test_Fail_WithoutZeroPoints( + /*gather_axis=*/0, /*quantize_axis=*/2, /*block_size=*/16, /*bits=*/3); +} +#endif + template void Test_GatherAxis0_WithZeroPoints(int bits = 4) { std::vector data = {-8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, diff --git a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc index ed65700ccd336..af095a141dfcd 100644 --- a/onnxruntime/test/contrib_ops/maxpool_mask_test.cc +++ b/onnxruntime/test/contrib_ops/maxpool_mask_test.cc @@ -161,5 +161,70 @@ TEST(ContribOpTest, MaxPoolWithMask_MaskEmptyBatchDim) { "Mask N and C dimensions must be greater than 0"); } +TEST(ContribOpTest, MaxPoolWithMask_KernelRankMismatch) { + OpTester test("MaxpoolWithMask", 1, onnxruntime::kMSDomain); + + // AddShapeToTensorData(false) omits input shape from the graph so ONNX shape inference is bypassed + // (convPoolShapeInference returns early when hasInputShape is false). This lets the model pass + // Graph::Resolve() and reach Compute() where the kernel-rank guard fires. + test.AddShapeToTensorData(false); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0}); + // 2D kernel_shape, but X has only one spatial dimension (rank 3). + test.AddAttribute("kernel_shape", std::vector{8, 8}); + + // Input X has shape {1, 1, 8} (rank 3 => one spatial dim) + std::vector x_dims = {1, 1, 8}; + std::vector x_vals(8, 1.0f); + + // Mask M matches X shape so the earlier spatial/rank guards pass and we reach the kernel-rank guard. + std::vector m_dims = {1, 1, 8}; + std::vector m_vals(8, 1); + + // Placeholder output shape and values (not validated since we expect failure) + std::vector expected_dims = {1, 1, 1, 1}; + std::vector expected_vals = {1.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddInput("M", m_dims, m_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(BaseTester::ExpectResult::kExpectFailure, + "Pooling kernel rank must equal input spatial rank"); +} + +TEST(ContribOpTest, MaxPoolWithMask_KernelRankTooLarge) { + OpTester test("MaxpoolWithMask", 1, onnxruntime::kMSDomain); + + // Bypass ONNX shape inference (see MaxPoolWithMask_KernelRankMismatch) so the model reaches Compute(). + test.AddShapeToTensorData(false); + + test.AddAttribute("auto_pad", ""); + test.AddAttribute("strides", std::vector{1, 1, 1, 1}); + test.AddAttribute("pads", std::vector{0, 0, 0, 0, 0, 0, 0, 0}); + // 4D kernel_shape with a matching rank-6 X (4 spatial dims): passes the kernel-rank equality guard + // but exceeds the supported 1D/2D/3D pooling ranks. + test.AddAttribute("kernel_shape", std::vector{2, 2, 2, 2}); + + // Input X has shape {1, 1, 2, 2, 2, 2} (rank 6 => four spatial dims) + std::vector x_dims = {1, 1, 2, 2, 2, 2}; + std::vector x_vals(16, 1.0f); + + // Mask M matches X so the earlier guards and the kernel-rank equality guard all pass. + std::vector m_dims = {1, 1, 2, 2, 2, 2}; + std::vector m_vals(16, 1); + + // Placeholder output shape and values (not validated since we expect failure) + std::vector expected_dims = {1, 1, 1, 1}; + std::vector expected_vals = {1.0f}; + + test.AddInput("X", x_dims, x_vals); + test.AddInput("M", m_dims, m_vals); + test.AddOutput("Y", expected_dims, expected_vals); + test.Run(BaseTester::ExpectResult::kExpectFailure, + "Only 1D, 2D, and 3D pooling are supported"); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc index 93a63ce19eb03..a0eceb8528b0a 100644 --- a/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc +++ b/onnxruntime/test/contrib_ops/quantize_lstm_op_test.cc @@ -525,5 +525,84 @@ TEST(DynamicQuantLSTMTest, SharedPrepackedWeights) { } #endif +// Builds a minimal per-tensor DynamicQuantizeLSTM and runs it expecting the kernel's input +// validation to reject the recurrence quantization parameters. The caller supplies the R_scale and +// R_zero_point shapes so a shape that is inconsistent with R (e.g. first dim != num_directions) can +// be exercised. The recurrence parameters must be validated symmetrically with the input (W) ones. +static void RunQuantLSTMExpectInvalidRecurrenceQuantParam(const std::vector& r_scale_dims, + const std::vector& r_zp_dims, + const std::string& expected_error) { + OpTester test("DynamicQuantizeLSTM", 1 /*opset_version*/, onnxruntime::kMSDomain /*domain*/); + + constexpr int64_t num_directions = 1; + constexpr int64_t input_size = 2; + constexpr int64_t hidden_size = 2; + constexpr int64_t batch_size = 1; + constexpr int64_t seq_len = 1; + + auto num_elements = [](const std::vector& dims) { + int64_t count = 1; + for (int64_t dim : dims) { + count *= dim; + } + return static_cast(count); + }; + + test.AddAttribute>("activations", {"sigmoid", "tanh", "tanh"}); + test.AddAttribute("direction", "forward"); + test.AddAttribute("hidden_size", hidden_size); + test.AddAttribute("input_forget", static_cast(0)); + + // X: [seq_length, batch_size, input_size] + test.AddInput("X", {seq_len, batch_size, input_size}, + std::vector(num_elements({seq_len, batch_size, input_size}), 0.0f)); + + // W / R quantized weight values are irrelevant: validation fails before any dequantization. + test.AddInput("W", {num_directions, input_size, 4 * hidden_size}, + std::vector(num_elements({num_directions, input_size, 4 * hidden_size}), 0)); + test.AddInput("R", {num_directions, hidden_size, 4 * hidden_size}, + std::vector(num_elements({num_directions, hidden_size, 4 * hidden_size}), 0)); + + test.AddOptionalInputEdge(); // B + test.AddOptionalInputEdge(); // sequence_lens + test.AddInput("initial_h", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddInput("initial_c", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOptionalInputEdge(); // P + + // Valid per-tensor quantization parameters for the input weights. + test.AddInput("W_scale", {num_directions}, std::vector(num_directions, 1.0f)); + test.AddInput("W_zero_point", {num_directions}, std::vector(num_directions, 0)); + + // Recurrence parameters with caller-supplied (possibly inconsistent) shapes. + test.AddInput("R_scale", r_scale_dims, std::vector(num_elements(r_scale_dims), 1.0f)); + test.AddInput("R_zero_point", r_zp_dims, std::vector(num_elements(r_zp_dims), 0)); + + // Placeholder outputs (not validated: the run fails during input validation). + test.AddOutput("Y", {seq_len, num_directions, batch_size, hidden_size}, + std::vector(num_elements({seq_len, num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOutput("Y_h", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + test.AddOutput("Y_c", {num_directions, batch_size, hidden_size}, + std::vector(num_elements({num_directions, batch_size, hidden_size}), 0.0f)); + + test.Run(OpTester::ExpectResult::kExpectFailure, expected_error); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceZeroPointShape) { + // R_zero_point's first dim must equal num_directions (1); {2} is inconsistent and must be rejected + // rather than silently validated against the input zero point's shape. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{1}, /*r_zp_dims=*/{2}, + "Input R_zero_point must have shape"); +} + +TEST(DynamicQuantLSTMTest, RejectsInconsistentRecurrenceScaleShape) { + // R_scale's first dim must equal num_directions (1); {2} is inconsistent and must be rejected + // rather than silently validated against the input scale's shape. + RunQuantLSTMExpectInvalidRecurrenceQuantParam(/*r_scale_dims=*/{2}, /*r_zp_dims=*/{1}, + "Input R_scale must have shape"); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/generator/range_test.cc b/onnxruntime/test/providers/cpu/generator/range_test.cc index 04664628ecdd1..4c04766ebf359 100644 --- a/onnxruntime/test/providers/cpu/generator/range_test.cc +++ b/onnxruntime/test/providers/cpu/generator/range_test.cc @@ -1,8 +1,13 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + +#include "core/graph/constants.h" +#include "core/session/inference_session.h" #include "gtest/gtest.h" #include "test/providers/provider_test_utils.h" +#include "test/test_environment.h" namespace onnxruntime { namespace test { @@ -88,5 +93,221 @@ TEST(RangeTest, AlmostSameStartAndLimitHighDelta) { RunTest(2.0f, 2.01f, 1000000.0f, {1}, {2.0f}); } +#ifndef DISABLE_CONTRIB_OPS + +namespace { +// Describes a single Range input supplied as a graph initializer: its declared dimensions +// and the exact raw_data bytes. Tests use this to craft both well-formed and truncated +// initializers (e.g. dims=[0] with empty raw_data) for the contrib Range schema. +struct RangeInputSpec { + std::vector dims; + std::string raw_data; +}; + +// Serializes a scalar value to its raw_data byte representation. +template +std::string ToRawData(T value) { + std::string bytes(sizeof(T), '\0'); + std::memcpy(bytes.data(), &value, sizeof(T)); + return bytes; +} + +// A correctly-sized scalar initializer (no dims) holding a single value. +template +RangeInputSpec ScalarInput(T value) { + return RangeInputSpec{{}, ToRawData(value)}; +} + +// A zero-element initializer (dims=[0]) whose raw_data is empty. The declared size matches +// the (empty) payload, so initializer size validation accepts it; shape inference, however, +// still attempts to read the first element. +RangeInputSpec EmptyInput() { + return RangeInputSpec{{0}, std::string{}}; +} + +void AddInitializer(ONNX_NAMESPACE::GraphProto* graph, const std::string& name, int data_type, + const RangeInputSpec& spec) { + auto* initializer = graph->add_initializer(); + initializer->set_name(name); + initializer->set_data_type(data_type); + for (int64_t dim : spec.dims) { + initializer->add_dims(dim); + } + initializer->set_raw_data(spec.raw_data); +} + +// Builds a single com.microsoft Range node model whose start/limit/delta inputs are supplied +// as graph initializers of the given element type, then loads and initializes it in a session. +common::Status BuildAndInitializeContribRangeModel(int data_type, + const RangeInputSpec& start, + const RangeInputSpec& limit, + const RangeInputSpec& delta, + InferenceSession& session_object) { + ONNX_NAMESPACE::ModelProto model; + model.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION); + model.add_opset_import()->set_version(11); + auto* ms_opset = model.add_opset_import(); + ms_opset->set_domain(kMSDomain); + ms_opset->set_version(1); + + auto* graph = model.mutable_graph(); + graph->set_name("ContribRangeGraph"); + + AddInitializer(graph, "start", data_type, start); + AddInitializer(graph, "limit", data_type, limit); + AddInitializer(graph, "delta", data_type, delta); + + auto* output = graph->add_output(); + output->set_name("Y"); + output->mutable_type()->mutable_tensor_type()->set_elem_type(data_type); + + auto* range_node = graph->add_node(); + range_node->set_name("Range"); + range_node->set_op_type("Range"); + range_node->set_domain(kMSDomain); + range_node->add_input("start"); + range_node->add_input("limit"); + range_node->add_input("delta"); + range_node->add_output("Y"); + + std::string serialized_model; + if (!model.SerializeToString(&serialized_model)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Failed to serialize test model."); + } + + std::stringstream model_stream(serialized_model); + ORT_RETURN_IF_ERROR(session_object.Load(model_stream)); + return session_object.Initialize(); +} + +// Returns the inferred dim_value of output Y after a successful Initialize, or -1 if absent. +int64_t GetInferredOutputDim(const InferenceSession& session_object) { + const auto outputs = session_object.GetModelOutputs(); + if (!outputs.first.IsOK() || outputs.second->size() != 1u) { + return -1; + } + const auto* shape = (*outputs.second)[0]->Shape(); + if (shape == nullptr || shape->dim_size() != 1 || !shape->dim(0).has_dim_value()) { + return -1; + } + return shape->dim(0).dim_value(); +} +} // namespace + +// Verifies that shape inference clamps an empty/backward range to a zero-sized dimension, +// matching the CPU kernel behavior, instead of emitting a negative dimension value. +TEST(RangeTest, ContribOp_BackwardRange_InfersZeroDim) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_BackwardRange_InfersZeroDim"; + InferenceSession session_object{so, GetEnvironment()}; + // start > limit with a positive delta yields an empty range. + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(5.0), ScalarInput(0.0), + ScalarInput(1.0), session_object); + ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 0); +} + +// Verifies that a correctly-sized raw_data initializer (exactly sizeof(T) bytes) loads +// successfully and produces the expected inferred dimension, so the length check does not +// over-reject valid models. +TEST(RangeTest, ContribOp_ExactSizeRawData_LoadsSuccessfully) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_ExactSizeRawData_LoadsSuccessfully"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(5.0), + ScalarInput(1.0), session_object); + ASSERT_STATUS_OK(status); + ASSERT_EQ(GetInferredOutputDim(session_object), 5); +} + +// The following tests exercise failure paths that surface via fail_shape_inference, which +// reports the error by throwing an inference error. They are excluded from no-exception +// builds where such a throw would abort rather than yield a failure Status. +#if !defined(ORT_NO_EXCEPTIONS) + +// Verifies that Range shape inference rejects an initializer whose raw_data holds fewer bytes +// than its declared data type requires (here: a zero-element initializer with empty raw_data, +// which initializer size validation accepts), returning a clean failure status instead of +// reading more bytes than the initializer declares. One test per input position. +TEST(RangeTest, ContribOp_TruncatedStartRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedStartRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, EmptyInput(), ScalarInput(5.0), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedLimitRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedLimitRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), EmptyInput(), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedDeltaRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedDeltaRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(5.0), + EmptyInput(), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Exercises the length check for different sizeof(T) template instantiations (float, int64). +TEST(RangeTest, ContribOp_TruncatedFloatRawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedFloatRawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_FLOAT, EmptyInput(), ScalarInput(5.0f), + ScalarInput(1.0f), session_object); + ASSERT_FALSE(status.IsOK()); +} + +TEST(RangeTest, ContribOp_TruncatedInt64RawData_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_TruncatedInt64RawData_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_INT64, EmptyInput(), ScalarInput(5), + ScalarInput(1), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Verifies that a zero delta is rejected cleanly by shape inference. +TEST(RangeTest, ContribOp_ZeroDelta_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_ZeroDelta_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(10.0), + ScalarInput(0.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +// Verifies that a finite element count that exceeds the int64 range is rejected cleanly, +// rather than reaching an out-of-range conversion (start=0, limit=1e19, delta=1). +TEST(RangeTest, ContribOp_CountExceedsInt64Range_FailsCleanly) { + SessionOptions so; + so.session_logid = "RangeTest.ContribOp_CountExceedsInt64Range_FailsCleanly"; + InferenceSession session_object{so, GetEnvironment()}; + const auto status = BuildAndInitializeContribRangeModel( + ONNX_NAMESPACE::TensorProto_DataType_DOUBLE, ScalarInput(0.0), ScalarInput(1e19), + ScalarInput(1.0), session_object); + ASSERT_FALSE(status.IsOK()); +} + +#endif // !defined(ORT_NO_EXCEPTIONS) + +#endif // DISABLE_CONTRIB_OPS + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc index dd23b76d8275d..43579d6217901 100644 --- a/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc +++ b/onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc @@ -50,6 +50,24 @@ std::vector> GetExecutionProviders() { return execution_providers; } +// Only the CPU and CUDA GridSample kernels currently validate that input spatial dimensions +// are non-empty. WebGPU/CoreML use separate kernels without that guard, so the expect-failure +// test below must exclude them to avoid a false failure on those builds. +std::vector> GetCpuAndCudaExecutionProviders() { + std::vector> execution_providers; + + execution_providers.emplace_back(DefaultCpuExecutionProvider()); + +#ifdef USE_CUDA + execution_providers.emplace_back(DefaultCudaExecutionProvider()); +#ifdef ENABLE_CUDA_NHWC_OPS + execution_providers.push_back(DefaultCudaNHWCExecutionProvider()); +#endif +#endif + + return execution_providers; +} + template void RunTests(T& test, std::vector>&& execution_providers) { for (size_t idx = 0; idx < execution_providers.size(); ++idx) { @@ -63,6 +81,17 @@ void RunCpuOnly(T& test) { RunTests(test, GetCpuExecutionProviders()); } +template +void RunTestsExpectFailure(T& test, const std::string& error_msg, + std::vector>&& execution_providers) { + for (size_t idx = 0; idx < execution_providers.size(); ++idx) { + test.Config(BaseTester::ExpectResult::kExpectFailure, error_msg) + .ConfigEp(std::move(execution_providers[idx])) + .RunWithConfig(); + } + execution_providers.clear(); +} + } // namespace template @@ -578,5 +607,32 @@ TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_mixed RunCpuOnly(test); } +TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_border_empty_spatial_rejected) { + // An input with a zero-size spatial dimension is rejected: the output is sized by the grid, + // so a non-empty grid skips the empty-output early return, and spatial dimensions must be + // non-empty for sampling to produce valid sample indices. + OpTester test("GridSample", 20); + test.AddAttribute("mode", std::string("nearest")); + test.AddAttribute("padding_mode", std::string("border")); + test.AddAttribute("align_corners", int64_t{0}); + + std::initializer_list X_shape{1, 1, 0, 5}; + std::initializer_list X_data{}; + std::initializer_list Grid_shape{1, 2, 2, 2}; + std::initializer_list Grid_data{ + TypeParam(0.0f), TypeParam(0.0f), TypeParam(-1.0f), TypeParam(-1.0f), + TypeParam(1.0f), TypeParam(1.0f), TypeParam(0.5f), TypeParam(-0.5f)}; + std::initializer_list Y_shape{1, 1, 2, 2}; + + test.AddInput("X", X_shape, X_data); + test.AddInput("Grid", Grid_shape, Grid_data); + test.AddOutput("Y", Y_shape, + {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)}); + // The zero-spatial-dim guard currently covers the CPU and CUDA kernels; other EPs (WebGPU/CoreML) + // are out of scope for this test, so restrict the expect-failure run to CPU + CUDA only. + RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling", + GetCpuAndCudaExecutionProviders()); +} + } // namespace test } // namespace onnxruntime