[Gopalakrishnan Nallasamy (GopalakrishnanN)]

Description

Hardens the CPU GridSample operator against undefined behavior caused by static_cast<int64_t> of NaN, ±Inf, or out-of-range floating-point coordinates. Tracks IcM 31000000575970.

Motivation and Context

GridSample previously fed denormalized grid coordinates directly into integer casts (std::floor, std::nearbyint) inside its 2D and 3D Compute paths, the bilinear fast path, and GsReflect. With pathological grid values (NaN, ±Inf, magnitudes beyond INT64_MAX), the casts invoke C++ undefined behavior. On some toolchains/sanitizers this corrupts indices, leading to out-of-bounds reads (a security concern) and on most platforms it just silently returns garbage indices that PixelAtGrid later clamps.

Changes

onnxruntime/core/providers/cpu/tensor/grid_sample.cc

• GsReflect: explicit guard that returns x_min when the input is non-finite or range <= 0 (covers align_corners=true with a 1×1 image where reflection range collapses to zero).
• New helper IsSafeForInt64Conversion<T>(v): rejects NaN/Inf and any magnitude > 2⁶².
• 2D Compute loop: sanitizes unsafe x/y to x_min/y_min before all static_cast<int64_t> calls (Nearest, Linear, Cubic).
• 3D Compute loop: same sanitization for x/y/z.
• PrecomputeBilinearSamplePlan2D (bilinear fast path): substitutes -0.5 for unsafe coords; the existing mask logic correctly rejects the resulting out-of-range neighbors.
• Adds <algorithm> and <cmath> includes.

onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.cc

• New standalone TU (renamed from a previously included .inc) with MIT copyright header.
• Adds 7 regression tests across <float, MLFloat16>:
  • NaN coordinates (nearest + reflection)
  • ±Inf coordinates (bilinear + reflection)
  • Extreme finite coordinates (bilinear + reflection)
  • Cubic + reflection with extreme coordinates (float-only by ONNX type constraint)
  • 3D / 5-D nearest + reflection with extreme coordinates
  • Bilinear + border with extreme coordinates
  • 1×1 image + align_corners=1 (zero reflection range)

cmake/onnxruntime_unittests.cmake

• Adds the new .cc test to the Emscripten exclusion list alongside the existing grid_sample_test.cc.

Validation

• Local Windows Release build: cmake --build . --target onnxruntime_provider_test.
• All *GridSample* tests pass: 116 (96 generated + 8 contrib op + 12 custom across float/MLFloat16).
• Other unrelated test failures observed in the full provider test run are pre-existing and unrelated to GridSample.

Related

Supersedes the work in #27975 (left open). Filed via origin remote per workflow guidance.

[Dmitri Smirnov (yuslepukhin)]

Gaps in test coverage:

• zeros padding + extreme/NaN/Inf coordinates: Not tested. While zeros mode is safe by design, a regression test would confirm the IsSafeForInt64Conversion sanitization doesn't break it.
• cubic + NaN/Inf: Only extreme-finite is tested for cubic. NaN and Inf paths through cubic interpolation are not covered.
• 3D (5D) + linear mode + extreme coordinates: Only nearest mode is tested for 3D. The 3D trilinear interpolation path with extreme coords is not covered.
• 3D + NaN/Inf coordinates: Not tested at all for the 3D path.
• align_corners=1 with non-degenerate images + extreme coordinates: Only the degenerate 1×1 case is tested with align_corners=1.
• Negative extreme coordinates asymmetry: Most tests use symmetric positive/negative extremes. A test with only-negative extreme coordinates in a non-constant image would verify correct reflection direction.
• Mixed normal + pathological coordinates in the same grid: All pathological tests use constant-valued images (all 1.0), making the expected output trivially deterministic. A test with a non-constant image and a mix of normal + pathological grid points would validate that normal sampling is unaffected by the sanitization.

[Dmitri Smirnov (yuslepukhin)]

CUDA is not fixed.

The CUDA implementation in grid_sample_impl.cu has the exact same vulnerability (this is a contrib op under onnxruntime::contrib::cuda):

Additionally:

• No NaN/Inf guard in GsReflect
• No safety clamp in PixelAtGrid after the reflection branch
• No IsSafeForInt64Conversion check before static_cast<int64_t> calls in the compute kernels
  The CUDA implementation needs an equivalent fix. While GPU UB is less dangerous than CPU UB (no ASAN-detectable heap corruption), it still produces garbage output silently and can crash the GPU context.

WebGPU/JS implementations are not vulnerable — WGSL integer conversions have defined saturation semantics.

File a follow-up for CUDA.

[github-actions]

You can commit the suggested changes from lintrunner.

[Gopalakrishnan Nallasamy (GopalakrishnanN)]

Work item to cover the same vulnerability in CUDA - https://aiinfra.visualstudio.com/ONNX%20Runtime/_workitems/edit/93864

[Dmitri Smirnov (yuslepukhin)]

LGTM