Another cupy branch

feectools/core/bsplines.py

@@ -16,6 +16,7 @@
 
 """
 import cunumpy as xp
+from cunumpy import PyccelKernel
 from cunumpy.xp import array_backend
 import numpy as np
 
@@ -38,6 +39,27 @@
                                           cell_index_p,
                                           basis_ders_on_irregular_grid_p)
 
+# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they
+# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel).
+find_span_p = PyccelKernel(find_span_p)
+find_spans_p = PyccelKernel(find_spans_p)
+basis_funs_p = PyccelKernel(basis_funs_p)
+basis_funs_array_p = PyccelKernel(basis_funs_array_p)
+basis_funs_1st_der_p = PyccelKernel(basis_funs_1st_der_p)
+basis_funs_all_ders_p = PyccelKernel(basis_funs_all_ders_p)
+collocation_matrix_p = PyccelKernel(collocation_matrix_p)
+histopolation_matrix_p = PyccelKernel(histopolation_matrix_p)
+greville_p = PyccelKernel(greville_p)
+breakpoints_p = PyccelKernel(breakpoints_p)
+elements_spans_p = PyccelKernel(elements_spans_p)
+make_knots_p = PyccelKernel(make_knots_p)
+elevate_knots_p = PyccelKernel(elevate_knots_p)
+quadrature_grid_p = PyccelKernel(quadrature_grid_p)
+basis_ders_on_quad_grid_p = PyccelKernel(basis_ders_on_quad_grid_p)
+basis_integrals_p = PyccelKernel(basis_integrals_p)
+cell_index_p = PyccelKernel(cell_index_p)
+basis_ders_on_irregular_grid_p = PyccelKernel(basis_ders_on_irregular_grid_p)
+
 __all__ = ('find_span',
            'find_spans',
            'basis_funs',
@@ -84,7 +106,7 @@ def find_span(knots, degree, x):
         Knot span index.
     """
     x = float(x)
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     return find_span_p(knots, degree, x)
 
 #==============================================================================
@@ -116,8 +138,8 @@ def find_spans(knots, degree, x, out=None):
     spans : array of ints
         Knots span indexes.
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
-    x = xp.ascontiguousarray(x, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
+    x = xp.ascontiguousarray(xp.asarray(x), dtype=float)
     if out is None:
         out = xp.zeros_like(x, dtype=int)
     else:
@@ -155,7 +177,7 @@ def basis_funs(knots, degree, x, span, out=None):
         1D array containing the values of ``degree + 1`` non-zero
         Bsplines at location ``x``.
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     # Get native float
     x = float(x)
     if out is None:
@@ -193,8 +215,8 @@ def basis_funs_array(knots, degree, span, x, out=None):
         2D array of shape ``(len(x), degree + 1)`` containing the values of ``degree + 1`` non-zero
         Bsplines at each location in ``x``.
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
-    x = xp.ascontiguousarray(x, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
+    x = xp.ascontiguousarray(xp.asarray(x), dtype=float)
     if out is None:
         out = xp.zeros(x.shape + (degree + 1,), dtype=float)
     else:
@@ -240,7 +262,7 @@ def basis_funs_1st_der(knots, degree, x, span, out=None):
     ----------
     .. [2] SELALIB, Semi-Lagrangian Library. http://selalib.gforge.inria.fr
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     # Get native float to work on windows
     x = float(x)
     if out is None:
@@ -291,7 +313,7 @@ def basis_funs_all_ders(knots, degree, x, span, n, normalization='B', out=None):
         ders[i,j] = (d/dx)^i B_k(x) with k=(span-degree+j),
         for 0 <= i <= n and 0 <= j <= degree+1.
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     # Get native float to work on windows
     x = float(x)
     if out is None:
@@ -346,8 +368,8 @@ def collocation_matrix(knots, degree, periodic, normalization, xgrid, out=None,
     if xgrid.size == 1:
         return xp.ones((1, 1), dtype=float)
 
-    knots = xp.ascontiguousarray(knots, dtype=float)
-    xgrid = xp.ascontiguousarray(xgrid, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
+    xgrid = xp.ascontiguousarray(xp.asarray(xgrid), dtype=float)
     if out is None:
         nb = len(knots) - degree - 1
         if periodic:
@@ -430,8 +452,8 @@ def histopolation_matrix(knots, degree, periodic, normalization, xgrid, multipli
     if not xp.all(xp.diff(xgrid) > 0):
         raise ValueError("Grid points must be ordered, with no repetitions: {}".format(xgrid))
 
-    knots = xp.ascontiguousarray(knots, dtype=float)
-    xgrid = xp.ascontiguousarray(xgrid, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
+    xgrid = xp.ascontiguousarray(xp.asarray(xgrid), dtype=float)
     elevated_knots = elevate_knots(knots, degree, periodic, multiplicity=multiplicity)
 
     normalization = normalization == "M"
@@ -477,7 +499,7 @@ def breakpoints(knots, degree, tol=1e-15, out=None):
     breaks : numpy.ndarray (1D)
         Abscissas of all breakpoints.
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     if out is None:
         out = xp.zeros(len(knots), dtype=float)
     else:
@@ -518,8 +540,7 @@ def greville(knots, degree, periodic, out=None, multiplicity=1):
     # 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 = xp.to_numpy(knots)
     knots = np.ascontiguousarray(knots, dtype=float)
     if out is None:
         n = len(knots) - 2 * degree - 2 + multiplicity if periodic else len(knots) - degree - 1
@@ -572,7 +593,7 @@ def elements_spans(knots, degree, out=None):
        spans  = xp.searchsorted( knots, breaks[:-1], side='right' ) - 1
 
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     if out is None:
         out = np.zeros(len(knots), dtype=xp.int64)
     else:
@@ -624,7 +645,7 @@ def make_knots(breaks, degree, periodic, multiplicity=1, out=None):
     # Consistency checks
     assert len(breaks) > 1
     # Convert to numpy for comparison since assertion needs Python bool
-    breaks_np = breaks.get() if hasattr(breaks, 'get') else breaks
+    breaks_np = xp.to_numpy(breaks)
     if isinstance(breaks_np, (list, tuple)):
         breaks_np = np.asarray(breaks_np)
     assert all( np.diff(breaks_np) > 0 )
@@ -638,8 +659,7 @@ def make_knots(breaks, degree, periodic, multiplicity=1, out=None):
 
     # 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 = xp.to_numpy(breaks)
     breaks = np.ascontiguousarray(breaks, dtype=float)
     if out is None:
         # Knots are index arrays, keep them on NumPy
@@ -693,8 +713,7 @@ def elevate_knots(knots, degree, periodic, multiplicity=1, tol=1e-15, out=None):
     multiplicity = int(multiplicity)
     if isinstance(knots, (list, tuple)):
         knots = np.asarray(knots, dtype=float)
-    if hasattr(knots, 'get'):
-        knots = knots.get()  # Convert CuPy to NumPy
+    knots = xp.to_numpy(knots)
     knots = np.ascontiguousarray(knots, dtype=float)
     if out is None:
         if periodic:
@@ -771,14 +790,13 @@ def quadrature_grid(breaks, quad_rule_x, quad_rule_w):
     assert max(quad_rule_x) <= +1
 
     # Convert breaks to numpy if CuPy (breaks/grids should stay on CPU)
-    if hasattr(breaks, 'get'):
-        breaks = breaks.get()
+    breaks = xp.to_numpy(breaks)
     breaks = np.ascontiguousarray(breaks, dtype=float)
 
     if array_backend.backend == "cupy":
         # 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 = xp.to_numpy(quad_rule_x)
+        quad_rule_w = xp.to_numpy(quad_rule_w)
 
     quad_rule_x = np.ascontiguousarray(quad_rule_x, dtype=float)
     quad_rule_w = np.ascontiguousarray(quad_rule_w, dtype=float)
@@ -848,8 +866,8 @@ def basis_ders_on_quad_grid(knots, degree, quad_grid, nders, normalization, offs
     """
     offset = int(offset)
     ne, nq = quad_grid.shape
-    knots = xp.ascontiguousarray(knots, dtype=float)
-    quad_grid = xp.ascontiguousarray(quad_grid, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
+    quad_grid = xp.ascontiguousarray(xp.asarray(quad_grid), dtype=float)
     if out is None:
         out = xp.zeros((ne, degree + 1, nders + 1, nq), dtype=float)
     else:
@@ -892,7 +910,7 @@ def basis_integrals(knots, degree, out=None):
     to (len(knots)-degree-1). In the periodic case the last (degree) values in
     the array are redundant, as they are a copy of the first (degree) values.
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
     if out is None:
         out = xp.zeros(len(knots) - degree - 1, dtype=float)
     else:
@@ -934,8 +952,8 @@ def cell_index(breaks, i_grid, tol=1e-15, out=None):
         ``cell_index[i]`` is the index of the cell in which
         ``i_grid[i]`` belong.
     """
-    breaks = xp.ascontiguousarray(breaks, dtype=float)
-    i_grid = xp.ascontiguousarray(i_grid, dtype=float)
+    breaks = xp.ascontiguousarray(xp.asarray(breaks), dtype=float)
+    i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float)
     if out is None:
         out = np.zeros_like(i_grid, dtype=xp.int64)
     else:
@@ -990,8 +1008,8 @@ def basis_ders_on_irregular_grid(knots, degree, i_grid, cell_index, nders, norma
         . il: local basis function   (0 <= il <= degree)
         . id: derivative             (0 <= id <= nders )
     """
-    knots = xp.ascontiguousarray(knots, dtype=float)
-    i_grid = xp.ascontiguousarray(i_grid, dtype=float)
+    knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float)
+    i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float)
     if out is None:
         nx = i_grid.shape[0]
         out = xp.zeros((nx, degree + 1, nders + 1), dtype=float)

feectools/ddm/blocking_data_exchanger.py

@@ -5,6 +5,7 @@
 from feectools.ddm.mpi import mpi as MPI
 
 from .cart import CartDecomposition, find_mpi_type
+from .device import synchronize_for_mpi
 from .basic import CartDataExchanger
 
 
@@ -82,6 +83,10 @@ def start_update_ghost_regions( self, array, requests ):
 
         assert isinstance( array, xp.ndarray )
 
+        # MPI reads/writes `array` directly; on a device backend the
+        # kernels that produced it must have finished first.
+        synchronize_for_mpi( array )
+
         # Shortcuts
         cart = self._cart
         comm = self._comm
@@ -123,6 +128,8 @@ def start_exchange_assembly_data( self, array ):
 
         assert isinstance( array, xp.ndarray )
 
+        synchronize_for_mpi( array )
+
         # Shortcuts
         cart  = self._cart
         comm  = self._comm

feectools/ddm/cart.py

@@ -1,19 +1,20 @@
 # coding: utf-8
 
 import os
+import cunumpy  # only for its backend-agnostic to_numpy(), see below -- not aliased to
+                 # xp here, since that alias is reserved for plain NumPy in this module.
 import numpy as np
-import cunumpy as xp
-from cunumpy.xp import array_backend
+import numpy as xp  # this module is host-only MPI/index bookkeeping, never device data
 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 cunumpy.xp import array_backend, to_numpy
+
+# Initialize the CUDA context before MPI if using CuPy backend, binding this
+# rank to its own GPU. Must stay above the feectools.ddm.mpi import, which
+# initialises MPI as a side effect.
+from feectools.ddm.device import bind_local_device
+
+bind_local_device()
 
 from feectools.ddm.mpi import mpi as MPI
 from feectools.ddm.mpi import MockMPI
@@ -482,6 +483,12 @@ class CartDecomposition():
     """
     def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads, shifts ):
 
+        # global_starts/global_ends are host-side decomposition metadata; callers
+        # may hand them in as CuPy arrays (e.g. built with cunumpy under the CuPy
+        # backend), so coerce them to NumPy up front.
+        global_starts = [ to_numpy(gs) for gs in global_starts ]
+        global_ends   = [ to_numpy(ge) for ge in global_ends   ]
+
         # Check input arguments
         # TODO: check that arguments are identical across all processes
         assert len( npts ) == len( global_starts ) == len( global_ends ) == len( pads ) == len(shifts)
@@ -494,8 +501,8 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads
         self._domain_decomposition = domain_decomposition
         self._npts          = tuple( npts    )
         # 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._global_starts = tuple( [ cunumpy.to_numpy(gs) for gs in global_starts]  )
+        self._global_ends   = tuple( [ cunumpy.to_numpy(ge) for ge in global_ends]    )
         self._pads          = tuple( pads    )
         self._shifts        = tuple( shifts  )
         self._periods       = domain_decomposition.periods
@@ -522,7 +529,11 @@ 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]
+        # cunumpy.to_numpy, not used here: self._coords may hold plain Python ints
+        # (mpi4py's Get_coords returns a plain list), and to_numpy would wrap those
+        # into 0-d NumPy arrays via np.asarray -- the wrong type to index a tuple of
+        # global_starts/ends with below. is_gpu leaves non-CuPy values untouched.
+        coords_np = [c.get() if cunumpy.is_gpu(c) 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), coords_np) )