[python-arrayfire] 234/250: Adding function to interop with pycuda

[Ghislain Vaillant]

This is an automated email from the git hooks/post-receive script.

ghisvail-guest pushed a commit to branch debian/master
in repository python-arrayfire.

commit 196f1103cb923fe77450e340a18a2f439f06d127
Author: Pavan Yalamanchili <pavan at arrayfire.com>
Date:   Sun Mar 20 01:03:38 2016 -0400

    Adding function to interop with pycuda
---
 arrayfire/__init__.py |  5 +++++
 arrayfire/interop.py  | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/arrayfire/__init__.py b/arrayfire/__init__.py
index 79959bb..40cc231 100644
--- a/arrayfire/__init__.py
+++ b/arrayfire/__init__.py
@@ -35,6 +35,11 @@ where name is one of 'cuda', 'opencl' or 'cpu'
 
 """
 
+try:
+    import pycuda.autoinit
+except:
+    pass
+
 from .library    import *
 from .array      import *
 from .data       import *
diff --git a/arrayfire/interop.py b/arrayfire/interop.py
index 8095629..deaaa99 100644
--- a/arrayfire/interop.py
+++ b/arrayfire/interop.py
@@ -13,9 +13,11 @@ Interop with other python packages.
 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 @@ try:
     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

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-science/packages/python-arrayfire.git