Check noise shape and contiguity in rrelu_with_noise_kernel

src/ATen/native/xpu/sycl/RreluWithNoiseKernels.cpp

@@ -9,6 +9,7 @@
  */
 
 #include <ATen/Dispatch.h>
+#include <ATen/MemoryOverlap.h>
 #include <ATen/native/Math.h>
 #include <ATen/native/Resize.h>
 #include <ATen/native/TensorIterator.h>
@@ -82,8 +83,8 @@ inline void _rrelu_with_noise_xpu_train(
     const Scalar& upper_,
     const std::optional<Generator>& 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<scalar_t, double> ? 2 : 4;
@@ -103,7 +104,7 @@ inline void _rrelu_with_noise_xpu_train(
   }
 
   const scalar_t* input_data = input.const_data_ptr<scalar_t>();
-  scalar_t* noise_data = noise.mutable_data_ptr<scalar_t>();
+  scalar_t* noise_data = tmp_noise.mutable_data_ptr<scalar_t>();
   scalar_t* output_data = tmp_output.mutable_data_ptr<scalar_t>();
 
   double lower = lower_.to<double>();
@@ -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(

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()