[RUNTIME][DSO] Improve TVMBackendPackedCFunc to allow return val

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 unloadd.
+    #
+    # 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

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