[VitisAI] fix deinit vitisai ep (#23725)

### Description
<!-- Describe your changes. -->
Removed the schema when unloading.


### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This fix the crash where the onnxruntime reloads vitisai ep.

Co-authored-by: Yueqing Zhang <[email protected]>

Author: 2 people

onnxruntime/core/providers/shared_library/provider_interfaces.h

@@ -604,6 +604,7 @@ struct ProviderHost {
   virtual ONNX_NAMESPACE::StringStringEntryProto* FunctionProto__add_metadata_props(ONNX_NAMESPACE::FunctionProto* p) = 0;
 
   virtual void RegisterSchema(const std::string& domain, const OrtCustomOp* op) = 0;
+  virtual void DeregisterSchema(const std::string& domain, const std::string& op_type, int version) = 0;
   virtual const ONNX_NAMESPACE::OpSchema* GetSchema(const std::string& name, const int maxInclusiveVersion, const std::string& domain) = 0;
   virtual const std::string& OpSchema__inputs__GetName(const ONNX_NAMESPACE::OpSchema* p, const size_t i) = 0;
   virtual const std::string& OpSchema__inputs__GetTypeStr(const ONNX_NAMESPACE::OpSchema* p, const size_t i) = 0;

onnxruntime/core/providers/vitisai/imp/global_api.cc

@@ -63,6 +63,7 @@ struct OrtVitisAIEpAPI {
   void (*profiler_collect)(
       std::vector<EventInfo>& api_events,
       std::vector<EventInfo>& kernel_events);
+  void (*deinitialize_onnxruntime_vitisai_ep)();
   void Ensure() {
     if (handle_)
       return;
@@ -91,6 +92,7 @@ struct OrtVitisAIEpAPI {
     ORT_THROW_IF_ERROR(env.GetSymbolFromLibrary(handle_, "create_ep_context_nodes", (void**)&create_ep_context_nodes));
     ORT_THROW_IF_ERROR(env.GetSymbolFromLibrary(handle_, "vitisai_ep_on_run_start", (void**)&vitisai_ep_on_run_start));
     ORT_THROW_IF_ERROR(env.GetSymbolFromLibrary(handle_, "vitisai_ep_set_ep_dynamic_options", (void**)&vitisai_ep_set_ep_dynamic_options));
+    std::ignore = env.GetSymbolFromLibrary(handle_, "deinitialize_onnxruntime_vitisai_ep", (void**)&deinitialize_onnxruntime_vitisai_ep);
   }
   void Clear() {
     if (handle_) {
@@ -192,7 +194,7 @@ struct MyCustomOpKernel : OpKernel {
   void* op_kernel_;
 };
 
-void create_kernel_registry(std::vector<OrtCustomOpDomain*> domains) {
+void create_kernel_registry(const std::vector<OrtCustomOpDomain*>& domains) {
   s_kernel_registry_vitisaiep = KernelRegistry::Create();
   for (const auto& domain : domains) {
     for (const auto* op : domain->custom_ops_) {
@@ -245,6 +247,7 @@ void create_kernel_registry(std::vector<OrtCustomOpDomain*> domains) {
     }
   }
 }
+
 void initialize_vitisai_ep() {
   s_library_vitisaiep.Ensure();
   s_domains_vitisaiep.reserve(100);
@@ -253,6 +256,18 @@ void initialize_vitisai_ep() {
   create_kernel_registry(s_domains_vitisaiep);
 }
 
+void deinitialize_vitisai_ep() {
+  if (s_library_vitisaiep.deinitialize_onnxruntime_vitisai_ep) {
+    s_library_vitisaiep.deinitialize_onnxruntime_vitisai_ep();
+  }
+  vaip::deregister_xir_ops(s_domains_vitisaiep);
+  // kernel registry would be repopulated, no need to delete kernel registry
+  s_domains_vitisaiep.clear();
+
+  s_library_vitisaiep.Clear();
+  s_kernel_registry_vitisaiep.reset();
+}
+
 static void set_version_info(vaip_core::OrtApiForVaip& api) {
   const char* magic = "VAIP";
   std::memcpy(reinterpret_cast<char*>(&api.magic), magic, sizeof(api.magic));
@@ -510,8 +525,3 @@ vaip_core::OrtApiForVaip* create_org_api_hook() {
     return &the_global_api;
   }
 }
-
-void deinitialize_vitisai_ep() {
-  s_library_vitisaiep.Clear();
-  s_kernel_registry_vitisaiep.reset();
-}

onnxruntime/core/providers/vitisai/imp/register_xir_ops.cc

@@ -12,11 +12,23 @@ namespace vaip {
 void register_xir_ops(const std::vector<OrtCustomOpDomain*>& domains) {
   for (auto domain : domains) {
     for (auto op : domain->custom_ops_) {
+      // skip dequant quant schema, but register kernel
       if (Provider_GetHost()->GetSchema(op->GetName(op), op->GetStartVersion(op), domain->domain_) == nullptr) {
         Provider_GetHost()->RegisterSchema(domain->domain_, op);
       }
     }
   }
 }
 
+void deregister_xir_ops(const std::vector<OrtCustomOpDomain*>& domains) {
+  for (auto domain : domains) {
+    if (domain->domain_ != "com.xilinx") continue;  // skip dequant quant schema
+    for (auto op : domain->custom_ops_) {
+      if (Provider_GetHost()->GetSchema(op->GetName(op), op->GetStartVersion(op), domain->domain_) != nullptr) {
+        Provider_GetHost()->DeregisterSchema(domain->domain_, op->GetName(op), op->GetStartVersion(op));
+      }
+    }
+  }
+}
+
 }  // namespace vaip

onnxruntime/core/providers/vitisai/imp/register_xir_ops.h

@@ -7,4 +7,5 @@
 
 namespace vaip {
 void register_xir_ops(const std::vector<OrtCustomOpDomain*>& domains);
-}
+void deregister_xir_ops(const std::vector<OrtCustomOpDomain*>& domains);
+}  // namespace vaip

onnxruntime/core/session/provider_bridge_ort.cc

@@ -761,6 +761,11 @@ struct ProviderHostImpl : ProviderHost {
     auto schema = CreateSchema(domain, {op});
     ONNX_NAMESPACE::RegisterSchema(schema, ORT_API_VERSION);
   }
+
+  void DeregisterSchema(const std::string& domain, const std::string& op_type, int version) override {
+    ONNX_NAMESPACE::DeregisterSchema(op_type, version, domain);
+  }
+
   const ONNX_NAMESPACE::OpSchema* GetSchema(const std::string& name, const int maxInclusiveVersion, const std::string& domain) override {
     return ONNX_NAMESPACE::OpSchemaRegistry::Instance()->GetSchema(name, maxInclusiveVersion, domain);
   }