diff --git a/apps/from_numpy/Makefile b/apps/from_numpy/Makefile new file mode 100644 index 0000000..17c32d9 --- /dev/null +++ b/apps/from_numpy/Makefile @@ -0,0 +1,2 @@ +all: numpy_dlpack.c + gcc -I../../include -shared -o libmain.so -fPIC numpy_dlpack.c diff --git a/apps/from_numpy/main.py b/apps/from_numpy/main.py new file mode 100644 index 0000000..011421a --- /dev/null +++ b/apps/from_numpy/main.py @@ -0,0 +1,103 @@ +from __future__ import print_function + +import numpy as np +import gc +import ctypes + +libmain = ctypes.cdll.LoadLibrary("libmain.so") + +class DLContext(ctypes.Structure): + _fields_ = [("device_type", ctypes.c_int), + ("device_id", ctypes.c_int)] + +class DLDataType(ctypes.Structure): + _fields_ = [("type_code", ctypes.c_uint8), + ("bits", ctypes.c_uint8), + ("lanes", ctypes.c_uint16)] + TYPE_MAP = { + "bool": (1, 1, 1), + "int32": (0, 32, 1), + "int64": (0, 64, 1), + "uint32": (1, 32, 1), + "uint64": (1, 64, 1), + "float32": (2, 32, 1), + "float64": (2, 64, 1), + } + +class DLTensor(ctypes.Structure): + _fields_ = [("data", ctypes.c_void_p), + ("ctx", DLContext), + ("ndim", ctypes.c_int), + ("dtype", DLDataType), + ("shape", ctypes.POINTER(ctypes.c_int64)), + ("strides", ctypes.POINTER(ctypes.c_int64)), + ("byte_offset", ctypes.c_uint64)] + +class DLManagedTensor(ctypes.Structure): + pass + +DLManagedTensorHandle = ctypes.POINTER(DLManagedTensor) + +DeleterFunc = ctypes.CFUNCTYPE(None, DLManagedTensorHandle) + +DLManagedTensor._fields_ = [("dl_tensor", DLTensor), + ("manager_ctx", ctypes.c_void_p), + ("deleter", DeleterFunc)] + +def display(array): + print("data =", hex(array.ctypes.data_as(ctypes.c_void_p).value)) + print("dtype =", array.dtype) + print("ndim =", array.ndim) + print("shape =", array.shape) + print("strides =", array.strides) + +def make_manager_ctx(obj): + pyobj = ctypes.py_object(obj) + void_p = ctypes.c_void_p.from_buffer(pyobj) + ctypes.pythonapi.Py_IncRef(pyobj) + return void_p + +# N.B.: In practice, one should ensure that this function +# is not destructed before the numpy array is destructed. +@DeleterFunc +def dl_managed_tensor_deleter(dl_managed_tensor_handle): + void_p = dl_managed_tensor_handle.contents.manager_ctx + pyobj = ctypes.cast(void_p, ctypes.py_object) + print("Deleting:") + display(pyobj.value) + ctypes.pythonapi.Py_DecRef(pyobj) + print("Done") + +def make_dl_tensor(array): + # You may check array.flags here, e.g. array.flags['C_CONTIGUOUS'] + ndim = array.ndim + dl_tensor = DLTensor() + dl_tensor.data = array.ctypes.data_as(ctypes.c_void_p) + dl_tensor.ctx = DLContext(1, 0) + dl_tensor.ndim = array.ndim + dl_tensor.dtype = DLDataType.TYPE_MAP[str(array.dtype)] + # For 0-dim ndarrays, strides and shape will be NULL + dl_tensor.shape = array.ctypes.shape_as(ctypes.c_int64) + dl_tensor.strides = array.ctypes.strides_as(ctypes.c_int64) + dl_tensor.byte_offset = 0 + return dl_tensor + +def main(): + array = np.random.rand(3, 1, 30).astype("float32") + print("Created:") + display(array) + c_obj = DLManagedTensor() + c_obj.dl_tensor = make_dl_tensor(array) + c_obj.manager_ctx = make_manager_ctx(array) + c_obj.deleter = dl_managed_tensor_deleter + print("-------------------------") + del array + gc.collect() + libmain.Give(c_obj) + print("-------------------------") + del c_obj + gc.collect() + libmain.Finalize() + +if __name__ == "__main__": + main() diff --git a/apps/from_numpy/numpy_dlpack.c b/apps/from_numpy/numpy_dlpack.c new file mode 100644 index 0000000..6e84453 --- /dev/null +++ b/apps/from_numpy/numpy_dlpack.c @@ -0,0 +1,45 @@ +#include +#include + +DLManagedTensor given; + +void display(DLManagedTensor a) { + puts("On C side:"); + int i; + int ndim = a.dl_tensor.ndim; + printf("data = %p\n", a.dl_tensor.data); + printf("ctx = (device_type = %d, device_id = %d)\n", + (int) a.dl_tensor.ctx.device_type, + (int) a.dl_tensor.ctx.device_id); + printf("dtype = (code = %d, bits = %d, lanes = %d)\n", + (int) a.dl_tensor.dtype.code, + (int) a.dl_tensor.dtype.bits, + (int) a.dl_tensor.dtype.lanes); + printf("ndim = %d\n", + (int) a.dl_tensor.ndim); + printf("shape = ("); + for (i = 0; i < ndim; ++i) { + if (i != 0) { + printf(", "); + } + printf("%d", (int) a.dl_tensor.shape[i]); + } + printf(")\n"); + printf("strides = ("); + for (i = 0; i < ndim; ++i) { + if (i != 0) { + printf(", "); + } + printf("%d", (int) a.dl_tensor.strides[i]); + } + printf(")\n"); +} + +void Give(DLManagedTensor dl_managed_tensor) { + display(dl_managed_tensor); + given = dl_managed_tensor; +} + +void Finalize() { + given.deleter(&given); +}