[TEST] Use device fixture for assert_helper.py and print_helper.py (#4643)

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

python/test/unit/language/assert_helper.py

@@ -51,8 +51,8 @@ def test_assert(func: str, device: str):
     N = 128  # This value should match with test_print in test_subprocess.py.
     num_warps = N // get_current_target_warp_size()
 
-    x = torch.arange(0, N, dtype=torch.int32, device='xpu')
-    y = torch.zeros((N, ), dtype=x.dtype, device="xpu")
+    x = torch.arange(0, N, dtype=torch.int32, device=device)
+    y = torch.zeros((N, ), dtype=x.dtype, device=device)
     if func == "device_assert":
         kernel_device_assert[(1, )](x, y, num_warps=num_warps, BLOCK=N)
     if func == "device_assert_passes":
@@ -80,7 +80,7 @@ def test_assert(func: str, device: str):
         kernel_assert_passes[(1, )](x, y, num_warps=num_warps, BLOCK=N)
     assert_close(y, x)
     # GPU/host synchronization before exiting the test.
-    torch.xpu.synchronize()
+    getattr(torch, device).synchronize()
 
 
 @triton.jit

python/test/unit/language/print_helper.py

@@ -131,19 +131,14 @@ def test_print(func: str, data_type: str, device: str):
     else:
         assert f"Unknown kernel: {func}"
 
-    if device == "xpu":
-        # FIXME: remove trigger to get output from kernel
-        repr(x)
-        repr(y)
-
     if func != "print_no_arg" and func != "no_arg_print" and func != "device_print_large" and \
        func != "print_multiple_args" and func != "device_print_multiple_args" and \
        func != "device_print_pointer" and func != "device_print_scalar":
         assert_close(y, x)
 
     # Wait until driver complete all the jobs for the device_print, especially test_subprocess
     # require this which captures stdout when child exits.
-    torch.xpu.synchronize()
+    getattr(torch, device).synchronize()
 
 
 if __name__ == "__main__":