diff --git a/apps/dso_plugin_module/.gitignore b/apps/dso_plugin_module/.gitignore new file mode 100644 index 0000000000000..a65b41774ad52 --- /dev/null +++ b/apps/dso_plugin_module/.gitignore @@ -0,0 +1 @@ +lib diff --git a/apps/dso_plugin_module/Makefile b/apps/dso_plugin_module/Makefile new file mode 100644 index 0000000000000..2ee6189e2876b --- /dev/null +++ b/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) diff --git a/apps/dso_plugin_module/README.md b/apps/dso_plugin_module/README.md new file mode 100644 index 0000000000000..c8803b106521f --- /dev/null +++ b/apps/dso_plugin_module/README.md @@ -0,0 +1,41 @@ + + + + + + + + + + + + + + + + + + +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. diff --git a/apps/dso_plugin_module/plugin_module.cc b/apps/dso_plugin_module/plugin_module.cc new file mode 100644 index 0000000000000..7c3c5accf1ec7 --- /dev/null +++ b/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 +#include +#include +#include + +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& sptr_to_self) final { + if (name == "add") { + return TypedPackedFunc([sptr_to_self, this](int value) { + return value_ + value; + }); + } else if (name == "mul") { + return TypedPackedFunc([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(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 diff --git a/apps/dso_plugin_module/test_plugin_module.py b/apps/dso_plugin_module/test_plugin_module.py new file mode 100644 index 0000000000000..d02ea07b10230 --- /dev/null +++ b/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() diff --git a/apps/extension/Makefile b/apps/extension/Makefile index 1680a003e06f2..e178b661f4038 100644 --- a/apps/extension/Makefile +++ b/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) diff --git a/include/tvm/ir_pass.h b/include/tvm/ir_pass.h index 6e1fed5a85427..5a81d590f7d78 100644 --- a/include/tvm/ir_pass.h +++ b/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 * diff --git a/include/tvm/runtime/c_backend_api.h b/include/tvm/runtime/c_backend_api.h index 05f774403fd0e..ffd13eccdac51 100644 --- a/include/tvm/runtime/c_backend_api.h +++ b/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). diff --git a/include/tvm/runtime/module.h b/include/tvm/runtime/module.h index d539934c3d3f8..686a6bf4ec0a3 100644 --- a/include/tvm/runtime/module.h +++ b/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() {} diff --git a/include/tvm/runtime/packed_func.h b/include/tvm/runtime/packed_func.h index 5219034acd720..19c9434ed64e5 100644 --- a/include/tvm/runtime/packed_func.h +++ b/include/tvm/runtime/packed_func.h @@ -771,6 +771,23 @@ class TVMRetValue : public TVMPODValue_ { *ret_type_code = type_code_; type_code_ = kNull; } + /*! + * \brief Construct a new TVMRetValue by + * moving from return value stored via C API. + * \param value the value. + * \param type_code The type code. + * \return The created TVMRetValue. + */ + static TVMRetValue MoveFromCHost(TVMValue value, + int type_code) { + // Can move POD and everything under the object system. + CHECK(type_code <= kFuncHandle || + type_code == kNDArrayContainer); + TVMRetValue ret; + ret.value_ = value; + ret.type_code_ = type_code; + return ret; + } /*! \return The value field, if the data is POD */ const TVMValue& value() const { CHECK(type_code_ != kObjectHandle && @@ -877,6 +894,104 @@ class TVMRetValue : public TVMPODValue_ { } }; +/*! + * \brief Export a function with the PackedFunc signature + * as a PackedFunc that can be loaded by LibraryModule. + * + * \param ExportName The symbol name to be exported. + * \param Function The function with PackedFunc signature. + * \sa PackedFunc + * + * \code + * + * void AddOne_(TVMArgs args, TVMRetValue* rv) { + * int value = args[0]; + * *rv = value + 1; + * } + * // Expose the function as "AddOne" + * TVM_DLL_EXPORT_PACKED_FUNC(AddOne, AddOne_); + * + * \endcode + */ +#define TVM_DLL_EXPORT_PACKED_FUNC(ExportName, Function) \ + extern "C" { \ + TVM_DLL int ExportName(TVMValue* args, \ + int* type_code, \ + int num_args, \ + TVMValue* out_value, \ + int* out_type_code) { \ + try { \ + ::tvm::runtime::TVMRetValue rv; \ + Function(::tvm::runtime::TVMArgs( \ + args, type_code, num_args), &rv); \ + rv.MoveToCHost(out_value, out_type_code); \ + return 0; \ + } catch (const ::std::runtime_error& _except_) { \ + TVMAPISetLastError(_except_.what()); \ + return -1; \ + } \ + } \ + } + +/*! + * \brief Export typed function as a PackedFunc + * that can be loaded by LibraryModule. + * + * \param ExportName The symbol name to be exported. + * \param Function The typed function. + * \note ExportName and Function must be different, + * see code examples below. + * + * \sa TypedPackedFunc + * + * \code + * + * int AddOne_(int x) { + * return x + 1; + * } + * + * // Expose the function as "AddOne" + * TVM_DLL_EXPORT_TYPED_FUNC(AddOne, AddOne_); + * + * // Expose the function as "SubOne" + * TVM_DLL_EXPORT_TYPED_FUNC(SubOne, [](int x) { + * return x - 1; + * }); + * + * // The following code will cause compilation error. + * // Because the same Function and ExortName + * // TVM_DLL_EXPORT_TYPED_FUNC(AddOne_, AddOne_); + * + * // The following code is OK, assuming the macro + * // is in a different namespace from xyz + * // TVM_DLL_EXPORT_TYPED_FUNC(AddOne_, xyz::AddOne_); + * + * \endcode + */ +#define TVM_DLL_EXPORT_TYPED_FUNC(ExportName, Function) \ + extern "C" { \ + TVM_DLL int ExportName(TVMValue* args, \ + int* type_code, \ + int num_args, \ + TVMValue* out_value, \ + int* out_type_code) { \ + try { \ + auto f = Function; \ + using FType = ::tvm::runtime::detail:: \ + function_signature::FType; \ + ::tvm::runtime::TVMRetValue rv; \ + ::tvm::runtime::detail::unpack_call_by_signature::run( \ + f, \ + ::tvm::runtime::TVMArgs(args, type_code, num_args), &rv); \ + rv.MoveToCHost(out_value, out_type_code); \ + return 0; \ + } catch (const ::std::runtime_error& _except_) { \ + TVMAPISetLastError(_except_.what()); \ + return -1; \ + } \ + } \ + } + // implementation details inline const char* TypeCode2Str(int type_code) { switch (type_code) { @@ -1218,6 +1333,20 @@ inline void unpack_call(const F& f, const TVMArgs& args, TVMRetValue* rv) { unpack_call_dispatcher::run(f, args, rv); } +template +struct unpack_call_by_signature { +}; + +template +struct unpack_call_by_signature { + template + static void run(const F& f, + const TVMArgs& args, + TVMRetValue* rv) { + unpack_call(f, args, rv); + } +}; + template inline R call_packed(const PackedFunc& pf, Args&& ...args) { return R(pf(std::forward(args)...)); diff --git a/src/codegen/codegen_source_base.cc b/src/codegen/codegen_source_base.cc index 7c4ed5b91c8b7..2b11d11281f54 100644 --- a/src/codegen/codegen_source_base.cc +++ b/src/codegen/codegen_source_base.cc @@ -6,9 +6,9 @@ * 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 diff --git a/src/codegen/llvm/llvm_module.cc b/src/codegen/llvm/llvm_module.cc index 937a1103eda47..933e78fc4f416 100644 --- a/src/codegen/llvm/llvm_module.cc +++ b/src/codegen/llvm/llvm_module.cc @@ -68,8 +68,8 @@ class LLVMModuleNode final : public runtime::ModuleNode { const std::string& fname = (name == runtime::symbol::tvm_module_main ? entry_func_ : name); - BackendPackedCFunc faddr = - reinterpret_cast(GetFunctionAddr(fname)); + TVMBackendPackedCFunc faddr = + reinterpret_cast(GetFunctionAddr(fname)); if (faddr == nullptr) return PackedFunc(); return WrapPackedFunc(faddr, sptr_to_self); } diff --git a/src/pass/make_api.cc b/src/pass/make_api.cc index f065502db6b4d..03a6035864aca 100644 --- a/src/pass/make_api.cc +++ b/src/pass/make_api.cc @@ -53,6 +53,8 @@ LoweredFunc MakeAPI(Stmt body, Var v_packed_args("args", DataType::Handle()); Var v_packed_arg_type_ids("arg_type_ids", DataType::Handle()); Var v_num_packed_args("num_args", DataType::Int(32)); + Var v_out_ret_value("out_ret_value", DataType::Handle()); + Var v_out_ret_tcode("out_ret_tcode", DataType::Handle()); // The arguments of the function. Array args; // The device context @@ -151,6 +153,15 @@ LoweredFunc MakeAPI(Stmt body, } } + // allow return value if the function is packed. + if (num_packed_args != 0) { + args.push_back(v_out_ret_value); + args.push_back(v_out_ret_tcode); + } + + size_t expected_nargs = num_unpacked_args + (num_packed_args != 0 ? 5 : 0); + CHECK_EQ(args.size(), expected_nargs); + // Arg definitions are defined before buffer binding to avoid the use before // def errors. // diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc index 9aaf5b9ad390c..2fabe1bec9525 100644 --- a/src/runtime/library_module.cc +++ b/src/runtime/library_module.cc @@ -28,7 +28,7 @@ #include #include #include -#include +#include #include "library_module.h" namespace tvm { @@ -48,15 +48,15 @@ class LibraryModuleNode final : public ModuleNode { PackedFunc GetFunction( const std::string& name, const ObjectPtr& sptr_to_self) final { - BackendPackedCFunc faddr; + TVMBackendPackedCFunc faddr; if (name == runtime::symbol::tvm_module_main) { const char* entry_name = reinterpret_cast( lib_->GetSymbol(runtime::symbol::tvm_module_main)); CHECK(entry_name!= nullptr) << "Symbol " << runtime::symbol::tvm_module_main << " is not presented"; - faddr = reinterpret_cast(lib_->GetSymbol(entry_name)); + faddr = reinterpret_cast(lib_->GetSymbol(entry_name)); } else { - faddr = reinterpret_cast(lib_->GetSymbol(name.c_str())); + faddr = reinterpret_cast(lib_->GetSymbol(name.c_str())); } if (faddr == nullptr) return PackedFunc(); return WrapPackedFunc(faddr, sptr_to_self); @@ -77,14 +77,21 @@ class ModuleInternal { } }; -PackedFunc WrapPackedFunc(BackendPackedCFunc faddr, +PackedFunc WrapPackedFunc(TVMBackendPackedCFunc faddr, const ObjectPtr& sptr_to_self) { return PackedFunc([faddr, sptr_to_self](TVMArgs args, TVMRetValue* rv) { + TVMValue ret_value; + int ret_type_code = kNull; int ret = (*faddr)( const_cast(args.values), const_cast(args.type_codes), - args.num_args); + args.num_args, + &ret_value, + &ret_type_code); CHECK_EQ(ret, 0) << TVMGetLastError(); + if (ret_type_code != kNull) { + *rv = TVMRetValue::MoveFromCHost(ret_value, ret_type_code); + } }); } diff --git a/src/runtime/library_module.h b/src/runtime/library_module.h index e5f5ad94893bb..8e84a4b48bd3c 100644 --- a/src/runtime/library_module.h +++ b/src/runtime/library_module.h @@ -29,13 +29,6 @@ #include #include -extern "C" { -// Function signature for generated packed function in shared library -typedef int (*BackendPackedCFunc)(void* args, - int* type_codes, - int num_args); -} // extern "C" - namespace tvm { namespace runtime { /*! @@ -58,11 +51,11 @@ class Library : public Object { }; /*! - * \brief Wrap a BackendPackedCFunc to packed function. + * \brief Wrap a TVMBackendPackedCFunc to packed function. * \param faddr The function address * \param mptr The module pointer node. */ -PackedFunc WrapPackedFunc(BackendPackedCFunc faddr, const ObjectPtr& mptr); +PackedFunc WrapPackedFunc(TVMBackendPackedCFunc faddr, const ObjectPtr& mptr); /*! * \brief Utility to initialize conext function symbols during startup diff --git a/src/runtime/runtime_base.h b/src/runtime/runtime_base.h index b490dbe113b8d..84fc3c462c3d6 100644 --- a/src/runtime/runtime_base.h +++ b/src/runtime/runtime_base.h @@ -6,9 +6,9 @@ * 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 diff --git a/tests/python/unittest/test_pass_makeapi.py b/tests/python/unittest/test_pass_makeapi.py index c506ccb799b39..77a97d8bffa85 100644 --- a/tests/python/unittest/test_pass_makeapi.py +++ b/tests/python/unittest/test_pass_makeapi.py @@ -37,7 +37,7 @@ def test_makeapi(): f = tvm.ir_pass.MakeAPI( stmt, "myadd", [n, Ab, Bb, Cb], num_unpacked_args, True) assert(f.handle_data_type[Ab.data].dtype == Ab.dtype) - assert(len(f.args) == 5) + assert(len(f.args) == 7) output_ssa = False diff --git a/tests/scripts/task_python_integration.sh b/tests/scripts/task_python_integration.sh index 491ce8b2bef2e..29ffb5f5f92a5 100755 --- a/tests/scripts/task_python_integration.sh +++ b/tests/scripts/task_python_integration.sh @@ -39,6 +39,15 @@ cd ../.. TVM_FFI=cython python3 -m pytest -v apps/extension/tests TVM_FFI=ctypes python3 -m pytest -v apps/extension/tests +# Test dso plugin +cd apps/dso_plugin_module +rm -rf lib +make +cd ../.. +TVM_FFI=cython python3 -m pytest -v apps/dso_plugin_module +TVM_FFI=ctypes python3 -m pytest -v apps/dso_plugin_module + + TVM_FFI=ctypes python3 -m pytest -v tests/python/integration TVM_FFI=ctypes python3 -m pytest -v tests/python/contrib