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

5 files changed, 70 insertions, 14 deletions

diff --git a/reference_model/src/command_line_utils.h b/reference_model/src/command_line_utils.h
index f8031d9..c1cc54c 100644
--- a/reference_model/src/command_line_utils.h
+++ b/reference_model/src/command_line_utils.h
@@ -70,6 +70,7 @@ int func_model_parse_cmd_line(
         ("l,loglevel", func_debug.get_debug_verbosity_help_string(), cxxopts::value<std::string>())
         ("o,logfile", "output log file", cxxopts::value<std::string>())
         ("d,debugmask", func_debug.get_debug_mask_help_string(), cxxopts::value<std::vector<std::string>>())
+        ("custom_op_lib_path", "Path to the shared lib for customOp evaluation", cxxopts::value<std::string>(func_config.custom_op_lib_path))
         ("h,help", "print help");
         // clang-format on
 
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];
diff --git a/reference_model/src/ops/custom.cc b/reference_model/src/ops/custom.cc
index cbc5742..39a6f87 100644
--- a/reference_model/src/ops/custom.cc
+++ b/reference_model/src/ops/custom.cc
@@ -1,5 +1,5 @@
 
-// Copyright (c) 2020, ARM Limited.
+// Copyright (c) 2020, 2023, ARM Limited.
 //
 //    Licensed under the Apache License, Version 2.0 (the "License");
 //    you may not use this file except in compliance with the License.
@@ -14,27 +14,45 @@
 //    limitations under the License.
 
 #include "custom.h"
+#include "attribute.h"
+
+#include "tensor.h"
 
 using namespace TosaReference;
 using namespace Eigen;
 using namespace tosa;
 
-OpCustom::OpCustom(SubgraphTraverser* sgt_, uint64_t id_)
+OpCustom::OpCustom(SubgraphTraverser* sgt_, TosaAttributeBase* attribute_, uint64_t id_)
     : GraphNode(sgt_, Op_CUSTOM, id_)
-{}
+{
+    // Init Attribute
+    if (auto p = dynamic_cast<TosaCustomAttribute*>(attribute_))
+        attribute = new TosaCustomAttribute(p);
+}
 
 OpCustom::~OpCustom()
 {}
 
 int OpCustom::checkTensorAttributes()
 {
+    // Get the pointer to customOp library
+    auto domain_name_vec   = attribute->domain_name();
+    auto operator_name_vec = attribute->operator_name();
+    std::string domain_name(domain_name_vec.begin(), domain_name_vec.end());
+    std::string operator_name(operator_name_vec.begin(), operator_name_vec.end());
+
+    auto getCustomNodeFunc = MasterRegistry::get_op(domain_name, operator_name);
+    ERROR_IF(getCustomNodeFunc == nullptr, "Can't find the custom shared library: %s::%s is not registered.",
+             domain_name.c_str(), operator_name.c_str());
+    this->custom_op_ptr = getCustomNodeFunc();
+
     return 0;
 }
 
 int OpCustom::eval()
 {
-    ERROR_IF(true, "not supported yet");
-
-    // Evaluation is trivial for constants
+    auto implementation_attrs_vec = attribute->implementation_attrs();
+    std::string implementation_attrs(implementation_attrs_vec.begin(), implementation_attrs_vec.end());
+    custom_op_ptr->eval(getInputs(), getOutputs(), implementation_attrs);
     return GraphNode::eval();
 }
diff --git a/reference_model/src/ops/custom.h b/reference_model/src/ops/custom.h
index d14c809..186d2c1 100644
--- a/reference_model/src/ops/custom.h
+++ b/reference_model/src/ops/custom.h
@@ -1,5 +1,5 @@
 
-// Copyright (c) 2020, ARM Limited.
+// Copyright (c) 2020, 2023 ARM Limited.
 //
 //    Licensed under the Apache License, Version 2.0 (the "License");
 //    you may not use this file except in compliance with the License.
@@ -16,23 +16,29 @@
 #ifndef OPS_CUSTOM_H
 #define OPS_CUSTOM_H
 
+#include "attribute.h"
+#include "custom_registry.h"
 #include "graph_node.h"
 
 using namespace tosa;
 
 namespace TosaReference
 {
-
 class OpCustom : public GraphNode
 {
 public:
-    OpCustom(SubgraphTraverser* sgt_, uint64_t id_);
-    virtual ~OpCustom();
+    OpCustom(SubgraphTraverser* sgt_, TosaAttributeBase* attribute_, uint64_t id_);
+    ~OpCustom();
 
-    virtual int checkTensorAttributes();
-    virtual int eval();
-};
+    int checkTensorAttributes();
+    int eval();
 
+protected:
+    TosaCustomAttribute* attribute;
+
+private:
+    CustomOpInterface* custom_op_ptr;
+};
 };    // namespace TosaReference
 
 #endif
diff --git a/reference_model/src/ops/op_factory.cc b/reference_model/src/ops/op_factory.cc
index d834b74..34db903 100644
--- a/reference_model/src/ops/op_factory.cc
+++ b/reference_model/src/ops/op_factory.cc
@@ -592,7 +592,7 @@ GraphNode* OpFactory::newOp(SubgraphTraverser* sgt,
 
         // custom
         case Op_CUSTOM:
-            return new OpCustom(sgt, id);
+            return new OpCustom(sgt, attribute, id);
 
         // control_flow
         case Op_COND_IF: