From 43c6f0794a0230efccce41f45cec9ffdfe041657 Mon Sep 17 00:00:00 2001 From: Pan Zhaowu Date: Thu, 14 Aug 2025 10:37:58 +0000 Subject: [PATCH 1/3] Fix namespace conflict issue between PIR and custom op, with style of override. --- paddle/fluid/eager/api/utils/global_utils.h | 7 ++++++- paddle/fluid/framework/custom_operator.cc | 9 ++++++--- paddle/fluid/framework/op_info.h | 11 ++++++----- paddle/fluid/pybind/pybind.cc | 11 +++++++++-- python/paddle/utils/cpp_extension/extension_utils.py | 6 +++++- 5 files changed, 32 insertions(+), 12 deletions(-) diff --git a/paddle/fluid/eager/api/utils/global_utils.h b/paddle/fluid/eager/api/utils/global_utils.h index 99287e66d5f825..1a5f90579f0bc7 100644 --- a/paddle/fluid/eager/api/utils/global_utils.h +++ b/paddle/fluid/eager/api/utils/global_utils.h @@ -101,7 +101,12 @@ class Controller { void MergeOpMetaInfoMap( const std::unordered_map>& map) { - op_meta_info_map_.insert(map.begin(), map.end()); + for (const auto& [key, value] : map) { + if (op_meta_info_map_.count(key)) { + VLOG(3) << "Replacing existing OpMetaInfo for op: " << key; + } + op_meta_info_map_[key] = value; + } } std::unordered_map& op_meta_infos, auto op_name = OpMetaInfoHelper::GetOpName(base_op_meta); if (OpInfoMap::Instance().Has(op_name)) { - LOG(WARNING) << "Operator (" << op_name << ") has been registered."; - return; + LOG(WARNING) << "Operator (" << op_name + << ") has been registered before as PIR op."; + LOG(WARNING) << "PIR Operator (" << op_name + << ") has been overridden by Custom op!."; } auto& op_inputs = OpMetaInfoHelper::GetInputs(base_op_meta); @@ -1282,7 +1284,8 @@ void RegisterOperatorWithMetaInfoMap( // Register PIR op - if (custom_dialect->HasRegistered(pair.first)) { + if (custom_dialect->HasRegistered(paddle::framework::kCustomDialectPrefix + + pair.first)) { VLOG(3) << "The operator `" << pair.first << "` has been registered. " "Therefore, we will not repeat the registration here."; diff --git a/paddle/fluid/framework/op_info.h b/paddle/fluid/framework/op_info.h index a48eb2edbcfccb..a23c7a06dcb597 100644 --- a/paddle/fluid/framework/op_info.h +++ b/paddle/fluid/framework/op_info.h @@ -138,11 +138,12 @@ class TEST_API OpInfoMap { } void Insert(const std::string& type, const OpInfo& info) { - PADDLE_ENFORCE_NE(Has(type), - true, - common::errors::AlreadyExists( - "Operator (%s) has been registered.", type)); - map_.insert({type, info}); + if (Has(type)) { + map_[type] = info; // override ops + VLOG(0) << "Overriding op: " << type; + } else { + map_.insert({type, info}); + } } const OpInfo& Get(const std::string& type) const { diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index 19034ba6459c13..e2e9c0a929a353 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/paddle/fluid/pybind/pybind.cc @@ -2969,8 +2969,15 @@ All parameter, weight, gradient are variables in Paddle. m.def("init_glog", framework::InitGLOG); m.def("init_memory_method", framework::InitMemoryMethod); m.def("load_op_meta_info_and_register_op", [](const std::string dso_name) { - egr::Controller::Instance().MergeOpMetaInfoMap( - framework::LoadOpMetaInfoAndRegisterOp(dso_name)); + const auto &new_op_meta_info_map = + framework::LoadOpMetaInfoAndRegisterOp(dso_name); + egr::Controller::Instance().MergeOpMetaInfoMap(new_op_meta_info_map); + + py::list key_list; + for (const auto &pair : new_op_meta_info_map) { + key_list.append(pair.first); + } + return key_list; }); m.def("init_devices", []() { framework::InitDevices(); }); m.def("init_default_kernel_signatures", diff --git a/python/paddle/utils/cpp_extension/extension_utils.py b/python/paddle/utils/cpp_extension/extension_utils.py index 2a2a84d0d736c0..3e21e63c008e26 100644 --- a/python/paddle/utils/cpp_extension/extension_utils.py +++ b/python/paddle/utils/cpp_extension/extension_utils.py @@ -164,7 +164,11 @@ def bootstrap_context(): def load_op_meta_info_and_register_op(lib_filename: str) -> list[str]: - core.load_op_meta_info_and_register_op(lib_filename) + # TODO(Pan Zhaowu): implement lib -> proto mapping, with proper update logic + # 1. register op in C++ even when conflict(override) + # 2. update op by get_lib_oplist method + # 3. return full op_names unless conflict inside. + new_list = core.load_op_meta_info_and_register_op(lib_filename) return OpProtoHolder.instance().update_op_proto() From a6a2d32b7a6132bdfef8d26859d5031620cb7e60 Mon Sep 17 00:00:00 2001 From: Pan Zhaowu Date: Thu, 14 Aug 2025 13:01:05 +0000 Subject: [PATCH 2/3] fix miscs. --- paddle/fluid/eager/api/utils/global_utils.h | 1 + paddle/fluid/framework/custom_operator.cc | 54 +++++++++++++++++-- paddle/fluid/framework/custom_operator.h | 7 +-- paddle/fluid/pybind/pybind.cc | 1 + python/paddle/base/framework.py | 4 +- .../utils/cpp_extension/extension_utils.py | 5 +- 6 files changed, 61 insertions(+), 11 deletions(-) diff --git a/paddle/fluid/eager/api/utils/global_utils.h b/paddle/fluid/eager/api/utils/global_utils.h index 1a5f90579f0bc7..2be972011101fe 100644 --- a/paddle/fluid/eager/api/utils/global_utils.h +++ b/paddle/fluid/eager/api/utils/global_utils.h @@ -105,6 +105,7 @@ class Controller { if (op_meta_info_map_.count(key)) { VLOG(3) << "Replacing existing OpMetaInfo for op: " << key; } + VLOG(3) << "Merging OpMetaInfo for op: " << key; op_meta_info_map_[key] = value; } } diff --git a/paddle/fluid/framework/custom_operator.cc b/paddle/fluid/framework/custom_operator.cc index bd3734a62cf5cf..afbe502c2749fe 100644 --- a/paddle/fluid/framework/custom_operator.cc +++ b/paddle/fluid/framework/custom_operator.cc @@ -1270,8 +1270,9 @@ void RegisterOperatorWithMetaInfo(const std::vector& op_meta_infos, OpInfoMap::Instance().Insert(cur_op_name, info); } -void RegisterOperatorWithMetaInfoMap( - const paddle::OpMetaInfoMap& op_meta_info_map, void* dso_handle) { +std::unordered_map> +RegisterOperatorWithMetaInfoMap(const paddle::OpMetaInfoMap& op_meta_info_map, + void* dso_handle) { auto& meta_info_map = op_meta_info_map.GetMap(); VLOG(3) << "Custom Operator: size of op meta info map - " << meta_info_map.size(); @@ -1279,6 +1280,7 @@ void RegisterOperatorWithMetaInfoMap( ::pir::IrContext* ctx = ::pir::IrContext::Instance(); auto* custom_dialect = ctx->GetOrRegisterDialect(); + std::unordered_map> diff_map; for (auto& pair : meta_info_map) { VLOG(3) << "Custom Operator: pair first -> op name: " << pair.first; @@ -1296,16 +1298,54 @@ void RegisterOperatorWithMetaInfoMap( << OpMetaInfoHelper::GetOpName(meta_info); custom_dialect->RegisterCustomOp(meta_info); } + diff_map[pair.first] = pair.second; + + // Register Fluid op + RegisterOperatorWithMetaInfo(pair.second, dso_handle); + } + return diff_map; +} + +/* +std::unordered_map> +RegisterOperatorWithMetaInfoMap( + const paddle::OpMetaInfoMap& op_meta_info_map, void* dso_handle) { + auto& meta_info_map = op_meta_info_map.GetMap(); + VLOG(3) << "Custom Operator: size of op meta info map - " + << meta_info_map.size(); + // pair: {op_type, OpMetaInfo} + ::pir::IrContext* ctx = ::pir::IrContext::Instance(); + auto* custom_dialect = + ctx->GetOrRegisterDialect(); + std::unordered_map> diff_map; + for (auto& pair : meta_info_map) { + VLOG(3) << "Custom Operator: pair first -> op name: " << pair.first; + // Register PIR op + if (custom_dialect->HasRegistered(paddle::framework::kCustomDialectPrefix + +pair.first)) { VLOG(3) << "The operator `" << pair.first + << "` has been registered. " + "Therefore, we will not repeat the registration here."; + diff_map[pair.first] = pair.second; + continue; + } + for (const auto& meta_info : pair.second) { + VLOG(3) << "register pir custom op :" + << OpMetaInfoHelper::GetOpName(meta_info); + custom_dialect->RegisterCustomOp(meta_info); + } + // Register Fluid op RegisterOperatorWithMetaInfo(pair.second, dso_handle); } + return diff_map; } +*/ ////////////////////// User APIs /////////////////////// // load op api -const std::unordered_map>& +std::unordered_map> LoadOpMetaInfoAndRegisterOp(const std::string& dso_name) { void* handle = phi::dynload::GetOpDsoHandle(dso_name); VLOG(3) << "load custom_op lib: " << dso_name; @@ -1313,8 +1353,12 @@ LoadOpMetaInfoAndRegisterOp(const std::string& dso_name) { auto* get_op_meta_info_map = detail::DynLoad(handle, "PD_GetOpMetaInfoMap"); auto& op_meta_info_map = get_op_meta_info_map(); - RegisterOperatorWithMetaInfoMap(op_meta_info_map, handle); - return op_meta_info_map.GetMap(); + auto diff_map = RegisterOperatorWithMetaInfoMap(op_meta_info_map, handle); + for (auto& pair : diff_map) { + VLOG(3) << "diff op name: " << pair.first; + } + // return op_meta_info_map.GetMap(); + return diff_map; } } // namespace paddle::framework diff --git a/paddle/fluid/framework/custom_operator.h b/paddle/fluid/framework/custom_operator.h index 1226be3df7564a..c779aa44aa8bf9 100644 --- a/paddle/fluid/framework/custom_operator.h +++ b/paddle/fluid/framework/custom_operator.h @@ -311,12 +311,13 @@ class CustomGradOpMaker }; // Load custom op api: register op after user compiled -const std::unordered_map>& +std::unordered_map> LoadOpMetaInfoAndRegisterOp(const std::string& dso_name); // Register custom op api: register op directly -void RegisterOperatorWithMetaInfoMap( - const paddle::OpMetaInfoMap& op_meta_info_map, void* dso_handle = nullptr); +std::unordered_map> +RegisterOperatorWithMetaInfoMap(const paddle::OpMetaInfoMap& op_meta_info_map, + void* dso_handle = nullptr); // Interface for selective register custom op. void RegisterOperatorWithMetaInfo(const std::vector& op_meta_infos, diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc index e2e9c0a929a353..805509a9b05890 100644 --- a/paddle/fluid/pybind/pybind.cc +++ b/paddle/fluid/pybind/pybind.cc @@ -2971,6 +2971,7 @@ All parameter, weight, gradient are variables in Paddle. m.def("load_op_meta_info_and_register_op", [](const std::string dso_name) { const auto &new_op_meta_info_map = framework::LoadOpMetaInfoAndRegisterOp(dso_name); + // Merging failed? egr::Controller::Instance().MergeOpMetaInfoMap(new_op_meta_info_map); py::list key_list; diff --git a/python/paddle/base/framework.py b/python/paddle/base/framework.py index 973063a331d007..4058b7a626cf57 100644 --- a/python/paddle/base/framework.py +++ b/python/paddle/base/framework.py @@ -3141,14 +3141,14 @@ def get_op_proto(self, type): raise ValueError(f'Operator "{type}" has not been registered.') return self.op_proto_map[type] - def update_op_proto(self): + def update_op_proto(self, new_op_list): op_protos = get_all_op_protos() custom_op_names = [] for proto in op_protos: if proto.type not in self.op_proto_map: self.op_proto_map[proto.type] = proto custom_op_names.append(proto.type) - + custom_op_names = list(set(custom_op_names).union(set(new_op_list))) return custom_op_names def has_op_proto(self, type): diff --git a/python/paddle/utils/cpp_extension/extension_utils.py b/python/paddle/utils/cpp_extension/extension_utils.py index 3e21e63c008e26..bae757fba747d1 100644 --- a/python/paddle/utils/cpp_extension/extension_utils.py +++ b/python/paddle/utils/cpp_extension/extension_utils.py @@ -169,7 +169,10 @@ def load_op_meta_info_and_register_op(lib_filename: str) -> list[str]: # 2. update op by get_lib_oplist method # 3. return full op_names unless conflict inside. new_list = core.load_op_meta_info_and_register_op(lib_filename) - return OpProtoHolder.instance().update_op_proto() + print("##### new list: ", new_list) + proto_sync_ops = OpProtoHolder.instance().update_op_proto(new_list) + print("# sync ops:", proto_sync_ops) + return proto_sync_ops def custom_write_stub(resource, pyfile): From 2b1c2f6d103170f8d5c911fda18f782fd90a7ad5 Mon Sep 17 00:00:00 2001 From: Pan Zhaowu Date: Thu, 14 Aug 2025 13:04:03 +0000 Subject: [PATCH 3/3] polish --- paddle/fluid/framework/custom_operator.cc | 36 ------------------- .../utils/cpp_extension/extension_utils.py | 6 ---- 2 files changed, 42 deletions(-) diff --git a/paddle/fluid/framework/custom_operator.cc b/paddle/fluid/framework/custom_operator.cc index afbe502c2749fe..47b228031f6848 100644 --- a/paddle/fluid/framework/custom_operator.cc +++ b/paddle/fluid/framework/custom_operator.cc @@ -1306,42 +1306,6 @@ RegisterOperatorWithMetaInfoMap(const paddle::OpMetaInfoMap& op_meta_info_map, return diff_map; } -/* -std::unordered_map> -RegisterOperatorWithMetaInfoMap( - const paddle::OpMetaInfoMap& op_meta_info_map, void* dso_handle) { - auto& meta_info_map = op_meta_info_map.GetMap(); - VLOG(3) << "Custom Operator: size of op meta info map - " - << meta_info_map.size(); - // pair: {op_type, OpMetaInfo} - ::pir::IrContext* ctx = ::pir::IrContext::Instance(); - auto* custom_dialect = - ctx->GetOrRegisterDialect(); - std::unordered_map> diff_map; - for (auto& pair : meta_info_map) { - VLOG(3) << "Custom Operator: pair first -> op name: " << pair.first; - // Register PIR op - if (custom_dialect->HasRegistered(paddle::framework::kCustomDialectPrefix + -pair.first)) { VLOG(3) << "The operator `" << pair.first - << "` has been registered. " - "Therefore, we will not repeat the registration here."; - diff_map[pair.first] = pair.second; - continue; - } - for (const auto& meta_info : pair.second) { - VLOG(3) << "register pir custom op :" - << OpMetaInfoHelper::GetOpName(meta_info); - custom_dialect->RegisterCustomOp(meta_info); - } - - - // Register Fluid op - RegisterOperatorWithMetaInfo(pair.second, dso_handle); - } - return diff_map; -} -*/ - ////////////////////// User APIs /////////////////////// // load op api diff --git a/python/paddle/utils/cpp_extension/extension_utils.py b/python/paddle/utils/cpp_extension/extension_utils.py index bae757fba747d1..6a9b1f40af7ae3 100644 --- a/python/paddle/utils/cpp_extension/extension_utils.py +++ b/python/paddle/utils/cpp_extension/extension_utils.py @@ -164,14 +164,8 @@ def bootstrap_context(): def load_op_meta_info_and_register_op(lib_filename: str) -> list[str]: - # TODO(Pan Zhaowu): implement lib -> proto mapping, with proper update logic - # 1. register op in C++ even when conflict(override) - # 2. update op by get_lib_oplist method - # 3. return full op_names unless conflict inside. new_list = core.load_op_meta_info_and_register_op(lib_filename) - print("##### new list: ", new_list) proto_sync_ops = OpProtoHolder.instance().update_op_proto(new_list) - print("# sync ops:", proto_sync_ops) return proto_sync_ops