Adding function to interop with pycuda

Author: pavanky

arrayfire/__init__.py

@@ -35,6 +35,11 @@
 
 """
 
+try:
+    import pycuda.autoinit
+except:
+    pass
+
 from .library    import *
 from .array      import *
 from .data       import *

arrayfire/interop.py

@@ -13,9 +13,11 @@
 This module provides interoperability with the following python packages.
 
      1. numpy
+     2. pycuda
 """
 
 from .array import *
+from .device import lock_array
 
 try:
     import numpy as np
@@ -60,3 +62,48 @@ def np_to_af_array(np_arr):
     from_ndarray = np_to_af_array
 except:
     AF_NUMPY_FOUND=False
+
+try:
+    import pycuda.gpuarray as CudaArray
+    AF_PYCUDA_FOUND=True
+
+    def pycuda_to_af_array(pycu_arr):
+        """
+        Convert pycuda.gpuarray to arrayfire.Array
+
+        Parameters
+        -----------
+        pycu_arr  : pycuda.GPUArray()
+
+        Returns
+        ----------
+        af_arr    : arrayfire.Array()
+        """
+        if (pycu_arr.flags.f_contiguous):
+            res = Array(pycu_arr.ptr, pycu_arr.shape, pycu_arr.dtype.char, is_device=True)
+            lock_array(res)
+            return res
+        elif (pycu_arr.flags.c_contiguous):
+            if pycu_arr.ndim == 1:
+                return Array(pycu_arr.ptr, pycu_arr.shape, pycu_arr.dtype.char, is_device=True)
+            elif pycu_arr.ndim == 2:
+                shape = (pycu_arr.shape[1], pycu_arr.shape[0])
+                res = Array(pycu_arr.ptr, shape, pycu_arr.dtype.char, is_device=True)
+                lock_array(res)
+                return reorder(res, 1, 0)
+            elif pycu_arr.ndim == 3:
+                shape = (pycu_arr.shape[2], pycu_arr.shape[1], pycu_arr.shape[0])
+                res = Array(pycu_arr.ptr, shape, pycu_arr.dtype.char, is_device=True)
+                lock_array(res)
+                return reorder(res, 2, 1, 0)
+            elif pycu_arr.ndim == 4:
+                shape = (pycu_arr.shape[3], pycu_arr.shape[2], pycu_arr.shape[1], pycu_arr.shape[0])
+                res = Array(pycu_arr.ptr, shape, pycu_arr.dtype.char, is_device=True)
+                lock_array(res)
+                return reorder(res, 3, 2, 1, 0)
+            else:
+                raise RuntimeError("Unsupported ndim")
+        else:
+            return pycuda_to_af_array(pycu_arr.copy())
+except:
+    AF_PYCUDA_FOUND=False