[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 <[email protected]> * Fix for rust ts and jvm * Update rust/tvm-rt/src/to_function.rs Co-authored-by: Junru Shao <[email protected]> Co-authored-by: Matthew Brookhart <[email protected]> Co-authored-by: Junru Shao <[email protected]>
apache · Oct 24, 2020 · fc69f68 · fc69f68
1 parent bb4179e
commit fc69f68
Show file tree

Hide file tree

Showing 7 changed files with 25 additions and 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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(