feectools 0.3.0__tar.gz → 0.3.1__tar.gz

{feectools-0.3.0/feectools.egg-info → feectools-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.3.0
+Version: 0.3.1
 Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
 Author-email: Psydac development team <psydac@googlegroups.com>
 Maintainer-email: Stefan Possanner <stefan.possanner@ipp.mpg.de>, Max Lindqvist <max.lindqvist@ipp.mpg.de>, Yaman Güçlü <yaman.guclu@gmail.com>, Martin Campos Pinto <martin.campos-pinto@ipp.mpg.de>, Ahmed Ratnani <ratnaniahmed@gmail.com>
@@ -42,7 +42,7 @@ Requires-Dist: matplotlib
 Requires-Dist: pyyaml>=5.1
 Requires-Dist: packaging
 Requires-Dist: pyevtk
-Requires-Dist: cunumpy
+Requires-Dist: cunumpy<0.6,>=0.5.0
 Requires-Dist: pyccel>=2.1.0
 Requires-Dist: h5py
 Requires-Dist: tblib

{feectools-0.3.0 → feectools-0.3.1}/feectools/core/bsplines.py

@@ -16,6 +16,7 @@ References:
 
 """
 import cunumpy as xp
+from cunumpy.kernels import PyccelKernel
 from cunumpy.xp import array_backend
 import numpy as np
 
@@ -38,6 +39,27 @@ from feectools.core.bsplines_kernels import (find_span_p,
                                           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.kernels.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-0.3.0 → feectools-0.3.1}/feectools/core/tests/test_bsplines.py

@@ -159,7 +159,7 @@ def test_histopolation_matrix(lims, nc, p, periodic, tol=1e-13):
 def test_cell_index(i_grid, expected):
     breaks = xp.array([0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.])
     out = cell_index(breaks, xp.asarray(i_grid))
-    assert xp.array_equal(expected, out)
+    assert xp.array_equal(xp.asarray(expected), out)
 
 #==============================================================================
 # SCRIPT FUNCTIONALITY: PLOT BASIS FUNCTIONS

{feectools-0.3.0 → feectools-0.3.1}/feectools/core/tests/test_bsplines_kernel.py

@@ -2,14 +2,16 @@
 
 import pytest
 import cunumpy as xp
+import numpy as np
 
 
 from feectools.core.bsplines_kernels import cell_index_p
 
 def test_cell_index_p():
-    breaks = xp.array([0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.])
-    breaks = xp.ascontiguousarray(breaks, dtype=float)
-    out = xp.zeros_like(breaks, dtype=xp.int64)
+    # This directly tests the raw Pyccel kernel, which intentionally accepts
+    # NumPy host arrays only; CuPy coverage belongs to the public wrapper.
+    breaks = np.ascontiguousarray([0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.], dtype=float)
+    out = np.zeros_like(breaks, dtype=np.int64)
     tol = 1e-15
 
     # limit case: code should decide wether point is in or out, not fall in infinite loop
@@ -26,15 +28,14 @@ def test_cell_index_p():
         assert status == expected_status
 
     # checking that the values match those of searchsorted (-1) for arbitrary grid points
-    i_grid = xp.array([0.14320482, 0.86569833, 0.77775327, 0.00895956, 0.074629  ,
+    i_grid = np.array([0.14320482, 0.86569833, 0.77775327, 0.00895956, 0.074629  ,
        0.45682646, 0.5384352 , 0.20915311, 0.73121977, 0.01057414,
        0.33756086, 0.17839759, 0.14023414, 0.09846206, 0.79970392,
        0.65330406, 0.82716552, 0.24185731, 0.24054685, 0.72466651,
        0.69125033, 0.3136558 , 0.64794089, 0.47975527, 0.99802844,
        0.64402598, 0.41263526, 0.28178414, 0.57274384, 0.73218562])
-    out = xp.zeros_like(i_grid, dtype=xp.int64)
+    out = np.zeros_like(i_grid, dtype=np.int64)
     status = cell_index_p(breaks, i_grid, tol, out)
     assert status == 0
-    nps = xp.searchsorted(breaks, i_grid)-1
-    assert xp.allclose(out, nps)
-
+    nps = np.searchsorted(breaks, i_grid)-1
+    assert np.allclose(out, nps)

{feectools-0.3.0 → feectools-0.3.1}/feectools/core/tests/test_bsplines_pyccel.py

@@ -187,8 +187,9 @@ def basis_funs_all_ders_true(knots, degree, x, span, n, normalization='B'):
 
     # Normalization to get M-Splines
     if normalization == 'M':
-        ders *= [(degree + 1) / (knots[i + degree + 1] - knots[i]) \
-                 for i in range(span - degree, span + 1)]
+        scaling = xp.asarray([(degree + 1) / (knots[i + degree + 1] - knots[i])
+                              for i in range(span - degree, span + 1)])
+        ders *= scaling
     return ders
 
 #==============================================================================
@@ -221,7 +222,15 @@ def collocation_matrix_true(knots, degree, periodic, normalization, xgrid):
     for i,x in enumerate( xgrid ):
         span  =  find_span_true( knots, degree, x )
         basis = basis_funs_true( knots, degree, x, span )
-        mat[i,js(span)] = normalize(basis, span)
+        values = normalize(basis, span)
+        if periodic:
+            # NumPy and CuPy differ for indexed assignment with repeated
+            # indices (which occurs when nb <= degree).  The production
+            # kernel assigns in loop order, so make the reference explicit.
+            for j, value in zip(js(span), values):
+                mat[i, j] = value
+        else:
+            mat[i, js(span)] = values
 
     # Mitigate round-off errors
     mat[abs(mat) < 1e-14] = 0.0
@@ -293,7 +302,7 @@ def histopolation_matrix_true(knots, degree, periodic, normalization, xgrid):
     # Compute span for each row (index of last non-zero basis function)
     # TODO: would be better to have this ready beforehand
     # TODO: use tolerance instead of comparing against zero
-    spans = [(row != 0).argmax() + (degree+1) for row in C]
+    spans = [int((row != 0).argmax()) + (degree+1) for row in C]
 
     # Compute histopolation matrix from collocation matrix of higher degree
     m = C.shape[0] - 1

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/blocking_data_exchanger.py

@@ -1,6 +1,7 @@
 # coding: utf-8
 
 import cunumpy as xp
+from cunumpy.mpi import synchronize_for_mpi
 import numpy as np
 from feectools.ddm.mpi import mpi as MPI
 
@@ -82,6 +83,10 @@ class BlockingCartDataExchanger(CartDataExchanger):
 
         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,10 @@ class BlockingCartDataExchanger(CartDataExchanger):
 
         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

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/cart.py

@@ -2,19 +2,21 @@
 
 import copy
 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
+
+# Bind this rank to its own GPU (by its rank within the node) and create the
+# CUDA context before MPI is initialized, as CUDA-aware MPI requires. Must stay
+# above the feectools.ddm.mpi import, which initializes MPI as a side effect.
+# A no-op on the NumPy backend.
+from cunumpy.cuda import bind_local_device
+
+bind_local_device()
 
 from feectools.ddm.mpi import mpi as MPI
 from feectools.ddm.mpi import MockMPI
@@ -541,6 +543,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)
@@ -553,8 +561,8 @@ class CartDecomposition():
         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( [ to_numpy(gs) for gs in global_starts]  )
+        self._global_ends   = tuple( [ to_numpy(ge) for ge in global_ends]    )
         self._pads          = tuple( pads    )
         self._shifts        = tuple( shifts  )
         self._periods       = domain_decomposition.periods
@@ -569,6 +577,11 @@ class CartDecomposition():
         self._shape         = (0,)*self._ndims
         self._parent_starts = (None,)*self._ndims
         self._parent_ends   = (None,)*self._ndims
+        # Serial decompositions have no neighbour exchanges, but exchange
+        # helpers still inspect these caches.  Define them before the early
+        # communicator exits so those helpers are backend-independent.
+        self._shift_info = {}
+        self._shift_info_non_blocking = {}
 
         if self._comm == MPI.COMM_NULL:
             return
@@ -581,7 +594,11 @@ class CartDecomposition():
         # 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) )
@@ -605,11 +622,6 @@ class CartDecomposition():
         # Create (N-1)-dimensional communicators within the Cartesian topology
         self._subcomm = domain_decomposition.subcomm
 
-        # dict to store information for communicating with neighbors
-        self._shift_info = {}
-
-#        # dict to store information for communicating with neighbors using non blocking communications
-        self._shift_info_non_blocking = {}
 
     #---------------------------------------------------------------------------
     # Global properties (same for each process)

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/interface_data_exchanger.py

@@ -1,5 +1,6 @@
 # coding: utf-8
 
+from cunumpy.mpi import synchronize_for_mpi
 from feectools.ddm.mpi import mpi as MPI
 
 from .cart import InterfaceCartDecomposition, find_mpi_type
@@ -48,6 +49,10 @@ class InterfaceCartDataExchanger:
 
     # ...
     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

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/mpi.py

@@ -187,11 +187,9 @@ def _mpi_disabled():
 
 if launched_under_mpi():
     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")
-
+        # MPI with the CuPy backend needs a CUDA-aware MPI library: device
+        # buffers are passed to MPI directly (after synchronize_for_mpi, see
+        # the data exchangers). A non-CUDA-aware MPI segfaults on them.
         if _mpi_disabled():
             raise ImportError("MPI disabled (feectools.use_mpi = False or FEECTOOLS_MPI=0)")
 

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/nonblocking_data_exchanger.py

@@ -1,6 +1,7 @@
 # coding: utf-8
 
 import cunumpy as xp
+from cunumpy.mpi import synchronize_for_mpi
 import numpy as np
 from itertools import product
 
@@ -98,6 +99,9 @@ class NonBlockingCartDataExchanger(CartDataExchanger):
         return tuple(requests)
 
     def start_update_ghost_regions(self, array, requests ):
+        # The persistent requests read/write `array` directly; on a device
+        # backend the kernels that produced it must have finished first.
+        synchronize_for_mpi( array )
         MPI.Prequest.Startall( requests )
 
     def end_update_ghost_regions(self, array, requests):
@@ -108,6 +112,10 @@ class NonBlockingCartDataExchanger(CartDataExchanger):
 
         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

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/partition.py

@@ -1,5 +1,4 @@
-import cunumpy as xp
-import numpy as np
+import numpy as xp  # this module is host-only MPI/index bookkeeping, never device data
 import numpy.ma as ma
 
 __all__ = ('compute_dims', 'partition_procs_per_patch')

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/petsc.py

@@ -1,6 +1,6 @@
 # coding: utf-8
 
-import cunumpy as xp
+import numpy as xp  # this module is host-only MPI/index bookkeeping, never device data
 from itertools import product
 
 import cunumpy as xp

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/tests/test_cart_1d.py

@@ -89,7 +89,7 @@ def run_cart_1d( data_exchanger_type, verbose=False ):
     #---------------------------------------------------------------------------
 
     # Fill in true domain with u[i1_loc]=i1_glob
-    u[p1:-p1] = [i1 for i1 in range(s1,e1+1)]
+    u[p1:-p1] = xp.asarray([i1 for i1 in range(s1,e1+1)])
 
     request = synchronizer.prepare_communications(u)
 
@@ -101,7 +101,7 @@ def run_cart_1d( data_exchanger_type, verbose=False ):
     # CHECK RESULTS
     #---------------------------------------------------------------------------
     # Verify that ghost cells contain correct data (note periodic domain!)
-    success = all( u[:] == [i1%n1 for i1 in range(s1-p1,e1+p1+1)] )
+    success = bool( (u[:] == xp.asarray([i1%n1 for i1 in range(s1-p1,e1+p1+1)])).all() )
 
     # MASTER only: collect information from all processes
     success_global = comm.reduce( success, op=MPI.LAND, root=0 )

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/tests/test_cart_2d.py

@@ -92,7 +92,9 @@ def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis=
     #---------------------------------------------------------------------------
 
     # Fill in true domain with u[i1_loc,i2_loc,:]=[i1_glob,i2_glob]
-    u[p1:-p1,p2:-p2,:] = [[(i1,i2) for i2 in range(s2,e2+1)] for i1 in range(s1,e1+1)]
+    u[p1:-p1,p2:-p2,:] = xp.asarray(
+        [[(i1, i2) for i2 in range(s2, e2 + 1)] for i1 in range(s1, e1 + 1)]
+    )
 
 
     request = synchronizer.prepare_communications(u)
@@ -109,7 +111,7 @@ def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis=
     val = lambda i1,i2: (i1%n1,i2) if 0<=i2<n2 else (0,0)
     uex = [[val(i1,i2) for i2 in range(s2-p2,e2+p2+1)] for i1 in range(s1-p1,e1+p1+1)]
 
-    success = (u == uex).all()
+    success = (u == xp.asarray(uex)).all()
 
     # MASTER only: collect information from all processes
     success_global = comm.reduce( success, op=MPI.LAND, root=0 )
@@ -161,4 +163,3 @@ if __name__=='__main__':
             print( "PASSED", end='\n\n', flush=True )
         else:
             print( "FAILED", end='\n\n', flush=True )
-

{feectools-0.3.0 → feectools-0.3.1}/feectools/ddm/tests/test_cart_3d.py

@@ -98,9 +98,11 @@ def run_cart_3d( data_exchanger_type, verbose=False ):
     #---------------------------------------------------------------------------
 
     # Fill in true domain with u[i1_loc,i2_loc,i3_loc,:]=[i1_glob,i2_glob,i3_glob]
-    u[p1:-p1,p2:-p2,p3:-p3,:] = [[[(i1,i2,i3) for i3 in range(s3,e3+1)] \
-                                              for i2 in range(s2,e2+1)] \
-                                              for i1 in range(s1,e1+1)]
+    u[p1:-p1,p2:-p2,p3:-p3,:] = xp.asarray(
+        [[[(i1, i2, i3) for i3 in range(s3, e3 + 1)]
+          for i2 in range(s2, e2 + 1)]
+         for i1 in range(s1, e1 + 1)]
+    )
 
     request = synchronizer.prepare_communications(u)
     # Update ghost regions
@@ -118,7 +120,7 @@ def run_cart_3d( data_exchanger_type, verbose=False ):
                            for i2 in range(s2-p2,e2+p2+1)] \
                            for i1 in range(s1-p1,e1+p1+1)]
 
-    success = (u == uex).all()
+    success = (u == xp.asarray(uex)).all()
 
     # MASTER only: collect information from all processes
     success_global = comm.reduce( success, op=MPI.LAND, root=0 )

feectools-0.3.1/feectools/ddm/tests/test_device_binding.py

@@ -0,0 +1,18 @@
+"""Each process is bound to its own GPU when feectools.ddm.cart is imported."""
+import cunumpy as xp
+import pytest
+from cunumpy.cuda import device_count
+from cunumpy.kernel_testing import requires_cupy
+from cunumpy.mpi import local_rank
+
+
+@requires_cupy
+def test_rank_is_bound_to_its_local_device():
+    import cupy as cp
+
+    import feectools.ddm.cart  # noqa: F401 -- binds the device on import
+
+    if xp.get_backend() != "cupy":
+        pytest.skip("device binding only happens on the CuPy backend")
+    expected = local_rank() % device_count()
+    assert cp.cuda.runtime.getDevice() == expected