GridSample: harden float->int64 casts against NaN/Inf/out-of-range coords

Merged
GO
6mo ago
+655-80

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.