private:
// Internal function for conversion.
- friend class runtime::TVMPODValue_;
+ friend struct runtime::PackedFuncValueConverter<PrimExpr>;
TVM_DLL static PrimExpr FromObject_(ObjectPtr<Object> ptr);
};
namespace tvm {
namespace runtime {
-// Additional implementattion overloads for PackedFunc.
-inline TVMPODValue_::operator tvm::PrimExpr() const {
- if (type_code_ == kTVMNullptr) return PrimExpr();
- if (type_code_ == kDLInt) {
- CHECK_LE(value_.v_int64, std::numeric_limits<int>::max());
- CHECK_GE(value_.v_int64, std::numeric_limits<int>::min());
- return PrimExpr(static_cast<int>(value_.v_int64));
+template<>
+struct PackedFuncValueConverter<PrimExpr> {
+ // common rule for both RetValue and ArgValue.
+ static PrimExpr From(const TVMPODValue_& val) {
+ if (val.type_code() == kTVMNullptr) {
+ return PrimExpr(ObjectPtr<Object>(nullptr));
+ }
+ 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 tvm
#endif // TVM_IR_EXPR_H_
#include <tvm/runtime/object.h>
#include <tvm/runtime/memory.h>
#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/container.h>
#include <type_traits>
#include <vector>
kTVMStr = 11U,
kTVMBytes = 12U,
kTVMNDArrayHandle = 13U,
+ kTVMObjectRValueRefArg = 14U,
// Extension codes for other frameworks to integrate TVM PackedFunc.
// To make sure each framework's id do not conflict, use first and
// last sections to mark ranges.
*
* \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.
#include <dmlc/logging.h>
#include <tvm/runtime/memory.h>
#include <tvm/runtime/object.h>
+#include <tvm/runtime/packed_func.h>
#include <cstring>
#include <initializer_list>
}
}
+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 tvm
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 class Object;
friend class ObjectRef;
friend class TVMArgsSetter;
friend class TVMRetValue;
friend class TVMArgValue;
+ friend class TVMMovableArgValue_;
template <typename RelayRefType, typename ObjType>
friend RelayRefType GetRef(const ObjType* ptr);
template <typename BaseType, typename ObjType>
bool unique() const {
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
* raw pointer of a corresponding type.
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/object.h>
-#include <tvm/runtime/container.h>
#include <functional>
#include <tuple>
#include <vector>
#endif
namespace tvm {
-// forward declarations
-class Integer;
-class PrimExpr;
-
namespace runtime {
// forward declarations
class TVMArgs;
class TVMArgValue;
+class TVMMovableArgValue_;
class TVMRetValue;
class TVMArgsSetter;
*/
inline TypedPackedFunc(const TVMArgValue& value); // NOLINT(*)
/*!
+ * \brief constructor from TVMMovableArgValue_
+ * \param value The TVMMovableArgValue_
+ */
+ inline TypedPackedFunc(TVMMovableArgValue_&& value); // NOLINT(*)
+ /*!
* \brief construct from a lambda function with the same signature.
*
* Example usage:
}
operator int() const {
TVM_CHECK_TYPE_CODE(type_code_, kDLInt);
- CHECK_LE(value_.v_int64,
- std::numeric_limits<int>::max());
+ CHECK_LE(value_.v_int64, std::numeric_limits<int>::max());
+ CHECK_GE(value_.v_int64, std::numeric_limits<int>::min());
return static_cast<int>(value_.v_int64);
}
operator bool() const {
inline bool IsObjectRef() const;
template<typename TObjectRef>
inline TObjectRef AsObjectRef() const;
- // ObjectRef Specializations
- inline operator tvm::PrimExpr() const;
- inline operator tvm::Integer() const;
protected:
friend class TVMArgsSetter;
using TVMPODValue_::operator Module;
using TVMPODValue_::IsObjectRef;
using TVMPODValue_::AsObjectRef;
- using TVMPODValue_::operator tvm::PrimExpr;
- using TVMPODValue_::operator tvm::Integer;
// conversion operator.
operator std::string() const {
return std::string(value_.v_str);
}
}
- operator tvm::runtime::String() const {
- if (IsObjectRef<tvm::runtime::String>()) {
- return AsObjectRef<tvm::runtime::String>();
- } else {
- return tvm::runtime::String(operator std::string());
- }
- }
operator DLDataType() const {
if (type_code_ == kTVMStr) {
return String2DLDataType(operator std::string());
const TVMValue& value() const {
return value_;
}
+
template<typename T,
typename = typename std::enable_if<
std::is_class<T>::value>::type>
};
/*!
+ * \brief Internal auxiliary struct for TypedPackedFunc to indicate a movable argument.
+ *
+ * We can only construct a movable argument once from a single argument position.
+ * If the argument is passed as RValue reference, the result will be moved.
+ * We should only construct a MovableArg from an argument once,
+ * as the result will can moved.
+ *
+ * \note For internal development purpose only.
+ */
+class TVMMovableArgValue_ : public TVMArgValue {
+ public:
+ TVMMovableArgValue_(TVMValue value, int type_code)
+ : TVMArgValue(value, type_code) {
+ }
+ // reuse converter from parent
+ using TVMArgValue::operator double;
+ using TVMArgValue::operator int64_t;
+ using TVMArgValue::operator uint64_t;
+ using TVMArgValue::operator int;
+ using TVMArgValue::operator bool;
+ using TVMArgValue::operator void*;
+ using TVMArgValue::operator DLTensor*;
+ using TVMArgValue::operator TVMContext;
+ using TVMArgValue::operator std::string;
+ using TVMArgValue::operator DLDataType;
+ using TVMArgValue::operator DataType;
+ using TVMArgValue::operator PackedFunc;
+ /*!
+ * \brief Helper converter function.
+ * Try to move out an argument if possible,
+ * fall back to normal argument conversion rule otherwise.
+ */
+ template<typename T,
+ typename = typename std::enable_if<
+ std::is_base_of<ObjectRef, T>::value>::type>
+ inline operator T() const;
+};
+
+/*!
* \brief Return Value container,
* Unlike TVMArgValue, which only holds reference and do not delete
* the underlying container during destruction.
using TVMPODValue_::operator Module;
using TVMPODValue_::IsObjectRef;
using TVMPODValue_::AsObjectRef;
- using TVMPODValue_::operator tvm::PrimExpr;
- using TVMPODValue_::operator tvm::Integer;
TVMRetValue(const TVMRetValue& other) : TVMPODValue_() {
this->Assign(other);
TVM_CHECK_TYPE_CODE(type_code_, kTVMStr);
return *ptr<std::string>();
}
- operator tvm::runtime::String() const {
- if (IsObjectRef<tvm::runtime::String>()) {
- return AsObjectRef<tvm::runtime::String>();
- } else {
- return tvm::runtime::String(operator std::string());
- }
- }
operator DLDataType() const {
if (type_code_ == kTVMStr) {
return String2DLDataType(operator std::string());
this->Assign(other);
return *this;
}
+ TVMRetValue& operator=(TVMMovableArgValue_&& other) {
+ this->Assign(other);
+ return *this;
+ }
/*!
* \brief Move the value back to front-end via C API.
* This marks the current container as null.
static_cast<Object*>(other.value_.v_handle)));
break;
}
+ case kTVMObjectRValueRefArg: {
+ operator=(other.operator ObjectRef());
+ break;
+ }
default: {
SwitchToPOD(other.type_code());
value_ = other.value_;
};
/*!
+ * \brief Type trait to specify special value conversion rules from
+ * TVMArgValue and TVMRetValue.
+ *
+ * The trait can be specialized to add type specific conversion logic
+ * from the TVMArgvalue and TVMRetValue.
+ *
+ * \tparam TObjectRef the specific ObjectRefType.
+ */
+template<typename TObjectRef>
+struct PackedFuncValueConverter {
+ /*!
+ * \brief Convert a TObjectRef from an argument value.
+ * \param val The argument value.
+ * \return the converted result.
+ */
+ static TObjectRef From(const TVMArgValue& val) {
+ return val.AsObjectRef<TObjectRef>();
+ }
+ /*!
+ * \brief Convert a TObjectRef from a return value.
+ * \param val The argument value.
+ * \return the converted result.
+ */
+ static TObjectRef From(const TVMRetValue& val) {
+ return val.AsObjectRef<TObjectRef>();
+ }
+};
+
+/*!
* \brief Export a function with the PackedFunc signature
* as a PackedFunc that can be loaded by LibraryModule.
*
// ObjectRef handling
template<typename TObjectRef,
typename = typename std::enable_if<
- std::is_base_of<ObjectRef, TObjectRef>::value>::type>
- inline void operator()(size_t i, const TObjectRef& value) const;
+ std::is_base_of<ObjectRef, TObjectRef>::value>
+ ::type>
+ void operator()(size_t i, const TObjectRef& value) const {
+ this->SetObject(i, value);
+ }
+
+ template<typename TObjectRef,
+ typename = typename std::enable_if<
+ std::is_base_of<ObjectRef,
+ typename std::remove_reference<TObjectRef>::type>::value>
+ ::type>
+ void operator()(size_t i, TObjectRef&& value) const {
+ this->SetObject(i, std::forward<TObjectRef>(value));
+ }
private:
+ template<typename TObjectRef>
+ inline void SetObject(size_t i, TObjectRef&& value) const;
/*! \brief The values fields */
TVMValue* values_;
/*! \brief The type code fields */
const TVMArgs& args_pack,
TVMRetValue* rv,
Args&&... unpacked_args) {
+ // construct a movable argument value
+ // which allows potential move of argument to the input of F.
unpack_call_dispatcher<R, nleft - 1, index + 1, F>
::run(f, args_pack, rv,
std::forward<Args>(unpacked_args)...,
- args_pack[index]);
+ TVMMovableArgValue_(args_pack.values[index],
+ args_pack.type_codes[index]));
}
};
: packed_(value.operator PackedFunc()) {}
template<typename R, typename ...Args>
+TypedPackedFunc<R(Args...)>::TypedPackedFunc(TVMMovableArgValue_&& value)
+ : packed_(value.operator PackedFunc()) {}
+
+template<typename R, typename ...Args>
template<typename FType>
inline void TypedPackedFunc<R(Args...)>::AssignTypedLambda(FType flambda) {
packed_ = PackedFunc([flambda](const TVMArgs& args, TVMRetValue* rv) {
// kTVMNDArrayHandle, kTVMModuleHandle, kTVMObjectHandle
//
// We use type traits to eliminate un-necessary checks.
-template<typename TObjectRef, typename>
-inline void TVMArgsSetter::operator()(size_t i, const TObjectRef& value) const {
+template<typename T>
+inline void TVMArgsSetter::SetObject(size_t i, T&& value) const {
+ using TObjectRef = typename std::remove_reference<T>::type;
if (value.defined()) {
Object* ptr = value.data_.data_;
if (std::is_base_of<NDArray, TObjectRef>::value ||
ptr->IsInstance<Module::ContainerType>())) {
values_[i].v_handle = ptr;
type_codes_[i] = kTVMModuleHandle;
+ } else if (std::is_rvalue_reference<T>::value) {
+ values_[i].v_handle = const_cast<Object**>(&(value.data_.data_));
+ type_codes_[i] = kTVMObjectRValueRefArg;
} else {
- values_[i].v_handle = ptr;
+ values_[i].v_handle = value.data_.data_;
type_codes_[i] = kTVMObjectHandle;
}
} else {
return type_code_ == kTVMModuleHandle &&
static_cast<Object*>(value_.v_handle)->IsInstance<ContainerType>();
}
+ // NOTE: we don't pass NDArray and runtime::Module as RValue ref.
+ if (type_code_ == kTVMObjectRValueRefArg) {
+ return ObjectTypeChecker<TObjectRef>::Check(
+ *static_cast<Object**>(value_.v_handle));
+ }
return
(std::is_base_of<TObjectRef, NDArray>::value && type_code_ == kTVMNDArrayHandle) ||
(std::is_base_of<TObjectRef, Module>::value && type_code_ == kTVMModuleHandle) ||
<< "Expect " << ObjectTypeChecker<TObjectRef>::TypeName()
<< " but get " << ptr->GetTypeKey();
return TObjectRef(GetObjectPtr<Object>(ptr));
+ } else if (type_code_ == kTVMObjectRValueRefArg) {
+ Object* ptr = *static_cast<Object**>(value_.v_handle);
+ CHECK(ObjectTypeChecker<TObjectRef>::Check(ptr))
+ << "Expect " << ObjectTypeChecker<TObjectRef>::TypeName()
+ << " but get " << ptr->GetTypeKey();
+ return TObjectRef(GetObjectPtr<Object>(ptr));
} else if (std::is_base_of<TObjectRef, NDArray>::value &&
type_code_ == kTVMNDArrayHandle) {
// Casting to a base class that NDArray can sub-class
return *this;
}
+
template<typename T, typename>
inline TVMArgValue::operator T() const {
- return AsObjectRef<T>();
+ return PackedFuncValueConverter<T>::From(*this);
+}
+
+template<typename T, typename>
+inline TVMMovableArgValue_::operator T() const {
+ if (type_code_ == kTVMObjectRValueRefArg) {
+ auto** ref = static_cast<Object**>(value_.v_handle);
+ if (ObjectTypeChecker<T>::Check(*ref)) {
+ return T(ObjectPtr<Object>::MoveFromRValueRefArg(ref));
+ }
+ }
+ // fallback
+ return PackedFuncValueConverter<T>::From(*this);
}
template<typename T, typename>
inline TVMRetValue::operator T() const {
- return AsObjectRef<T>();
+ return PackedFuncValueConverter<T>::From(*this);
}
inline PackedFunc Module::GetFunction(const std::string& name, bool query_imports) {
namespace tvm {
namespace runtime {
// Additional implementattion overloads for PackedFunc.
-inline TVMPODValue_::operator tvm::Integer() const {
- if (type_code_ == kTVMNullptr) return Integer();
- if (type_code_ == kDLInt) {
- CHECK_LE(value_.v_int64, std::numeric_limits<int>::max());
- CHECK_GE(value_.v_int64, std::numeric_limits<int>::min());
- return Integer(static_cast<int>(value_.v_int64));
+
+template<>
+struct PackedFuncValueConverter<tvm::Integer> {
+ // common rule for RetValue and ArgValue
+ static tvm::Integer From(const TVMPODValue_& val) {
+ 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 tvm
int tcode = typeCodes[i];
if (tcode == kTVMObjectHandle ||
tcode == kTVMPackedFuncHandle ||
+ tcode == kTVMObjectRValueRefArg ||
tcode == kTVMModuleHandle) {
- TVMCbArgToReturn(&arg, tcode);
+ TVMCbArgToReturn(&arg, &tcode);
}
jobject jarg = tvmRetValueToJava(env, arg, tcode);
env->SetObjectArrayElement(jargs, i, jarg);
C_TO_PY_ARG_SWITCH[TypeCode.OBJECT_HANDLE] = _wrap_arg_func(
_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):
"""Base object for all object types"""
from ..base import _LIB, get_last_ffi_error, py2cerror, check_call
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 .ndarray import NDArrayBase, _make_array
from .types import TVMValue, TypeCode
elif isinstance(arg, ctypes.c_void_p):
values[i].v_handle = arg
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):
arg = convert_to_tvm_func(arg)
values[i].v_handle = arg.handle
def _wrap_arg_func(return_f, type_code):
- tcode = ctypes.c_int(type_code)
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 _wrap_func
kTVMStr = 11
kTVMBytes = 12
kTVMNDArrayHandle = 13
+ kTVMObjectRefArg = 14
kTVMExtBegin = 15
cdef extern from "tvm/runtime/c_runtime_api.h":
void* resource_handle,
TVMPackedCFuncFinalizer fin,
TVMPackedFuncHandle *out)
- int TVMCbArgToReturn(TVMValue* value, int code)
+ int TVMCbArgToReturn(TVMValue* value, int* code)
int TVMArrayAlloc(tvm_index_t* shape,
tvm_index_t ndim,
DLDataType dtype,
property handle:
def __get__(self):
- if self.chandle == NULL:
- return None
- else:
- return ctypes_handle(self.chandle)
+ return ctypes_handle(self.chandle)
def __set__(self, value):
self._set_handle(value)
from cpython cimport Py_INCREF, Py_DECREF
from numbers import Number, Integral
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):
if (tcode == kTVMObjectHandle or
tcode == kTVMPackedFuncHandle or
tcode == kTVMModuleHandle or
+ tcode == kTVMObjectRefArg or
tcode > kTVMExtBegin):
- CALL(TVMCbArgToReturn(&value, tcode))
+ CALL(TVMCbArgToReturn(&value, &tcode))
if tcode != kTVMDLTensorHandle:
pyargs.append(make_ret(value, tcode))
elif isinstance(arg, ctypes.c_void_p):
value[0].v_handle = c_handle(arg)
tcode[0] = kTVMOpaqueHandle
+ elif isinstance(arg, ObjectRValueRef):
+ value[0].v_handle = &((<ObjectBase>(arg.obj)).chandle)
+ tcode[0] = kTVMObjectRefArg
elif callable(arg):
arg = convert_to_tvm_func(arg)
value[0].v_handle = (<PackedFuncBase>arg).chandle
STR = 11
BYTES = 12
NDARRAY_HANDLE = 13
+ OBJECT_RVALUE_REF_ARG = 14
EXT_BEGIN = 15
("strides", ctypes.POINTER(tvm_shape_index_t)),
("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)
import ctypes
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
try:
else:
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)
# pylint: disable=unused-import, invalid-name
from numbers import Number, Integral
from tvm._ffi.base import string_types
+from tvm._ffi.runtime_ctypes import ObjectRValueRef
from . import _ffi_node_api, _ffi_api
from .object import ObjectBase, _set_class_object_generic
raise NotImplementedError()
-ObjectTypes = (ObjectBase, NDArrayBase, Module)
+ObjectTypes = (ObjectBase, NDArrayBase, Module, ObjectRValueRef)
def convert_to_object(value):
|| tcode == ffi::TVMTypeCode_kTVMPackedFuncHandle 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));
}
.set_body_typed(
[](runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func,
PassInfo pass_info) {
- return ModulePass(pass_func, pass_info);
+ return ModulePass(pass_func, pass_info);
});
TVM_REGISTER_GLOBAL("transform.RunPass")
Object* ptr = static_cast<Object*>(args[0].value().v_handle);
if (ptr->IsInstance<MapNode>()) {
- CHECK(args[1].type_code() == kTVMObjectHandle);
auto* n = static_cast<const MapNode*>(ptr);
auto it = n->data.find(args[1].operator ObjectRef());
CHECK(it != n->data.end())
API_END();
}
-int TVMCbArgToReturn(TVMValue* value, int code) {
+int TVMCbArgToReturn(TVMValue* value, int* code) {
API_BEGIN();
tvm::runtime::TVMRetValue rv;
- rv = tvm::runtime::TVMArgValue(*value, code);
- int tcode;
- rv.MoveToCHost(value, &tcode);
- CHECK_EQ(tcode, code);
+ rv = tvm::runtime::TVMMovableArgValue_(*value, *code);
+ rv.MoveToCHost(value, code);
API_END();
}
.set_body_typed(ErrorTest);
// 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) {
- runtime::NDArray nd = args[0];
- // substract the current one
- *ret = (nd.use_count() - 1);
+ runtime::ObjectRef obj = args[0];
+ // substract the current one because we always copy
+ // and get another value.
+ *ret = (obj.use_count() - 1);
});
-
} // namespace tvm
#include <dmlc/logging.h>
#include <gtest/gtest.h>
#include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/container.h>
#include <tvm/runtime/registry.h>
#include <tvm/tir/expr.h>
Var x;
Var t = PackedFunc([&](TVMArgs args, TVMRetValue* rv) {
CHECK(args.num_args == 1);
- CHECK(args.type_codes[0] == kTVMObjectHandle);
+ CHECK(args[0].IsObjectRef<ObjectRef>());
Var b = args[0];
CHECK(x.same_as(b));
*rv = b;
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) {
testing::InitGoogleTest(&argc, argv);
testing::FLAGS_gtest_death_test_style = "threadsafe";
x = tvm.testing.context_test(x, x.device_type, x.device_id)
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():
n = 2
x = te.placeholder((n,n,n), name="X", dtype="float32")
for t in ["float64", "float32"]:
check_assign(t)
+
+
if __name__ == "__main__":
+ test_rvalue_ref()
+ exit(0)
test_empty_array()
test_get_global()
test_get_callback_with_node()
if name == "get_arr":
return lambda : nd
elif name == "ref_count":
- return lambda : tvm.testing.ndarray_use_count(nd)
+ return lambda : tvm.testing.object_use_count(nd)
elif name == "get_elem":
return lambda idx: nd.asnumpy()[idx]
elif name == "get_arr_elem":
var kTVMPackedFuncHandle = 10;
var kTVMStr = 11;
var kTVMBytes = 12;
+ var kTVMObjectRValueRefArg = 14;
//-----------------------------------------
// TVM CWrap library
// ----------------------------------------
("TVMCbArgToReturn",
"number",
["number", // TVMValue* value
- "number" // int code
+ "number" // int* code
]);
var TVMFuncCreateFromCFunc = Module.cwrap
var args = [];
for (var i = 0; i < nargs; ++i) {
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 ||
+ tcode == kTVMObjectRValueRefArg ||
tcode == kTVMPackedFuncHandle ||
tcode == kTVMModuleHandle) {
- TVM_CALL(TVMCbArgToReturn(vptr, tcode));
+ TVM_CALL(TVMCbArgToReturn(vptr, tcodeptr));
}
+ tcode = Module.getValue(tcodeptr, "i32");
args.push(TVMRetValueToJS(vptr, tcode));
}
var rv = funcTable[handle].apply(null, args);