From 0a55cbd7527016688075c976232b97f7df92b670 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 30 Sep 2026 23:53:54 +0200 Subject: [PATCH 1/3] CUDA step 1: run feectools with the CuPy backend Make feectools work when cunumpy's backend is CuPy, without any device kernels yet: kernels that stay on the host still copy their arrays. - Wrap all Pyccel kernels (stencil, B-splines, field evaluation, DOF kernels) in cunumpy.PyccelKernel, so they accept CuPy arrays. - Keep host-only metadata on NumPy: MPI/index bookkeeping in ddm (cart, partition, petsc) and fem.partitioning, Kronecker solver sizes, and index arithmetic with Python ints (compute_diag_len, math.prod). - Stage data for host-only libraries by array, not by global backend: LAPACK/SuperLU direct solvers, SciPy FFT, SciPy sparse products. - Fix calls that ran host kernels on device arrays: the second stencil2coo call in StencilMatrix.tosparse and the conjugate transpose. - Vectorize the construction of the 1D collocation matrices in the global projectors (element-wise indexing was a device round trip per entry: 334 s of a 348 s Derham setup on the GPU). - GMRES: take real scalars from CuPy views before modifying them. - Fix StencilMatrix._update_ghost_regions_serial: the ghost region is pads * shifts wide (wrong whenever shifts > 1, on both backends). - Tests: work with CuPy arrays; skip PETSc tests without petsc4py. Serial tests pass on both backends (core, ddm, fem, linalg). Co-Authored-By: Claude Opus 5.5 --- conftest.py | 5 + feectools/core/bsplines.py | 82 ++++++----- feectools/core/tests/test_bsplines.py | 2 +- feectools/core/tests/test_bsplines_kernel.py | 17 +-- feectools/core/tests/test_bsplines_pyccel.py | 17 ++- feectools/ddm/cart.py | 33 +++-- feectools/ddm/partition.py | 3 +- feectools/ddm/petsc.py | 2 +- feectools/ddm/tests/test_cart_2d.py | 7 +- feectools/ddm/tests/test_cart_3d.py | 10 +- feectools/feec/derivatives.py | 22 ++- feectools/feec/global_geometric_projectors.py | 134 ++++++++++-------- feectools/fem/partitioning.py | 13 +- feectools/fem/splines.py | 15 +- feectools/fem/tensor.py | 21 ++- feectools/fem/tests/analytical_profiles_1d.py | 2 + feectools/fem/tests/utilities.py | 4 + feectools/linalg/direct_solvers.py | 47 +++--- feectools/linalg/fft.py | 9 +- feectools/linalg/kron.py | 42 +++--- feectools/linalg/solvers.py | 12 +- feectools/linalg/sparse.py | 13 +- feectools/linalg/stencil.py | 89 +++++++----- feectools/linalg/tests/test_fft.py | 3 +- .../linalg/tests/test_kron_stencil_matrix.py | 3 +- feectools/linalg/tests/test_linalg.py | 57 ++++---- .../tests/test_stencil_interface_matrix.py | 14 +- feectools/utilities/utils.py | 24 ++-- 28 files changed, 419 insertions(+), 283 deletions(-) diff --git a/conftest.py b/conftest.py index 6efc25b37..c7c2f8d42 100644 --- a/conftest.py +++ b/conftest.py @@ -1,6 +1,7 @@ """Root-level pytest configuration.""" import pytest import sys +import importlib.util from pathlib import Path @@ -32,6 +33,7 @@ def pytest_collection_modifyitems(config, items): items_to_remove = [] skip = pytest.mark.skip(reason="Requires optional dependency (sympde)") + petsc_available = importlib.util.find_spec("petsc4py") is not None for item in items: # Skip if module is in skip list @@ -39,6 +41,9 @@ def pytest_collection_modifyitems(config, items): items_to_remove.append(item) continue + if item.get_closest_marker("petsc") and not petsc_available: + item.add_marker(pytest.mark.skip(reason="petsc4py is not installed")) + # If running with xdist, automatically skip mpi and petsc tests if config.pluginmanager.has_plugin("xdist"): if item.get_closest_marker("mpi") or item.get_closest_marker("petsc"): diff --git a/feectools/core/bsplines.py b/feectools/core/bsplines.py index 189d95f0e..2ec9d7cd5 100644 --- a/feectools/core/bsplines.py +++ b/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) diff --git a/feectools/core/tests/test_bsplines.py b/feectools/core/tests/test_bsplines.py index 5ec11e777..92a1431e2 100644 --- a/feectools/core/tests/test_bsplines.py +++ b/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 diff --git a/feectools/core/tests/test_bsplines_kernel.py b/feectools/core/tests/test_bsplines_kernel.py index 441057eb3..ef6e148a4 100644 --- a/feectools/core/tests/test_bsplines_kernel.py +++ b/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) diff --git a/feectools/core/tests/test_bsplines_pyccel.py b/feectools/core/tests/test_bsplines_pyccel.py index e39bd370e..4be129e6d 100644 --- a/feectools/core/tests/test_bsplines_pyccel.py +++ b/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 diff --git a/feectools/ddm/cart.py b/feectools/ddm/cart.py index 2b2b58b41..48b6942e1 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -1,11 +1,14 @@ # 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 +from cunumpy.xp import array_backend, to_numpy + # Initialize CUDA context before MPI if using CuPy backend if array_backend.backend == "cupy": try: @@ -482,6 +485,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 +503,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 @@ -510,6 +519,11 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads 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 +536,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) ) @@ -546,11 +564,6 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads # 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) diff --git a/feectools/ddm/partition.py b/feectools/ddm/partition.py index 538a58cc9..a187557a5 100644 --- a/feectools/ddm/partition.py +++ b/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') diff --git a/feectools/ddm/petsc.py b/feectools/ddm/petsc.py index 4f3a8b9c3..d03d3a449 100644 --- a/feectools/ddm/petsc.py +++ b/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 diff --git a/feectools/ddm/tests/test_cart_2d.py b/feectools/ddm/tests/test_cart_2d.py index 66ec3fba1..d9ae40853 100644 --- a/feectools/ddm/tests/test_cart_2d.py +++ b/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<=i2cupy conversion, + # cupy->numpy needs an explicit .get()/xp.to_numpy()) -- these + # diagonals are tiny and only ever feed this one-time host-side + # sparse assembly, never a device computation. + import numpy as np + + maindiag = np.ones(domain_local) * (-sign) + adddiag = np.ones(domain_local) * sign # handle special case with not self.domain.parallel and not with_pads and periodic if self.domain.periods[d] and not self.domain.parallel and not with_pads: # then: add element to other side of the array - adddiagcirc = xp.array([sign]) + adddiagcirc = np.array([sign]) offsets = (-codomain_local+1, 0, 1) diags = (adddiagcirc, maindiag, adddiag) else: diff --git a/feectools/feec/global_geometric_projectors.py b/feectools/feec/global_geometric_projectors.py index 3dc7752be..126995ed3 100644 --- a/feectools/feec/global_geometric_projectors.py +++ b/feectools/feec/global_geometric_projectors.py @@ -12,6 +12,20 @@ from feectools.fem.basic import FemField from feectools.feec import dof_kernels +from cunumpy import PyccelKernel + +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +for _name in ( + 'evaluate_dofs_1d_0form', 'evaluate_dofs_1d_1form', + 'evaluate_dofs_2d_0form', 'evaluate_dofs_2d_1form_hcurl', 'evaluate_dofs_2d_1form_hdiv', + 'evaluate_dofs_2d_2form', 'evaluate_dofs_2d_vec', + 'evaluate_dofs_3d_0form', 'evaluate_dofs_3d_1form', 'evaluate_dofs_3d_2form', + 'evaluate_dofs_3d_3form', 'evaluate_dofs_3d_vec', +): + setattr(dof_kernels, _name, PyccelKernel(getattr(dof_kernels, _name))) +del _name + from feectools.fem.tensor import TensorFemSpace from feectools.fem.vector import VectorFemSpace, MultipatchFemSpace @@ -29,15 +43,44 @@ def _to_numpy_for_kernel(*args): """Convert CuPy arrays to NumPy for compiled kernel calls.""" - result = [] - for arg in args: - if hasattr(arg, 'get'): # CuPy array - result.append(arg.get()) - else: - result.append(arg) + result = [arg.get() if xp.is_gpu(arg) else arg for arg in args] return result if len(result) > 1 else result[0] +#============================================================================== +def _fill_collocation_stencil(M, mat, s, e, p, m): + """Scatter the non-zeros of a 1D collocation matrix into stencil storage. + + `mat` is the dense 1D interpolation/histopolation matrix; only its rows + owned by this process (`s <= row <= e`) are written, at the stencil + offsets `(row - s + m*p, (col + p - row) % ncols)`. + + Both the index arithmetic and the value lookup happen on the host, and the + result reaches `M._data` in a single assignment. Doing it the obvious way + -- `for row_i, col_i in zip(*xp.nonzero(mat))` with a `row_i in range(...)` + membership test -- pulls every index and every value off the array one + element at a time, which on a device backend is a synchronizing round trip + apiece: it cost 334 s of a 348 s Derham build on a 128 x 1024 x 1 grid, + against 1.5 s for the whole build on NumPy. + + Returns the host copy of `mat`, which the callers reuse for their + correctness assertion. + """ + mat_host = xp.to_numpy(mat) + + rows, cols = np.nonzero(mat_host) + + owned = (rows >= int(s)) & (rows <= int(e)) + rows = rows[owned] + cols = cols[owned] + + data_rows = rows - int(s) + m * p + data_cols = (cols + p - rows) % mat_host.shape[1] + + M._data[xp.asarray(data_rows), xp.asarray(data_cols)] = xp.asarray(mat_host[rows, cols]) + + return mat_host + #============================================================================== class GlobalGeometricProjector(metaclass=ABCMeta): """ @@ -169,7 +212,7 @@ def __init__(self, space, nquads = None): if cell == 'I': # interpolation case if intp_x[j] is None: - intp_x[j] = V.greville[s:e+1] + intp_x[j] = xp.asarray(V.greville[s:e+1]) # V.greville is always NumPy local_intp_x = intp_x[j] # for the grids, make interpolation appear like quadrature @@ -178,26 +221,10 @@ def __init__(self, space, nquads = None): solvercells += [V._interpolator] # make 1D collocation matrix in stencil format - if array_backend.backend == "cupy": - V_imat = xp.asarray(V.imat) # converts to cupy array if not already - else: - V_imat = V.imat - - - row_indices, col_indices = xp.nonzero(V_imat) - - for row_i, col_i in zip(row_indices, col_indices): - - # only consider row indices on process - if row_i in range(int(V_cart.starts[0]), int(V_cart.ends[0]) + 1): - row_i_loc = row_i - s - - - M._data[row_i_loc + m*p, (col_i + p - row_i)%V.imat.shape[1]] = V_imat[row_i, col_i] + imat_host = _fill_collocation_stencil(M, V.imat, s, e, p, m) # check if stencil matrix was built correctly - # assert xp.allclose(M.toarray()[s:e + 1], V_imat[s:e + 1]) - assert xp.allclose(M.toarray()[int(s):int(e) + 1], V_imat[int(s):int(e) + 1]) + assert np.allclose(xp.to_numpy(M.toarray())[int(s):int(e) + 1], imat_host[int(s):int(e) + 1]) # TODO Fix toarray() for multiplicity m > 1 matrixcells += [M.copy()] @@ -220,21 +247,10 @@ def __init__(self, space, nquads = None): solvercells += [V._histopolator] # make 1D collocation matrix in stencil format - # Always use NumPy for indices since they're used for indexing/comparison - if array_backend.backend == "cupy": - row_indices, col_indices = np.nonzero(np.asarray(V.hmat)) - else: - row_indices, col_indices = xp.nonzero(V.hmat) - - for row_i, col_i in zip(row_indices, col_indices): - - # only consider row indices on process - if row_i in range(int(V_cart.starts[0]), int(V_cart.ends[0]) + 1): - row_i_loc = row_i - s - M._data[row_i_loc + m*p, (col_i + p - row_i)%V.hmat.shape[1]] = V.hmat[int(row_i), int(col_i)] + hmat_host = _fill_collocation_stencil(M, V.hmat, s, e, p, m) # check if stencil matrix was built correctly - assert xp.allclose(M.toarray()[int(s):int(e) + 1], V.hmat[int(s):int(e) + 1]) + assert np.allclose(xp.to_numpy(M.toarray())[int(s):int(e) + 1], hmat_host[int(s):int(e) + 1]) matrixcells += [M.copy()] @@ -816,7 +832,7 @@ def evaluate_dofs_1d_0form( F_temp_np, f_pts_np = _to_numpy_for_kernel(F_temp, f_pts) dof_kernels.evaluate_dofs_1d_0form(F_temp_np, f_pts_np) - if hasattr(F_temp, 'get'): + if xp.is_gpu(F_temp): F_temp[:] = xp.asarray(F_temp_np) F[:] = F_temp @@ -838,7 +854,7 @@ def evaluate_dofs_1d_1form( quad_w1_np, F_temp_np, f_pts_np = _to_numpy_for_kernel(quad_w1, F_temp, f_pts) dof_kernels.evaluate_dofs_1d_1form(quad_w1_np, F_temp_np, f_pts_np) - if hasattr(F_temp, 'get'): + if xp.is_gpu(F_temp): F_temp[:] = xp.asarray(F_temp_np) F[:] = F_temp @@ -864,7 +880,7 @@ def evaluate_dofs_2d_0form( F_temp_np, f_pts_np = _to_numpy_for_kernel(F_temp, f_pts) dof_kernels.evaluate_dofs_2d_0form(F_temp_np, f_pts_np) - if hasattr(F_temp, 'get'): + if xp.is_gpu(F_temp): F_temp[:] = xp.asarray(F_temp_np) F[:, :] = F_temp @@ -894,9 +910,9 @@ def evaluate_dofs_2d_1form_hcurl( quad_w1_np, quad_w2_np, F1_temp_np, F2_temp_np, f1_pts_np, f2_pts_np = _to_numpy_for_kernel(quad_w1, quad_w2, F1_temp, F2_temp, f1_pts, f2_pts) dof_kernels.evaluate_dofs_2d_1form_hcurl(quad_w1_np, quad_w2_np, F1_temp_np, F2_temp_np, f1_pts_np, f2_pts_np) - if hasattr(F1_temp, 'get'): + if xp.is_gpu(F1_temp): F1_temp[:] = xp.asarray(F1_temp_np) - if hasattr(F2_temp, 'get'): + if xp.is_gpu(F2_temp): F2_temp[:] = xp.asarray(F2_temp_np) F1[:, :] = F1_temp @@ -927,9 +943,9 @@ def evaluate_dofs_2d_1form_hdiv( quad_w1_np, quad_w2_np, F1_temp_np, F2_temp_np, f1_pts_np, f2_pts_np = _to_numpy_for_kernel(quad_w1, quad_w2, F1_temp, F2_temp, f1_pts, f2_pts) dof_kernels.evaluate_dofs_2d_1form_hdiv(quad_w1_np, quad_w2_np, F1_temp_np, F2_temp_np, f1_pts_np, f2_pts_np) - if hasattr(F1_temp, 'get'): + if xp.is_gpu(F1_temp): F1_temp[:] = xp.asarray(F1_temp_np) - if hasattr(F2_temp, 'get'): + if xp.is_gpu(F2_temp): F2_temp[:] = xp.asarray(F2_temp_np) F1[:, :] = F1_temp @@ -952,7 +968,7 @@ def evaluate_dofs_2d_2form( quad_w1_np, quad_w2_np, F_temp_np, f_pts_np = _to_numpy_for_kernel(quad_w1, quad_w2, F_temp, f_pts) dof_kernels.evaluate_dofs_2d_2form(quad_w1_np, quad_w2_np, F_temp_np, f_pts_np) - if hasattr(F_temp, 'get'): + if xp.is_gpu(F_temp): F_temp[:] = xp.asarray(F_temp_np) F[:, :] = F_temp @@ -978,9 +994,9 @@ def evaluate_dofs_2d_vec( F1_temp_np, F2_temp_np, f1_pts_np, f2_pts_np = _to_numpy_for_kernel(F1_temp, F2_temp, f1_pts, f2_pts) dof_kernels.evaluate_dofs_2d_vec(F1_temp_np, F2_temp_np, f1_pts_np, f2_pts_np) - if hasattr(F1_temp, 'get'): + if xp.is_gpu(F1_temp): F1_temp[:] = xp.asarray(F1_temp_np) - if hasattr(F2_temp, 'get'): + if xp.is_gpu(F2_temp): F2_temp[:] = xp.asarray(F2_temp_np) F1[:, :] = F1_temp @@ -1008,7 +1024,7 @@ def evaluate_dofs_3d_0form( F_temp_np, f_pts_np = _to_numpy_for_kernel(F_temp, f_pts) dof_kernels.evaluate_dofs_3d_0form(F_temp_np, f_pts_np) - if hasattr(F_temp, 'get'): + if xp.is_gpu(F_temp): F_temp[:] = xp.asarray(F_temp_np) F[:, :, :] = F_temp @@ -1043,11 +1059,11 @@ def evaluate_dofs_3d_1form( quad_w1_np, quad_w2_np, quad_w3_np, F1_temp_np, F2_temp_np, F3_temp_np, f1_pts_np, f2_pts_np, f3_pts_np = _to_numpy_for_kernel(quad_w1, quad_w2, quad_w3, F1_temp, F2_temp, F3_temp, f1_pts, f2_pts, f3_pts) dof_kernels.evaluate_dofs_3d_1form(quad_w1_np, quad_w2_np, quad_w3_np, F1_temp_np, F2_temp_np, F3_temp_np, f1_pts_np, f2_pts_np, f3_pts_np) - if hasattr(F1_temp, 'get'): + if xp.is_gpu(F1_temp): F1_temp[:] = xp.asarray(F1_temp_np) - if hasattr(F2_temp, 'get'): + if xp.is_gpu(F2_temp): F2_temp[:] = xp.asarray(F2_temp_np) - if hasattr(F3_temp, 'get'): + if xp.is_gpu(F3_temp): F3_temp[:] = xp.asarray(F3_temp_np) F1[:, :, :] = F1_temp @@ -1084,11 +1100,11 @@ def evaluate_dofs_3d_2form( quad_w1_np, quad_w2_np, quad_w3_np, F1_temp_np, F2_temp_np, F3_temp_np, f1_pts_np, f2_pts_np, f3_pts_np = _to_numpy_for_kernel(quad_w1, quad_w2, quad_w3, F1_temp, F2_temp, F3_temp, f1_pts, f2_pts, f3_pts) dof_kernels.evaluate_dofs_3d_2form(quad_w1_np, quad_w2_np, quad_w3_np, F1_temp_np, F2_temp_np, F3_temp_np, f1_pts_np, f2_pts_np, f3_pts_np) - if hasattr(F1_temp, 'get'): + if xp.is_gpu(F1_temp): F1_temp[:] = xp.asarray(F1_temp_np) - if hasattr(F2_temp, 'get'): + if xp.is_gpu(F2_temp): F2_temp[:] = xp.asarray(F2_temp_np) - if hasattr(F3_temp, 'get'): + if xp.is_gpu(F3_temp): F3_temp[:] = xp.asarray(F3_temp_np) F1[:, :, :] = F1_temp @@ -1112,7 +1128,7 @@ def evaluate_dofs_3d_3form( quad_w1_np, quad_w2_np, quad_w3_np, F_temp_np, f_pts_np = _to_numpy_for_kernel(quad_w1, quad_w2, quad_w3, F_temp, f_pts) dof_kernels.evaluate_dofs_3d_3form(quad_w1_np, quad_w2_np, quad_w3_np, F_temp_np, f_pts_np) - if hasattr(F_temp, 'get'): + if xp.is_gpu(F_temp): F_temp[:] = xp.asarray(F_temp_np) F[:, :, :] = F_temp @@ -1141,11 +1157,11 @@ def evaluate_dofs_3d_vec( F1_temp_np, F2_temp_np, F3_temp_np, f1_pts_np, f2_pts_np, f3_pts_np = _to_numpy_for_kernel(F1_temp, F2_temp, F3_temp, f1_pts, f2_pts, f3_pts) dof_kernels.evaluate_dofs_3d_vec(F1_temp_np, F2_temp_np, F3_temp_np, f1_pts_np, f2_pts_np, f3_pts_np) - if hasattr(F1_temp, 'get'): + if xp.is_gpu(F1_temp): F1_temp[:] = xp.asarray(F1_temp_np) - if hasattr(F2_temp, 'get'): + if xp.is_gpu(F2_temp): F2_temp[:] = xp.asarray(F2_temp_np) - if hasattr(F3_temp, 'get'): + if xp.is_gpu(F3_temp): F3_temp[:] = xp.asarray(F3_temp_np) F1[:, :, :] = F1_temp diff --git a/feectools/fem/partitioning.py b/feectools/fem/partitioning.py index 4df8bae0b..51987884b 100644 --- a/feectools/fem/partitioning.py +++ b/feectools/fem/partitioning.py @@ -1,8 +1,10 @@ # -*- 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 +import numpy as xp # this module is host-only MPI/index bookkeeping, never device data from feectools.ddm.cart import CartDecomposition, InterfaceCartDecomposition, create_interfaces_cart from feectools.core.bsplines import elements_spans @@ -58,9 +60,8 @@ def partition_coefficients(domain_decomposition, spaces, min_blocks=None): m = multiplicity[axis] # Convert to numpy if CuPy (needed for MPI operations later) - if hasattr(ee, 'get'): - ee = ee.get() - + ee = cunumpy.to_numpy(ee) + global_ends [axis] = m*(ee+1)-1 global_ends [axis][-1] = npts[axis]-1 global_starts[axis] = xp.array([0] + (global_ends[axis][:-1]+1).tolist()) @@ -69,8 +70,8 @@ def partition_coefficients(domain_decomposition, spaces, min_blocks=None): min_blocks = [None] * ndims for s, e, V, mb in zip(global_starts, global_ends, spaces, min_blocks): - s_host = s.get() if hasattr(s, 'get') else np.asarray(s) - e_host = e.get() if hasattr(e, 'get') else np.asarray(e) + s_host = cunumpy.to_numpy(s) + e_host = cunumpy.to_numpy(e) local_sizes = e_host - s_host + 1 if V.periodic or mb is None: diff --git a/feectools/fem/splines.py b/feectools/fem/splines.py index ac62f9fa7..b292cfa2c 100644 --- a/feectools/fem/splines.py +++ b/feectools/fem/splines.py @@ -189,19 +189,13 @@ def init_interpolation( self, dtype=float ): # Convert to CSC format and compute sparse LU decomposition # Convert to LAPACK banded format (see DGBTRF function) - if hasattr(imat, 'get'): - imat = imat.get() - else: - imat = _np.asanyarray(imat) + imat = xp.to_numpy(imat) self._interpolator = SparseSolver( csc_matrix( imat ) ) else: # Convert to LAPACK banded format (see DGBTRF function) - if array_backend.backend == "cupy": - imat = imat.get() - else: - imat = _np.asanyarray(imat) + imat = xp.to_numpy(imat) dmat = dia_matrix( imat ) l = abs( dmat.offsets.min() ) u = dmat.offsets.max() @@ -231,10 +225,7 @@ def init_histopolation( self, dtype=float): xgrid = self.ext_greville, multiplicity = self._multiplicity ) - if hasattr(imat, 'get'): - imat = imat.get() - else: - imat = _np.asanyarray(imat) + imat = xp.to_numpy(imat) self.hmat= imat if self.periodic: diff --git a/feectools/fem/tensor.py b/feectools/fem/tensor.py index 513635c52..f5b9ed2eb 100644 --- a/feectools/fem/tensor.py +++ b/feectools/fem/tensor.py @@ -44,6 +44,23 @@ eval_fields_3d_weighted, eval_fields_3d_irregular_weighted) +from cunumpy import PyccelKernel + +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +eval_fields_1d_no_weights = PyccelKernel(eval_fields_1d_no_weights) +eval_fields_1d_irregular_no_weights = PyccelKernel(eval_fields_1d_irregular_no_weights) +eval_fields_1d_weighted = PyccelKernel(eval_fields_1d_weighted) +eval_fields_1d_irregular_weighted = PyccelKernel(eval_fields_1d_irregular_weighted) +eval_fields_2d_no_weights = PyccelKernel(eval_fields_2d_no_weights) +eval_fields_2d_irregular_no_weights = PyccelKernel(eval_fields_2d_irregular_no_weights) +eval_fields_2d_weighted = PyccelKernel(eval_fields_2d_weighted) +eval_fields_2d_irregular_weighted = PyccelKernel(eval_fields_2d_irregular_weighted) +eval_fields_3d_no_weights = PyccelKernel(eval_fields_3d_no_weights) +eval_fields_3d_irregular_no_weights = PyccelKernel(eval_fields_3d_irregular_no_weights) +eval_fields_3d_weighted = PyccelKernel(eval_fields_3d_weighted) +eval_fields_3d_irregular_weighted = PyccelKernel(eval_fields_3d_irregular_weighted) + __all__ = ('TensorFemSpace',) #=============================================================================== @@ -500,7 +517,7 @@ def eval_fields(self, grid, *fields, weights=None, npts_per_cell=None, overlap=0 # -> grid is tensor-product, but npts_per_cell is not the same in each cell elif grid[0].ndim == 1 and npts_per_cell is None: out_fields = self.eval_fields_irregular_tensor_grid(grid, *fields, weights=weights, overlap=overlap) - return [xp.ascontiguousarray(out_fields[..., i]) for i in range(len(fields))] + return [xp.ascontiguousarray(xp.asarray(out_fields[..., i])) for i in range(len(fields))] # Case 3. 1D arrays of coordinates and npts_per_cell is a tuple or an integer # -> grid is tensor-product, and each cell has the same number of evaluation points @@ -512,7 +529,7 @@ def eval_fields(self, grid, *fields, weights=None, npts_per_cell=None, overlap=0 grid[i] = xp.reshape(grid[i], (ncells_i, npts_per_cell[i])) out_fields = self.eval_fields_regular_tensor_grid(grid, *fields, weights=weights, overlap=overlap) # return a list - return [xp.ascontiguousarray(out_fields[..., i]) for i in range(len(fields))] + return [xp.ascontiguousarray(xp.asarray(out_fields[..., i])) for i in range(len(fields))] # Case 4. (self.ldim)D arrays of coordinates and no npts_per_cell # -> unstructured grid diff --git a/feectools/fem/tests/analytical_profiles_1d.py b/feectools/fem/tests/analytical_profiles_1d.py index 9966399da..7921260d3 100644 --- a/feectools/fem/tests/analytical_profiles_1d.py +++ b/feectools/fem/tests/analytical_profiles_1d.py @@ -32,6 +32,7 @@ def poly_order( self ): return -1 def eval( self, x, diff=0 ): + x = xp.asarray(x) return self._k**diff * xp.cos( 0.5*math.pi*diff + self._k*x + self._phi ) def max_norm( self, diff=0 ): @@ -58,6 +59,7 @@ def poly_order( self ): return -1 def eval( self, x, diff=0 ): + x = xp.asarray(x) return self._k**diff * xp.sin( 0.5*math.pi*diff + self._k*x + self._phi ) def max_norm( self, diff=0 ): diff --git a/feectools/fem/tests/utilities.py b/feectools/fem/tests/utilities.py index fa1cd0d26..46596f3e2 100644 --- a/feectools/fem/tests/utilities.py +++ b/feectools/fem/tests/utilities.py @@ -8,6 +8,10 @@ def horner( x, *poly_coeffs ): """ Use Horner's Scheme to evaluate a polynomial of coefficients *poly_coeffs at location x. """ + # Spline metadata (notably Greville abscissas) is intentionally host-side. + # Convert it at the numerical API boundary so CuPy coefficients and NumPy + # coordinates can be combined just as they can under NumPy. + x = xp.asarray(x) p = 0 for c in poly_coeffs[::-1]: p = p*x + c diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index da932c149..f7c79418a 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -70,11 +70,8 @@ def __init__(self, u, l, bmat, transposed=False): else: msg = f'Cannot create a BandedSolver for bmat.dtype = {bmat.dtype}' raise NotImplementedError(msg) - # print(f"{bmat = } {type(bmat) = }") - if hasattr(bmat, "get"): # CuPy array - bmat = bmat.get() - else: - bmat = xp.asanyarray(bmat) + # LAPACK is host-only: the factorization always lives on the host. + bmat = xp.to_numpy(bmat) self._bmat, self._ipiv, self._finfo = self._factor_function(bmat, l, u) self._sinfo = None @@ -144,9 +141,10 @@ def solve(self, rhs, out=None): transposed = self._transposed if out is None: - preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, rhs.T, self._ipiv, - trans=transposed) - out = preout.T + # LAPACK is host-only: solve on the host, return on the caller's backend. + preout, self._sinfo = self._solver_function(self._bmat, self._l, self._u, xp.to_numpy(rhs).T, + self._ipiv, trans=transposed) + out = xp.asarray(preout.T) if xp.is_gpu(rhs) else preout.T else: assert out.shape == rhs.shape @@ -158,17 +156,18 @@ def solve(self, rhs, out=None): # TODO: handle non-contiguous views? - # we want FORTRAN-contiguous data (default is assumed to be C contiguous) - from cunumpy.xp import array_backend - if array_backend.backend == "numpy": - _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True, + # we want FORTRAN-contiguous data (default is assumed to be C contiguous). + # LAPACK is host-only: a device array is solved in a host copy. Decided by + # the array itself, not the global backend, since host arrays may be passed + # on the CuPy backend too. + if xp.is_gpu(out): + out_cpu = xp.to_numpy(out) + _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out_cpu.T, self._ipiv, overwrite_b=True, trans=transposed) + out[...] = xp.asarray(out_cpu) else: - # GPU - out_cpu = out.get() - _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out_cpu.T, self._ipiv, overwrite_b=True, + _, self._sinfo = self._solver_function(self._bmat, self._l, self._u, out.T, self._ipiv, overwrite_b=True, trans=transposed) - out.set(out_cpu) return out @@ -231,18 +230,18 @@ def solve(self, rhs, out=None): transposed = self._transposed if out is None: - out = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T + # SuperLU is host-only: solve on the host, return on the caller's backend. + out = self._splu.solve(xp.to_numpy(rhs).T, trans='T' if transposed else 'N').T + if xp.is_gpu(rhs): + out = xp.asarray(out) else: assert out.shape == rhs.shape assert out.dtype == rhs.dtype - # currently no in-place solve exposed - if array_backend.backend == "numpy": - out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T - else: - rhs_cpu = rhs.get() - result_cpu = self._splu.solve(rhs_cpu.T, trans='T' if transposed else 'N').T - out[:] = xp.asarray(result_cpu) + # currently no in-place solve exposed. SuperLU is host-only; decided by + # the arrays themselves, not the global backend. + result = self._splu.solve(xp.to_numpy(rhs).T, trans='T' if transposed else 'N').T + out[:] = xp.asarray(result) if xp.is_gpu(out) else result return out diff --git a/feectools/linalg/fft.py b/feectools/linalg/fft.py index b98268cd6..7a2ece10e 100644 --- a/feectools/linalg/fft.py +++ b/feectools/linalg/fft.py @@ -59,7 +59,14 @@ def solve(self, rhs, out=None): if out is not rhs: out[:] = rhs - self._function(out) + # SciPy FFT routines are host-only. Preserve the vector backend + # by staging only the local work array when called with CuPy. + if xp.is_gpu(out): + out_cpu = xp.to_numpy(out) + self._function(out_cpu) + out[...] = xp.asarray(out_cpu) + else: + self._function(out) return out diff --git a/feectools/linalg/kron.py b/feectools/linalg/kron.py index 80b013f14..cf413ad75 100644 --- a/feectools/linalg/kron.py +++ b/feectools/linalg/kron.py @@ -1,6 +1,7 @@ #coding = utf-8 from functools import reduce +import numpy as np import cunumpy as xp from scipy.sparse import kron from scipy.sparse import coo_matrix @@ -113,7 +114,10 @@ def dot(self, x, out=None): for jj in xp.ndindex(*pnrows): i_mats = [mat._data[s, j] for s,j,mat in zip(xx, jj, mats)] ii_jj = tuple(i+j+(s-1)*p for i,j,p,s in zip(ii, jj, pads, shifts)) - v += x._data[ii_jj] * xp.prod(i_mats) + # ``array_api_compat.cupy`` does not accept a Python list in + # ``prod``; multiplying the scalar factors also avoids a + # temporary device array in this innermost loop. + v += x._data[ii_jj] * reduce(lambda a, b: a * b, i_mats, 1) out._data[xx] = v @@ -151,7 +155,7 @@ def __getitem__(self, key): cols = key[self.ndim:] mats = self.mats elements = [A[i,j] for A,i,j in zip(mats, rows, cols)] - return xp.prod(elements) + return reduce(lambda a, b: a * b, elements, 1) def tostencil(self): @@ -189,7 +193,7 @@ def _tostencil(M, mats, nrows, nrows_extra, pads, xpads): for kk in xp.ndindex( *ndiags ): values = [mat[i,k] for mat,i,k in zip(mats, ii, kk)] - M[(*ii, *kk)] = xp.prod(values) + M[(*ii, *kk)] = reduce(lambda a, b: a * b, values, 1) # handle partly-multiplied rows new_nrows = nrows.copy() @@ -212,7 +216,7 @@ def _tostencil(M, mats, nrows, nrows_extra, pads, xpads): for kk in xp.ndindex( *ndiags ): values = [mat[i,k] for mat,i,k in zip(mats, ii, kk)] - M[(*ii, *kk)] = xp.prod(values) + M[(*ii, *kk)] = reduce(lambda a, b: a * b, values, 1) new_nrows[d] += er def tosparse(self): @@ -440,15 +444,19 @@ def _setup_solvers(self): Computes the distribution of elements and sets up the solvers (which potentially utilize MPI). """ - # slice sizes - starts = xp.array(self._domain.starts) - ends = xp.array(self._domain.ends) + 1 + # slice sizes -- domain-decomposition bookkeeping (sizes/starts/ends), + # always host-resident regardless of the active array backend, matching + # self._domain.starts/ends (which are already plain host ints, see + # CartDecomposition in feectools.ddm.cart) and the MPI sizes/displacements + # computed from them below in KroneckerSolverParallelPass. + starts = np.array(self._domain.starts) + ends = np.array(self._domain.ends) + 1 self._slice = tuple([slice(s, e) for s,e in zip(starts, ends)]) # local and global sizes nglobals = self._domain.npts nlocals = ends - starts - self._localsize = xp.prod(nlocals) + self._localsize = np.prod(nlocals) mglobals = self._localsize // nlocals self._nlocals = nlocals @@ -491,7 +499,9 @@ def _setup_permutations(self): # we use a single permutation for all steps # it is: (n, 1, 2, ..., n-1) - self._perm = xp.arange(self._ndim) + # host bookkeeping (ndim-sized), like self._shapes/self._nlocals it + # indexes -- see the note on _setup_solvers above. + self._perm = np.arange(self._ndim) self._perm[1:] = self._perm[:-1] self._perm[0] = self._ndim - 1 @@ -793,20 +803,20 @@ def __init__(self, solver, mpi_type, i, cart, mglobal, nglobal, nlocal, localsiz cartstart = cart.global_starts[i] cartsize = cartend - cartstart - # source MPI sizes and disps - # distribute the data like - # (N+1, N+1, ..., N+1, N, N, ...) - # where N = floor(mglobaldata / comm.size) + # source MPI sizes and disps -- these are passed straight to + # mpi4py's Alltoallv as counts/displacements, which (like + # cart.global_starts/global_ends above) must be host int arrays + # regardless of backend; keep this whole computation on numpy. mlocal_pre = mglobal // comm.size mlocal_add = mglobal % comm.size - sourcesizes = xp.full((comm.size,), mlocal_pre, dtype=int) + sourcesizes = np.full((comm.size,), mlocal_pre, dtype=int) sourcesizes[:mlocal_add] += 1 mlocal = sourcesizes[comm.rank] sourcesizes *= nlocal # disps, created from the sizes - sourcedisps = xp.zeros((comm.size+1,), dtype=int) - xp.cumsum(sourcesizes, out=sourcedisps[1:]) + sourcedisps = np.zeros((comm.size+1,), dtype=int) + np.cumsum(sourcesizes, out=sourcedisps[1:]) sourcedisps = sourcedisps[:-1] # target MPI sizes and disps diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index d2e673a5f..e7278d619 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -1897,7 +1897,9 @@ def apply_givens_rotation(self, k, sn, cn): h = self._H[:k+2, k] for i in range(k): - h_i_prev = h[i] + # On CuPy, scalar indexing can retain a view into ``h``. Keep a + # genuine scalar before modifying that entry in-place. + h_i_prev = h[i].item() if xp.is_gpu(h) else h[i] h[i] *= cn[i] h[i] += sn[i] * h[i+1] @@ -1905,9 +1907,11 @@ def apply_givens_rotation(self, k, sn, cn): h[i+1] *= cn[i] h[i+1] -= sn[i] * h_i_prev - mod = (h[k]**2 + h[k+1]**2)**0.5 - cn.append( h[k] / mod ) - sn.append( h[k+1] / mod ) + h_k = h[k].item() if xp.is_gpu(h) else h[k] + h_k1 = h[k+1].item() if xp.is_gpu(h) else h[k+1] + mod = (h_k**2 + h_k1**2)**0.5 + cn.append(h_k / mod) + sn.append(h_k1 / mod) h[k] *= cn[k] h[k] += sn[k] * h[k+1] diff --git a/feectools/linalg/sparse.py b/feectools/linalg/sparse.py index d14a9cb55..3af174dd5 100644 --- a/feectools/linalg/sparse.py +++ b/feectools/linalg/sparse.py @@ -2,6 +2,7 @@ from scipy.sparse import sparray, csr_array, bsr_array from scipy.sparse import spmatrix, csr_matrix, bsr_matrix +import cunumpy as xp from feectools.linalg.basic import LinearOperator from feectools.linalg.basic import VectorSpace, Vector, LinearOperator @@ -97,7 +98,17 @@ def _dot_recursive(self, v, out, ind_V=0, ind_W=0): dim_W = W.dimension dim_V = V.dimension - out[index_global_W].flat += self._matrix[ind_W:ind_W+dim_W, ind_V:ind_V+dim_V] @ v[index_global_V].flat + matrix = self._matrix[ind_W:ind_W+dim_W, ind_V:ind_V+dim_V] + values = v[index_global_V].flat + # SciPy sparse matrices are host-only. Stage just this local + # vector slice when the active backend is CuPy, then assign the + # result back through the backend-neutral vector interface. + if xp.is_gpu(v._data): + product = matrix @ xp.to_numpy(v[index_global_V]).ravel() + target = out[index_global_W] + target[...] += xp.asarray(product).reshape(target.shape) + else: + out[index_global_W].flat += matrix @ values elif isinstance(v, BlockVector): diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 4848595c3..fc53b124e 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -3,11 +3,14 @@ # LICENSE file or go to https://github.com/pyccel/psydac/blob/devel/LICENSE # # for full license details. # #---------------------------------------------------------------------------# + +import math import os import warnings from types import MappingProxyType import cunumpy as xp +from cunumpy import PyccelKernel from cunumpy.xp import array_backend from scipy.sparse import coo_matrix, diags as sp_diags @@ -38,17 +41,11 @@ def _to_numpy_int64(val): """Convert CuPy or NumPy scalar/array to numpy int64.""" import numpy as _np - if hasattr(val, 'get'): - # CuPy array - convert to NumPy first - val = val.get() - return _np.int64(val) + return _np.int64(val.get() if xp.is_gpu(val) else val) def _to_numpy_array(val): """Convert CuPy array to NumPy array, preserving dtype. Return as-is if already NumPy.""" - if hasattr(val, 'get'): - # CuPy array - convert to NumPy - return val.get() - return val + return val.get() if xp.is_gpu(val) else val #========================================================================# Dictionary used to select correct kernel functions based on dimensionality kernels = { @@ -61,6 +58,24 @@ def _to_numpy_array(val): 'C': (None, stencil2coo_1d_C, stencil2coo_2d_C, stencil2coo_3d_C)} } +#======================================================================== + +def _wrap_kernel_table(table): + """Wrap every Pyccel kernel in `table` with PyccelKernel, recursively, + so StencilMatrix/StencilVector operations also work with CuPy arrays + (Pyccel kernels only understand NumPy arrays, see cunumpy.kernel). + """ + if table is None: + return None + if isinstance(table, dict): + return {k: _wrap_kernel_table(v) for k, v in table.items()} + if isinstance(table, tuple): + return tuple(_wrap_kernel_table(v) for v in table) + return PyccelKernel(table) + + +kernels = _wrap_kernel_table(kernels) + #======================================================================== def compute_diag_len(pads, shifts_domain, shifts_codomain, return_padding=False): """ @@ -89,13 +104,15 @@ def compute_diag_len(pads, shifts_domain, shifts_codomain, return_padding=False) ep : (int) Padding that constitutes the starting index of the non zero elements. """ - n = ((xp.ceil((pads+1)/shifts_codomain)-1)*shifts_domain).astype('int') - ep = -xp.minimum(0, n-pads) + # pads/shifts are plain Python ints (per-direction metadata), not device + # arrays, so this is computed with builtins rather than the array backend. + n = int((math.ceil((pads+1)/shifts_codomain)-1)*shifts_domain) + ep = -min(0, n-pads) n = n + ep + pads + 1 if return_padding: - return n.astype('int'), ep.astype('int') + return int(n), int(ep) else: - return n.astype('int') + return int(n) #======================================================================== class StencilVectorSpace(VectorSpace): @@ -214,7 +231,7 @@ def dimension(self): """ The dimension of a vector space V is the cardinality (i.e. the number of vectors) of a basis of V over its base field. """ - return xp.prod(self._npts) + return math.prod(self._npts) # ... @property @@ -1116,7 +1133,7 @@ def dot(self, v, out=None): self._func(self_data_np, v_data_np, out_data_np, **args_np) # Copy result back to CuPy array if needed - if hasattr(out._data, 'get'): + if xp.is_gpu(out._data): import cupy as cp out._data[:] = cp.asarray(out_data_np) else: @@ -1175,7 +1192,7 @@ def vdot( self, v, out=None): self._func(self_data_np, v_data_conj_np, out_data_np, **args_np) # Copy result back to CuPy array if needed - if hasattr(out._data, 'get'): + if xp.is_gpu(out._data): import cupy as cp out_data_conj = cp.conjugate(cp.asarray(out_data_np)) out._data[:] = out_data_conj @@ -1223,8 +1240,9 @@ def transpose(self, conjugate=False, out=None): out_data_np = _to_numpy_array(out._data) if conjugate: - self._transpose_func(_to_numpy_array(xp.conjugate(M_data_np)), out_data_np, **self._transpose_args) - self._transpose_func(xp.conjugate(M._data), out._data, **self._transpose_args) + # This kernel is host-backed. Conjugate the staged host array, + # rather than passing it through CuPy's ufunc dispatcher. + self._transpose_func(M_data_np.conj(), out_data_np, **self._transpose_args) else: self._transpose_func(M_data_np, out_data_np, **self._transpose_args) @@ -1687,7 +1705,7 @@ def tocoo_local(self, order='C'): M = coo_matrix( (data,(rows,cols)), - shape = [xp.prod(nr),xp.prod(nc)], + shape = [math.prod(nr),math.prod(nc)], dtype = self._domain.dtype ) @@ -1749,23 +1767,12 @@ def _tocoo_no_pads(self , order='C'): data[:ind] = cp.asarray(data_np[:ind]) rows[:ind] = cp.asarray(rows_np[:ind]) cols[:ind] = cp.asarray(cols_np[:ind]) - nrl = [_np.int64(e-s+1) for s,e in zip(self.codomain.starts, self.codomain.ends)] - ncl = [_np.int64(i) for i in self._data.shape[nd:]] - ss = [_np.int64(i) for i in ss] - nr = [_np.int64(i) for i in nr] - nc = [_np.int64(i) for i in nc] - dm = [_np.int64(i) for i in dm] - cm = [_np.int64(i) for i in cm] - cpads = [_np.int64(i) for i in cpads] - pp = [_np.int64(i) for i in pp] - stencil2coo = kernels['stencil2coo'][order][nd] - ind = stencil2coo(self._data, data, rows, cols, *nrl, *ncl, *ss, *nr, *nc, *dm, *cm, *cpads, *pp) - - if array_backend.backend == "cupy": + def _host(a): + return xp.to_numpy(a) M = coo_matrix( - (data[:ind].get(), (rows[:ind].get(), cols[:ind].get())), + (_host(data[:ind]), (_host(rows[:ind]), _host(cols[:ind]))), shape=[int(_np.prod(nr)), int(_np.prod(nc))], dtype=self.dtype ) @@ -1854,7 +1861,7 @@ def _tocoo_parallel_with_pads(self , order='C'): # Create Scipy COO matrix M = coo_matrix( (data,(rows,cols)), - shape = [xp.prod(nr), xp.prod(nc)], + shape = [math.prod(nr), math.prod(nc)], dtype = self._domain.dtype ) @@ -1881,7 +1888,10 @@ def _update_ghost_regions_serial(self): for direction in range(self._codomain.ndim): periodic = self._codomain.periods[direction] - p = self._codomain.pads [direction] + # The ghost region is `shifts` copies of `pads` wide, exactly as + # in StencilVector._update_ghost_regions_serial; using `pads` + # alone silently mismatched the MPI exchange whenever shifts > 1. + p = self._codomain.pads[direction] * self._codomain.shifts[direction] if p == 0: continue @@ -2008,7 +2018,7 @@ def set_backend(self, backend, precompiled): # matvec kernel dot_func_name = 'matvec_' + str(self._ndim) + 'd_kernel' - self._func = getattr(stencil_dot_kernels, dot_func_name) + self._func = PyccelKernel(getattr(stencil_dot_kernels, dot_func_name)) # parameter for rectangular matrices add = [int(end_in >= end_out) for end_in, end_out in zip(self.domain.ends, self.codomain.ends)] @@ -2032,7 +2042,7 @@ def set_backend(self, backend, precompiled): # transpose kernel transp_func_name = 'transpose_' + str(self._ndim) + 'd_kernel' - self._transpose_func = getattr(stencil_transpose_kernels, transp_func_name) + self._transpose_func = PyccelKernel(getattr(stencil_transpose_kernels, transp_func_name)) # parameter for rectangular matrices add = [int(end_out >= end_in) for end_in, end_out in zip(self.domain.ends, self.codomain.ends)] @@ -2167,7 +2177,7 @@ def _get_diagonal_indices(self): nrows = [e - s + 1 for s, e in zip(self.codomain.starts, self.codomain.ends)] ndim = self.domain.ndim - indices = [xp.zeros(xp.prod(nrows), dtype=int) for _ in range(2 * ndim)] + indices = [xp.zeros(math.prod(nrows), dtype=int) for _ in range(2 * ndim)] for l, xx in enumerate(xp.ndindex(*nrows)): ii = [m * p + x for m, p, x in zip(dm, dp, xx)] @@ -2239,7 +2249,8 @@ def nbytes(self): return int(self._data.nbytes) def tosparse(self): - return sp_diags(self._data.ravel()) + # scipy.sparse.diags expects a host sequence of diagonal arrays. + return sp_diags([xp.to_numpy(self._data).ravel()], [0]) def toarray(self): return self._data.copy() @@ -2927,7 +2938,7 @@ def _tocoo_no_pads(self): M = coo_matrix( (data,(rows,cols)), - shape = [xp.prod(nr),xp.prod(nc)], + shape = [math.prod(nr),math.prod(nc)], dtype = self.domain.dtype) return M diff --git a/feectools/linalg/tests/test_fft.py b/feectools/linalg/tests/test_fft.py index db772ad49..e261e483f 100644 --- a/feectools/linalg/tests/test_fft.py +++ b/feectools/linalg/tests/test_fft.py @@ -84,7 +84,8 @@ def method_test(seed, comm, config, dtype, classtype, comparison, verbose=False) if verbose: print(f'[{rank}] Vector built', flush=True) - X_glob = comparison(Y_glob) + X_glob = comparison(xp.to_numpy(Y_glob)) if xp.is_gpu(Y_glob) else comparison(Y_glob) + X_glob = xp.asarray(X_glob) compare = classtype(V) X = compare.dot(Y) diff --git a/feectools/linalg/tests/test_kron_stencil_matrix.py b/feectools/linalg/tests/test_kron_stencil_matrix.py index 59d192019..d160394fb 100644 --- a/feectools/linalg/tests/test_kron_stencil_matrix.py +++ b/feectools/linalg/tests/test_kron_stencil_matrix.py @@ -111,4 +111,5 @@ def test_KroneckerStencilMatrix(dtype, npts, pads, periodic): assert (M_sp.T - M.T.tosparse().tocsr()).count_nonzero() == 0 # Test dot product - assert xp.array_equal(M_sp.dot(w.toarray()), M.dot(w).toarray()) + expected = M_sp.dot(xp.to_numpy(w.toarray())) + assert xp.array_equal(xp.asarray(expected), M.dot(w).toarray()) diff --git a/feectools/linalg/tests/test_linalg.py b/feectools/linalg/tests/test_linalg.py index d259f9b0c..e6d1a6d28 100644 --- a/feectools/linalg/tests/test_linalg.py +++ b/feectools/linalg/tests/test_linalg.py @@ -1,5 +1,6 @@ import pytest import cunumpy as xp +import numpy as np from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace from feectools.linalg.basic import LinearOperator, ZeroOperator, IdentityOperator, ComposedLinearOperator, SumLinearOperator, PowerLinearOperator, ScaledLinearOperator @@ -24,7 +25,9 @@ def sparse_equal(a, b): def assert_pos_def(A): assert isinstance(A, LinearOperator) A_array = A.toarray() - assert xp.all(xp.linalg.eigvals(A_array) > 0) + # array-api-compat's CuPy linalg namespace does not expose eigvals. + eigvals = np.linalg.eigvals(xp.to_numpy(A_array)) + assert np.all(eigvals > 0) def compute_global_starts_ends(domain_decomposition, npts): ndims = len(npts) @@ -126,7 +129,7 @@ def test_square_stencil_basic(n1, n2, p1, p2, P1=False, P2=False): for k2 in range(-p2,p2+1): S[:,:,k1,k2] = nonzero_values[k1,k2] S.remove_spurious_entries() - Sa = S.toarray() + Sa = xp.asarray(S.toarray()) nonzero_values1 = dict() for k1 in range(-p1,p1+1): @@ -142,7 +145,7 @@ def test_square_stencil_basic(n1, n2, p1, p2, P1=False, P2=False): for k2 in range(-p2,p2+1): S1[:,:,k1,k2] = nonzero_values1[k1,k2] S1.remove_spurious_entries() - S1a = S1.toarray() + S1a = xp.asarray(S1.toarray()) nonzero_values2 = dict() for k1 in range(-p1,p1+1): @@ -159,7 +162,7 @@ def test_square_stencil_basic(n1, n2, p1, p2, P1=False, P2=False): for k2 in range(-p2,p2+1): S2[:,:,k1,k2] = nonzero_values2[k1,k2] S2.remove_spurious_entries() - S2a = S2.toarray() + S2a = xp.asarray(S2.toarray()) # Construct exact matrices by hand A1 = xp.zeros( S.shape ) @@ -232,8 +235,8 @@ def test_square_stencil_basic(n1, n2, p1, p2, P1=False, P2=False): assert not xp.array_equal(S2a, S2a.T) # using a nonsymmetric matrix throughout assert isinstance(S2.T, StencilMatrix) - assert xp.array_equal(S2.T.toarray(), S2a.T) - assert xp.array_equal(S2.T.T.toarray(), S2a) + assert xp.array_equal(xp.asarray(S2.T.toarray()), xp.asarray(S2a.T)) + assert xp.array_equal(xp.asarray(S2.T.T.toarray()), xp.asarray(S2a)) ### ### 3. Test special cases @@ -478,9 +481,9 @@ def test_in_place_operations(n1, n2, p1, p2, P1=False, P2=False): assert isinstance(I1, ZeroOperator) assert isinstance(I2, IdentityOperator) assert isinstance(I3, ScaledLinearOperator) - assert xp.array_equal(v3.toarray(), xp.dot(v_array, 3)) + assert xp.array_equal(v3.toarray(), v_array * 3) assert isinstance(I4, ScaledLinearOperator) - assert xp.array_equal(v4.toarray(), xp.dot(v_array, 3j)) + assert xp.array_equal(v4.toarray(), v_array * 3j) # testing __iadd__ and __isub__ although not explicitly implemented (in the LinearOperator class) @@ -506,7 +509,7 @@ def test_in_place_operations(n1, n2, p1, p2, P1=False, P2=False): S[:,:,k1,k2] = nonzero_values1[k1,k2] S.remove_spurious_entries() T = S.copy() - Sa = S.toarray() + Sa = xp.asarray(S.toarray()) Z1 += S S += Z2 @@ -519,7 +522,7 @@ def test_in_place_operations(n1, n2, p1, p2, P1=False, P2=False): w = S.dot(v) assert isinstance(S, StencilMatrix) - assert xp.array_equal(w.toarray(), xp.dot(xp.dot(2, Sa), v_array)) + assert xp.array_equal(w.toarray(), 2 * Sa @ v_array) Z3 -= T T -= Z2 @@ -529,7 +532,7 @@ def test_in_place_operations(n1, n2, p1, p2, P1=False, P2=False): assert isinstance(Z3, StencilMatrix) assert isinstance(T, StencilMatrix) - assert xp.array_equal(w2.toarray(), xp.dot(xp.dot(2, Sa), v_array)) + assert xp.array_equal(w2.toarray(), 2 * Sa @ v_array) #=============================================================================== @pytest.mark.parametrize('n1', n1array) @@ -630,7 +633,7 @@ def test_inverse_transpose_interaction(n1, n2, p1, p2, P1=False, P2=False): scaled_matrix = B * xp.random.random() # Ensure the diagonal elements != 1 diagonal_values = scaled_matrix.diagonal(sqrt=False).toarray() sqrt_diagonal_values = scaled_matrix.diagonal(sqrt=True).toarray() - assert xp.array_equal(sqrt_diagonal_values, xp.sqrt(diagonal_values)) + assert xp.array_equal(xp.asarray(sqrt_diagonal_values), xp.sqrt(xp.asarray(diagonal_values))) tol = 1e-5 C = inverse(B, 'cg', tol=tol) @@ -786,20 +789,22 @@ def test_operator_evaluation(n1, n2, p1, p2): b1 = ( B**1 @ u ).toarray() b2 = ( B**2 @ u ).toarray() assert xp.array_equal(uarr, b0) - assert xp.linalg.norm( xp.dot(Bmat, uarr) - b1 ) < 1e-10 - assert xp.linalg.norm( xp.dot(Bmat, xp.dot(Bmat, uarr)) - b2 ) < 1e-10 + assert xp.linalg.norm(xp.asarray(Bmat) @ uarr - b1) < 1e-10 + Bmat_xp = xp.asarray(Bmat) + assert xp.linalg.norm(Bmat_xp @ (Bmat_xp @ uarr) - b2) < 1e-10 bi0 = ( B_ILO**0 @ u ).toarray() bi1 = ( B_ILO**1 @ u ).toarray() bi2 = ( B_ILO**2 @ u ).toarray() - B_inv_mat = xp.linalg.inv(Bmat) - b_inv_arr = xp.matrix.flatten(B_inv_mat) - error_est = 2 + n1 * n2 * xp.max( [ xp.abs(b_inv_arr[i]) for i in range(len(b_inv_arr)) ] ) + Bmat_xp = xp.asarray(Bmat) + B_inv_mat = xp.linalg.inv(Bmat_xp) + b_inv_arr = xp.reshape(B_inv_mat, (-1,)) + error_est = 2 + n1 * n2 * xp.max(xp.abs(b_inv_arr)) assert xp.array_equal(uarr, bi0) - bi12 = xp.linalg.solve(Bmat, uarr) - bi22 = xp.linalg.solve(Bmat, bi12) - assert xp.linalg.norm( (Bmat @ bi12) - uarr ) < tol - assert xp.linalg.norm( (Bmat @ bi22) - bi12 ) < error_est * tol + bi12 = xp.linalg.solve(Bmat_xp, uarr) + bi22 = xp.linalg.solve(Bmat_xp, bi12) + assert xp.linalg.norm( (Bmat_xp @ bi12) - uarr ) < tol + assert xp.linalg.norm( (Bmat_xp @ bi22) - bi12 ) < error_est * tol zeros = U.zeros().toarray() z0 = ( Z**0 @ u ).toarray() @@ -809,22 +814,22 @@ def test_operator_evaluation(n1, n2, p1, p2): assert xp.array_equal(zeros, z1) assert xp.array_equal(zeros, z2) - Smat = S.toarray() + Smat = xp.asarray(S.toarray()) assert_pos_def(S) varr = v.toarray() s0 = ( S**0 @ v ).toarray() s1 = ( S**1 @ v ).toarray() s2 = ( S**2 @ v ).toarray() assert xp.array_equal(varr, s0) - assert xp.linalg.norm( xp.dot(Smat, varr) - s1 ) < 1e-10 - assert xp.linalg.norm( xp.dot(Smat, xp.dot(Smat, varr)) - s2 ) < 1e-10 + assert xp.linalg.norm(Smat @ varr - s1) < 1e-10 + assert xp.linalg.norm(Smat @ (Smat @ varr) - s2) < 1e-10 si0 = ( S_ILO**0 @ v ).toarray() si1 = ( S_ILO**1 @ v ).toarray() si2 = ( S_ILO**2 @ v ).toarray() S_inv_mat = xp.linalg.inv(Smat) - s_inv_arr = xp.matrix.flatten(S_inv_mat) - error_est = 2 + n1 * n2 * xp.max( [ xp.abs(s_inv_arr[i]) for i in range(len(s_inv_arr)) ] ) + s_inv_arr = xp.reshape(S_inv_mat, (-1,)) + error_est = 2 + n1 * n2 * xp.max(xp.abs(s_inv_arr)) assert xp.array_equal(varr, si0) si12 = xp.linalg.solve(Smat, varr) si22 = xp.linalg.solve(Smat, si12) diff --git a/feectools/linalg/tests/test_stencil_interface_matrix.py b/feectools/linalg/tests/test_stencil_interface_matrix.py index cfb755f2a..8406fd350 100644 --- a/feectools/linalg/tests/test_stencil_interface_matrix.py +++ b/feectools/linalg/tests/test_stencil_interface_matrix.py @@ -5,6 +5,7 @@ #---------------------------------------------------------------------------# import pytest import cunumpy as xp +import numpy as np from random import random from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix, StencilInterfaceMatrix @@ -22,7 +23,8 @@ def compute_global_starts_ends(domain_decomposition, npts, pads): global_ends [axis] = ee.copy() global_ends [axis][-1] = npts[axis]-1 - global_starts[axis] = xp.array([0] + (global_ends[axis][:-1]+1).tolist()) + # Cartesian partition data is host metadata, including on CuPy. + global_starts[axis] = np.array([0] + (global_ends[axis][:-1]+1).tolist()) for s, e, p in zip(global_starts, global_ends, pads): assert all(e - s + 1 >= p) @@ -94,7 +96,7 @@ def test_stencil_interface_matrix_1d_serial_init(dtype, n1, p1, s1, axis, ext, P assert M.domain_start == (0,) * M.dim assert M.codomain_start == (0,) * M.dim assert M.flip == (1,) * M.dim - assert xp.array_equal(M.permutation, [0]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([0])) assert M.pads == (p1,) assert M.backend == None assert M._data.shape == (p1 + 1 + 2 * p1 * s1, 1 + 2 * p1) @@ -149,9 +151,9 @@ def test_stencil_interface_matrix_2d_serial_init(dtype, n1, n2, p1, p2, s1, s2, elif axis2 == 1: assert M._data.shape == (n1 + 2 * p1 * s1, p2 + 1 + 2 * p2 * s2, 1 + 2 * p1, 1 + 2 * p2) if axis1 == axis2: - assert xp.array_equal(M.permutation, [0, 1]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([0, 1])) else: - assert xp.array_equal(M.permutation, [1, 0]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([1, 0])) assert M.shape == (n1 * n2, n1 * n2) # =============================================================================== @@ -211,11 +213,11 @@ def test_stencil_interface_matrix_3d_serial_init(dtype, n1, n2, n3, p1, p2, p3, assert M._data.shape == ( n1 + 2 * p1 * s1, n2 + 2 * p2 * s2, p3 + 1 + 2 * p3 * s3, 1 + 2 * p1, 1 + 2 * p2, 1 + 2 * p3) if axis1 == axis2: - assert xp.array_equal(M.permutation, [0, 1, 2]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([0, 1, 2])) else: permutation = [0, 1, 2] permutation[axis1], permutation[axis2] = permutation[axis2], permutation[axis1] - assert xp.array_equal(M.permutation, permutation) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray(permutation)) assert M.shape == (n1 * n2 * n3, n1 * n2 * n3) #=============================================================================== # Parallel TESTS diff --git a/feectools/utilities/utils.py b/feectools/utilities/utils.py index b0043b171..dca6e0f10 100644 --- a/feectools/utilities/utils.py +++ b/feectools/utilities/utils.py @@ -73,9 +73,7 @@ def unroll_edges(domain, xgrid): xA, xB = domain # Convert to numpy if needed (grid arrays should be on CPU) - if hasattr(xgrid, 'get'): - xgrid = xgrid.get() - xgrid = np.asarray(xgrid) + xgrid = xp.to_numpy(xgrid) # Convert to numpy for comparison assert all(np.diff(xgrid) >= 0) @@ -105,19 +103,15 @@ def roll_edges(domain, points): assert xA < xB # Convert domain bounds to same backend as points to ensure compatibility - # First, normalize xA and xB to Python float or correct backend - if hasattr(xA, 'get'): - xA = float(xA.get()) - elif hasattr(xA, '__array__'): - xA = float(xA) - - if hasattr(xB, 'get'): - xB = float(xB.get()) - elif hasattr(xB, '__array__'): - xB = float(xB) - + # First, normalize xA and xB to Python float or correct backend. xp.to_numpy + # handles a CuPy array, a NumPy array/scalar, or a plain Python float uniformly + # (all become something float() accepts), replacing the previous hasattr-based + # get()/__array__ branching. + xA = float(xp.to_numpy(xA)) + xB = float(xp.to_numpy(xB)) + # Now convert to backend of points if needed - if hasattr(points, 'get'): # CuPy array + if xp.is_gpu(points): xA = xp.asarray(xA) xB = xp.asarray(xB) From ccb39f8f64c22b412dd05c0f61f0d52c36bb20e6 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 1 Oct 2026 00:03:17 +0200 Subject: [PATCH 2/3] CUDA step 2: MPI with device buffers Allow MPI on the CuPy backend and make it correct with device buffers (requires a CUDA-aware MPI library). - feectools.ddm.mpi no longer disables MPI when ARRAY_BACKEND=cupy; the segfaults it guarded against come from MPI libraries that are not CUDA-aware. - Call cunumpy.synchronize_for_mpi before every MPI call on device buffers: CuPy kernels run asynchronously and MPI does not know about CUDA streams, so a buffer still being written would be sent silently wrong. Covers the blocking, non-blocking and interface data exchangers, the Allreduce in StencilVectorSpace.inner and the Alltoallv calls of the parallel Kronecker solver. Requires cunumpy >= 0.3.0. - Fix CuPy incompatibilities reached only by the MPI tests: xp.dot/vdot on .flat iterators in StencilInterfaceMatrix._dot and the pure-Python inner product, and test_cart_1d assigning Python lists to CuPy arrays. - Add test_mpi_device.py: distributed results against global references, and that the exchangers synchronize before MPI. With 2 MPI ranks and a CUDA-aware Open MPI, all MPI tests in ddm and linalg pass on both backends; serial tests are unchanged. Co-Authored-By: Claude Opus 5.5 --- feectools/ddm/blocking_data_exchanger.py | 9 + feectools/ddm/interface_data_exchanger.py | 5 + feectools/ddm/mpi.py | 8 +- feectools/ddm/nonblocking_data_exchanger.py | 8 + feectools/ddm/tests/test_cart_1d.py | 4 +- feectools/linalg/kron.py | 5 + feectools/linalg/stencil.py | 11 +- feectools/linalg/tests/test_mpi_device.py | 251 ++++++++++++++++++++ pyproject.toml | 2 +- 9 files changed, 291 insertions(+), 12 deletions(-) create mode 100644 feectools/linalg/tests/test_mpi_device.py diff --git a/feectools/ddm/blocking_data_exchanger.py b/feectools/ddm/blocking_data_exchanger.py index b8be40bf3..a6ce09980 100644 --- a/feectools/ddm/blocking_data_exchanger.py +++ b/feectools/ddm/blocking_data_exchanger.py @@ -1,6 +1,7 @@ # coding: utf-8 import cunumpy as xp +from cunumpy import synchronize_for_mpi import numpy as np from feectools.ddm.mpi import mpi as MPI @@ -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,10 @@ def start_exchange_assembly_data( self, array ): 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 diff --git a/feectools/ddm/interface_data_exchanger.py b/feectools/ddm/interface_data_exchanger.py index 9e7a59e14..e25f8d057 100644 --- a/feectools/ddm/interface_data_exchanger.py +++ b/feectools/ddm/interface_data_exchanger.py @@ -1,5 +1,6 @@ # coding: utf-8 +from cunumpy import synchronize_for_mpi from feectools.ddm.mpi import mpi as MPI from .cart import InterfaceCartDecomposition, find_mpi_type @@ -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 diff --git a/feectools/ddm/mpi.py b/feectools/ddm/mpi.py index 242bd4ea3..90e7f6b28 100644 --- a/feectools/ddm/mpi.py +++ b/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)") diff --git a/feectools/ddm/nonblocking_data_exchanger.py b/feectools/ddm/nonblocking_data_exchanger.py index ea6139a50..df7624138 100644 --- a/feectools/ddm/nonblocking_data_exchanger.py +++ b/feectools/ddm/nonblocking_data_exchanger.py @@ -1,6 +1,7 @@ # coding: utf-8 import cunumpy as xp +from cunumpy import synchronize_for_mpi import numpy as np from itertools import product @@ -98,6 +99,9 @@ def prepare_communications(self, u): 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 @@ def start_exchange_assembly_data( self, array ): 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 diff --git a/feectools/ddm/tests/test_cart_1d.py b/feectools/ddm/tests/test_cart_1d.py index c7bea0cee..474426dbb 100644 --- a/feectools/ddm/tests/test_cart_1d.py +++ b/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 ) diff --git a/feectools/linalg/kron.py b/feectools/linalg/kron.py index cf413ad75..672c54963 100644 --- a/feectools/linalg/kron.py +++ b/feectools/linalg/kron.py @@ -3,6 +3,7 @@ import numpy as np import cunumpy as xp +from cunumpy import synchronize_for_mpi from scipy.sparse import kron from scipy.sparse import coo_matrix @@ -885,6 +886,9 @@ def solve_pass(self, workmem, tempmem): targetargs = [tempmem[:self._datasize], self._target_transfer, self._mpi_type] # parts of stripes -> blocked stripes + # (MPI reads/writes the work arrays directly; on a device backend the + # kernels that produced them must have finished first.) + synchronize_for_mpi(workmem, tempmem) self._comm.Alltoallv(sourceargs, targetargs) # blocked stripes -> ordered stripes @@ -897,6 +901,7 @@ def solve_pass(self, workmem, tempmem): self._contiguous_to_blocked(workmem, tempmem) # blocked stripes -> parts of stripes + synchronize_for_mpi(workmem, tempmem) self._comm.Alltoallv(targetargs, sourceargs) #============================================================================== diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index fc53b124e..120593eff 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -10,7 +10,7 @@ from types import MappingProxyType import cunumpy as xp -from cunumpy import PyccelKernel +from cunumpy import PyccelKernel, synchronize_for_mpi from cunumpy.xp import array_backend from scipy.sparse import coo_matrix, diags as sp_diags @@ -221,7 +221,8 @@ def _axpy_python(self, a, x, y): @staticmethod def _inner_python(v1, v2, nghost): index = tuple(slice(ng, -ng) for ng in nghost) - return xp.vdot(v1[index].flat, v2[index].flat) + # ravel, not flat: CuPy's flatiter cannot be used as an array + return xp.vdot(v1[index].ravel(), v2[index].ravel()) #-------------------------------------- # Abstract interface @@ -294,6 +295,8 @@ def inner(self, x, y): if self.parallel: # Sometimes in the parallel case, we can get an empty vector that breaks our kernel x._dot_send_data[0] = 0 if x._data.shape[0] == 0 else inner_func(*inner_args) + # The send buffer was just written on the device; MPI reads it directly. + synchronize_for_mpi(x._dot_send_data, x._dot_recv_data) self.cart.global_comm.Allreduce((x._dot_send_data, self.mpi_type), (x._dot_recv_data, self.mpi_type), op=MPI.SUM ) @@ -2594,7 +2597,7 @@ def _dot(mat, v, out, starts, nrows, nrows_extra, gpads, pads, dm, cm, c_axis, d ii_kk = tuple( ii + kk ) ii[c_axis] += c_start - out[tuple(ii)] = xp.dot( mat[ii_kk].flat, v[jj].flat ) + out[tuple(ii)] = xp.dot( mat[ii_kk].ravel(), v[jj].ravel() ) new_nrows = nrows.copy() @@ -2616,7 +2619,7 @@ def _dot(mat, v, out, starts, nrows, nrows_extra, gpads, pads, dm, cm, c_axis, d kk = [slice(None,n-e) for n,e in zip(ndiags, ee)] ii_kk = tuple( ii + kk ) ii[c_axis] += c_start - out[tuple(ii)] = xp.dot( mat[ii_kk].flat, v[jj].flat ) + out[tuple(ii)] = xp.dot( mat[ii_kk].ravel(), v[jj].ravel() ) new_nrows[d] += er diff --git a/feectools/linalg/tests/test_mpi_device.py b/feectools/linalg/tests/test_mpi_device.py new file mode 100644 index 000000000..1173adda2 --- /dev/null +++ b/feectools/linalg/tests/test_mpi_device.py @@ -0,0 +1,251 @@ +#---------------------------------------------------------------------------# +# 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. # +#---------------------------------------------------------------------------# +""" +Tests that distributed results are *absolutely* correct, not merely +self-consistent. + +Comparing a distributed run against another distributed run of the same code +hides whole classes of bug: if both sides are wrong in the same way they still +agree. In particular, CuPy kernels run asynchronously while MPI knows nothing +about the CuPy stream, so a ghost exchange started before the producing kernels +finish sends stale data -- and every rank agrees on the wrong answer. The tests +below therefore pin distributed results to values computed from the global +field, and check that they do not depend on the decomposition. + +Run with, e.g.:: + + mpirun -np 4 python -m pytest test_mpi_device.py --with-mpi +""" +import numpy as np +import pytest +import cunumpy as xp + +from feectools.ddm.mpi import mpi as MPI +from feectools.ddm.cart import DomainDecomposition, CartDecomposition +from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix + +pytestmark = pytest.mark.mpi + +NPTS = (16, 12) +PADS = (1, 2) + + +# =============================================================================== +def make_space(npts, pads, dtype=float, comm=None): + ndim = len(npts) + D = DomainDecomposition(list(npts), periods=[True] * ndim, comm=comm) + gs, ge = [], [] + for axis in range(ndim): + ee = D.global_element_ends[axis].copy() + ee[-1] = npts[axis] - 1 + ge.append(ee) + gs.append(xp.array([0] + (ee[:-1] + 1).tolist())) + C = CartDecomposition(D, list(npts), gs, ge, pads=list(pads), + shifts=[1] * ndim) + return StencilVectorSpace(C, dtype=dtype) + + +def global_field(npts): + """A deterministic global field, independent of any decomposition.""" + i1 = np.arange(npts[0])[:, None] + i2 = np.arange(npts[1])[None, :] + return ((i1 + 2 * i2) % 17 + 1).astype(float) + + +def scatter(V, glob): + """Put this rank's part of the global field into a new vector.""" + v = StencilVector(V) + owned = tuple(slice(int(s), int(e) + 1) + for s, e in zip(V.starts, V.ends)) + local = tuple(slice(int(p), int(p) + sl.stop - sl.start) + for p, sl in zip(V.pads, owned)) + v._data[local] = xp.asarray(glob[owned]) + v.update_ghost_regions() + return v + + +def laplacian(V, diag=4.5): + A = StencilMatrix(V, V) + A[:, :, 0, 0] = diag + for axis in range(2): + for shift in (-1, 1): + key = [slice(None)] * 2 + [0, 0] + key[2 + axis] = shift + A[tuple(key)] = -1.0 + A.remove_spurious_entries() + return A + + +def reference_apply(glob, diag=4.5): + """The same periodic stencil applied to the global field.""" + out = diag * glob + for axis in range(2): + for shift in (-1, 1): + out = out - np.roll(glob, -shift, axis=axis) + return out + + +# =============================================================================== +def test_ghost_regions_have_the_right_values(): + """Every entry of the local array, ghosts included, must equal the global + field at the corresponding (periodic) global index.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + v = scatter(V, glob) + + data = v._data + data = xp.to_numpy(data) + s0, s1 = int(V.starts[0]), int(V.starts[1]) + p0, p1 = int(V.pads[0]), int(V.pads[1]) + + expected = np.empty_like(data) + for k0 in range(data.shape[0]): + for k1 in range(data.shape[1]): + expected[k0, k1] = glob[(s0 - p0 + k0) % NPTS[0], + (s1 - p1 + k1) % NPTS[1]] + + assert np.allclose(data, expected, rtol=0.0, atol=1e-14) + + +# =============================================================================== +def test_matvec_matches_global_reference(): + """A @ v must equal the stencil applied to the global field, whatever the + decomposition. This is the check that catches an unsynchronized ghost + exchange: a self-consistency check between two distributed runs does not, + because both would be wrong identically.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + v = scatter(V, glob) + A = laplacian(V) + + w = A.dot(v) + ref = reference_apply(glob) + + # Compare through a global reduction, so the check is decomposition-free. + got = float(w.inner(v)) + expected = float((ref * glob).sum()) + assert abs(got - expected) <= 1e-9 * abs(expected) + + # And entry by entry on the rows this rank owns + data = w._data + data = xp.to_numpy(data) + p0, p1 = int(V.pads[0]), int(V.pads[1]) + for i1 in range(int(V.starts[0]), int(V.ends[0]) + 1): + for i2 in range(int(V.starts[1]), int(V.ends[1]) + 1): + k0 = p0 + i1 - int(V.starts[0]) + k1 = p1 + i2 - int(V.starts[1]) + assert abs(data[k0, k1] - ref[i1, i2]) <= 1e-12 + + +# =============================================================================== +def test_matvec_after_device_kernels_without_explicit_sync(): + """The exchange must be safe when the vector was just written by kernels + and the ghost update happens implicitly inside `A.dot` -- the ordering the + PCG loop produces. + + A missing synchronization here is a data race, so the test has to force it + rather than hope for it: a long chain of asynchronous work on the vector is + queued and the exchange is triggered immediately afterwards, leaving the + stream busy while MPI reads the buffer. + """ + comm = MPI.COMM_WORLD + npts, pads = (512, 512), (1, 2) + V = make_space(npts, pads, comm=comm) + glob = global_field(npts) + A = laplacian(V) + + r = scatter(V, glob) + + # Queue work that writes r, mathematically the identity so the expected + # result is unchanged. On a device the arrays are large and the chain long + # enough that kernels are still queued when the exchange starts -- which is + # what makes the race reproducible rather than occasional. There is no race + # on the host, so one pass is enough there. + passes = 200 if xp.is_gpu(r._data) else 1 + for _ in range(passes): + r._data *= 2.0 + r._data *= 0.5 + + r.ghost_regions_in_sync = False + w = A.dot(r) # triggers the implicit ghost update + + ref = reference_apply(glob) + got = float(w.inner(r)) + expected = float((ref * glob).sum()) + assert abs(got - expected) <= 1e-9 * abs(expected) + + +# =============================================================================== +def test_ghost_exchange_synchronizes_before_mpi(monkeypatch): + """ + The exchangers must call `synchronize_for_mpi` before giving a buffer to + MPI. + + This is checked structurally rather than by observing corrupted data, + because the underlying race is not deterministic: whether MPI actually + reads a half-written buffer depends on which internal protocol it picks for + the message, and some of those happen to synchronize with the CuPy stream + by accident. Relying on that accident is exactly the bug, so the contract + is what gets tested. + """ + import feectools.ddm.blocking_data_exchanger as blocking + import feectools.ddm.nonblocking_data_exchanger as nonblocking + + calls = [] + for module in (blocking, nonblocking): + monkeypatch.setattr(module, 'synchronize_for_mpi', + lambda *args: calls.append(args)) + + V = make_space(NPTS, PADS, comm=MPI.COMM_WORLD) + v = StencilVector(V) + v.ghost_regions_in_sync = False + v.update_ghost_regions() + + assert calls, 'ghost exchange handed a buffer to MPI without synchronizing' + assert any(v._data is arg for args in calls for arg in args), \ + 'the synchronized buffer was not the one being exchanged' + + +# =============================================================================== +def test_axpy_then_matvec_is_correct(): + """`mul_iadd` writes on the device; the following exchange must see it.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + A = laplacian(V) + + x = scatter(V, glob) + y = scatter(V, glob) + x.mul_iadd(2.0, y) # x = 3 * glob + w = A.dot(x) + + ref = reference_apply(3.0 * glob) + got = float(w.inner(x)) + expected = float((ref * (3.0 * glob)).sum()) + assert abs(got - expected) <= 1e-9 * abs(expected) + + +# =============================================================================== +def test_inner_matches_global_reference(): + """Reductions must equal the value computed from the global field.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + other = np.flipud(glob).copy() + + x = scatter(V, glob) + y = scatter(V, other) + + assert abs(float(x.inner(y)) - float((glob * other).sum())) <= 1e-9 + assert abs(float(x.inner(x)) - float((glob * glob).sum())) <= 1e-9 + + +# =============================================================================== +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, '-v', '--with-mpi'])) diff --git a/pyproject.toml b/pyproject.toml index ac462afe6..08cfb9ed0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ dependencies = [ 'pyyaml >= 5.1', 'packaging', 'pyevtk', - 'cunumpy', + 'cunumpy>=0.3.0', # Our packages from PyPi 'pyccel >= 2.1.0', From 3a11ad17715c97e1954ab1dfc6863944d9839b19 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 1 Oct 2026 00:06:07 +0200 Subject: [PATCH 3/3] CUDA step 3: bind each MPI rank to its own GPU Replace the initialization that always used GPU 0 by cunumpy.bind_local_device(): each process uses GPU local_rank % device_count, chosen from the node-local rank that the MPI launcher exports, and its CUDA context is created before MPI is initialized (as CUDA-aware MPI requires). No-op on the NumPy backend. Co-Authored-By: Claude Opus 5.5 --- feectools/ddm/cart.py | 15 +++++++-------- feectools/ddm/tests/test_device_binding.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+), 8 deletions(-) create mode 100644 feectools/ddm/tests/test_device_binding.py diff --git a/feectools/ddm/cart.py b/feectools/ddm/cart.py index 48b6942e1..a0e99dd51 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -9,14 +9,13 @@ from cunumpy.xp import array_backend, to_numpy -# 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 +# 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 import bind_local_device + +bind_local_device() from feectools.ddm.mpi import mpi as MPI from feectools.ddm.mpi import MockMPI diff --git a/feectools/ddm/tests/test_device_binding.py b/feectools/ddm/tests/test_device_binding.py new file mode 100644 index 000000000..fc3bca71a --- /dev/null +++ b/feectools/ddm/tests/test_device_binding.py @@ -0,0 +1,15 @@ +"""Each process is bound to its own GPU when feectools.ddm.cart is imported.""" +import pytest +import cunumpy + + +@pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") +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 cunumpy.get_backend() != "cupy": + pytest.skip("device binding only happens on the CuPy backend") + expected = cunumpy.local_rank() % cunumpy.device_count() + assert cp.cuda.runtime.getDevice() == expected