saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Simplify some CustomOp code (#1206)

Browse source 
test_inference.cc reformatted with clang-format
This commit is contained in:
Ryan Hill 2019-06-11 17:00:04 -07:00 committed by GitHub
parent 87d65389e6
commit 38963d81eb
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 16 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
include/onnxruntime/core/session/onnxruntime_c_api.h
View file

@ -580,8 +580,7 @@ struct OrtCustomOp {
// Returns the name of the op
const char*(ORT_API_CALL* GetName)(_In_ struct OrtCustomOp* op);
// Returns the type of the execution provider
// If the function pointer is null, use CPU execution provider by default
// Returns the type of the execution provider, return nullptr to use CPU execution provider
const char*(ORT_API_CALL* GetExecutionProviderType)(_In_ struct OrtCustomOp* op);
// Returns the count and types of the input & output tensors

33
include/onnxruntime/core/session/onnxruntime_cxx_api.h
View file

@ -278,33 +278,6 @@ struct CustomOpApi {
const OrtCustomOpApi& api_;
};
namespace CustomOpImpl {
// SFINAE definition to determine whether class T as GetExecutionProviderType interface
template <typename T, typename... Args>
class HasProvider {
template <typename U = T,
typename = decltype(std::declval<U>().GetExecutionProviderType(std::declval<Args>()...))>
static std::true_type test(int);
template <typename U = T>
static std::false_type test(...);
public:
static constexpr bool value = decltype(test(0))::value;
};
// SFINAE definitions to get the execution provider of op.
// If type T has GetExecutionProviderType, use the result of op->GetExecutionProviderType().
// Otherwise, use kCpuExecutionProvider by default.
template <typename T>
std::enable_if_t<HasProvider<T>::value, const char*> GetProvider(T* op) {
return op->GetExecutionProviderType();
}
template <typename T>
std::enable_if_t<!HasProvider<T>::value, const char*> GetProvider(T*) {
return "CPUExecutionProvider";
}
} // namespace CustomOpImpl
template <typename TOp, typename TKernel>
struct CustomOpBase : OrtCustomOp {
CustomOpBase() {
@ -312,8 +285,7 @@ struct CustomOpBase : OrtCustomOp {
OrtCustomOp::CreateKernel = [](OrtCustomOp* this_, const OrtCustomOpApi* api, const OrtKernelInfo* info) { return static_cast<TOp*>(this_)->CreateKernel(*api, info); };
OrtCustomOp::GetName = [](OrtCustomOp* this_) { return static_cast<TOp*>(this_)->GetName(); };
// If OrtCustomOp does not have a definition of GetExecutionProviderType, use CPUExecutorProvider by default
OrtCustomOp::GetExecutionProviderType = [](OrtCustomOp* this_) { return CustomOpImpl::GetProvider<TOp>(static_cast<TOp*>(this_)); };
OrtCustomOp::GetExecutionProviderType = [](OrtCustomOp* this_) { return static_cast<TOp*>(this_)->GetExecutionProviderType(); };
OrtCustomOp::GetInputTypeCount = [](OrtCustomOp* this_) { return static_cast<TOp*>(this_)->GetInputTypeCount(); };
OrtCustomOp::GetInputType = [](OrtCustomOp* this_, size_t index) { return static_cast<TOp*>(this_)->GetInputType(index); };
@ -324,6 +296,9 @@ struct CustomOpBase : OrtCustomOp {
OrtCustomOp::KernelCompute = [](void* op_kernel, OrtKernelContext* context) { static_cast<TKernel*>(op_kernel)->Compute(context); };
OrtCustomOp::KernelDestroy = [](void* op_kernel) { delete static_cast<TKernel*>(op_kernel); };
}
// Default implementation of GetExecutionProviderType that returns nullptr to default to the CPU provider
const char* GetExecutionProviderType() const { return nullptr; }
};
} // namespace Ort

7
onnxruntime/core/session/custom_ops.cc
View file

@ -135,11 +135,10 @@ common::Status CreateCustomRegistry(const std::vector<OrtCustomOpDomain*>& op_do
def_builder.SetName(op->GetName(op))
.SetDomain(domain->domain_)
.SinceVersion(1);
if (op->GetExecutionProviderType) {
def_builder.Provider(op->GetExecutionProviderType(op));
} else {
if (const char* provider_type = op->GetExecutionProviderType(op))
def_builder.Provider(provider_type);
else
def_builder.Provider(onnxruntime::kCpuExecutionProvider);
}
KernelCreateFn kernel_create_fn = [&op](const OpKernelInfo& info) -> OpKernel* { return new CustomOpKernel(info, *op); };
KernelCreateInfo create_info(def_builder.Build(), kernel_create_fn);

12
onnxruntime/test/shared_lib/test_inference.cc
View file

@ -223,10 +223,14 @@ TEST_F(CApiTest, test_pyop) {
std::cout << "Test model with pyop" << std::endl;
std::ofstream module("mymodule.py");
module << "class MyKernel:" << std::endl;
module << "\t" << "def __init__(self,A,B,C):" << std::endl;
module << "\t\t" << "self.a,self.b,self.c = A,B,C" << std::endl;
module << "\t" << "def compute(self,x):" << std::endl;
module << "\t\t" << "return x*2" << std::endl;
module << "\t"
<< "def __init__(self,A,B,C):" << std::endl;
module << "\t\t"
<< "self.a,self.b,self.c = A,B,C" << std::endl;
module << "\t"
<< "def compute(self,x):" << std::endl;
module << "\t\t"
<< "return x*2" << std::endl;
module.close();
std::vector<Input> inputs(1);
Input& input = inputs[0];