Add custom device properties api (#73281)

YqGe585 · web-flow · commit 4bcc5b66c8f9 · 2025-06-15T20:21:35.000+08:00
* Add custom device properties api

* fix example

* small fix

* fix static check

* fix static check
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -171,6 +171,7 @@ limitations under the License. */
 
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
 #include "paddle/fluid/platform/profiler/custom_device/custom_tracer.h"
+#include "paddle/phi/backends/device_base.h"
 #include "paddle/phi/capi/capi.h"
 #include "paddle/phi/core/platform/collective_helper.h"
 #include "paddle/phi/core/platform/device/custom/custom_device_resource_pool.h"
@@ -2924,6 +2925,43 @@ All parameter, weight, gradient are variables in Paddle.
 #endif
 #endif
 
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  m.def(
+      "get_device_properties",
+      [](std::string dev_type, int id) -> const phi::DeviceProp & {
+        return phi::DeviceManager::GetDeviceProperties(dev_type, id);
+      },
+      py::return_value_policy::copy);
+
+  py::class_<phi::DeviceProp>(m, "_customDeviceProperties", py::module_local())
+      .def_property_readonly(
+          "name", [](const phi::DeviceProp &prop) { return prop.name; })
+      .def_property_readonly(
+          "major", [](const phi::DeviceProp &prop) { return prop.major; })
+      .def_property_readonly(
+          "minor", [](const phi::DeviceProp &prop) { return prop.minor; })
+      .def_property_readonly(
+          "total_memory",
+          [](const phi::DeviceProp &prop) { return prop.totalGlobalMem; })
+      .def_property_readonly(
+          "multi_processor_count",
+          [](const phi::DeviceProp &prop) { return prop.multiProcessorCount; })
+      .def_property_readonly(
+          "is_multi_gpu_board",
+          [](const phi::DeviceProp &prop) { return prop.isMultiGpuBoard; })
+      .def_property_readonly(
+          "is_integrated",
+          [](const phi::DeviceProp &prop) { return prop.integrated; })
+      .def("__repr__", [](const phi::DeviceProp &prop) {
+        std::stringstream ostr;
+        ostr << "_customDeviceProperties(name='" << prop.name
+             << "', major=" << prop.major << ", minor=" << prop.minor
+             << ", total_memory=" << prop.totalGlobalMem / (1024 * 1024)
+             << "MB, multi_processor_count=" << prop.multiProcessorCount << ")";
+        return ostr.str();
+      });
+#endif
+
 #ifdef PADDLE_WITH_IPU
   m.def("get_ipu_device_count", platform::GetIPUDeviceCount);
 #endif
diff --git a/paddle/phi/backends/custom/custom_device.cc b/paddle/phi/backends/custom/custom_device.cc
@@ -551,6 +551,20 @@ class CustomDevice : public DeviceInterface {
     return compute_capability;
   }
 
+  DeviceProp& GetDeviceProperties(size_t dev_id) override {
+    const auto device = &devices_pool[dev_id];
+    static DeviceProp prop;
+    if (pimpl_->get_device_properties) {
+      pimpl_->get_device_properties(device, &prop);
+    }
+    VLOG(10) << Type() << " get device properties"
+             << "DeviceProperties(name='" << prop.name
+             << "', major=" << prop.major << ", minor=" << prop.minor
+             << ", total_memory=" << prop.totalGlobalMem / (1024 * 1024)
+             << "MB, multi_processor_count=" << prop.multiProcessorCount << ")";
+    return prop;
+  }
+
   size_t GetRuntimeVersion(size_t dev_id) override {
     const auto device = &devices_pool[dev_id];
     size_t version = 0;
diff --git a/paddle/phi/backends/device_base.cc b/paddle/phi/backends/device_base.cc
@@ -36,6 +36,12 @@ size_t DeviceInterface::GetComputeCapability(size_t dev_id) {
   return 0;
 }
 
+DeviceProp& DeviceInterface::GetDeviceProperties(size_t dev_id) {
+  DeviceProp prop;
+  VLOG(10) << Type() << " get device properties " << 0;
+  return prop;
+}
+
 size_t DeviceInterface::GetRuntimeVersion(size_t dev_id) {
   VLOG(10) << Type() << " get runtime version " << 0;
   return 0;
diff --git a/paddle/phi/backends/device_base.h b/paddle/phi/backends/device_base.h
@@ -23,6 +23,33 @@
 
 namespace phi {
 
+struct DeviceProp {
+  std::string name;
+  int major = 0;
+  int minor = 0;
+  size_t totalGlobalMem = 0;
+  int multiProcessorCount = 0;
+  bool isMultiGpuBoard = false;
+  bool integrated = false;
+
+  DeviceProp() = default;
+
+  DeviceProp(const std::string& name_,
+             int major_,
+             int minor_,
+             size_t totalGlobalMem_,
+             int multiProcessorCount_,
+             bool isMultiGpuBoard_,
+             bool integrated_)
+      : name(name_),
+        major(major_),
+        minor(minor_),
+        totalGlobalMem(totalGlobalMem_),
+        multiProcessorCount(multiProcessorCount_),
+        isMultiGpuBoard(isMultiGpuBoard_),
+        integrated(integrated_) {}
+};
+
 class TraceEventCollector;
 
 class DeviceInterface {  // Driver / Runtime
@@ -38,6 +65,8 @@ class DeviceInterface {  // Driver / Runtime
   // Info
   virtual size_t GetComputeCapability(size_t dev_id);
 
+  virtual DeviceProp& GetDeviceProperties(size_t dev_id);
+
   virtual size_t GetRuntimeVersion(size_t dev_id);
 
   virtual size_t GetDriverVersion(size_t dev_id);
diff --git a/paddle/phi/backends/device_ext.h b/paddle/phi/backends/device_ext.h
@@ -537,6 +537,14 @@ struct C_DeviceInterface {
   C_Status (*get_compute_capability)(const C_Device device,
                                      size_t* compute_capability);
 
+  /**
+   * @brief Get device properties
+   *
+   * @param[void*]    device_properties
+   */
+  C_Status (*get_device_properties)(const C_Device device,
+                                    void* device_properties);
+
   /**
    * @brief Get runtime version
    *
diff --git a/paddle/phi/backends/device_manager.cc b/paddle/phi/backends/device_manager.cc
@@ -474,6 +474,12 @@ size_t DeviceManager::GetComputeCapability(const Place& place) {
   return dev_impl->GetComputeCapability(device_id);
 }
 
+phi::DeviceProp& DeviceManager::GetDeviceProperties(
+    const std::string& device_type, size_t device_id) {
+  auto dev_impl = GetDeviceInterfaceWithType(device_type);
+  return dev_impl->GetDeviceProperties(device_id);
+}
+
 size_t DeviceManager::GetRuntimeVersion(const Place& place) {
   auto device_type = place.GetDeviceType();
   auto device_id = place.GetDeviceId();
diff --git a/paddle/phi/backends/device_manager.h b/paddle/phi/backends/device_manager.h
@@ -171,6 +171,9 @@ class DeviceManager {
 
   static size_t GetComputeCapability(const Place& place);
 
+  static phi::DeviceProp& GetDeviceProperties(const std::string& device_type,
+                                              size_t device_id);
+
   static size_t GetRuntimeVersion(const Place& place);
 
   static size_t GetDriverVersion(const Place& place);
diff --git a/python/paddle/device/__init__.py b/python/paddle/device/__init__.py
@@ -20,6 +20,8 @@
 import re
 from typing import TYPE_CHECKING, Union
 
+from typing_extensions import TypeAlias
+
 import paddle
 from paddle.base import core, framework
 from paddle.base.framework import (
@@ -43,6 +45,15 @@
     _InitStreamBase = Union[core.CUDAStream, core.CustomDeviceStream]
     _InitEventBase = Union[core.CUDAEvent, core.CustomDeviceEvent]
 
+    from paddle import CustomPlace
+    from paddle.base.libpaddle import _customDeviceProperties
+
+    _CustomPlaceLike: TypeAlias = Union[
+        CustomPlace,
+        str,  # some string like "iluvatar_gpu" "metax_gpu:0", etc.
+        int,  # some int like 0, 1, etc.
+    ]
+
 __all__ = [
     'get_cudnn_version',
     'set_device',
@@ -60,6 +71,7 @@
     'get_all_custom_device_type',
     'get_available_device',
     'get_available_custom_device',
+    'get_device_properties',
     'Stream',
     'Event',
     'current_stream',
@@ -460,6 +472,80 @@ def get_available_custom_device() -> list[str] | None:
     return core.get_available_custom_device()
 
 
+def get_device_properties(
+    device: _CustomPlaceLike | None = None,
+) -> _customDeviceProperties:
+    """
+
+    Return the properties of given device.
+
+    Args:
+        device(|paddle.CustomPlace|int|str|None, optional): The device, the id of the device or
+            the string name of device like npu:x' which to get the properties of the
+            device from. If device is None, the device is the current device.
+            Default: None.
+
+    Returns:
+       _customDeviceProperties: The properties of the device which include ASCII string
+        identifying device, major compute capability, minor compute capability, global
+        memory available and the number of multiprocessors on the device.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # import paddle
+            >>> # paddle.device.set_device('npu')
+            >>> # paddle.device.get_device_properties('npu:0')
+            >>> # _customDeviceProperties(name='', major=0, minor=0, total_memory=0MB, multi_processor_count=0)
+
+            >>> # paddle.device.get_device_properties('npu')
+            >>> # _customDeviceProperties(name='', major=0, minor=0, total_memory=0MB, multi_processor_count=0)
+    """
+    device_name = None
+
+    if device is not None:
+        if isinstance(device, str):
+            colon_idx = device.rfind(':')
+            if colon_idx == -1:
+                device_name = device
+                device_id = 0
+            else:
+                device_name = device[:colon_idx]
+                device_id_str = device[colon_idx + 1 :]
+
+                if not device_id_str.isdigit():
+                    raise ValueError(
+                        f"Invalid device ID '{device_id_str}'. "
+                        f"After colon must be digits only. "
+                        "Example: 'metax_gpu:0'"
+                    )
+
+                device_id = int(device_id_str)
+
+        else:
+            raise ValueError(
+                f"The input: {device} is not expected. Because paddle.device."
+                "get_device_properties only support str. "
+                "Please input appropriate device again!"
+                "Example: 'metax_gpu:0'"
+            )
+    else:
+        raise ValueError(
+            f"The input: {device} is not expected. Because paddle.device."
+            "get_device_properties only support str. "
+            "Please input appropriate device again!"
+            "Example: 'metax_gpu:0'"
+        )
+    if not core.is_compiled_with_custom_device(device_name):
+        raise ValueError(
+            f"PaddlePaddle is not compiled with support for '{device_name}' device. "
+            "Please reinstall PaddlePaddle with Custom Device support "
+            "to call this API."
+        )
+
+    return core.get_device_properties(device_name, device_id)
+
+
 class Event:
     '''
 
