saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-10 17:37:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

CustomRegistry should use composition instead of inheritence

Browse source 
CustomRegistry should use composition instead of inheritence
This commit is contained in:
utsabsingharoy 2019-04-06 02:44:10 +05:30 committed by Changming Sun
parent 867e961ee8
commit 36ed91ee9f
7 changed files with 46 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
include/onnxruntime/core/framework/customregistry.h
View file

@ -15,10 +15,11 @@ namespace onnxruntime {
/**
Represents a registry that contains both custom kernels and custom schemas.
*/
class CustomRegistry : public KernelRegistry, public onnxruntime::OnnxRuntimeOpSchemaRegistry {
class CustomRegistry final {
public:
CustomRegistry() = default;
~CustomRegistry() override = default;
CustomRegistry() :
kernel_registry_(std::make_shared<KernelRegistry>()),
opschema_registry_(std::make_shared<onnxruntime::OnnxRuntimeOpSchemaRegistry>()) {}
/**
* Register a kernel definition together with kernel factory method to this session.
@ -31,8 +32,18 @@ class CustomRegistry : public KernelRegistry, public onnxruntime::OnnxRuntimeOpS
common::Status RegisterCustomKernel(KernelCreateInfo&);
common::Status RegisterOpSet(std::vector<ONNX_NAMESPACE::OpSchema>& schemas, const std::string& domain,
int baseline_opset_version, int opset_version);
const std::shared_ptr<KernelRegistry>& GetKernelRegistry();
const std::shared_ptr<onnxruntime::OnnxRuntimeOpSchemaRegistry>& GetOpschemaRegistry();
private:
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(CustomRegistry);
std::shared_ptr<KernelRegistry> kernel_registry_;
std::shared_ptr<onnxruntime::OnnxRuntimeOpSchemaRegistry> opschema_registry_;
};
} // namespace onnxruntime

21
onnxruntime/core/framework/customregistry.cc
View file

@ -5,11 +5,28 @@
namespace onnxruntime {
common::Status CustomRegistry::RegisterCustomKernel(KernelDefBuilder& kernel_def_builder, const KernelCreateFn& kernel_creator) {
return Register(kernel_def_builder, kernel_creator);
return kernel_registry_->Register(kernel_def_builder, kernel_creator);
}
common::Status CustomRegistry::RegisterCustomKernel(KernelCreateInfo& create_info) {
return Register(std::move(create_info));
return kernel_registry_->Register(std::move(create_info));
}
common::Status CustomRegistry::RegisterOpSet(
std::vector<ONNX_NAMESPACE::OpSchema>& schemas,
const std::string& domain,
int baseline_opset_version,
int opset_version) {
return opschema_registry_->RegisterOpSet(schemas, domain, baseline_opset_version, opset_version);
}
const std::shared_ptr<KernelRegistry>& CustomRegistry::GetKernelRegistry() {
return kernel_registry_;
}
const std::shared_ptr<onnxruntime::OnnxRuntimeOpSchemaRegistry>& CustomRegistry::GetOpschemaRegistry() {
return opschema_registry_;
}
} // namespace onnxruntime

1
onnxruntime/core/framework/kernel_registry_manager.h
View file

@ -9,6 +9,7 @@
#include "core/common/status.h"
#include "core/platform/ort_mutex.h"
#include "core/graph/graph_viewer.h"
#include "core/framework/customregistry.h"
namespace onnxruntime {
struct KernelCreateInfo;

6
onnxruntime/core/session/inference_session.cc
View file

@ -294,11 +294,13 @@ common::Status InferenceSession::RegisterCustomRegistry(std::shared_ptr<CustomRe
return Status(common::ONNXRUNTIME, common::FAIL, "Received nullptr for custom registry");
}
custom_registries_.push_back(custom_registry);
// Insert session-level customized kernel registry.
kernel_registry_manager_.RegisterKernelRegistry(custom_registry);
kernel_registry_manager_.RegisterKernelRegistry(custom_registry->GetKernelRegistry());
// if (custom_schema_registries_.empty())
// custom_schema_registries_.push_back();
custom_schema_registries_.push_back(custom_registry);
custom_schema_registries_.push_back(custom_registry->GetOpschemaRegistry());
return Status::OK();
}

5
onnxruntime/core/session/inference_session.h
View file

@ -417,5 +417,10 @@ class InferenceSession {
bool is_inited_ = false; // GUARDED_BY(session_mutex_)
InsertCastTransformer insert_cast_transformer_;
//CustomRegistry objects own the corresponding KernelRegistry and OnnxRuntimeOpSchemaRegistry objects.
//So its lifetime should be same as its constituents. This vector is to extend the lifetime of the owner.
std::vector<std::shared_ptr<CustomRegistry>> custom_registries_;
};
} // namespace onnxruntime

2
onnxruntime/test/framework/opaque_kernels_test.cc
View file

@ -296,7 +296,7 @@ TEST_F(OpaqueTypeTests, RunModel) {
auto shape_def = ConstructFetchSparseShape();
EXPECT_TRUE(registry->RegisterCustomKernel(shape_def, [](const OpKernelInfo& info) { return new FetchSparseTensorShape(info); }).IsOK());
IOnnxRuntimeOpSchemaRegistryList custom_schema_registries_ = {registry};
IOnnxRuntimeOpSchemaRegistryList custom_schema_registries_ = {registry->GetOpschemaRegistry()};
std::unordered_map<std::string, int> domain_to_version = {{onnxruntime::kMLDomain, 8}};
Model model("SparseTensorTest", false, ModelMetaData(), custom_schema_registries_, domain_to_version);

3
onnxruntime/test/providers/provider_test_utils.h
View file

@ -221,7 +221,8 @@ class OpTester {
void AddCustomOpRegistry(std::shared_ptr<CustomRegistry> registry) {
// need to do some static casting so we can easily use this later
custom_schema_registries_.push_back(std::static_pointer_cast<IOnnxRuntimeOpSchemaCollection>(registry));
//UTScustom_schema_registries_.push_back(std::static_pointer_cast<IOnnxRuntimeOpSchemaCollection>(registry->GetOpschemaRegistry()));
custom_schema_registries_.push_back(registry->GetOpschemaRegistry());
custom_session_registries_.push_back(std::static_pointer_cast<CustomRegistry>(registry));
}