From 4b2942365628d8f2926a906e89768d95ce257fd1 Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Mon, 19 Oct 2020 17:14:31 -0700 Subject: [PATCH] Re-enable custom op shared library test for debug builds (#5475) --- onnxruntime/test/shared_lib/test_inference.cc | 28 +++++++++++++---- .../custom_op_library/custom_op_library.cc | 31 ++++++++++++++++--- 2 files changed, 48 insertions(+), 11 deletions(-) diff --git a/onnxruntime/test/shared_lib/test_inference.cc b/onnxruntime/test/shared_lib/test_inference.cc index d1fae73651..9b99f3b448 100644 --- a/onnxruntime/test/shared_lib/test_inference.cc +++ b/onnxruntime/test/shared_lib/test_inference.cc @@ -17,6 +17,12 @@ #include "test_allocator.h" #include "test_fixture.h" +#ifdef _WIN32 +#include +#else +#include +#endif + struct Input { const char* name = nullptr; std::vector dims; @@ -68,6 +74,7 @@ void TestInference(Ort::Env& env, T model_uri, int provider_type, OrtCustomOpDomain* custom_op_domain_ptr, const char* custom_op_library_filename, + void** library_handle = nullptr, bool test_session_creation_only = false) { Ort::SessionOptions session_options; @@ -100,8 +107,7 @@ void TestInference(Ort::Env& env, T model_uri, } if (custom_op_library_filename) { - void* library_handle = nullptr; // leak this, no harm. - Ort::ThrowOnError(Ort::GetApi().RegisterCustomOpsLibrary((OrtSessionOptions*)session_options, custom_op_library_filename, &library_handle)); + Ort::ThrowOnError(Ort::GetApi().RegisterCustomOpsLibrary((OrtSessionOptions*)session_options, custom_op_library_filename, library_handle)); } // if session creation passes, model loads fine @@ -289,7 +295,7 @@ TEST(CApiTest, custom_op_handler) { // It is enough to test for successful session creation because if the custom node wasn't assigned an EP, // the session creation would fail. Since the custom node is only tied to the CUDA EP (in CUDA-enabled builds), // if the session creation succeeds, it is assumed that the node got assigned to the CUDA EP. - TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 1, custom_op_domain, nullptr, true); + TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 1, custom_op_domain, nullptr, nullptr, true); #else TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, expected_values_y, 0, custom_op_domain, nullptr); #endif @@ -317,12 +323,12 @@ TEST(CApiTest, RegisterCustomOpForCPUAndCUDA) { custom_op_domain.Add(&custom_op_cuda); TestInference(*ort_env, CUSTOM_OP_MODEL_URI, inputs, "Y", expected_dims_y, - expected_values_y, 1, custom_op_domain, nullptr, true); + expected_values_y, 1, custom_op_domain, nullptr, nullptr, true); } #endif //It has memory leak. The OrtCustomOpDomain created in custom_op_library.cc:RegisterCustomOps function was not freed -#if defined(__ANDROID__) || defined(ONNXRUNTIME_ENABLE_MEMLEAK_CHECK) +#if defined(__ANDROID__) TEST(CApiTest, DISABLED_test_custom_op_library) { #else TEST(CApiTest, test_custom_op_library) { @@ -357,7 +363,17 @@ TEST(CApiTest, test_custom_op_library) { lib_name = "./libcustom_op_library.so"; #endif - TestInference(*ort_env, CUSTOM_OP_LIBRARY_TEST_MODEL_URI, inputs, "output", expected_dims_y, expected_values_y, 0, nullptr, lib_name.c_str()); + void* library_handle = nullptr; + TestInference(*ort_env, CUSTOM_OP_LIBRARY_TEST_MODEL_URI, inputs, "output", expected_dims_y, + expected_values_y, 0, nullptr, lib_name.c_str(), &library_handle); + +#ifdef _WIN32 + bool success = ::FreeLibrary(reinterpret_cast(library_handle)); + ORT_ENFORCE(success, "Error while closing custom op shared library"); +#else + int retval = dlclose(library_handle); + ORT_ENFORCE(retval == 0, "Error while closing custom op shared library"); +#endif } #if defined(ENABLE_LANGUAGE_INTEROP_OPS) diff --git a/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc b/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc index 055acda9be..c408733a39 100644 --- a/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc +++ b/onnxruntime/test/testdata/custom_op_library/custom_op_library.cc @@ -6,9 +6,31 @@ #include #include +#include static const char* c_OpDomain = "test.customop"; +struct OrtCustomOpDomainDeleter { + explicit OrtCustomOpDomainDeleter(const OrtApi* ort_api) { + ort_api_ = ort_api; + } + void operator()(OrtCustomOpDomain* domain) const { + ort_api_->ReleaseCustomOpDomain(domain); + } + + const OrtApi* ort_api_; +}; + +using OrtCustomOpDomainUniquePtr = std::unique_ptr; +static std::vector ort_custom_op_domain_container; +static std::mutex ort_custom_op_domain_mutex; + +static void AddOrtCustomOpDomainToContainer(OrtCustomOpDomain* domain, const OrtApi* ort_api) { + std::lock_guard lock(ort_custom_op_domain_mutex); + auto ptr = std::unique_ptr(domain, OrtCustomOpDomainDeleter(ort_api)); + ort_custom_op_domain_container.push_back(std::move(ptr)); +} + struct OrtTensorDimensions : std::vector { OrtTensorDimensions(Ort::CustomOpApi ort, const OrtValue* value) { OrtTensorTypeAndShapeInfo* info = ort.GetTensorTypeAndShape(value); @@ -17,12 +39,10 @@ struct OrtTensorDimensions : std::vector { } }; - struct KernelOne { KernelOne(OrtApi api) - :api_(api), - ort_(api_) - { + : api_(api), + ort_(api_) { } void Compute(OrtKernelContext* context) { @@ -115,7 +135,6 @@ struct CustomOpTwo : Ort::CustomOpBase { } c_CustomOpTwo; - OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtApiBase* api) { OrtCustomOpDomain* domain = nullptr; const OrtApi* ortApi = api->GetApi(ORT_API_VERSION); @@ -124,6 +143,8 @@ OrtStatus* ORT_API_CALL RegisterCustomOps(OrtSessionOptions* options, const OrtA return status; } + AddOrtCustomOpDomainToContainer(domain, ortApi); + if (auto status = ortApi->CustomOpDomain_Add(domain, &c_CustomOpOne)) { return status; }