Add device guard for xpu conv on multi device (pytorch#153345)

Add device guard for xpu conv on multi device (pytorch#153067)

# Motivation
fixes pytorch#153022
The root cause is that the XPU backend registers the convolution op using `m.impl`, which bypasses the device guard logic typically added by the code generation system. This can lead to unexpected behavior if the current device isn't explicitly set.

# Additional Context
run the following script
```python
import torch
import torchvision.models as models

torch.manual_seed(0)

model = models.resnet50(weights="ResNet50_Weights.DEFAULT")
model.eval()
data = torch.rand(1, 3, 224, 224)

device = torch.device('xpu:1')  # 'xpu:0'
model = model.to(device=device, dtype=torch.float16)
data = data.to(device, dtype=torch.float16)

with torch.no_grad():
    ret = model(data)
    print(ret)

print("Execution finished")
```
The output is
```bash
         -9.2102e-02, -7.7588e-01, -1.4111e+00, -9.2383e-01,  6.4551e-01,
         -6.0730e-03, -7.8271e-01, -1.1904e+00, -4.1602e-01,  3.2715e-02,
         -4.9854e-01, -6.3623e-01, -8.5107e-01, -6.8555e-01, -9.4434e-01,
         -8.8672e-01, -6.7969e-01, -6.9824e-01, -2.8882e-01,  2.0312e+00]],
       device='xpu:1', dtype=torch.float16)
Execution finished

```

Pull Request resolved: pytorch#153067
Approved by: https://github.com/albanD, https://github.com/EikanWang

(cherry picked from commit e06a080)

Co-authored-by: Yu, Guangye <[email protected]>

Author: pytorchbot

aten/src/ATen/native/mkldnn/xpu/Conv.cpp

@@ -401,6 +401,11 @@ Tensor _convolution_out(
     int64_t groups_,
     Attr attr,
     IntArrayRef pad_nd = IntArrayRef({})) {
+  CheckedFrom c = "xpu_convolution";
+  TensorArg input_t{input_r, "input", 1}, weight_t{weight_r, "weight", 2};
+  checkAllSameType(c, {input_t, weight_t});
+  checkAllSameGPU(c, {input_t, weight_t});
+  c10::DeviceGuard device_guard(input_r.device());
   auto ndim = input_r.ndimension();
   TORCH_CHECK(
       3 == ndim || 4 == ndim || 5 == ndim,
@@ -611,6 +616,8 @@ std::tuple<Tensor, Tensor, Tensor> convolution_backward_overrideable(
     IntArrayRef output_padding,
     int64_t groups,
     std::array<bool, 3> output_mask) {
+  CheckedFrom c = "xpu_convolution_backward";
+  c10::DeviceGuard device_guard(grad_output.device());
   auto ndim = input.ndimension();
   TORCH_CHECK(
       3 == ndim || 4 == ndim || 5 == ndim,
@@ -675,6 +682,10 @@ std::tuple<Tensor, Tensor, Tensor> convolution_backward_overrideable(
     grad_bias = at::empty({grad_output_.size(1)}, opt);
 
   if (output_mask[0]) {
+    TensorArg grad_output_t{grad_output, "grad_output", 1},
+        input_t{input, "input", 2};
+    checkAllSameType(c, {grad_output_t, input_t});
+    checkAllSameGPU(c, {grad_output_t, input_t});
     if (input.numel() > 0) {
       if (transposed_) {
         onednn::deconvolution_backward_data(
@@ -701,6 +712,10 @@ std::tuple<Tensor, Tensor, Tensor> convolution_backward_overrideable(
     }
   }
   if (output_mask[1] || output_mask[2]) {
+    TensorArg grad_output_t{grad_output, "grad_output", 1},
+        weight_t{weight, "weight", 2};
+    checkAllSameType(c, {grad_output_t, weight_t});
+    checkAllSameGPU(c, {grad_output_t, weight_t});
     if (input.numel() > 0) {
       if (transposed_) {
         onednn::deconvolution_backward_weights(

test/xpu/test_conv.py

@@ -1,5 +1,6 @@
 # Owner(s): ["module: intel"]
 
+import copy
 import itertools
 import math
 import unittest
@@ -1191,6 +1192,25 @@ def test_conv2d_no_grad(self, device, dtype):
                 output = m(input)
                 self.assertEqual(output, output_ng, rtol=1e-2, atol=1e-5)
 
+    @unittest.skipIf(torch.xpu.device_count() < 2, "only one GPU detected")
+    @dtypes(torch.double, torch.float, torch.half)
+    def test_conv2d_on_multi_device(self, dtype):
+        input = torch.randn(3, 256, 224, 224, dtype=dtype, requires_grad=True)
+        conv = torch.nn.Conv2d(256, 256, kernel_size=3, padding=1, dtype=dtype)
+        output_grad = torch.randn(3, 256, 224, 224, dtype=dtype)
+        input_0 = input.to(device="xpu:0")
+        conv_0 = copy.deepcopy(conv).to(device="xpu:0")
+        output_0 = conv_0(input_0)
+        input_1 = input.to(device="xpu:1")
+        conv_1 = copy.deepcopy(conv).to(device="xpu:1")
+        output_1 = conv_1(input_1)
+        self.assertEqual(output_0.cpu(), output_1.cpu())
+        output_grad_0 = output_grad.to(device="xpu:0")
+        output_0.backward(output_grad_0)
+        output_grad_1 = output_grad.to(device="xpu:1")
+        output_1.backward(output_grad_1)
+        self.assertEqual(output_grad_0.cpu(), output_grad_1.cpu())
+
     def test_conv_double_backward_strided_with_3D_input_and_weight(self, device):
         input = torch.randn(2, 3, 6, device=device)
         weight = torch.randn(3, 3, 3, device=device)