Generalize test_print and test_tensor_descriptor to support different accelerators (#801)

Author: EikanWang

test/test_print.py

@@ -38,8 +38,8 @@ def run_kernel_and_capture_output(self, kernel_fn, args):
             code, result = code_and_output(kernel_fn, args)
 
             # Wait for any device prints to reach the host
-            if hasattr(result, "device") and result.device.type == "cuda":
-                torch.cuda.synchronize()
+            if hasattr(result, "device") and result.device.type == DEVICE.type:
+                torch.accelerator.synchronize()
 
             # Grab what pytest captured:  stdout + stderr
             out, err = self._capfd.readouterr()
@@ -69,8 +69,8 @@ def run_kernel_and_capture_output(self, kernel_fn, args):
                 code, result = code_and_output(kernel_fn, args)
 
                 # Force GPU synchronization to ensure all device prints complete
-                if hasattr(result, "device") and result.device.type == "cuda":
-                    torch.cuda.synchronize()
+                if hasattr(result, "device") and result.device.type == DEVICE.type:
+                    torch.accelerator.synchronize()
 
                 # Ensure all output is flushed
                 sys.stdout.flush()

test/test_tensor_descriptor.py

@@ -240,7 +240,7 @@ def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         y = torch.randn((64, 64), device=DEVICE, dtype=torch.float16)
 
         code, result = code_and_output(matmul, (x, y))
-        torch.cuda.synchronize()
+        torch.accelerator.synchronize()
         expected = torch.matmul(x, y)
         torch.testing.assert_close(result, expected, atol=1e-2, rtol=1e-2)