Skip to content

Commit

Permalink
[RUNTIME][DSO] Improve TVMBackendPackedCFunc to allow return value.
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
tqchen committed Jan 6, 2020
1 parent dc30880 commit f00cdf0
Show file tree
Hide file tree
Showing 16 changed files with 392 additions and 23 deletions.
1 change: 1 addition & 0 deletions apps/dso_plugin_module/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
lib
33 changes: 33 additions & 0 deletions apps/dso_plugin_module/Makefile
Original file line number Diff line number Diff line change
@@ -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)
41 changes: 41 additions & 0 deletions apps/dso_plugin_module/README.md
Original file line number Diff line number Diff line change
@@ -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.
82 changes: 82 additions & 0 deletions apps/dso_plugin_module/plugin_module.cc
Original file line number Diff line number Diff line change
@@ -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
47 changes: 47 additions & 0 deletions apps/dso_plugin_module/test_plugin_module.py
Original file line number Diff line number Diff line change
@@ -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()
1 change: 1 addition & 0 deletions apps/extension/Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
3 changes: 2 additions & 1 deletion include/tvm/ir_pass.h
Original file line number Diff line number Diff line change
Expand Up @@ -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
*
Expand Down
18 changes: 17 additions & 1 deletion include/tvm/runtime/c_backend_api.h
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
Loading

0 comments on commit f00cdf0

Please sign in to comment.