Line data    Source code

       1              : /**
       2              :  * Copyright (c) 2026 Huawei Technologies Co., Ltd.
       3              :  * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
       4              :  * CANN Open Software License Agreement Version 2.0 (the "License").
       5              :  * Please refer to the License for details. You may not use this file except in compliance with the License.
       6              :  * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
       7              :  * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
       8              :  * See LICENSE in the root of the software repository for the full text of the License.
       9              :  */
      10              : 
      11              : #include "pass_plugin_loader.h"
      12              : 
      13              : #include <cstdlib>
      14              : #include <mutex>
      15              : 
      16              : #include "framework/common/debug/ge_log.h"
      17              : #include "register/custom_pass_helper.h"
      18              : #include "python_pass_pybind_bridge.h"
      19              : 
      20              : namespace ge {
      21              : namespace fusion {
      22              : namespace {
      23              : constexpr const char *kEnvPythonPassPath = "ASCEND_GE_PY_PASS_PATH";
      24              : 
      25              : bool NeedLoadPythonPasses() {
      26              :   const char *env_value = std::getenv(kEnvPythonPassPath);
      27              :   return (env_value != nullptr) && (env_value[0] != '\0');
      28              : }
      29              : 
      30              : class PassPluginLoader {
      31              :  public:
      32              :   static PassPluginLoader &GetInstance() {
      33              :     static PassPluginLoader instance;
      34              :     return instance;
      35              :   }
      36              : 
      37              :   Status Load() {
      38              :     std::lock_guard<std::mutex> lock(mutex_);
      39              :     if (active_users_ == 0U) {
      40              :       if (!cpp_pass_loaded_) {
      41              :         const auto ret = CustomPassHelper::Instance().Load();
      42              :         if (ret != SUCCESS) {
      43              :           GELOGE(ret, "Load C++ custom pass plugins failed.");
      44              :           return ret;
      45              :         }
      46              :         cpp_pass_loaded_ = true;
      47              :       }
      48              :       if ((!python_pass_loaded_) && NeedLoadPythonPasses()) {
      49              :         const auto ret = RegisterPythonPassesFromPlugin();
      50              :         if (ret != SUCCESS) {
      51              :           GELOGE(ret, "Load Python fusion pass plugins failed.");
      52              :           (void)CustomPassHelper::Instance().Unload();
      53              :           cpp_pass_loaded_ = false;
      54              :           return ret;
      55              :         }
      56              :         python_pass_loaded_ = true;
      57              :       }
      58              :     }
      59              :     active_users_++;
      60         1139 :     GELOGD("LoadPassPlugins active_users_=%zu", active_users_);
      61              :     return SUCCESS;
      62              :   }
      63              : 
      64              :   Status Unload() {
      65              :     std::lock_guard<std::mutex> lock(mutex_);
      66              :     if (active_users_ == 0U) {
      67              :       GELOGW("UnloadPassPlugins called with no active users, possible reference leak.");
      68              :       return SUCCESS;
      69              :     }
      70              :     active_users_--;
      71         1134 :     GELOGD("UnloadPassPlugins active_users_=%zu", active_users_);
      72              :     if (active_users_ == 0U) {
      73              :       if (python_pass_loaded_) {
      74              :         UnloadPythonPasses();
      75              :         python_pass_loaded_ = false;
      76              :       }
      77              :       if (cpp_pass_loaded_) {
      78              :         cpp_pass_loaded_ = false;
      79              :         (void)CustomPassHelper::Instance().Unload();
      80              :       }
      81              :       if (!shutdown_done_) {
      82              :         shutdown_done_ = true;
      83              :         ShutdownPythonPassesForProcess();
      84              :         GELOGI("[PythonPass] ShutdownPythonPassesForProcess done.");
      85              :       }
      86              :     }
      87              :     return SUCCESS;
      88              :   }
      89              : 
      90              :  private:
      91              :   std::mutex mutex_;
      92              :   size_t active_users_{0U};
      93              :   bool cpp_pass_loaded_{false};
      94              :   bool python_pass_loaded_{false};
      95              :   bool shutdown_done_{false};
      96              : };
      97              : }  // namespace
      98              : 
      99              : Status LoadPassPlugins() {
     100              :   return PassPluginLoader::GetInstance().Load();
     101              : }
     102              : 
     103              : Status UnloadPassPlugins() {
     104              :   return PassPluginLoader::GetInstance().Unload();
     105              : }
     106              : }  // namespace fusion
     107              : }  // namespace ge