From 5637a8606bc3caeec3c590350de770c7fcec8dd7 Mon Sep 17 00:00:00 2001
From: Jerry Ge <jerry.ge@arm.com>
Date: Mon, 30 Oct 2023 10:18:45 -0700
Subject: Support loading shared libraries for custom operators

- Add a new command line option to allow users to specify a custom
  defined dll library
- Add a custom registry to store all registered libraries
- Add a dummy example (custom_op_example.cpp) for demonstrating this new
  feature

Signed-off-by: Jerry Ge <jerry.ge@arm.com>
Change-Id: I7c360835933f77e33fcbd772cabfe01d82282d47
---
 reference_model/src/main.cpp | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

(limited to 'reference_model/src/main.cpp')

diff --git a/reference_model/src/main.cpp b/reference_model/src/main.cpp
index 80125ee..24784b5 100644
--- a/reference_model/src/main.cpp
+++ b/reference_model/src/main.cpp
@@ -18,6 +18,8 @@
 
 #include "arith_util.h"
 #include "command_line_utils.h"
+#include "custom_op_interface.h"
+#include "custom_registry.h"
 #include "ops/op_factory.h"
 #include "subgraph_traverser.h"
 #include "tosa_serialization_handler.h"
@@ -38,6 +40,7 @@ int readInputTensors(SubgraphTraverser& gt, json& test_desc);
 int writeFinalTensors(SubgraphTraverser& gt, json& test_desc, const std::string& filename_prefix);
 int readVariableTensors(SubgraphTraverser& gt, json test_desc);
 int writeVariableTensors(SubgraphTraverser& gt, json test_desc);
+int loadSharedLibs(std::string& custom_op_lib_path);
 int loadGraph(TosaSerializationHandler& tsh, json& test_desc);
 void parse_value(const std::string& text, tosa_level_t& value);
 const std::string getResultFilenamePrefix();
@@ -83,6 +86,15 @@ int main(int argc, char** argv)
         FATAL_ERROR("Unable to load test json");
     }
 
+    // load shared libs if specified
+    if (g_func_config.custom_op_lib_path != "")
+    {
+        if (loadSharedLibs(g_func_config.custom_op_lib_path))
+        {
+            FATAL_ERROR("Shared library specified but not loaded successfully");
+        }
+    }
+
     if (loadGraph(tsh, test_desc))
     {
         FATAL_ERROR("Unable to load graph");
@@ -236,6 +248,25 @@ int main(int argc, char** argv)
     return (int)status;
 }
 
+int loadSharedLibs(std::string& custom_op_lib_path)
+{
+    // Load the shared_lib
+    void* lib_handle = dlopen(custom_op_lib_path.c_str(), RTLD_LAZY);
+    if (lib_handle == nullptr)
+    {
+        FATAL_ERROR("Library %s does not exist\n", custom_op_lib_path.c_str());
+    }
+
+    typedef int (*get_customOp_function_t)(registration_callback_t registration_func);
+    auto get_customOp_creation_funcs = (get_customOp_function_t)dlsym(lib_handle, "getCustomOpCreationFuncs");
+    if (get_customOp_creation_funcs == nullptr)
+    {
+        FATAL_ERROR("Can't find the getCustomOpCreationFuncs \n");
+    }
+
+    return get_customOp_creation_funcs(&MasterRegistry::register_function);
+}
+
 int loadGraph(TosaSerializationHandler& tsh, json& test_desc)
 {
     char graph_fullname[1024];
-- 
cgit v1.2.1