diff --git a/src/ATen/native/xpu/sycl/RreluWithNoiseKernels.cpp b/src/ATen/native/xpu/sycl/RreluWithNoiseKernels.cpp index 3d2e2e5368a..8519115636f 100644 --- a/src/ATen/native/xpu/sycl/RreluWithNoiseKernels.cpp +++ b/src/ATen/native/xpu/sycl/RreluWithNoiseKernels.cpp @@ -9,6 +9,7 @@ */ #include +#include #include #include #include @@ -82,8 +83,8 @@ inline void _rrelu_with_noise_xpu_train( const Scalar& upper_, const std::optional& generator) { auto input = input_.contiguous(); - auto noise = noise_.contiguous(); Tensor tmp_output = output.contiguous(); + Tensor tmp_noise = noise_.contiguous(); int64_t numel = input.numel(); constexpr int unroll_factor = std::is_same_v ? 2 : 4; @@ -103,7 +104,7 @@ inline void _rrelu_with_noise_xpu_train( } const scalar_t* input_data = input.const_data_ptr(); - scalar_t* noise_data = noise.mutable_data_ptr(); + scalar_t* noise_data = tmp_noise.mutable_data_ptr(); scalar_t* output_data = tmp_output.mutable_data_ptr(); double lower = lower_.to(); @@ -149,6 +150,10 @@ inline void _rrelu_with_noise_xpu_train( if (!output.is_contiguous()) { output.copy_(tmp_output); } + if (!noise_.is_contiguous()) { + at::assert_no_internal_overlap(noise_); + noise_.copy_(tmp_noise); + } } Tensor& rrelu_with_noise_kernel( @@ -169,6 +174,12 @@ Tensor& rrelu_with_noise_kernel( output_arg{output, "output", 3}; checkAllSameGPU( "rrelu_with_noise_out_xpu", {self_arg, noise_arg, output_arg}); + TORCH_CHECK( + self.sym_sizes() == noise.sym_sizes(), + "noise tensor shape must match self tensor shape. Got self.shape = ", + self.sym_sizes(), + " noise.shape = ", + noise.sym_sizes()); if (training) { AT_DISPATCH_FLOATING_TYPES_AND2( diff --git a/test/regressions/test_rrelu_with_noise.py b/test/regressions/test_rrelu_with_noise.py new file mode 100644 index 00000000000..d48e49d8c26 --- /dev/null +++ b/test/regressions/test_rrelu_with_noise.py @@ -0,0 +1,68 @@ +# Copyright 2020-2026 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 + +# Owner(s): ["module: intel"] +""" +Regression tests for the noise shape check and non-contiguous noise +handling in the XPU rrelu_with_noise kernel. + +The training kernel writes one noise value per element of self directly +into the caller's noise tensor. Nothing used to check that noise had the +same shape as self or was contiguous, so a non-contiguous noise (e.g. an +expanded view) was silently written through a discarded contiguous +temporary: the returned output was correct but the caller's noise kept +its old values. Non-contiguous noise now materializes, the same way the +output does. +""" + +import torch +from torch.testing._internal.common_utils import run_tests, TestCase + +xpu_device = torch.device("xpu") + + +class TestRreluWithNoise(TestCase): + def test_noise_shape_mismatch(self): + x = torch.randn(4, 4, device=xpu_device) + for noise in [ + torch.empty(2, 8, device=xpu_device), + torch.empty(4, 1, device=xpu_device), + ]: + with self.assertRaisesRegex( + RuntimeError, "noise tensor shape must match self tensor shape" + ): + torch._C._nn.rrelu_with_noise(x, noise, training=True) + + def test_noise_non_contiguous(self): + x = torch.randn(4, 8, device=xpu_device) + noise = torch.randn(4, 16, device=xpu_device)[:, ::2] + self.assertFalse(noise.is_contiguous()) + out = torch._C._nn.rrelu_with_noise(x, noise, training=True) + self.assertEqual(out, x * noise) + + # An expanded view: its storage holds one row; writing through its + # data pointer used to go past the storage. Distinct per-element + # noise values cannot land there, so the copy back rejects it. + noise = torch.zeros(1, 8, device=xpu_device).expand(4, 8) + with self.assertRaisesRegex(RuntimeError, "refers to a single memory location"): + torch._C._nn.rrelu_with_noise(x, noise, training=True) + + def test_noise_written(self): + # Positive path: the kernel must write into the caller's noise. + x = torch.randn(64, 64, device=xpu_device) + noise = torch.zeros_like(x) + out = torch._C._nn.rrelu_with_noise(x, noise, 0.1, 0.9, training=True) + pos, neg = x > 0, x <= 0 + self.assertTrue(torch.all(noise[pos] == 1)) + self.assertTrue(torch.all(noise[neg] >= 0.1)) + self.assertTrue(torch.all(noise[neg] <= 0.9)) + self.assertEqual(out, x * noise) + + +if __name__ == "__main__": + run_tests()