add xpu accelerator in 3x (#1754)

xin3he · web-flow · commit 368ba5293ab2 · 2024-04-28T10:38:19.000+08:00
Signed-off-by: xin3he &lt;xin3.he@intel.com&gt;
diff --git a/neural_compressor/torch/utils/auto_accelerator.py b/neural_compressor/torch/utils/auto_accelerator.py
@@ -35,8 +35,9 @@
 htcore = LazyImport("habana_frameworks.torch.core")
 
 PRIORITY_HPU = 100
-PRIORITY_CUDA = 95
-PRIORITY_CPU = 90
+PRIORITY_XPU = 95
+PRIORITY_CUDA = 90
+PRIORITY_CPU = 80
 
 
 class AcceleratorRegistry:
@@ -213,8 +214,44 @@ def device(self, device_index=None):
     def empty_cache(self):
         return torch.cuda.empty_cache()
 
-    def mark_step(self):
-        pass
+
+@register_accelerator(name="xpu", priority=PRIORITY_XPU)
+class XPU_Accelerator(Auto_Accelerator):
+    def __init__(self) -> None:
+        self._name = "xpu"
+
+    def name(self) -> str:
+        return self._name
+
+    @classmethod
+    def is_available(cls) -> bool:
+        if hasattr(torch, "xpu") and torch.xpu.is_available():
+            return True
+        else:
+            return False
+
+    def device_name(self, device_indx) -> str:
+        if device_indx is None:
+            return "xpu"
+        return f"xpu:{device_indx}"
+
+    def synchronize(self):
+        return torch.xpu.synchronize()
+
+    def set_device(self, device_index):
+        return torch.xpu.set_device(device_index)
+
+    def current_device(self):
+        return torch.xpu.current_device()
+
+    def current_device_name(self):
+        return "xpu:{}".format(torch.xpu.current_device())
+
+    def device(self, device_index=None):
+        return torch.xpu.device(device_index)
+
+    def empty_cache(self):
+        return torch.xpu.empty_cache()
 
 
 @register_accelerator(name="hpu", priority=PRIORITY_HPU)
