feectools 0.2.0__tar.gz → 0.3.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.2.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.2.0 → feectools-0.3.1}/feectools/api/essential_bc.py

@@ -72,6 +72,9 @@ def apply_essential_bc_stencil(a, *, axis, ext, order, identity=False):
     if isinstance(a, StencilVector):
         V = a.space
         n = V.ndim
+        # Boundary entries may be ghost entries of neighbouring processes, which
+        # all call this function: their ghost regions are no longer up to date.
+        a.ghost_regions_in_sync = False
     elif isinstance(a, StencilMatrix):
         V = a.codomain
         n = V.ndim * 2

{feectools-0.2.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.2.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.2.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.2.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.2.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.2.0 → feectools-0.3.1}/feectools/ddm/cart.py

@@ -1,19 +1,22 @@
 # coding: utf-8
 
+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
@@ -409,42 +412,100 @@ class DomainDecomposition:
     def refine(self, ncells, global_element_starts, global_element_ends):
         """ Create the new Cartesian decomposition of the refined domain.
 
+        The process topology (and its communicators) is shared with ``self``.
+
         Parameters
         ----------
         ncells : list or tuple of int
             Number of cells of refined space.
 
-        global_starts: list of list of int
-            The starts of the coefficients for every process along each direction.
+        global_element_starts : list of list of int
+            The element starts for every process along each direction.
 
-        global_ends: list of list of int
-            The ends of the coefficients for every process along each direction.
+        global_element_ends : list of list of int
+            The element ends for every process along each direction.
 
         Returns
         -------
-        domain : CartDecomposition
-            Cartesian decomposition of the refined domain.
+        domain : DomainDecomposition
+            Domain decomposition of the refined domain.
         """
 
         # Check input arguments
         assert len( ncells ) == len( self.ncells )
         assert all(nc>=snc for nc, snc in zip(ncells, self.ncells))
 
-        domain         = DomainDecomposition(self.ncells, self.periods, comm=self.comm,
-                                            global_comm=self.global_comm, num_threads=self.num_threads,
-                                            size=self.size)
-        domain._ncells = tuple ( ncells )
+        return self._with_element_partition(ncells, global_element_starts, global_element_ends)
+
+    def coarsen(self, factors):
+        """ Create the Cartesian decomposition of a coarsened domain, aligned with ``self``.
+
+        Along axis ``i`` every ``factors[i]`` consecutive cells are merged into one coarse cell.
+        Each process owns the coarse cells covering exactly its fine cells, so the process
+        topology (and its communicators) is shared with ``self``. This requires that the
+        element starts and ends+1 of every process are divisible by ``factors[i]``.
+
+        Parameters
+        ----------
+        factors : list or tuple of int
+            Coarsening factor (>= 1) along each direction.
+
+        Returns
+        -------
+        domain : DomainDecomposition
+            Domain decomposition of the coarse domain.
+        """
+
+        assert len( factors ) == self.ndim
+        assert all( isinstance(f, (int, np.integer)) and f >= 1 for f in factors )
+
+        ncells = []
+        global_element_starts = []
+        global_element_ends   = []
+        for axis, f in enumerate(factors):
+            gs = xp.asarray(self._global_element_starts[axis])
+            ge = xp.asarray(self._global_element_ends  [axis])
+            if self._ncells[axis] % f != 0 or xp.any(gs % f != 0) or xp.any((ge + 1) % f != 0):
+                raise ValueError(
+                    f"Cannot coarsen axis {axis} by a factor {f}: ncells={self._ncells[axis]}, "
+                    f"element starts={gs.tolist()}, ends={ge.tolist()} are not all aligned."
+                )
+            ncells.append(self._ncells[axis] // f)
+            global_element_starts.append(xp.array(gs // f))
+            global_element_ends  .append(xp.array((ge + 1) // f - 1))
+
+        return self._with_element_partition(ncells, global_element_starts, global_element_ends)
+
+    def _with_element_partition(self, ncells, global_element_starts, global_element_ends):
+        """ Return a copy of ``self`` with the same process topology but a new element partition.
+
+        Communicators are shared (not duplicated), hence this method is not collective.
+        """
+
+        assert len( ncells ) == self.ndim
+        for axis in range(self.ndim):
+            gs = xp.asarray(global_element_starts[axis])
+            ge = xp.asarray(global_element_ends  [axis])
+            assert len(gs) == len(ge) == self._nprocs[axis], \
+                f"Axis {axis}: need one block per process ({self._nprocs[axis]}), got {len(gs)}."
+            assert gs[0] == 0 and ge[-1] == ncells[axis] - 1, \
+                f"Axis {axis}: blocks must cover [0, {ncells[axis] - 1}]."
+            assert xp.all(ge >= gs), f"Axis {axis}: empty blocks are not allowed."
+            assert xp.all(gs[1:] == ge[:-1] + 1), f"Axis {axis}: blocks must be contiguous."
+
+        domain = copy.copy(self)
+        domain._ncells = tuple( int(n) for n in ncells )
 
         # Store arrays with all the starts and ends along each direction for every process
-        domain._global_element_starts = tuple(global_element_starts)
-        domain._global_element_ends   = tuple(global_element_ends)
+        domain._global_element_starts = list(global_element_starts)
+        domain._global_element_ends   = list(global_element_ends)
         if self.is_comm_null:return domain
 
         # Start/end values of global indices (without ghost regions)
         domain._starts = tuple( domain._global_element_starts[axis][c] for axis,c in zip(range(self._ndims), self._coords) )
         domain._ends   = tuple( domain._global_element_ends  [axis][c] for axis,c in zip(range(self._ndims), self._coords) )
 
-        domain._local_ncells = tuple(e-s+1 for s,e in zip(self._starts, self._ends))
+        domain._local_ncells = tuple(e-s+1 for s,e in zip(domain._starts, domain._ends))
         return domain
 
 #==================================================================================
@@ -482,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)
@@ -494,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
@@ -510,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
@@ -522,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) )
@@ -546,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.2.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.2.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.2.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.2.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.2.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