Support paddle.set_flags and paddle.get_flags in custom device (#73079)

YqGe585 · web-flow · commit 2647bc087ff6 · 2025-06-07T16:50:57.000+08:00
* CustomDevice memory support

* modified global_value_getter_setter.cc to support set_flag and get_flag in custom device

* small fix

* small fix

* fix dcu test error

* small fix
diff --git a/paddle/fluid/pybind/global_value_getter_setter.cc b/paddle/fluid/pybind/global_value_getter_setter.cc
@@ -48,6 +48,8 @@ namespace paddle::pybind {
 
 namespace py = pybind11;
 
+static void RegisterGlobalVarGetterSetter();
+
 class PYBIND11_HIDDEN GlobalVarGetterSetterRegistry {
   DISABLE_COPY_AND_ASSIGN(GlobalVarGetterSetterRegistry);
 
@@ -199,6 +201,9 @@ class PYBIND11_HIDDEN GlobalVarGetterSetterRegistry {
   }
 
   bool IsPublic(const std::string &name) const {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+    RegisterGlobalVarGetterSetter();
+#endif
     return var_infos_.count(name) > 0 && var_infos_.at(name).is_public;
   }
 
@@ -219,8 +224,6 @@ class PYBIND11_HIDDEN GlobalVarGetterSetterRegistry {
 
 GlobalVarGetterSetterRegistry GlobalVarGetterSetterRegistry::instance_;
 
-static void RegisterGlobalVarGetterSetter();
-
 void BindGlobalValueGetterSetter(pybind11::module *module) {
   RegisterGlobalVarGetterSetter();
 
@@ -290,7 +293,27 @@ struct RegisterGetterSetterVisitor {
   bool is_writable_;
   void *value_ptr_;
 };
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+static void RegisterGlobalVarGetterSetter() {
+  static std::unordered_set<std::string> registered_flags;
+  const auto &flag_map = phi::GetExportedFlagInfoMap();
+  for (const auto &pair : flag_map) {
+    const std::string &name = pair.second.name;
+
+    if (registered_flags.count(name)) {
+      continue;
+    }
+    registered_flags.insert(name);
 
+    bool is_writable = pair.second.is_writable;
+    void *value_ptr = pair.second.value_ptr;
+    const auto &default_value = pair.second.default_value;
+    RegisterGetterSetterVisitor visitor(
+        "FLAGS_" + name, is_writable, value_ptr);
+    paddle::visit(visitor, default_value);
+  }
+}
+#else
 static void RegisterGlobalVarGetterSetter() {
   const auto &flag_map = phi::GetExportedFlagInfoMap();
   for (const auto &pair : flag_map) {
@@ -303,5 +326,5 @@ static void RegisterGlobalVarGetterSetter() {
     paddle::visit(visitor, default_value);
   }
 }
-
+#endif
 }  // namespace paddle::pybind
