add set_device dcu (#70308)

xiaoguoguo626807 · web-flow · commit 8ddf04fe712f · 2024-12-23T09:37:23.000+08:00
diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py
@@ -8145,11 +8145,11 @@ def device_guard(device: str | None = None) -> Generator[None, None, None]:
         if device == "cpu":
             raise ValueError("Should not set device id for cpu.")
     if (
-        device not in ["cpu", "gpu", "xpu", "", None]
+        device not in ["cpu", "gpu", "dcu", "xpu", "", None]
         and device not in core.get_all_custom_device_type()
     ):
         raise ValueError(
-            "The Attr(device) should be 'cpu', 'xpu', 'gpu' or custom device, and it can also be empty string or None "
+            "The Attr(device) should be 'cpu', 'xpu', 'dcu', 'gpu' or custom device, and it can also be empty string or None "
             f"when there is no need to specify device. But received {device}"
         )
     if index:
@@ -8226,15 +8226,20 @@ def _get_paddle_place(place):
 
     # GPU
     available_gpu_place = re.match(r"gpu:\d+", place)
-    if place == "gpu_pinned" or place == "gpu" or available_gpu_place:
+    if (
+        place == "gpu_pinned"
+        or place == "gpu"
+        or place == "dcu"
+        or available_gpu_place
+    ):
         if not core.is_compiled_with_cuda():
             raise ValueError(
                 f"The device should not be {available_gpu_place.group()}, since PaddlePaddle is "
                 "not compiled with CUDA"
             )
         if place == "gpu_pinned":
             return core.CUDAPinnedPlace()
-        elif place == "gpu":
+        elif place == "gpu" or place == "dcu":
             return core.CUDAPlace(0)
         else:
             place_info_list = place.split(":", 1)
diff --git a/python/paddle/device/__init__.py b/python/paddle/device/__init__.py
@@ -208,7 +208,7 @@ def _convert_to_place(device: str) -> PlaceLike:
         place = core.CustomPlace(device, device_id)
     elif lower_device == 'cpu':
         place = core.CPUPlace()
-    elif lower_device == 'gpu':
+    elif lower_device == 'gpu' or lower_device == 'dcu':
         if not core.is_compiled_with_cuda():
             raise ValueError(
                 "The device should not be 'gpu', "
@@ -232,7 +232,9 @@ def _convert_to_place(device: str) -> PlaceLike:
             )
         place = core.IPUPlace()
     else:
-        available_gpu_device = re.match(r'gpu:\d+', lower_device)
+        available_gpu_device = re.match(r'gpu:\d+', lower_device) or re.match(
+            r'dcu:\d+', lower_device
+        )
         available_xpu_device = re.match(r'xpu:\d+', lower_device)
         if available_gpu_device:
             if not core.is_compiled_with_cuda():
@@ -268,6 +270,7 @@ def _convert_to_place(device: str) -> PlaceLike:
                             f"'{x}', '{x}:x'"
                             for x in [
                                 'gpu',
+                                'dcu',
                                 'xpu',
                                 'npu',
                                 *core.get_all_custom_device_type(),
