[XPU] support pytthon memory api in XPU

Author: cqulilujia

paddle/fluid/pybind/pybind.cc

@@ -2917,6 +2917,7 @@ All parameter, weight, gradient are variables in Paddle.
 
 #ifdef PADDLE_WITH_XPU
   m.def("get_xpu_device_count", platform::GetXPUDeviceCount);
+  m.def("get_xpu_current_device_id", &platform::GetXPUCurrentDeviceId);
   m.def("xpu_empty_cache", platform::EmptyCache);
   m.def("get_xpu_device_utilization_rate",
         platform::GetXPUDeviceUtilizationRate);

paddle/phi/backends/xpu/xpu_info.cc

@@ -249,7 +249,7 @@ int GetXPUDeviceTotalMemory(int dev_id) {
   xpumlMemory_t dev_mem_info;
   PADDLE_ENFORCE_XPUML_SUCCESS(
       xpumlDeviceGetMemoryInfo(dev_handle, &dev_mem_info));
-  return dev_mem_info.totalGlobalMemory / 1024 / 1024;  // MB
+  return dev_mem_info.totalGlobalMemory;  // with Byte
 }
 
 int GetXPUDeviceUsedMemory(int dev_id) {
@@ -264,7 +264,7 @@ int GetXPUDeviceUsedMemory(int dev_id) {
   xpumlMemory_t dev_mem_info;
   PADDLE_ENFORCE_XPUML_SUCCESS(
       xpumlDeviceGetMemoryInfo(dev_handle, &dev_mem_info));
-  return dev_mem_info.usedGlobalMemory / 1024 / 1024;  // MB
+  return dev_mem_info.usedGlobalMemory;  // with Byte
 }
 
 XPUVersion get_xpu_version(int dev_id) {

python/paddle/device/xpu/__init__.py

@@ -25,9 +25,66 @@
 
     _XPUPlaceLike: TypeAlias = Union[
         XPUPlace,
+        str,  # some str like 'xpu:0', 'xpu:1', etc.
         int,  # some int like 0, 1, etc.
     ]
-__all__ = ['synchronize', 'empty_cache']
+
+__all__ = [
+    'synchronize',
+    'device_count',
+    'set_debug_level',
+    'empty_cache',
+    'max_memory_allocated',
+    'max_memory_reserved',
+    'reset_max_memory_allocated',
+    'reset_max_memory_reserved',
+    'memory_allocated',
+    'memory_reserved',
+    'memory_total',  # memory maneged by runtime, not paddle
+    'memory_used',  # memory maneged by runtime, not paddle
+]
+
+
+def extract_xpu_device_id(device: _XPUPlaceLike, op_name: str) -> int:
+    '''
+    Return the id of the given xpu device. It is just a utility that will not be exposed to users.
+
+    Args:
+        device(paddle.XPUPlace or int or str): The device, the id of the device or
+            the string name of device like 'xpu:x'.
+            Default: None.
+
+    Return:
+        int: The id of the given device. If device is None, return the id of current device.
+    '''
+    if device is None:
+        return core.get_xpu_current_device_id()
+
+    if isinstance(device, int):
+        device_id = device
+    elif isinstance(device, core.XPUPlace):
+        device_id = device.get_device_id()
+    elif isinstance(device, str):
+        if device.startswith('xpu:'):
+            device_id = int(device[4:])
+        else:
+            raise ValueError(
+                f"The current string {device} is not expected. Because {op_name} only support string which is like 'xpu:x'. "
+                "Please input appropriate string again!"
+            )
+    else:
+        raise ValueError(
+            f"The device type {device} is not expected. Because {op_name} only support int, str or paddle.XPUPlace. "
+            "Please input appropriate device again!"
+        )
+
+    assert (
+        device_id >= 0
+    ), f"The device id must be not less than 0, but got id = {device_id}."
+    assert (
+        device_id < device_count()
+    ), f"The device id {device_id} exceeds xpu card number {device_count()}"
+    return device_id
 
 
 @deprecated(
@@ -135,6 +192,270 @@ def empty_cache() -> None:
             >>> del tensor
             >>> paddle.device.xpu.empty_cache()
     '''
-
-    if core.is_compiled_with_xpu():
+    name = "paddle.device.xpu.empty_cache"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    else:
         core.xpu_empty_cache()
+
+
+def max_memory_allocated(device: _XPUPlaceLike | None = None) -> int:
+    '''
+    Return the peak size of xpu memory that is allocated to tensor of the given device.
+
+    Note:
+        The size of XPU memory allocated to tensor is 256-byte aligned in Paddle, which may larger than the memory size that tensor actually need.
+        For instance, a float32 0-D Tensor with shape [] in XPU will take up 256 bytes memory, even though storing a float32 data requires only 4 bytes.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Return:
+        int: The peak size of xpu memory that is allocated to tensor of the given device, in bytes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated(paddle.XPUPlace(0))
+            >>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated(0)
+            >>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated("xpu:0")
+    '''
+    name = "paddle.device.xpu.max_memory_allocated"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    return core.device_memory_stat_peak_value("Allocated", device_id)
+
+
+def max_memory_reserved(device: _XPUPlaceLike | None = None) -> int:
+    '''
+    Return the peak size of XPU memory that is held by the allocator of the given device.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Return:
+        int: The peak size of XPU memory that is held by the allocator of the given device, in bytes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved(paddle.XPUPlace(0))
+            >>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved(0)
+            >>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved("xpu:0")
+    '''
+    name = "paddle.device.xpu.max_memory_reserved"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    return core.device_memory_stat_peak_value("Reserved", device_id)
+
+
+def reset_max_memory_allocated(device: _XPUPlaceLike | None = None) -> None:
+    '''
+    Reset the peak size of XPU memory that is allocated to tensor of the given device.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> paddle.device.xpu.reset_max_memory_allocated(paddle.XPUPlace(0))
+            >>> paddle.device.xpu.reset_max_memory_allocated(0)
+            >>> paddle.device.xpu.reset_max_memory_allocated("xpu:0")
+    '''
+
+    name = "paddle.device.xpu.reset_max_memory_allocated"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    core.device_memory_stat_reset_peak_value("Allocated", device_id)
+
+
+def reset_max_memory_reserved(device: _XPUPlaceLike | None = None) -> None:
+    '''
+    Reset the peak size of XPU memory that is held by the allocator of the given device.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> paddle.device.xpu.reset_max_memory_reserved(paddle.XPUPlace(0))
+            >>> paddle.device.xpu.reset_max_memory_reserved(0)
+            >>> paddle.device.xpu.reset_max_memory_reserved("xpu:0")
+    '''
+
+    name = "paddle.device.xpu.reset_max_memory_reserved"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    core.device_memory_stat_reset_peak_value("Reserved", device_id)
+
+
+def memory_allocated(device: _XPUPlaceLike | None = None) -> int:
+    '''
+    Return the current size of xpu memory that is allocated to tensor of the given device.
+
+    Note:
+        The size of XPU memory allocated to tensor is 256-byte aligned in Paddle, which may be larger than the memory size that tensor actually need.
+        For instance, a float32 0-D Tensor with shape [] in XPU will take up 256 bytes memory, even though storing a float32 data requires only 4 bytes.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Return:
+        int: The current size of xpu memory that is allocated to tensor of the given device, in bytes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> memory_allocated_size = paddle.device.xpu.memory_allocated(paddle.XPUPlace(0))
+            >>> memory_allocated_size = paddle.device.xpu.memory_allocated(0)
+            >>> memory_allocated_size = paddle.device.xpu.memory_allocated("xpu:0")
+    '''
+    name = "paddle.device.xpu.memory_allocated"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    return core.device_memory_stat_current_value("Allocated", device_id)
+
+
+def memory_reserved(device: _XPUPlaceLike | None = None) -> int:
+    '''
+    Return the current size of XPU memory that is held by the allocator of the given device.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Return:
+        int: The current size of XPU memory that is held by the allocator of the given device, in bytes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> memory_reserved_size = paddle.device.xpu.memory_reserved(paddle.XPUPlace(0))
+            >>> memory_reserved_size = paddle.device.xpu.memory_reserved(0)
+            >>> memory_reserved_size = paddle.device.xpu.memory_reserved("xpu:0")
+    '''
+    name = "paddle.device.xpu.memory_reserved"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    return core.device_memory_stat_current_value("Reserved", device_id)
+
+
+def memory_total(device: _XPUPlaceLike | None = None) -> int:
+    '''
+    Return the total size of XPU memory of the given device that is held by the XPU Runtime.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Return:
+        int: The total size of XPU memory of the given device that is held by the XPU Runtime, in bytes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> memory_total_size = paddle.device.xpu.memory_total(paddle.XPUPlace(0))
+            >>> memory_total_size = paddle.device.xpu.memory_total(0)
+            >>> memory_total_size = paddle.device.xpu.memory_total("xpu:0")
+    '''
+    name = "paddle.device.xpu.memory_total"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    return core.get_xpu_device_total_memory(device_id)
+
+
+def memory_used(device: _XPUPlaceLike | None = None) -> int:
+    '''
+    Return the used size of XPU memory of the given device that is held by the XPU Runtime.
+
+    Args:
+        device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like 'xpu:x'. If device is None, the device is the current device.
+            Default: None.
+
+    Return:
+        int: The used size of XPU memory of the given device that is held by the XPU Runtime, in bytes.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:XPU)
+            >>> import paddle
+            >>> paddle.device.set_device('xpu')
+
+            >>> memory_used_size = paddle.device.xpu.memory_used(paddle.XPUPlace(0))
+            >>> memory_used_size = paddle.device.xpu.memory_used(0)
+            >>> memory_used_size = paddle.device.xpu.memory_used("xpu:0")
+    '''
+    name = "paddle.device.xpu.memory_used"
+    if not core.is_compiled_with_xpu():
+        raise ValueError(
+            f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
+        )
+    device_id = extract_xpu_device_id(device, op_name=name)
+    return core.get_xpu_device_used_memory(device_id)

python/paddle/distributed/launch/utils/nvsmi.py

@@ -157,8 +157,12 @@ def query_xpu_smi(query=None, index=None, dtype=None):
 
     for dev_id in range(core.get_xpu_device_count()):
         utilization_xpu = core.get_xpu_device_utilization_rate(dev_id)
-        mem_total = core.get_xpu_device_total_memory(dev_id)
-        mem_used = core.get_xpu_device_used_memory(dev_id)
+        mem_total = (
+            core.get_xpu_device_total_memory(dev_id) / 1024 / 1024
+        )  # with MB
+        mem_used = (
+            core.get_xpu_device_used_memory(dev_id) / 1024 / 1024
+        )  # with MB
         result = [
             dev_id,
             utilization_xpu,