Start work on CFloat32Data3DGPU to allow persistent/external GPU memory

5 files changed, 113 insertions, 38 deletions

diff --git a/python/astra/PyIncludes.pxd b/python/astra/PyIncludes.pxd
index 39f9039..d5da17c 100644
--- a/python/astra/PyIncludes.pxd
+++ b/python/astra/PyIncludes.pxd
@@ -172,6 +172,13 @@ IF HAVE_CUDA==True:
 	cdef extern from "astra/CudaProjector2D.h" namespace "astra":
 		cdef cppclass CCudaProjector2D
 
+	cdef extern from "astra/Float32Data3DGPU.h" namespace "astraCUDA3d":
+		cdef cppclass MemHandle3D:
+			pass
+
+	cdef extern from "astra/Float32Data3DGPU.h" namespace "astraCUDA3d":
+		cdef MemHandle3D wrapHandle(float *D_ptr, unsigned int x, unsigned int y, unsigned int z, unsigned int pitch)
+
 
 cdef extern from "astra/SparseMatrix.h" namespace "astra":
 	cdef cppclass CSparseMatrix:
@@ -184,14 +191,18 @@ cdef extern from "astra/SparseMatrix.h" namespace "astra":
 		unsigned int* m_piColIndices
 		unsigned long* m_plRowStarts
 
-cdef extern from "astra/Float32Data3DMemory.h" namespace "astra":
-	cdef cppclass CFloat32Data3DMemory:
-		CFloat32Data3DMemory()
+cdef extern from "astra/Float32Data3D.h" namespace "astra":
+	cdef cppclass CFloat32Data3D:
 		bool isInitialized()
 		int getSize()
 		int getWidth()
 		int getHeight()
 		int getDepth()
+
+
+cdef extern from "astra/Float32Data3DMemory.h" namespace "astra":
+	cdef cppclass CFloat32Data3DMemory(CFloat32Data3D):
+		CFloat32Data3DMemory()
 		void updateStatistics()
 		float32 *getData()
 		float32 ***getData3D()
@@ -228,8 +239,6 @@ cdef extern from "astra/Float32VolumeData3DMemory.h" namespace "astra":
 		int getSliceCount()
 		bool isInitialized()
 
-
-
 cdef extern from "astra/ParallelProjectionGeometry3D.h" namespace "astra":
 	cdef cppclass CParallelProjectionGeometry3D:
 		CParallelProjectionGeometry3D()
@@ -260,12 +269,23 @@ cdef extern from "astra/Float32ProjectionData3DMemory.h" namespace "astra":
 		int getAngleCount()
 		bool isInitialized()
 
-cdef extern from "astra/Float32Data3D.h" namespace "astra":
-	cdef cppclass CFloat32Data3D:
-		CFloat32Data3D()
-		bool isInitialized()
-		int getSize()
-		int getWidth()
-		int getHeight()
-		int getDepth()
-		void updateStatistics()
+IF HAVE_CUDA==True:
+	cdef extern from "astra/Float32VolumeData3DGPU.h" namespace "astra":
+		cdef cppclass CFloat32VolumeData3DGPU:
+			CFloat32VolumeData3DGPU(CVolumeGeometry3D*, MemHandle3D)
+			CVolumeGeometry3D* getGeometry()
+			void changeGeometry(CVolumeGeometry3D*)
+			int getRowCount()
+			int getColCount()
+			int getSliceCount()
+			bool isInitialized()
+
+	cdef extern from "astra/Float32ProjectionData3DGPU.h" namespace "astra":
+		cdef cppclass CFloat32ProjectionData3DGPU:
+			CFloat32ProjectionData3DGPU(CProjectionGeometry3D*, MemHandle3D)
+			CProjectionGeometry3D* getGeometry()
+			void changeGeometry(CProjectionGeometry3D*)
+			int getRowCount()
+			int getColCount()
+			int getSliceCount()
+			bool isInitialized()
diff --git a/python/astra/data3d.py b/python/astra/data3d.py
index 9c2bea4..a825700 100644
--- a/python/astra/data3d.py
+++ b/python/astra/data3d.py
@@ -26,6 +26,8 @@
 from . import data3d_c as d
 import numpy as np
 
+from .pythonutils import GPULink
+
 def create(datatype,geometry,data=None):
     """Create a 3D object.
 
@@ -52,11 +54,11 @@ def link(datatype, geometry, data):
     :returns: :class:`int` -- the ID of the constructed object.
     
     """
-    if not isinstance(data,np.ndarray):
+    if not isinstance(data,np.ndarray) and not isinstance(data,GPULink):
         raise ValueError("Input should be a numpy array")
-    if not data.dtype==np.float32:
+    if not isinstance(data,GPULink) and not data.dtype==np.float32:
         raise ValueError("Numpy array should be float32")
-    if not (data.flags['C_CONTIGUOUS'] and data.flags['ALIGNED']):
+    if not isinstance(data,GPULink) and not (data.flags['C_CONTIGUOUS'] and data.flags['ALIGNED']):
         raise ValueError("Numpy array should be C_CONTIGUOUS and ALIGNED")
     return d.create(datatype,geometry,data,True)
 
diff --git a/python/astra/data3d_c.pyx b/python/astra/data3d_c.pyx
index 3934f22..56247de 100644
--- a/python/astra/data3d_c.pyx
+++ b/python/astra/data3d_c.pyx
@@ -45,12 +45,17 @@ from .PyXMLDocument cimport XMLDocument
 cimport utils
 from .utils import wrap_from_bytes
 
-from .pythonutils import geom_size
+from .pythonutils import geom_size, GPULink
 
 import operator
 
 from six.moves import reduce
 
+include "config.pxi"
+
+cdef extern from "Python.h":
+    void* PyLong_AsVoidPtr(object)
+
 
 cdef CData3DManager * man3d = <CData3DManager * >PyData3DManager.getSingletonPtr()
 
@@ -65,12 +70,19 @@ def create(datatype,geometry,data=None, link=False):
     cdef Config *cfg
     cdef CVolumeGeometry3D * pGeometry
     cdef CProjectionGeometry3D * ppGeometry
-    cdef CFloat32Data3DMemory * pDataObject3D
+    cdef CFloat32Data3D * pDataObject3D
     cdef CConeProjectionGeometry3D* pppGeometry
-    cdef CFloat32CustomMemory * pCustom
-
-    if link and data.shape!=geom_size(geometry):
-        raise Exception("The dimensions of the data do not match those specified in the geometry.")
+    cdef CFloat32CustomMemory * pCustom = NULL
+    IF HAVE_CUDA==True:
+        cdef MemHandle3D hnd
+
+    if link:
+        if isinstance(data, GPULink):
+            s = geom_size(geometry)
+            if geom_size(geometry) != ( data.z, data.y, data.x ):
+                raise Exception("The dimensions of the data do not match those specified in the geometry.")
+        elif data.shape!=geom_size(geometry):
+            raise Exception("The dimensions of the data do not match those specified in the geometry.")
 
     if datatype == '-vol':
         cfg = utils.dictToConfig(six.b('VolumeGeometry'), geometry)
@@ -80,10 +92,18 @@ def create(datatype,geometry,data=None, link=False):
             del pGeometry
             raise Exception('Geometry class not initialized.')
         if link:
-            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
-            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32VolumeData3DMemory(pGeometry, pCustom)
+            if isinstance(data, GPULink):
+                IF HAVE_CUDA==True:
+                    s = geom_size(geometry)
+                    hnd = wrapHandle(<float*>PyLong_AsVoidPtr(data.ptr), data.x, data.y, data.z, data.pitch/4)
+                    pDataObject3D = <CFloat32Data3D * > new CFloat32VolumeData3DGPU(pGeometry, hnd)
+                ELSE:
+                    raise NotImplementedError("CUDA support is not enabled in ASTRA")
+            else:
+                pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+                pDataObject3D = <CFloat32Data3D * > new CFloat32VolumeData3DMemory(pGeometry, pCustom)
         else:
-            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32VolumeData3DMemory(pGeometry)
+            pDataObject3D = <CFloat32Data3D * > new CFloat32VolumeData3DMemory(pGeometry)
         del cfg
         del pGeometry
     elif datatype == '-sino' or datatype == '-proj3d' or datatype == '-sinocone':
@@ -105,8 +125,16 @@ def create(datatype,geometry,data=None, link=False):
             del ppGeometry
             raise Exception('Geometry class not initialized.')
         if link:
-            pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
-            pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(ppGeometry, pCustom)
+            if isinstance(data, GPULink):
+                IF HAVE_CUDA==True:
+                    s = geom_size(geometry)
+                    hnd = wrapHandle(<float*>PyLong_AsVoidPtr(data.ptr), data.x, data.y, data.z, data.pitch/4)
+                    pDataObject3D = <CFloat32Data3D * > new CFloat32ProjectionData3DGPU(ppGeometry, hnd)
+                ELSE:
+                    raise NotImplementedError("CUDA support is not enabled in ASTRA")
+            else:
+                pCustom = <CFloat32CustomMemory*> new CFloat32CustomPython(data)
+                pDataObject3D = <CFloat32Data3D * > new CFloat32ProjectionData3DMemory(ppGeometry, pCustom)
         else:
             pDataObject3D = <CFloat32Data3DMemory * > new CFloat32ProjectionData3DMemory(ppGeometry)
         del ppGeometry
@@ -118,8 +146,7 @@ def create(datatype,geometry,data=None, link=False):
         del pDataObject3D
         raise Exception("Couldn't initialize data object.")
 
-    if not link: fillDataObject(pDataObject3D, data)
-
+    if not link: fillDataObject(dynamic_cast_mem(pDataObject3D), data)
 
     return man3d.store(<CFloat32Data3D*>pDataObject3D)
 
diff --git a/python/astra/experimental.pyx b/python/astra/experimental.pyx
index b6c7881..0af3118 100644
--- a/python/astra/experimental.pyx
+++ b/python/astra/experimental.pyx
@@ -36,12 +36,20 @@ IF HAVE_CUDA==True:
 
     cdef extern from "astra/CompositeGeometryManager.h" namespace "astra":
         cdef cppclass CCompositeGeometryManager:
-            bool doFP(CProjector3D *, vector[CFloat32VolumeData3DMemory *], vector[CFloat32ProjectionData3DMemory *])
-            bool doBP(CProjector3D *, vector[CFloat32VolumeData3DMemory *], vector[CFloat32ProjectionData3DMemory *])
+            bool doFP(CProjector3D *, vector[CFloat32VolumeData3D *], vector[CFloat32ProjectionData3D *])
+            bool doBP(CProjector3D *, vector[CFloat32VolumeData3D *], vector[CFloat32ProjectionData3D *])
 
     cdef extern from *:
-        CFloat32VolumeData3DMemory * dynamic_cast_vol_mem "dynamic_cast<astra::CFloat32VolumeData3DMemory*>" (CFloat32Data3D * ) except NULL
-        CFloat32ProjectionData3DMemory * dynamic_cast_proj_mem "dynamic_cast<astra::CFloat32ProjectionData3DMemory*>" (CFloat32Data3D * ) except NULL
+        CFloat32VolumeData3D * dynamic_cast_vol_mem "dynamic_cast<astra::CFloat32VolumeData3D*>" (CFloat32Data3D * ) except NULL
+        CFloat32ProjectionData3D * dynamic_cast_proj_mem "dynamic_cast<astra::CFloat32ProjectionData3D*>" (CFloat32Data3D * ) except NULL
+
+    cdef extern from "astra/Float32ProjectionData3D.h" namespace "astra":
+        cdef cppclass CFloat32ProjectionData3D:
+            bool isInitialized()
+    cdef extern from "astra/Float32VolumeData3D.h" namespace "astra":
+        cdef cppclass CFloat32VolumeData3D:
+            bool isInitialized()
+
 
     cimport PyProjector3DManager
     from .PyProjector3DManager cimport CProjector3DManager
@@ -52,9 +60,9 @@ IF HAVE_CUDA==True:
     cdef CData3DManager * man3d = <CData3DManager * >PyData3DManager.getSingletonPtr()
 
     def do_composite(projector_id, vol_ids, proj_ids, t):
-        cdef vector[CFloat32VolumeData3DMemory *] vol
-        cdef CFloat32VolumeData3DMemory * pVolObject
-        cdef CFloat32ProjectionData3DMemory * pProjObject
+        cdef vector[CFloat32VolumeData3D *] vol
+        cdef CFloat32VolumeData3D * pVolObject
+        cdef CFloat32ProjectionData3D * pProjObject
         for v in vol_ids:
             pVolObject = dynamic_cast_vol_mem(man3d.get(v))
             if pVolObject == NULL:
@@ -62,7 +70,7 @@ IF HAVE_CUDA==True:
             if not pVolObject.isInitialized():
                 raise Exception("Data object not initialized properly")
             vol.push_back(pVolObject)
-        cdef vector[CFloat32ProjectionData3DMemory *] proj
+        cdef vector[CFloat32ProjectionData3D *] proj
         for v in proj_ids:
             pProjObject = dynamic_cast_proj_mem(man3d.get(v))
             if pProjObject == NULL:
diff --git a/python/astra/pythonutils.py b/python/astra/pythonutils.py
index 3bd3321..27fa8fd 100644
--- a/python/astra/pythonutils.py
+++ b/python/astra/pythonutils.py
@@ -61,3 +61,21 @@ def geom_size(geom, dim=None):
         s = s[dim]
 
     return s
+
+class GPULink(object):
+    """Utility class for astra.data3d.link with a CUDA pointer
+
+    The CUDA pointer ptr must point to an array of floats.
+
+    x is the fastest-changing coordinate, z the slowest-changing.
+
+    pitch is the width in bytes of the memory block. For a contiguous
+    memory block, pitch is equal to sizeof(float) * x. For a memory block
+    allocated by cudaMalloc3D, pitch is the pitch as returned by cudaMalloc3D.
+    """
+    def __init__(self, ptr, x, y, z, pitch):
+        self.ptr = ptr
+        self.x = x
+        self.y = y
+        self.z = z
+        self.pitch = pitch