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/ops/custom.cc     | 30 ++++++++++++++++++++++++------
 reference_model/src/ops/custom.h      | 20 +++++++++++++-------
 reference_model/src/ops/op_factory.cc |  2 +-
 3 files changed, 38 insertions(+), 14 deletions(-)

(limited to 'reference_model/src/ops')

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:
-- 
cgit v1.2.1