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:
@@ -572,7 +594,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:
@@ -848,8 +870,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 +914,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 +956,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 +1012,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

@@ -2,18 +2,17 @@
 
 import os
 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 +481,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)

feectools/ddm/device.py

@@ -0,0 +1,105 @@
+"""
+Binding of MPI processes to GPUs.
+
+Kept free of any MPI import on purpose: the CUDA context should exist before
+``MPI_Init`` runs, and importing :mod:`feectools.ddm.mpi` initialises MPI as a
+side effect. The rank of the process within its node is therefore taken from
+the environment variables the launcher sets, which are available before
+``MPI_Init``, rather than from a communicator.
+"""
+
+import os
+
+from cunumpy.xp import array_backend
+
+__all__ = ('local_rank', 'bind_local_device', 'synchronize_for_mpi')
+
+# Node-local rank, as exported by the common launchers.
+_LOCAL_RANK_VARS = (
+    'OMPI_COMM_WORLD_LOCAL_RANK',   # Open MPI
+    'MV2_COMM_WORLD_LOCAL_RANK',    # MVAPICH2
+    'MPI_LOCALRANKID',              # Intel MPI
+    'PMI_LOCAL_RANK',               # MPICH / PMI
+    'SLURM_LOCALID',                # Slurm
+)
+
+
+def synchronize_for_mpi(*arrays):
+    """
+    Wait for pending device work before MPI touches `arrays`.
+
+    CuPy launches kernels asynchronously on the current stream; MPI knows
+    nothing about that stream. Handing it a device buffer that a kernel is
+    still writing lets it send whatever happens to be in memory at that
+    moment, which shows up as silently wrong ghost regions rather than as an
+    error. Every MPI call that reads or writes device memory must therefore be
+    preceded by this.
+
+    The reverse direction needs no barrier: MPI completes its own transfers
+    before the corresponding wait returns, so kernels launched afterwards see
+    the received data.
+
+    Parameters
+    ----------
+    *arrays : array | None
+        The buffers about to be given to MPI. Synchronization happens only if
+        at least one of them lives on a device, so host-only exchanges (and the
+        whole NumPy backend) pay nothing.
+    """
+    if not any(hasattr(a, 'get') for a in arrays if a is not None):
+        return
+
+    import cupy as cp
+
+    cp.cuda.get_current_stream().synchronize()
+
+
+def local_rank():
+    """
+    The rank of this process within its node, or 0 if no launcher told us
+    (which is the right answer for a serial run).
+    """
+    for var in _LOCAL_RANK_VARS:
+        value = os.environ.get(var)
+        if value is None:
+            continue
+        try:
+            return int(value)
+        except ValueError:
+            continue
+    return 0
+
+
+def bind_local_device():
+    """
+    Bind this process to one GPU, chosen round-robin by its node-local rank, and
+    create the CUDA context.
+
+    Without this every rank on a node would share GPU 0: they would contend for
+    one device while the others idled, and one device's memory would have to
+    hold every rank's data. `CUDA_VISIBLE_DEVICES` still applies first, so a
+    launcher that already hands each rank its own device keeps working (each
+    process then sees a single device and picks index 0).
+
+    Returns
+    -------
+    int | None
+        The index of the device that was selected, or None if the CuPy backend
+        is not active or no device is available.
+    """
+    if array_backend.backend != 'cupy':
+        return None
+
+    try:
+        import cupy as cp
+
+        count = cp.cuda.runtime.getDeviceCount()
+        if count == 0:
+            return None
+
+        device = local_rank() % count
+        cp.cuda.Device(device).use()
+        cp.cuda.Stream.null.synchronize()
+        return device
+    except Exception:  # noqa: BLE001 - a driver/runtime failure must not be fatal
+        return None

feectools/ddm/interface_data_exchanger.py

@@ -3,6 +3,7 @@
 from feectools.ddm.mpi import mpi as MPI
 
 from .cart import InterfaceCartDecomposition, find_mpi_type
+from .device import synchronize_for_mpi
 
 __all__ = ('InterfaceCartDataExchanger',)
 
@@ -48,6 +49,10 @@ def update_ghost_regions( self, array_minus=None, array_plus=None ):
 
     # ...
     def start_update_ghost_regions( self, array_minus=None, array_plus=None ):
+        # MPI reads/writes these buffers directly; on a device backend the
+        # kernels that produced them must have finished first.
+        synchronize_for_mpi( array_minus, array_plus )
+
         send_req = []
         recv_req = []
         cart      = self._cart