Validate GridSample input spatial dimensions are non-empty

[Ti-Tai Wang (titaiwangms)]

Description

GridSample sizes its output from the grid, so an input with a zero-size spatial dimension skips the empty-output early return and reaches interpolation with an invalid spatial extent. This change validates that every input spatial dimension is non-empty in both the CPU and CUDA kernels, returning a clear error for such malformed inputs.

Changes

• CPU GridSample::Compute: reject any zero-size input spatial dimension with a descriptive error.
• CUDA GridSample::ComputeInternal: same validation, layout-aware for NCHW and NHWC, covering both 4-D (H, W) and 5-D (D, H, W) inputs. Messages are consistent across both EPs.
• Add an expect-failure test (zero-size spatial dim, padding_mode=border) scoped to the CPU and CUDA execution providers, which are the kernels that carry the guard.

Motivation

Improves input validation and error diagnostics for malformed GridSample inputs. Covers both the CPU and CUDA execution providers; no behavior change for valid inputs.

Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com

[Copilot]

Pull request overview

This PR tightens input validation for the GridSample operator by rejecting malformed inputs where any input spatial dimension is zero, which previously could bypass the “empty output” early return (since output shape is derived from the grid) and reach interpolation with invalid spatial extents.

Changes:

• Add input-shape validation in CPU GridSample::Compute to reject any zero-size spatial dimension.
• Add equivalent validation in CUDA GridSample::ComputeInternal, handling both NCHW and NHWC layouts for 4-D and 5-D inputs.
• Add an expect-failure unit test for the zero-spatial-dimension case, scoped to CPU and CUDA EPs.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
onnxruntime/core/providers/cpu/tensor/grid_sample.cc	Adds runtime validation to reject zero-size spatial dimensions before sampling/interpolation.
onnxruntime/core/providers/cuda/tensor/grid_sample.cc	Adds the same validation for CUDA (layout-aware for NCHW/NHWC, 4-D/5-D).
onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc	Adds an expect-failure test and helpers to run it only on EPs that implement the guard (CPU/CUDA).

[Dmitri Smirnov (yuslepukhin)]

Issues found
Inconsistent error handling style — CPU kernel (Minor)

In grid_sample.cc:374-382, the existing checks use ORT_ENFORCE (throws on failure), but the new check uses ORT_RETURN_IF_NOT (returns error Status). This is actually the correct choice for the new code — ORT_RETURN_IF_NOT is preferable for input validation since it returns a clean Status rather than throwing. However, it means that this kernel mixes error-handling styles: the preceding rank/grid-consistency checks all ORT_ENFORCE, which will abort if exceptions are disabled. The new code is fine as-is; the old code is the problem, not this PR.

2. WebGPU EP is also vulnerable (Not addressed) The WebGPU GridSample::ComputeInternal at grid_sample.cc:197-240:

Derives output from grid shape (H_out, W_out)
Has an early return only when output_size == 0

Passes H_in/W_in as shader uniforms — if they are 0, the gs_pixel function's bounds checks (r >= i32(uniforms.H_in)) will treat every coordinate as OOB, which for padding_mode=zeros returns 0 (safe), for padding_mode=border clamps to clamp(r, 0, -1) → undefined behavior in WGSL (or at least incorrect), and for padding_mode=reflection divides by 0 in range = v_max - v_min (e.g., f32(0) - 1.0 = -1.0, then dv / range with negative range).

The PR description acknowledges this with "WebGPU/CoreML use separate kernels without that guard, so the expect-failure test below must exclude them." This is acceptable as a scoped fix but the WebGPU EP should get the same guard in a follow-up.

Missing tests

Zero W instead of zero H — e.g., {1,1,5,0} to verify the loop catches non-first spatial dimensions
All spatial dims zero — {1,1,0,0} to confirm no double-error or crash
5-D with non-first spatial dim zero — e.g., {1,1,2,0,3} or {1,1,2,3,0}
NHWC layout for CUDA — The test runs on DefaultCudaNHWCExecutionProvider() but uses NCHW-shaped input {1,1,0,5}. An explicit NHWC-shaped input (e.g., {1,0,5,1} for the NHWC interpretation) would exercise the IsNHWC ? 1 : 2 / IsNHWC ? size()-1 : size() code path more meaningfully.

padding_mode=zeros — The tests only use border. Adding a zeros variant ensures the guard fires before the interpolation regardless of padding mode.
padding_mode=reflection — Same reason; this is the most dangerous mode (division by zero in reflect logic).
Valid input sanity test — A test with non-zero spatial dims that still passes, confirming no false rejection.

[Tianlei Wu (tianleiwu)]

Overall this is a clean, well-tested input-validation hardening: the non-empty spatial-dimension guard is applied consistently across the CPU, CUDA, and WebGPU kernels with a shared, descriptive error message, and the expect-failure tests cover 4-D/5-D, leading/trailing/middle zero dims, and all padding modes, plus a positive sanity case. A few non-blocking notes below.

1. Placement relative to the empty-output early return (inline on CPU/CUDA): the guard runs before the existing if (Y.Shape().Size() == 0) return Status::OK(); short-circuit, so a zero-size input spatial dim combined with a zero-size grid/output — which previously returned OK — now errors. Since no sampling happens when the output is empty, consider placing the guard after that early return to preserve the existing zero-size-output behavior. Non-blocking.

2. PR description scope: the "Changes"/"Motivation" sections mention only the CPU and CUDA kernels, but this PR also adds the same guard to onnxruntime/core/providers/webgpu/tensor/grid_sample.cc and a WebGPU expect-failure test. Worth updating the description to reflect the WebGPU change so it matches the diff.

(The earlier automated suggestion about returning INVALID_ARGUMENT explicitly vs. ORT_RETURN_IF_NOT is still a reasonable consistency point but is left to author discretion.)