Skip to content

T
tic
Commit 433756b9 by Tianqi Chen Committed by GitHub
Options

Revert "[RUNTIME] Refactor extension type handling, now it is header only (#924)" (#925) 

This reverts commit 12d15704d7f5d30cff7540f1fd16be64c6baca68.
parent a6b4a219
Showing with 127 additions and 166 deletions
apps/extension/src/tvm_ext.cc
......@@ -22,10 +22,13 @@ struct extension_class_info<tvm_ext::IntVector> {
} // namespace tvm
} // namespace runtime
namespace tvm_ext {
using namespace tvm;
using namespace tvm::runtime;
namespace tvm_ext {
TVM_REGISTER_EXT_TYPE(IntVector);
TVM_REGISTER_GLOBAL("tvm_ext.ivec_create")
.set_body([](TVMArgs args, TVMRetValue *rv) {
......@@ -63,18 +66,3 @@ TVM_REGISTER_GLOBAL("device_api.ext_dev")
*rv = (*tvm::runtime::Registry::Get("device_api.cpu"))();
});
} // namespace tvm_ext
// This callback approach allows extension allows tvm to extract
// This way can be helpful when we want to use a header only
// minimum version of TVM Runtime.
extern "C" int TVMExtDeclare(TVMFunctionHandle pregister) {
const PackedFunc& fregister =
*static_cast<PackedFunc*>(pregister);
auto mul = [](TVMArgs args, TVMRetValue *rv) {
int x = args[0];
int y = args[1];
*rv = x * y;
};
fregister("mul", PackedFunc(mul));
return 0;
}
apps/extension/tests/test_ext.py
......@@ -44,14 +44,8 @@ def test_ext_vec():
tvm.convert(ivec_cb)(ivec)
def test_extract_ext():
fdict = tvm.extract_ext_funcs(tvm_ext._LIB.TVMExtDeclare)
assert fdict["mul"](3, 4) == 12
if __name__ == "__main__":
test_ext_dev()
test_ext_vec()
test_bind_add()
test_sym_add()
test_extract_ext()
include/tvm/runtime/c_runtime_api.h
......@@ -24,13 +24,6 @@
#define TVM_EXTERN_C
#endif
// Macros to do weak linking
#ifdef _MSC_VER
#define TVM_WEAK __declspec(selectany)
#else
#define TVM_WEAK __attribute__((weak))
#endif
#ifdef __EMSCRIPTEN__
#include <emscripten/emscripten.h>
#define TVM_DLL EMSCRIPTEN_KEEPALIVE
......@@ -321,17 +314,6 @@ typedef int (*TVMPackedCFunc)(
typedef void (*TVMPackedCFuncFinalizer)(void* resource_handle);
/*!
* \brief Signature for extension function declarer.
*
* TVM call this function to get the extension functions
* The declarer will call register_func to register function and their name.
*
* \param resource_func_handle The register function
* \return 0 if success, -1 if failure happens
*/
typedef int (*TVMExtensionFuncDeclarer)(TVMFunctionHandle register_func_handle);
/*!
* \brief Wrap a TVMPackedCFunc to become a FunctionHandle.
*
* The resource_handle will be managed by TVM API, until the function is no longer used.
......
include/tvm/runtime/module.h
......@@ -38,14 +38,8 @@ class Module {
* \param query_imports Whether also query dependency modules.
* \return The result function.
* This function will return PackedFunc(nullptr) if function do not exist.
* \note Implemented in packed_func.cc
*/
inline PackedFunc GetFunction(const std::string& name, bool query_imports = false);
/*! \return internal container */
inline ModuleNode* operator->();
/*! \return internal container */
inline const ModuleNode* operator->() const;
// The following functions requires link with runtime.
TVM_DLL PackedFunc GetFunction(const std::string& name, bool query_imports = false);
/*!
* \brief Import another module into this module.
* \param other The module to be imported.
......@@ -63,6 +57,10 @@ class Module {
*/
TVM_DLL static Module LoadFromFile(const std::string& file_name,
const std::string& format = "");
/*! \return internal container */
inline ModuleNode* operator->();
/*! \return internal container */
inline const ModuleNode* operator->() const;
private:
std::shared_ptr<ModuleNode> node_;
......
include/tvm/runtime/packed_func.h
......@@ -183,17 +183,31 @@ struct extension_class_info {
};
/*!
* \brief Capsule structure holding extension types
* Capsule is self-contained and include
* all the information to clone and destroy the type.
* \brief Runtime function table about extension type.
*/
struct TVMExtTypeCapsule {
/*! \brief The pointer to the object */
void* ptr;
class ExtTypeVTable {
public:
/*! \brief function to be called to delete a handle */
void (*destroy)(void* handle);
/*! \brief function to be called when clone a handle */
void* (*clone)(void* handle);
/*!
* \brief Register type
* \tparam T The type to be register.
* \return The registered vtable.
*/
template <typename T>
static inline ExtTypeVTable* Register_();
/*!
* \brief Get a vtable based on type code.
* \param type_code The type code
* \return The registered vtable.
*/
TVM_DLL static ExtTypeVTable* Get(int type_code);
private:
// Internal registration function.
TVM_DLL static ExtTypeVTable* RegisterInternal(int type_code, const ExtTypeVTable& vt);
};
/*!
......@@ -241,9 +255,8 @@ class TVMPODValue_ {
}
template<typename TExtension>
const TExtension& AsExtension() const {
CHECK_EQ(type_code_, extension_class_info<TExtension>::code);
return static_cast<TExtension*>(
static_cast<TVMExtTypeCapsule*>(value_.v_handle)->ptr)[0];
CHECK_LT(type_code_, kExtEnd);
return static_cast<TExtension*>(value_.v_handle)[0];
}
int type_code() const {
return type_code_;
......@@ -475,6 +488,14 @@ class TVMRetValue : public TVMPODValue_ {
this->Assign(other);
return *this;
}
template<typename T,
typename = typename std::enable_if<
extension_class_info<T>::code != 0>::type>
TVMRetValue& operator=(const T& other) {
this->SwitchToClass<T>(
extension_class_info<T>::code, other);
return *this;
}
/*!
* \brief Move the value back to front-end via C API.
* This marks the current container as null.
......@@ -500,11 +521,6 @@ class TVMRetValue : public TVMPODValue_ {
type_code_ != kStr) << "TVMRetValue.value can only be used for POD data";
return value_;
}
// assign extension
template<typename T,
typename = typename std::enable_if<
extension_class_info<T>::code != 0>::type>
inline TVMRetValue& operator=(const T& other);
// NodeRef related extenstions: in tvm/packed_func_ext.h
template<typename T,
typename = typename std::enable_if<
......@@ -548,9 +564,11 @@ class TVMRetValue : public TVMPODValue_ {
SwitchToPOD(other.type_code());
value_ = other.value_;
} else {
TVMExtTypeCapsule cap = *other.template ptr<TVMExtTypeCapsule>();
cap.ptr = cap.clone(cap.ptr);
SwitchToClass<TVMExtTypeCapsule>(other.type_code(), cap);
this->Clear();
type_code_ = other.type_code();
value_.v_handle =
(*(ExtTypeVTable::Get(other.type_code())->clone))(
other.value().v_handle);
}
break;
}
......@@ -582,9 +600,7 @@ class TVMRetValue : public TVMPODValue_ {
case kNodeHandle: delete ptr<std::shared_ptr<Node> >(); break;
}
if (type_code_ > kExtBegin) {
TVMExtTypeCapsule *cap = ptr<TVMExtTypeCapsule>();
cap->destroy(cap->ptr);
delete cap;
(*(ExtTypeVTable::Get(type_code_)->destroy))(value_.v_handle);
}
type_code_ = kNull;
}
......@@ -700,10 +716,8 @@ inline void for_each(const F& f, Args&&... args) { // NOLINT(*)
/* \brief argument settter to PackedFunc */
class TVMArgsSetter {
public:
TVMArgsSetter(TVMValue* values,
int* type_codes,
TVMExtTypeCapsule* exts)
: values_(values), type_codes_(type_codes), exts_(exts) {}
TVMArgsSetter(TVMValue* values, int* type_codes)
: values_(values), type_codes_(type_codes) {}
// setters for POD types
template<typename T,
typename = typename std::enable_if<
......@@ -793,21 +807,15 @@ class TVMArgsSetter {
TVMValue* values_;
/*! \brief The type code fields */
int* type_codes_;
/*! \brief Temporary storage for extension types */
TVMExtTypeCapsule* exts_;
};
template<typename... Args>
inline TVMRetValue PackedFunc::operator()(Args&& ...args) const {
const int kNumArgs = sizeof...(Args);
// Compiler will remove an static array when it is not touched.
const int kArraySize = kNumArgs > 0 ? kNumArgs : 1;
TVMValue values[kArraySize];
int type_codes[kArraySize];
// If the function call does not contain extension type,
// exts will get optimized away by compiler.
TVMExtTypeCapsule exts[kArraySize];
detail::for_each(TVMArgsSetter(values, type_codes, exts),
detail::for_each(TVMArgsSetter(values, type_codes),
std::forward<Args>(args)...);
TVMRetValue rv;
body_(TVMArgs(values, type_codes, kNumArgs), &rv);
......@@ -845,6 +853,14 @@ inline TVMRetValue::operator T() const {
::Apply(this);
}
template<typename T, typename>
inline void TVMArgsSetter::operator()(size_t i, const T& value) const {
static_assert(extension_class_info<T>::code != 0,
"Need to have extesion code");
type_codes_[i] = extension_class_info<T>::code;
values_[i].v_handle = const_cast<T*>(&value);
}
// extension type handling
template<typename T>
struct ExtTypeInfo {
......@@ -856,42 +872,16 @@ struct ExtTypeInfo {
}
};
template<typename T, typename>
inline TVMRetValue& TVMRetValue::operator=(const T& other) {
TVMExtTypeCapsule cap;
cap.clone = ExtTypeInfo<T>::clone;
cap.destroy = ExtTypeInfo<T>::destroy;
cap.ptr = new T(other);
SwitchToClass<TVMExtTypeCapsule>(
extension_class_info<T>::code, cap);
return *this;
}
template<typename T, typename>
inline void TVMArgsSetter::operator()(size_t i, const T& value) const {
static_assert(extension_class_info<T>::code != 0,
"Need to have extesion code");
type_codes_[i] = extension_class_info<T>::code;
exts_[i].clone = ExtTypeInfo<T>::clone;
exts_[i].destroy = ExtTypeInfo<T>::destroy;
exts_[i].ptr = const_cast<T*>(&value);
values_[i].v_handle = &exts_[i];
}
// Implement Module::GetFunction
// Put implementation in this file so we have seen the PackedFunc
inline PackedFunc Module::GetFunction(const std::string& name, bool query_imports) {
PackedFunc pf = node_->GetFunction(name, node_);
if (pf != nullptr) return pf;
if (query_imports) {
for (const Module& m : node_->imports_) {
pf = m.node_->GetFunction(name, m.node_);
if (pf != nullptr) return pf;
}
}
return pf;
template<typename T>
inline ExtTypeVTable* ExtTypeVTable::Register_() {
const int code = extension_class_info<T>::code;
static_assert(code != 0,
"require extension_class_info traits to be declared with non-zero code");
ExtTypeVTable vt;
vt.clone = ExtTypeInfo<T>::clone;
vt.destroy = ExtTypeInfo<T>::destroy;
return ExtTypeVTable::RegisterInternal(code, vt);
}
} // namespace runtime
} // namespace tvm
#endif // TVM_RUNTIME_PACKED_FUNC_H_
python/tvm/_ffi/function.py
......@@ -234,31 +234,6 @@ def list_global_func_names():
return fnames
def extract_ext_funcs(finit):
"""
Extract the extension PackedFuncs from a C module.
Parameters
----------
finit : ctypes function
a ctypes that takes signature of TVMExtensionDeclarer
Returns
-------
fdict : dict of str to Function
The extracted functions
"""
fdict = {}
def _list(name, func):
fdict[name] = func
myf = convert_to_tvm_func(_list)
ret = finit(myf.handle)
_ = myf
if ret != 0:
raise RuntimeError("cannot initialize with %s" % finit)
return fdict
def _get_api(f):
flocal = f
flocal.is_global = True
......
python/tvm/api.py
......@@ -8,7 +8,7 @@ from ._ffi.base import string_types
from ._ffi.node import register_node, NodeBase
from ._ffi.node import convert_to_node as _convert_to_node
from ._ffi.function import Function
from ._ffi.function import _init_api, register_func, get_global_func, extract_ext_funcs
from ._ffi.function import _init_api, register_func, get_global_func
from ._ffi.function import convert_to_tvm_func as _convert_tvm_func
from ._ffi.runtime_ctypes import TVMType
from . import _api_internal
......
python/tvm/build_module.py
......@@ -23,16 +23,16 @@ from . import target as _target
from . import make
class DumpIR(object):
"""
Dump IR for each pass.
With it, you can dump ir just like gcc/llvm.
"""Dump IR for each pass.
With it, you can dump ir just like gcc/llvm.
How to use:
-----------
.. code-block:: python
How to use:
-----------
.. code-block:: python
with tvm.build_config(dump_pass_ir=True)
run()
with tvm.build_config(dump_pass_ir=True)
run()
"""
scope_level = 0
def __init__(self):
......@@ -40,9 +40,9 @@ class DumpIR(object):
self._recover_list = []
def decorate(self, func):
""" decorate the pass function"""
''' decorate the pass function'''
def dump(*args, **kwargs):
"""dump function"""
'''dump function'''
retv = func(*args, **kwargs)
if not isinstance(retv, (_stmt.Stmt, container.LoweredFunc, container.Array)):
return retv
......@@ -59,7 +59,7 @@ class DumpIR(object):
return dump
def decorate_irpass(self):
"""decorate ir_pass and ScheduleOps"""
'''decorate ir_pass and ScheduleOps'''
self._old_sgpass = schedule.ScheduleOps
schedule.ScheduleOps = self.decorate(schedule.ScheduleOps)
vset = vars(ir_pass)
......@@ -71,7 +71,7 @@ class DumpIR(object):
vset[k] = self.decorate(v) if isinstance(v, types.FunctionType) else v
def decorate_custompass(self):
""" decorate add_lower_pass pass in BuildConfig"""
''' decorate add_lower_pass pass in BuildConfig'''
cfg = BuildConfig.current
self._old_custom_pass = cfg.add_lower_pass
custom_pass = cfg.add_lower_pass if cfg.add_lower_pass else []
......@@ -79,7 +79,7 @@ class DumpIR(object):
BuildConfig.current.add_lower_pass = pass_list
def enter(self):
"""only decorate outermost nest"""
'''only decorate outermost nest'''
if DumpIR.scope_level > 0:
return
self.decorate_irpass()
......@@ -88,7 +88,7 @@ class DumpIR(object):
DumpIR.scope_level += 1
def exit(self):
"""recover outermost nest"""
'''recover outermost nest'''
if DumpIR.scope_level > 1:
return
# recover decorated functions
......@@ -163,7 +163,6 @@ class BuildConfig(NodeBase):
"'%s' object cannot set attribute '%s'" % (str(type(self)), name))
return super(BuildConfig, self).__setattr__(name, value)
def build_config(**kwargs):
"""Configure the build behavior by setting config variables.
......@@ -227,7 +226,6 @@ def build_config(**kwargs):
setattr(config, k, kwargs[k])
return config
if not _RUNTIME_ONLY:
# BuildConfig is not available in tvm_runtime
BuildConfig.current = build_config()
......@@ -354,10 +352,8 @@ def lower(sch,
stmt = f(stmt)
if simple_mode:
return stmt
return ir_pass.MakeAPI(stmt, name, arg_list, 0, cfg.restricted_func)
def build(sch,
args=None,
target=None,
......
src/runtime/c_runtime_api.cc
......@@ -347,14 +347,6 @@ int TVMFuncCreateFromCFunc(TVMPackedCFunc func,
API_END();
}
int TVMExtTypeFree(void* handle, int type_code) {
API_BEGIN();
TVMExtTypeCapsule* cap = static_cast<TVMExtTypeCapsule*>(handle);
cap->destroy(cap->ptr);
delete cap;
API_END();
}
int TVMArrayAlloc(const tvm_index_t* shape,
int ndim,
int dtype_code,
......
src/runtime/module.cc
......@@ -13,6 +13,19 @@
namespace tvm {
namespace runtime {
PackedFunc Module::GetFunction(
const std::string& name, bool query_imports) {
PackedFunc pf = node_->GetFunction(name, node_);
if (pf != nullptr) return pf;
if (query_imports) {
for (const Module& m : node_->imports_) {
pf = m.node_->GetFunction(name, m.node_);
if (pf != nullptr) return pf;
}
}
return pf;
}
void Module::Import(Module other) {
// specially handle rpc
if (!std::strcmp((*this)->type_key(), "rpc")) {
......
src/runtime/registry.cc
......@@ -22,10 +22,15 @@ struct Registry::Manager {
// and the resource can become invalid because of indeterminstic order of destruction.
// The resources will only be recycled during program exit.
std::unordered_map<std::string, Registry*> fmap;
// vtable for extension type
std::array<ExtTypeVTable, kExtEnd> ext_vtable;
// mutex
std::mutex mutex;
Manager() {
for (auto& x : ext_vtable) {
x.destroy = nullptr;
}
}
static Manager* Global() {
......@@ -83,6 +88,24 @@ std::vector<std::string> Registry::ListNames() {
return keys;
}
ExtTypeVTable* ExtTypeVTable::Get(int type_code) {
CHECK(type_code > kExtBegin && type_code < kExtEnd);
Registry::Manager* m = Registry::Manager::Global();
ExtTypeVTable* vt = &(m->ext_vtable[type_code]);
CHECK(vt->destroy != nullptr)
<< "Extension type not registered";
return vt;
}
ExtTypeVTable* ExtTypeVTable::RegisterInternal(
int type_code, const ExtTypeVTable& vt) {
CHECK(type_code > kExtBegin && type_code < kExtEnd);
Registry::Manager* m = Registry::Manager::Global();
std::lock_guard<std::mutex>(m->mutex);
ExtTypeVTable* pvt = &(m->ext_vtable[type_code]);
pvt[0] = vt;
return pvt;
}
} // namespace runtime
} // namespace tvm
......@@ -97,6 +120,12 @@ struct TVMFuncThreadLocalEntry {
/*! \brief Thread local store that can be used to hold return values. */
typedef dmlc::ThreadLocalStore<TVMFuncThreadLocalEntry> TVMFuncThreadLocalStore;
int TVMExtTypeFree(void* handle, int type_code) {
API_BEGIN();
tvm::runtime::ExtTypeVTable::Get(type_code)->destroy(handle);
API_END();
}
int TVMFuncRegisterGlobal(
const char* name, TVMFunctionHandle f, int override) {
API_BEGIN();
......
tests/cpp/packed_func_test.cc
......@@ -126,6 +126,9 @@ struct extension_class_info<test::IntVector> {
} // runtime
} // tvm
// do registration, this need to be in cc file
TVM_REGISTER_EXT_TYPE(test::IntVector);
TEST(PackedFunc, ExtensionType) {
using namespace tvm;
using namespace tvm::runtime;
......
tests/scripts/task_python_integration.sh
......@@ -6,7 +6,6 @@ rm -rf python/tvm/*.pyc python/tvm/*/*.pyc
# Test TVM
make cython || exit -1
make cython3 || exit -1
# Test extern package package
cd apps/extension
......
topi/src/topi.cc
......@@ -54,6 +54,8 @@ namespace topi {
using namespace tvm;
using namespace tvm::runtime;
TVM_REGISTER_EXT_TYPE(tvm::Target);
/*! \brief Canonicalize an argument that may be Array<Expr> or int to Array<Expr> */
Array<Expr> ArrayOrInt(TVMArgValue arg) {
if (arg.type_code() == kDLInt || arg.type_code() == kDLUInt) {
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment