Validate GridSample input spatial dimensions are non-empty

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

@@ -379,6 +379,14 @@ Status GridSample<T>::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);

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

@@ -124,6 +124,17 @@ Status GridSample<T, IsNHWC>::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<IsNHWC>;
 
   TensorShapeVector dims_output(dims_input.size());

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

@@ -204,6 +204,15 @@ Status GridSample::ComputeInternal(ComputeContext& context) const {
   ORT_RETURN_IF_NOT(grid_shape.NumDimensions() == 4, "grid must be 4-D");
   ORT_RETURN_IF_NOT(grid_shape[3] == 2, "grid last dimension must be 2");
 
+  // 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.
+  // This kernel is NCHW and 4-D, so the spatial dimensions are indices [2, rank).
+  for (size_t i = 2; i < X_shape.NumDimensions(); ++i) {
+    ORT_RETURN_IF_NOT(X_shape[i] > 0,
+                      "Input spatial dimensions must be non-empty for sampling. Dimension ", i,
+                      " has size ", X_shape[i]);
+  }
+
   const int64_t N = X_shape[0];
   const int64_t C = X_shape[1];
   const int64_t H_in = X_shape[2];

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

@@ -63,6 +63,17 @@ void RunCpuOnly(T& test) {
   RunTests(test, GetCpuExecutionProviders());
 }
 
+template <typename T>
+void RunTestsExpectFailure(T& test, const std::string& error_msg,
+                           std::vector<std::unique_ptr<IExecutionProvider>>&& 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();
+}
+
 // Execution providers targeted by the float->int64 cast hardening in the CUDA GridSample
 // kernels (onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu). CoreML and WebGPU
 // are intentionally excluded here: their integer-conversion semantics differ from the
@@ -602,10 +613,245 @@ TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_reflection_mixed
   RunCpuOnly(test);
 }
 
+// -----------------------------------------------------------------------------------------
+// Non-empty input spatial-dimension validation tests: a zero-size spatial dim must be
+// rejected before sampling (guard fires host-side across CPU/CUDA/WebGPU).
+// -----------------------------------------------------------------------------------------
+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<int64_t> X_shape{1, 1, 0, 5};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
+  std::initializer_list<TypeParam> 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<int64_t> Y_shape{1, 1, 2, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("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());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_5D_nearest_border_empty_spatial_rejected) {
+  // Same rejection as the 4-D case, but for a 5-D input with a zero-size spatial dimension. This keeps
+  // the 5-D branch of the guard (and the CUDA NCHW/NHWC layouts) covered: 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<int64_t> X_shape{1, 1, 0, 2, 2};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 1, 1, 2, 3};
+  std::initializer_list<TypeParam> Grid_data{
+      TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f),
+      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f)};
+  std::initializer_list<int64_t> Y_shape{1, 1, 1, 1, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(0.0f), TypeParam(0.0f)});
+  // Restrict the expect-failure run to CPU + CUDA only, matching the 4-D case above.
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        GetCpuAndCudaExecutionProviders());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_border_empty_last_spatial_rejected) {
+  // Zero-size trailing spatial dimension (W = 0), confirming the guard loop catches a non-first
+  // spatial dimension, not just the leading one. This runs on the CUDA NHWC provider as well: for a
+  // logical NCHW {1, 1, 5, 0} input the internal NHWC layout transform places the zero at a spatial
+  // index in the (N, H, W, C) layout, so it exercises the layout-aware spatial range on both kernels.
+  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<int64_t> X_shape{1, 1, 5, 0};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
+  std::initializer_list<TypeParam> 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<int64_t> Y_shape{1, 1, 2, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape,
+                            {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)});
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        GetCpuAndCudaExecutionProviders());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_border_all_spatial_zero_rejected) {
+  // Both spatial dimensions are zero ({1, 1, 0, 0}); the guard must reject on the first zero dimension
+  // without double-reporting or crashing.
+  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<int64_t> X_shape{1, 1, 0, 0};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
+  std::initializer_list<TypeParam> 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<int64_t> Y_shape{1, 1, 2, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape,
+                            {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)});
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        GetCpuAndCudaExecutionProviders());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_5D_nearest_border_middle_spatial_zero_rejected) {
+  // 5-D input with a zero in a non-first spatial dimension ({1, 1, 2, 0, 3}), confirming the guard
+  // spans the full spatial range of a 5-D tensor rather than only checking the first spatial index.
+  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<int64_t> X_shape{1, 1, 2, 0, 3};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 1, 1, 2, 3};
+  std::initializer_list<TypeParam> Grid_data{
+      TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f),
+      TypeParam(1.0f), TypeParam(1.0f), TypeParam(1.0f)};
+  std::initializer_list<int64_t> Y_shape{1, 1, 1, 1, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape, {TypeParam(0.0f), TypeParam(0.0f)});
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        GetCpuAndCudaExecutionProviders());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_linear_zeros_empty_spatial_rejected) {
+  // The guard must fire before interpolation regardless of padding mode. This covers padding_mode
+  // "zeros" (the mode where reads outside the image are otherwise silently treated as 0).
+  OpTester test("GridSample", 20);
+  test.AddAttribute("mode", std::string("linear"));
+  test.AddAttribute("padding_mode", std::string("zeros"));
+  test.AddAttribute("align_corners", int64_t{0});
+
+  std::initializer_list<int64_t> X_shape{1, 1, 0, 5};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
+  std::initializer_list<TypeParam> 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<int64_t> Y_shape{1, 1, 2, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape,
+                            {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)});
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        GetCpuAndCudaExecutionProviders());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_linear_reflection_empty_spatial_rejected) {
+  // Padding mode "reflection" is the most sensitive to a zero-size spatial dimension because the
+  // reflect logic divides by (length - 1); the guard must reject the input before that computation.
+  OpTester test("GridSample", 20);
+  test.AddAttribute("mode", std::string("linear"));
+  test.AddAttribute("padding_mode", std::string("reflection"));
+  test.AddAttribute("align_corners", int64_t{0});
+
+  std::initializer_list<int64_t> X_shape{1, 1, 0, 5};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
+  std::initializer_list<TypeParam> 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<int64_t> Y_shape{1, 1, 2, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape,
+                            {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)});
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        GetCpuAndCudaExecutionProviders());
+}
+
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_20_4D_nearest_border_nonzero_spatial_accepted) {
+  // Sanity check that the guard does not reject valid, non-empty spatial dimensions: a normal 2x2
+  // input with align_corners=1 samples its four corners exactly and must succeed on both kernels.
+  OpTester test("GridSample", 20);
+  test.AddAttribute("mode", std::string("nearest"));
+  test.AddAttribute("padding_mode", std::string("border"));
+  test.AddAttribute("align_corners", int64_t{1});
+
+  std::initializer_list<int64_t> X_shape{1, 1, 2, 2};
+  std::initializer_list<TypeParam> X_data{TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f), TypeParam(4.0f)};
+  std::initializer_list<int64_t> Grid_shape{1, 1, 4, 2};
+  // align_corners=1 on a 2x2 image maps normalized [-1, 1] -> pixel coords [0, 1]; sample all corners.
+  std::initializer_list<TypeParam> Grid_data{
+      TypeParam(-1.0f), TypeParam(-1.0f),  // -> pixel(row 0, col 0) = 1
+      TypeParam(1.0f), TypeParam(-1.0f),   // -> pixel(row 0, col 1) = 2
+      TypeParam(-1.0f), TypeParam(1.0f),   // -> pixel(row 1, col 0) = 3
+      TypeParam(1.0f), TypeParam(1.0f)};   // -> pixel(row 1, col 1) = 4
+  std::initializer_list<int64_t> Y_shape{1, 1, 1, 4};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape,
+                            {TypeParam(1.0f), TypeParam(2.0f), TypeParam(3.0f), TypeParam(4.0f)});
+  RunTests(test, GetCpuAndCudaExecutionProviders());
+}
+
+#ifdef USE_WEBGPU
+TYPED_TEST(GridSampleCustomTest, test_grid_sample_16_4D_webgpu_empty_spatial_rejected) {
+  // The WebGPU GridSample kernel is registered for opset 16-19 (NCHW, 4-D only) and now carries the
+  // same non-empty spatial-dimension guard as the CPU and CUDA kernels. Exercise it directly on the
+  // WebGPU provider with a zero-size spatial dimension; the guard rejects host-side before any shader
+  // dispatch, so this validates the WebGPU path without depending on a particular GPU adapter.
+  OpTester test("GridSample", 16);
+  test.AddAttribute("mode", std::string("nearest"));
+  test.AddAttribute("padding_mode", std::string("border"));
+  test.AddAttribute("align_corners", int64_t{0});
+
+  std::initializer_list<int64_t> X_shape{1, 1, 0, 5};
+  std::initializer_list<TypeParam> X_data{};
+  std::initializer_list<int64_t> Grid_shape{1, 2, 2, 2};
+  std::initializer_list<TypeParam> 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<int64_t> Y_shape{1, 1, 2, 2};
+
+  test.AddInput<TypeParam>("X", X_shape, X_data);
+  test.AddInput<TypeParam>("Grid", Grid_shape, Grid_data);
+  test.AddOutput<TypeParam>("Y", Y_shape,
+                            {TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f), TypeParam(0.0f)});
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultWebGpuExecutionProvider());
+  RunTestsExpectFailure(test, "Input spatial dimensions must be non-empty for sampling",
+                        std::move(execution_providers));
+}
+#endif
+
 // -----------------------------------------------------------------------------------------
 // CUDA hardening regression tests.
 //
-// These mirror the CPU-only cases above but also run on the CUDA execution provider (when
+// These mirror the coordinate-hardening cases above but also run on the CUDA execution provider (when
 // built) to cover the equivalent float->int64 cast hardening in
 // onnxruntime/core/providers/cuda/tensor/grid_sample_impl.cu. They are float-only because
 // the CUDA GridSample kernel is registered for float only, and they use constant-valued