From a8b9dbbb4359fe82385384a9834c5f7088f191bf Mon Sep 17 00:00:00 2001 From: Piotr Bielak Date: Thu, 1 Oct 2026 15:14:37 +0000 Subject: [PATCH 1/2] Validate `grad_output` in `grid_sampler_{2d,3d}_backward` XPU was the only backend still skipping the `grad_output` shape check, so a mismatched `grad_output` made both backward kernels read out of bounds. CPU, CUDA and MPS already use the hardened helpers added upstream by https://github.com/pytorch/pytorch/pull/191915. - 2d: switch to `check_grid_sampler_2d_backward` - 3d: add `check_grid_sampler_backward`; the bicubic-rejecting 3-arg `check_grid_sampler_3d` overload is kept until 5-D bicubic lands (#5511) - port `test_grid_sample_backward_error_checking` from upstream `test_nn.py` Requested in review of https://github.com/intel/torch-xpu-ops/pull/5511 --- src/ATen/native/xpu/sycl/GridSampler.cpp | 12 ++++++--- test/xpu/test_nn_xpu.py | 31 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/src/ATen/native/xpu/sycl/GridSampler.cpp b/src/ATen/native/xpu/sycl/GridSampler.cpp index 4e4c094b3f9..349b92cea5c 100644 --- a/src/ATen/native/xpu/sycl/GridSampler.cpp +++ b/src/ATen/native/xpu/sycl/GridSampler.cpp @@ -682,8 +682,7 @@ void grid_sampler_2d_backward_kernel( int64_t padding_mode, bool align_corners, std::array output_mask) { - check_grid_sampler_common(input, grid); - check_grid_sampler_2d(input, grid); + check_grid_sampler_2d_backward(input, grid, grad_output); globalContext().alertNotDeterministic("grid_sampler_2d_backward_xpu"); auto N = input.size(0); @@ -1491,8 +1490,13 @@ void grid_sampler_3d_backward_kernel( std::array output_mask) { // See NOTE [ grid_sampler Native Functions ]. // Add checks here in case this is called instead of grid_sampler. - check_grid_sampler_common(input, grid); - check_grid_sampler_3d(input, grid, interpolation_mode); + check_grid_sampler_3d_backward(input, grid, grad_output); + TORCH_CHECK_NOT_IMPLEMENTED( + static_cast(interpolation_mode) != + GridSamplerInterpolation::Bicubic, + "grid_sampler(): bicubic interpolation with 5D input is not " + "implemented for ", + input.device().type()); globalContext().alertNotDeterministic("grid_sampler_3d_backward_xpu"); auto input_requires_grad = output_mask[0]; diff --git a/test/xpu/test_nn_xpu.py b/test/xpu/test_nn_xpu.py index 80491dbb1af..35ac0066dff 100644 --- a/test/xpu/test_nn_xpu.py +++ b/test/xpu/test_nn_xpu.py @@ -14201,6 +14201,37 @@ def test_save_lstm_compatibility(self, device): result = loaded_model(x) self.assertEqual(result, expected) + def test_grid_sample_backward_error_checking(self, device): + input = torch.empty(1, 1, 2, 2, device=device) + grid = torch.empty(1, 1, 1, 2, device=device) + grad_output = torch.empty(1, 1, 1, 1, device=device) + + # assert no error + torch.ops.aten.grid_sampler_2d_backward( + grad_output, input, grid, 0, 0, False, (True, True) + ) + + with self.assertRaisesRegex(ValueError, "expected grad_output to have sizes"): + invalid_grad_output = torch.empty(2, 1, 1, 1, device=device) + torch.ops.aten.grid_sampler_2d_backward( + invalid_grad_output, input, grid, 0, 0, False, (True, True) + ) + + input = torch.empty(1, 1, 2, 2, 2, device=device) + grid = torch.empty(1, 1, 1, 1, 3, device=device) + grad_output = torch.empty(1, 1, 1, 1, 1, device=device) + + # assert no error + torch.ops.aten.grid_sampler_3d_backward( + grad_output, input, grid, 0, 0, False, (True, True) + ) + + with self.assertRaisesRegex(ValueError, "expected grad_output to have sizes"): + invalid_grad_output = torch.empty(2, 1, 1, 1, 1, device=device) + torch.ops.aten.grid_sampler_3d_backward( + invalid_grad_output, input, grid, 0, 0, False, (True, True) + ) + @onlyOn(["cuda", "xpu"]) @tf32_on_and_off(0.005) def test_grid_sample_large(self, device): From 7f2d04436aac979de1d04720c788353adf985472 Mon Sep 17 00:00:00 2001 From: Piotr Bielak Date: Fri, 2 Oct 2026 09:37:18 +0000 Subject: [PATCH 2/2] Fix dispatch name --- src/ATen/native/xpu/sycl/GridSampler.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ATen/native/xpu/sycl/GridSampler.cpp b/src/ATen/native/xpu/sycl/GridSampler.cpp index 349b92cea5c..b31154c85bf 100644 --- a/src/ATen/native/xpu/sycl/GridSampler.cpp +++ b/src/ATen/native/xpu/sycl/GridSampler.cpp @@ -1510,7 +1510,7 @@ void grid_sampler_3d_backward_kernel( at::ScalarType::BFloat16, at::ScalarType::Half, input.scalar_type(), - "grid_sampler_2d_backward_xpu", + "grid_sampler_3d_backward_xpu", [&] { if (canUse32BitIndexMath(input) && canUse32BitIndexMath(grid) && canUse32BitIndexMath(grad_output)) {