Use nanobind::ndarray for axpby interface

Author: purepani

Wrappers/Python/cil/framework/data_container.py

@@ -553,8 +553,6 @@ def _axpby(self, a, b, y, out, dtype=numpy.float32, num_threads=NUM_THREADS):
         :type num_threads: int, optional, default 1/2 CPU of the system
         '''
 
-        c_float_p = ctypes.POINTER(ctypes.c_float)
-        c_double_p = ctypes.POINTER(ctypes.c_double)
 
         #convert a and b to numpy arrays and get the reference to the data (length = 1 or ndx.size)
         try:
@@ -593,19 +591,19 @@ def _axpby(self, a, b, y, out, dtype=numpy.float32, num_threads=NUM_THREADS):
             ndb = ndb.astype(dtype, casting='same_kind')
 
         if dtype == numpy.float32:
-            x_p = ndx.ctypes.data_as(c_float_p)
-            y_p = ndy.ctypes.data_as(c_float_p)
-            out_p = ndout.ctypes.data_as(c_float_p)
-            a_p = nda.ctypes.data_as(c_float_p)
-            b_p = ndb.ctypes.data_as(c_float_p)
+            x_p = ndx
+            y_p = ndy
+            out_p = ndout
+            a_p = nda
+            b_p = ndb
             f = cilacc.saxpby
 
         elif dtype == numpy.float64:
-            x_p = ndx.ctypes.data_as(c_double_p)
-            y_p = ndy.ctypes.data_as(c_double_p)
-            out_p = ndout.ctypes.data_as(c_double_p)
-            a_p = nda.ctypes.data_as(c_double_p)
-            b_p = ndb.ctypes.data_as(c_double_p)
+            x_p = ndx
+            y_p = ndy
+            out_p = ndout
+            a_p = nda
+            b_p = ndb
             f = cilacc.daxpby
 
         else:

Wrappers/Python/cil/optimisation/operators/GradientOperator.py

@@ -269,60 +269,8 @@ def adjoint(self, x, out=None):
 
 import ctypes, platform
 from ctypes import util
-# check for the extension
-if platform.system() == 'Linux':
-    dll = 'libcilacc.so'
-elif platform.system() == 'Windows':
-    dll_file = 'cilacc.dll'
-    dll = util.find_library(dll_file)
-elif platform.system() == 'Darwin':
-    dll = 'libcilacc.dylib'
-else:
-    raise ValueError('Not supported platform, ', platform.system())
-
-cilacc = ctypes.cdll.LoadLibrary(dll)
-
-c_float_p = ctypes.POINTER(ctypes.c_float)
-
-cilacc.openMPtest.restypes = ctypes.c_int32
-cilacc.openMPtest.argtypes = [ctypes.c_int32]
-
-cilacc.fdiff4D.restype = ctypes.c_int32
-cilacc.fdiff4D.argtypes = [ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.c_size_t,
-                       ctypes.c_size_t,
-                       ctypes.c_size_t,
-                       ctypes.c_size_t,
-                       ctypes.c_int32,
-                       ctypes.c_int32,
-                       ctypes.c_int32]
-
-cilacc.fdiff3D.restype = ctypes.c_int32
-cilacc.fdiff3D.argtypes = [ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.c_size_t,
-                       ctypes.c_size_t,
-                       ctypes.c_size_t,
-                       ctypes.c_int32,
-                       ctypes.c_int32,
-                       ctypes.c_int32]
-
-cilacc.fdiff2D.restype = ctypes.c_int32
-cilacc.fdiff2D.argtypes = [ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.POINTER(ctypes.c_float),
-                       ctypes.c_size_t,
-                       ctypes.c_size_t,
-                       ctypes.c_int32,
-                       ctypes.c_int32,
-                       ctypes.c_int32]
 
+import cil.cilacc as cilacc
 
 class Gradient_C(LinearOperator):
 
@@ -391,7 +339,7 @@ def ndarray_as_c_pointer(ndx):
     def direct(self, x, out=None):
 
         ndx = np.asarray(x.as_array(), dtype=np.float32, order='C')
-        x_p = Gradient_C.ndarray_as_c_pointer(ndx)
+        x_p = ndx
 
         if out is None:
             out = self.range_geometry().allocate(None)
@@ -404,7 +352,7 @@ def direct(self, x, out=None):
             ndout.insert(ind, out.get_item(0).as_array()) #insert channels dc at correct point for channel data
 
         #pass list of all arguments
-        arg1 = [Gradient_C.ndarray_as_c_pointer(ndout[i]) for i in range(len(ndout))]
+        arg1 = [ndout[i] for i in range(len(ndout))]
         arg2 = [el for el in self.domain_shape]
         args = arg1 + arg2 + [self.bnd_cond, 1, self.num_threads]
         status = self.fd(x_p, *args)
@@ -437,7 +385,7 @@ def adjoint(self, x, out=None):
             out = self.domain_geometry().allocate(None)
 
         ndout = np.asarray(out.as_array(), dtype=np.float32, order='C')
-        out_p = Gradient_C.ndarray_as_c_pointer(ndout)
+        out_p = ndout
 
         if self.split is False:
             ndx = [el.as_array() for el in x.containers]
@@ -450,7 +398,7 @@ def adjoint(self, x, out=None):
             if el != 1:
                 ndx[i]/=el
 
-        arg1 = [Gradient_C.ndarray_as_c_pointer(ndx[i]) for i in range(self.ndim)]
+        arg1 = [ndx[i] for i in range(self.ndim)]
         arg2 = [el for el in self.domain_shape]
         args = arg1 + arg2 + [self.bnd_cond, 0, self.num_threads]
 

src/cilacc/axpby.cpp

@@ -104,7 +104,13 @@ int daxpby_asbv(const double * x, const double * y, double * out, double a, cons
 	}
 	return 0;
 }
-int saxpby(const float * x, const float * y, float * out, const float *a, int a_type, const float* b, int b_type, int64 size, int nThreads)
+
+int saxpby(DataFloatInput x, DataFloatInput y, 
+		DataFloatOutput  out, 
+		DataFloatInput a, int type_a, 
+		DataFloatInput b, int type_b, 
+		int64 size, int nThreads
+	  )
 {
 	//type = 0 float
 	//type = 1 array of floats
@@ -114,20 +120,24 @@ int saxpby(const float * x, const float * y, float * out, const float *a, int a_
 	int nThreads_initial;
 	threads_setup(nThreads, &nThreads_initial);
 
-	if (a_type == 0 && b_type == 0)
-		saxpby_asbs(x, y, out, *a, *b, size, nThreads);
-	else if (a_type == 1 && b_type == 1)
-		saxpby_avbv(x, y, out, a, b, size, nThreads);
-	else if (a_type == 0 && b_type == 1)
-		saxpby_asbv(x, y, out, *a, b, size, nThreads);
-	else if (a_type == 1 && b_type == 0)
-		saxpby_asbv(y, x, out, *b, a, size, nThreads);
+	if (type_a == 0 && type_b == 0)
+		saxpby_asbs(x.data(), y.data(), out.data(), *a.data(), *b.data(), size, nThreads);
+	else if (type_a == 1 && type_b == 1)
+		saxpby_avbv(x.data(), y.data(), out.data(), a.data(), b.data(), size, nThreads);
+	else if (type_a == 0 && type_b == 1)
+		saxpby_asbv(x.data(), y.data(), out.data(), *a.data(), b.data(), size, nThreads);
+	else if (type_a == 1 && type_b == 0)
+		saxpby_asbv(y.data(), x.data(), out.data(), *b.data(), a.data(), size, nThreads);
 
 	omp_set_num_threads(nThreads_initial);
 
 	return 0;
 }
-int daxpby(const double * x, const double * y, double * out, const double *a, int a_type, const double* b, int b_type, int64 size, int nThreads)
+int daxpby(DataDoubleInput x, DataDoubleInput y, 
+		DataDoubleOutput out, 
+		DataDoubleInput a, int type_a, 
+		DataDoubleInput b, int type_b, 
+		int64 size, int nThreads)
 {
 	//type = 0 double
 	//type = 1 array of double
@@ -137,14 +147,14 @@ int daxpby(const double * x, const double * y, double * out, const double *a, in
 	int nThreads_initial;
 	threads_setup(nThreads, &nThreads_initial);
 
-	if (a_type == 0 && b_type == 0)
-		daxpby_asbs(x, y, out, *a, *b, size, nThreads);
-	else if (a_type == 1 && b_type == 1)
-		daxpby_avbv(x, y, out, a, b, size, nThreads);
-	else if (a_type == 0 && b_type == 1)
-		daxpby_asbv(x, y, out, *a, b, size, nThreads);
-	else if (a_type == 1 && b_type == 0)
-		daxpby_asbv(y, x, out, *b, a, size, nThreads);
+	if (type_a == 0 && type_b == 0)
+		daxpby_asbs(x.data(), y.data(), out.data(), *a.data(), *b.data(), size, nThreads);
+	else if (type_a == 1 && type_b == 1)
+		daxpby_avbv(x.data(), y.data(), out.data(), a.data(), b.data(), size, nThreads);
+	else if (type_a == 0 && type_b == 1)
+		daxpby_asbv(x.data(), y.data(), out.data(), *a.data(), b.data(), size, nThreads);
+	else if (type_a == 1 && type_b == 0)
+		daxpby_asbv(y.data(), x.data(), out.data(), *b.data(), a.data(), size, nThreads);
 
 	omp_set_num_threads(nThreads_initial);
 

src/cilacc/include/axpby.h

@@ -20,16 +20,34 @@
 #include <stdio.h>
 #include <omp.h>
 #include "utilities.h"
+#include <nanobind/ndarray.h>
+
+namespace nb = nanobind;
 
 using int64 = long long;
 
+using DataFloatInput = nb::ndarray<const float>;
+using DataDoubleInput = nb::ndarray<const double>;
+
+using DataFloatOutput = nb::ndarray<float>;
+using DataDoubleOutput = nb::ndarray<double>;
+
+
 int saxpby_asbs(const float * x, const float * y, float * out, float a, float b, int64 size, int nThreads);
 int saxpby_avbv(const float * x, const float * y, float * out, const float * a, const float * b, int64 size, int nThreads);
 int saxpby_asbv(const float * x, const float * y, float * out, float a, const float * b, int64 size, int nThreads);
 int daxpby_asbs(const double * x, const double * y, double * out, double a, double b, int64 size, int nThreads);
 int daxpby_avbv(const double * x, const double * y, double * out, const double * a, const double * b, int64 size, int nThreads);
 int daxpby_asbv(const double * x, const double * y, double * out, double a, const double * b, int64 size, int nThreads);
 
-int saxpby(const float * x, const float * y, float * out, const float * a, int type_a, const float * b, int type_b, int64 size, int nThreads);
-int daxpby(const double * x, const double * y, double * out, const double * a, int type_a, const double * b, int type_b, int64 size, int nThreads);
+int saxpby(DataFloatInput x, DataFloatInput y, 
+		DataFloatOutput  out, 
+		DataFloatInput a, int type_a, 
+		DataFloatInput b, int type_b, 
+		int64 size, int nThreads);
+int daxpby(DataDoubleInput x, DataDoubleInput y, 
+		DataDoubleOutput out, 
+		DataDoubleInput a, int type_a, 
+		DataDoubleInput b, int type_b, 
+		int64 size, int nThreads);