FEAT: Adding the ability to get, lock, and unlock raw device pointers

Author: pavanky

arrayfire/array.py

@@ -20,11 +20,15 @@
 from .index import *
 from .index import _Index4
 
-def _create_array(buf, numdims, idims, dtype):
+def _create_array(buf, numdims, idims, dtype, is_device):
     out_arr = ct.c_void_p(0)
     c_dims = dim4(idims[0], idims[1], idims[2], idims[3])
-    safe_call(backend.get().af_create_array(ct.pointer(out_arr), ct.c_void_p(buf),
-                                            numdims, ct.pointer(c_dims), dtype.value))
+    if (not is_device):
+        safe_call(backend.get().af_create_array(ct.pointer(out_arr), ct.c_void_p(buf),
+                                                numdims, ct.pointer(c_dims), dtype.value))
+    else:
+        safe_call(backend.get().af_device_array(ct.pointer(out_arr), ct.c_void_p(buf),
+                                                numdims, ct.pointer(c_dims), dtype.value))
     return out_arr
 
 def _create_empty_array(numdims, idims, dtype):
@@ -348,7 +352,7 @@ class Array(BaseArray):
 
     """
 
-    def __init__(self, src=None, dims=(0,), dtype=None):
+    def __init__(self, src=None, dims=(0,), dtype=None, is_device=False):
 
         super(Array, self).__init__()
 
@@ -385,7 +389,8 @@ def __init__(self, src=None, dims=(0,), dtype=None):
                 _type_char = tmp.typecode
                 numdims, idims = _get_info(dims, buf_len)
             elif isinstance(src, int) or isinstance(src, ct.c_void_p):
-                buf = src
+                buf = src if not isinstance(src, ct.c_void_p) else src.value
+
                 numdims, idims = _get_info(dims, buf_len)
 
                 elements = 1
@@ -407,7 +412,7 @@ def __init__(self, src=None, dims=(0,), dtype=None):
                 type_char != _type_char):
                 raise TypeError("Can not create array of requested type from input data type")
 
-            self.arr = _create_array(buf, numdims, idims, to_dtype[_type_char])
+            self.arr = _create_array(buf, numdims, idims, to_dtype[_type_char], is_device)
 
         else:
 

arrayfire/device.py

@@ -138,3 +138,55 @@ def device_gc():
     Ask the garbage collector to free all unlocked memory
     """
     safe_call(backend.get().af_device_gc())
+
+def get_device_ptr(a):
+    """
+    Get the raw device pointer of an array
+
+    Parameters
+    ----------
+    a: af.Array
+       - A multi dimensional arrayfire array.
+
+    Returns
+    -------
+        - internal device pointer held by a
+
+    Note
+    -----
+        - The device pointer of `a` is not freed by memory manager until `unlock_device_ptr()` is called.
+        - This function enables the user to interoperate arrayfire with other CUDA/OpenCL/C libraries.
+
+    """
+    ptr = ct.c_void_p(0)
+    safe_call(backend.get().af_get_device_ptr(ct.pointer(ptr), a.arr))
+    return ptr
+
+def lock_device_ptr(a):
+    """
+    Ask arrayfire to not perform garbage collection on raw data held by an array.
+
+    Parameters
+    ----------
+    a: af.Array
+       - A multi dimensional arrayfire array.
+
+    Note
+    -----
+        - The device pointer of `a` is not freed by memory manager until `unlock_device_ptr()` is called.
+    """
+    ptr = ct.c_void_p(0)
+    safe_call(backend.get().af_lock_device_ptr(a.arr))
+
+def unlock_device_ptr(a):
+    """
+    Tell arrayfire to resume garbage collection on raw data held by an array.
+
+    Parameters
+    ----------
+    a: af.Array
+       - A multi dimensional arrayfire array.
+
+    """
+    ptr = ct.c_void_p(0)
+    safe_call(backend.get().af_unlock_device_ptr(a.arr))

tests/simple/device.py

@@ -40,4 +40,13 @@ def simple_device(verbose=False):
 
     af.set_device(dev)
 
+    a = af.randu(10,10)
+    display_func(a)
+    dev_ptr = af.get_device_ptr(a)
+    print_func(dev_ptr)
+    b = af.Array(src=dev_ptr, dims=a.dims(), dtype=a.dtype(), is_device=True)
+    display_func(b)
+    af.lock_device_ptr(b)
+    af.unlock_device_ptr(b)
+
 _util.tests['device'] = simple_device