wehub-resource-sync
94 lines
3.1 KiB
C++
94 lines
3.1 KiB
C++
//
|
|
// PassRegistry.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on b'2020/12/07'.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <string>
|
|
#include <unordered_map>
|
|
|
|
#include "MNN/MNNDefine.h"
|
|
#include "converter/source/optimizer/passes/PassRegistry.hpp"
|
|
|
|
namespace MNN {
|
|
namespace passes {
|
|
|
|
// All registered passes.
|
|
static std::unordered_map<std::string, std::unique_ptr<Pass>>*
|
|
AllRegisteredPasses() {
|
|
static std::unordered_map<std::string, \
|
|
std::unique_ptr<Pass>> g_registered_passes;
|
|
return &g_registered_passes;
|
|
}
|
|
// All registered pass managers.
|
|
static std::vector<std::unique_ptr<PassManager>>* AllRegisteredPassManagers() {
|
|
static std::vector<std::unique_ptr<PassManager>> g_registered_pass_managers;
|
|
return &g_registered_pass_managers;
|
|
}
|
|
|
|
/*static*/ PassManager* PassManagerRegistry::GetPassManager(int index) {
|
|
auto* g_registered_pass_managers = AllRegisteredPassManagers();
|
|
MNN_CHECK(index < g_registered_pass_managers->size(),
|
|
"The pass manager index is out of bounds.");
|
|
return (*g_registered_pass_managers)[index].get();
|
|
}
|
|
|
|
/*static*/ std::vector<PassManager*> PassManagerRegistry::GetAllPassManagers() {
|
|
std::vector<PassManager*> pass_managers;
|
|
for (auto& pm : *(AllRegisteredPassManagers())) {
|
|
pass_managers.push_back(pm.get());
|
|
}
|
|
return pass_managers;
|
|
}
|
|
|
|
/*static*/ void PassManagerRegistry::AddPassManager(const PassManager& pm) {
|
|
auto* g_registered_pass_managers = AllRegisteredPassManagers();
|
|
g_registered_pass_managers->emplace_back(new PassManager(pm));
|
|
}
|
|
|
|
/*static*/ void PassRegistry::AddPass(std::unique_ptr<Pass>&& pass) {
|
|
auto* g_registered_passes = AllRegisteredPasses();
|
|
g_registered_passes->emplace(pass->name(), std::move(pass));
|
|
}
|
|
|
|
/*static*/ Pass* PassRegistry::GetPass(const std::string& pass_name) {
|
|
auto* g_registered_passes = AllRegisteredPasses();
|
|
const auto& it = g_registered_passes->find(pass_name);
|
|
if (it != g_registered_passes->end()) {
|
|
return it->second.get();
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
RewritePassRegistry::RewritePassRegistry(const std::string& pass_name)
|
|
: pass_name_(pass_name) {
|
|
std::unique_ptr<Pass> pass(new RewritePass(pass_name));
|
|
PassRegistry::AddPass(std::move(pass));
|
|
}
|
|
|
|
RewritePass* GetRewritePassByName(const std::string& pass_name) {
|
|
Pass* pass = PassRegistry::GetPass(pass_name);
|
|
MNN_CHECK(pass, "Pass has not been setup.");
|
|
if (pass->type() != Pass::PassType::kRewrite) {
|
|
MNN_ERROR("Pass %s is registered but not rewrite pass.",
|
|
pass_name.c_str());
|
|
}
|
|
RewritePass *rewrite_pass = static_cast<RewritePass*>(pass);
|
|
return rewrite_pass;
|
|
}
|
|
|
|
void RewritePassRegistry::SetVerify(RewritePassRegistry::FuncType verify_fn) {
|
|
auto *rewrite_pass = GetRewritePassByName(pass_name_);
|
|
rewrite_pass->SetVerify(verify_fn);
|
|
}
|
|
|
|
void RewritePassRegistry::SetRewrite(RewritePassRegistry::FuncType rewrite_fn) {
|
|
auto *rewrite_pass = GetRewritePassByName(pass_name_);
|
|
rewrite_pass->SetRewrite(rewrite_fn);
|
|
}
|
|
|
|
} // namespace passes
|
|
} // namespace MNN
|