Feectools with cupy

feectools/core/bsplines.py

@@ -17,6 +17,7 @@
 """
 import cunumpy as xp
 from cunumpy.xp import array_backend
+import numpy as np
 
 from feectools.core.bsplines_kernels import (find_span_p,
                                           find_spans_p,
@@ -352,9 +353,9 @@ def collocation_matrix(knots, degree, periodic, normalization, xgrid, out=None,
         if periodic:
             nb -= degree + 1 - multiplicity
 
-        out = xp.zeros((int(xgrid.shape[0]), int(nb)), dtype=float)
+        out = np.zeros((int(xgrid.shape[0]), int(nb)), dtype=float)
     else:
-        assert out.shape == ((int(xgrid.shape[0]), int(nb))) and out.dtype == xp.dtype('float')
+        assert out.shape == ((int(xgrid.shape[0]), int(nb))) and out.dtype == np.dtype('float')
 
     bool_normalization = normalization == "M"
     multiplicity = int(multiplicity)
@@ -437,9 +438,9 @@ def histopolation_matrix(knots, degree, periodic, normalization, xgrid, multipli
 
     if out is None:
         if periodic:
-            out = xp.zeros((len(xgrid), len(knots) - 2 * int(degree) - 2 + int(multiplicity)), dtype=float)
+            out = np.zeros((len(xgrid), len(knots) - 2 * int(degree) - 2 + int(multiplicity)), dtype=float)
         else:
-            out = xp.zeros((len(xgrid) - 1, len(elevated_knots) - (int(degree) + 1) - 1 - 1), dtype=float)
+            out = np.zeros((len(xgrid) - 1, len(elevated_knots) - (int(degree) + 1) - 1 - 1), dtype=float)
     else:
         if periodic:
             assert out.shape == (len(xgrid), len(knots) - 2 * degree - 2 + multiplicity)
@@ -514,10 +515,15 @@ def greville(knots, degree, periodic, out=None, multiplicity=1):
         Abscissas of all Greville points.
 
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    # Greville points are index arrays, keep on NumPy
+    if isinstance(knots, (list, tuple)):
+        knots = np.asarray(knots, dtype=float)
+    if hasattr(knots, 'get'):
+        knots = knots.get()  # Convert CuPy to NumPy
+    knots = np.ascontiguousarray(knots, dtype=float)
     if out is None:
         n = len(knots) - 2 * degree - 2 + multiplicity if periodic else len(knots) - degree - 1
-        out = xp.zeros(int(n))
+        out = np.zeros(int(n))
     multiplicity = int(multiplicity)
     greville_p(knots, degree, periodic, out, multiplicity)
     return out
@@ -568,7 +574,7 @@ def elements_spans(knots, degree, out=None):
     """
     knots = xp.ascontiguousarray(knots, dtype=float)
     if out is None:
-        out = xp.zeros(len(knots), dtype=xp.int64)
+        out = np.zeros(len(knots), dtype=xp.int64)
     else:
         assert out.shape == knots.shape and out.dtype == xp.dtype('int64')
     i_final = elements_spans_p(knots, degree, out)
@@ -617,7 +623,11 @@ def make_knots(breaks, degree, periodic, multiplicity=1, out=None):
 
     # Consistency checks
     assert len(breaks) > 1
-    assert all( xp.diff(breaks) > 0 )
+    # Convert to numpy for comparison since assertion needs Python bool
+    breaks_np = breaks.get() if hasattr(breaks, 'get') else breaks
+    if isinstance(breaks_np, (list, tuple)):
+        breaks_np = np.asarray(breaks_np)
+    assert all( np.diff(breaks_np) > 0 )
     assert degree >= 0
     assert 1 <= multiplicity and multiplicity <= degree + 1
     # Cast potential numpy.int64 into python native int
@@ -626,9 +636,14 @@ def make_knots(breaks, degree, periodic, multiplicity=1, out=None):
     if periodic:
         assert len(breaks) > degree
 
-    breaks = xp.ascontiguousarray(breaks, dtype=float)
+    # Keep breaks on NumPy for initialization - knots are index arrays needed for CPU operations
+    breaks = np.asarray(breaks, dtype=float) if isinstance(breaks, (list, tuple)) else breaks
+    if hasattr(breaks, 'get'):
+        breaks = breaks.get()  # Convert CuPy to NumPy
+    breaks = np.ascontiguousarray(breaks, dtype=float)
     if out is None:
-        out = xp.zeros(multiplicity * len(breaks[1:-1]) + 2 + 2 * degree)
+        # Knots are index arrays, keep them on NumPy
+        out = np.zeros(multiplicity * len(breaks[1:-1]) + 2 + 2 * degree)
     else:
         assert out.shape == (multiplicity * len(breaks[1:-1]) + 2 + 2 * degree,) \
             and out.dtype == xp.dtype('float')
@@ -676,25 +691,29 @@ def elevate_knots(knots, degree, periodic, multiplicity=1, tol=1e-15, out=None):
         Knots sequence of spline space of degree p+1.
     """
     multiplicity = int(multiplicity)
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    if isinstance(knots, (list, tuple)):
+        knots = np.asarray(knots, dtype=float)
+    if hasattr(knots, 'get'):
+        knots = knots.get()  # Convert CuPy to NumPy
+    knots = np.ascontiguousarray(knots, dtype=float)
     if out is None:
         if periodic:
-            out = xp.zeros(knots.shape[0] + 2, dtype=float)
+            out = np.zeros(knots.shape[0] + 2, dtype=float)
         else:
             shape = 2*(degree + 2)
             if len(knots) - 2 * (degree + 1) > 0:
-                uniques = (xp.diff(knots[degree + 1:-degree - 1]) > tol).nonzero()
+                uniques = (np.diff(knots[degree + 1:-degree - 1]) > tol).nonzero()
                 shape += multiplicity * (1 + uniques[0].shape[0])
-            out = xp.zeros(shape, dtype=float)
+            out = np.zeros(shape, dtype=float)
     else:
         if periodic:
-            assert out.shape == (knots.shape[0] + 2,) and out.dtype == xp.dtype('float')
+            assert out.shape == (knots.shape[0] + 2,) and out.dtype == np.dtype('float')
         else:
             shape = 2*(degree + 2)
             if len(knots) - 2 * (degree + 1) > 0:
-                uniques = (xp.diff(knots[degree + 1:-degree - 1]) > tol).nonzero()
+                uniques = (np.diff(knots[degree + 1:-degree - 1]) > tol).nonzero()
                 shape += multiplicity * (1 + uniques[0].shape[0])
-            assert out.shape == shape and out.dtype == xp.dtype('float')
+            assert out.shape == shape and out.dtype == np.dtype('float')
 
     elevate_knots_p(knots, degree, periodic, out, multiplicity, tol)
     return out
@@ -751,14 +770,18 @@ def quadrature_grid(breaks, quad_rule_x, quad_rule_w):
     assert min(quad_rule_x) >= -1
     assert max(quad_rule_x) <= +1
 
-    breaks = xp.ascontiguousarray(breaks, dtype=float)
+    # Convert breaks to numpy if CuPy (breaks/grids should stay on CPU)
+    if hasattr(breaks, 'get'):
+        breaks = breaks.get()
+    breaks = np.ascontiguousarray(breaks, dtype=float)
 
     if array_backend.backend == "cupy":
-        quad_rule_x = xp.ascontiguousarray(xp.array(quad_rule_x), dtype=float)
-        quad_rule_w = xp.ascontiguousarray( xp.array(quad_rule_w), dtype=float )
-    else:
-        quad_rule_x = xp.ascontiguousarray(quad_rule_x, dtype=float)
-        quad_rule_w = xp.ascontiguousarray( quad_rule_w, dtype=float )
+        # Convert CuPy arrays to NumPy
+        quad_rule_x = quad_rule_x.get() if hasattr(quad_rule_x, 'get') else quad_rule_x
+        quad_rule_w = quad_rule_w.get() if hasattr(quad_rule_w, 'get') else quad_rule_w
+
+    quad_rule_x = np.ascontiguousarray(quad_rule_x, dtype=float)
+    quad_rule_w = np.ascontiguousarray(quad_rule_w, dtype=float)
 
 
     out1 = xp.zeros((len(breaks) - 1, len(quad_rule_x)))
@@ -914,7 +937,7 @@ def cell_index(breaks, i_grid, tol=1e-15, out=None):
     breaks = xp.ascontiguousarray(breaks, dtype=float)
     i_grid = xp.ascontiguousarray(i_grid, dtype=float)
     if out is None:
-        out = xp.zeros_like(i_grid, dtype=xp.int64)
+        out = np.zeros_like(i_grid, dtype=xp.int64)
     else:
         assert out.shape == i_grid.shape and out.dtype == xp.dtype('int64')
     status = cell_index_p(breaks, i_grid, tol, out)

feectools/core/bsplines_kernels.py

@@ -7,6 +7,8 @@
 # This file holds the pyccelisable versions of the functions in bsplines.py
 # This will be changed once pyccel can return arrays and can get out=None arguments
 # like Numpy functions.
+# NOTE: This file must use ONLY numpy for pyccel compilation compatibility.
+# Backend conversion (NumPy/CuPy) happens at the Python wrapper level.
 
 from pyccel.decorators import pure
 from numpy import shape, abs
@@ -401,6 +403,9 @@ def basis_funs_all_ders_p(knots: 'float[:]', degree: int, x: float, span: int, n
     .. [1] L. Piegl and W. Tiller. The NURBS Book, 2nd ed.,
         Springer-Verlag Berlin Heidelberg GmbH, 1997.
     """
+    # Detect backend from output array
+    # Backend array operations removed - always use numpy
+
     sh_a  = np.empty(2)
     sh_b  = np.empty(2)
     left  = np.empty(degree)
@@ -556,6 +561,9 @@ def collocation_matrix_p(knots: 'float[:]', degree: int, periodic: bool, normali
         multiplicity applies to each interior knot.
 
     """
+    # Detect backend from output array
+    # Backend array operations removed - always use numpy
+
     # Number of basis functions (in periodic case remove degree repeated elements)
     nb = len(knots)-degree-1
     if periodic:
@@ -565,7 +573,7 @@ def collocation_matrix_p(knots: 'float[:]', degree: int, periodic: bool, normali
     nx = len(xgrid)
 
     basis = np.zeros((nx, degree + 1))
-    spans = np.zeros(nx, dtype=int)
+    spans = np.zeros(nx, dtype=int)  # Keep indices on CPU
     find_spans_p(knots, degree, xgrid, spans)
     basis_funs_array_p(knots, degree, xgrid, spans, basis)
 
@@ -644,6 +652,9 @@ def histopolation_matrix_p(knots: 'float[:]', degree: int, periodic: bool, norma
     contains the integrals of each B-spline basis function :math:`B_j` between
     two successive grid points.
     """
+    # Detect backend from output array
+    # Backend array operations removed - always use numpy
+
     nb = len(knots) - degree - 1
     if periodic:
         nb -= degree + 1 - multiplicity
@@ -700,7 +711,7 @@ def histopolation_matrix_p(knots: 'float[:]', degree: int, periodic: bool, norma
     m = colloc.shape[0] - 1
     n = colloc.shape[1] - 1
 
-    spans = np.zeros(colloc.shape[0], dtype=int)
+    spans = np.zeros(colloc.shape[0], dtype=int)  # Keep indices on CPU
     for i in range(colloc.shape[0]):
         local_span = 0
         for j in range(colloc.shape[1]):
@@ -767,7 +778,10 @@ def merge_sort(a: 'float[:]'):
     """
     if len(a) != 1 and len(a) != 0:
         n = len(a)
-
+
+        # Detect backend and use the appropriate array module
+        # Backend array operations removed - always use numpy
+
         a1 = np.zeros(n // 2)
         a1[:] = a[:n // 2]
         a2 = np.zeros(n - n // 2)
@@ -1167,12 +1181,16 @@ def basis_ders_on_quad_grid_p(knots: 'float[:]', degree: int, quad_grid: 'float[
     """
     ne = quad_grid.shape[0]
     nq = quad_grid.shape[1]
+
+    # Detect backend from output array
+    # Backend array operations removed - always use numpy
+
     if normalization:
         integrals = np.zeros(knots.shape[0] - degree - 1)
         basis_integrals_p(knots, degree, integrals)
         scaling = 1.0 /integrals
 
-    temp_spans = np.zeros(len(knots), dtype=int)
+    temp_spans = np.zeros(len(knots), dtype=int)  # Keep indices on CPU
     actual_index = elements_spans_p(knots, degree, temp_spans)
     spans = temp_spans[:actual_index]
 

feectools/ddm/blocking_data_exchanger.py

@@ -1,8 +1,5 @@
-#---------------------------------------------------------------------------#
-# This file is part of PSYDAC which is released under MIT License. See the  #
-# LICENSE file or go to https://github.com/pyccel/psydac/blob/devel/LICENSE #
-# for full license details.                                                 #
-#---------------------------------------------------------------------------#
+# coding: utf-8
+
 import cunumpy as xp
 import numpy as np
 from feectools.ddm.mpi import mpi as MPI

feectools/ddm/cart.py

@@ -1,9 +1,20 @@
 # coding: utf-8
 
 import os
+import numpy as np
 import cunumpy as xp
+from cunumpy.xp import array_backend
 from itertools import product
 
+# Initialize CUDA context before MPI if using CuPy backend
+if array_backend.backend == "cupy":
+    try:
+        import cupy as cp
+        cp.cuda.Device(0).use()
+        cp.cuda.Stream.null.synchronize()
+    except Exception:
+        pass
+
 from feectools.ddm.mpi import mpi as MPI
 from feectools.ddm.mpi import MockMPI
 from feectools.ddm.partition import compute_dims, partition_procs_per_patch
@@ -43,6 +54,17 @@ def find_mpi_type( dtype ):
 
     return mpi_type
 
+#===============================================================================
+def _cuda_sync_before_mpi():
+    """Synchronize CUDA before MPI operations to avoid conflicts."""
+    if array_backend.backend == "cupy":
+        try:
+            import cupy as cp
+            cp.cuda.Stream.null.synchronize()
+        except Exception:
+            pass
+
+#===============================================================================
 class MultiPatchDomainDecomposition:
     """
     Cartesian decomposition of multiple N-Cube grids.
@@ -260,11 +282,14 @@ def __init__(self, ncells, periods, comm=None, global_comm=None, num_threads=Non
 
         if comm is None:
             # compute the coords for all processes
-            self._global_coords = xp.array([xp.unravel_index(xp.int64(rank), nprocs) for rank in range(self._size)])
+            self._global_coords = np.array([np.unravel_index(int(rank), nprocs) for rank in range(self._size)])
             self._coords        = self._global_coords[self._rank]
             self._rank_in_topo  = 0
             self._ranks_in_topo = xp.array([0])
         else:
+            # Synchronize CUDA before MPI operations
+            _cuda_sync_before_mpi()
+
             # Create a MPI cart
             self._comm_cart = comm.Create_cart(
                 dims    = self._nprocs,
@@ -288,6 +313,7 @@ def __init__(self, ncells, periods, comm=None, global_comm=None, num_threads=Non
         # Create (N-1)-dimensional communicators within the Cartesian topology
         self._subcomm = [None]*self._ndims
         for i in range(self._ndims):
+            _cuda_sync_before_mpi()  # Synchronize before each Sub() call
             remain_dims     = [i==j for j in range( self._ndims )]
             self._subcomm[i] = self._comm_cart.Sub( remain_dims )
 
@@ -467,8 +493,9 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads
         # Store input arguments
         self._domain_decomposition = domain_decomposition
         self._npts          = tuple( npts    )
-        self._global_starts = tuple( [ xp.asarray(gs) for gs in global_starts]  )
-        self._global_ends   = tuple( [ xp.asarray(ge) for ge in global_ends]    )
+        # Convert to NumPy arrays for MPI compatibility (MPI can't handle CuPy arrays)
+        self._global_starts = tuple( [ np.asarray(gs.get() if hasattr(gs, 'get') else gs) for gs in global_starts]  )
+        self._global_ends   = tuple( [ np.asarray(ge.get() if hasattr(ge, 'get') else ge) for ge in global_ends]    )
         self._pads          = tuple( pads    )
         self._shifts        = tuple( shifts  )
         self._periods       = domain_decomposition.periods
@@ -494,10 +521,12 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads
 
         # Know my coordinates in the topology
         self._coords = domain_decomposition.coords
+        # Convert coords to NumPy for indexing (MPI coords should be on CPU)
+        coords_np = [c.get() if hasattr(c, 'get') else c for c in self._coords]
 
         # Start/end values of global indices (without ghost regions)
-        self._starts = tuple( self._global_starts[axis][c] for axis,c in zip(range(self._ndims), self._coords) )
-        self._ends   = tuple( self._global_ends  [axis][c] for axis,c in zip(range(self._ndims), self._coords) )
+        self._starts = tuple( self._global_starts[axis][c] for axis,c in zip(range(self._ndims), coords_np) )
+        self._ends   = tuple( self._global_ends  [axis][c] for axis,c in zip(range(self._ndims), coords_np) )
 
         # List of 1D global indices (without ghost regions)
         # self._grids = tuple( range(s,e+1) for s,e in zip( self._starts, self._ends ) )
@@ -906,7 +935,7 @@ def _compute_shift_info_non_blocking( self, shift ):
         if len([i for i in shift if i==0]) == 2 and rank_dest != MPI.PROC_NULL:
             direction = [i for i,s in enumerate(shift) if s != 0][0]
             comm = self._subcomm[direction]
-            # local_dest_rank = self._comm_cart.group.Translate_ranks(xp.array([rank_dest]), comm.group)[0]
+            # local_dest_rank = self._comm_cart.group.Translate_ranks(np.array([rank_dest]), comm.group)[0]
             local_dest_rank = self._comm_cart.group.Translate_ranks([int(rank_dest)], comm.group)[0]
 
         else:
@@ -921,7 +950,7 @@ def _compute_shift_info_non_blocking( self, shift ):
         if len([i for i in shift if i==0]) == 2 and rank_source != MPI.PROC_NULL:
             direction = [i for i,s in enumerate(shift) if s != 0][0]
             comm = self._subcomm[direction]
-            # local_source_rank = self._comm_cart.group.Translate_ranks(xp.array([rank_source]), comm.group)[0]
+            # local_source_rank = self._comm_cart.group.Translate_ranks(np.array([rank_source]), comm.group)[0]
             local_source_rank = self._comm_cart.group.Translate_ranks([int(rank_source)], comm.group)[0]
         else:
             local_source_rank = rank_source

feectools/ddm/mpi.py

@@ -81,6 +81,11 @@ def COMM_WORLD(self):
 
 
 try:
+    # Disable MPI when using CuPy due to known segfault issues with OpenMPI + CUDA
+    import os
+    if os.environ.get('ARRAY_BACKEND') == 'cupy':
+        raise ImportError("MPI disabled when using CuPy backend")
+
     from mpi4py import MPI
 
     _comm = MPI.COMM_WORLD

feectools/ddm/partition.py

@@ -1,9 +1,5 @@
-#---------------------------------------------------------------------------#
-# This file is part of PSYDAC which is released under MIT License. See the  #
-# LICENSE file or go to https://github.com/pyccel/psydac/blob/devel/LICENSE #
-# for full license details.                                                 #
-#---------------------------------------------------------------------------#
 import cunumpy as xp
+import numpy as np
 import numpy.ma as ma
 
 from sympy.ntheory import factorint