From fc69f68f05e164140ad36f8a43e2e95c1d5cab10 Mon Sep 17 00:00:00 2001 From: Tianqi Chen Date: Sat, 24 Oct 2020 09:56:31 -0400 Subject: [PATCH] [FFI][BUGFIX] Fix memory leak when Pac callback argument is NDArray (#6744) * [FFI][BUGFIX] Fix leak when Packed callback arg is ndarray. Co-authored-by: Matthew Brookhart * Fix for rust ts and jvm * Update rust/tvm-rt/src/to_function.rs Co-authored-by: Junru Shao Co-authored-by: Matthew Brookhart Co-authored-by: Junru Shao --- .../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 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/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/rust/tvm-rt/src/to_function.rs b/rust/tvm-rt/src/to_function.rs index a89652b0378c..affd81b0e7ed 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_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 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() 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(