[RUNTIME][DSO] Improve TVMBackendPackedCFunc to allow return val (apa…

…che#4637)

* [RUNTIME][DSO] Improve TVMBackendPackedCFunc to allow return value.

Previously the signature of LibraryModule's PackedFunc does not support return value.
This wasn't a limitation for our current usecase but could become one
as we start to generate more interesting functions.

This feature also start to get interesting as we move towards unified
object protocol and start to pass object around.
This PR enhances the function signature to allow return values.

We also created two macros TVM_DLL_EXPORT_PACKED_FUNC and TVM_DLL_EXPORT_TYPED_FUNC
to allow manual creation of functions that can be loaded by a LibraryModule.

Examples are added in apps/dso_plugin_module.
The change to TVMBackendPackedCFunc is backward compatible,
as previous function will simply ignore the return value field.

* address review comments

apps/dso_plugin_module/.gitignore

@@ -0,0 +1 @@
+lib

apps/dso_plugin_module/Makefile

@@ -0,0 +1,33 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+
+TVM_ROOT=$(shell cd ../..; pwd)
+PKG_CFLAGS = -std=c++11 -O2 -fPIC\
+	-I${TVM_ROOT}/include\
+	-I${TVM_ROOT}/3rdparty/dmlc-core/include\
+	-I${TVM_ROOT}/3rdparty/dlpack/include
+
+PKG_LDFLAGS =-L${TVM_ROOT}/build
+UNAME_S := $(shell uname -s)
+
+ifeq ($(UNAME_S), Darwin)
+	PKG_LDFLAGS += -undefined dynamic_lookup
+endif
+
+lib/plugin_module.so: plugin_module.cc
+	@mkdir -p $(@D)
+	$(CXX) $(PKG_CFLAGS) -shared -o $@ $^ $(PKG_LDFLAGS)

apps/dso_plugin_module/README.md

@@ -0,0 +1,41 @@
+<!--- Licensed to the Apache Software Foundation (ASF) under one -->
+<!--- or more contributor license agreements.  See the NOTICE file -->
+<!--- distributed with this work for additional information -->
+<!--- regarding copyright ownership.  The ASF licenses this file -->
+<!--- to you 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. -->
+
+
+Example Plugin Module
+=====================
+This folder contains an example that implements a C++ module
+that can be directly loaded as TVM's DSOModule (via tvm.module.load)
+
+## Guideline
+
+When possible, we always recommend exposing
+functions that modifies memory passed by the caller,
+and calls into the runtime API for memory allocations.
+
+## Advanced Usecases
+
+In advanced usecases, we do allow the plugin module to
+create and return managed objects.
+However, there are several restrictions to keep in mind:
+
+- If the module returns an object, we need to make sure
+  that the object get destructed before the module get unloaded.
+  Otherwise segfault can happen because of calling into an unloaded destructor.
+- If the module returns a PackedFunc, then
+  we need to ensure that the libc of the DLL and tvm runtime matches.
+  Otherwise segfault can happen due to incompatibility of std::function.

apps/dso_plugin_module/plugin_module.cc

@@ -0,0 +1,82 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+/*!
+ * \brief Example code that can be compiled and loaded by TVM runtime.
+ * \file plugin_module.cc
+ */
+#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/module.h>
+#include <tvm/runtime/registry.h>
+#include <tvm/runtime/ndarray.h>
+
+namespace tvm_dso_plugin {
+
+using namespace tvm::runtime;
+
+class MyModuleNode : public ModuleNode {
+ public:
+  explicit MyModuleNode(int value)
+      : value_(value) {}
+
+  virtual const char* type_key() const final {
+    return "MyModule";
+  }
+
+  virtual PackedFunc GetFunction(
+      const std::string& name,
+      const ObjectPtr<Object>& sptr_to_self) final {
+    if (name == "add") {
+      return TypedPackedFunc<int(int)>([sptr_to_self, this](int value) {
+          return value_ + value;
+        });
+    } else if (name == "mul") {
+      return TypedPackedFunc<int(int)>([sptr_to_self, this](int value) {
+          return value_ * value;
+        });
+    } else {
+      LOG(FATAL) << "unknown function " << name;
+      return PackedFunc();
+    }
+  }
+
+ private:
+  int value_;
+};
+
+void CreateMyModule_(TVMArgs args, TVMRetValue* rv) {
+  int value = args[0];
+  *rv = Module(make_object<MyModuleNode>(value));
+}
+
+int SubOne_(int x) {
+  return x - 1;
+}
+
+// USE TVM_DLL_EXPORT_TYPED_PACKED_FUNC to export a
+// typed function as packed function.
+TVM_DLL_EXPORT_TYPED_FUNC(SubOne, SubOne_);
+
+// TVM_DLL_EXPORT_TYPED_PACKED_FUNC also works for lambda.
+TVM_DLL_EXPORT_TYPED_FUNC(AddOne, [](int x) -> int {
+  return x + 1;
+});
+
+// Use TVM_EXPORT_PACKED_FUNC to export a function with
+TVM_DLL_EXPORT_PACKED_FUNC(CreateMyModule, tvm_dso_plugin::CreateMyModule_);
+}  // namespace tvm_dso_plugin

apps/dso_plugin_module/test_plugin_module.py

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+import tvm
+import os
+
+def test_plugin_module():
+    curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
+    mod = tvm.module.load(os.path.join(curr_path, "lib", "plugin_module.so"))
+    # NOTE: we need to make sure all managed resources returned
+    # from mod get destructed before mod get unloaded.
+    #
+    # Failure mode we want to prevent from:
+    # We retain an object X whose destructor is within mod.
+    # The program will segfault if X get destructed after mod,
+    # because the destructor function has already been unloaded.
+    #
+    # The easiest way to achieve this is to wrap the
+    # logics related to mod inside a function.
+    def run_module(mod):
+        # normal functions
+        assert mod["AddOne"](10) == 11
+        assert mod["SubOne"](10) == 9
+        # advanced usecase: return a module
+        mymod = mod["CreateMyModule"](10);
+        fadd = mymod["add"]
+        assert fadd(10) == 20
+        assert mymod["mul"](10) == 100
+
+    run_module(mod)
+
+
+if __name__ == "__main__":
+    test_plugin_module()

apps/extension/Makefile

@@ -22,6 +22,7 @@ PKG_CFLAGS = -std=c++11 -O2 -fPIC\
 	-I${TVM_ROOT}/3rdparty/dmlc-core/include\
 	-I${TVM_ROOT}/3rdparty/dlpack/include
 
+
 PKG_LDFLAGS =-L${TVM_ROOT}/build
 UNAME_S := $(shell uname -s)
 

include/tvm/ir_pass.h

@@ -410,7 +410,8 @@ Stmt HoistIfThenElse(Stmt stmt);
  *
  *  if num_packed_args is not zero:
  *       f(TVMArg* packed_args, int* packed_arg_type_ids, int num_packed_args,
- *         api_arg_k, api_arg_k+1, ... api_arg_n)
+ *         api_arg_k, api_arg_k+1, ... api_arg_n,
+ *         TVMValue* out_ret_val, int* out_ret_tcode)
  *
  *       where n == len(api_args), k == num_packed_args
  *

include/tvm/runtime/c_backend_api.h

@@ -34,7 +34,23 @@
 extern "C" {
 #endif
 
-// Backend related functions.
+/*!
+ * \brief Signature for backend functions exported as DLL.
+ *
+ * \param args The arguments
+ * \param type_codes The type codes of the arguments
+ * \param num_args Number of arguments.
+ * \param out_ret_value The output value of the the return value.
+ * \param out_ret_tcode The output type code of the return value.
+ *
+ * \return 0 if success, -1 if failure happens, set error via TVMAPISetLastError.
+ */
+typedef int (*TVMBackendPackedCFunc)(TVMValue* args,
+                                     int* type_codes,
+                                     int num_args,
+                                     TVMValue* out_ret_value,
+                                     int* out_ret_tcode);
+
 /*!
  * \brief Backend function for modules to get function
  *  from its environment mod_node (its imports and global function).

include/tvm/runtime/module.h

@@ -120,7 +120,7 @@ class Module : public ObjectRef {
  *
  * \endcode
  */
-class ModuleNode : public Object {
+class TVM_DLL ModuleNode : public Object {
  public:
   /*! \brief virtual destructor */
   TVM_DLL virtual ~ModuleNode() {}