unify device_interface

hoshibara · hoshibara · commit b7bbb96a5b7b · 2025-04-29T10:06:41.000Z
remove duplicated wa
add format ignore flag
diff --git a/torch/_inductor/autotune_process.py b/torch/_inductor/autotune_process.py
@@ -547,20 +547,11 @@ def do_bench(
         if len(device_idx_set) == 1:
             device_idx = next(iter(device_idx_set))
         else:
-            if torch.cuda.is_available():
-                device_idx = torch.cuda.current_device()
-            elif torch.xpu.is_available():
-                device_idx = torch.xpu.current_device()
-
-        
-        if torch.cuda.is_available():
-            with torch.cuda.device(device_idx):
-                out = benchmarker.benchmark_gpu(fn)
-                torch.cuda.synchronize() # shake out any CUDA errors
-        elif torch.xpu.is_available():
-            with torch.xpu.device(device_idx):
-                out = benchmarker.benchmark_gpu(fn)
-                torch.xpu.synchronize() # shake out any XPU errors
+            device_idx = device_interface.current_device()
+
+        with device_interface.device(device_idx):  # type: ignore[attr-defined]
+            out = benchmarker.benchmark_gpu(fn)
+            device_interface.synchronize()  # shake out any GPU errors
 
         return out
 
diff --git a/torch/_ops.py b/torch/_ops.py
@@ -256,7 +256,7 @@ def resolve_key(op: OperatorBase, k: DispatchKey):  # type: ignore[valid-type]
     DispatchKey.BackendSelect,
     DispatchKey.AutocastCPU,  # type: ignore[attr-defined]
     DispatchKey.AutocastCUDA,  # type: ignore[attr-defined]
-    DispatchKey.AutocastXPU,
+    DispatchKey.AutocastXPU,  # type: ignore[attr-defined]
 ]
 
 
diff --git a/torch/testing/_internal/common_device_type.py b/torch/testing/_internal/common_device_type.py
@@ -1339,8 +1339,7 @@ def dep_fn(self, *args, **kwargs):
             if inductor and torch._inductor.config.cpp_wrapper and _device != "cpu":
                 size_bytes *= 2
 
-            # TODO: Memory availability checks for Intel GPU
-            if _device != "xpu" and not _has_sufficient_memory(_device, size_bytes):
+            if not _has_sufficient_memory(_device, size_bytes):
                 raise unittest.SkipTest(f"Insufficient {_device} memory")
 
             return fn(self, *args, **kwargs)

Original file line number	Diff line number	Diff line change
`@@ -256,7 +256,7 @@ def resolve_key(op: OperatorBase, k: DispatchKey): # type: ignore[valid-type]`
`256`	`256`	`DispatchKey.BackendSelect,`
`257`	`257`	`DispatchKey.AutocastCPU, # type: ignore[attr-defined]`
`258`	`258`	`DispatchKey.AutocastCUDA, # type: ignore[attr-defined]`
`259`		`- DispatchKey.AutocastXPU,`
	`259`	`+ DispatchKey.AutocastXPU, # type: ignore[attr-defined]`
`260`	`260`	`]`
`261`	`261`
`262`	`262`