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

12 files changed, 335 insertions, 14 deletions

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6556b27..117bc9f 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -20,3 +20,5 @@ add_subdirectory(thirdparty)
 if(TOSA_TOOLS_BUILD_REFERENCE_MODEL)
     add_subdirectory(reference_model)
 endif()
+
+add_subdirectory(reference_model/custom_op_example)
diff --git a/reference_model/CMakeLists.txt b/reference_model/CMakeLists.txt
index d36167a..5c40ce5 100644
--- a/reference_model/CMakeLists.txt
+++ b/reference_model/CMakeLists.txt
@@ -131,6 +131,8 @@ target_link_libraries(tosa_reference_model_lib
 
 set(PUBLIC_HEADERS)
 list(APPEND PUBLIC_HEADERS
+    include/custom_op_interface.h
+    include/custom_registry.h
     include/debug_modes.def
     include/debug_types.h
     include/dtype.h
@@ -202,6 +204,7 @@ if(BUILD_TOSA_REFERENCE_MODEL_EXECUTABLE)
       ${SERIALIZATION_LIB}
       nlohmann_json::nlohmann_json
       cxxopts
+      ${CMAKE_DL_LIBS}
   )
 
   install(TARGETS tosa_reference_model DESTINATION bin)
diff --git a/reference_model/custom_op_example/CMakeLists.txt b/reference_model/custom_op_example/CMakeLists.txt
new file mode 100644
index 0000000..aea8071
--- /dev/null
+++ b/reference_model/custom_op_example/CMakeLists.txt
@@ -0,0 +1,39 @@
+cmake_minimum_required (VERSION 3.4)
+
+# Copyright (c) 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.
+#    You may obtain a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+set (CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
+
+add_library(tosa_example_plugin SHARED
+    custom_op_example.cpp
+)
+
+target_include_directories(tosa_example_plugin
+    PUBLIC
+        ${PUBLIC_INCLUDE_DIRS}
+        "../../thirdparty/serialization_lib/include"
+        "../../thirdparty/serialization_lib/third_party/flatbuffers/include"
+        "../../thirdparty/serialization_lib/third_party/half/include"
+        "../../thirdparty/eigen/"
+        "../../thirdparty/eigen/unsupported"
+        "../include"
+        "../src"
+        "../src/ops"
+    PRIVATE
+        ${PRIVATE_INCLUDE_DIRS}
+)
\ No newline at end of file
diff --git a/reference_model/custom_op_example/custom_op_example.cpp b/reference_model/custom_op_example/custom_op_example.cpp
new file mode 100644
index 0000000..27f30d4
--- /dev/null
+++ b/reference_model/custom_op_example/custom_op_example.cpp
@@ -0,0 +1,94 @@
+// Copyright (c) 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.
+//    You may obtain a copy of the License at
+//
+//         http://www.apache.org/licenses/LICENSE-2.0
+//
+//    Unless required by applicable law or agreed to in writing, software
+//    distributed under the License is distributed on an "AS IS" BASIS,
+//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//    See the License for the specific language governing permissions and
+//    limitations under the License.
+
+#include "custom_op_interface.h"
+#include "custom_registry.h"
+#include <vector>
+
+using namespace tosa;
+
+namespace TosaReference
+{
+class CustomOpExample : public CustomOpInterface
+{
+public:
+    CustomOpExample() = default;
+    CustomOpExample(std::string& domain_name, std::string& operator_name, std::string& version)
+        : _domain_name(domain_name)
+        , _operator_name(operator_name)
+        , _version(version)
+    {}
+    int eval(std::vector<TosaReference::Tensor*>& input_tensors,
+             std::vector<TosaReference::Tensor*>& output_tensors,
+             const std::string& implementation_attrs) override
+    {
+        auto input_tensor_ptr  = input_tensors[0];
+        auto output_tensor_ptr = output_tensors[0];
+
+        // down_cast to EigenTensor
+        using TIn  = Eigen::Tensor<float, 1>;
+        using TOut = Eigen::Tensor<float, 1>;
+
+        auto eigenInputTensor  = reinterpret_cast<TosaReference::TensorTemplate<TIn>*>(input_tensor_ptr);
+        auto eigenOutputTensor = reinterpret_cast<TosaReference::TensorTemplate<TIn>*>(output_tensor_ptr);
+
+        // Assign the input to output as an example
+        // This is plug-in implementation specific
+        auto fcn                       = [](float a) -> float { return a; };
+        eigenOutputTensor->getTensor() = eigenInputTensor->getTensor().unaryExpr(fcn);
+
+        return 0;
+    };
+
+    std::string getDomainName() const override
+    {
+        return this->_domain_name;
+    }
+
+    std::string getOperatorName() const override
+    {
+        return this->_operator_name;
+    }
+
+    std::string getVersion() const override
+    {
+        return this->_version;
+    }
+
+    ~CustomOpExample(){};
+
+private:
+    std::string _domain_name;
+    std::string _operator_name;
+    std::string _version;
+};
+
+CustomOpInterface* customOpExample()
+{
+    std::string domain_name         = "ExampleDomain";
+    std::string operator_name       = "ExampleOp";
+    std::string version             = "1.0";
+    CustomOpInterface* customOp_ptr = new CustomOpExample(domain_name, operator_name, version);
+
+    return customOp_ptr;
+}
+
+extern "C" int getCustomOpCreationFuncs(registration_callback_t registration_func)
+{
+    std::string domain_name   = "ExampleDomain";
+    std::string operator_name = "ExampleOp";
+    return registration_func(domain_name, operator_name, &customOpExample);
+}
+
+}    // namespace TosaReference
diff --git a/reference_model/include/custom_op_interface.h b/reference_model/include/custom_op_interface.h
new file mode 100644
index 0000000..aea9086
--- /dev/null
+++ b/reference_model/include/custom_op_interface.h
@@ -0,0 +1,38 @@
+// Copyright (c) 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.
+//    You may obtain a copy of the License at
+//
+//         http://www.apache.org/licenses/LICENSE-2.0
+//
+//    Unless required by applicable law or agreed to in writing, software
+//    distributed under the License is distributed on an "AS IS" BASIS,
+//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//    See the License for the specific language governing permissions and
+//    limitations under the License.
+
+#ifndef CUSTOMOPINTERFACE_H
+#define CUSTOMOPINTERFACE_H
+
+#include "tensor.h"
+#include <vector>
+
+using namespace tosa;
+
+namespace TosaReference
+{
+class CustomOpInterface
+{
+public:
+    CustomOpInterface()                                       = default;
+    virtual std::string getDomainName() const                 = 0;
+    virtual std::string getOperatorName() const               = 0;
+    virtual int eval(std::vector<TosaReference::Tensor*>& input_tensors,
+                     std::vector<TosaReference::Tensor*>& output_tensors,
+                     const std::string& implementation_attrs) = 0;
+    virtual std::string getVersion() const                    = 0;
+};
+}    // namespace TosaReference
+
+#endif
diff --git a/reference_model/include/custom_registry.h b/reference_model/include/custom_registry.h
new file mode 100644
index 0000000..f1a9b8c
--- /dev/null
+++ b/reference_model/include/custom_registry.h
@@ -0,0 +1,88 @@
+// Copyright (c) 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.
+//    You may obtain a copy of the License at
+//
+//         http://www.apache.org/licenses/LICENSE-2.0
+//
+//    Unless required by applicable law or agreed to in writing, software
+//    distributed under the License is distributed on an "AS IS" BASIS,
+//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//    See the License for the specific language governing permissions and
+//    limitations under the License.
+
+#ifndef CUSTOMREGISTRY_H
+#define CUSTOMREGISTRY_H
+
+#include "custom_op_interface.h"
+#include <dlfcn.h>
+#include <unordered_map>
+
+using namespace tosa;
+
+namespace TosaReference
+{
+
+typedef CustomOpInterface* (*op_creation_function_t)();
+typedef int (*registration_callback_t)(const std::string& domain_name,
+                                       const std::string& operator_name,
+                                       const op_creation_function_t& op_creation_function);
+
+class MasterRegistry
+{
+public:
+    static int register_function(const std::string& domain_name,
+                                 const std::string& operator_name,
+                                 const op_creation_function_t& op_creation_function)
+    {
+        std::string unique_id    = domain_name + "::" + operator_name;
+        MasterRegistry& instance = get_instance();
+        if (instance.op_creation_map.find(unique_id) != instance.op_creation_map.end())
+        {
+            std::cout << std::endl;
+            printf("domain_name: %s and operator_name: %s pair has already been registered", domain_name.c_str(),
+                   operator_name.c_str());
+            return 1;
+        }
+        instance.op_creation_map[unique_id] = op_creation_function;
+        return 0;
+    }
+
+    static MasterRegistry& get_instance()
+    {
+        static MasterRegistry instance;
+        return instance;
+    }
+
+    MasterRegistry(const MasterRegistry&) = delete;
+    void operator=(const MasterRegistry&) = delete;
+
+    std::unordered_map<std::string, op_creation_function_t> get_ops() const
+    {
+        return op_creation_map;
+    }
+
+    static op_creation_function_t get_op(const std::string& domain_name, const std::string& operator_name)
+    {
+        std::string unique_id    = domain_name + "::" + operator_name;
+        MasterRegistry& instance = get_instance();
+        auto all_ops_map         = instance.get_ops();
+        if (all_ops_map.find(unique_id) == all_ops_map.end())
+        {
+            return nullptr;
+        }
+        else
+        {
+            op_creation_function_t& op_creation_function = all_ops_map[unique_id];
+            return op_creation_function;
+        }
+    }
+
+private:
+    MasterRegistry() = default;
+    std::unordered_map<std::string, op_creation_function_t> op_creation_map;
+};
+}    // namespace TosaReference
+
+#endif
diff --git a/reference_model/include/func_config.h b/reference_model/include/func_config.h
index 22e7e2c..97afa82 100644
--- a/reference_model/include/func_config.h
+++ b/reference_model/include/func_config.h
@@ -54,6 +54,7 @@ struct func_config_t
     uint32_t dump_intermediates                    = 0;
     uint32_t initialize_variable_tensor_from_numpy = 0;
     std::string fp_format                          = "0.5";
+    std::string custom_op_lib_path                 = "";
     uint32_t precise_mode                          = 0;
     bool abs_mode                                  = 0;        // set in main as second run of precise_mode
     bool float_is_big_endian                       = false;    // Set in arith_util.h by float_is_big_endian()
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: