def test_allocate(self): name = "{}_test_allocate".format(self.__class__.__name__) source = r""" #include <stdio.h> #include "numba/core/runtime/nrt_external.h" NRT_MemInfo* test_nrt_api(NRT_api_functions *nrt, size_t n) { size_t *data = NULL; NRT_MemInfo *mi = nrt->allocate(n); data = nrt->get_data(mi); data[0] = 0xded; data[1] = 0xabc; data[2] = 0xdef; return mi; } """ cdef = "void* test_nrt_api(void *nrt, size_t n);" ffi, mod = self.compile_cffi_module(name, source, cdef) table = self.get_nrt_api_table() numbytes = 3 * np.dtype(np.intp).itemsize out = mod.lib.test_nrt_api(table, numbytes) mi_addr = int(ffi.cast("size_t", out)) mi = nrt.MemInfo(mi_addr) self.assertEqual(mi.refcount, 1) buffer = ffi.buffer(ffi.cast("char [{}]".format(numbytes), mi.data)) arr = np.ndarray(shape=(3, ), dtype=np.intp, buffer=buffer) np.testing.assert_equal(arr, [0xded, 0xabc, 0xdef])
def test_manage_memory(self): name = "{}_test_manage_memory".format(self.__class__.__name__) source = r""" #include <stdio.h> #include "numba/core/runtime/nrt_external.h" int status = 0; void my_dtor(void *ptr) { free(ptr); status = 0xdead; } NRT_MemInfo* test_nrt_api(NRT_api_functions *nrt) { void * data = malloc(10); NRT_MemInfo *mi = nrt->manage_memory(data, my_dtor); nrt->acquire(mi); nrt->release(mi); status = 0xa110c; return mi; } """ cdef = """ void* test_nrt_api(void *nrt); int status; """ ffi, mod = self.compile_cffi_module(name, source, cdef) # Init status is 0 self.assertEqual(mod.lib.status, 0) table = self.get_nrt_api_table() out = mod.lib.test_nrt_api(table) # status is now 0xa110c self.assertEqual(mod.lib.status, 0xa110c) mi_addr = int(ffi.cast("size_t", out)) mi = nrt.MemInfo(mi_addr) self.assertEqual(mi.refcount, 1) del mi # force deallocation on mi # status is now 0xdead self.assertEqual(mod.lib.status, 0xdead)