Unverified Commit b72dd9d9 by Tianqi Chen Committed by GitHub

[RUNTIME] Introduce RValue reference(move) support to TypedPackedFunc (#5271)

* [RUNTIME] Introduce RValue reference(move) support to TypedPackedFunc

This PR introduces RValue reference support the PackedFunc calling convention to address the above issue.
Specifically, when an argument is a r-value reference, we will use a assign a different type code(`kObjectRValueRefArg`),
and pass `Object**`  (the address to the Object pointer) instead through the values array.
The callee can choose to move out this Object pointer and set the original Object pointer from the caller side to be nullptr.

We also add an experimental move support to the python side(marked as _move so to indicate the dev nature).
This enhancement will enable copy on write optimizations through out the TVM stack.

* Address review comments

* fix compilation
parent 575d5369
...@@ -123,7 +123,7 @@ class PrimExpr : public BaseExpr { ...@@ -123,7 +123,7 @@ class PrimExpr : public BaseExpr {
private: private:
// Internal function for conversion. // Internal function for conversion.
friend class runtime::TVMPODValue_; friend struct runtime::PackedFuncValueConverter<PrimExpr>;
TVM_DLL static PrimExpr FromObject_(ObjectPtr<Object> ptr); TVM_DLL static PrimExpr FromObject_(ObjectPtr<Object> ptr);
}; };
...@@ -451,22 +451,24 @@ inline const TTypeNode* RelayExprNode::type_as() const { ...@@ -451,22 +451,24 @@ inline const TTypeNode* RelayExprNode::type_as() const {
namespace tvm { namespace tvm {
namespace runtime { namespace runtime {
// Additional implementattion overloads for PackedFunc. template<>
inline TVMPODValue_::operator tvm::PrimExpr() const { struct PackedFuncValueConverter<PrimExpr> {
if (type_code_ == kTVMNullptr) return PrimExpr(); // common rule for both RetValue and ArgValue.
if (type_code_ == kDLInt) { static PrimExpr From(const TVMPODValue_& val) {
CHECK_LE(value_.v_int64, std::numeric_limits<int>::max()); if (val.type_code() == kTVMNullptr) {
CHECK_GE(value_.v_int64, std::numeric_limits<int>::min()); return PrimExpr(ObjectPtr<Object>(nullptr));
return PrimExpr(static_cast<int>(value_.v_int64)); }
if (val.type_code() == kDLInt) {
return PrimExpr(val.operator int());
}
if (val.type_code() == kDLFloat) {
return PrimExpr(static_cast<float>(val.operator double()));
}
TVM_CHECK_TYPE_CODE(val.type_code(), kTVMObjectHandle);
Object* ptr = val.ptr<Object>();
return PrimExpr::FromObject_(GetObjectPtr<Object>(ptr));
} }
if (type_code_ == kDLFloat) { };
return PrimExpr(static_cast<float>(value_.v_float64));
}
TVM_CHECK_TYPE_CODE(type_code_, kTVMObjectHandle);
Object* ptr = static_cast<Object*>(value_.v_handle);
return PrimExpr::FromObject_(ObjectPtr<Object>(ptr));
}
} // namespace runtime } // namespace runtime
} // namespace tvm } // namespace tvm
#endif // TVM_IR_EXPR_H_ #endif // TVM_IR_EXPR_H_
...@@ -26,6 +26,7 @@ ...@@ -26,6 +26,7 @@
#include <tvm/runtime/object.h> #include <tvm/runtime/object.h>
#include <tvm/runtime/memory.h> #include <tvm/runtime/memory.h>
#include <tvm/runtime/packed_func.h> #include <tvm/runtime/packed_func.h>
#include <tvm/runtime/container.h>
#include <type_traits> #include <type_traits>
#include <vector> #include <vector>
......
...@@ -104,6 +104,7 @@ typedef enum { ...@@ -104,6 +104,7 @@ typedef enum {
kTVMStr = 11U, kTVMStr = 11U,
kTVMBytes = 12U, kTVMBytes = 12U,
kTVMNDArrayHandle = 13U, kTVMNDArrayHandle = 13U,
kTVMObjectRValueRefArg = 14U,
// Extension codes for other frameworks to integrate TVM PackedFunc. // Extension codes for other frameworks to integrate TVM PackedFunc.
// To make sure each framework's id do not conflict, use first and // To make sure each framework's id do not conflict, use first and
// last sections to mark ranges. // last sections to mark ranges.
...@@ -290,7 +291,7 @@ TVM_DLL int TVMCFuncSetReturn(TVMRetValueHandle ret, ...@@ -290,7 +291,7 @@ TVM_DLL int TVMCFuncSetReturn(TVMRetValueHandle ret,
* *
* \return 0 when success, -1 when failure happens. * \return 0 when success, -1 when failure happens.
*/ */
TVM_DLL int TVMCbArgToReturn(TVMValue* value, int code); TVM_DLL int TVMCbArgToReturn(TVMValue* value, int* code);
/*! /*!
* \brief C type of packed function. * \brief C type of packed function.
......
...@@ -27,6 +27,7 @@ ...@@ -27,6 +27,7 @@
#include <dmlc/logging.h> #include <dmlc/logging.h>
#include <tvm/runtime/memory.h> #include <tvm/runtime/memory.h>
#include <tvm/runtime/object.h> #include <tvm/runtime/object.h>
#include <tvm/runtime/packed_func.h>
#include <cstring> #include <cstring>
#include <initializer_list> #include <initializer_list>
...@@ -590,6 +591,25 @@ inline int String::memncmp(const char* lhs, const char* rhs, size_t lhs_count, ...@@ -590,6 +591,25 @@ inline int String::memncmp(const char* lhs, const char* rhs, size_t lhs_count,
} }
} }
template<>
struct PackedFuncValueConverter<::tvm::runtime::String> {
static String From(const TVMArgValue& val) {
if (val.IsObjectRef<tvm::runtime::String>()) {
return val.AsObjectRef<tvm::runtime::String>();
} else {
return tvm::runtime::String(val.operator std::string());
}
}
static String From(const TVMRetValue& val) {
if (val.IsObjectRef<tvm::runtime::String>()) {
return val.AsObjectRef<tvm::runtime::String>();
} else {
return tvm::runtime::String(val.operator std::string());
}
}
};
} // namespace runtime } // namespace runtime
} // namespace tvm } // namespace tvm
......
...@@ -477,6 +477,17 @@ class ObjectPtr { ...@@ -477,6 +477,17 @@ class ObjectPtr {
data_->IncRef(); data_->IncRef();
} }
} }
/*!
* \brief Move an ObjectPtr from an RValueRef argument.
* \param ref The rvalue reference.
* \return the moved result.
*/
static ObjectPtr<T> MoveFromRValueRefArg(Object** ref) {
ObjectPtr<T> ptr;
ptr.data_ = *ref;
*ref = nullptr;
return ptr;
}
// friend classes // friend classes
friend class Object; friend class Object;
friend class ObjectRef; friend class ObjectRef;
...@@ -489,6 +500,7 @@ class ObjectPtr { ...@@ -489,6 +500,7 @@ class ObjectPtr {
friend class TVMArgsSetter; friend class TVMArgsSetter;
friend class TVMRetValue; friend class TVMRetValue;
friend class TVMArgValue; friend class TVMArgValue;
friend class TVMMovableArgValue_;
template <typename RelayRefType, typename ObjType> template <typename RelayRefType, typename ObjType>
friend RelayRefType GetRef(const ObjType* ptr); friend RelayRefType GetRef(const ObjType* ptr);
template <typename BaseType, typename ObjType> template <typename BaseType, typename ObjType>
...@@ -550,6 +562,10 @@ class ObjectRef { ...@@ -550,6 +562,10 @@ class ObjectRef {
bool unique() const { bool unique() const {
return data_.unique(); return data_.unique();
} }
/*! \return The use count of the ptr, for debug purposes */
int use_count() const {
return data_.use_count();
}
/*! /*!
* \brief Try to downcast the internal Object to a * \brief Try to downcast the internal Object to a
* raw pointer of a corresponding type. * raw pointer of a corresponding type.
......
...@@ -1338,20 +1338,20 @@ enum TVMStructFieldKind : int { ...@@ -1338,20 +1338,20 @@ enum TVMStructFieldKind : int {
namespace tvm { namespace tvm {
namespace runtime { namespace runtime {
// Additional implementattion overloads for PackedFunc. // Additional implementattion overloads for PackedFunc.
inline TVMPODValue_::operator tvm::Integer() const {
if (type_code_ == kTVMNullptr) return Integer(); template<>
if (type_code_ == kDLInt) { struct PackedFuncValueConverter<tvm::Integer> {
CHECK_LE(value_.v_int64, std::numeric_limits<int>::max()); // common rule for RetValue and ArgValue
CHECK_GE(value_.v_int64, std::numeric_limits<int>::min()); static tvm::Integer From(const TVMPODValue_& val) {
return Integer(static_cast<int>(value_.v_int64)); if (val.type_code() == kTVMNullptr) {
return Integer(ObjectPtr<Object>(nullptr));
}
if (val.type_code() == kDLInt) {
return Integer(val.operator int());
}
return val.AsObjectRef<tvm::Integer>();
} }
TVM_CHECK_TYPE_CODE(type_code_, kTVMObjectHandle); };
Object* ptr = static_cast<Object*>(value_.v_handle);
CHECK(ObjectTypeChecker<Integer>::Check(ptr))
<< "Expect type " << ObjectTypeChecker<PrimExpr>::TypeName()
<< " but get " << ptr->GetTypeKey();
return Integer(ObjectPtr<Object>(ptr));
}
} // namespace runtime } // namespace runtime
} // namespace tvm } // namespace tvm
......
...@@ -244,8 +244,9 @@ extern "C" int funcInvokeCallback(TVMValue *args, ...@@ -244,8 +244,9 @@ extern "C" int funcInvokeCallback(TVMValue *args,
int tcode = typeCodes[i]; int tcode = typeCodes[i];
if (tcode == kTVMObjectHandle || if (tcode == kTVMObjectHandle ||
tcode == kTVMPackedFuncHandle || tcode == kTVMPackedFuncHandle ||
tcode == kTVMObjectRValueRefArg ||
tcode == kTVMModuleHandle) { tcode == kTVMModuleHandle) {
TVMCbArgToReturn(&arg, tcode); TVMCbArgToReturn(&arg, &tcode);
} }
jobject jarg = tvmRetValueToJava(env, arg, tcode); jobject jarg = tvmRetValueToJava(env, arg, tcode);
env->SetObjectArrayElement(jargs, i, jarg); env->SetObjectArrayElement(jargs, i, jarg);
......
...@@ -60,6 +60,9 @@ RETURN_SWITCH[TypeCode.OBJECT_HANDLE] = _return_object ...@@ -60,6 +60,9 @@ RETURN_SWITCH[TypeCode.OBJECT_HANDLE] = _return_object
C_TO_PY_ARG_SWITCH[TypeCode.OBJECT_HANDLE] = _wrap_arg_func( C_TO_PY_ARG_SWITCH[TypeCode.OBJECT_HANDLE] = _wrap_arg_func(
_return_object, TypeCode.OBJECT_HANDLE) _return_object, TypeCode.OBJECT_HANDLE)
C_TO_PY_ARG_SWITCH[TypeCode.OBJECT_RVALUE_REF_ARG] = _wrap_arg_func(
_return_object, TypeCode.OBJECT_RVALUE_REF_ARG)
class ObjectBase(object): class ObjectBase(object):
"""Base object for all object types""" """Base object for all object types"""
......
...@@ -23,7 +23,7 @@ from numbers import Number, Integral ...@@ -23,7 +23,7 @@ from numbers import Number, Integral
from ..base import _LIB, get_last_ffi_error, py2cerror, check_call from ..base import _LIB, get_last_ffi_error, py2cerror, check_call
from ..base import c_str, string_types from ..base import c_str, string_types
from ..runtime_ctypes import DataType, TVMByteArray, TVMContext from ..runtime_ctypes import DataType, TVMByteArray, TVMContext, ObjectRValueRef
from . import ndarray as _nd from . import ndarray as _nd
from .ndarray import NDArrayBase, _make_array from .ndarray import NDArrayBase, _make_array
from .types import TVMValue, TypeCode from .types import TVMValue, TypeCode
...@@ -164,6 +164,9 @@ def _make_tvm_args(args, temp_args): ...@@ -164,6 +164,9 @@ def _make_tvm_args(args, temp_args):
elif isinstance(arg, ctypes.c_void_p): elif isinstance(arg, ctypes.c_void_p):
values[i].v_handle = arg values[i].v_handle = arg
type_codes[i] = TypeCode.HANDLE type_codes[i] = TypeCode.HANDLE
elif isinstance(arg, ObjectRValueRef):
values[i].v_handle = ctypes.cast(ctypes.byref(arg.obj.handle), ctypes.c_void_p)
type_codes[i] = TypeCode.OBJECT_RVALUE_REF_ARG
elif callable(arg): elif callable(arg):
arg = convert_to_tvm_func(arg) arg = convert_to_tvm_func(arg)
values[i].v_handle = arg.handle values[i].v_handle = arg.handle
......
...@@ -73,9 +73,9 @@ def _return_context(value): ...@@ -73,9 +73,9 @@ def _return_context(value):
def _wrap_arg_func(return_f, type_code): def _wrap_arg_func(return_f, type_code):
tcode = ctypes.c_int(type_code)
def _wrap_func(x): def _wrap_func(x):
check_call(_LIB.TVMCbArgToReturn(ctypes.byref(x), tcode)) tcode = ctypes.c_int(type_code)
check_call(_LIB.TVMCbArgToReturn(ctypes.byref(x), ctypes.byref(tcode)))
return return_f(x) return return_f(x)
return _wrap_func return _wrap_func
......
...@@ -37,6 +37,7 @@ cdef enum TVMTypeCode: ...@@ -37,6 +37,7 @@ cdef enum TVMTypeCode:
kTVMStr = 11 kTVMStr = 11
kTVMBytes = 12 kTVMBytes = 12
kTVMNDArrayHandle = 13 kTVMNDArrayHandle = 13
kTVMObjectRefArg = 14
kTVMExtBegin = 15 kTVMExtBegin = 15
cdef extern from "tvm/runtime/c_runtime_api.h": cdef extern from "tvm/runtime/c_runtime_api.h":
...@@ -113,7 +114,7 @@ cdef extern from "tvm/runtime/c_runtime_api.h": ...@@ -113,7 +114,7 @@ cdef extern from "tvm/runtime/c_runtime_api.h":
void* resource_handle, void* resource_handle,
TVMPackedCFuncFinalizer fin, TVMPackedCFuncFinalizer fin,
TVMPackedFuncHandle *out) TVMPackedFuncHandle *out)
int TVMCbArgToReturn(TVMValue* value, int code) int TVMCbArgToReturn(TVMValue* value, int* code)
int TVMArrayAlloc(tvm_index_t* shape, int TVMArrayAlloc(tvm_index_t* shape,
tvm_index_t ndim, tvm_index_t ndim,
DLDataType dtype, DLDataType dtype,
......
...@@ -64,10 +64,7 @@ cdef class ObjectBase: ...@@ -64,10 +64,7 @@ cdef class ObjectBase:
property handle: property handle:
def __get__(self): def __get__(self):
if self.chandle == NULL: return ctypes_handle(self.chandle)
return None
else:
return ctypes_handle(self.chandle)
def __set__(self, value): def __set__(self, value):
self._set_handle(value) self._set_handle(value)
......
...@@ -20,7 +20,7 @@ import traceback ...@@ -20,7 +20,7 @@ import traceback
from cpython cimport Py_INCREF, Py_DECREF from cpython cimport Py_INCREF, Py_DECREF
from numbers import Number, Integral from numbers import Number, Integral
from ..base import string_types, py2cerror from ..base import string_types, py2cerror
from ..runtime_ctypes import DataType, TVMContext, TVMByteArray from ..runtime_ctypes import DataType, TVMContext, TVMByteArray, ObjectRValueRef
cdef void tvm_callback_finalize(void* fhandle): cdef void tvm_callback_finalize(void* fhandle):
...@@ -43,8 +43,9 @@ cdef int tvm_callback(TVMValue* args, ...@@ -43,8 +43,9 @@ cdef int tvm_callback(TVMValue* args,
if (tcode == kTVMObjectHandle or if (tcode == kTVMObjectHandle or
tcode == kTVMPackedFuncHandle or tcode == kTVMPackedFuncHandle or
tcode == kTVMModuleHandle or tcode == kTVMModuleHandle or
tcode == kTVMObjectRefArg or
tcode > kTVMExtBegin): tcode > kTVMExtBegin):
CALL(TVMCbArgToReturn(&value, tcode)) CALL(TVMCbArgToReturn(&value, &tcode))
if tcode != kTVMDLTensorHandle: if tcode != kTVMDLTensorHandle:
pyargs.append(make_ret(value, tcode)) pyargs.append(make_ret(value, tcode))
...@@ -167,6 +168,9 @@ cdef inline int make_arg(object arg, ...@@ -167,6 +168,9 @@ cdef inline int make_arg(object arg,
elif isinstance(arg, ctypes.c_void_p): elif isinstance(arg, ctypes.c_void_p):
value[0].v_handle = c_handle(arg) value[0].v_handle = c_handle(arg)
tcode[0] = kTVMOpaqueHandle tcode[0] = kTVMOpaqueHandle
elif isinstance(arg, ObjectRValueRef):
value[0].v_handle = &((<ObjectBase>(arg.obj)).chandle)
tcode[0] = kTVMObjectRefArg
elif callable(arg): elif callable(arg):
arg = convert_to_tvm_func(arg) arg = convert_to_tvm_func(arg)
value[0].v_handle = (<PackedFuncBase>arg).chandle value[0].v_handle = (<PackedFuncBase>arg).chandle
......
...@@ -39,6 +39,7 @@ class TypeCode(object): ...@@ -39,6 +39,7 @@ class TypeCode(object):
STR = 11 STR = 11
BYTES = 12 BYTES = 12
NDARRAY_HANDLE = 13 NDARRAY_HANDLE = 13
OBJECT_RVALUE_REF_ARG = 14
EXT_BEGIN = 15 EXT_BEGIN = 15
...@@ -281,4 +282,18 @@ class TVMArray(ctypes.Structure): ...@@ -281,4 +282,18 @@ class TVMArray(ctypes.Structure):
("strides", ctypes.POINTER(tvm_shape_index_t)), ("strides", ctypes.POINTER(tvm_shape_index_t)),
("byte_offset", ctypes.c_uint64)] ("byte_offset", ctypes.c_uint64)]
class ObjectRValueRef:
"""Represent an RValue ref to an object that can be moved.
Parameters
----------
obj : tvm.runtime.Object
The object that this value refers to
"""
__slots__ = ["obj"]
def __init__(self, obj):
self.obj = obj
TVMArrayHandle = ctypes.POINTER(TVMArray) TVMArrayHandle = ctypes.POINTER(TVMArray)
...@@ -19,6 +19,7 @@ ...@@ -19,6 +19,7 @@
import ctypes import ctypes
from tvm._ffi.base import _FFI_MODE, _RUNTIME_ONLY, check_call, _LIB, c_str from tvm._ffi.base import _FFI_MODE, _RUNTIME_ONLY, check_call, _LIB, c_str
from tvm._ffi.runtime_ctypes import ObjectRValueRef
from . import _ffi_api, _ffi_node_api from . import _ffi_api, _ffi_node_api
try: try:
...@@ -85,5 +86,35 @@ class Object(ObjectBase): ...@@ -85,5 +86,35 @@ class Object(ObjectBase):
else: else:
self.handle = None self.handle = None
def _move(self):
"""Create an RValue reference to the object and mark the object as moved.
This is a advanced developer API that can be useful when passing an
unique reference to an Object that you no longer needed to a function.
A unique reference can trigger copy on write optimization that avoids
copy when we transform an object.
Note
----
All the reference of the object becomes invalid after it is moved.
Be very careful when using this feature.
Examples
--------
.. code-block:: python
x = tvm.tir.Var("x", "int32")
x0 = x
some_packed_func(x._move())
# both x0 and x will points to None after the function call.
Returns
-------
rvalue : The rvalue reference.
"""
return ObjectRValueRef(self)
_set_class_object(Object) _set_class_object(Object)
...@@ -18,6 +18,7 @@ ...@@ -18,6 +18,7 @@
# pylint: disable=unused-import, invalid-name # pylint: disable=unused-import, invalid-name
from numbers import Number, Integral from numbers import Number, Integral
from tvm._ffi.base import string_types from tvm._ffi.base import string_types
from tvm._ffi.runtime_ctypes import ObjectRValueRef
from . import _ffi_node_api, _ffi_api from . import _ffi_node_api, _ffi_api
from .object import ObjectBase, _set_class_object_generic from .object import ObjectBase, _set_class_object_generic
...@@ -33,7 +34,7 @@ class ObjectGeneric(object): ...@@ -33,7 +34,7 @@ class ObjectGeneric(object):
raise NotImplementedError() raise NotImplementedError()
ObjectTypes = (ObjectBase, NDArrayBase, Module) ObjectTypes = (ObjectBase, NDArrayBase, Module, ObjectRValueRef)
def convert_to_object(value): def convert_to_object(value):
......
...@@ -261,7 +261,7 @@ unsafe extern "C" fn tvm_callback( ...@@ -261,7 +261,7 @@ unsafe extern "C" fn tvm_callback(
|| tcode == ffi::TVMTypeCode_kTVMPackedFuncHandle as c_int || tcode == ffi::TVMTypeCode_kTVMPackedFuncHandle as c_int
|| tcode == ffi::TVMTypeCode_kTVMModuleHandle as c_int || tcode == ffi::TVMTypeCode_kTVMModuleHandle as c_int
{ {
check_call!(ffi::TVMCbArgToReturn(&mut value as *mut _, tcode)); check_call!(ffi::TVMCbArgToReturn(&mut value as *mut _, &mut tcode as *mut _));
} }
local_args.push(TVMArgValue::from_tvm_value(value, tcode as u32)); local_args.push(TVMArgValue::from_tvm_value(value, tcode as u32));
} }
......
...@@ -371,7 +371,7 @@ TVM_REGISTER_GLOBAL("transform.MakeModulePass") ...@@ -371,7 +371,7 @@ TVM_REGISTER_GLOBAL("transform.MakeModulePass")
.set_body_typed( .set_body_typed(
[](runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func, [](runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func,
PassInfo pass_info) { PassInfo pass_info) {
return ModulePass(pass_func, pass_info); return ModulePass(pass_func, pass_info);
}); });
TVM_REGISTER_GLOBAL("transform.RunPass") TVM_REGISTER_GLOBAL("transform.RunPass")
......
...@@ -370,7 +370,6 @@ TVM_REGISTER_GLOBAL("node.MapGetItem") ...@@ -370,7 +370,6 @@ TVM_REGISTER_GLOBAL("node.MapGetItem")
Object* ptr = static_cast<Object*>(args[0].value().v_handle); Object* ptr = static_cast<Object*>(args[0].value().v_handle);
if (ptr->IsInstance<MapNode>()) { if (ptr->IsInstance<MapNode>()) {
CHECK(args[1].type_code() == kTVMObjectHandle);
auto* n = static_cast<const MapNode*>(ptr); auto* n = static_cast<const MapNode*>(ptr);
auto it = n->data.find(args[1].operator ObjectRef()); auto it = n->data.find(args[1].operator ObjectRef());
CHECK(it != n->data.end()) CHECK(it != n->data.end())
......
...@@ -577,13 +577,11 @@ int TVMStreamStreamSynchronize(int device_type, ...@@ -577,13 +577,11 @@ int TVMStreamStreamSynchronize(int device_type,
API_END(); API_END();
} }
int TVMCbArgToReturn(TVMValue* value, int code) { int TVMCbArgToReturn(TVMValue* value, int* code) {
API_BEGIN(); API_BEGIN();
tvm::runtime::TVMRetValue rv; tvm::runtime::TVMRetValue rv;
rv = tvm::runtime::TVMArgValue(*value, code); rv = tvm::runtime::TVMMovableArgValue_(*value, *code);
int tcode; rv.MoveToCHost(value, code);
rv.MoveToCHost(value, &tcode);
CHECK_EQ(tcode, code);
API_END(); API_END();
} }
......
...@@ -107,11 +107,11 @@ TVM_REGISTER_GLOBAL("testing.ErrorTest") ...@@ -107,11 +107,11 @@ TVM_REGISTER_GLOBAL("testing.ErrorTest")
.set_body_typed(ErrorTest); .set_body_typed(ErrorTest);
// internal function used for debug and testing purposes // internal function used for debug and testing purposes
TVM_REGISTER_GLOBAL("testing.ndarray_use_count") TVM_REGISTER_GLOBAL("testing.object_use_count")
.set_body([](TVMArgs args, TVMRetValue *ret) { .set_body([](TVMArgs args, TVMRetValue *ret) {
runtime::NDArray nd = args[0]; runtime::ObjectRef obj = args[0];
// substract the current one // substract the current one because we always copy
*ret = (nd.use_count() - 1); // and get another value.
*ret = (obj.use_count() - 1);
}); });
} // namespace tvm } // namespace tvm
...@@ -20,6 +20,7 @@ ...@@ -20,6 +20,7 @@
#include <dmlc/logging.h> #include <dmlc/logging.h>
#include <gtest/gtest.h> #include <gtest/gtest.h>
#include <tvm/runtime/packed_func.h> #include <tvm/runtime/packed_func.h>
#include <tvm/runtime/container.h>
#include <tvm/runtime/registry.h> #include <tvm/runtime/registry.h>
#include <tvm/tir/expr.h> #include <tvm/tir/expr.h>
...@@ -51,7 +52,7 @@ TEST(PackedFunc, Node) { ...@@ -51,7 +52,7 @@ TEST(PackedFunc, Node) {
Var x; Var x;
Var t = PackedFunc([&](TVMArgs args, TVMRetValue* rv) { Var t = PackedFunc([&](TVMArgs args, TVMRetValue* rv) {
CHECK(args.num_args == 1); CHECK(args.num_args == 1);
CHECK(args.type_codes[0] == kTVMObjectHandle); CHECK(args[0].IsObjectRef<ObjectRef>());
Var b = args[0]; Var b = args[0];
CHECK(x.same_as(b)); CHECK(x.same_as(b));
*rv = b; *rv = b;
...@@ -269,6 +270,50 @@ TEST(PackedFunc, ObjectConversion) { ...@@ -269,6 +270,50 @@ TEST(PackedFunc, ObjectConversion) {
pf2(ObjectRef(m), Module()); pf2(ObjectRef(m), Module());
} }
TEST(TypedPackedFunc, RValue) {
using namespace tvm;
using namespace tvm::runtime;
{
auto f = [](tir::Var x, bool move) {
if (move) {
CHECK(x.unique());
} else {
CHECK(!x.unique());
}
CHECK(x->name_hint == "x");
return x;
};
TypedPackedFunc<tir::Var(tir::Var, bool)> tf(f);
tir::Var var("x");
CHECK(var.unique());
f(var, false);
// move the result to the function.
tir::Var ret = f(std::move(var), true);
CHECK(!var.defined());
}
{
// pass child class.
auto f = [](PrimExpr x, bool move) {
if (move) {
CHECK(x.unique());
} else {
CHECK(!x.unique());
}
return x;
};
TypedPackedFunc<PrimExpr(PrimExpr, bool)> tf(f);
tir::Var var("x");
CHECK(var.unique());
f(var, false);
f(std::move(var), true);
// auto conversion.
f(1, true);
}
}
int main(int argc, char ** argv) { int main(int argc, char ** argv) {
testing::InitGoogleTest(&argc, argv); testing::InitGoogleTest(&argc, argv);
testing::FLAGS_gtest_death_test_style = "threadsafe"; testing::FLAGS_gtest_death_test_style = "threadsafe";
......
...@@ -98,6 +98,33 @@ def test_ctx(): ...@@ -98,6 +98,33 @@ def test_ctx():
x = tvm.testing.context_test(x, x.device_type, x.device_id) x = tvm.testing.context_test(x, x.device_type, x.device_id)
assert x == tvm.opencl(10) assert x == tvm.opencl(10)
def test_rvalue_ref():
def callback(x, expected_count):
assert expected_count == tvm.testing.object_use_count(x)
return x
f = tvm.runtime.convert(callback)
def check0():
x = tvm.tir.Var("x", "int32")
assert tvm.testing.object_use_count(x) == 1
f(x, 2)
y = f(x._move(), 1)
assert x.handle.value == None
def check1():
x = tvm.tir.Var("x", "int32")
assert tvm.testing.object_use_count(x) == 1
y = f(x, 2)
z = f(x._move(), 2)
assert x.handle.value == None
assert y.handle.value is not None
check0()
check1()
def test_trace_default_action(): def test_trace_default_action():
n = 2 n = 2
x = te.placeholder((n,n,n), name="X", dtype="float32") x = te.placeholder((n,n,n), name="X", dtype="float32")
...@@ -269,7 +296,11 @@ def test_trace_can_change_traced_value_float(): ...@@ -269,7 +296,11 @@ def test_trace_can_change_traced_value_float():
for t in ["float64", "float32"]: for t in ["float64", "float32"]:
check_assign(t) check_assign(t)
if __name__ == "__main__": if __name__ == "__main__":
test_rvalue_ref()
exit(0)
test_empty_array() test_empty_array()
test_get_global() test_get_global()
test_get_callback_with_node() test_get_callback_with_node()
......
...@@ -212,7 +212,7 @@ def test_rpc_return_ndarray(): ...@@ -212,7 +212,7 @@ def test_rpc_return_ndarray():
if name == "get_arr": if name == "get_arr":
return lambda : nd return lambda : nd
elif name == "ref_count": elif name == "ref_count":
return lambda : tvm.testing.ndarray_use_count(nd) return lambda : tvm.testing.object_use_count(nd)
elif name == "get_elem": elif name == "get_elem":
return lambda idx: nd.asnumpy()[idx] return lambda idx: nd.asnumpy()[idx]
elif name == "get_arr_elem": elif name == "get_arr_elem":
......
...@@ -105,6 +105,7 @@ var tvm_runtime = tvm_runtime || {}; ...@@ -105,6 +105,7 @@ var tvm_runtime = tvm_runtime || {};
var kTVMPackedFuncHandle = 10; var kTVMPackedFuncHandle = 10;
var kTVMStr = 11; var kTVMStr = 11;
var kTVMBytes = 12; var kTVMBytes = 12;
var kTVMObjectRValueRefArg = 14;
//----------------------------------------- //-----------------------------------------
// TVM CWrap library // TVM CWrap library
// ---------------------------------------- // ----------------------------------------
...@@ -171,7 +172,7 @@ var tvm_runtime = tvm_runtime || {}; ...@@ -171,7 +172,7 @@ var tvm_runtime = tvm_runtime || {};
("TVMCbArgToReturn", ("TVMCbArgToReturn",
"number", "number",
["number", // TVMValue* value ["number", // TVMValue* value
"number" // int code "number" // int* code
]); ]);
var TVMFuncCreateFromCFunc = Module.cwrap var TVMFuncCreateFromCFunc = Module.cwrap
...@@ -496,12 +497,15 @@ var tvm_runtime = tvm_runtime || {}; ...@@ -496,12 +497,15 @@ var tvm_runtime = tvm_runtime || {};
var args = []; var args = [];
for (var i = 0; i < nargs; ++i) { for (var i = 0; i < nargs; ++i) {
var vptr = arg_value + i * SIZEOF_TVMVALUE; var vptr = arg_value + i * SIZEOF_TVMVALUE;
var tcode = Module.getValue(arg_tcode + i * SIZEOF_INT, "i32"); var tcodeptr = arg_tcode + i * SIZEOF_INT;
var tcode = Module.getValue(tcodeptr, "i32");
if (tcode == kTVMObjectHandle || if (tcode == kTVMObjectHandle ||
tcode == kTVMObjectRValueRefArg ||
tcode == kTVMPackedFuncHandle || tcode == kTVMPackedFuncHandle ||
tcode == kTVMModuleHandle) { tcode == kTVMModuleHandle) {
TVM_CALL(TVMCbArgToReturn(vptr, tcode)); TVM_CALL(TVMCbArgToReturn(vptr, tcodeptr));
} }
tcode = Module.getValue(tcodeptr, "i32");
args.push(TVMRetValueToJS(vptr, tcode)); args.push(TVMRetValueToJS(vptr, tcode));
} }
var rv = funcTable[handle].apply(null, args); var rv = funcTable[handle].apply(null, args);
......
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