From 2a003fdbeafe9a5580959e384e7ece4393e7dd7f Mon Sep 17 00:00:00 2001 From: tqchen Date: Fri, 23 Oct 2020 13:54:58 -0400 Subject: [PATCH 1/3] [FFI][BUGFIX] Fix leak when Packed callback arg is ndarray. Co-authored-by: Matthew Brookhart --- python/tvm/_ffi/_ctypes/packed_func.py | 4 +++- python/tvm/_ffi/_cython/ndarray.pxi | 4 ++++ python/tvm/_ffi/_cython/packed_func.pxi | 1 + tests/python/unittest/test_runtime_packed_func.py | 12 ++++++++++++ 4 files changed, 20 insertions(+), 1 deletion(-) diff --git a/python/tvm/_ffi/_ctypes/packed_func.py b/python/tvm/_ffi/_ctypes/packed_func.py index acf9776d9b8b..fd82b263e2dd 100644 --- a/python/tvm/_ffi/_ctypes/packed_func.py +++ b/python/tvm/_ffi/_ctypes/packed_func.py @@ -306,7 +306,9 @@ def _get_global_func(name, allow_missing=False): _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 9fd3aa43841f..e671ef626205 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 16b146119f0d..00585659ab76 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/tests/python/unittest/test_runtime_packed_func.py b/tests/python/unittest/test_runtime_packed_func.py index 718fe03d5c16..b681e4fc25d7 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() From fb96d5cf52cd3328a734014fe8948d2ac08f4738 Mon Sep 17 00:00:00 2001 From: tqchen Date: Fri, 23 Oct 2020 14:17:59 -0400 Subject: [PATCH 2/3] Fix for rust ts and jvm --- jvm/native/src/main/native/org_apache_tvm_native_c_api.cc | 3 ++- rust/tvm-rt/src/to_function.rs | 2 ++ web/src/runtime.ts | 1 + 3 files changed, 5 insertions(+), 1 deletion(-) 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 6fc316ca8739..e3ea4b9c3766 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/rust/tvm-rt/src/to_function.rs b/rust/tvm-rt/src/to_function.rs index a89652b0378c..2f542e18a9b3 100644 --- a/rust/tvm-rt/src/to_function.rs +++ b/rust/tvm-rt/src/to_function.rs @@ -103,8 +103,10 @@ pub trait ToFunction: Sized { value = args_list[i]; tcode = type_codes_list[i]; if tcode == ffi::TVMArgTypeCode_kTVMObjectHandle as c_int + || tcode == ffi::TVMArgTypeCode_kkTVMObjectRValueRefArg as c_int || tcode == ffi::TVMArgTypeCode_kTVMPackedFuncHandle as c_int || tcode == ffi::TVMArgTypeCode_kTVMModuleHandle as c_int + || tcode == ffi::TVMArgTypeCode_kTVMNDArrayFuncHandle as c_int { check_call!(ffi::TVMCbArgToReturn( &mut value as *mut _, diff --git a/web/src/runtime.ts b/web/src/runtime.ts index 5c9b9d8181d7..80e7d71f06ad 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( From 55c6e9d88b9045ed916b3eaa9c7a4213c7bbdf86 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Fri, 23 Oct 2020 19:30:47 -0400 Subject: [PATCH 3/3] Update rust/tvm-rt/src/to_function.rs Co-authored-by: Junru Shao --- rust/tvm-rt/src/to_function.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/tvm-rt/src/to_function.rs b/rust/tvm-rt/src/to_function.rs index 2f542e18a9b3..affd81b0e7ed 100644 --- a/rust/tvm-rt/src/to_function.rs +++ b/rust/tvm-rt/src/to_function.rs @@ -103,10 +103,10 @@ pub trait ToFunction: Sized { value = args_list[i]; tcode = type_codes_list[i]; if tcode == ffi::TVMArgTypeCode_kTVMObjectHandle as c_int - || tcode == ffi::TVMArgTypeCode_kkTVMObjectRValueRefArg 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_kTVMNDArrayFuncHandle as c_int + || tcode == ffi::TVMArgTypeCode_kTVMNDArrayHandle as c_int { check_call!(ffi::TVMCbArgToReturn( &mut value as *mut _,