This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch v0.7
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git
The following commit(s) were added to refs/heads/v0.7 by this push:
new 091a427 [FFI][BUGFIX] Fix memory leak when Pac callback argument is
NDArray (#6744) (#6821)
091a427 is described below
commit 091a427e347aab92474941431f8a775aaeb2c1eb
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Nov 2 08:55:36 2020 -0500
[FFI][BUGFIX] Fix memory leak when Pac callback argument is NDArray (#6744)
(#6821)
Co-authored-by: Matthew Brookhart <[email protected]>
Co-authored-by: Junru Shao <[email protected]>
---
jvm/native/src/main/native/org_apache_tvm_native_c_api.cc | 3 ++-
python/tvm/_ffi/_ctypes/packed_func.py | 4 +++-
python/tvm/_ffi/_cython/ndarray.pxi | 4 ++++
python/tvm/_ffi/_cython/packed_func.pxi | 1 +
rust/tvm-rt/src/to_function.rs | 2 ++
tests/python/unittest/test_runtime_packed_func.py | 12 ++++++++++++
web/src/runtime.ts | 1 +
7 files changed, 25 insertions(+), 2 deletions(-)
diff --git a/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
b/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
index 6fc316c..e3ea4b9 100644
--- a/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
+++ b/jvm/native/src/main/native/org_apache_tvm_native_c_api.cc
@@ -243,7 +243,8 @@ extern "C" int funcInvokeCallback(TVMValue* args, int*
typeCodes, int numArgs,
TVMValue arg = args[i];
int tcode = typeCodes[i];
if (tcode == kTVMObjectHandle || tcode == kTVMPackedFuncHandle ||
- tcode == kTVMObjectRValueRefArg || tcode == kTVMModuleHandle) {
+ tcode == kTVMObjectRValueRefArg || tcode == kTVMModuleHandle ||
+ tcode == kTVMNDArrayHandle) {
TVMCbArgToReturn(&arg, &tcode);
}
jobject jarg = tvmRetValueToJava(env, arg, tcode);
diff --git a/python/tvm/_ffi/_ctypes/packed_func.py
b/python/tvm/_ffi/_ctypes/packed_func.py
index acf9776..fd82b26 100644
--- a/python/tvm/_ffi/_ctypes/packed_func.py
+++ b/python/tvm/_ffi/_ctypes/packed_func.py
@@ -306,7 +306,9 @@ C_TO_PY_ARG_SWITCH[ArgTypeCode.MODULE_HANDLE] =
_wrap_arg_func(
_return_module, ArgTypeCode.MODULE_HANDLE
)
C_TO_PY_ARG_SWITCH[ArgTypeCode.DLTENSOR_HANDLE] = lambda x:
_make_array(x.v_handle, True, False)
-C_TO_PY_ARG_SWITCH[ArgTypeCode.NDARRAY_HANDLE] = lambda x:
_make_array(x.v_handle, False, True)
+C_TO_PY_ARG_SWITCH[ArgTypeCode.NDARRAY_HANDLE] = _wrap_arg_func(
+ lambda x: _make_array(x.v_handle, False, True), ArgTypeCode.NDARRAY_HANDLE
+)
_CLASS_MODULE = None
_CLASS_PACKED_FUNC = None
diff --git a/python/tvm/_ffi/_cython/ndarray.pxi
b/python/tvm/_ffi/_cython/ndarray.pxi
index 9fd3aa4..e671ef6 100644
--- a/python/tvm/_ffi/_cython/ndarray.pxi
+++ b/python/tvm/_ffi/_cython/ndarray.pxi
@@ -68,6 +68,10 @@ cdef class NDArrayBase:
def __set__(self, value):
self._set_handle(value)
+ property is_view:
+ def __get__(self):
+ return self.c_is_view != 0
+
@property
def shape(self):
"""Shape of this array"""
diff --git a/python/tvm/_ffi/_cython/packed_func.pxi
b/python/tvm/_ffi/_cython/packed_func.pxi
index 16b1461..0058565 100644
--- a/python/tvm/_ffi/_cython/packed_func.pxi
+++ b/python/tvm/_ffi/_cython/packed_func.pxi
@@ -43,6 +43,7 @@ cdef int tvm_callback(TVMValue* args,
if (tcode == kTVMObjectHandle or
tcode == kTVMPackedFuncHandle or
tcode == kTVMModuleHandle or
+ tcode == kTVMNDArrayHandle or
tcode == kTVMObjectRefArg or
tcode > kTVMExtBegin):
CALL(TVMCbArgToReturn(&value, &tcode))
diff --git a/rust/tvm-rt/src/to_function.rs b/rust/tvm-rt/src/to_function.rs
index a89652b..affd81b 100644
--- a/rust/tvm-rt/src/to_function.rs
+++ b/rust/tvm-rt/src/to_function.rs
@@ -103,8 +103,10 @@ pub trait ToFunction<I, O>: Sized {
value = args_list[i];
tcode = type_codes_list[i];
if tcode == ffi::TVMArgTypeCode_kTVMObjectHandle as c_int
+ || tcode == ffi::TVMArgTypeCode_kTVMObjectRValueRefArg as
c_int
|| tcode == ffi::TVMArgTypeCode_kTVMPackedFuncHandle as
c_int
|| tcode == ffi::TVMArgTypeCode_kTVMModuleHandle as c_int
+ || tcode == ffi::TVMArgTypeCode_kTVMNDArrayHandle as c_int
{
check_call!(ffi::TVMCbArgToReturn(
&mut value as *mut _,
diff --git a/tests/python/unittest/test_runtime_packed_func.py
b/tests/python/unittest/test_runtime_packed_func.py
index 718fe03..b681e4f 100644
--- a/tests/python/unittest/test_runtime_packed_func.py
+++ b/tests/python/unittest/test_runtime_packed_func.py
@@ -333,7 +333,19 @@ def test_numpy_scalar():
assert tvm.testing.echo(np.int64(maxint)) == maxint
+def test_ndarray_args():
+ def check(arr):
+ assert not arr.is_view
+ assert tvm.testing.object_use_count(arr) == 2
+
+ fcheck = tvm.runtime.convert(check)
+ x = tvm.nd.array([1, 2, 3])
+ fcheck(x)
+ assert tvm.testing.object_use_count(x) == 1
+
+
if __name__ == "__main__":
+ test_ndarray_args()
test_numpy_scalar()
test_rvalue_ref()
test_empty_array()
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index 5c9b9d8..80e7d71 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -1216,6 +1216,7 @@ export class Instance implements Disposable {
tcode == ArgTypeCode.TVMObjectHandle ||
tcode == ArgTypeCode.TVMObjectRValueRefArg ||
tcode == ArgTypeCode.TVMPackedFuncHandle ||
+ tcode == ArgTypeCode.TVMNDArrayHandle ||
tcode == ArgTypeCode.TVMModuleHandle
) {
lib.checkCall(
