From dd72e6def3dd69aed8b86bfca5b36723df41f62a Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 11 Aug 2026 10:42:37 +0200 Subject: [PATCH 01/23] mr conflicts --- feectools/core/bsplines.py | 22 ++++++++ feectools/ddm/cart.py | 12 +++-- feectools/ddm/partition.py | 3 +- feectools/ddm/petsc.py | 2 +- feectools/feec/global_geometric_projectors.py | 14 +++++ feectools/fem/partitioning.py | 2 +- feectools/fem/tensor.py | 17 ++++++ feectools/linalg/stencil.py | 52 ++++++++++++++----- 8 files changed, 105 insertions(+), 19 deletions(-) diff --git a/feectools/core/bsplines.py b/feectools/core/bsplines.py index 189d95f0e..4a2d6fef8 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', diff --git a/feectools/ddm/cart.py b/feectools/ddm/cart.py index 2b2b58b41..3af494462 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -1,11 +1,11 @@ # coding: utf-8 import os -import numpy as np -import cunumpy as xp -from cunumpy.xp import array_backend +import numpy as xp # this module is host-only MPI/index bookkeeping, never device data from itertools import product +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 +482,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) diff --git a/feectools/ddm/partition.py b/feectools/ddm/partition.py index 8b2b0d3b7..1db4d64d2 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 from sympy.ntheory import factorint 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/feec/global_geometric_projectors.py b/feectools/feec/global_geometric_projectors.py index 3dc7752be..639080d73 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 diff --git a/feectools/fem/partitioning.py b/feectools/fem/partitioning.py index 4df8bae0b..ff6a1577e 100644 --- a/feectools/fem/partitioning.py +++ b/feectools/fem/partitioning.py @@ -2,7 +2,7 @@ import os 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 diff --git a/feectools/fem/tensor.py b/feectools/fem/tensor.py index 513635c52..a942e2b8f 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',) #=============================================================================== diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 4848595c3..c74b4d43c 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 @@ -62,6 +65,24 @@ def _to_numpy_array(val): } #======================================================================== + +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): """ Compute the diagonal length and the padding of the stencil matrix for each direction, @@ -89,13 +110,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 +237,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 @@ -280,7 +303,12 @@ def inner(self, x, y): self.cart.global_comm.Allreduce((x._dot_send_data, self.mpi_type), (x._dot_recv_data, self.mpi_type), op=MPI.SUM ) - return x._dot_recv_data[0] + # _dot_recv_data is a persistent per-vector scratch buffer reused + # across calls; under CuPy, basic indexing (`arr[0]`) returns a + # *view* rather than an independent scalar (unlike NumPy), so a + # caller holding on to this result would see it silently change + # on the vector's next .inner() call. .item() forces a real copy. + return x._dot_recv_data[0].item() else: return inner_func(*inner_args) @@ -1687,7 +1715,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 ) @@ -1854,7 +1882,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 ) @@ -2008,7 +2036,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 +2060,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 +2195,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)] @@ -2927,7 +2955,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 From f850aefd988f188e7599064a2325d8221b4e8f25 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 11 Aug 2026 11:14:24 +0200 Subject: [PATCH 02/23] Moved asarray changes --- feectools/core/bsplines.py | 42 ++++++++++++++++++------------------- feectools/ddm/cart.py | 1 + feectools/fem/splines.py | 2 +- feectools/fem/tensor.py | 4 ++-- feectools/linalg/stencil.py | 4 +++- 5 files changed, 28 insertions(+), 25 deletions(-) diff --git a/feectools/core/bsplines.py b/feectools/core/bsplines.py index 4a2d6fef8..22fcfcbf5 100644 --- a/feectools/core/bsplines.py +++ b/feectools/core/bsplines.py @@ -106,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) #============================================================================== @@ -138,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: @@ -177,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: @@ -215,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: @@ -262,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: @@ -313,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: @@ -368,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: @@ -452,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" @@ -499,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: @@ -594,7 +594,7 @@ def elements_spans(knots, degree, out=None): spans = xp.searchsorted( knots, breaks[:-1], side='right' ) - 1 """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = np.zeros(len(knots), dtype=xp.int64) else: @@ -870,8 +870,8 @@ def basis_ders_on_quad_grid(knots, degree, quad_grid, nders, normalization, offs """ offset = int(offset) ne, nq = quad_grid.shape - knots = xp.ascontiguousarray(knots, dtype=float) - quad_grid = xp.ascontiguousarray(quad_grid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + quad_grid = xp.ascontiguousarray(xp.asarray(quad_grid), dtype=float) if out is None: out = xp.zeros((ne, degree + 1, nders + 1, nq), dtype=float) else: @@ -914,7 +914,7 @@ def basis_integrals(knots, degree, out=None): to (len(knots)-degree-1). In the periodic case the last (degree) values in the array are redundant, as they are a copy of the first (degree) values. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = xp.zeros(len(knots) - degree - 1, dtype=float) else: @@ -956,8 +956,8 @@ def cell_index(breaks, i_grid, tol=1e-15, out=None): ``cell_index[i]`` is the index of the cell in which ``i_grid[i]`` belong. """ - breaks = xp.ascontiguousarray(breaks, dtype=float) - i_grid = xp.ascontiguousarray(i_grid, dtype=float) + breaks = xp.ascontiguousarray(xp.asarray(breaks), dtype=float) + i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float) if out is None: out = np.zeros_like(i_grid, dtype=xp.int64) else: @@ -1012,8 +1012,8 @@ def basis_ders_on_irregular_grid(knots, degree, i_grid, cell_index, nders, norma . il: local basis function (0 <= il <= degree) . id: derivative (0 <= id <= nders ) """ - knots = xp.ascontiguousarray(knots, dtype=float) - i_grid = xp.ascontiguousarray(i_grid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float) if out is None: nx = i_grid.shape[0] out = xp.zeros((nx, degree + 1, nders + 1), dtype=float) diff --git a/feectools/ddm/cart.py b/feectools/ddm/cart.py index 3af494462..f3ccf7f87 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -1,6 +1,7 @@ # coding: utf-8 import os +import numpy as np import numpy as xp # this module is host-only MPI/index bookkeeping, never device data from itertools import product diff --git a/feectools/fem/splines.py b/feectools/fem/splines.py index ac62f9fa7..6f06cf74b 100644 --- a/feectools/fem/splines.py +++ b/feectools/fem/splines.py @@ -198,7 +198,7 @@ def init_interpolation( self, dtype=float ): else: # Convert to LAPACK banded format (see DGBTRF function) - if array_backend.backend == "cupy": + if hasattr(imat, 'get'): imat = imat.get() else: imat = _np.asanyarray(imat) diff --git a/feectools/fem/tensor.py b/feectools/fem/tensor.py index a942e2b8f..f5b9ed2eb 100644 --- a/feectools/fem/tensor.py +++ b/feectools/fem/tensor.py @@ -517,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 @@ -529,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/linalg/stencil.py b/feectools/linalg/stencil.py index c74b4d43c..bdd515b66 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -1792,8 +1792,10 @@ def _tocoo_no_pads(self , order='C'): if array_backend.backend == "cupy": + def _host(a): + return a.get() if hasattr(a, 'get') else 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 ) From 77c92a6e708f9e5c594d891bb12898c1a4fce045 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 11 Aug 2026 11:25:17 +0200 Subject: [PATCH 03/23] more asarray --- feectools/feec/global_geometric_projectors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools/feec/global_geometric_projectors.py b/feectools/feec/global_geometric_projectors.py index 639080d73..159e2d8e2 100644 --- a/feectools/feec/global_geometric_projectors.py +++ b/feectools/feec/global_geometric_projectors.py @@ -183,7 +183,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 From 3606f27e450909efa895bcd921d31fc5a2d3c31d Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 12 Aug 2026 10:10:55 +0200 Subject: [PATCH 04/23] Added inner_many to improve GPU performance --- feectools/linalg/basic.py | 163 +++++++++ feectools/linalg/block.py | 121 +++++- feectools/linalg/solvers.py | 13 +- feectools/linalg/stencil.py | 150 +++++++- feectools/linalg/tests/test_inner_many.py | 426 ++++++++++++++++++++++ 5 files changed, 846 insertions(+), 27 deletions(-) create mode 100644 feectools/linalg/tests/test_inner_many.py diff --git a/feectools/linalg/basic.py b/feectools/linalg/basic.py index 266172d00..c98772bda 100644 --- a/feectools/linalg/basic.py +++ b/feectools/linalg/basic.py @@ -12,12 +12,15 @@ from inspect import signature import cunumpy as xp +import numpy as np from scipy.sparse import coo_matrix +from feectools.ddm.mpi import mpi as MPI from feectools.utilities.utils import is_real __all__ = ( 'VectorSpace', + 'ReductionWorkspace', 'Vector', 'LinearOperator', 'ZeroOperator', @@ -98,6 +101,32 @@ def inner(self, x, y): """ + def inner_many(self, *pairs): + """ + Evaluate several inner products of this space V in one go. + + Semantically identical to ``tuple(self.inner(x, y) for x, y in pairs)``, + but subclasses are free to fuse the work: on a distributed space the + local partial sums of all pairs are reduced with a *single* collective, + and on a device backend all results are brought back to the host with a + *single* transfer. Krylov solvers, which need several global scalar + products per iteration, should prefer this over repeated `inner` calls. + + This base implementation is the unfused fallback. + + Parameters + ---------- + *pairs : tuple[Vector, Vector] + The (x, y) pairs to evaluate. As for `inner`, the first vector of + each pair is the conjugated one in the complex case. + + Returns + ------- + tuple[float | complex, ...] + One scalar per pair, in the order the pairs were given. + """ + return tuple(self.inner(x, y) for x, y in pairs) + @abstractmethod def axpy(self, a, x, y): """ @@ -117,6 +146,119 @@ def axpy(self, a, x, y): The vector modified by this function (incremented by a * x). """ +#=============================================================================== +class ReductionWorkspace: + """ + Mixin giving a vector space reusable scratch for a fused multi-scalar + reduction: a buffer to accumulate the process-local partial sums in, and, + on a device backend, a host mirror to reduce and read them through. + + The scratch lives on the space rather than on the vectors, so a solver + holding a handful of temporaries does not pay for one reduction buffer per + vector. Buffers are grown on demand and then reused, so no allocation + happens in a Krylov iteration once the first one is done. + """ + + __slots__ = () + + def _reduction_send(self, n): + """ + Get a contiguous buffer for `n` locally-computed scalars of the dtype + of this space, living wherever the vector data lives. It aliases + persistent scratch, so callers must consume it before the next call. + """ + buf = getattr(self, '_reduce_send_buf', None) + if buf is None or buf.size < n: + buf = xp.zeros((max(n, 8),), dtype=self.dtype) + self._reduce_send_buf = buf + # A prefix slice stays contiguous, which MPI requires of a raw buffer. + return buf[:n] + + def _host_reduction_buffers(self, n): + """ + Get a pair of host (NumPy) buffers for `n` scalars, in page-locked + memory when available so that the device-to-host copy of the partial + sums is as cheap as it can be. + """ + buffers = getattr(self, '_reduce_host_bufs', None) + if buffers is None or buffers[0].size < n: + buffers = (_pinned_empty(max(n, 8), self.dtype), + _pinned_empty(max(n, 8), self.dtype)) + self._reduce_host_bufs = buffers + return buffers[0][:n], buffers[1][:n] + + def _reduce_to_host(self, send, comm, mpi_type): + """ + Globally sum the local partial sums in `send` and return them on the + host, as a tuple of NumPy scalars. + + When `send` is a device buffer it is first copied to the host and the + reduction is done there. The solver needs these scalars on the host + anyway (to divide by them and to test convergence), so the copy is not + extra work -- it just moves the one unavoidable synchronization ahead + of the collective. In exchange, MPI reduces a handful of bytes of host + memory on its fastest path, no result has to be copied back to the + device, and nothing here depends on the MPI build being CUDA-aware. + + Parameters + ---------- + send : array + Buffer holding this process' partial sums, one per scalar. + + comm : MPI communicator | None + Communicator to reduce over, or None if the space is not + distributed (in which case the partial sums are already the + answer). + + mpi_type : MPI datatype + Datatype matching the dtype of this space. + + Returns + ------- + tuple + One NumPy scalar per entry of `send`. NumPy scalars rather than + Python ones, so the results keep carrying the dtype of the space; + they are independent copies, so they stay valid once the scratch + is overwritten by the next reduction. + """ + n = send.size + + if hasattr(send, 'get'): # device buffer: one D2H copy for the batch + host_send, host_recv = self._host_reduction_buffers(n) + send.get(out=host_send) + if comm is None: + return tuple(host_send) + comm.Allreduce((host_send, mpi_type), (host_recv, mpi_type), + op=MPI.SUM) + return tuple(host_recv) + + if comm is None: + return tuple(send) + + _, host_recv = self._host_reduction_buffers(n) + comm.Allreduce((send, mpi_type), (host_recv, mpi_type), op=MPI.SUM) + return tuple(host_recv) + + +#=============================================================================== +def _pinned_empty(n, dtype): + """ + Allocate an uninitialised 1D host array of `n` entries, in page-locked + (pinned) memory if CuPy is in use, otherwise ordinary host memory. + """ + try: + import cupy as cp + except ImportError: + return np.empty(n, dtype=dtype) + + dtype = np.dtype(dtype) + try: + mem = cp.cuda.alloc_pinned_memory(n * dtype.itemsize) + except Exception: # noqa: BLE001 - no pinned memory is not an error + return np.empty(n, dtype=dtype) + return np.frombuffer(mem, dtype=dtype, count=n) + + #=============================================================================== class Vector(ABC): """ @@ -147,6 +289,27 @@ def inner(self, v): assert self.space is v.space return self.space.inner(self, v) + def inner_many(self, *vectors): + """ + Evaluate the scalar products of self with several vectors of the same + space, fusing them into a single reduction. Shorthand for + ``self.space.inner_many(*((self, v) for v in vectors))``. + + Parameters + ---------- + *vectors : Vector + Vectors belonging to the same space as self. As in `inner`, self is + the conjugated argument in the complex case. + + Returns + ------- + tuple[float | complex, ...] + One scalar per vector, in the order the vectors were given. + """ + assert all(isinstance(v, Vector) and self.space is v.space + for v in vectors) + return self.space.inner_many(*((self, v) for v in vectors)) + def mul_iadd(self, a, v): """ Compute self += a * v, where v is another vector of the same space. diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py index f4085edc2..6858615ad 100644 --- a/feectools/linalg/block.py +++ b/feectools/linalg/block.py @@ -8,14 +8,15 @@ from scipy.sparse import bmat, lil_matrix from feectools.linalg.basic import VectorSpace, Vector, LinearOperator +from feectools.linalg.basic import ReductionWorkspace from feectools.linalg.stencil import StencilMatrix -from feectools.ddm.cart import InterfaceCartDecomposition +from feectools.ddm.cart import InterfaceCartDecomposition, find_mpi_type from feectools.ddm.utilities import get_data_exchanger __all__ = ('BlockVectorSpace', 'BlockVector', 'BlockLinearOperator') #=============================================================================== -class BlockVectorSpace(VectorSpace): +class BlockVectorSpace(ReductionWorkspace, VectorSpace): """ Product Vector Space V of two Vector Spaces (V1,V2) or more. @@ -53,6 +54,9 @@ def __init__(self, *spaces, connectivity=None): else: raise NotImplementedError("The matrices domains don't have the same data type.") + # MPI datatype used by the fused reduction of `inner_many` + self._mpi_dtype = find_mpi_type(self._dtype) + self._connectivity = connectivity or {} self._connectivity_readonly = MappingProxyType(self._connectivity) @@ -122,7 +126,118 @@ def inner(self, x, y): assert isinstance(y, BlockVector) assert x.space is self assert y.space is self - return sum(Vi.inner(xi, yi) for Vi, xi, yi in zip(self.spaces, x.blocks, y.blocks)) + return self.inner_many((x, y))[0] + + #... + def inner_many(self, *pairs): + """ + Evaluate several inner products of this product space in one go, see + :meth:`feectools.linalg.basic.VectorSpace.inner_many`. + + The blocks are summed into the reduction buffer locally, so the whole + batch costs one collective for all pairs *and* all blocks, instead of + one per (pair, block) combination as repeated `inner` calls would. + + Parameters + ---------- + *pairs : tuple[BlockVector, BlockVector] + The (x, y) pairs to evaluate; x is the conjugated one. + + Returns + ------- + tuple[float | complex, ...] + One scalar per pair, in the order the pairs were given. + """ + n = len(pairs) + if n == 0: + return () + + for x, y in pairs: + assert isinstance(x, BlockVector) + assert isinstance(y, BlockVector) + assert x.space is self + assert y.space is self + + comms = self._reduction_comms() + if comms is None or not all(hasattr(Vi, '_inner_local_into') + for Vi in self._spaces): + # A sub-space we do not know how to get local partial sums out of, + # or blocks that do not all reduce over the same communicator: let + # every block reduce for itself. Note this must not go through + # `self.inner`, which delegates back here. + return tuple(self._inner_unfused(x, y) for x, y in pairs) + + if self._reduction_is_trivial(pairs[0][0]): + # Serial and on the host: there is neither a collective nor a + # transfer to amortize, so the plain per-block sum is cheaper than + # routing everything through the reduction scratch. + return tuple(self._inner_unfused(x, y) for x, y in pairs) + + send = self._reduction_send(n) + self._inner_local_into(pairs, send) + return self._reduce_to_host(send, comms[0] if comms else None, + self._mpi_dtype) + + #... + def _inner_unfused(self, x, y): + """Inner product as the sum of the inner products of the blocks, each + reduced on its own.""" + return sum(Vi.inner(xi, yi) + for Vi, xi, yi in zip(self.spaces, x.blocks, y.blocks)) + + #... + def _reduction_is_trivial(self, x): + """ + Whether a reduction over this space has nothing to do beyond the local + sums, i.e. that holds for every block. See + :meth:`feectools.linalg.stencil.StencilVectorSpace._reduction_is_trivial`. + """ + for Vj, xj in zip(self._spaces, x.blocks): + predicate = getattr(Vj, '_reduction_is_trivial', None) + if predicate is None or not predicate(xj): + return False + return True + + #... + def _inner_local_into(self, pairs, out, accumulate=False): + """ + Sum the process-local (unreduced) inner products of each pair over the + blocks of this space, writing one entry per pair into `out`. See + :meth:`feectools.linalg.stencil.StencilVectorSpace._inner_local_into`. + """ + for j, Vj in enumerate(self._spaces): + Vj._inner_local_into( + [(x.blocks[j], y.blocks[j]) for x, y in pairs], + out, + accumulate=accumulate or j > 0, + ) + + #... + def _reduction_comms(self): + """ + The communicators a fused reduction over this space goes through, or + None if the blocks cannot share a single collective. + + All blocks must agree exactly: either all are serial, or all reduce + over the same communicator. Anything else -- blocks on different + communicators, or a mix of serial and distributed blocks, where summing + the local contributions first would reduce the serial ones once per + rank -- disqualifies the fused path. + """ + agreed = None + for Vj in self._spaces: + getter = getattr(Vj, '_reduction_comms', None) + if getter is None: + return None + sub = getter() + if sub is None: + return None + if agreed is None: + agreed = sub + elif len(sub) != len(agreed) or any(a is not b for a, b + in zip(sub, agreed)): + return None + return () if agreed is None else agreed #... def axpy(self, a, x, y): diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index d2e673a5f..4e375650b 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -355,9 +355,11 @@ def solve(self, b, out=None): A.dot(x, out=v) b.copy(out=r) r -= v - nrmr_sqr = r.inner(r).real pc.dot(r, out=s) - am = s.inner(r) + # (r, r) and (s, r) are reduced together: one collective and, on a GPU + # backend, one device-to-host synchronization instead of two. + nrmr_sqr, am = r.space.inner_many((r, r), (s, r)) + nrmr_sqr = nrmr_sqr.real s.copy(out=p) tol_sqr = tol**2 @@ -383,10 +385,13 @@ def solve(self, b, out=None): x.mul_iadd(l, p) # this is x += l*p r.mul_iadd(-l, v) # this is r -= l*v - nrmr_sqr = r.inner(r).real pc.dot(r, out=s) - am1 = s.inner(r) + # As above, the residual norm rides along in the reduction that the + # recurrence needs anyway, so the convergence criterion stays the + # Euclidean one and costs no extra collective. + nrmr_sqr, am1 = r.space.inner_many((r, r), (s, r)) + nrmr_sqr = nrmr_sqr.real # we are computing p = (am1 / am) * p + s by using axpy on s and exchanging the arrays s.mul_iadd((am1/am), p) diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index bdd515b66..1d9ff0971 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -16,6 +16,7 @@ from feectools.ddm.mpi import mpi as MPI from feectools.linalg.basic import VectorSpace, Vector, LinearOperator +from feectools.linalg.basic import ReductionWorkspace from feectools.linalg.memory import stencil_matrix_memory from feectools.ddm.cart import find_mpi_type, CartDecomposition, InterfaceCartDecomposition from feectools.ddm.utilities import get_data_exchanger @@ -53,6 +54,10 @@ def _to_numpy_array(val): return val.get() return val +def _is_device_array(val): + """Whether `val` lives on a device (CuPy) rather than on the host.""" + return hasattr(val, 'get') + #========================================================================# Dictionary used to select correct kernel functions based on dimensionality kernels = { 'axpy' : (None, axpy_1d, axpy_2d, axpy_3d), @@ -121,7 +126,7 @@ def compute_diag_len(pads, shifts_domain, shifts_codomain, return_padding=False) return int(n) #======================================================================== -class StencilVectorSpace(VectorSpace): +class StencilVectorSpace(ReductionWorkspace, VectorSpace): """ Vector space for n-dimensional stencil format. Two different initializations are possible: @@ -212,6 +217,18 @@ def __init__(self, cart, dtype=float): import numpy as np self._inner_consts = tuple(np.int64(p) * np.int64(s) for p, s in zip(self._pads, self._shifts)) + # Index expression selecting the owned (non-ghost) part of the data + # array, matching the loop bounds of the kernels above. Written as + # `slice(ng, n - ng)` rather than `slice(ng, -ng)` because the latter + # selects nothing when a direction has no ghost cells at all. + self._inner_index = tuple(slice(int(ng), int(n) - int(ng)) + for ng, n in zip(self._inner_consts, self._shape)) + + # Number of owned entries: zero means this rank holds no data, in which + # case the compiled kernels must not be called at all. + self._inner_size = math.prod(max(0, s.stop - s.start) + for s in self._inner_index) + # TODO [YG, 06.09.2023]: print warning if pure Python functions are used @@ -289,28 +306,118 @@ def inner(self, x, y): """ + if self._reduction_is_trivial(x): + return self._inner_local(x, y) + + return self.inner_many((x, y))[0] + + # ... + def inner_many(self, *pairs): + """ + Evaluate several inner products of this space in one go, see + :meth:`feectools.linalg.basic.VectorSpace.inner_many`. + + All local partial sums are computed first, then reduced across the + communicator with a single Allreduce, and finally brought to the host + with a single transfer. Compared to calling `inner` once per pair this + saves (n - 1) collectives and, on a device backend, (n - 1) + device-to-host synchronizations. + + Parameters + ---------- + *pairs : tuple[StencilVector, StencilVector] + The (x, y) pairs to evaluate; x is the conjugated one. + + Returns + ------- + tuple[float | complex, ...] + One scalar per pair, in the order the pairs were given. + """ + if len(pairs) == 0: + return () + + if self._reduction_is_trivial(pairs[0][0]): + return tuple(self._inner_local(x, y) for x, y in pairs) + + send = self._reduction_send(len(pairs)) + self._inner_local_into(pairs, send) + comms = self._reduction_comms() + return self._reduce_to_host(send, comms[0] if comms else None, + self.mpi_type) + + # ... + def _reduction_is_trivial(self, x): + """ + Whether a reduction over this space has nothing to do beyond the local + sums: the space is serial (no collective) and its data is on the host + (no transfer). Then the local values are already the answer and the + reduction scratch can be skipped entirely, which is worth doing because + on a small space the bookkeeping is visible next to the kernel itself. + """ + return (not self.parallel + and self._inner_size != 0 + and not _is_device_array(x._data)) + + # ... + def _inner_local(self, x, y): + """ + The process-local (unreduced) inner product of `x` and `y`, left + wherever it was computed: a host scalar on the NumPy backend, a device + scalar on the CuPy one. + """ assert isinstance(x, StencilVector) assert isinstance(y, StencilVector) assert x.space is self assert y.space is self - inner_func = self._inner_func - inner_args = (x._data, y._data, *self._inner_consts) - - 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) - self.cart.global_comm.Allreduce((x._dot_send_data, self.mpi_type), - (x._dot_recv_data, self.mpi_type), - op=MPI.SUM ) - # _dot_recv_data is a persistent per-vector scratch buffer reused - # across calls; under CuPy, basic indexing (`arr[0]`) returns a - # *view* rather than an independent scalar (unlike NumPy), so a - # caller holding on to this result would see it silently change - # on the vector's next .inner() call. .item() forces a real copy. - return x._dot_recv_data[0].item() - else: - return inner_func(*inner_args) + if self._inner_size == 0: + # This rank owns no coefficients; the kernels cannot be called on + # an empty array, and the local contribution is zero. + return 0 + + if _is_device_array(x._data): + # The compiled kernels are host code, so feeding them device arrays + # would copy both operands off the device and run a serial loop on + # the CPU. Reduce on the device instead. + index = self._inner_index + return xp.sum(xp.conj(x._data[index]) * y._data[index], + dtype=self._dtype) + + return self._inner_func(x._data, y._data, *self._inner_consts) + + # ... + def _inner_local_into(self, pairs, out, accumulate=False): + """ + Compute the process-local (unreduced) inner product of each pair and + write it into `out`, one entry per pair. With `accumulate=True` the + values are added to what `out` already holds, which is how a + BlockVectorSpace sums the contributions of its blocks into a single + buffer before reducing once. + + Parameters + ---------- + pairs : sequence[tuple[StencilVector, StencilVector]] + The (x, y) pairs to evaluate; x is the conjugated one. + + out : array + Buffer of at least len(pairs) entries, of the dtype of this space. + + accumulate : bool + Add to `out` instead of overwriting it. + """ + for i, (x, y) in enumerate(pairs): + if accumulate: + out[i] += self._inner_local(x, y) + else: + out[i] = self._inner_local(x, y) + + # ... + def _reduction_comms(self): + """ + The distinct communicators a fused reduction over this space has to go + through: empty in serial, one entry when distributed. + """ + return (self.cart.global_comm,) if self.parallel else () # ... def axpy(self, a, x, y): @@ -505,8 +612,11 @@ def __init__(self, V): self._ndim = len(V.npts) # self._data = xp.zeros(V.shape, dtype=V.dtype) self._data = xp.zeros(tuple(int(s) for s in V.shape), dtype=V.dtype) - self._dot_send_data = xp.zeros((1,), dtype=V.dtype) - self._dot_recv_data = xp.zeros((1,), dtype=V.dtype) + # NOTE: the scratch buffers backing the reduction in `inner`/`inner_many` + # used to live here, one pair per vector. They now belong to the space + # (see ReductionWorkspace), which both avoids duplicating them across + # the temporaries a Krylov solver holds and lets several scalars share + # a single collective. self._interface_data = {} self._requests = None diff --git a/feectools/linalg/tests/test_inner_many.py b/feectools/linalg/tests/test_inner_many.py new file mode 100644 index 000000000..64d7d06d3 --- /dev/null +++ b/feectools/linalg/tests/test_inner_many.py @@ -0,0 +1,426 @@ +#---------------------------------------------------------------------------# +# 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 for the fused multi-scalar reduction `VectorSpace.inner_many`, and for +the equivalence of the `inner_many`-based PCG with the textbook recurrence it +replaces. + +These run on whichever array backend is active (NumPy or CuPy, selected with +the ARRAY_BACKEND environment variable), serially and under MPI. +""" +from math import sqrt + +import pytest +import cunumpy as xp + +from feectools.ddm.mpi import mpi as MPI +from feectools.ddm.cart import DomainDecomposition, CartDecomposition +from feectools.linalg.basic import IdentityOperator, MatrixFreeLinearOperator +from feectools.linalg.block import BlockVectorSpace, BlockVector +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix + +# =============================================================================== +def compute_global_starts_ends(domain_decomposition, npts): + ndims = len(npts) + global_starts = [None] * ndims + global_ends = [None] * ndims + + for axis in range(ndims): + ee = domain_decomposition.global_element_ends[axis] + global_ends[axis] = ee.copy() + global_ends[axis][-1] = npts[axis] - 1 + global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) + + return global_starts, global_ends + + +# =============================================================================== +def make_space(npts, pads, dtype, comm=None): + """Build a StencilVectorSpace over `npts` points, distributed if `comm`.""" + ndim = len(npts) + D = DomainDecomposition(list(npts), periods=[True] * ndim, comm=comm) + global_starts, global_ends = compute_global_starts_ends(D, list(npts)) + C = CartDecomposition(D, list(npts), global_starts, global_ends, + pads=list(pads), shifts=[1] * ndim) + return StencilVectorSpace(C, dtype=dtype) + + +# =============================================================================== +def fill(v, seed): + """Fill the owned coefficients of `v` with reproducible values.""" + V = v.space + ranges = [range(int(s), int(e) + 1) for s, e in zip(V.starts, V.ends)] + + def value(idx): + r = sum((k + 1) * (i + seed) for k, i in enumerate(idx)) % 17 + 1 + return r + 1j * (r % 5 - 2) if V.dtype == complex else float(r) + + if len(ranges) == 1: + for i1 in ranges[0]: + v[i1] = value((i1,)) + elif len(ranges) == 2: + for i1 in ranges[0]: + for i2 in ranges[1]: + v[i1, i2] = value((i1, i2)) + else: + for i1 in ranges[0]: + for i2 in ranges[1]: + for i3 in ranges[2]: + v[i1, i2, i3] = value((i1, i2, i3)) + v.update_ghost_regions() + return v + + +# =============================================================================== +def assert_same(got, expected): + """Compare two scalars that must agree to the last bit: `inner_many` sums + exactly the same terms in the same order as `inner`.""" + assert complex(got) == complex(expected) + + +# =============================================================================== +# SERIAL TESTS +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +@pytest.mark.parametrize('npts, pads', [((15,), (2,)), + ((8, 11), (2, 3)), + ((7, 6, 9), (1, 2, 3))]) +def test_inner_many_serial(dtype, npts, pads): + """inner_many agrees with the same inner products taken one at a time.""" + V = make_space(npts, pads, dtype) + x, y, z = (fill(StencilVector(V), s) for s in (1, 2, 3)) + + expected = (V.inner(x, x), V.inner(x, y), V.inner(z, y), V.inner(y, z)) + got = V.inner_many((x, x), (x, y), (z, y), (y, z)) + + assert len(got) == len(expected) + for g, e in zip(got, expected): + assert_same(g, e) + + # The dtype of the space is carried by the results, as it is for `inner` + for g in got: + assert xp.dtype(type(g)) == xp.dtype(dtype) + + # Degenerate and single-pair cases + assert V.inner_many() == () + assert_same(V.inner_many((x, y))[0], V.inner(x, y)) + + +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +def test_inner_many_is_conjugating(dtype): + """The first argument of each pair is conjugated, as for `inner`: a plain + (non-conjugating) dot product would pass the real case but fail here.""" + V = make_space((6, 5), (1, 2), dtype) + x, y = (fill(StencilVector(V), s) for s in (1, 2)) + + xy, yx = V.inner_many((x, y), (y, x)) + + assert_same(xy, V.inner(x, y)) + assert_same(yx, complex(xy).conjugate()) + # inner(x, x) is real and positive for a non-zero vector + assert complex(V.inner_many((x, x))[0]).imag == 0.0 + assert complex(V.inner_many((x, x))[0]).real > 0.0 + + +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +def test_inner_many_results_are_independent(dtype): + """Results must survive later calls, which recycle the same scratch.""" + V = make_space((6, 7), (2, 1), dtype) + x, y = (fill(StencilVector(V), s) for s in (1, 2)) + + first = V.inner_many((x, x), (x, y)) + kept = tuple(complex(c) for c in first) + + for _ in range(3): + V.inner_many((y, y), (y, x), (x, x)) + + assert tuple(complex(c) for c in first) == kept + + +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +def test_vector_inner_many(dtype): + """The Vector-level shorthand matches the space-level call.""" + V = make_space((5, 6), (1, 1), dtype) + x, y, z = (fill(StencilVector(V), s) for s in (1, 2, 3)) + + got = x.inner_many(x, y, z) + for g, e in zip(got, (x.inner(x), x.inner(y), x.inner(z))): + assert_same(g, e) + + +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +def test_block_inner_many_serial(dtype): + """A BlockVectorSpace sums its blocks into one fused reduction.""" + V1 = make_space((6, 5), (1, 2), dtype) + V2 = make_space((4, 7), (2, 1), dtype) + W = BlockVectorSpace(V1, V2) + + def block(seeds): + return BlockVector(W, [fill(StencilVector(V1), seeds[0]), + fill(StencilVector(V2), seeds[1])]) + + x, y = block((1, 2)), block((3, 4)) + + expected = (W.inner(x, x), W.inner(x, y), W.inner(y, x)) + got = W.inner_many((x, x), (x, y), (y, x)) + for g, e in zip(got, expected): + assert_same(g, e) + + # Nested product spaces reduce through the same single collective + WW = BlockVectorSpace(W, W) + xx = BlockVector(WW, [x, y]) + yy = BlockVector(WW, [y, x]) + assert_same(WW.inner_many((xx, yy))[0], WW.inner(xx, yy)) + + +# =============================================================================== +def _pcg_reference(A, b, pc, x0, tol, maxiter): + """The PCG recurrence as it was written before `inner_many`, used as the + reference the optimized solver must reproduce.""" + x = x0.copy() + v = b.space.zeros() + r = b.space.zeros() + + A.dot(x, out=v) + b.copy(out=r) + r -= v + nrmr_sqr = r.inner(r).real + s = pc.dot(r) + am = s.inner(r) + p = s.copy() + + tol_sqr = tol ** 2 + for k in range(2, maxiter + 1): + if nrmr_sqr < tol_sqr: + k -= 1 + break + v = A.dot(p, out=v) + l = am / v.inner(p) + x.mul_iadd(l, p) + r.mul_iadd(-l, v) + nrmr_sqr = r.inner(r).real + s = pc.dot(r, out=s) + am1 = s.inner(r) + s.mul_iadd((am1 / am), p) + s, p = p, s + am = am1 + + return x, {'niter': k, 'success': nrmr_sqr < tol_sqr, + 'res_norm': sqrt(nrmr_sqr)} + + +# =============================================================================== +def _laplacian(V): + """Symmetric positive-definite stencil matrix on the space V.""" + ndim = len(V.npts) + A = StencilMatrix(V, V) + center = [slice(None)] * ndim + [0] * ndim + A[tuple(center)] = 2.0 * ndim + 0.5 + for axis in range(ndim): + for shift in (-1, 1): + key = [slice(None)] * ndim + [0] * ndim + key[ndim + axis] = shift + A[tuple(key)] = -1.0 + A.remove_spurious_entries() + return A + + +# =============================================================================== +def _shifted_identity(V, shift): + """A Hermitian positive-definite operator built from vector operations + only. Used to exercise the complex case, which the compiled stencil matvec + kernel does not support (it is typed on float64).""" + def dot(v, out=None): + w = v.copy(out=out) + w *= V.dtype(shift) + return w + + return MatrixFreeLinearOperator(domain=V, codomain=V, dot=dot, + dot_transpose=dot) + + +# =============================================================================== +@pytest.mark.parametrize('npts, pads', [((12,), (1,)), ((7, 9), (1, 2))]) +def test_pcg_matches_reference(npts, pads): + """The solver reproduces the reference recurrence: same solution, same + iteration count, same reported residual.""" + V = make_space(npts, pads, float) + A = _laplacian(V) + b = fill(StencilVector(V), 5) + x0 = StencilVector(V) + pc = IdentityOperator(V) + + tol, maxiter = 1e-12, 200 + x_ref, info_ref = _pcg_reference(A, b, pc, x0, tol, maxiter) + + solver = inverse(A, 'pcg', pc=pc, x0=x0, tol=tol, maxiter=maxiter) + x_new = solver.solve(b) + info_new = solver.get_info() + + assert info_new['niter'] == info_ref['niter'] + assert info_new['success'] == info_ref['success'] + assert info_new['success'] + assert abs(info_new['res_norm'] - info_ref['res_norm']) <= 1e-10 * max( + 1.0, info_ref['res_norm']) + + diff = x_new - x_ref + assert sqrt(abs(complex(diff.inner(diff)))) <= 1e-10 * sqrt( + abs(complex(x_ref.inner(x_ref)))) + + # And it really solved the system + res = b - A.dot(x_new) + assert sqrt(abs(complex(res.inner(res)))) <= 1e-6 + + +# =============================================================================== +def test_pcg_complex_matches_reference(): + """PCG on a complex space: the Hermitian inner products keep the recurrence + real where it has to be, and the fused reductions change nothing.""" + V = make_space((6, 7), (1, 2), complex) + A = _shifted_identity(V, 3.0) + b = fill(StencilVector(V), 5) + x0 = StencilVector(V) + pc = IdentityOperator(V) + + tol, maxiter = 1e-13, 100 + x_ref, info_ref = _pcg_reference(A, b, pc, x0, tol, maxiter) + + solver = inverse(A, 'pcg', pc=pc, x0=x0, tol=tol, maxiter=maxiter) + x_new = solver.solve(b) + + assert solver.get_info()['niter'] == info_ref['niter'] + assert solver.get_info()['success'] + + # A = 3*I, so the solution is b/3 -- known in closed form + expected = b.copy() + expected *= complex(1.0 / 3.0) + diff = x_new - expected + assert sqrt(abs(complex(diff.inner(diff)))) <= 1e-10 + + +# =============================================================================== +def test_pcg_recycle_and_out(): + """`recycle` and `out=` keep working with the fused reductions.""" + V = make_space((8, 8), (1, 1), float) + A = _laplacian(V) + b = fill(StencilVector(V), 5) + x0 = StencilVector(V) + + solver = inverse(A, 'pcg', x0=x0, tol=1e-12, maxiter=200, recycle=True) + + out = StencilVector(V) + returned = solver.solve(b, out=out) + assert returned is out + niter_first = solver.get_info()['niter'] + + # With `recycle` the solution was stored as the next initial guess, so + # solving the same system again converges immediately. + solver.solve(b) + assert solver.get_info()['niter'] <= niter_first + assert solver.get_info()['success'] + + +# =============================================================================== +# PARALLEL TESTS +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +@pytest.mark.parametrize('npts, pads', [((16, 12), (2, 3)), + ((10, 9, 8), (1, 2, 1))]) +@pytest.mark.mpi +def test_inner_many_parallel(dtype, npts, pads): + """Under MPI, the fused reduction agrees with the unfused one.""" + comm = MPI.COMM_WORLD + V = make_space(npts, pads, dtype, comm=comm) + x, y, z = (fill(StencilVector(V), s) for s in (1, 2, 3)) + + expected = (V.inner(x, x), V.inner(x, y), V.inner(z, y)) + got = V.inner_many((x, x), (x, y), (z, y)) + for g, e in zip(got, expected): + assert_same(g, e) + + # Every rank must come out of the collective with the same values + per_rank = comm.allgather(tuple(complex(g) for g in got)) + assert all(vals == per_rank[0] for vals in per_rank) + + +# =============================================================================== +@pytest.mark.parametrize('dtype', [float, complex]) +@pytest.mark.mpi +def test_block_inner_many_parallel(dtype): + """Blocks distributed over the same communicator share one collective.""" + comm = MPI.COMM_WORLD + V1 = make_space((12, 10), (1, 2), dtype, comm=comm) + V2 = make_space((8, 14), (2, 1), dtype, comm=comm) + W = BlockVectorSpace(V1, V2) + + x = BlockVector(W, [fill(StencilVector(V1), 1), fill(StencilVector(V2), 2)]) + y = BlockVector(W, [fill(StencilVector(V1), 3), fill(StencilVector(V2), 4)]) + + expected = (W.inner(x, x), W.inner(x, y)) + got = W.inner_many((x, x), (x, y)) + for g, e in zip(got, expected): + assert_same(g, e) + + +# =============================================================================== +@pytest.mark.mpi +def test_block_inner_many_mixed_serial_and_parallel(): + """A product of a distributed block and a replicated (serial) one cannot + share one collective: summing the blocks locally first would count the + serial block once per rank. The result must still be right.""" + comm = MPI.COMM_WORLD + V_par = make_space((12, 10), (1, 2), float, comm=comm) + V_ser = make_space((6, 6), (1, 1), float) + W = BlockVectorSpace(V_par, V_ser) + + assert W._reduction_comms() is None # fused path correctly declined + + x = BlockVector(W, [fill(StencilVector(V_par), 1), + fill(StencilVector(V_ser), 2)]) + y = BlockVector(W, [fill(StencilVector(V_par), 3), + fill(StencilVector(V_ser), 4)]) + + expected = V_par.inner(x.blocks[0], y.blocks[0]) \ + + V_ser.inner(x.blocks[1], y.blocks[1]) + assert_same(W.inner_many((x, y))[0], expected) + assert_same(W.inner(x, y), expected) + + +# =============================================================================== +@pytest.mark.mpi +def test_pcg_matches_reference_parallel(): + """The distributed solver reproduces the reference recurrence too.""" + comm = MPI.COMM_WORLD + V = make_space((16, 12), (1, 2), float, comm=comm) + A = _laplacian(V) + b = fill(StencilVector(V), 5) + x0 = StencilVector(V) + pc = IdentityOperator(V) + + tol, maxiter = 1e-12, 300 + x_ref, info_ref = _pcg_reference(A, b, pc, x0, tol, maxiter) + + solver = inverse(A, 'pcg', pc=pc, x0=x0, tol=tol, maxiter=maxiter) + x_new = solver.solve(b) + info_new = solver.get_info() + + assert info_new['niter'] == info_ref['niter'] + assert info_new['success'] == info_ref['success'] + + diff = x_new - x_ref + assert sqrt(abs(complex(diff.inner(diff)))) <= 1e-10 * sqrt( + abs(complex(x_ref.inner(x_ref)))) + + +# =============================================================================== +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, '-v'])) From 0f16476be2255f290e0c9e275b06388310865e25 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 12 Aug 2026 10:29:46 +0200 Subject: [PATCH 05/23] added device_matvec which generates and caches a cuda kernel per ndim, dtype --- feectools/linalg/kernels/device_matvec.py | 185 +++++++++++++++ feectools/linalg/stencil.py | 69 +++++- feectools/linalg/tests/test_device_matvec.py | 225 +++++++++++++++++++ 3 files changed, 470 insertions(+), 9 deletions(-) create mode 100644 feectools/linalg/kernels/device_matvec.py create mode 100644 feectools/linalg/tests/test_device_matvec.py diff --git a/feectools/linalg/kernels/device_matvec.py b/feectools/linalg/kernels/device_matvec.py new file mode 100644 index 000000000..21bc5d0f6 --- /dev/null +++ b/feectools/linalg/kernels/device_matvec.py @@ -0,0 +1,185 @@ +""" +Device (CUDA) counterpart of the compiled stencil matrix-vector kernels in +``feectools.linalg.stencil_dot_kernels``. + +Those kernels are host code. Handing them CuPy arrays makes +:class:`cunumpy.PyccelKernel` copy the matrix *and* the vector off the device, +run a serial loop on the CPU, and copy the result back -- which costs far more +than the product itself. This module computes the same thing on the device. + +The operation is the stencil product + + out[i] = sum_d mat[i, d] * x[i + d - s_in] + +over the owned rows ``i`` of the codomain, where ``d`` runs over the +``2 * p_in + 1`` diagonals in each direction. The last owned row along a +direction uses ``2 * p_in + add`` diagonals instead, which is how the compiled +kernels handle a rectangular matrix (for a square one ``add == 1`` and the +distinction disappears). + +One thread computes one output point and loops over the diagonals internally, +so a product costs a single kernel launch regardless of the bandwidth. +""" + +import numpy as np + +__all__ = ('device_matvec', 'supports') + +# Dtypes the generated kernels cover, mapped to their CUDA C type. +_CTYPES = { + np.dtype(np.float64): 'double', + np.dtype(np.complex128): 'complex', +} + +# Cache of compiled kernels, keyed by (ndim, dtype). +_KERNELS = {} + +_THREADS_PER_BLOCK = 256 + + +def supports(ndim, dtype): + """Whether a device kernel exists for this dimensionality and dtype.""" + return ndim in (1, 2, 3) and np.dtype(dtype) in _CTYPES + + +def _source(ndim, ctype): + """Generate the CUDA C source of the matvec kernel for `ndim` dimensions. + + Indices are named per direction so the generated code stays close to the + compiled kernels it mirrors: + + * ``i{k}`` -- local row index along direction k, in [0, n{k}) + * ``d{k}`` -- diagonal index along direction k + * ``po{k}`` -- padding of the codomain, the offset of row 0 in `mat`/`out` + * ``of{k}`` -- s_out - s_in, the offset of row 0 in `x` + """ + dims = range(ndim) + + params = ', '.join( + [f'const int n{k}' for k in dims] + + [f'const int nd{k}' for k in dims] # diagonals, interior rows + + [f'const int na{k}' for k in dims] # diagonals, last row + + [f'const long ms{k}' for k in dims] # mat strides, row axes + + [f'const long md{k}' for k in dims] # mat strides, diagonal axes + + [f'const long xs{k}' for k in dims] + + [f'const long os{k}' for k in dims] + + [f'const int po{k}' for k in dims] + + [f'const int of{k}' for k in dims] + ) + + # Unflatten the thread id into one index per direction (last varies fastest) + total = ' * '.join(f'(long)n{k}' for k in dims) + unflatten = [] + for k in reversed(list(dims)): + divisor = ' * '.join(f'(long)n{j}' for j in range(k + 1, ndim)) + if k == 0: + unflatten.append(f' int i0 = (int)(tid / ({divisor}));' + if divisor else ' int i0 = (int)tid;') + elif divisor: + unflatten.append(f' int i{k} = (int)((tid / ({divisor})) % n{k});') + else: + unflatten.append(f' int i{k} = (int)(tid % n{k});') + unflatten = '\n'.join(reversed(unflatten)) + + mat_base = ' + '.join(f'(long)(po{k} + i{k}) * ms{k}' for k in dims) + x_base = ' + '.join(f'(long)(of{k} + i{k}) * xs{k}' for k in dims) + out_index = ' + '.join(f'(long)(po{k} + i{k}) * os{k}' for k in dims) + + # The last row along a direction uses a different number of diagonals. + bounds = '\n'.join( + f' const int b{k} = (i{k} == n{k} - 1) ? na{k} : nd{k};' for k in dims + ) + + loops = '' + for k in dims: + loops += ' ' * (k + 1) + f'for (int d{k} = 0; d{k} < b{k}; ++d{k})\n' + body_indent = ' ' * (ndim + 1) + mat_off = ' + '.join(f'(long)d{k} * md{k}' for k in dims) + x_off = ' + '.join(f'(long)d{k} * xs{k}' for k in dims) + loops += (f'{body_indent}val += mat[mbase + {mat_off}]\n' + f'{body_indent} * x[xbase + {x_off}];\n') + + return f''' +#include + +extern "C" __global__ +void stencil_matvec(const {ctype}* __restrict__ mat, + const {ctype}* __restrict__ x, + {ctype}* __restrict__ out, + {params}) +{{ + long tid = blockIdx.x * (long)blockDim.x + threadIdx.x; + if (tid >= {total}) return; + +{unflatten} + +{bounds} + + const long mbase = {mat_base}; + const long xbase = {x_base}; + + {ctype} val = {ctype}(0); +{loops} + out[{out_index}] = val; +}} +''' + + +def _kernel(ndim, dtype): + """Compile (once) and return the kernel for this dimensionality/dtype.""" + key = (ndim, np.dtype(dtype)) + if key not in _KERNELS: + import cupy as cp + _KERNELS[key] = cp.RawKernel(_source(ndim, _CTYPES[key[1]]), + 'stencil_matvec') + return _KERNELS[key] + + +def _strides(arr, axes): + """Strides of `arr` along `axes`, in elements rather than bytes.""" + return [arr.strides[a] // arr.itemsize for a in axes] + + +def device_matvec(mat, x, out, s_in, p_in, add, s_out, e_out, p_out): + """ + Compute ``out = mat @ x`` on the device, in place. + + Only the owned rows of `out` are written, exactly as the compiled kernels + do; the caller is responsible for the state of the padding. + + Parameters + ---------- + mat : cupy.ndarray + Matrix data, of shape (rows..., diagonals...) -- 2 * ndim axes. + + x, out : cupy.ndarray + Domain and codomain vector data, each of ndim axes. + + s_in, p_in, add, s_out, e_out, p_out : sequence[int] + Per-direction start of the domain, padding of the domain, rectangular + correction, and start/end/padding of the codomain -- the same values + the compiled kernels take. + """ + ndim = x.ndim + + n = [int(e) - int(s) + 1 for s, e in zip(s_out, e_out)] + nd = [2 * int(p) + 1 for p in p_in] + na = [2 * int(p) + int(a) for p, a in zip(p_in, add)] + off = [int(so) - int(si) for so, si in zip(s_out, s_in)] + po = [int(p) for p in p_out] + + args = (mat, x, out, + *n, *nd, *na, + *_strides(mat, range(ndim)), + *_strides(mat, range(ndim, 2 * ndim)), + *_strides(x, range(ndim)), + *_strides(out, range(ndim)), + *po, *off) + + total = 1 + for k in n: + total *= k + blocks = (total + _THREADS_PER_BLOCK - 1) // _THREADS_PER_BLOCK + + _kernel(ndim, x.dtype)((blocks,), (_THREADS_PER_BLOCK,), args) + return out diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 1d9ff0971..8e76376a9 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -22,6 +22,8 @@ from feectools.ddm.utilities import get_data_exchanger from feectools.api.settings import PSYDAC_BACKENDS +from feectools.linalg.kernels.device_matvec import device_matvec +from feectools.linalg.kernels.device_matvec import supports as device_matvec_supports from feectools.linalg.kernels.axpy_kernels import axpy_1d, axpy_2d, axpy_3d from feectools.linalg.kernels.inner_kernels import inner_1d, inner_2d, inner_3d from feectools.linalg.kernels.matvec_kernels import matvec_1d, matvec_2d, matvec_3d @@ -450,22 +452,24 @@ def axpy(self, a, x, y): else: a = float(a) + if _is_device_array(y._data): + # The compiled kernel is host code; on the device this is just a + # scaled add over the whole array (ghost regions included), which + # is what the kernel does too. + y._data += a * x._data + for axis, ext in self.interfaces: + y._interface_data[axis, ext] += a * x._interface_data[axis, ext] + x._sync = x._sync and y._sync + return + x_data_np = _to_numpy_array(x._data) y_data_np = _to_numpy_array(y._data) self._axpy_func(a, x_data_np, y_data_np) - # Copy result back if CuPy - if hasattr(y._data, 'get'): - import cupy as cp - y._data[:] = cp.asarray(y_data_np) for axis, ext in self.interfaces: x_int_np = _to_numpy_array(x._interface_data[axis, ext]) y_int_np = _to_numpy_array(y._interface_data[axis, ext]) self._axpy_func(a, x_int_np, y_int_np) - # Copy result back if CuPy - if hasattr(y._interface_data[axis, ext], 'get'): - import cupy as cp - y._interface_data[axis, ext][:] = cp.asarray(y_int_np) x._sync = x._sync and y._sync @@ -1237,6 +1241,23 @@ def dot(self, v, out=None): if not v.ghost_regions_in_sync: v.update_ghost_regions() + if (self._device_matvec_args() is not None + and _is_device_array(self._data) + and _is_device_array(v._data) + and _is_device_array(out._data)): + # Data is on the device: run the product there. Going through the + # compiled (host) kernel would copy the matrix and the vector off + # the device and reduce serially on the CPU. + # zeros, not empty: the kernel only writes the interior + # (non-padding) region, and the host path leaves the padding zeroed. + out._data[...] = 0 + device_matvec(self._data, v._data, out._data, + **self._device_matvec_args()) + + # IMPORTANT: flag that ghost regions are not up-to-date + out.ghost_regions_in_sync = False + return out + # Convert arrays for compiled kernel - create NumPy output import numpy as _np self_data_np = _to_numpy_array(self._data) @@ -1252,7 +1273,7 @@ def dot(self, v, out=None): args_np[key] = _to_numpy_array(val) 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'): import cupy as cp @@ -1264,6 +1285,36 @@ def dot(self, v, out=None): out.ghost_regions_in_sync = False return out + # ... + def _device_matvec_args(self): + """ + The arguments for :func:`device_matvec`, or None if this matrix cannot + use it. + + The device kernel mirrors the *precompiled* stencil matvec, which is + the one selected by `set_backend(..., precompiled=True)` and is + recognised by the parameters it takes. Any other backend (in + particular the pure-Python `_dot`, which is parametrised differently) + falls back to the host path. + """ + cached = getattr(self, '_device_matvec_args_cache', False) + if cached is not False: + return cached + + keys = ('s_in', 'p_in', 'add', 's_out', 'e_out', 'p_out') + if (not device_matvec_supports(self._ndim, self.dtype) + or set(self._args) != set(keys)): + args = None + else: + # For ndim == 1 these are plain ints, otherwise arrays; the device + # helper wants a sequence per direction either way. + args = {k: (_to_numpy_array(self._args[k]).tolist() + if self._ndim > 1 else [int(self._args[k])]) + for k in keys} + + self._device_matvec_args_cache = args + return args + # ... def vdot( self, v, out=None): """ diff --git a/feectools/linalg/tests/test_device_matvec.py b/feectools/linalg/tests/test_device_matvec.py new file mode 100644 index 000000000..32ca5116a --- /dev/null +++ b/feectools/linalg/tests/test_device_matvec.py @@ -0,0 +1,225 @@ +#---------------------------------------------------------------------------# +# 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 for the device (CUDA) stencil matrix-vector product used by +`StencilMatrix.dot` when the data lives on a GPU. + +The reference is the same stencil sum expressed with shifted array views. It is +backend-independent, so on the NumPy backend these tests check the reference +against the compiled host kernel, and on the CuPy backend they check the device +kernel against the reference -- which pins the device kernel to the compiled one +by transitivity. +""" +import itertools + +import numpy as np +import pytest +import cunumpy as xp +from cunumpy.xp import array_backend + +from feectools.ddm.cart import DomainDecomposition, CartDecomposition +from feectools.linalg.kernels.device_matvec import supports as device_supports +from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix + +ON_CUPY = array_backend.backend == "cupy" + + +# =============================================================================== +def make_space(npts, pads, dtype): + ndim = len(npts) + D = DomainDecomposition(list(npts), periods=[True] * ndim) + global_starts, global_ends = [], [] + for axis in range(ndim): + ee = D.global_element_ends[axis].copy() + ee[-1] = npts[axis] - 1 + global_ends.append(ee) + global_starts.append(xp.array([0] + (ee[:-1] + 1).tolist())) + C = CartDecomposition(D, list(npts), global_starts, global_ends, + pads=list(pads), shifts=[1] * ndim) + return StencilVectorSpace(C, dtype=dtype) + + +# =============================================================================== +def dot_args(A, ndim): + """The per-direction matvec parameters of `A`, as lists of ints.""" + def seq(key): + val = A._args[key] + if ndim == 1: + return [int(val)] + return [int(k) for k in (val.get() if hasattr(val, 'get') else val)] + + return {k: seq(k) for k in + ('s_in', 'p_in', 'add', 's_out', 'e_out', 'p_out')} + + +# =============================================================================== +def reference_matvec(A, v, out): + """ + out = A @ v, as a sum of shifted elementwise products. + + Interior rows use 2 * p_in + 1 diagonals along each direction, the last row + along a direction uses 2 * p_in + add, so every combination of + (interior, last) over the directions is accumulated separately. + """ + ndim = v.space.ndim + a = dot_args(A, ndim) + n = [e - s + 1 for s, e in zip(a['s_out'], a['e_out'])] + off = [so - si for so, si in zip(a['s_out'], a['s_in'])] + p_out, p_in, add = a['p_out'], a['p_in'], a['add'] + + out._data[...] = 0 + + for last in itertools.product([False, True], repeat=ndim): + rows = [] + for k in range(ndim): + if last[k]: + rows.append(slice(p_out[k] + n[k] - 1, p_out[k] + n[k])) + else: + rows.append(slice(p_out[k], p_out[k] + n[k] - 1)) + if any(r.stop <= r.start for r in rows): + continue + + nrow = [r.stop - r.start for r in rows] + base = [r.start - p_out[k] for k, r in enumerate(rows)] + bounds = [2 * p_in[k] + (add[k] if last[k] else 1) for k in range(ndim)] + + for d in itertools.product(*[range(b) for b in bounds]): + src = tuple(slice(off[k] + base[k] + d[k], + off[k] + base[k] + d[k] + nrow[k]) + for k in range(ndim)) + out._data[tuple(rows)] += A._data[tuple(rows) + tuple(d)] * v._data[src] + + return out + + +# =============================================================================== +def fill_like(arr, seed): + rng = np.random.default_rng(seed) + shape = tuple(int(s) for s in arr.shape) + values = rng.random(shape) + if np.dtype(arr.dtype).kind == 'c': + values = values + 1j * rng.random(shape) + return xp.asarray(values.astype(arr.dtype)) + + +# =============================================================================== +def build(npts_domain, npts_codomain, pads, dtype): + V = make_space(npts_domain, pads, dtype) + W = V if npts_domain == npts_codomain else make_space(npts_codomain, pads, dtype) + + A = StencilMatrix(V, W) + A._data[...] = fill_like(A._data, 1) + A.remove_spurious_entries() + + v = StencilVector(V) + v._data[...] = fill_like(v._data, 2) + v.update_ghost_regions() + + return V, W, A, v + + +# =============================================================================== +# Square matrices, and rectangular ones whose spaces differ by one point in a +# direction -- the case that makes `add` zero there, as for derivative operators. +CASES = [ + ('1d-square', (24,), (24,), (2,)), + ('1d-rect', (23,), (24,), (2,)), + ('2d-square', (10, 12), (10, 12), (2, 3)), + ('2d-rect', (11, 10), (12, 10), (2, 2)), + ('2d-rect-both', (11, 9), (12, 10), (1, 2)), + ('3d-square', (7, 8, 9), (7, 8, 9), (1, 2, 3)), + ('3d-rect', (8, 8, 9), (9, 8, 10), (1, 2, 2)), +] + + +@pytest.mark.parametrize('name, npts_d, npts_c, pads', CASES, + ids=[c[0] for c in CASES]) +def test_matvec_matches_reference(name, npts_d, npts_c, pads): + """`StencilMatrix.dot` agrees with the shifted-view reference, whichever + kernel the active backend selects.""" + V, W, A, v = build(npts_d, npts_c, pads, float) + + got = A.dot(v, out=StencilVector(W)) + ref = reference_matvec(A, v, StencilVector(W)) + + assert xp.allclose(got._data, ref._data, rtol=0.0, atol=1e-12) + + +# =============================================================================== +@pytest.mark.parametrize('name, npts_d, npts_c, pads', CASES, + ids=[c[0] for c in CASES]) +def test_matvec_complex(name, npts_d, npts_c, pads): + """Complex matvec. The compiled host kernel is typed on float64 and cannot + do this at all, so it is only checked where the device kernel runs.""" + if not (ON_CUPY and device_supports(len(npts_d), complex)): + pytest.skip('complex matvec needs the device kernel') + + V, W, A, v = build(npts_d, npts_c, pads, complex) + + got = A.dot(v, out=StencilVector(W)) + ref = reference_matvec(A, v, StencilVector(W)) + + assert xp.allclose(got._data, ref._data, rtol=0.0, atol=1e-12) + + +# =============================================================================== +def test_matvec_leaves_padding_zeroed(): + """The kernel writes only the owned rows; the padding of `out` must come + out zeroed, as it does on the host path, even when `out` is reused.""" + V, W, A, v = build((7, 8), (7, 8), (2, 3), float) + + out = StencilVector(W) + out._data[...] = fill_like(out._data, 7) # dirty the buffer, padding too + A.dot(v, out=out) + + a = dot_args(A, 2) + n = [e - s + 1 for s, e in zip(a['s_out'], a['e_out'])] + interior = tuple(slice(a['p_out'][k], a['p_out'][k] + n[k]) for k in range(2)) + + mask = xp.ones(tuple(int(s) for s in out._data.shape), dtype=bool) + mask[interior] = False + assert not bool(xp.any(out._data[mask] != 0)) + + +# =============================================================================== +def test_matvec_out_and_repeated_calls_agree(): + """Reusing an `out` vector gives the same answer as a fresh one.""" + V, W, A, v = build((7, 8, 9), (7, 8, 9), (1, 2, 2), float) + + fresh = A.dot(v) + reused = StencilVector(W) + for _ in range(3): + A.dot(v, out=reused) + + assert xp.allclose(fresh._data, reused._data, rtol=0.0, atol=1e-14) + assert not fresh.ghost_regions_in_sync + + +# =============================================================================== +@pytest.mark.skipif(not ON_CUPY, reason='device kernel requires the CuPy backend') +def test_device_kernel_is_actually_used(): + """Guard against the device path silently falling back to the host one, + which would still be correct but would undo the point of the kernel.""" + V, W, A, v = build((7, 8, 9), (7, 8, 9), (1, 2, 2), float) + assert A._device_matvec_args() is not None + + +# =============================================================================== +def test_unsupported_dtype_falls_back(): + """A dtype without a device kernel must decline the fast path rather than + produce a wrong answer.""" + from feectools.linalg.kernels.device_matvec import supports + + assert supports(3, np.float64) + assert supports(3, np.complex128) + assert not supports(3, np.float32) + assert not supports(4, np.float64) + + +# =============================================================================== +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, '-v'])) From 70e445f15c3fcd5e9d9981c51d9dce2e13394bff Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 12 Aug 2026 11:50:35 +0200 Subject: [PATCH 06/23] gpu improvements --- feectools/ddm/blocking_data_exchanger.py | 7 + feectools/ddm/cart.py | 14 +- feectools/ddm/device.py | 105 ++++++++ feectools/ddm/interface_data_exchanger.py | 5 + feectools/ddm/mpi.py | 36 ++- feectools/ddm/nonblocking_data_exchanger.py | 6 + feectools/linalg/tests/test_mpi_device.py | 253 ++++++++++++++++++++ 7 files changed, 412 insertions(+), 14 deletions(-) create mode 100644 feectools/ddm/device.py 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..225ff1554 100644 --- a/feectools/ddm/blocking_data_exchanger.py +++ b/feectools/ddm/blocking_data_exchanger.py @@ -5,6 +5,7 @@ from feectools.ddm.mpi import mpi as MPI from .cart import CartDecomposition, find_mpi_type +from .device import synchronize_for_mpi from .basic import CartDataExchanger @@ -82,6 +83,10 @@ def start_update_ghost_regions( self, array, requests ): assert isinstance( array, xp.ndarray ) + # MPI reads/writes `array` directly; on a device backend the + # kernels that produced it must have finished first. + synchronize_for_mpi( array ) + # Shortcuts cart = self._cart comm = self._comm @@ -123,6 +128,8 @@ def start_exchange_assembly_data( self, array ): assert isinstance( array, xp.ndarray ) + synchronize_for_mpi( array ) + # Shortcuts cart = self._cart comm = self._comm diff --git a/feectools/ddm/cart.py b/feectools/ddm/cart.py index f3ccf7f87..837220c4e 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -7,14 +7,12 @@ 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 +# Initialize the CUDA context before MPI if using CuPy backend, binding this +# rank to its own GPU. Must stay above the feectools.ddm.mpi import, which +# initialises MPI as a side effect. +from feectools.ddm.device import bind_local_device + +bind_local_device() from feectools.ddm.mpi import mpi as MPI from feectools.ddm.mpi import MockMPI diff --git a/feectools/ddm/device.py b/feectools/ddm/device.py new file mode 100644 index 000000000..51f342649 --- /dev/null +++ b/feectools/ddm/device.py @@ -0,0 +1,105 @@ +""" +Binding of MPI processes to GPUs. + +Kept free of any MPI import on purpose: the CUDA context should exist before +``MPI_Init`` runs, and importing :mod:`feectools.ddm.mpi` initialises MPI as a +side effect. The rank of the process within its node is therefore taken from +the environment variables the launcher sets, which are available before +``MPI_Init``, rather than from a communicator. +""" + +import os + +from cunumpy.xp import array_backend + +__all__ = ('local_rank', 'bind_local_device', 'synchronize_for_mpi') + +# Node-local rank, as exported by the common launchers. +_LOCAL_RANK_VARS = ( + 'OMPI_COMM_WORLD_LOCAL_RANK', # Open MPI + 'MV2_COMM_WORLD_LOCAL_RANK', # MVAPICH2 + 'MPI_LOCALRANKID', # Intel MPI + 'PMI_LOCAL_RANK', # MPICH / PMI + 'SLURM_LOCALID', # Slurm +) + + +def synchronize_for_mpi(*arrays): + """ + Wait for pending device work before MPI touches `arrays`. + + CuPy launches kernels asynchronously on the current stream; MPI knows + nothing about that stream. Handing it a device buffer that a kernel is + still writing lets it send whatever happens to be in memory at that + moment, which shows up as silently wrong ghost regions rather than as an + error. Every MPI call that reads or writes device memory must therefore be + preceded by this. + + The reverse direction needs no barrier: MPI completes its own transfers + before the corresponding wait returns, so kernels launched afterwards see + the received data. + + Parameters + ---------- + *arrays : array | None + The buffers about to be given to MPI. Synchronization happens only if + at least one of them lives on a device, so host-only exchanges (and the + whole NumPy backend) pay nothing. + """ + if not any(hasattr(a, 'get') for a in arrays if a is not None): + return + + import cupy as cp + + cp.cuda.get_current_stream().synchronize() + + +def local_rank(): + """ + The rank of this process within its node, or 0 if no launcher told us + (which is the right answer for a serial run). + """ + for var in _LOCAL_RANK_VARS: + value = os.environ.get(var) + if value is None: + continue + try: + return int(value) + except ValueError: + continue + return 0 + + +def bind_local_device(): + """ + Bind this process to one GPU, chosen round-robin by its node-local rank, and + create the CUDA context. + + Without this every rank on a node would share GPU 0: they would contend for + one device while the others idled, and one device's memory would have to + hold every rank's data. `CUDA_VISIBLE_DEVICES` still applies first, so a + launcher that already hands each rank its own device keeps working (each + process then sees a single device and picks index 0). + + Returns + ------- + int | None + The index of the device that was selected, or None if the CuPy backend + is not active or no device is available. + """ + if array_backend.backend != 'cupy': + return None + + try: + import cupy as cp + + count = cp.cuda.runtime.getDeviceCount() + if count == 0: + return None + + device = local_rank() % count + cp.cuda.Device(device).use() + cp.cuda.Stream.null.synchronize() + return device + except Exception: # noqa: BLE001 - a driver/runtime failure must not be fatal + return None diff --git a/feectools/ddm/interface_data_exchanger.py b/feectools/ddm/interface_data_exchanger.py index 9e7a59e14..3ec78d52a 100644 --- a/feectools/ddm/interface_data_exchanger.py +++ b/feectools/ddm/interface_data_exchanger.py @@ -3,6 +3,7 @@ from feectools.ddm.mpi import mpi as MPI from .cart import InterfaceCartDecomposition, find_mpi_type +from .device import synchronize_for_mpi __all__ = ('InterfaceCartDataExchanger',) @@ -48,6 +49,10 @@ def update_ghost_regions( self, array_minus=None, array_plus=None ): # ... def start_update_ghost_regions( self, array_minus=None, array_plus=None ): + # MPI reads/writes these buffers directly; on a device backend the + # kernels that produced them must have finished first. + synchronize_for_mpi( array_minus, array_plus ) + send_req = [] recv_req = [] cart = self._cart diff --git a/feectools/ddm/mpi.py b/feectools/ddm/mpi.py index 9b6caf23d..fb3a13ee9 100644 --- a/feectools/ddm/mpi.py +++ b/feectools/ddm/mpi.py @@ -80,12 +80,36 @@ def COMM_WORLD(self): # return 1 +import os + +def _enabled(name, default=False): + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() not in ('', '0', 'false', 'no') + + 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 is off by default on the CuPy backend, and on by default otherwise. + # + # It is no longer *incorrect* to combine the two -- the reductions in + # feectools.linalg stage their (tiny) buffers through the host, the ghost + # exchangers synchronize the device before handing it a buffer, and each + # rank binds to its own GPU. It is, however, still slow: a ghost exchange + # of device memory through MPI derived datatypes costs milliseconds, so a + # single-GPU run pays several times over for communication it does not + # need. Until that is addressed, opt in explicitly: + # + # FEECTOOLS_ENABLE_MPI=1 use MPI on the CuPy backend + # FEECTOOLS_DISABLE_MPI=1 force the serial path on any backend + if _enabled('FEECTOOLS_DISABLE_MPI'): + raise ImportError('MPI disabled by FEECTOOLS_DISABLE_MPI') + + if os.environ.get('ARRAY_BACKEND', '').lower() == 'cupy' \ + and not _enabled('FEECTOOLS_ENABLE_MPI'): + raise ImportError('MPI off by default on the CuPy backend; ' + 'set FEECTOOLS_ENABLE_MPI=1 to use it') + from mpi4py import MPI _comm = MPI.COMM_WORLD @@ -93,7 +117,7 @@ def COMM_WORLD(self): # size = _comm.Get_size() mpi_enabled = True except ImportError: - # mpi4py not installed + # mpi4py not installed, or disabled on purpose mpi_enabled = False except Exception: # mpi4py installed but not running under mpirun diff --git a/feectools/ddm/nonblocking_data_exchanger.py b/feectools/ddm/nonblocking_data_exchanger.py index ea6139a50..7facaa93c 100644 --- a/feectools/ddm/nonblocking_data_exchanger.py +++ b/feectools/ddm/nonblocking_data_exchanger.py @@ -6,6 +6,7 @@ from feectools.ddm.mpi import mpi as MPI from .cart import CartDecomposition, find_mpi_type +from .device import synchronize_for_mpi from .basic import CartDataExchanger __all__ = ('NonBlockingCartDataExchanger',) @@ -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,8 @@ def start_exchange_assembly_data( self, array ): assert isinstance( array, xp.ndarray ) + synchronize_for_mpi( array ) + # Shortcuts cart = self._cart comm = self._comm diff --git a/feectools/linalg/tests/test_mpi_device.py b/feectools/linalg/tests/test_mpi_device.py new file mode 100644 index 000000000..69753a8ea --- /dev/null +++ b/feectools/linalg/tests/test_mpi_device.py @@ -0,0 +1,253 @@ +#---------------------------------------------------------------------------# +# 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 = data.get() if hasattr(data, 'get') else 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 = data.get() if hasattr(data, 'get') else 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 hasattr(r._data, 'get') 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 + a, b = V.inner_many((x, x), (x, y)) + assert abs(float(a) - float((glob * glob).sum())) <= 1e-9 + assert abs(float(b) - float((glob * other).sum())) <= 1e-9 + + +# =============================================================================== +if __name__ == "__main__": + import sys + sys.exit(pytest.main([__file__, '-v', '--with-mpi'])) From 920695172104fa2ce225063ce9b9d1d73e14a839 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Sat, 15 Aug 2026 12:10:33 +0200 Subject: [PATCH 07/23] Fix MPI --- feectools/linalg/kron.py | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/feectools/linalg/kron.py b/feectools/linalg/kron.py index 80b013f14..833e3fcbe 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 @@ -440,15 +441,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 +496,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 +800,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 From 01e9f1da28557f7afe1a9485e9c098fc8a33c8c9 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 17 Aug 2026 22:45:16 +0200 Subject: [PATCH 08/23] Use xp.to_numpy instead of .get() manually --- feectools/core/bsplines.py | 18 +++---- feectools/ddm/cart.py | 12 +++-- feectools/ddm/device.py | 3 +- feectools/feec/global_geometric_projectors.py | 49 +++++++++---------- feectools/fem/partitioning.py | 11 +++-- feectools/fem/splines.py | 15 ++---- feectools/linalg/basic.py | 2 +- feectools/linalg/direct_solvers.py | 5 +- feectools/linalg/stencil.py | 18 +++---- feectools/linalg/tests/test_device_matvec.py | 2 +- feectools/linalg/tests/test_mpi_device.py | 6 +-- feectools/utilities/utils.py | 24 ++++----- 12 files changed, 70 insertions(+), 95 deletions(-) diff --git a/feectools/core/bsplines.py b/feectools/core/bsplines.py index 22fcfcbf5..2ec9d7cd5 100644 --- a/feectools/core/bsplines.py +++ b/feectools/core/bsplines.py @@ -540,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 @@ -646,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 ) @@ -660,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 @@ -715,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: @@ -793,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) diff --git a/feectools/ddm/cart.py b/feectools/ddm/cart.py index 837220c4e..af4c4013c 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -1,6 +1,8 @@ # 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 numpy as xp # this module is host-only MPI/index bookkeeping, never device data from itertools import product @@ -499,8 +501,8 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads self._domain_decomposition = domain_decomposition self._npts = tuple( npts ) # Convert to NumPy arrays for MPI compatibility (MPI can't handle CuPy arrays) - self._global_starts = tuple( [ np.asarray(gs.get() if hasattr(gs, 'get') else gs) for gs in global_starts] ) - self._global_ends = tuple( [ np.asarray(ge.get() if hasattr(ge, 'get') else ge) for ge in global_ends] ) + self._global_starts = tuple( [ cunumpy.to_numpy(gs) for gs in global_starts] ) + self._global_ends = tuple( [ cunumpy.to_numpy(ge) for ge in global_ends] ) self._pads = tuple( pads ) self._shifts = tuple( shifts ) self._periods = domain_decomposition.periods @@ -527,7 +529,11 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads # Know my coordinates in the topology self._coords = domain_decomposition.coords # Convert coords to NumPy for indexing (MPI coords should be on CPU) - coords_np = [c.get() if hasattr(c, 'get') else c for c in self._coords] + # cunumpy.to_numpy, not used here: self._coords may hold plain Python ints + # (mpi4py's Get_coords returns a plain list), and to_numpy would wrap those + # into 0-d NumPy arrays via np.asarray -- the wrong type to index a tuple of + # global_starts/ends with below. is_gpu leaves non-CuPy values untouched. + coords_np = [c.get() if cunumpy.is_gpu(c) else c for c in self._coords] # Start/end values of global indices (without ghost regions) self._starts = tuple( self._global_starts[axis][c] for axis,c in zip(range(self._ndims), coords_np) ) diff --git a/feectools/ddm/device.py b/feectools/ddm/device.py index 51f342649..b5cc715bc 100644 --- a/feectools/ddm/device.py +++ b/feectools/ddm/device.py @@ -10,6 +10,7 @@ import os +import cunumpy as xp from cunumpy.xp import array_backend __all__ = ('local_rank', 'bind_local_device', 'synchronize_for_mpi') @@ -46,7 +47,7 @@ def synchronize_for_mpi(*arrays): at least one of them lives on a device, so host-only exchanges (and the whole NumPy backend) pay nothing. """ - if not any(hasattr(a, 'get') for a in arrays if a is not None): + if not any(xp.is_gpu(a) for a in arrays if a is not None): return import cupy as cp diff --git a/feectools/feec/global_geometric_projectors.py b/feectools/feec/global_geometric_projectors.py index 159e2d8e2..8ddf73b14 100644 --- a/feectools/feec/global_geometric_projectors.py +++ b/feectools/feec/global_geometric_projectors.py @@ -43,12 +43,7 @@ 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] @@ -830,7 +825,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 @@ -852,7 +847,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 @@ -878,7 +873,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 @@ -908,9 +903,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 @@ -941,9 +936,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 @@ -966,7 +961,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 @@ -992,9 +987,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 @@ -1022,7 +1017,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 @@ -1057,11 +1052,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 @@ -1098,11 +1093,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 @@ -1126,7 +1121,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 @@ -1155,11 +1150,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 ff6a1577e..51987884b 100644 --- a/feectools/fem/partitioning.py +++ b/feectools/fem/partitioning.py @@ -1,6 +1,8 @@ # -*- 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 numpy as xp # this module is host-only MPI/index bookkeeping, never device data @@ -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 6f06cf74b..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 hasattr(imat, 'get'): - 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/linalg/basic.py b/feectools/linalg/basic.py index c98772bda..b5990d297 100644 --- a/feectools/linalg/basic.py +++ b/feectools/linalg/basic.py @@ -223,7 +223,7 @@ def _reduce_to_host(self, send, comm, mpi_type): """ n = send.size - if hasattr(send, 'get'): # device buffer: one D2H copy for the batch + if xp.is_gpu(send): # device buffer: one D2H copy for the batch host_send, host_recv = self._host_reduction_buffers(n) send.get(out=host_send) if comm is None: diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 1b6096d12..b8bf2fe20 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -70,10 +70,7 @@ def __init__(self, u, l, bmat, transposed=False): 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) + bmat = xp.to_numpy(bmat) self._bmat, self._ipiv, self._finfo = self._factor_function(bmat, l, u) self._sinfo = None diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 8e76376a9..2cb3af997 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -44,21 +44,15 @@ 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 def _is_device_array(val): """Whether `val` lives on a device (CuPy) rather than on the host.""" - return hasattr(val, 'get') + return xp.is_gpu(val) #========================================================================# Dictionary used to select correct kernel functions based on dimensionality kernels = { @@ -1275,7 +1269,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: @@ -1364,7 +1358,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 @@ -1954,7 +1948,7 @@ def _tocoo_no_pads(self , order='C'): if array_backend.backend == "cupy": def _host(a): - return a.get() if hasattr(a, 'get') else a + return xp.to_numpy(a) M = coo_matrix( (_host(data[:ind]), (_host(rows[:ind]), _host(cols[:ind]))), shape=[int(_np.prod(nr)), int(_np.prod(nc))], diff --git a/feectools/linalg/tests/test_device_matvec.py b/feectools/linalg/tests/test_device_matvec.py index 32ca5116a..01a293e54 100644 --- a/feectools/linalg/tests/test_device_matvec.py +++ b/feectools/linalg/tests/test_device_matvec.py @@ -49,7 +49,7 @@ def seq(key): val = A._args[key] if ndim == 1: return [int(val)] - return [int(k) for k in (val.get() if hasattr(val, 'get') else val)] + return [int(k) for k in xp.to_numpy(val)] return {k: seq(k) for k in ('s_in', 'p_in', 'add', 's_out', 'e_out', 'p_out')} diff --git a/feectools/linalg/tests/test_mpi_device.py b/feectools/linalg/tests/test_mpi_device.py index 69753a8ea..aa55737f2 100644 --- a/feectools/linalg/tests/test_mpi_device.py +++ b/feectools/linalg/tests/test_mpi_device.py @@ -98,7 +98,7 @@ def test_ghost_regions_have_the_right_values(): v = scatter(V, glob) data = v._data - data = data.get() if hasattr(data, 'get') else 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]) @@ -133,7 +133,7 @@ def test_matvec_matches_global_reference(): # And entry by entry on the rows this rank owns data = w._data - data = data.get() if hasattr(data, 'get') else 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): @@ -166,7 +166,7 @@ def test_matvec_after_device_kernels_without_explicit_sync(): # 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 hasattr(r._data, 'get') else 1 + passes = 200 if xp.is_gpu(r._data) else 1 for _ in range(passes): r._data *= 2.0 r._data *= 0.5 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 9a8807ee97cfba3ece24ddc5a877a3550a497b02 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 18 Aug 2026 13:56:31 +0200 Subject: [PATCH 09/23] Added DirectSolver(InverseLinearOperator) --- feectools/feec/derivatives.py | 22 ++++- feectools/linalg/direct_solvers.py | 15 ++- feectools/linalg/solvers.py | 147 +++++++++++++++++++++++++++++ feectools/linalg/stencil.py | 13 --- 4 files changed, 174 insertions(+), 23 deletions(-) diff --git a/feectools/feec/derivatives.py b/feectools/feec/derivatives.py index f50ce3c64..6b12a7a43 100644 --- a/feectools/feec/derivatives.py +++ b/feectools/feec/derivatives.py @@ -290,8 +290,13 @@ def tosparse(self, **kwargs): with_pads = kwargs.pop('with_pads', False) - # avoid this case (no pads, but parallel) - assert not (self.domain.parallel and not with_pads) + # avoid this case (no pads, but genuinely decomposed across more than one rank): + # `.parallel` only means "an MPI communicator is attached", true even at 1 rank + # (e.g. under `srun -n 1`), where the no-pads local range already *is* the full + # global range and this restriction does not apply -- so check the rank count + # (`cart.nprocs`) directly rather than `.parallel`. + if self.domain.parallel: + assert with_pads or all(n == 1 for n in self._spaceV.cart.nprocs) # begin with a 1×1 matrix matrix = spa.identity(1, format='coo') @@ -315,13 +320,20 @@ def tosparse(self, **kwargs): directional_matrix = spa.coo_array((codomain_local, domain_local)) else: - maindiag = xp.ones(domain_local) * (-sign) - adddiag = xp.ones(domain_local) * sign + # Plain NumPy, not xp: scipy.sparse.diags is host-only and rejects a + # CuPy array outright (unlike an implicit numpy->cupy 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/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index b8bf2fe20..c6197cbd7 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -231,12 +231,17 @@ def solve(self, rhs, out=None): 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() + # currently no in-place solve exposed. Branch on whether `rhs` itself is a + # device array (not the global `array_backend.backend` flag): the LU + # factorization always lives on the host regardless of backend, and a caller + # may deliberately pass an already-host `rhs`/`out` pair even while the + # active backend is CuPy (see feectools.linalg.solvers.DirectSolver), in + # which case `.get()`-ing a plain NumPy array would fail outright. + if xp.is_gpu(rhs): + rhs_cpu = xp.to_numpy(rhs) result_cpu = self._splu.solve(rhs_cpu.T, trans='T' if transposed else 'N').T out[:] = xp.asarray(result_cpu) + else: + out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T return out diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 4e375650b..e4f491836 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -4,6 +4,7 @@ """ import cunumpy as xp +import numpy as np from math import sqrt, inf from feectools.utilities.utils import is_real @@ -17,6 +18,7 @@ 'inverse', 'ConjugateGradient', 'PConjugateGradient', + 'DirectSolver', 'BiConjugateGradient', 'BiConjugateGradientStabilized', 'PBiConjugateGradientStabilized', @@ -60,6 +62,7 @@ def inverse(A, solver, **kwargs): solvers_dict = { 'cg' : ConjugateGradient, 'pcg' : PConjugateGradient, + 'direct' : DirectSolver, 'bicg' : BiConjugateGradient, 'bicgstab' : BiConjugateGradientStabilized, 'pbicgstab': PBiConjugateGradientStabilized, @@ -416,6 +419,150 @@ def solve(self, b, out=None): def dot(self, b, out=None): return self.solve(b, out=out) +#=============================================================================== +class DirectSolver(InverseLinearOperator): + """ + Exact sparse-direct solve, for linear systems whose left-hand-side operator A does + not actually change across repeated `solve()` calls -- e.g. a time-independent field + operator solved once per time step with only the right-hand side changing (see + `struphy.propagators.implicit_diffusion.ImplicitDiffusion`, whose LHS is constant + whenever `divide_by_dt=False`). A single sparse LU factorization + (`feectools.linalg.direct_solvers.SparseSolver`) then serves every call, instead of + an iterative method repeating (in the worst case, all the way to `maxiter`) every + single call. + + The factorization is built lazily, on the first `solve()` call, and then reused by + every later call without ever re-examining `A` again -- including through a `.linop` + reassignment, e.g. `ImplicitDiffusion.__call__` unconditionally reassigns `.linop` to + a freshly *built* operator every step, regardless of whether its *values* actually + changed. This is a deliberate, cheap-by-construction design, not a value comparison: + `A.tosparse()` is not assumed to be cheap (composed operators can include a + basis-vector sweep, see e.g. `AverageOperator.tosparse`/`BoundaryOperator.tosparse` + in `struphy.feec.mass`/`struphy.feec.linear_operators`), so re-deriving and comparing + it on every call would undo most of the point of factorizing once. The caller is + therefore responsible for knowing that `A`'s *values* are actually constant across + calls (true whenever `ImplicitDiffusion.divide_by_dt=False`, since neither `epsilon` + nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and + the factorization must be rebuilt on the next `solve()`. + + Only supports a serial (non-MPI-parallel) `A`/domain/codomain: a distributed + sparse-direct solve would need its own implementation, which + `feectools.linalg.direct_solvers.SparseSolver` (and therefore this class) does not + attempt. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix A of the linear system. Must support `.tosparse()` and + have a serial (non-parallel) domain/codomain. + + pc, tol, maxiter, verbose : ignored + Accepted only so this class is a drop-in alternative to the iterative solvers + behind the same `solvers.inverse(A, solver, ...)` call site; a direct solve has + no preconditioner, iteration count, or convergence tolerance. + + x0 : feectools.linalg.basic.Vector, optional + Ignored for solving (a direct solve needs no initial guess); if `recycle=True`, + still receives a copy of each solution, for interface consistency with the + iterative solvers (some callers read `x0` back out directly). + + recycle : bool + If True, a copy of the output is stored in x0, as the iterative solvers do. + """ + + def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False, recycle=False): + + self._options = {"x0": x0, "pc": pc, "tol": tol, "maxiter": maxiter, "verbose": verbose, "recycle": recycle} + + super().__init__(A, **self._options) + + # `.parallel` only means "an MPI communicator is attached", true even at 1 rank + # (e.g. under `srun -n 1`); what actually matters for a local sparse-direct + # solve is the rank *count*, so check `cart.nprocs` (per-direction process + # counts) directly rather than `.parallel`. + if self.domain.parallel: + assert all(n == 1 for n in self.domain.cart.nprocs), \ + "DirectSolver only supports a single MPI rank; SparseSolver has no distributed factorization." + + self._sparse_solver = None + self._info = None + + def _check_options(self, **kwargs): + # tol/maxiter/verbose are meaningless for a direct solve (see class docstring); + # only x0, if given, is worth the base class's type/space check. + x0 = kwargs.get("x0") + if x0 is not None: + assert isinstance(x0, Vector), "x0 must be a Vector or None" + assert x0.space == self.codomain, "x0 belongs to the wrong VectorSpace" + + def invalidate(self): + """Force the next `solve()` call to rebuild the factorization from `A`. + + Call this after actually changing `A` (in place, or via the `.linop` setter with + a numerically different operator) -- see the class docstring for why this is not + detected automatically. + """ + self._sparse_solver = None + + def _ensure_factorized(self): + if self._sparse_solver is None: + from feectools.linalg.direct_solvers import SparseSolver + + self._sparse_solver = SparseSolver(self._A.tosparse().tocsc()) + + def solve(self, b, out=None): + """ + Solve A x = b exactly via the cached sparse LU factorization. + + Parameters + ---------- + b : feectools.linalg.stencil.StencilVector + Right-hand-side vector of the linear system. + + out : feectools.linalg.basic.Vector | NoneType + The output vector, or None (optional). + + Returns + ------- + x : feectools.linalg.basic.Vector + The exact (up to factorization round-off) solution of the linear system. + """ + assert isinstance(b, Vector) + assert b.space is self.domain + + self._ensure_factorized() + + # SparseSolver's factorization always lives on the host (scipy splu); the + # host round trip here is one flat vector of the field-solve's DOF count, not + # the particle arrays, so it is cheap relative to the iterations it replaces. + b_flat = xp.to_numpy(b.toarray()) + x_flat = np.empty_like(b_flat) + self._sparse_solver.solve(b_flat, out=x_flat) + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + + # Same local/no-pad interior slice StencilVector.toarray_local() reads from, + # see feectools.linalg.stencil.StencilVector.toarray_local. + idx = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(out.pads, out.space.shifts) + ) + out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) + + self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} + + if self._options.get("recycle") and self._options.get("x0") is not None: + out.copy(out=self._options["x0"]) + + return out + + def dot(self, b, out=None): + return self.solve(b, out=out) + #=============================================================================== class BiConjugateGradient(InverseLinearOperator): """ diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 2cb3af997..b19d08ea3 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -1932,20 +1932,7 @@ 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) From 39202a32777c72cbce4fc4b3a41e253555b6bfcb Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 18 Aug 2026 16:46:04 +0200 Subject: [PATCH 10/23] Added DirectSolver(InverseLinearOperator) --- feectools/feec/derivatives.py | 22 ++++- feectools/linalg/direct_solvers.py | 15 ++- feectools/linalg/solvers.py | 147 +++++++++++++++++++++++++++++ feectools/linalg/stencil.py | 13 --- 4 files changed, 174 insertions(+), 23 deletions(-) diff --git a/feectools/feec/derivatives.py b/feectools/feec/derivatives.py index f50ce3c64..6b12a7a43 100644 --- a/feectools/feec/derivatives.py +++ b/feectools/feec/derivatives.py @@ -290,8 +290,13 @@ def tosparse(self, **kwargs): with_pads = kwargs.pop('with_pads', False) - # avoid this case (no pads, but parallel) - assert not (self.domain.parallel and not with_pads) + # avoid this case (no pads, but genuinely decomposed across more than one rank): + # `.parallel` only means "an MPI communicator is attached", true even at 1 rank + # (e.g. under `srun -n 1`), where the no-pads local range already *is* the full + # global range and this restriction does not apply -- so check the rank count + # (`cart.nprocs`) directly rather than `.parallel`. + if self.domain.parallel: + assert with_pads or all(n == 1 for n in self._spaceV.cart.nprocs) # begin with a 1×1 matrix matrix = spa.identity(1, format='coo') @@ -315,13 +320,20 @@ def tosparse(self, **kwargs): directional_matrix = spa.coo_array((codomain_local, domain_local)) else: - maindiag = xp.ones(domain_local) * (-sign) - adddiag = xp.ones(domain_local) * sign + # Plain NumPy, not xp: scipy.sparse.diags is host-only and rejects a + # CuPy array outright (unlike an implicit numpy->cupy 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/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 1b6096d12..0baa0c8b9 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -234,12 +234,17 @@ def solve(self, rhs, out=None): 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() + # currently no in-place solve exposed. Branch on whether `rhs` itself is a + # device array (not the global `array_backend.backend` flag): the LU + # factorization always lives on the host regardless of backend, and a caller + # may deliberately pass an already-host `rhs`/`out` pair even while the + # active backend is CuPy (see feectools.linalg.solvers.DirectSolver), in + # which case `.get()`-ing a plain NumPy array would fail outright. + if xp.is_gpu(rhs): + rhs_cpu = xp.to_numpy(rhs) result_cpu = self._splu.solve(rhs_cpu.T, trans='T' if transposed else 'N').T out[:] = xp.asarray(result_cpu) + else: + out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T return out diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index d2e673a5f..2d2650085 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -4,6 +4,7 @@ """ import cunumpy as xp +import numpy as np from math import sqrt, inf from feectools.utilities.utils import is_real @@ -17,6 +18,7 @@ 'inverse', 'ConjugateGradient', 'PConjugateGradient', + 'DirectSolver', 'BiConjugateGradient', 'BiConjugateGradientStabilized', 'PBiConjugateGradientStabilized', @@ -60,6 +62,7 @@ def inverse(A, solver, **kwargs): solvers_dict = { 'cg' : ConjugateGradient, 'pcg' : PConjugateGradient, + 'direct' : DirectSolver, 'bicg' : BiConjugateGradient, 'bicgstab' : BiConjugateGradientStabilized, 'pbicgstab': PBiConjugateGradientStabilized, @@ -411,6 +414,150 @@ def solve(self, b, out=None): def dot(self, b, out=None): return self.solve(b, out=out) +#=============================================================================== +class DirectSolver(InverseLinearOperator): + """ + Exact sparse-direct solve, for linear systems whose left-hand-side operator A does + not actually change across repeated `solve()` calls -- e.g. a time-independent field + operator solved once per time step with only the right-hand side changing (see + `struphy.propagators.implicit_diffusion.ImplicitDiffusion`, whose LHS is constant + whenever `divide_by_dt=False`). A single sparse LU factorization + (`feectools.linalg.direct_solvers.SparseSolver`) then serves every call, instead of + an iterative method repeating (in the worst case, all the way to `maxiter`) every + single call. + + The factorization is built lazily, on the first `solve()` call, and then reused by + every later call without ever re-examining `A` again -- including through a `.linop` + reassignment, e.g. `ImplicitDiffusion.__call__` unconditionally reassigns `.linop` to + a freshly *built* operator every step, regardless of whether its *values* actually + changed. This is a deliberate, cheap-by-construction design, not a value comparison: + `A.tosparse()` is not assumed to be cheap (composed operators can include a + basis-vector sweep, see e.g. `AverageOperator.tosparse`/`BoundaryOperator.tosparse` + in `struphy.feec.mass`/`struphy.feec.linear_operators`), so re-deriving and comparing + it on every call would undo most of the point of factorizing once. The caller is + therefore responsible for knowing that `A`'s *values* are actually constant across + calls (true whenever `ImplicitDiffusion.divide_by_dt=False`, since neither `epsilon` + nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and + the factorization must be rebuilt on the next `solve()`. + + Only supports a serial (non-MPI-parallel) `A`/domain/codomain: a distributed + sparse-direct solve would need its own implementation, which + `feectools.linalg.direct_solvers.SparseSolver` (and therefore this class) does not + attempt. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix A of the linear system. Must support `.tosparse()` and + have a serial (non-parallel) domain/codomain. + + pc, tol, maxiter, verbose : ignored + Accepted only so this class is a drop-in alternative to the iterative solvers + behind the same `solvers.inverse(A, solver, ...)` call site; a direct solve has + no preconditioner, iteration count, or convergence tolerance. + + x0 : feectools.linalg.basic.Vector, optional + Ignored for solving (a direct solve needs no initial guess); if `recycle=True`, + still receives a copy of each solution, for interface consistency with the + iterative solvers (some callers read `x0` back out directly). + + recycle : bool + If True, a copy of the output is stored in x0, as the iterative solvers do. + """ + + def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False, recycle=False): + + self._options = {"x0": x0, "pc": pc, "tol": tol, "maxiter": maxiter, "verbose": verbose, "recycle": recycle} + + super().__init__(A, **self._options) + + # `.parallel` only means "an MPI communicator is attached", true even at 1 rank + # (e.g. under `srun -n 1`); what actually matters for a local sparse-direct + # solve is the rank *count*, so check `cart.nprocs` (per-direction process + # counts) directly rather than `.parallel`. + if self.domain.parallel: + assert all(n == 1 for n in self.domain.cart.nprocs), \ + "DirectSolver only supports a single MPI rank; SparseSolver has no distributed factorization." + + self._sparse_solver = None + self._info = None + + def _check_options(self, **kwargs): + # tol/maxiter/verbose are meaningless for a direct solve (see class docstring); + # only x0, if given, is worth the base class's type/space check. + x0 = kwargs.get("x0") + if x0 is not None: + assert isinstance(x0, Vector), "x0 must be a Vector or None" + assert x0.space == self.codomain, "x0 belongs to the wrong VectorSpace" + + def invalidate(self): + """Force the next `solve()` call to rebuild the factorization from `A`. + + Call this after actually changing `A` (in place, or via the `.linop` setter with + a numerically different operator) -- see the class docstring for why this is not + detected automatically. + """ + self._sparse_solver = None + + def _ensure_factorized(self): + if self._sparse_solver is None: + from feectools.linalg.direct_solvers import SparseSolver + + self._sparse_solver = SparseSolver(self._A.tosparse().tocsc()) + + def solve(self, b, out=None): + """ + Solve A x = b exactly via the cached sparse LU factorization. + + Parameters + ---------- + b : feectools.linalg.stencil.StencilVector + Right-hand-side vector of the linear system. + + out : feectools.linalg.basic.Vector | NoneType + The output vector, or None (optional). + + Returns + ------- + x : feectools.linalg.basic.Vector + The exact (up to factorization round-off) solution of the linear system. + """ + assert isinstance(b, Vector) + assert b.space is self.domain + + self._ensure_factorized() + + # SparseSolver's factorization always lives on the host (scipy splu); the + # host round trip here is one flat vector of the field-solve's DOF count, not + # the particle arrays, so it is cheap relative to the iterations it replaces. + b_flat = xp.to_numpy(b.toarray()) + x_flat = np.empty_like(b_flat) + self._sparse_solver.solve(b_flat, out=x_flat) + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + + # Same local/no-pad interior slice StencilVector.toarray_local() reads from, + # see feectools.linalg.stencil.StencilVector.toarray_local. + idx = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(out.pads, out.space.shifts) + ) + out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) + + self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} + + if self._options.get("recycle") and self._options.get("x0") is not None: + out.copy(out=self._options["x0"]) + + return out + + def dot(self, b, out=None): + return self.solve(b, out=out) + #=============================================================================== class BiConjugateGradient(InverseLinearOperator): """ diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 4848595c3..2dd9f9148 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -1749,20 +1749,7 @@ 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": M = coo_matrix( (data[:ind].get(), (rows[:ind].get(), cols[:ind].get())), From e2d6ce6eec784a87eb3d995ada8516cfc0438dfe Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 18 Aug 2026 16:47:25 +0200 Subject: [PATCH 11/23] Updated version number --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0794c0d0f..ac462afe6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "feectools" -version = "0.1.10" +version = "0.1.11" description = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies." readme = "README.md" requires-python = ">= 3.10" From 2f0cd1651c8f8dfbd7fca81aa53b207857eac848 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 12:10:48 +0200 Subject: [PATCH 12/23] Fix for multi-rank direct solver --- feectools/linalg/solvers.py | 108 ++++++++++---- feectools/linalg/tests/test_solvers.py | 62 +++++++- feectools/linalg/utilities.py | 188 ++++++++++++++++++++++++- 3 files changed, 331 insertions(+), 27 deletions(-) diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index e4f491836..e6d175173 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -7,6 +7,8 @@ import numpy as np from math import sqrt, inf +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI from feectools.utilities.utils import is_real from feectools.linalg.utilities import _sym_ortho from feectools.linalg.basic import (Vector, LinearOperator, @@ -445,16 +447,25 @@ class DirectSolver(InverseLinearOperator): nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and the factorization must be rebuilt on the next `solve()`. - Only supports a serial (non-MPI-parallel) `A`/domain/codomain: a distributed - sparse-direct solve would need its own implementation, which - `feectools.linalg.direct_solvers.SparseSolver` (and therefore this class) does not - attempt. + At `nprocs > 1`, this factorizes a *replicated* copy of the full global matrix on + every rank (assembled once via `feectools.linalg.utilities.tosparse_via_matvec`, + which applies `A` to every global unit vector through its own -- already + MPI-correct -- `.dot()`, since `A.tosparse()` itself is only correct in serial for + several composed/derivative operators), rather than attempting an actual + distributed factorization. Every rank redundantly solves the same full system and + keeps only its own slice of the result -- correct and simple, but each rank does + `O(A.domain.dimension)` work per solve instead of `O(A.domain.dimension / nprocs)`, + and the one-time assembly is `O(A.domain.dimension)` *collective* `.dot()` calls that + do not get cheaper with more ranks. This trade only makes sense for problems small + enough that `splu` and this redundant work stay cheap (e.g. the few-thousand-DOF + field solves this class targets); a genuinely distributed sparse-direct solve (e.g. + via PETSc/MUMPS) would need its own implementation. Parameters ---------- A : feectools.linalg.basic.LinearOperator - Left-hand-side matrix A of the linear system. Must support `.tosparse()` and - have a serial (non-parallel) domain/codomain. + Left-hand-side matrix A of the linear system. Must support `.tosparse()` (serial) + or `.dot()` (parallel, via `tosparse_via_matvec`). pc, tol, maxiter, verbose : ignored Accepted only so this class is a drop-in alternative to the iterative solvers @@ -477,12 +488,12 @@ def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False super().__init__(A, **self._options) # `.parallel` only means "an MPI communicator is attached", true even at 1 rank - # (e.g. under `srun -n 1`); what actually matters for a local sparse-direct - # solve is the rank *count*, so check `cart.nprocs` (per-direction process - # counts) directly rather than `.parallel`. - if self.domain.parallel: - assert all(n == 1 for n in self.domain.cart.nprocs), \ - "DirectSolver only supports a single MPI rank; SparseSolver has no distributed factorization." + # (e.g. under `srun -n 1`), where the serial `.tosparse()` path is already + # correct (local range == global range) and faster than the replicated-assembly + # path -- so check the rank *count* (`cart.nprocs`) directly. + cart = self.domain.spaces[0].cart if isinstance(self.domain, BlockVectorSpace) else self.domain.cart + self._parallel = self.domain.parallel and any(n != 1 for n in cart.nprocs) + self._comm = cart.comm if self._parallel else None self._sparse_solver = None self._info = None @@ -508,7 +519,30 @@ def _ensure_factorized(self): if self._sparse_solver is None: from feectools.linalg.direct_solvers import SparseSolver - self._sparse_solver = SparseSolver(self._A.tosparse().tocsc()) + if self._parallel: + from feectools.linalg.utilities import tosparse_via_matvec + + mat = tosparse_via_matvec(self._A, format="csr") + else: + mat = self._A.tosparse().tocsr() + + # `A` can be exactly singular at essential-BC-masked DOFs: an operator + # built through a BoundaryOperator zero-masks both the input and output at + # those rows by design (struphy.feec.linear_operators.BoundaryOperator.dot, + # via apply_essential_bc_to_array) -- fine for an iterative solver, which + # never inverts A directly, as long as `b` is masked the same way (true for + # every caller here: e.g. ImplicitDiffusion.__call__ builds `rhs` via the + # same BoundaryOperator-wrapped `.dot()`, so `b` is already 0 at these rows + # too). A direct factorization needs those rows regularized to identity so + # `x = 1^{-1} * 0 = 0` comes out right there instead of `splu` raising + # "Factor is exactly singular" -- a zero row is unsolvable on its own even + # though the underlying (masked) system is perfectly well posed. + zero_rows = np.flatnonzero(mat.getnnz(axis=1) == 0) + if zero_rows.size: + mat = mat.tolil() + mat[zero_rows, zero_rows] = 1.0 + + self._sparse_solver = SparseSolver(mat.tocsc()) def solve(self, b, out=None): """ @@ -536,22 +570,46 @@ def solve(self, b, out=None): # host round trip here is one flat vector of the field-solve's DOF count, not # the particle arrays, so it is cheap relative to the iterations it replaces. b_flat = xp.to_numpy(b.toarray()) + + if self._parallel: + # `b.toarray()` in parallel already returns the full global-shape array with + # only this rank's own (disjoint) entries filled in -- see + # `StencilVector._toarray_parallel_no_pads` -- so summing every rank's copy + # assembles the true global right-hand side. + if isinstance(self._comm, MockComm): + b_global = b_flat + else: + b_global = np.empty_like(b_flat) + self._comm.Allreduce(b_flat, b_global, op=MPI.SUM) + b_flat = b_global + x_flat = np.empty_like(b_flat) self._sparse_solver.solve(b_flat, out=x_flat) - if out is None: - out = self.codomain.zeros() + if self._parallel: + from feectools.linalg.utilities import array_to_psydac + + x_vec = array_to_psydac(x_flat, self.codomain) + if out is None: + out = x_vec + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + x_vec.copy(out=out) else: - assert isinstance(out, Vector) - assert out.space is self.codomain - - # Same local/no-pad interior slice StencilVector.toarray_local() reads from, - # see feectools.linalg.stencil.StencilVector.toarray_local. - idx = tuple( - slice(m * p, -m * p) if p != 0 else slice(0, None) - for p, m in zip(out.pads, out.space.shifts) - ) - out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + + # Same local/no-pad interior slice StencilVector.toarray_local() reads from, + # see feectools.linalg.stencil.StencilVector.toarray_local. + idx = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(out.pads, out.space.shifts) + ) + out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} diff --git a/feectools/linalg/tests/test_solvers.py b/feectools/linalg/tests/test_solvers.py index e877888f4..bc3baa8cb 100644 --- a/feectools/linalg/tests/test_solvers.py +++ b/feectools/linalg/tests/test_solvers.py @@ -1,10 +1,11 @@ import cunumpy as xp import pytest -from feectools.linalg.solvers import inverse +from feectools.linalg.solvers import inverse, DirectSolver from feectools.linalg.stencil import StencilVectorSpace, StencilMatrix, StencilVector from feectools.linalg.basic import LinearSolver from feectools.ddm.cart import DomainDecomposition, CartDecomposition +from feectools.ddm.mpi import mpi as MPI def define_data_hermitian(n, p, dtype=float): @@ -204,6 +205,65 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False): assert errh_norm < tol assert solver == 'pcg' or errc_norm < tol +#=============================================================================== +def _compute_global_starts_ends(domain_decomposition, npts): + # Same as feectools.linalg.tests.test_block.compute_global_starts_ends. + global_starts = [None] * len(npts) + global_ends = [None] * len(npts) + for axis in range(len(npts)): + ee = domain_decomposition.global_element_ends[axis] + global_ends[axis] = ee.copy() + global_ends[axis][-1] = npts[axis] - 1 + global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) + return global_starts, global_ends + + +@pytest.mark.parametrize('n1', [8, 16]) +@pytest.mark.parametrize('n2', [8, 12]) +@pytest.mark.parametrize('p1', [1, 2]) +@pytest.mark.parallel +def test_direct_solver_parallel(n1, n2, p1, verbose=False): + """`DirectSolver` at nprocs > 1 must recover the exact solution, same as serial.""" + p2 = 1 + + comm = MPI.COMM_WORLD + D = DomainDecomposition([n1, n2], periods=[False, False], comm=comm) + npts = [n1, n2] + global_starts, global_ends = _compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, global_starts, global_ends, pads=[p1, p2], shifts=[1, 1]) + + V = StencilVectorSpace(cart, dtype=float) + A = StencilMatrix(V, V) + + # Diagonally dominant (hence nonsingular) stencil: -1 on every off-diagonal, enough + # on the main diagonal to dominate the row sum -- same style as define_data_hermitian + # above, extended to 2D. + n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 + for k1 in range(-p1, p1 + 1): + for k2 in range(-p2, p2 + 1): + A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 + A[:, :, 0, 0] = n_offdiag + 1.0 + A.remove_spurious_entries() + + s1, s2 = V.starts + e1, e2 = V.ends + xe = StencilVector(V) + for i1 in range(s1, e1 + 1): + for i2 in range(s2, e2 + 1): + xe[i1, i2] = xp.random.random() + xe.update_ghost_regions() + + be = A @ xe + + solv = DirectSolver(A) + x = solv.solve(be) + + err_norm = xp.linalg.norm((x - xe).toarray()) + if verbose: + print(f"n1={n1} n2={n2} p1={p1} p2={p2} nprocs={comm.Get_size()} err_norm={err_norm:.2e}") + assert err_norm < 1e-9 + + # =============================================================================== # SCRIPT FUNCTIONALITY #=============================================================================== diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 42d5636b7..131eb2181 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -1,8 +1,14 @@ # coding: utf-8 -import cunumpy as xp +import itertools from math import sqrt +import cunumpy as xp +import numpy as np +from scipy import sparse + +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI from feectools.linalg.basic import Vector from feectools.linalg.stencil import StencilVector, StencilVectorSpace from feectools.linalg.block import BlockVector, BlockVectorSpace @@ -11,6 +17,7 @@ __all__ = ( 'array_to_psydac', 'petsc_to_psydac', + 'tosparse_via_matvec', '_sym_ortho', ) @@ -73,6 +80,185 @@ def _array_to_psydac_recursive(x, u): else: raise NotImplementedError(f'Can only handle StencilVector or BlockVector spaces, got {type(V)} instead') +#============================================================================== +def tosparse_via_matvec(op, format="csc"): + """ + Assemble the full global sparse matrix of a `LinearOperator` by applying it to every + global unit vector via `.dot()`, rather than via `.tosparse()`. + + Every operator's `.dot()` is already exercised (and therefore correct, including + cross-rank ghost/boundary coupling) every time it is actually used, unlike + `.tosparse()`, which several composed/derivative operators only implement correctly + in serial (see e.g. `feectools.feec.derivatives.DirectionalDerivativeOperator.tosparse`). + This is a port of `struphy.feec.linear_operators.LinOpWithTransp.toarray_struphy`'s + `is_sparse=True` branch into feectools (which `DirectSolver` -- the caller this exists + for -- must not import struphy from): same Allgather-starts/ends plus + unit-vector-`dot()` plus gather/broadcast-triplets algorithm, so every rank ends up + with an identical copy of the full global matrix (a "replicated" assembly, not a + distributed one -- deliberate, see `feectools.linalg.solvers.DirectSolver`). + + Cost: O(N) collective `.dot()` calls, N = `op.domain.dimension` -- does not shrink + with rank count (every call needs every rank's participation), so this is only + appropriate as a one-time, cached setup cost, not something to call every step. + + Parameters + ---------- + op : feectools.linalg.basic.LinearOperator + Operator to assemble. `op.domain`/`op.codomain` must each be a + `StencilVectorSpace` or `BlockVectorSpace`. + + format : str + scipy.sparse matrix format of the result ("csr", "csc", "coo", ...). + + Returns + ------- + out : scipy.sparse matrix + The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on + every rank. + """ + v = op.domain.zeros() + tmp2 = op.codomain.zeros() + + if isinstance(op.domain, BlockVectorSpace): + comm = op.domain.spaces[0].cart.comm + elif isinstance(op.domain, StencilVectorSpace): + comm = op.domain.cart.comm + else: + raise NotImplementedError( + f'tosparse_via_matvec only supports StencilVectorSpace/BlockVectorSpace domains, got {type(op.domain)}', + ) + + if comm is None or isinstance(comm, MockComm): + rank = 0 + size = 1 + else: + rank = comm.Get_rank() + size = comm.Get_size() + + numrows = op.codomain.dimension + numcols = op.domain.dimension + data, row, col = [], [], [] + + if isinstance(op.domain, BlockVectorSpace): + starts = [vi.starts for vi in v] + ends = [vi.ends for vi in v] + npts = [sp.npts for sp in op.domain.spaces] + nsp = len(op.domain.spaces) + ndim = [sp.ndim for sp in op.domain.spaces] + + # Plain NumPy throughout: this is tiny host-side index bookkeeping (rank + # starts/ends, a running column count), never device compute -- `xp.array` + # under the CuPy backend would produce 0-d CuPy scalars that `range()` (and + # plain Python int arithmetic below) cannot consume, the same class of + # NumPy-vs-CuPy scalar-typing trap documented for `AdhocTorus`/`xp.sqrt`. + startsarr = np.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + allstarts = np.empty(size * len(startsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allstarts = startsarr + else: + comm.Allgather(startsarr, allstarts) + allstarts = allstarts.reshape((size, len(startsarr))) + + endsarr = np.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + allends = np.empty(size * len(endsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allends = endsarr + else: + comm.Allgather(endsarr, allends) + allends = allends.reshape((size, len(endsarr))) + + for currentrank in range(size): + spoint = 0 + npredim = 0 + for h in range(nsp): + iterables = [ + range(int(allstarts[currentrank][i + npredim]), int(allends[currentrank][i + npredim]) + 1) + for i in range(ndim[h]) + ] + for i in itertools.product(*iterables): + if rank == currentrank: + v[h][i] = 1.0 + v[h].update_ghost_regions() + tmp2 *= 0.0 + op.dot(v, out=tmp2) + c = spoint + int(np.ravel_multi_index(i, npts[h])) + aux = xp.to_numpy(tmp2.toarray()) + for r in np.nonzero(aux)[0]: + data.append(aux[r]) + col.append(c) + row.append(int(r)) + if rank == currentrank: + v[h][i] = 0.0 + v[h].update_ghost_regions() + cumulative = 1 + for i in range(ndim[h]): + cumulative *= npts[h][i] + spoint += cumulative + npredim += ndim[h] + + else: + starts = v.starts + ends = v.ends + npts = op.domain.npts + ndim = op.domain.ndim + + # Plain NumPy, same reasoning as the BlockVectorSpace branch above. + startsarr = np.array([starts[j] for j in range(ndim)], dtype=int) + allstarts = np.empty(size * len(startsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allstarts = startsarr + else: + comm.Allgather(startsarr, allstarts) + allstarts = allstarts.reshape((size, len(startsarr))) + + endsarr = np.array([ends[j] for j in range(ndim)], dtype=int) + allends = np.empty(size * len(endsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allends = endsarr + else: + comm.Allgather(endsarr, allends) + allends = allends.reshape((size, len(endsarr))) + + for currentrank in range(size): + iterables = [ + range(int(allstarts[currentrank][i]), int(allends[currentrank][i]) + 1) for i in range(ndim) + ] + for i in itertools.product(*iterables): + if rank == currentrank: + v[i] = 1.0 + v.update_ghost_regions() + op.dot(v, out=tmp2) + c = int(np.ravel_multi_index(i, npts)) + aux = xp.to_numpy(tmp2.toarray()) + for r in np.nonzero(aux)[0]: + data.append(aux[r]) + col.append(c) + row.append(int(r)) + if rank == currentrank: + v[i] = 0.0 + v.update_ghost_regions() + + if comm is None or isinstance(comm, MockComm): + all_rows, all_cols, all_data = row, col, data + else: + gathered_rows = comm.gather(row, root=0) + gathered_cols = comm.gather(col, root=0) + gathered_data = comm.gather(data, root=0) + if rank == 0: + all_rows = [item for sublist in gathered_rows for item in sublist] + all_cols = [item for sublist in gathered_cols for item in sublist] + all_data = [item for sublist in gathered_data for item in sublist] + comm.bcast(all_rows, root=0) + comm.bcast(all_cols, root=0) + comm.bcast(all_data, root=0) + else: + all_rows = comm.bcast(None, root=0) + all_cols = comm.bcast(None, root=0) + all_data = comm.bcast(None, root=0) + + mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) + return mat.asformat(format) + #============================================================================== def petsc_to_psydac(x, Xh, out=None): """ From 8a175c5dca280285939320c28819aefa8bd28c71 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 12:26:44 +0200 Subject: [PATCH 13/23] x_vec to numpy --- feectools/linalg/solvers.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index e6d175173..703a99f20 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -589,7 +589,9 @@ def solve(self, b, out=None): if self._parallel: from feectools.linalg.utilities import array_to_psydac - x_vec = array_to_psydac(x_flat, self.codomain) + # x_flat should be a numpy array since SparseSolver's factorization + # is on host + x_vec = array_to_psydac(xp.asarray(x_flat), self.codomain) if out is None: out = x_vec else: @@ -2341,4 +2343,4 @@ def solve(self, b, out=None): return BlockVector(self.domain, blocks=[block_u, p]) def dot(self, b, out=None): - return self.solve(b, out=out) \ No newline at end of file + return self.solve(b, out=out) From da50cf7b98f25d1e2fe7d1d47b912a9e59be10f0 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 13:10:08 +0200 Subject: [PATCH 14/23] Fixes for nprocs > 1 --- feectools/linalg/solvers.py | 110 ++++++++++---- feectools/linalg/tests/test_solvers.py | 61 +++++++- feectools/linalg/utilities.py | 190 ++++++++++++++++++++++++- 3 files changed, 333 insertions(+), 28 deletions(-) diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 2d2650085..883c05e3e 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -7,6 +7,8 @@ import numpy as np from math import sqrt, inf +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI from feectools.utilities.utils import is_real from feectools.linalg.utilities import _sym_ortho from feectools.linalg.basic import (Vector, LinearOperator, @@ -440,16 +442,25 @@ class DirectSolver(InverseLinearOperator): nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and the factorization must be rebuilt on the next `solve()`. - Only supports a serial (non-MPI-parallel) `A`/domain/codomain: a distributed - sparse-direct solve would need its own implementation, which - `feectools.linalg.direct_solvers.SparseSolver` (and therefore this class) does not - attempt. + At `nprocs > 1`, this factorizes a *replicated* copy of the full global matrix on + every rank (assembled once via `feectools.linalg.utilities.tosparse_via_matvec`, + which applies `A` to every global unit vector through its own -- already + MPI-correct -- `.dot()`, since `A.tosparse()` itself is only correct in serial for + several composed/derivative operators), rather than attempting an actual + distributed factorization. Every rank redundantly solves the same full system and + keeps only its own slice of the result -- correct and simple, but each rank does + `O(A.domain.dimension)` work per solve instead of `O(A.domain.dimension / nprocs)`, + and the one-time assembly is `O(A.domain.dimension)` *collective* `.dot()` calls that + do not get cheaper with more ranks. This trade only makes sense for problems small + enough that `splu` and this redundant work stay cheap (e.g. the few-thousand-DOF + field solves this class targets); a genuinely distributed sparse-direct solve (e.g. + via PETSc/MUMPS) would need its own implementation. Parameters ---------- A : feectools.linalg.basic.LinearOperator - Left-hand-side matrix A of the linear system. Must support `.tosparse()` and - have a serial (non-parallel) domain/codomain. + Left-hand-side matrix A of the linear system. Must support `.tosparse()` (serial) + or `.dot()` (parallel, via `tosparse_via_matvec`). pc, tol, maxiter, verbose : ignored Accepted only so this class is a drop-in alternative to the iterative solvers @@ -472,12 +483,12 @@ def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False super().__init__(A, **self._options) # `.parallel` only means "an MPI communicator is attached", true even at 1 rank - # (e.g. under `srun -n 1`); what actually matters for a local sparse-direct - # solve is the rank *count*, so check `cart.nprocs` (per-direction process - # counts) directly rather than `.parallel`. - if self.domain.parallel: - assert all(n == 1 for n in self.domain.cart.nprocs), \ - "DirectSolver only supports a single MPI rank; SparseSolver has no distributed factorization." + # (e.g. under `srun -n 1`), where the serial `.tosparse()` path is already + # correct (local range == global range) and faster than the replicated-assembly + # path -- so check the rank *count* (`cart.nprocs`) directly. + cart = self.domain.spaces[0].cart if isinstance(self.domain, BlockVectorSpace) else self.domain.cart + self._parallel = self.domain.parallel and any(n != 1 for n in cart.nprocs) + self._comm = cart.comm if self._parallel else None self._sparse_solver = None self._info = None @@ -503,7 +514,30 @@ def _ensure_factorized(self): if self._sparse_solver is None: from feectools.linalg.direct_solvers import SparseSolver - self._sparse_solver = SparseSolver(self._A.tosparse().tocsc()) + if self._parallel: + from feectools.linalg.utilities import tosparse_via_matvec + + mat = tosparse_via_matvec(self._A, format="csr") + else: + mat = self._A.tosparse().tocsr() + + # `A` can be exactly singular at essential-BC-masked DOFs: an operator + # built through a BoundaryOperator zero-masks both the input and output at + # those rows by design (struphy.feec.linear_operators.BoundaryOperator.dot, + # via apply_essential_bc_to_array) -- fine for an iterative solver, which + # never inverts A directly, as long as `b` is masked the same way (true for + # every caller here: e.g. ImplicitDiffusion.__call__ builds `rhs` via the + # same BoundaryOperator-wrapped `.dot()`, so `b` is already 0 at these rows + # too). A direct factorization needs those rows regularized to identity so + # `x = 1^{-1} * 0 = 0` comes out right there instead of `splu` raising + # "Factor is exactly singular" -- a zero row is unsolvable on its own even + # though the underlying (masked) system is perfectly well posed. + zero_rows = np.flatnonzero(mat.getnnz(axis=1) == 0) + if zero_rows.size: + mat = mat.tolil() + mat[zero_rows, zero_rows] = 1.0 + + self._sparse_solver = SparseSolver(mat.tocsc()) def solve(self, b, out=None): """ @@ -531,22 +565,48 @@ def solve(self, b, out=None): # host round trip here is one flat vector of the field-solve's DOF count, not # the particle arrays, so it is cheap relative to the iterations it replaces. b_flat = xp.to_numpy(b.toarray()) + + if self._parallel: + # `b.toarray()` in parallel already returns the full global-shape array with + # only this rank's own (disjoint) entries filled in -- see + # `StencilVector._toarray_parallel_no_pads` -- so summing every rank's copy + # assembles the true global right-hand side. + if isinstance(self._comm, MockComm): + b_global = b_flat + else: + b_global = np.empty_like(b_flat) + self._comm.Allreduce(b_flat, b_global, op=MPI.SUM) + b_flat = b_global + x_flat = np.empty_like(b_flat) self._sparse_solver.solve(b_flat, out=x_flat) - if out is None: - out = self.codomain.zeros() + if self._parallel: + from feectools.linalg.utilities import array_to_psydac + + # x_flat should be a numpy array since SparseSolver's factorization + # is on host + x_vec = array_to_psydac(xp.asarray(x_flat), self.codomain) + if out is None: + out = x_vec + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + x_vec.copy(out=out) else: - assert isinstance(out, Vector) - assert out.space is self.codomain - - # Same local/no-pad interior slice StencilVector.toarray_local() reads from, - # see feectools.linalg.stencil.StencilVector.toarray_local. - idx = tuple( - slice(m * p, -m * p) if p != 0 else slice(0, None) - for p, m in zip(out.pads, out.space.shifts) - ) - out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + + # Same local/no-pad interior slice StencilVector.toarray_local() reads from, + # see feectools.linalg.stencil.StencilVector.toarray_local. + idx = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(out.pads, out.space.shifts) + ) + out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} diff --git a/feectools/linalg/tests/test_solvers.py b/feectools/linalg/tests/test_solvers.py index e877888f4..02bc39159 100644 --- a/feectools/linalg/tests/test_solvers.py +++ b/feectools/linalg/tests/test_solvers.py @@ -1,10 +1,11 @@ import cunumpy as xp import pytest -from feectools.linalg.solvers import inverse +from feectools.linalg.solvers import inverse, DirectSolver from feectools.linalg.stencil import StencilVectorSpace, StencilMatrix, StencilVector from feectools.linalg.basic import LinearSolver from feectools.ddm.cart import DomainDecomposition, CartDecomposition +from feectools.ddm.mpi import mpi as MPI def define_data_hermitian(n, p, dtype=float): @@ -204,6 +205,64 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False): assert errh_norm < tol assert solver == 'pcg' or errc_norm < tol +#=============================================================================== +def _compute_global_starts_ends(domain_decomposition, npts): + # Same as feectools.linalg.tests.test_block.compute_global_starts_ends. + global_starts = [None] * len(npts) + global_ends = [None] * len(npts) + for axis in range(len(npts)): + ee = domain_decomposition.global_element_ends[axis] + global_ends[axis] = ee.copy() + global_ends[axis][-1] = npts[axis] - 1 + global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) + return global_starts, global_ends + + +@pytest.mark.parametrize('n1', [8, 16]) +@pytest.mark.parametrize('n2', [8, 12]) +@pytest.mark.parametrize('p1', [1, 2]) +@pytest.mark.parallel +def test_direct_solver_parallel(n1, n2, p1, verbose=False): + """`DirectSolver` at nprocs > 1 must recover the exact solution, same as serial.""" + p2 = 1 + + comm = MPI.COMM_WORLD + D = DomainDecomposition([n1, n2], periods=[False, False], comm=comm) + npts = [n1, n2] + global_starts, global_ends = _compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, global_starts, global_ends, pads=[p1, p2], shifts=[1, 1]) + + V = StencilVectorSpace(cart, dtype=float) + A = StencilMatrix(V, V) + + # Diagonally dominant (hence nonsingular) stencil: -1 on every off-diagonal, enough + # on the main diagonal to dominate the row sum -- same style as define_data_hermitian + # above, extended to 2D. + n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 + for k1 in range(-p1, p1 + 1): + for k2 in range(-p2, p2 + 1): + A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 + A[:, :, 0, 0] = n_offdiag + 1.0 + A.remove_spurious_entries() + + s1, s2 = V.starts + e1, e2 = V.ends + xe = StencilVector(V) + for i1 in range(s1, e1 + 1): + for i2 in range(s2, e2 + 1): + xe[i1, i2] = xp.random.random() + xe.update_ghost_regions() + + be = A @ xe + + solv = DirectSolver(A) + x = solv.solve(be) + + err_norm = xp.linalg.norm((x - xe).toarray()) + if verbose: + print(f"n1={n1} n2={n2} p1={p1} p2={p2} nprocs={comm.Get_size()} err_norm={err_norm:.2e}") + assert err_norm < 1e-9 + # =============================================================================== # SCRIPT FUNCTIONALITY #=============================================================================== diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 42d5636b7..e7d47a233 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -1,8 +1,14 @@ # coding: utf-8 -import cunumpy as xp +import itertools from math import sqrt +import cunumpy as xp +import numpy as np +from scipy import sparse + +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI from feectools.linalg.basic import Vector from feectools.linalg.stencil import StencilVector, StencilVectorSpace from feectools.linalg.block import BlockVector, BlockVectorSpace @@ -11,6 +17,7 @@ __all__ = ( 'array_to_psydac', 'petsc_to_psydac', + 'tosparse_via_matvec', '_sym_ortho', ) @@ -72,7 +79,186 @@ def _array_to_psydac_recursive(x, u): else: raise NotImplementedError(f'Can only handle StencilVector or BlockVector spaces, got {type(V)} instead') - + +#============================================================================== +def tosparse_via_matvec(op, format="csc"): + """ + Assemble the full global sparse matrix of a `LinearOperator` by applying it to every + global unit vector via `.dot()`, rather than via `.tosparse()`. + + Every operator's `.dot()` is already exercised (and therefore correct, including + cross-rank ghost/boundary coupling) every time it is actually used, unlike + `.tosparse()`, which several composed/derivative operators only implement correctly + in serial (see e.g. `feectools.feec.derivatives.DirectionalDerivativeOperator.tosparse`). + This is a port of `struphy.feec.linear_operators.LinOpWithTransp.toarray_struphy`'s + `is_sparse=True` branch into feectools (which `DirectSolver` -- the caller this exists + for -- must not import struphy from): same Allgather-starts/ends plus + unit-vector-`dot()` plus gather/broadcast-triplets algorithm, so every rank ends up + with an identical copy of the full global matrix (a "replicated" assembly, not a + distributed one -- deliberate, see `feectools.linalg.solvers.DirectSolver`). + + Cost: O(N) collective `.dot()` calls, N = `op.domain.dimension` -- does not shrink + with rank count (every call needs every rank's participation), so this is only + appropriate as a one-time, cached setup cost, not something to call every step. + + Parameters + ---------- + op : feectools.linalg.basic.LinearOperator + Operator to assemble. `op.domain`/`op.codomain` must each be a + `StencilVectorSpace` or `BlockVectorSpace`. + + format : str + scipy.sparse matrix format of the result ("csr", "csc", "coo", ...). + + Returns + ------- + out : scipy.sparse matrix + The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on + every rank. + """ + v = op.domain.zeros() + tmp2 = op.codomain.zeros() + + if isinstance(op.domain, BlockVectorSpace): + comm = op.domain.spaces[0].cart.comm + elif isinstance(op.domain, StencilVectorSpace): + comm = op.domain.cart.comm + else: + raise NotImplementedError( + f'tosparse_via_matvec only supports StencilVectorSpace/BlockVectorSpace domains, got {type(op.domain)}', + ) + + if comm is None or isinstance(comm, MockComm): + rank = 0 + size = 1 + else: + rank = comm.Get_rank() + size = comm.Get_size() + + numrows = op.codomain.dimension + numcols = op.domain.dimension + data, row, col = [], [], [] + + if isinstance(op.domain, BlockVectorSpace): + starts = [vi.starts for vi in v] + ends = [vi.ends for vi in v] + npts = [sp.npts for sp in op.domain.spaces] + nsp = len(op.domain.spaces) + ndim = [sp.ndim for sp in op.domain.spaces] + + # Plain NumPy throughout: this is tiny host-side index bookkeeping (rank + # starts/ends, a running column count), never device compute -- `xp.array` + # under the CuPy backend would produce 0-d CuPy scalars that `range()` (and + # plain Python int arithmetic below) cannot consume, the same class of + # NumPy-vs-CuPy scalar-typing trap documented for `AdhocTorus`/`xp.sqrt`. + startsarr = np.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + allstarts = np.empty(size * len(startsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allstarts = startsarr + else: + comm.Allgather(startsarr, allstarts) + allstarts = allstarts.reshape((size, len(startsarr))) + + endsarr = np.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + allends = np.empty(size * len(endsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allends = endsarr + else: + comm.Allgather(endsarr, allends) + allends = allends.reshape((size, len(endsarr))) + + for currentrank in range(size): + spoint = 0 + npredim = 0 + for h in range(nsp): + iterables = [ + range(int(allstarts[currentrank][i + npredim]), int(allends[currentrank][i + npredim]) + 1) + for i in range(ndim[h]) + ] + for i in itertools.product(*iterables): + if rank == currentrank: + v[h][i] = 1.0 + v[h].update_ghost_regions() + tmp2 *= 0.0 + op.dot(v, out=tmp2) + c = spoint + int(np.ravel_multi_index(i, npts[h])) + aux = xp.to_numpy(tmp2.toarray()) + for r in np.nonzero(aux)[0]: + data.append(aux[r]) + col.append(c) + row.append(int(r)) + if rank == currentrank: + v[h][i] = 0.0 + v[h].update_ghost_regions() + cumulative = 1 + for i in range(ndim[h]): + cumulative *= npts[h][i] + spoint += cumulative + npredim += ndim[h] + + else: + starts = v.starts + ends = v.ends + npts = op.domain.npts + ndim = op.domain.ndim + + # Plain NumPy, same reasoning as the BlockVectorSpace branch above. + startsarr = np.array([starts[j] for j in range(ndim)], dtype=int) + allstarts = np.empty(size * len(startsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allstarts = startsarr + else: + comm.Allgather(startsarr, allstarts) + allstarts = allstarts.reshape((size, len(startsarr))) + + endsarr = np.array([ends[j] for j in range(ndim)], dtype=int) + allends = np.empty(size * len(endsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allends = endsarr + else: + comm.Allgather(endsarr, allends) + allends = allends.reshape((size, len(endsarr))) + + for currentrank in range(size): + iterables = [ + range(int(allstarts[currentrank][i]), int(allends[currentrank][i]) + 1) for i in range(ndim) + ] + for i in itertools.product(*iterables): + if rank == currentrank: + v[i] = 1.0 + v.update_ghost_regions() + op.dot(v, out=tmp2) + c = int(np.ravel_multi_index(i, npts)) + aux = xp.to_numpy(tmp2.toarray()) + for r in np.nonzero(aux)[0]: + data.append(aux[r]) + col.append(c) + row.append(int(r)) + if rank == currentrank: + v[i] = 0.0 + v.update_ghost_regions() + + if comm is None or isinstance(comm, MockComm): + all_rows, all_cols, all_data = row, col, data + else: + gathered_rows = comm.gather(row, root=0) + gathered_cols = comm.gather(col, root=0) + gathered_data = comm.gather(data, root=0) + if rank == 0: + all_rows = [item for sublist in gathered_rows for item in sublist] + all_cols = [item for sublist in gathered_cols for item in sublist] + all_data = [item for sublist in gathered_data for item in sublist] + comm.bcast(all_rows, root=0) + comm.bcast(all_cols, root=0) + comm.bcast(all_data, root=0) + else: + all_rows = comm.bcast(None, root=0) + all_cols = comm.bcast(None, root=0) + all_data = comm.bcast(None, root=0) + + mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) + return mat.asformat(format) + #============================================================================== def petsc_to_psydac(x, Xh, out=None): """ From 737d440c6f534e3a5cf789855b052955a9ab11a0 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 13:12:13 +0200 Subject: [PATCH 15/23] batch mpi calls --- feectools/linalg/utilities.py | 78 +++++++++++++++++++++++++++++------ 1 file changed, 66 insertions(+), 12 deletions(-) diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 131eb2181..c69b22c4f 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -139,6 +139,43 @@ def tosparse_via_matvec(op, format="csc"): numcols = op.domain.dimension data, row, col = [], [], [] + def local_nonzero_rows(stencil_vec, row_offset): + """(global_row_indices, values) for `stencil_vec`'s LOCAL interior data only. + + Avoids `stencil_vec.toarray()`: under the parallel branch that allocates a + fresh, full-`codomain.dimension`-sized array and device->host-transfers it in + full, every single call -- the dominant cost under CuPy (a device alloc, a + device-side scatter-write kernel, and a full-size device->host copy per unit + vector, even though a stencil operator's column is actually sparse/local). + Reading only the local interior slice and adding `starts` to get the global row + index (same approach `StencilMatrix._tocoo_no_pads` already uses for columns) + transfers only the local, typically-mostly-zero data instead. + """ + space = stencil_vec.space + idx_local = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(stencil_vec.pads, space.shifts) + ) + local_data = xp.to_numpy(stencil_vec._data[idx_local]) + nz = np.nonzero(local_data) + starts = space.starts + global_multi = tuple(nz[d] + int(starts[d]) for d in range(len(nz))) + rows = row_offset + np.ravel_multi_index(global_multi, space.npts) + return rows, local_data[nz] + + def codomain_local_nonzero_rows(vec): + """`local_nonzero_rows`, dispatched over `op.codomain`'s type.""" + if isinstance(op.codomain, BlockVectorSpace): + all_rows, all_vals = [], [] + row_offset = 0 + for b, sp in enumerate(op.codomain.spaces): + r, val = local_nonzero_rows(vec[b], row_offset) + all_rows.append(r) + all_vals.append(val) + row_offset += sp.dimension + return np.concatenate(all_rows), np.concatenate(all_vals) + return local_nonzero_rows(vec, 0) + if isinstance(op.domain, BlockVectorSpace): starts = [vi.starts for vi in v] ends = [vi.ends for vi in v] @@ -182,14 +219,19 @@ def tosparse_via_matvec(op, format="csc"): tmp2 *= 0.0 op.dot(v, out=tmp2) c = spoint + int(np.ravel_multi_index(i, npts[h])) - aux = xp.to_numpy(tmp2.toarray()) - for r in np.nonzero(aux)[0]: - data.append(aux[r]) - col.append(c) - row.append(int(r)) + rs, vals = codomain_local_nonzero_rows(tmp2) + row.append(rs) + col.append(np.full(rs.shape, c)) + data.append(vals) if rank == currentrank: + # No `update_ghost_regions()` here: resetting this rank's own + # entry back to 0 only needs to be visible to neighbors before + # their *own* next `dot()` call, which is exactly what the + # `update_ghost_regions()` at the top of every iteration (run by + # every rank, every iteration, whether or not it owns that + # iteration's unit vector) already provides -- an extra call + # here would just be the same synchronization done twice. v[h][i] = 0.0 - v[h].update_ghost_regions() cumulative = 1 for i in range(ndim[h]): cumulative *= npts[h][i] @@ -229,14 +271,15 @@ def tosparse_via_matvec(op, format="csc"): v.update_ghost_regions() op.dot(v, out=tmp2) c = int(np.ravel_multi_index(i, npts)) - aux = xp.to_numpy(tmp2.toarray()) - for r in np.nonzero(aux)[0]: - data.append(aux[r]) - col.append(c) - row.append(int(r)) + rs, vals = codomain_local_nonzero_rows(tmp2) + row.append(rs) + col.append(np.full(rs.shape, c)) + data.append(vals) if rank == currentrank: + # See the matching comment in the BlockVectorSpace branch above -- + # no `update_ghost_regions()` needed here, the one at the top of the + # next iteration already covers it. v[i] = 0.0 - v.update_ghost_regions() if comm is None or isinstance(comm, MockComm): all_rows, all_cols, all_data = row, col, data @@ -256,6 +299,17 @@ def tosparse_via_matvec(op, format="csc"): all_cols = comm.bcast(None, root=0) all_data = comm.bcast(None, root=0) + # `row`/`col`/`data` (and therefore `all_rows`/`all_cols`/`all_data`) are lists of + # small per-iteration arrays -- one nonzero-entries batch per unit vector, from + # `codomain_local_nonzero_rows` -- not lists of scalars, so concatenate before + # handing them to `coo_matrix`, which expects flat 1D array-likes. + if all_rows: + all_rows = np.concatenate(all_rows) + all_cols = np.concatenate(all_cols) + all_data = np.concatenate(all_data) + else: + all_rows = all_cols = all_data = np.empty(0, dtype=int) + mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) return mat.asformat(format) From d972105bb68b9511b5dcc00f7c3226497aef542e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 19:06:51 +0200 Subject: [PATCH 16/23] Added some utilities --- feectools/linalg/solvers.py | 16 +- feectools/linalg/tests/test_utilities.py | 428 +++++++++++++++++++++ feectools/linalg/utilities.py | 453 +++++++++++++++++++++++ 3 files changed, 894 insertions(+), 3 deletions(-) create mode 100644 feectools/linalg/tests/test_utilities.py diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 883c05e3e..32a8821ee 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -515,9 +515,19 @@ def _ensure_factorized(self): from feectools.linalg.direct_solvers import SparseSolver if self._parallel: - from feectools.linalg.utilities import tosparse_via_matvec - - mat = tosparse_via_matvec(self._A, format="csr") + from feectools.linalg.utilities import FastAssemblyUnavailable, parallel_tosparse, tosparse_via_matvec + + # `parallel_tosparse` assembles in O(1) collective rounds (one per leaf + # operator) instead of `tosparse_via_matvec`'s O(A.domain.dimension) + # rounds (one per global DOF) -- a difference of several orders of + # magnitude for a field-solve-sized system (see its docstring for how). + # It only recognizes a subset of operator shapes, self-verified against + # `A`'s own `.dot()`; fall back to the always-correct (if much slower) + # sweep when it can't. + try: + mat = parallel_tosparse(self._A, self._comm, format="csr") + except FastAssemblyUnavailable: + mat = tosparse_via_matvec(self._A, format="csr") else: mat = self._A.tosparse().tocsr() diff --git a/feectools/linalg/tests/test_utilities.py b/feectools/linalg/tests/test_utilities.py new file mode 100644 index 000000000..94e511c5b --- /dev/null +++ b/feectools/linalg/tests/test_utilities.py @@ -0,0 +1,428 @@ +import cunumpy as xp +import numpy as np +import pytest +from scipy import sparse + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ComposedLinearOperator, IdentityOperator, LinearOperator, ScaledLinearOperator, SumLinearOperator +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace +from feectools.linalg.stencil import StencilMatrix, StencilVectorSpace +from feectools.linalg.utilities import ( + FastAssemblyUnavailable, + _get_entry, + _local_flat_entries, + _set_entry, + parallel_tosparse, + tosparse_via_matvec, +) + + +def compute_global_starts_ends(domain_decomposition, npts): + global_starts = [None] * len(npts) + global_ends = [None] * len(npts) + for axis in range(len(npts)): + ee = domain_decomposition.global_element_ends[axis] + global_ends[axis] = ee.copy() + global_ends[axis][-1] = npts[axis] - 1 + global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) + return global_starts, global_ends + + +def make_space(n1, n2, p1, p2, comm, periodic=True): + D = DomainDecomposition([n1, n2], periods=[periodic, False], comm=comm) + npts = [n1, n2] + gs, ge = compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, gs, ge, pads=[p1, p2], shifts=[1, 1]) + return StencilVectorSpace(cart, dtype=float) + + +def make_stencil_matrix(V, p1, p2, scale=1.0): + A = StencilMatrix(V, V) + n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 + for k1 in range(-p1, p1 + 1): + for k2 in range(-p2, p2 + 1): + A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 * scale + A[:, :, 0, 0] = (n_offdiag + 1.0) * scale + A.remove_spurious_entries() + return A + + +class FakeBoundaryOperator(LinearOperator): + """Mimics struphy.feec.linear_operators.BoundaryOperator well enough to exercise + `parallel_tosparse`'s diagonal-probe fallback: a diagonal 0/1 mask (zeroes the + first local DOF of each block on each rank), with a `.tosparse()` that is only + valid in serial (returns a small *locally*-indexed diagonal, exactly as struphy's + real implementation does -- see `struphy.feec.linear_operators. + BoundaryOperator.tosparse`) -- so the fast path must detect the mismatch via + validation and retry with the probe strategy instead of trusting `.tosparse()` + blindly. Works for both a plain StencilVectorSpace (V) and a BlockVectorSpace + (e.g. Hcurl, the codomain a real `grad` is wrapped in) domain, matching either + shape BoundaryOperator actually appears in. + """ + + def __init__(self, V): + self._V = V + self._entries = _local_flat_entries(V) + + @property + def domain(self): + return self._V + + @property + def codomain(self): + return self._V + + @property + def dtype(self): + return float + + def tosparse(self): + # Deliberately wrong at nprocs > 1: local indices, not global ones. + n = len(self._entries) + diag = np.ones(n) + diag[0] = 0.0 + return sparse.diags(diag, format="csr") + + def toarray(self): + return self.tosparse().toarray() + + def transpose(self, conjugate=False): + return self + + def dot(self, v, out=None): + if out is None: + out = self.codomain.zeros() + else: + out *= 0.0 + for k, (setter, _) in enumerate(self._entries): + _set_entry(out, setter, 0.0 if k == 0 else _get_entry(v, setter)) + out.update_ghost_regions() + return out + + +class UnsupportedOperator(LinearOperator): + """A leaf `parallel_tosparse` cannot possibly handle: no `.tosparse()`, and + domain is not codomain (so the diagonal-probe strategy doesn't apply either).""" + + def __init__(self, V_in, V_out): + self._Vin = V_in + self._Vout = V_out + + @property + def domain(self): + return self._Vin + + @property + def codomain(self): + return self._Vout + + @property + def dtype(self): + return float + + def tosparse(self): + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + def transpose(self, conjugate=False): + return UnsupportedOperator(self._Vout, self._Vin) + + def dot(self, v, out=None): + if out is None: + out = self.codomain.zeros() + return out + + +class FakeWeightedMassOperator(LinearOperator): + """Mimics struphy.feec.mass.WeightedMassOperator closely enough to exercise + `parallel_tosparse`'s duck-typed `._mat`-composition unwrap: wraps an inner + `._mat` behind identity extraction ops (trivial, as struphy's own serial + `.tosparse()` requires) but *non-trivial* boundary ops (masks the first and last + local DOF on each rank) -- i.e. `.dot()` is genuinely not the same as `._mat.dot()` + alone, which is exactly the bug this class was written to catch (found via a real + Struphy run, not anticipated up front -- see the git history of + `parallel_tosparse`'s `._mat`-unwrap branch). + """ + + def __init__(self, V, mat, mask_first=True, mask_last=True): + self._V = V + self._mat = mat + self._V_extraction_op = IdentityOperator(V) + self._W_extraction_op = IdentityOperator(V) + self._V_boundary_op = _MaskBoundaryOperator(V, mask_first, mask_last) + self._W_boundary_op = _MaskBoundaryOperator(V, mask_first, mask_last) + self._transposed = False + + @property + def domain(self): + return self._V + + @property + def codomain(self): + return self._V + + @property + def dtype(self): + return float + + def tosparse(self): + # Deliberately unusable at nprocs > 1, exactly like struphy's real + # WeightedMassOperator.tosparse() when boundary masking is actually active + # (it asserts outright there); here it just raises, which + # `parallel_tosparse` must also handle gracefully. + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + def transpose(self, conjugate=False): + raise NotImplementedError + + def dot(self, v, out=None): + tmp = self._V_boundary_op.transpose().dot(v) + tmp = self._V_extraction_op.transpose().dot(tmp) + tmp = self._mat.dot(tmp) + tmp = self._W_extraction_op.dot(tmp) + return self._W_boundary_op.dot(tmp, out=out) + + +class _MaskBoundaryOperator(LinearOperator): + """A minimal BoundaryOperator stand-in: zeroes the first and/or last local DOF on + each rank (self-adjoint, so `.transpose()` returns itself).""" + + def __init__(self, V, mask_first, mask_last): + self._V = V + self._entries = _local_flat_entries(V) + self._mask_first = mask_first + self._mask_last = mask_last + + @property + def domain(self): + return self._V + + @property + def codomain(self): + return self._V + + @property + def dtype(self): + return float + + def tosparse(self): + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + def transpose(self, conjugate=False): + return self + + def dot(self, v, out=None): + if out is None: + out = self.codomain.zeros() + else: + out *= 0.0 + n = len(self._entries) + for k, (setter, _) in enumerate(self._entries): + masked = (self._mask_first and k == 0) or (self._mask_last and k == n - 1) + _set_entry(out, setter, 0.0 if masked else _get_entry(v, setter)) + out.update_ghost_regions() + return out + + +@pytest.mark.parametrize('n1', [8, 16]) +@pytest.mark.parametrize('p1', [1, 2]) +@pytest.mark.parallel +def test_parallel_tosparse_matches_matvec_stencil_matrix(n1, p1, verbose=False): + """A plain StencilMatrix (no wrapper) must assemble identically via the fast + (`parallel_tosparse`) and slow (`tosparse_via_matvec`) paths.""" + n2, p2 = 8, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm) + A = make_stencil_matrix(V, p1, p2) + + fast = parallel_tosparse(A, comm) + slow = tosparse_via_matvec(A, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"n1={n1} p1={p1} nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-10 + + +@pytest.mark.parallel +def test_parallel_tosparse_composed_and_sum(verbose=False): + """A Sum-of-Scaled-and-Composed operator tree (the shape ImplicitDiffusion's + left-hand side actually takes: sigma*M + G^T @ D @ G) must also match the slow + reference path.""" + n1, n2, p1, p2 = 10, 8, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm) + M = make_stencil_matrix(V, p1, p2, scale=1.0) + G = make_stencil_matrix(V, p1, p2, scale=0.5) + D = make_stencil_matrix(V, p1, p2, scale=2.0) + + composed = ComposedLinearOperator(V, V, G, D, G) + scaled = ScaledLinearOperator(V, V, c=3.0, A=M) + total = SumLinearOperator(V, V, scaled, composed) + + fast = parallel_tosparse(total, comm) + slow = tosparse_via_matvec(total, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_block_vector_space(verbose=False): + """A BlockLinearOperator (grad's actual shape, e.g. H1 -> Hcurl's 3 stacked + components) must also assemble identically via the fast and slow paths.""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + W = BlockVectorSpace(V, V) + B = BlockLinearOperator(W, W) + B[0, 0] = make_stencil_matrix(V, p1, p2, scale=1.0) + B[1, 1] = make_stencil_matrix(V, p1, p2, scale=2.0) + B[0, 1] = make_stencil_matrix(V, p1, p2, scale=0.3) + + fast = parallel_tosparse(B, comm) + slow = tosparse_via_matvec(B, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_diagonal_probe_fallback_block_vector_space(verbose=False): + """The same diagonal-mask-on-both-sides shape as + `test_parallel_tosparse_diagonal_probe_fallback`, but over a BlockVectorSpace -- + the actual shape `BoundaryOperator ∘ grad ∘ BoundaryOperator` takes in the real + Poisson benchmark (grad's codomain, Hcurl, has 3 stacked components).""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + W = BlockVectorSpace(V, V) + B = BlockLinearOperator(W, W) + B[0, 0] = make_stencil_matrix(V, p1, p2, scale=1.0) + B[1, 1] = make_stencil_matrix(V, p1, p2, scale=2.0) + mask = FakeBoundaryOperator(W) + + total = ComposedLinearOperator(W, W, mask, B, mask) + + fast = parallel_tosparse(total, comm) + slow = tosparse_via_matvec(total, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_diagonal_probe_fallback(verbose=False): + """A BoundaryOperator-like diagonal mask, wired in on both sides of a + StencilMatrix (the actual shape struphy's BC-wrapped operators take), must be + correctly recovered by the diagonal-probe fallback -- not silently misassembled + from its serial-only `.tosparse()`.""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + A = make_stencil_matrix(V, p1, p2) + mask = FakeBoundaryOperator(V) + + total = ComposedLinearOperator(V, V, mask, A, mask) + + fast = parallel_tosparse(total, comm) + slow = tosparse_via_matvec(total, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_raises_for_unsupported_operator(verbose=False): + """An operator this module genuinely cannot handle (no `.tosparse()`, not + diagonal-shaped) must raise `FastAssemblyUnavailable` -- identically on every + rank, so callers can fall back to `tosparse_via_matvec` without any risk of a + partial/divergent collective-call sequence.""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + Vin = make_space(n1, n2, p1, p2, comm, periodic=False) + Vout = make_space(n1, n2, p1, p2, comm, periodic=False) + op = UnsupportedOperator(Vin, Vout) + + with pytest.raises(FastAssemblyUnavailable): + parallel_tosparse(op, comm) + + +def make_deriv_space(n1, n2, p1, p2, comm, periodic): + # dim 0 always periodic (matches make_space's default); dim 1 (the + # differentiation direction in test_parallel_tosparse_directional_derivative) + # uses `periodic`, since that's the one whose periodicity actually changes the + # matrix structure (wraparound coupling vs. a one-fewer-point boundary). + D = DomainDecomposition([n1, n2], periods=[True, periodic], comm=comm) + npts = [n1, n2] + gs, ge = compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, gs, ge, pads=[p1, p2], shifts=[1, 1]) + return StencilVectorSpace(cart, dtype=float) + + +@pytest.mark.parametrize('diffdir_periodic', [False, True]) +@pytest.mark.parametrize('transposed', [False, True]) +@pytest.mark.parallel +def test_parallel_tosparse_directional_derivative(diffdir_periodic, transposed, verbose=False): + """`DirectionalDerivativeOperator`'s default `.tosparse()` asserts outright at + nprocs > 1 (see `_directional_derivative_triples`'s docstring) -- the closed-form + reconstruction it falls back to instead must match the slow reference path, for + both a periodic and a non-periodic differentiation direction, transposed or not. + """ + p1, p2 = 1, 1 + comm = MPI.COMM_WORLD + n1, n2 = 8, 8 + V = make_deriv_space(n1, n2, p1, p2, comm, periodic=diffdir_periodic) + W = make_deriv_space(n1, n2 if diffdir_periodic else n2 - 1, p1, p2, comm, periodic=diffdir_periodic) + + op = DirectionalDerivativeOperator(V, W, diffdir=1, negative=False, transposed=transposed) + + fast = parallel_tosparse(op, comm) + slow = tosparse_via_matvec(op, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"periodic={diffdir_periodic} transposed={transposed} nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-10 + + +@pytest.mark.parallel +def test_parallel_tosparse_weighted_mass_operator_boundary_composition(verbose=False): + """`FakeWeightedMassOperator` wraps a StencilMatrix behind trivial extraction ops + but *non-trivial* boundary masking (`.dot()` != `._mat.dot()` alone) -- exactly + the shape that caused a real, silent ~8% numeric mismatch against the naive + "just unwrap `._mat`" shortcut on an actual Struphy run (before + `parallel_tosparse`'s `._mat`-composition unwrap rebuilt the *whole* + boundary/extraction chain instead). Must match the slow reference path. + """ + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + inner = make_stencil_matrix(V, p1, p2) + op = FakeWeightedMassOperator(V, inner) + + fast = parallel_tosparse(op, comm) + slow = tosparse_via_matvec(op, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index e7d47a233..79dd35121 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -18,6 +18,8 @@ 'array_to_psydac', 'petsc_to_psydac', 'tosparse_via_matvec', + 'parallel_tosparse', + 'FastAssemblyUnavailable', '_sym_ortho', ) @@ -259,6 +261,457 @@ def tosparse_via_matvec(op, format="csc"): mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) return mat.asformat(format) +#============================================================================== +class FastAssemblyUnavailable(Exception): + """Raised by `parallel_tosparse` when the operator tree could not be assembled via + the fast (O(1)-communication-round) path -- see its docstring. Callers should catch + this and fall back to `tosparse_via_matvec`.""" + + +def _local_flat_entries(V): + """ + List of (setter, global_flat_index) pairs, one per DOF *owned* by this rank (no + ghost/pad region), for a StencilVectorSpace or BlockVectorSpace V. + + `setter` is `('b', block_index, multi_index)` (usable as `vec[h][idx] = ...` for a + BlockVector) or `('s', multi_index)` (usable as `vec[idx] = ...` otherwise) -- + tagged rather than inferred from shape, since a StencilVectorSpace's own + `multi_index` can itself start with an int indistinguishable from a block index. + `global_flat_index` uses the same block-major, + `numpy.ravel_multi_index`-against-global-`npts` convention as + `tosparse_via_matvec` and `StencilMatrix.tosparse()` (`_tocoo_no_pads`) -- the same + one `Vector.toarray()` flattens to, which every caller of this module (e.g. + `DirectSolver.solve`'s `b.toarray()`/`x_flat.reshape(...)` round trip) already + relies on. All three MUST agree, since results from this function are combined + with plain `StencilMatrix`/`BlockLinearOperator` sparse matrices in the same + right-hand-side/solution vectors. + """ + if isinstance(V, BlockVectorSpace): + entries = [] + spoint = 0 + for h, sp in enumerate(V.spaces): + npts = sp.npts + iterables = [range(s, e + 1) for s, e in zip(sp.starts, sp.ends)] + for idx in itertools.product(*iterables): + flat = spoint + int(np.ravel_multi_index(idx, npts)) + entries.append((('b', h, idx), flat)) + spoint += int(np.prod(npts)) + return entries + elif isinstance(V, StencilVectorSpace): + npts = V.npts + iterables = [range(s, e + 1) for s, e in zip(V.starts, V.ends)] + return [(('s', idx), int(np.ravel_multi_index(idx, npts))) for idx in itertools.product(*iterables)] + else: + raise FastAssemblyUnavailable( + f'_local_flat_entries only supports StencilVectorSpace/BlockVectorSpace, got {type(V)}', + ) + + +def _set_entry(vec, setter, value): + if setter[0] == 'b': + _, h, idx = setter + vec[h][idx] = value + else: + _, idx = setter + vec[idx] = value + + +def _get_entry(vec, setter): + if setter[0] == 'b': + _, h, idx = setter + return vec[h][idx] + else: + _, idx = setter + return vec[idx] + + +def _replicate_triples(rows, cols, vals, shape, comm, dtype): + """Gather (rows, cols, vals) COO triples -- assumed *local* to this rank -- from + every rank and sum-combine them (matching duplicates, e.g. periodic wraparound, + exactly as `scipy.sparse.coo_matrix` does on `.tocsr()`) into one matrix identical + on every rank. One collective round, regardless of `shape`. + + `rows`/`cols`/`vals` may be plain sequences or arrays; always gathered and + concatenated as numpy arrays (`comm.allgather` pickles a numpy array through + mpi4py's out-of-band buffer protocol, and `np.concatenate` is vectorized) rather + than as Python lists -- for a leaf with real FEM bandwidth (tens to hundreds of + thousands of local nonzeros, e.g. a 3D mass matrix), converting through + element-by-element Python lists first was the dominant cost, dwarfing the O(1) + round-count win this function exists for. + """ + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + vals = np.asarray(vals, dtype=dtype) + + if comm is None or isinstance(comm, MockComm): + all_rows, all_cols, all_vals = rows, cols, vals + else: + gathered = comm.allgather((rows, cols, vals)) + all_rows = np.concatenate([g[0] for g in gathered]) + all_cols = np.concatenate([g[1] for g in gathered]) + all_vals = np.concatenate([g[2] for g in gathered]) + return sparse.coo_matrix((all_vals, (all_rows, all_cols)), shape=shape, dtype=dtype).tocsr() + + +def _probe_vector(V, entries, offset): + """A deterministic, reproducible-across-ranks probe Vector of space V: each owned + DOF gets a distinct nonzero value derived from its global flat index (never 0, and + never equal across two different `offset`s), so an operator's actual coupling + structure is very unlikely to accidentally look diagonal/masking by coincidence.""" + v = V.zeros() + for setter, flat in entries: + _set_entry(v, setter, 1.0 + 0.618033988749895 * ((flat + offset) % 104729)) + v.update_ghost_regions() + return v + + +def _validate_against_dot(node, candidate, comm, entries_domain, entries_codomain, seed): + """Check `candidate @ p == node.dot(p)` (this rank's owned output entries only) for + one probe vector `p`built from `_probe_vector`. `entries_domain` values already + carry each entry's global flat column index; `entries_codomain` likewise for rows. + Returns a local bool -- the caller combines these across ranks (see + `parallel_tosparse`) before trusting `candidate`.""" + V_domain = node.domain + p = _probe_vector(V_domain, entries_domain, seed) + p_flat_local = {flat: _get_entry(p, setter) for setter, flat in entries_domain} + + if comm is None or isinstance(comm, MockComm): + p_flat_full = dict(p_flat_local) + else: + gathered = comm.allgather(p_flat_local) + p_flat_full = {} + for d in gathered: + p_flat_full.update(d) + + p_full = np.zeros(candidate.shape[1], dtype=candidate.dtype) + for flat, val in p_flat_full.items(): + p_full[flat] = val + + q_candidate = candidate @ p_full + q_true = node.dot(p) + + # An aggregate (L2-norm) check, not a per-entry one: a wide-bandwidth FEM operator + # (e.g. a mass matrix with a degree-3 spline direction) sums many terms per row, + # in a different order than `candidate`'s (scipy's own summation order for the + # sparse matvec) -- individual output entries can then legitimately differ by much + # more than a tight per-entry relative tolerance even when `candidate` is exactly + # right, especially where terms partially cancel. Comparing the whole local output + # vector's norm to the whole error vector's norm is robust to that per-entry + # cancellation while still easily catching a genuinely wrong `candidate` (which + # differs at O(1) relative scale, not at rounding-error scale). + true_local = np.fromiter((_get_entry(q_true, setter) for setter, _ in entries_codomain), dtype=float) + cand_local = np.fromiter((q_candidate[flat] for _, flat in entries_codomain), dtype=float) + err = float(np.linalg.norm(true_local - cand_local)) + scale = float(np.linalg.norm(true_local)) + return err <= 1e-8 * scale + 1e-10 + + +def _directional_derivative_triples(op): + """Closed-form local (row, col, value) triples for a + `feectools.feec.derivatives.DirectionalDerivativeOperator` -- `.tosparse()`'s + default (no-pads) form isn't valid at nprocs > 1 for this operator (it asserts), + and its `with_pads=True` form returns a small *local* matrix in a totally + different (ghost-inclusive, non-globally-indexed) convention this module's + block-major global indexing can't reuse -- so this reconstructs the same bidiagonal + difference-operator matrix its serial `.tosparse()` builds, directly from the + operator's own definition (`out[i] = sign * (in[i + e_d] - in[i])` along direction + `d = op._diffdir`, `e_d` wrapped modulo `V.npts[d]` when periodic), one local + (globally-indexed) row at a time -- no basis-vector sweep, no padding subtleties. + Built in the "V -> W" (non-transposed) sense regardless of `op._transposed`; + `parallel_tosparse` transposes the result back if needed, exactly as the operator's + own serial `.tosparse()` does. + """ + V, W, d = op._spaceV, op._spaceW, op._diffdir + sign = -1.0 if op._negative else 1.0 + periodic = V.periods[d] + + if V.npts[d] == 1 and W.npts[d] == 1 and periodic: + return [], [], [] # degenerate single-cell-periodic case: the zero matrix + + rows, cols, vals = [], [], [] + for idx, row_flat in _local_flat_entries(W): + _, ii = idx + jj = ii + jj_next = list(ii) + jj_next[d] = (ii[d] + 1) % V.npts[d] if periodic else ii[d] + 1 + col_flat = int(np.ravel_multi_index(jj, V.npts)) + rows.append(row_flat) + cols.append(col_flat) + vals.append(-sign) + if periodic or jj_next[d] < V.npts[d]: + col_next_flat = int(np.ravel_multi_index(jj_next, V.npts)) + rows.append(row_flat) + cols.append(col_next_flat) + vals.append(sign) + return rows, cols, vals + + +def parallel_tosparse(op, comm, format="csr"): + """ + Assemble the full global sparse matrix of a `LinearOperator` tree using O(1) + collective-communication rounds (one per leaf node, roughly), instead of + `tosparse_via_matvec`'s O(`op.domain.dimension`) rounds (one basis vector per + global DOF) -- for the same replicated-on-every-rank result. + + Walks the operator tree using only types `feectools` itself defines + (`SumLinearOperator`, `ScaledLinearOperator`, `ComposedLinearOperator`, + `IdentityOperator`, `ZeroOperator`, `BlockLinearOperator`, + `DirectionalDerivativeOperator`): the first five compose exactly the way + `.tosparse()` already does in serial, just with the *leaves* below assembled + without a per-DOF basis-vector sweep; `BlockLinearOperator` is recursed into + block-by-block (not treated as one leaf) since a real Derham `grad`/`grad.T` can + have `DirectionalDerivativeOperator` blocks, whose own `.tosparse()` is unusable + in parallel (see `_directional_derivative_triples`, which reconstructs it in + closed form instead). For anything else -- most of it defined outside feectools + (`struphy.feec.mass.WeightedMassOperator`, `struphy.feec.linear_operators. + BoundaryOperator`, ...), which this module must not import -- one of three + O(1)-round strategies applies, tried in order: + + 1. Duck-typed unwrap: if the node has a `._mat` plus the same + `._V_extraction_op`/`._W_extraction_op`/`._V_boundary_op`/`._W_boundary_op`/ + `._transposed` attributes `struphy.feec.mass.WeightedMassOperator` has, + rebuild the exact composition its own `.dot()` applies (boundary and + extraction maps included, not just `._mat` alone -- an earlier version of + this function assumed trivial extraction ops meant `._mat` alone was enough, + which a real Struphy run showed is false whenever the boundary masks are + non-trivial) from parts each recursed into via `build()` in turn. + + 2. If the leaf has its own `.tosparse()` (true of `StencilMatrix` and + `BlockLinearOperator`): call it *locally* (no communication -- the same call + `A.tosparse()` already makes in serial, just once per rank instead of once + globally) and `allgather`-sum the local fragments into the replicated global + matrix. + + 3. Otherwise, if `leaf.domain is leaf.codomain` (a necessary condition to act as + a diagonal map): probe it with a value-tagged vector and check whether the + output is consistent with a per-DOF diagonal scaling (this is exactly what + essential-BC masking operators like `BoundaryOperator` are). Two `.dot()` + calls plus one `allgather`. + + Every leaf's result is *always* cross-checked against the operator's own `.dot()` + on a probe vector (strategies 1 and 2 both feed their candidate through the same + `_validate_against_dot` check that strategy 3 uses to detect diagonality in the + first place); an unrecognized leaf type (none of the three strategies applicable + or valid) is treated the same as a failed check. Whether any of this happens is a + pure function of operator *types*, identical on every rank by construction (same + model, same run) -- so every rank always issues the same sequence of collective + calls regardless of any individual check's pass/fail outcome; only *after* the + full tree is walked does one final `allreduce(MPI.LAND)` combine every check + across every rank into a single decision, so a data-dependent failure on one rank + cannot leave another rank waiting on a collective call that rank never issues (no + deadlock risk from divergent control flow). If that combined decision is False -- + or the tree contains a node type this function does not know how to handle at all + (e.g. `MatrixFreeLinearOperator`, always a deterministic, type-only decision, so + still consistent across ranks) -- `FastAssemblyUnavailable` is raised (on every + rank, identically) and the caller should fall back to `tosparse_via_matvec`. + + Parameters + ---------- + op : feectools.linalg.basic.LinearOperator + Operator to assemble. + + comm : MPI.Comm | feectools.ddm.mpi.MockComm | None + Communicator spanning every rank that owns a piece of `op`. + + format : str + scipy.sparse matrix format of the result. + + Returns + ------- + out : scipy.sparse matrix + The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on + every rank. + """ + # Imported here, not at module scope: these are feectools types this function + # checks via isinstance, kept local to make the "only feectools composite types + # are special-cased" contract easy to audit at a glance. + from feectools.linalg.basic import ComposedLinearOperator, IdentityOperator, ScaledLinearOperator, SumLinearOperator, ZeroOperator + from feectools.linalg.block import BlockLinearOperator + from feectools.feec.derivatives import DirectionalDerivativeOperator + + checks = [] + probe_seed = [1000003] # mutable cell; a fresh seed per leaf keeps probes independent + + def build(node): + if isinstance(node, ScaledLinearOperator): + return node._scalar * build(node._operator) + if isinstance(node, SumLinearOperator): + mats = [build(a) for a in node._addends] + out = mats[0] + for m in mats[1:]: + out = out + m + return out + if isinstance(node, ComposedLinearOperator): + mats = [build(m) for m in node._multiplicants] + out = mats[0] + for m in mats[1:]: + out = out @ m + return out + if isinstance(node, IdentityOperator): + return sparse.identity(node.domain.dimension, format="csr", dtype=node.dtype or float) + if isinstance(node, ZeroOperator): + return sparse.csr_matrix(node.shape, dtype=node.dtype or float) + if isinstance(node, BlockLinearOperator): + # Recurse into each block individually rather than calling + # `node.tosparse()` on the whole thing: a real Derham `grad`/`grad.T` is a + # BlockLinearOperator whose blocks can themselves be + # `DirectionalDerivativeOperator`s, whose *own* default `.tosparse()` + # asserts outright at nprocs > 1 (see `_directional_derivative_triples`) + # -- one such block would otherwise make the whole (possibly mostly + # StencilMatrix) BlockLinearOperator's `.tosparse()` raise. + nrows, ncols = node.n_block_rows, node.n_block_cols + block_domain = (lambda j: node.domain[j]) if ncols > 1 else (lambda j: node.domain) + block_codomain = (lambda i: node.codomain[i]) if nrows > 1 else (lambda i: node.codomain) + grid = [[None for _ in range(ncols)] for _ in range(nrows)] + for i in range(nrows): + for j in range(ncols): + if (i, j) in node._blocks: + grid[i][j] = build(node._blocks[i, j]) + else: + grid[i][j] = sparse.csr_matrix((block_codomain(i).dimension, block_domain(j).dimension)) + return sparse.bmat(grid, format="csr") + if isinstance(node, DirectionalDerivativeOperator): + # No generic strategy below applies (not diagonal-shaped in general, and + # its own `.tosparse()` is unusable here -- see + # `_directional_derivative_triples`); its structure is simple and fixed + # enough to reconstruct in closed form directly, still validated below + # like everything else. + V, W = node._spaceV, node._spaceW + rows, cols, vals = _directional_derivative_triples(node) + mat_vw = _replicate_triples(rows, cols, vals, (W.dimension, V.dimension), comm, node.dtype or float) + candidate = mat_vw.T.tocsr() if node._transposed else mat_vw + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + probe_seed[0] += 97 + ok = _validate_against_dot(node, candidate, comm, entries_domain, entries_codomain, probe_seed[0]) + checks.append(ok) + return candidate + + # Leaf: not one of the composite types above. Every strategy applicable to + # this leaf's *type* is always attempted, on every rank, regardless of any + # other rank's or strategy's data-dependent validation outcome -- see the + # docstring's "same collective calls on every rank" invariant. Only the first + # strategy that actually validates is kept. + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + shape = (node.codomain.dimension, node.domain.dimension) + dtype = node.dtype or float + probe_seed[0] += 97 + + # Duck-typed unwrap: struphy's `WeightedMassOperator` (M0, M1, ...) computes + # `V_boundary_op @ V_extraction_op @ _mat @ W_extraction_op.T @ W_boundary_op.T` + # (or the mirrored order when `._transposed`) on every `.dot()` call, by + # default with the boundary masks actually applied (`apply_bc=True`) -- *not* + # just `._mat` alone, even when both extraction ops are the identity (its own + # boundary masks can still be non-trivial, e.g. Dirichlet-BC-adjacent DOFs on + # a component of an Hcurl mass matrix -- discovered by this function's own + # validation rejecting the naive "just `._mat`" shortcut on exactly such a + # case, not by inspecting struphy's BC configuration). Rebuilding that same + # composition from its parts -- each recursed into via `build()`, so a + # boundary mask that itself needs the diagonal-probe strategy below still + # gets it -- is exact when every part is present (duck-typed by attribute, + # not `isinstance`, since none of these types live in feectools); still + # validated below regardless, as insurance against this composition itself + # being incomplete for some other struphy wrapper shaped differently. + parts = [getattr(node, name, "missing") for name in ( + "_mat", "_V_extraction_op", "_W_extraction_op", "_V_boundary_op", "_W_boundary_op", + )] + if "missing" not in parts: + inner_mat, v_ext, w_ext, v_bnd, w_bnd = parts + transposed = bool(getattr(node, "_transposed", False)) + try: + # Matches struphy.feec.mass.WeightedMassOperator.dot's own step + # sequence exactly (apply_bc=True, its default): non-transposed + # applies V_boundary_op.T, then V_extraction_op.T, then `._mat`, then + # W_extraction_op, then W_boundary_op, in that order (v -> out); the + # composed *matrix* is those same maps in reverse (rightmost applied + # first). `._transposed` mirrors V and W throughout. + if not transposed: + order = [w_bnd, w_ext, inner_mat, v_ext.transpose(), v_bnd.transpose()] + else: + order = [v_bnd, v_ext, inner_mat, w_ext.transpose(), w_bnd.transpose()] + mats = [build(m) for m in order] + candidate_inner = mats[0] + for m in mats[1:]: + candidate_inner = candidate_inner @ m + except Exception: + candidate_inner = None + if candidate_inner is not None and _validate_against_dot( + node, candidate_inner, comm, entries_domain, entries_codomain, probe_seed[0], + ): + checks.append(True) + return candidate_inner + + candidate_tosparse = None + try: + local_coo = node.tosparse().tocoo() + candidate_tosparse = _replicate_triples( + local_coo.row, local_coo.col, local_coo.data, + shape, comm, dtype, + ) + except Exception: + pass # this leaf's .tosparse() -- if it has one -- doesn't work here (e.g. + # raises, or -- as for struphy's BoundaryOperator -- succeeds but is + # documented serial-only and produces locally- rather than + # globally-indexed rows/cols at nprocs > 1); validation below (or, failing + # that, the diagonal-probe strategy) is what actually decides trust, not + # whether this call happened to raise. + + if candidate_tosparse is not None and _validate_against_dot( + node, candidate_tosparse, comm, entries_domain, entries_codomain, probe_seed[0], + ): + checks.append(True) + return candidate_tosparse + + if node.domain is not node.codomain: + # Not diagonal-shaped, and the .tosparse() attempt above (if any) didn't + # validate: nothing left to try for this leaf. + checks.append(False) + return candidate_tosparse if candidate_tosparse is not None else sparse.csr_matrix(shape, dtype=dtype) + + # Diagonal-probe strategy: two independent value-tagged probes; a true + # diagonal map reproduces (scaled by a per-DOF constant) or zeroes each one, + # consistently between the two -- exactly what an essential-BC mask does. + p1 = _probe_vector(node.domain, entries_domain, probe_seed[0]) + p2 = _probe_vector(node.domain, entries_domain, probe_seed[0] + 50000) + try: + o1 = node.dot(p1) + o2 = node.dot(p2) + except Exception: + checks.append(False) + return candidate_tosparse if candidate_tosparse is not None else sparse.csr_matrix(shape, dtype=dtype) + + rows, cols, vals = [], [], [] + ok = True + for setter, flat in entries_domain: + v1 = _get_entry(p1, setter) + v2 = _get_entry(p2, setter) + a1 = _get_entry(o1, setter) + a2 = _get_entry(o2, setter) + is_zero = abs(a1) < 1e-300 and abs(a2) < 1e-300 + if is_zero: + continue + d1 = a1 / v1 + d2 = a2 / v2 + if abs(d1 - d2) > 1e-8 * max(1.0, abs(d1)): + ok = False + break + rows.append(flat) + cols.append(flat) + vals.append(d1) + checks.append(ok) + return _replicate_triples(rows, cols, vals, shape, comm, dtype) + + result = build(op) + + all_ok = all(checks) + if comm is not None and not isinstance(comm, MockComm): + all_ok = comm.allreduce(all_ok, op=MPI.LAND) + if not all_ok: + raise FastAssemblyUnavailable('one or more leaf operators could not be verified') + + return result.asformat(format) + #============================================================================== def petsc_to_psydac(x, Xh, out=None): """ From 8b5ccf2cf97e0fc0eb2d04e6164f3d169f319c23 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 19:18:17 +0200 Subject: [PATCH 17/23] Added codomain_local_nonzero_rows utility --- feectools/linalg/direct_solvers.py | 6 +--- feectools/linalg/utilities.py | 54 +++++++++++++++++++++++------- 2 files changed, 43 insertions(+), 17 deletions(-) diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 0baa0c8b9..785ce0107 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -69,11 +69,7 @@ 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) + bmat = xp.to_numpy(bmat) self._bmat, self._ipiv, self._finfo = self._factor_function(bmat, l, u) self._sinfo = None diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 79dd35121..41407d9c5 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -141,6 +141,33 @@ def tosparse_via_matvec(op, format="csc"): numcols = op.domain.dimension data, row, col = [], [], [] + def local_nonzero_rows(stencil_vec, row_offset): + """(global_row_indices, values) for `stencil_vec`'s local interior data.""" + space = stencil_vec.space + idx_local = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(stencil_vec.pads, space.shifts) + ) + local_data = xp.to_numpy(stencil_vec._data[idx_local]) + nz = np.nonzero(local_data) + starts = space.starts + global_multi = tuple(nz[d] + int(starts[d]) for d in range(len(nz))) + rows = row_offset + np.ravel_multi_index(global_multi, space.npts) + return rows, local_data[nz] + + def codomain_local_nonzero_rows(vec): + """`local_nonzero_rows`, dispatched over `op.codomain`'s type.""" + if isinstance(op.codomain, BlockVectorSpace): + all_rows, all_vals = [], [] + row_offset = 0 + for b, sp in enumerate(op.codomain.spaces): + r, val = local_nonzero_rows(vec[b], row_offset) + all_rows.append(r) + all_vals.append(val) + row_offset += sp.dimension + return np.concatenate(all_rows), np.concatenate(all_vals) + return local_nonzero_rows(vec, 0) + if isinstance(op.domain, BlockVectorSpace): starts = [vi.starts for vi in v] ends = [vi.ends for vi in v] @@ -184,14 +211,12 @@ def tosparse_via_matvec(op, format="csc"): tmp2 *= 0.0 op.dot(v, out=tmp2) c = spoint + int(np.ravel_multi_index(i, npts[h])) - aux = xp.to_numpy(tmp2.toarray()) - for r in np.nonzero(aux)[0]: - data.append(aux[r]) - col.append(c) - row.append(int(r)) + rs, vals = codomain_local_nonzero_rows(tmp2) + row.append(rs) + col.append(np.full(rs.shape, c)) + data.append(vals) if rank == currentrank: v[h][i] = 0.0 - v[h].update_ghost_regions() cumulative = 1 for i in range(ndim[h]): cumulative *= npts[h][i] @@ -231,14 +256,12 @@ def tosparse_via_matvec(op, format="csc"): v.update_ghost_regions() op.dot(v, out=tmp2) c = int(np.ravel_multi_index(i, npts)) - aux = xp.to_numpy(tmp2.toarray()) - for r in np.nonzero(aux)[0]: - data.append(aux[r]) - col.append(c) - row.append(int(r)) + rs, vals = codomain_local_nonzero_rows(tmp2) + row.append(rs) + col.append(np.full(rs.shape, c)) + data.append(vals) if rank == currentrank: v[i] = 0.0 - v.update_ghost_regions() if comm is None or isinstance(comm, MockComm): all_rows, all_cols, all_data = row, col, data @@ -258,6 +281,13 @@ def tosparse_via_matvec(op, format="csc"): all_cols = comm.bcast(None, root=0) all_data = comm.bcast(None, root=0) + if all_rows: + all_rows = np.concatenate(all_rows) + all_cols = np.concatenate(all_cols) + all_data = np.concatenate(all_data) + else: + all_rows = all_cols = all_data = np.empty(0, dtype=int) + mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) return mat.asformat(format) From 6b7b090c71b1acdbbd8ce733e30cf3737237bf20 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 20 Aug 2026 12:48:04 +0200 Subject: [PATCH 18/23] bugfix --- feectools/linalg/utilities.py | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index ea029f950..ba5315af6 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -589,6 +589,7 @@ def build(node): # asserts outright at nprocs > 1 (see `_directional_derivative_triples`) # -- one such block would otherwise make the whole (possibly mostly # StencilMatrix) BlockLinearOperator's `.tosparse()` raise. + checks_before = len(checks) nrows, ncols = node.n_block_rows, node.n_block_cols block_domain = (lambda j: node.domain[j]) if ncols > 1 else (lambda j: node.domain) block_codomain = (lambda i: node.codomain[i]) if nrows > 1 else (lambda i: node.codomain) @@ -599,7 +600,39 @@ def build(node): grid[i][j] = build(node._blocks[i, j]) else: grid[i][j] = sparse.csr_matrix((block_codomain(i).dimension, block_domain(j).dimension)) - return sparse.bmat(grid, format="csr") + candidate = sparse.bmat(grid, format="csr") + + # A block matrix can only be trusted if the already-validated child + # blocks also agree with the parent BlockLinearOperator's own indexing + # convention. Repeated component spaces (e.g. BlockVectorSpace(V, V)) are + # especially easy to stitch incorrectly while every scalar block still + # validates in isolation. + if not all(checks[checks_before:]): + return candidate + + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + probe_seed[0] += 97 + block_ok = _validate_against_dot( + node, candidate, comm, entries_domain, entries_codomain, probe_seed[0], + ) + if comm is not None and not isinstance(comm, MockComm): + block_ok = comm.allreduce(block_ok, op=MPI.LAND) + if block_ok: + checks.append(True) + return candidate + + # Preserve correctness for valid child blocks even when the assembled + # block offsets do not match the parent operator. This is slower, but + # still deterministic across ranks and keeps callers from receiving a + # silently wrong sparse matrix. + try: + fallback = tosparse_via_matvec(node, format="csr") + except Exception: + checks.append(False) + return candidate + checks.append(True) + return fallback if isinstance(node, DirectionalDerivativeOperator): # No generic strategy below applies (not diagonal-shaped in general, and # its own `.tosparse()` is unusable here -- see From 06f95fcc62f095e43b1cd2d46eb46452c287318f Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 20 Aug 2026 12:53:22 +0200 Subject: [PATCH 19/23] bugfix --- feectools/linalg/utilities.py | 36 +++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 41407d9c5..8dcc54221 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -589,6 +589,7 @@ def build(node): # asserts outright at nprocs > 1 (see `_directional_derivative_triples`) # -- one such block would otherwise make the whole (possibly mostly # StencilMatrix) BlockLinearOperator's `.tosparse()` raise. + checks_before = len(checks) nrows, ncols = node.n_block_rows, node.n_block_cols block_domain = (lambda j: node.domain[j]) if ncols > 1 else (lambda j: node.domain) block_codomain = (lambda i: node.codomain[i]) if nrows > 1 else (lambda i: node.codomain) @@ -599,7 +600,34 @@ def build(node): grid[i][j] = build(node._blocks[i, j]) else: grid[i][j] = sparse.csr_matrix((block_codomain(i).dimension, block_domain(j).dimension)) - return sparse.bmat(grid, format="csr") + candidate = sparse.bmat(grid, format="csr") + + children_ok = all(checks[checks_before:]) + if comm is not None and not isinstance(comm, MockComm): + children_ok = comm.allreduce(children_ok, op=MPI.LAND) + block_ok = False + if children_ok: + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + probe_seed[0] += 97 + block_ok = _validate_against_dot( + node, candidate, comm, entries_domain, entries_codomain, probe_seed[0], + ) + if comm is not None and not isinstance(comm, MockComm): + block_ok = comm.allreduce(block_ok, op=MPI.LAND) + if block_ok: + checks.append(True) + return candidate + + try: + fallback = tosparse_via_matvec(node, format="csr") + except Exception: + checks.append(False) + return candidate + + del checks[checks_before:] + checks.append(True) + return fallback if isinstance(node, DirectionalDerivativeOperator): # No generic strategy below applies (not diagonal-shaped in general, and # its own `.tosparse()` is unusable here -- see @@ -729,8 +757,12 @@ def build(node): rows.append(flat) cols.append(flat) vals.append(d1) + candidate_diag = _replicate_triples(rows, cols, vals, shape, comm, dtype) + ok = ok and _validate_against_dot( + node, candidate_diag, comm, entries_domain, entries_codomain, probe_seed[0], + ) checks.append(ok) - return _replicate_triples(rows, cols, vals, shape, comm, dtype) + return candidate_diag result = build(op) From cefae919ade2bd8b339d74a11b5d59e1c88fa900 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 20 Aug 2026 13:33:17 +0200 Subject: [PATCH 20/23] Added v[h].update_ghost_regions() --- feectools/linalg/utilities.py | 1 + 1 file changed, 1 insertion(+) diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 8dcc54221..bd5556141 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -217,6 +217,7 @@ def codomain_local_nonzero_rows(vec): data.append(vals) if rank == currentrank: v[h][i] = 0.0 + v[h].update_ghost_regions() cumulative = 1 for i in range(ndim[h]): cumulative *= npts[h][i] From 5aeb2c2260a1a15cbfc41f73640efde064a8f6fa Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 20 Aug 2026 14:10:01 +0200 Subject: [PATCH 21/23] Removed the DirectSolver --- feectools/linalg/direct_solvers.py | 3 +- feectools/linalg/solvers.py | 217 ----- feectools/linalg/tests/test_solvers.py | 61 +- feectools/linalg/tests/test_utilities.py | 428 --------- feectools/linalg/utilities.py | 1070 +--------------------- 5 files changed, 48 insertions(+), 1731 deletions(-) delete mode 100644 feectools/linalg/tests/test_utilities.py diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 785ce0107..8d02a6b21 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -234,8 +234,7 @@ def solve(self, rhs, out=None): # device array (not the global `array_backend.backend` flag): the LU # factorization always lives on the host regardless of backend, and a caller # may deliberately pass an already-host `rhs`/`out` pair even while the - # active backend is CuPy (see feectools.linalg.solvers.DirectSolver), in - # which case `.get()`-ing a plain NumPy array would fail outright. + # active backend is CuPy. if xp.is_gpu(rhs): rhs_cpu = xp.to_numpy(rhs) result_cpu = self._splu.solve(rhs_cpu.T, trans='T' if transposed else 'N').T diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index a90ff24fc..56ca82743 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -4,11 +4,8 @@ """ import cunumpy as xp -import numpy as np from math import sqrt, inf -from feectools.ddm.mpi import MockComm -from feectools.ddm.mpi import mpi as MPI from feectools.utilities.utils import is_real from feectools.linalg.utilities import _sym_ortho from feectools.linalg.basic import (Vector, LinearOperator, @@ -20,7 +17,6 @@ 'inverse', 'ConjugateGradient', 'PConjugateGradient', - 'DirectSolver', 'BiConjugateGradient', 'BiConjugateGradientStabilized', 'PBiConjugateGradientStabilized', @@ -64,7 +60,6 @@ def inverse(A, solver, **kwargs): solvers_dict = { 'cg' : ConjugateGradient, 'pcg' : PConjugateGradient, - 'direct' : DirectSolver, 'bicg' : BiConjugateGradient, 'bicgstab' : BiConjugateGradientStabilized, 'pbicgstab': PBiConjugateGradientStabilized, @@ -421,218 +416,6 @@ def solve(self, b, out=None): def dot(self, b, out=None): return self.solve(b, out=out) -#=============================================================================== -class DirectSolver(InverseLinearOperator): - """ - Exact sparse-direct solve, for linear systems whose left-hand-side operator A does - not actually change across repeated `solve()` calls -- e.g. a time-independent field - operator solved once per time step with only the right-hand side changing (see - `struphy.propagators.implicit_diffusion.ImplicitDiffusion`, whose LHS is constant - whenever `divide_by_dt=False`). A single sparse LU factorization - (`feectools.linalg.direct_solvers.SparseSolver`) then serves every call, instead of - an iterative method repeating (in the worst case, all the way to `maxiter`) every - single call. - - The factorization is built lazily, on the first `solve()` call, and then reused by - every later call without ever re-examining `A` again -- including through a `.linop` - reassignment, e.g. `ImplicitDiffusion.__call__` unconditionally reassigns `.linop` to - a freshly *built* operator every step, regardless of whether its *values* actually - changed. This is a deliberate, cheap-by-construction design, not a value comparison: - `A.tosparse()` is not assumed to be cheap (composed operators can include a - basis-vector sweep, see e.g. `AverageOperator.tosparse`/`BoundaryOperator.tosparse` - in `struphy.feec.mass`/`struphy.feec.linear_operators`), so re-deriving and comparing - it on every call would undo most of the point of factorizing once. The caller is - therefore responsible for knowing that `A`'s *values* are actually constant across - calls (true whenever `ImplicitDiffusion.divide_by_dt=False`, since neither `epsilon` - nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and - the factorization must be rebuilt on the next `solve()`. - - At `nprocs > 1`, this factorizes a *replicated* copy of the full global matrix on - every rank (assembled once via `feectools.linalg.utilities.tosparse_via_matvec`, - which applies `A` to every global unit vector through its own -- already - MPI-correct -- `.dot()`, since `A.tosparse()` itself is only correct in serial for - several composed/derivative operators), rather than attempting an actual - distributed factorization. Every rank redundantly solves the same full system and - keeps only its own slice of the result -- correct and simple, but each rank does - `O(A.domain.dimension)` work per solve instead of `O(A.domain.dimension / nprocs)`, - and the one-time assembly is `O(A.domain.dimension)` *collective* `.dot()` calls that - do not get cheaper with more ranks. This trade only makes sense for problems small - enough that `splu` and this redundant work stay cheap (e.g. the few-thousand-DOF - field solves this class targets); a genuinely distributed sparse-direct solve (e.g. - via PETSc/MUMPS) would need its own implementation. - - Parameters - ---------- - A : feectools.linalg.basic.LinearOperator - Left-hand-side matrix A of the linear system. Must support `.tosparse()` (serial) - or `.dot()` (parallel, via `tosparse_via_matvec`). - - pc, tol, maxiter, verbose : ignored - Accepted only so this class is a drop-in alternative to the iterative solvers - behind the same `solvers.inverse(A, solver, ...)` call site; a direct solve has - no preconditioner, iteration count, or convergence tolerance. - - x0 : feectools.linalg.basic.Vector, optional - Ignored for solving (a direct solve needs no initial guess); if `recycle=True`, - still receives a copy of each solution, for interface consistency with the - iterative solvers (some callers read `x0` back out directly). - - recycle : bool - If True, a copy of the output is stored in x0, as the iterative solvers do. - """ - - def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False, recycle=False): - - self._options = {"x0": x0, "pc": pc, "tol": tol, "maxiter": maxiter, "verbose": verbose, "recycle": recycle} - - super().__init__(A, **self._options) - - # `.parallel` only means "an MPI communicator is attached", true even at 1 rank - # (e.g. under `srun -n 1`), where the serial `.tosparse()` path is already - # correct (local range == global range) and faster than the replicated-assembly - # path -- so check the rank *count* (`cart.nprocs`) directly. - cart = self.domain.spaces[0].cart if isinstance(self.domain, BlockVectorSpace) else self.domain.cart - self._parallel = self.domain.parallel and any(n != 1 for n in cart.nprocs) - self._comm = cart.comm if self._parallel else None - - self._sparse_solver = None - self._info = None - - def _check_options(self, **kwargs): - # tol/maxiter/verbose are meaningless for a direct solve (see class docstring); - # only x0, if given, is worth the base class's type/space check. - x0 = kwargs.get("x0") - if x0 is not None: - assert isinstance(x0, Vector), "x0 must be a Vector or None" - assert x0.space == self.codomain, "x0 belongs to the wrong VectorSpace" - - def invalidate(self): - """Force the next `solve()` call to rebuild the factorization from `A`. - - Call this after actually changing `A` (in place, or via the `.linop` setter with - a numerically different operator) -- see the class docstring for why this is not - detected automatically. - """ - self._sparse_solver = None - - def _ensure_factorized(self): - if self._sparse_solver is None: - from feectools.linalg.direct_solvers import SparseSolver - - if self._parallel: - from feectools.linalg.utilities import FastAssemblyUnavailable, parallel_tosparse, tosparse_via_matvec - - # `parallel_tosparse` assembles in O(1) collective rounds (one per leaf - # operator) instead of `tosparse_via_matvec`'s O(A.domain.dimension) - # rounds (one per global DOF) -- a difference of several orders of - # magnitude for a field-solve-sized system (see its docstring for how). - # It only recognizes a subset of operator shapes, self-verified against - # `A`'s own `.dot()`; fall back to the always-correct (if much slower) - # sweep when it can't. - try: - mat = parallel_tosparse(self._A, self._comm, format="csr") - except FastAssemblyUnavailable: - mat = tosparse_via_matvec(self._A, format="csr") - else: - mat = self._A.tosparse().tocsr() - - # `A` can be exactly singular at essential-BC-masked DOFs: an operator - # built through a BoundaryOperator zero-masks both the input and output at - # those rows by design (struphy.feec.linear_operators.BoundaryOperator.dot, - # via apply_essential_bc_to_array) -- fine for an iterative solver, which - # never inverts A directly, as long as `b` is masked the same way (true for - # every caller here: e.g. ImplicitDiffusion.__call__ builds `rhs` via the - # same BoundaryOperator-wrapped `.dot()`, so `b` is already 0 at these rows - # too). A direct factorization needs those rows regularized to identity so - # `x = 1^{-1} * 0 = 0` comes out right there instead of `splu` raising - # "Factor is exactly singular" -- a zero row is unsolvable on its own even - # though the underlying (masked) system is perfectly well posed. - zero_rows = np.flatnonzero(mat.getnnz(axis=1) == 0) - if zero_rows.size: - mat = mat.tolil() - mat[zero_rows, zero_rows] = 1.0 - - self._sparse_solver = SparseSolver(mat.tocsc()) - - def solve(self, b, out=None): - """ - Solve A x = b exactly via the cached sparse LU factorization. - - Parameters - ---------- - b : feectools.linalg.stencil.StencilVector - Right-hand-side vector of the linear system. - - out : feectools.linalg.basic.Vector | NoneType - The output vector, or None (optional). - - Returns - ------- - x : feectools.linalg.basic.Vector - The exact (up to factorization round-off) solution of the linear system. - """ - assert isinstance(b, Vector) - assert b.space is self.domain - - self._ensure_factorized() - - # SparseSolver's factorization always lives on the host (scipy splu); the - # host round trip here is one flat vector of the field-solve's DOF count, not - # the particle arrays, so it is cheap relative to the iterations it replaces. - b_flat = xp.to_numpy(b.toarray()) - - if self._parallel: - # `b.toarray()` in parallel already returns the full global-shape array with - # only this rank's own (disjoint) entries filled in -- see - # `StencilVector._toarray_parallel_no_pads` -- so summing every rank's copy - # assembles the true global right-hand side. - if isinstance(self._comm, MockComm): - b_global = b_flat - else: - b_global = np.empty_like(b_flat) - self._comm.Allreduce(b_flat, b_global, op=MPI.SUM) - b_flat = b_global - - x_flat = np.empty_like(b_flat) - self._sparse_solver.solve(b_flat, out=x_flat) - - if self._parallel: - from feectools.linalg.utilities import array_to_psydac - - # x_flat should be a numpy array since SparseSolver's factorization - # is on host - x_vec = array_to_psydac(xp.asarray(x_flat), self.codomain) - if out is None: - out = x_vec - else: - assert isinstance(out, Vector) - assert out.space is self.codomain - x_vec.copy(out=out) - else: - if out is None: - out = self.codomain.zeros() - else: - assert isinstance(out, Vector) - assert out.space is self.codomain - - # Same local/no-pad interior slice StencilVector.toarray_local() reads from, - # see feectools.linalg.stencil.StencilVector.toarray_local. - idx = tuple( - slice(m * p, -m * p) if p != 0 else slice(0, None) - for p, m in zip(out.pads, out.space.shifts) - ) - out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) - - self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} - - if self._options.get("recycle") and self._options.get("x0") is not None: - out.copy(out=self._options["x0"]) - - return out - - def dot(self, b, out=None): - return self.solve(b, out=out) - #=============================================================================== class BiConjugateGradient(InverseLinearOperator): """ diff --git a/feectools/linalg/tests/test_solvers.py b/feectools/linalg/tests/test_solvers.py index 02bc39159..e877888f4 100644 --- a/feectools/linalg/tests/test_solvers.py +++ b/feectools/linalg/tests/test_solvers.py @@ -1,11 +1,10 @@ import cunumpy as xp import pytest -from feectools.linalg.solvers import inverse, DirectSolver +from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVectorSpace, StencilMatrix, StencilVector from feectools.linalg.basic import LinearSolver from feectools.ddm.cart import DomainDecomposition, CartDecomposition -from feectools.ddm.mpi import mpi as MPI def define_data_hermitian(n, p, dtype=float): @@ -205,64 +204,6 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False): assert errh_norm < tol assert solver == 'pcg' or errc_norm < tol -#=============================================================================== -def _compute_global_starts_ends(domain_decomposition, npts): - # Same as feectools.linalg.tests.test_block.compute_global_starts_ends. - global_starts = [None] * len(npts) - global_ends = [None] * len(npts) - for axis in range(len(npts)): - ee = domain_decomposition.global_element_ends[axis] - global_ends[axis] = ee.copy() - global_ends[axis][-1] = npts[axis] - 1 - global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) - return global_starts, global_ends - - -@pytest.mark.parametrize('n1', [8, 16]) -@pytest.mark.parametrize('n2', [8, 12]) -@pytest.mark.parametrize('p1', [1, 2]) -@pytest.mark.parallel -def test_direct_solver_parallel(n1, n2, p1, verbose=False): - """`DirectSolver` at nprocs > 1 must recover the exact solution, same as serial.""" - p2 = 1 - - comm = MPI.COMM_WORLD - D = DomainDecomposition([n1, n2], periods=[False, False], comm=comm) - npts = [n1, n2] - global_starts, global_ends = _compute_global_starts_ends(D, npts) - cart = CartDecomposition(D, npts, global_starts, global_ends, pads=[p1, p2], shifts=[1, 1]) - - V = StencilVectorSpace(cart, dtype=float) - A = StencilMatrix(V, V) - - # Diagonally dominant (hence nonsingular) stencil: -1 on every off-diagonal, enough - # on the main diagonal to dominate the row sum -- same style as define_data_hermitian - # above, extended to 2D. - n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 - for k1 in range(-p1, p1 + 1): - for k2 in range(-p2, p2 + 1): - A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 - A[:, :, 0, 0] = n_offdiag + 1.0 - A.remove_spurious_entries() - - s1, s2 = V.starts - e1, e2 = V.ends - xe = StencilVector(V) - for i1 in range(s1, e1 + 1): - for i2 in range(s2, e2 + 1): - xe[i1, i2] = xp.random.random() - xe.update_ghost_regions() - - be = A @ xe - - solv = DirectSolver(A) - x = solv.solve(be) - - err_norm = xp.linalg.norm((x - xe).toarray()) - if verbose: - print(f"n1={n1} n2={n2} p1={p1} p2={p2} nprocs={comm.Get_size()} err_norm={err_norm:.2e}") - assert err_norm < 1e-9 - # =============================================================================== # SCRIPT FUNCTIONALITY #=============================================================================== diff --git a/feectools/linalg/tests/test_utilities.py b/feectools/linalg/tests/test_utilities.py deleted file mode 100644 index 94e511c5b..000000000 --- a/feectools/linalg/tests/test_utilities.py +++ /dev/null @@ -1,428 +0,0 @@ -import cunumpy as xp -import numpy as np -import pytest -from scipy import sparse - -from feectools.ddm.cart import CartDecomposition, DomainDecomposition -from feectools.ddm.mpi import mpi as MPI -from feectools.feec.derivatives import DirectionalDerivativeOperator -from feectools.linalg.basic import ComposedLinearOperator, IdentityOperator, LinearOperator, ScaledLinearOperator, SumLinearOperator -from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace -from feectools.linalg.stencil import StencilMatrix, StencilVectorSpace -from feectools.linalg.utilities import ( - FastAssemblyUnavailable, - _get_entry, - _local_flat_entries, - _set_entry, - parallel_tosparse, - tosparse_via_matvec, -) - - -def compute_global_starts_ends(domain_decomposition, npts): - global_starts = [None] * len(npts) - global_ends = [None] * len(npts) - for axis in range(len(npts)): - ee = domain_decomposition.global_element_ends[axis] - global_ends[axis] = ee.copy() - global_ends[axis][-1] = npts[axis] - 1 - global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) - return global_starts, global_ends - - -def make_space(n1, n2, p1, p2, comm, periodic=True): - D = DomainDecomposition([n1, n2], periods=[periodic, False], comm=comm) - npts = [n1, n2] - gs, ge = compute_global_starts_ends(D, npts) - cart = CartDecomposition(D, npts, gs, ge, pads=[p1, p2], shifts=[1, 1]) - return StencilVectorSpace(cart, dtype=float) - - -def make_stencil_matrix(V, p1, p2, scale=1.0): - A = StencilMatrix(V, V) - n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 - for k1 in range(-p1, p1 + 1): - for k2 in range(-p2, p2 + 1): - A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 * scale - A[:, :, 0, 0] = (n_offdiag + 1.0) * scale - A.remove_spurious_entries() - return A - - -class FakeBoundaryOperator(LinearOperator): - """Mimics struphy.feec.linear_operators.BoundaryOperator well enough to exercise - `parallel_tosparse`'s diagonal-probe fallback: a diagonal 0/1 mask (zeroes the - first local DOF of each block on each rank), with a `.tosparse()` that is only - valid in serial (returns a small *locally*-indexed diagonal, exactly as struphy's - real implementation does -- see `struphy.feec.linear_operators. - BoundaryOperator.tosparse`) -- so the fast path must detect the mismatch via - validation and retry with the probe strategy instead of trusting `.tosparse()` - blindly. Works for both a plain StencilVectorSpace (V) and a BlockVectorSpace - (e.g. Hcurl, the codomain a real `grad` is wrapped in) domain, matching either - shape BoundaryOperator actually appears in. - """ - - def __init__(self, V): - self._V = V - self._entries = _local_flat_entries(V) - - @property - def domain(self): - return self._V - - @property - def codomain(self): - return self._V - - @property - def dtype(self): - return float - - def tosparse(self): - # Deliberately wrong at nprocs > 1: local indices, not global ones. - n = len(self._entries) - diag = np.ones(n) - diag[0] = 0.0 - return sparse.diags(diag, format="csr") - - def toarray(self): - return self.tosparse().toarray() - - def transpose(self, conjugate=False): - return self - - def dot(self, v, out=None): - if out is None: - out = self.codomain.zeros() - else: - out *= 0.0 - for k, (setter, _) in enumerate(self._entries): - _set_entry(out, setter, 0.0 if k == 0 else _get_entry(v, setter)) - out.update_ghost_regions() - return out - - -class UnsupportedOperator(LinearOperator): - """A leaf `parallel_tosparse` cannot possibly handle: no `.tosparse()`, and - domain is not codomain (so the diagonal-probe strategy doesn't apply either).""" - - def __init__(self, V_in, V_out): - self._Vin = V_in - self._Vout = V_out - - @property - def domain(self): - return self._Vin - - @property - def codomain(self): - return self._Vout - - @property - def dtype(self): - return float - - def tosparse(self): - raise NotImplementedError - - def toarray(self): - raise NotImplementedError - - def transpose(self, conjugate=False): - return UnsupportedOperator(self._Vout, self._Vin) - - def dot(self, v, out=None): - if out is None: - out = self.codomain.zeros() - return out - - -class FakeWeightedMassOperator(LinearOperator): - """Mimics struphy.feec.mass.WeightedMassOperator closely enough to exercise - `parallel_tosparse`'s duck-typed `._mat`-composition unwrap: wraps an inner - `._mat` behind identity extraction ops (trivial, as struphy's own serial - `.tosparse()` requires) but *non-trivial* boundary ops (masks the first and last - local DOF on each rank) -- i.e. `.dot()` is genuinely not the same as `._mat.dot()` - alone, which is exactly the bug this class was written to catch (found via a real - Struphy run, not anticipated up front -- see the git history of - `parallel_tosparse`'s `._mat`-unwrap branch). - """ - - def __init__(self, V, mat, mask_first=True, mask_last=True): - self._V = V - self._mat = mat - self._V_extraction_op = IdentityOperator(V) - self._W_extraction_op = IdentityOperator(V) - self._V_boundary_op = _MaskBoundaryOperator(V, mask_first, mask_last) - self._W_boundary_op = _MaskBoundaryOperator(V, mask_first, mask_last) - self._transposed = False - - @property - def domain(self): - return self._V - - @property - def codomain(self): - return self._V - - @property - def dtype(self): - return float - - def tosparse(self): - # Deliberately unusable at nprocs > 1, exactly like struphy's real - # WeightedMassOperator.tosparse() when boundary masking is actually active - # (it asserts outright there); here it just raises, which - # `parallel_tosparse` must also handle gracefully. - raise NotImplementedError - - def toarray(self): - raise NotImplementedError - - def transpose(self, conjugate=False): - raise NotImplementedError - - def dot(self, v, out=None): - tmp = self._V_boundary_op.transpose().dot(v) - tmp = self._V_extraction_op.transpose().dot(tmp) - tmp = self._mat.dot(tmp) - tmp = self._W_extraction_op.dot(tmp) - return self._W_boundary_op.dot(tmp, out=out) - - -class _MaskBoundaryOperator(LinearOperator): - """A minimal BoundaryOperator stand-in: zeroes the first and/or last local DOF on - each rank (self-adjoint, so `.transpose()` returns itself).""" - - def __init__(self, V, mask_first, mask_last): - self._V = V - self._entries = _local_flat_entries(V) - self._mask_first = mask_first - self._mask_last = mask_last - - @property - def domain(self): - return self._V - - @property - def codomain(self): - return self._V - - @property - def dtype(self): - return float - - def tosparse(self): - raise NotImplementedError - - def toarray(self): - raise NotImplementedError - - def transpose(self, conjugate=False): - return self - - def dot(self, v, out=None): - if out is None: - out = self.codomain.zeros() - else: - out *= 0.0 - n = len(self._entries) - for k, (setter, _) in enumerate(self._entries): - masked = (self._mask_first and k == 0) or (self._mask_last and k == n - 1) - _set_entry(out, setter, 0.0 if masked else _get_entry(v, setter)) - out.update_ghost_regions() - return out - - -@pytest.mark.parametrize('n1', [8, 16]) -@pytest.mark.parametrize('p1', [1, 2]) -@pytest.mark.parallel -def test_parallel_tosparse_matches_matvec_stencil_matrix(n1, p1, verbose=False): - """A plain StencilMatrix (no wrapper) must assemble identically via the fast - (`parallel_tosparse`) and slow (`tosparse_via_matvec`) paths.""" - n2, p2 = 8, 1 - comm = MPI.COMM_WORLD - V = make_space(n1, n2, p1, p2, comm) - A = make_stencil_matrix(V, p1, p2) - - fast = parallel_tosparse(A, comm) - slow = tosparse_via_matvec(A, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"n1={n1} p1={p1} nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-10 - - -@pytest.mark.parallel -def test_parallel_tosparse_composed_and_sum(verbose=False): - """A Sum-of-Scaled-and-Composed operator tree (the shape ImplicitDiffusion's - left-hand side actually takes: sigma*M + G^T @ D @ G) must also match the slow - reference path.""" - n1, n2, p1, p2 = 10, 8, 1, 1 - comm = MPI.COMM_WORLD - V = make_space(n1, n2, p1, p2, comm) - M = make_stencil_matrix(V, p1, p2, scale=1.0) - G = make_stencil_matrix(V, p1, p2, scale=0.5) - D = make_stencil_matrix(V, p1, p2, scale=2.0) - - composed = ComposedLinearOperator(V, V, G, D, G) - scaled = ScaledLinearOperator(V, V, c=3.0, A=M) - total = SumLinearOperator(V, V, scaled, composed) - - fast = parallel_tosparse(total, comm) - slow = tosparse_via_matvec(total, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-8 - - -@pytest.mark.parallel -def test_parallel_tosparse_block_vector_space(verbose=False): - """A BlockLinearOperator (grad's actual shape, e.g. H1 -> Hcurl's 3 stacked - components) must also assemble identically via the fast and slow paths.""" - n1, n2, p1, p2 = 8, 6, 1, 1 - comm = MPI.COMM_WORLD - V = make_space(n1, n2, p1, p2, comm, periodic=False) - W = BlockVectorSpace(V, V) - B = BlockLinearOperator(W, W) - B[0, 0] = make_stencil_matrix(V, p1, p2, scale=1.0) - B[1, 1] = make_stencil_matrix(V, p1, p2, scale=2.0) - B[0, 1] = make_stencil_matrix(V, p1, p2, scale=0.3) - - fast = parallel_tosparse(B, comm) - slow = tosparse_via_matvec(B, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-8 - - -@pytest.mark.parallel -def test_parallel_tosparse_diagonal_probe_fallback_block_vector_space(verbose=False): - """The same diagonal-mask-on-both-sides shape as - `test_parallel_tosparse_diagonal_probe_fallback`, but over a BlockVectorSpace -- - the actual shape `BoundaryOperator ∘ grad ∘ BoundaryOperator` takes in the real - Poisson benchmark (grad's codomain, Hcurl, has 3 stacked components).""" - n1, n2, p1, p2 = 8, 6, 1, 1 - comm = MPI.COMM_WORLD - V = make_space(n1, n2, p1, p2, comm, periodic=False) - W = BlockVectorSpace(V, V) - B = BlockLinearOperator(W, W) - B[0, 0] = make_stencil_matrix(V, p1, p2, scale=1.0) - B[1, 1] = make_stencil_matrix(V, p1, p2, scale=2.0) - mask = FakeBoundaryOperator(W) - - total = ComposedLinearOperator(W, W, mask, B, mask) - - fast = parallel_tosparse(total, comm) - slow = tosparse_via_matvec(total, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-8 - - -@pytest.mark.parallel -def test_parallel_tosparse_diagonal_probe_fallback(verbose=False): - """A BoundaryOperator-like diagonal mask, wired in on both sides of a - StencilMatrix (the actual shape struphy's BC-wrapped operators take), must be - correctly recovered by the diagonal-probe fallback -- not silently misassembled - from its serial-only `.tosparse()`.""" - n1, n2, p1, p2 = 8, 6, 1, 1 - comm = MPI.COMM_WORLD - V = make_space(n1, n2, p1, p2, comm, periodic=False) - A = make_stencil_matrix(V, p1, p2) - mask = FakeBoundaryOperator(V) - - total = ComposedLinearOperator(V, V, mask, A, mask) - - fast = parallel_tosparse(total, comm) - slow = tosparse_via_matvec(total, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-8 - - -@pytest.mark.parallel -def test_parallel_tosparse_raises_for_unsupported_operator(verbose=False): - """An operator this module genuinely cannot handle (no `.tosparse()`, not - diagonal-shaped) must raise `FastAssemblyUnavailable` -- identically on every - rank, so callers can fall back to `tosparse_via_matvec` without any risk of a - partial/divergent collective-call sequence.""" - n1, n2, p1, p2 = 8, 6, 1, 1 - comm = MPI.COMM_WORLD - Vin = make_space(n1, n2, p1, p2, comm, periodic=False) - Vout = make_space(n1, n2, p1, p2, comm, periodic=False) - op = UnsupportedOperator(Vin, Vout) - - with pytest.raises(FastAssemblyUnavailable): - parallel_tosparse(op, comm) - - -def make_deriv_space(n1, n2, p1, p2, comm, periodic): - # dim 0 always periodic (matches make_space's default); dim 1 (the - # differentiation direction in test_parallel_tosparse_directional_derivative) - # uses `periodic`, since that's the one whose periodicity actually changes the - # matrix structure (wraparound coupling vs. a one-fewer-point boundary). - D = DomainDecomposition([n1, n2], periods=[True, periodic], comm=comm) - npts = [n1, n2] - gs, ge = compute_global_starts_ends(D, npts) - cart = CartDecomposition(D, npts, gs, ge, pads=[p1, p2], shifts=[1, 1]) - return StencilVectorSpace(cart, dtype=float) - - -@pytest.mark.parametrize('diffdir_periodic', [False, True]) -@pytest.mark.parametrize('transposed', [False, True]) -@pytest.mark.parallel -def test_parallel_tosparse_directional_derivative(diffdir_periodic, transposed, verbose=False): - """`DirectionalDerivativeOperator`'s default `.tosparse()` asserts outright at - nprocs > 1 (see `_directional_derivative_triples`'s docstring) -- the closed-form - reconstruction it falls back to instead must match the slow reference path, for - both a periodic and a non-periodic differentiation direction, transposed or not. - """ - p1, p2 = 1, 1 - comm = MPI.COMM_WORLD - n1, n2 = 8, 8 - V = make_deriv_space(n1, n2, p1, p2, comm, periodic=diffdir_periodic) - W = make_deriv_space(n1, n2 if diffdir_periodic else n2 - 1, p1, p2, comm, periodic=diffdir_periodic) - - op = DirectionalDerivativeOperator(V, W, diffdir=1, negative=False, transposed=transposed) - - fast = parallel_tosparse(op, comm) - slow = tosparse_via_matvec(op, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"periodic={diffdir_periodic} transposed={transposed} nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-10 - - -@pytest.mark.parallel -def test_parallel_tosparse_weighted_mass_operator_boundary_composition(verbose=False): - """`FakeWeightedMassOperator` wraps a StencilMatrix behind trivial extraction ops - but *non-trivial* boundary masking (`.dot()` != `._mat.dot()` alone) -- exactly - the shape that caused a real, silent ~8% numeric mismatch against the naive - "just unwrap `._mat`" shortcut on an actual Struphy run (before - `parallel_tosparse`'s `._mat`-composition unwrap rebuilt the *whole* - boundary/extraction chain instead). Must match the slow reference path. - """ - n1, n2, p1, p2 = 8, 6, 1, 1 - comm = MPI.COMM_WORLD - V = make_space(n1, n2, p1, p2, comm, periodic=False) - inner = make_stencil_matrix(V, p1, p2) - op = FakeWeightedMassOperator(V, inner) - - fast = parallel_tosparse(op, comm) - slow = tosparse_via_matvec(op, format="csr") - - diff = abs(fast - slow).max() - if verbose: - print(f"nprocs={comm.Get_size()} diff={diff:.2e}") - assert diff < 1e-8 diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index f771424af..fa64a0802 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -1,52 +1,33 @@ # coding: utf-8 -import itertools from math import sqrt import cunumpy as xp -import numpy as np -from scipy import sparse -from feectools.ddm.mpi import MockComm -from feectools.ddm.mpi import mpi as MPI -from feectools.linalg.basic import Vector +from feectools.linalg.basic import Vector +from feectools.linalg.block import BlockVector, BlockVectorSpace from feectools.linalg.stencil import StencilVector, StencilVectorSpace -from feectools.linalg.block import BlockVector, BlockVectorSpace -from feectools.linalg.topetsc import petsc_local_to_psydac, get_npts_per_block +from feectools.linalg.topetsc import get_npts_per_block, petsc_local_to_psydac __all__ = ( 'array_to_psydac', 'petsc_to_psydac', - 'tosparse_via_matvec', - 'parallel_tosparse', - 'FastAssemblyUnavailable', '_sym_ortho', ) + #============================================================================== def array_to_psydac(x, V): - """ - Convert a NumPy array to a Vector of the space V. This function is designed to be the inverse of the method .toarray() of the class Vector. - Note: This function works in parallel but it is very costly and should be avoided if performance is a priority. - - Parameters - ---------- - x : numpy.ndarray - Array to be converted. It only contains the true data, the ghost regions must not be included. - - V : feectools.linalg.stencil.StencilVectorSpace or feectools.linalg.block.BlockVectorSpace - Space of the final Psydac Vector. - - Returns - ------- - u : feectools.linalg.stencil.StencilVector or feectools.linalg.block.BlockVector - Element of space V, the coefficients of which (excluding ghost regions) are the entries of x. The ghost regions of u are up to date. - + """ + Convert a NumPy array to a Vector of the space V. This function is designed to be + the inverse of the method .toarray() of the class Vector. + Note: This function works in parallel but it is very costly and should be avoided + if performance is a priority. """ assert x.ndim == 1, 'Array must be 1D.' - if x.dtype==complex: - assert V.dtype==complex, 'Complex array cannot be converted to a real StencilVector' + if x.dtype == complex: + assert V.dtype == complex, 'Complex array cannot be converted to a real StencilVector' assert x.size == V.dimension, 'Array must have the same global size as the space.' u = V.zeros() @@ -57,984 +38,44 @@ def array_to_psydac(x, V): def _array_to_psydac_recursive(x, u): - """ - Recursive function filling in the coefficients of each block of u. - """ + """Recursive function filling in the coefficients of each block of u.""" + assert isinstance(u, Vector) V = u.space assert x.ndim == 1, 'Array must be 1D.' - if x.dtype==complex: - assert V.dtype==complex, 'Complex array cannot be converted to a real StencilVector' - assert x.size == V.dimension, 'Array must have the same global size as the space.' + if x.dtype == complex: + assert V.dtype == complex, 'Complex array cannot be converted to a real StencilVector' + assert x.size == V.dimension, 'Array must have the same global size as the space.' if isinstance(V, BlockVectorSpace): for i, V_i in enumerate(V.spaces): x_i = x[:V_i.dimension] - x = x[V_i.dimension:] + x = x[V_i.dimension:] u_i = u[i] _array_to_psydac_recursive(x_i, u_i) elif isinstance(V, StencilVectorSpace): - index_global = tuple(slice(s, e+1) for s, e in zip(V.starts, V.ends)) + index_global = tuple(slice(s, e + 1) for s, e in zip(V.starts, V.ends)) u[index_global] = x.reshape(V.npts)[index_global] else: raise NotImplementedError(f'Can only handle StencilVector or BlockVector spaces, got {type(V)} instead') -#============================================================================== -def tosparse_via_matvec(op, format="csc"): - """ - Assemble the full global sparse matrix of a `LinearOperator` by applying it to every - global unit vector via `.dot()`, rather than via `.tosparse()`. - - Every operator's `.dot()` is already exercised (and therefore correct, including - cross-rank ghost/boundary coupling) every time it is actually used, unlike - `.tosparse()`, which several composed/derivative operators only implement correctly - in serial (see e.g. `feectools.feec.derivatives.DirectionalDerivativeOperator.tosparse`). - This is a port of `struphy.feec.linear_operators.LinOpWithTransp.toarray_struphy`'s - `is_sparse=True` branch into feectools (which `DirectSolver` -- the caller this exists - for -- must not import struphy from): same Allgather-starts/ends plus - unit-vector-`dot()` plus gather/broadcast-triplets algorithm, so every rank ends up - with an identical copy of the full global matrix (a "replicated" assembly, not a - distributed one -- deliberate, see `feectools.linalg.solvers.DirectSolver`). - - Cost: O(N) collective `.dot()` calls, N = `op.domain.dimension` -- does not shrink - with rank count (every call needs every rank's participation), so this is only - appropriate as a one-time, cached setup cost, not something to call every step. - - Parameters - ---------- - op : feectools.linalg.basic.LinearOperator - Operator to assemble. `op.domain`/`op.codomain` must each be a - `StencilVectorSpace` or `BlockVectorSpace`. - - format : str - scipy.sparse matrix format of the result ("csr", "csc", "coo", ...). - - Returns - ------- - out : scipy.sparse matrix - The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on - every rank. - """ - v = op.domain.zeros() - tmp2 = op.codomain.zeros() - - if isinstance(op.domain, BlockVectorSpace): - comm = op.domain.spaces[0].cart.comm - elif isinstance(op.domain, StencilVectorSpace): - comm = op.domain.cart.comm - else: - raise NotImplementedError( - f'tosparse_via_matvec only supports StencilVectorSpace/BlockVectorSpace domains, got {type(op.domain)}', - ) - - if comm is None or isinstance(comm, MockComm): - rank = 0 - size = 1 - else: - rank = comm.Get_rank() - size = comm.Get_size() - - numrows = op.codomain.dimension - numcols = op.domain.dimension - data, row, col = [], [], [] - - def local_nonzero_rows(stencil_vec, row_offset): - """(global_row_indices, values) for `stencil_vec`'s local interior data.""" - space = stencil_vec.space - idx_local = tuple( - slice(m * p, -m * p) if p != 0 else slice(0, None) - for p, m in zip(stencil_vec.pads, space.shifts) - ) - local_data = xp.to_numpy(stencil_vec._data[idx_local]) - nz = np.nonzero(local_data) - starts = space.starts - global_multi = tuple(nz[d] + int(starts[d]) for d in range(len(nz))) - rows = row_offset + np.ravel_multi_index(global_multi, space.npts) - return rows, local_data[nz] - - def codomain_local_nonzero_rows(vec): - """`local_nonzero_rows`, dispatched over `op.codomain`'s type.""" - if isinstance(op.codomain, BlockVectorSpace): - all_rows, all_vals = [], [] - row_offset = 0 - for b, sp in enumerate(op.codomain.spaces): - r, val = local_nonzero_rows(vec[b], row_offset) - all_rows.append(r) - all_vals.append(val) - row_offset += sp.dimension - return np.concatenate(all_rows), np.concatenate(all_vals) - return local_nonzero_rows(vec, 0) - - if isinstance(op.domain, BlockVectorSpace): - starts = [vi.starts for vi in v] - ends = [vi.ends for vi in v] - npts = [sp.npts for sp in op.domain.spaces] - nsp = len(op.domain.spaces) - ndim = [sp.ndim for sp in op.domain.spaces] - - # Plain NumPy throughout: this is tiny host-side index bookkeeping (rank - # starts/ends, a running column count), never device compute -- `xp.array` - # under the CuPy backend would produce 0-d CuPy scalars that `range()` (and - # plain Python int arithmetic below) cannot consume, the same class of - # NumPy-vs-CuPy scalar-typing trap documented for `AdhocTorus`/`xp.sqrt`. - startsarr = np.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) - allstarts = np.empty(size * len(startsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allstarts = startsarr - else: - comm.Allgather(startsarr, allstarts) - allstarts = allstarts.reshape((size, len(startsarr))) - - endsarr = np.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) - allends = np.empty(size * len(endsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allends = endsarr - else: - comm.Allgather(endsarr, allends) - allends = allends.reshape((size, len(endsarr))) - - for currentrank in range(size): - spoint = 0 - npredim = 0 - for h in range(nsp): - iterables = [ - range(int(allstarts[currentrank][i + npredim]), int(allends[currentrank][i + npredim]) + 1) - for i in range(ndim[h]) - ] - for i in itertools.product(*iterables): - if rank == currentrank: - v[h][i] = 1.0 - v[h].update_ghost_regions() - tmp2 *= 0.0 - op.dot(v, out=tmp2) - c = spoint + int(np.ravel_multi_index(i, npts[h])) - rs, vals = codomain_local_nonzero_rows(tmp2) - row.append(rs) - col.append(np.full(rs.shape, c)) - data.append(vals) - if rank == currentrank: - v[h][i] = 0.0 - v[h].update_ghost_regions() - cumulative = 1 - for i in range(ndim[h]): - cumulative *= npts[h][i] - spoint += cumulative - npredim += ndim[h] - - else: - starts = v.starts - ends = v.ends - npts = op.domain.npts - ndim = op.domain.ndim - - # Plain NumPy, same reasoning as the BlockVectorSpace branch above. - startsarr = np.array([starts[j] for j in range(ndim)], dtype=int) - allstarts = np.empty(size * len(startsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allstarts = startsarr - else: - comm.Allgather(startsarr, allstarts) - allstarts = allstarts.reshape((size, len(startsarr))) - - endsarr = np.array([ends[j] for j in range(ndim)], dtype=int) - allends = np.empty(size * len(endsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allends = endsarr - else: - comm.Allgather(endsarr, allends) - allends = allends.reshape((size, len(endsarr))) - - for currentrank in range(size): - iterables = [ - range(int(allstarts[currentrank][i]), int(allends[currentrank][i]) + 1) for i in range(ndim) - ] - for i in itertools.product(*iterables): - if rank == currentrank: - v[i] = 1.0 - v.update_ghost_regions() - op.dot(v, out=tmp2) - c = int(np.ravel_multi_index(i, npts)) - rs, vals = codomain_local_nonzero_rows(tmp2) - row.append(rs) - col.append(np.full(rs.shape, c)) - data.append(vals) - if rank == currentrank: - v[i] = 0.0 - - if comm is None or isinstance(comm, MockComm): - all_rows, all_cols, all_data = row, col, data - else: - gathered_rows = comm.gather(row, root=0) - gathered_cols = comm.gather(col, root=0) - gathered_data = comm.gather(data, root=0) - if rank == 0: - all_rows = [item for sublist in gathered_rows for item in sublist] - all_cols = [item for sublist in gathered_cols for item in sublist] - all_data = [item for sublist in gathered_data for item in sublist] - comm.bcast(all_rows, root=0) - comm.bcast(all_cols, root=0) - comm.bcast(all_data, root=0) - else: - all_rows = comm.bcast(None, root=0) - all_cols = comm.bcast(None, root=0) - all_data = comm.bcast(None, root=0) - - if all_rows: - all_rows = np.concatenate(all_rows) - all_cols = np.concatenate(all_cols) - all_data = np.concatenate(all_data) - else: - all_rows = all_cols = all_data = np.empty(0, dtype=int) - - mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) - return mat.asformat(format) - -#============================================================================== -class FastAssemblyUnavailable(Exception): - """Raised by `parallel_tosparse` when the operator tree could not be assembled via - the fast (O(1)-communication-round) path -- see its docstring. Callers should catch - this and fall back to `tosparse_via_matvec`.""" - - -def _local_flat_entries(V): - """ - List of (setter, global_flat_index) pairs, one per DOF *owned* by this rank (no - ghost/pad region), for a StencilVectorSpace or BlockVectorSpace V. - - `setter` is `('b', block_index, multi_index)` (usable as `vec[h][idx] = ...` for a - BlockVector) or `('s', multi_index)` (usable as `vec[idx] = ...` otherwise) -- - tagged rather than inferred from shape, since a StencilVectorSpace's own - `multi_index` can itself start with an int indistinguishable from a block index. - `global_flat_index` uses the same block-major, - `numpy.ravel_multi_index`-against-global-`npts` convention as - `tosparse_via_matvec` and `StencilMatrix.tosparse()` (`_tocoo_no_pads`) -- the same - one `Vector.toarray()` flattens to, which every caller of this module (e.g. - `DirectSolver.solve`'s `b.toarray()`/`x_flat.reshape(...)` round trip) already - relies on. All three MUST agree, since results from this function are combined - with plain `StencilMatrix`/`BlockLinearOperator` sparse matrices in the same - right-hand-side/solution vectors. - """ - if isinstance(V, BlockVectorSpace): - entries = [] - spoint = 0 - for h, sp in enumerate(V.spaces): - npts = sp.npts - iterables = [range(s, e + 1) for s, e in zip(sp.starts, sp.ends)] - for idx in itertools.product(*iterables): - flat = spoint + int(np.ravel_multi_index(idx, npts)) - entries.append((('b', h, idx), flat)) - spoint += int(np.prod(npts)) - return entries - elif isinstance(V, StencilVectorSpace): - npts = V.npts - iterables = [range(s, e + 1) for s, e in zip(V.starts, V.ends)] - return [(('s', idx), int(np.ravel_multi_index(idx, npts))) for idx in itertools.product(*iterables)] - else: - raise FastAssemblyUnavailable( - f'_local_flat_entries only supports StencilVectorSpace/BlockVectorSpace, got {type(V)}', - ) - - -def _set_entry(vec, setter, value): - if setter[0] == 'b': - _, h, idx = setter - vec[h][idx] = value - else: - _, idx = setter - vec[idx] = value - - -def _get_entry(vec, setter): - if setter[0] == 'b': - _, h, idx = setter - return vec[h][idx] - else: - _, idx = setter - return vec[idx] - - -def _replicate_triples(rows, cols, vals, shape, comm, dtype): - """Gather (rows, cols, vals) COO triples -- assumed *local* to this rank -- from - every rank and sum-combine them (matching duplicates, e.g. periodic wraparound, - exactly as `scipy.sparse.coo_matrix` does on `.tocsr()`) into one matrix identical - on every rank. One collective round, regardless of `shape`. - - `rows`/`cols`/`vals` may be plain sequences or arrays; always gathered and - concatenated as numpy arrays (`comm.allgather` pickles a numpy array through - mpi4py's out-of-band buffer protocol, and `np.concatenate` is vectorized) rather - than as Python lists -- for a leaf with real FEM bandwidth (tens to hundreds of - thousands of local nonzeros, e.g. a 3D mass matrix), converting through - element-by-element Python lists first was the dominant cost, dwarfing the O(1) - round-count win this function exists for. - """ - rows = np.asarray(rows, dtype=np.int64) - cols = np.asarray(cols, dtype=np.int64) - vals = np.asarray(vals, dtype=dtype) - - if comm is None or isinstance(comm, MockComm): - all_rows, all_cols, all_vals = rows, cols, vals - else: - gathered = comm.allgather((rows, cols, vals)) - all_rows = np.concatenate([g[0] for g in gathered]) - all_cols = np.concatenate([g[1] for g in gathered]) - all_vals = np.concatenate([g[2] for g in gathered]) - return sparse.coo_matrix((all_vals, (all_rows, all_cols)), shape=shape, dtype=dtype).tocsr() - - -def _probe_vector(V, entries, offset): - """A deterministic, reproducible-across-ranks probe Vector of space V: each owned - DOF gets a distinct nonzero value derived from its global flat index (never 0, and - never equal across two different `offset`s), so an operator's actual coupling - structure is very unlikely to accidentally look diagonal/masking by coincidence.""" - v = V.zeros() - for setter, flat in entries: - _set_entry(v, setter, 1.0 + 0.618033988749895 * ((flat + offset) % 104729)) - v.update_ghost_regions() - return v - - -def _validate_against_dot(node, candidate, comm, entries_domain, entries_codomain, seed): - """Check `candidate @ p == node.dot(p)` (this rank's owned output entries only) for - one probe vector `p`built from `_probe_vector`. `entries_domain` values already - carry each entry's global flat column index; `entries_codomain` likewise for rows. - Returns a local bool -- the caller combines these across ranks (see - `parallel_tosparse`) before trusting `candidate`.""" - V_domain = node.domain - p = _probe_vector(V_domain, entries_domain, seed) - p_flat_local = {flat: _get_entry(p, setter) for setter, flat in entries_domain} - - if comm is None or isinstance(comm, MockComm): - p_flat_full = dict(p_flat_local) - else: - gathered = comm.allgather(p_flat_local) - p_flat_full = {} - for d in gathered: - p_flat_full.update(d) - - p_full = np.zeros(candidate.shape[1], dtype=candidate.dtype) - for flat, val in p_flat_full.items(): - p_full[flat] = val - - q_candidate = candidate @ p_full - q_true = node.dot(p) - - # An aggregate (L2-norm) check, not a per-entry one: a wide-bandwidth FEM operator - # (e.g. a mass matrix with a degree-3 spline direction) sums many terms per row, - # in a different order than `candidate`'s (scipy's own summation order for the - # sparse matvec) -- individual output entries can then legitimately differ by much - # more than a tight per-entry relative tolerance even when `candidate` is exactly - # right, especially where terms partially cancel. Comparing the whole local output - # vector's norm to the whole error vector's norm is robust to that per-entry - # cancellation while still easily catching a genuinely wrong `candidate` (which - # differs at O(1) relative scale, not at rounding-error scale). - true_local = np.fromiter((_get_entry(q_true, setter) for setter, _ in entries_codomain), dtype=float) - cand_local = np.fromiter((q_candidate[flat] for _, flat in entries_codomain), dtype=float) - err = float(np.linalg.norm(true_local - cand_local)) - scale = float(np.linalg.norm(true_local)) - return err <= 1e-8 * scale + 1e-10 - - -def _directional_derivative_triples(op): - """Closed-form local (row, col, value) triples for a - `feectools.feec.derivatives.DirectionalDerivativeOperator` -- `.tosparse()`'s - default (no-pads) form isn't valid at nprocs > 1 for this operator (it asserts), - and its `with_pads=True` form returns a small *local* matrix in a totally - different (ghost-inclusive, non-globally-indexed) convention this module's - block-major global indexing can't reuse -- so this reconstructs the same bidiagonal - difference-operator matrix its serial `.tosparse()` builds, directly from the - operator's own definition (`out[i] = sign * (in[i + e_d] - in[i])` along direction - `d = op._diffdir`, `e_d` wrapped modulo `V.npts[d]` when periodic), one local - (globally-indexed) row at a time -- no basis-vector sweep, no padding subtleties. - Built in the "V -> W" (non-transposed) sense regardless of `op._transposed`; - `parallel_tosparse` transposes the result back if needed, exactly as the operator's - own serial `.tosparse()` does. - """ - V, W, d = op._spaceV, op._spaceW, op._diffdir - sign = -1.0 if op._negative else 1.0 - periodic = V.periods[d] - - if V.npts[d] == 1 and W.npts[d] == 1 and periodic: - return [], [], [] # degenerate single-cell-periodic case: the zero matrix - - rows, cols, vals = [], [], [] - for idx, row_flat in _local_flat_entries(W): - _, ii = idx - jj = ii - jj_next = list(ii) - jj_next[d] = (ii[d] + 1) % V.npts[d] if periodic else ii[d] + 1 - col_flat = int(np.ravel_multi_index(jj, V.npts)) - rows.append(row_flat) - cols.append(col_flat) - vals.append(-sign) - if periodic or jj_next[d] < V.npts[d]: - col_next_flat = int(np.ravel_multi_index(jj_next, V.npts)) - rows.append(row_flat) - cols.append(col_next_flat) - vals.append(sign) - return rows, cols, vals - - -def parallel_tosparse(op, comm, format="csr"): - """ - Assemble the full global sparse matrix of a `LinearOperator` tree using O(1) - collective-communication rounds (one per leaf node, roughly), instead of - `tosparse_via_matvec`'s O(`op.domain.dimension`) rounds (one basis vector per - global DOF) -- for the same replicated-on-every-rank result. - - Walks the operator tree using only types `feectools` itself defines - (`SumLinearOperator`, `ScaledLinearOperator`, `ComposedLinearOperator`, - `IdentityOperator`, `ZeroOperator`, `BlockLinearOperator`, - `DirectionalDerivativeOperator`): the first five compose exactly the way - `.tosparse()` already does in serial, just with the *leaves* below assembled - without a per-DOF basis-vector sweep; `BlockLinearOperator` is recursed into - block-by-block (not treated as one leaf) since a real Derham `grad`/`grad.T` can - have `DirectionalDerivativeOperator` blocks, whose own `.tosparse()` is unusable - in parallel (see `_directional_derivative_triples`, which reconstructs it in - closed form instead). For anything else -- most of it defined outside feectools - (`struphy.feec.mass.WeightedMassOperator`, `struphy.feec.linear_operators. - BoundaryOperator`, ...), which this module must not import -- one of three - O(1)-round strategies applies, tried in order: - - 1. Duck-typed unwrap: if the node has a `._mat` plus the same - `._V_extraction_op`/`._W_extraction_op`/`._V_boundary_op`/`._W_boundary_op`/ - `._transposed` attributes `struphy.feec.mass.WeightedMassOperator` has, - rebuild the exact composition its own `.dot()` applies (boundary and - extraction maps included, not just `._mat` alone -- an earlier version of - this function assumed trivial extraction ops meant `._mat` alone was enough, - which a real Struphy run showed is false whenever the boundary masks are - non-trivial) from parts each recursed into via `build()` in turn. - - 2. If the leaf has its own `.tosparse()` (true of `StencilMatrix` and - `BlockLinearOperator`): call it *locally* (no communication -- the same call - `A.tosparse()` already makes in serial, just once per rank instead of once - globally) and `allgather`-sum the local fragments into the replicated global - matrix. - - 3. Otherwise, if `leaf.domain is leaf.codomain` (a necessary condition to act as - a diagonal map): probe it with a value-tagged vector and check whether the - output is consistent with a per-DOF diagonal scaling (this is exactly what - essential-BC masking operators like `BoundaryOperator` are). Two `.dot()` - calls plus one `allgather`. - - Every leaf's result is *always* cross-checked against the operator's own `.dot()` - on a probe vector (strategies 1 and 2 both feed their candidate through the same - `_validate_against_dot` check that strategy 3 uses to detect diagonality in the - first place); an unrecognized leaf type (none of the three strategies applicable - or valid) is treated the same as a failed check. Whether any of this happens is a - pure function of operator *types*, identical on every rank by construction (same - model, same run) -- so every rank always issues the same sequence of collective - calls regardless of any individual check's pass/fail outcome; only *after* the - full tree is walked does one final `allreduce(MPI.LAND)` combine every check - across every rank into a single decision, so a data-dependent failure on one rank - cannot leave another rank waiting on a collective call that rank never issues (no - deadlock risk from divergent control flow). If that combined decision is False -- - or the tree contains a node type this function does not know how to handle at all - (e.g. `MatrixFreeLinearOperator`, always a deterministic, type-only decision, so - still consistent across ranks) -- `FastAssemblyUnavailable` is raised (on every - rank, identically) and the caller should fall back to `tosparse_via_matvec`. - - Parameters - ---------- - op : feectools.linalg.basic.LinearOperator - Operator to assemble. - - comm : MPI.Comm | feectools.ddm.mpi.MockComm | None - Communicator spanning every rank that owns a piece of `op`. - - format : str - scipy.sparse matrix format of the result. - - Returns - ------- - out : scipy.sparse matrix - The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on - every rank. - """ - # Imported here, not at module scope: these are feectools types this function - # checks via isinstance, kept local to make the "only feectools composite types - # are special-cased" contract easy to audit at a glance. - from feectools.linalg.basic import ComposedLinearOperator, IdentityOperator, ScaledLinearOperator, SumLinearOperator, ZeroOperator - from feectools.linalg.block import BlockLinearOperator - from feectools.feec.derivatives import DirectionalDerivativeOperator - - checks = [] - probe_seed = [1000003] # mutable cell; a fresh seed per leaf keeps probes independent - - def build(node): - if isinstance(node, ScaledLinearOperator): - return node._scalar * build(node._operator) - if isinstance(node, SumLinearOperator): - mats = [build(a) for a in node._addends] - out = mats[0] - for m in mats[1:]: - out = out + m - return out - if isinstance(node, ComposedLinearOperator): - mats = [build(m) for m in node._multiplicants] - out = mats[0] - for m in mats[1:]: - out = out @ m - return out - if isinstance(node, IdentityOperator): - return sparse.identity(node.domain.dimension, format="csr", dtype=node.dtype or float) - if isinstance(node, ZeroOperator): - return sparse.csr_matrix(node.shape, dtype=node.dtype or float) - if isinstance(node, BlockLinearOperator): - # Recurse into each block individually rather than calling - # `node.tosparse()` on the whole thing: a real Derham `grad`/`grad.T` is a - # BlockLinearOperator whose blocks can themselves be - # `DirectionalDerivativeOperator`s, whose *own* default `.tosparse()` - # asserts outright at nprocs > 1 (see `_directional_derivative_triples`) - # -- one such block would otherwise make the whole (possibly mostly - # StencilMatrix) BlockLinearOperator's `.tosparse()` raise. - checks_before = len(checks) - nrows, ncols = node.n_block_rows, node.n_block_cols - block_domain = (lambda j: node.domain[j]) if ncols > 1 else (lambda j: node.domain) - block_codomain = (lambda i: node.codomain[i]) if nrows > 1 else (lambda i: node.codomain) - grid = [[None for _ in range(ncols)] for _ in range(nrows)] - for i in range(nrows): - for j in range(ncols): - if (i, j) in node._blocks: - grid[i][j] = build(node._blocks[i, j]) - else: - grid[i][j] = sparse.csr_matrix((block_codomain(i).dimension, block_domain(j).dimension)) - candidate = sparse.bmat(grid, format="csr") - - children_ok = all(checks[checks_before:]) - if comm is not None and not isinstance(comm, MockComm): - children_ok = comm.allreduce(children_ok, op=MPI.LAND) - block_ok = False - if children_ok: - entries_domain = _local_flat_entries(node.domain) - entries_codomain = _local_flat_entries(node.codomain) - probe_seed[0] += 97 - block_ok = _validate_against_dot( - node, candidate, comm, entries_domain, entries_codomain, probe_seed[0], - ) - if comm is not None and not isinstance(comm, MockComm): - block_ok = comm.allreduce(block_ok, op=MPI.LAND) - if block_ok: - checks.append(True) - return candidate - - try: - fallback = tosparse_via_matvec(node, format="csr") - except Exception: - checks.append(False) - return candidate - - del checks[checks_before:] - checks.append(True) - return fallback - if isinstance(node, DirectionalDerivativeOperator): - # No generic strategy below applies (not diagonal-shaped in general, and - # its own `.tosparse()` is unusable here -- see - # `_directional_derivative_triples`); its structure is simple and fixed - # enough to reconstruct in closed form directly, still validated below - # like everything else. - V, W = node._spaceV, node._spaceW - rows, cols, vals = _directional_derivative_triples(node) - mat_vw = _replicate_triples(rows, cols, vals, (W.dimension, V.dimension), comm, node.dtype or float) - candidate = mat_vw.T.tocsr() if node._transposed else mat_vw - entries_domain = _local_flat_entries(node.domain) - entries_codomain = _local_flat_entries(node.codomain) - probe_seed[0] += 97 - ok = _validate_against_dot(node, candidate, comm, entries_domain, entries_codomain, probe_seed[0]) - checks.append(ok) - return candidate - - # Leaf: not one of the composite types above. Every strategy applicable to - # this leaf's *type* is always attempted, on every rank, regardless of any - # other rank's or strategy's data-dependent validation outcome -- see the - # docstring's "same collective calls on every rank" invariant. Only the first - # strategy that actually validates is kept. - entries_domain = _local_flat_entries(node.domain) - entries_codomain = _local_flat_entries(node.codomain) - shape = (node.codomain.dimension, node.domain.dimension) - dtype = node.dtype or float - probe_seed[0] += 97 - - # Duck-typed unwrap: struphy's `WeightedMassOperator` (M0, M1, ...) computes - # `V_boundary_op @ V_extraction_op @ _mat @ W_extraction_op.T @ W_boundary_op.T` - # (or the mirrored order when `._transposed`) on every `.dot()` call, by - # default with the boundary masks actually applied (`apply_bc=True`) -- *not* - # just `._mat` alone, even when both extraction ops are the identity (its own - # boundary masks can still be non-trivial, e.g. Dirichlet-BC-adjacent DOFs on - # a component of an Hcurl mass matrix -- discovered by this function's own - # validation rejecting the naive "just `._mat`" shortcut on exactly such a - # case, not by inspecting struphy's BC configuration). Rebuilding that same - # composition from its parts -- each recursed into via `build()`, so a - # boundary mask that itself needs the diagonal-probe strategy below still - # gets it -- is exact when every part is present (duck-typed by attribute, - # not `isinstance`, since none of these types live in feectools); still - # validated below regardless, as insurance against this composition itself - # being incomplete for some other struphy wrapper shaped differently. - parts = [getattr(node, name, "missing") for name in ( - "_mat", "_V_extraction_op", "_W_extraction_op", "_V_boundary_op", "_W_boundary_op", - )] - if "missing" not in parts: - inner_mat, v_ext, w_ext, v_bnd, w_bnd = parts - transposed = bool(getattr(node, "_transposed", False)) - try: - # Matches struphy.feec.mass.WeightedMassOperator.dot's own step - # sequence exactly (apply_bc=True, its default): non-transposed - # applies V_boundary_op.T, then V_extraction_op.T, then `._mat`, then - # W_extraction_op, then W_boundary_op, in that order (v -> out); the - # composed *matrix* is those same maps in reverse (rightmost applied - # first). `._transposed` mirrors V and W throughout. - if not transposed: - order = [w_bnd, w_ext, inner_mat, v_ext.transpose(), v_bnd.transpose()] - else: - order = [v_bnd, v_ext, inner_mat, w_ext.transpose(), w_bnd.transpose()] - mats = [build(m) for m in order] - candidate_inner = mats[0] - for m in mats[1:]: - candidate_inner = candidate_inner @ m - except Exception: - candidate_inner = None - if candidate_inner is not None and _validate_against_dot( - node, candidate_inner, comm, entries_domain, entries_codomain, probe_seed[0], - ): - checks.append(True) - return candidate_inner - - candidate_tosparse = None - try: - local_coo = node.tosparse().tocoo() - candidate_tosparse = _replicate_triples( - local_coo.row, local_coo.col, local_coo.data, - shape, comm, dtype, - ) - except Exception: - pass # this leaf's .tosparse() -- if it has one -- doesn't work here (e.g. - # raises, or -- as for struphy's BoundaryOperator -- succeeds but is - # documented serial-only and produces locally- rather than - # globally-indexed rows/cols at nprocs > 1); validation below (or, failing - # that, the diagonal-probe strategy) is what actually decides trust, not - # whether this call happened to raise. - - if candidate_tosparse is not None and _validate_against_dot( - node, candidate_tosparse, comm, entries_domain, entries_codomain, probe_seed[0], - ): - checks.append(True) - return candidate_tosparse - - if node.domain is not node.codomain: - # Not diagonal-shaped, and the .tosparse() attempt above (if any) didn't - # validate: nothing left to try for this leaf. - checks.append(False) - return candidate_tosparse if candidate_tosparse is not None else sparse.csr_matrix(shape, dtype=dtype) - - # Diagonal-probe strategy: two independent value-tagged probes; a true - # diagonal map reproduces (scaled by a per-DOF constant) or zeroes each one, - # consistently between the two -- exactly what an essential-BC mask does. - p1 = _probe_vector(node.domain, entries_domain, probe_seed[0]) - p2 = _probe_vector(node.domain, entries_domain, probe_seed[0] + 50000) - try: - o1 = node.dot(p1) - o2 = node.dot(p2) - except Exception: - checks.append(False) - return candidate_tosparse if candidate_tosparse is not None else sparse.csr_matrix(shape, dtype=dtype) - - rows, cols, vals = [], [], [] - ok = True - for setter, flat in entries_domain: - v1 = _get_entry(p1, setter) - v2 = _get_entry(p2, setter) - a1 = _get_entry(o1, setter) - a2 = _get_entry(o2, setter) - is_zero = abs(a1) < 1e-300 and abs(a2) < 1e-300 - if is_zero: - continue - d1 = a1 / v1 - d2 = a2 / v2 - if abs(d1 - d2) > 1e-8 * max(1.0, abs(d1)): - ok = False - break - rows.append(flat) - cols.append(flat) - vals.append(d1) - candidate_diag = _replicate_triples(rows, cols, vals, shape, comm, dtype) - ok = ok and _validate_against_dot( - node, candidate_diag, comm, entries_domain, entries_codomain, probe_seed[0], - ) - checks.append(ok) - return candidate_diag - - result = build(op) - - all_ok = all(checks) - if comm is not None and not isinstance(comm, MockComm): - all_ok = comm.allreduce(all_ok, op=MPI.LAND) - if not all_ok: - raise FastAssemblyUnavailable('one or more leaf operators could not be verified') - - return result.asformat(format) - -#============================================================================== -def tosparse_via_matvec(op, format="csc"): - """ - Assemble the full global sparse matrix of a `LinearOperator` by applying it to every - global unit vector via `.dot()`, rather than via `.tosparse()`. - - Every operator's `.dot()` is already exercised (and therefore correct, including - cross-rank ghost/boundary coupling) every time it is actually used, unlike - `.tosparse()`, which several composed/derivative operators only implement correctly - in serial (see e.g. `feectools.feec.derivatives.DirectionalDerivativeOperator.tosparse`). - This is a port of `struphy.feec.linear_operators.LinOpWithTransp.toarray_struphy`'s - `is_sparse=True` branch into feectools (which `DirectSolver` -- the caller this exists - for -- must not import struphy from): same Allgather-starts/ends plus - unit-vector-`dot()` plus gather/broadcast-triplets algorithm, so every rank ends up - with an identical copy of the full global matrix (a "replicated" assembly, not a - distributed one -- deliberate, see `feectools.linalg.solvers.DirectSolver`). - - Cost: O(N) collective `.dot()` calls, N = `op.domain.dimension` -- does not shrink - with rank count (every call needs every rank's participation), so this is only - appropriate as a one-time, cached setup cost, not something to call every step. - - Parameters - ---------- - op : feectools.linalg.basic.LinearOperator - Operator to assemble. `op.domain`/`op.codomain` must each be a - `StencilVectorSpace` or `BlockVectorSpace`. - - format : str - scipy.sparse matrix format of the result ("csr", "csc", "coo", ...). - - Returns - ------- - out : scipy.sparse matrix - The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on - every rank. - """ - v = op.domain.zeros() - tmp2 = op.codomain.zeros() - - if isinstance(op.domain, BlockVectorSpace): - comm = op.domain.spaces[0].cart.comm - elif isinstance(op.domain, StencilVectorSpace): - comm = op.domain.cart.comm - else: - raise NotImplementedError( - f'tosparse_via_matvec only supports StencilVectorSpace/BlockVectorSpace domains, got {type(op.domain)}', - ) - - if comm is None or isinstance(comm, MockComm): - rank = 0 - size = 1 - else: - rank = comm.Get_rank() - size = comm.Get_size() - - numrows = op.codomain.dimension - numcols = op.domain.dimension - data, row, col = [], [], [] - - def local_nonzero_rows(stencil_vec, row_offset): - """(global_row_indices, values) for `stencil_vec`'s LOCAL interior data only. - - Avoids `stencil_vec.toarray()`: under the parallel branch that allocates a - fresh, full-`codomain.dimension`-sized array and device->host-transfers it in - full, every single call -- the dominant cost under CuPy (a device alloc, a - device-side scatter-write kernel, and a full-size device->host copy per unit - vector, even though a stencil operator's column is actually sparse/local). - Reading only the local interior slice and adding `starts` to get the global row - index (same approach `StencilMatrix._tocoo_no_pads` already uses for columns) - transfers only the local, typically-mostly-zero data instead. - """ - space = stencil_vec.space - idx_local = tuple( - slice(m * p, -m * p) if p != 0 else slice(0, None) - for p, m in zip(stencil_vec.pads, space.shifts) - ) - local_data = xp.to_numpy(stencil_vec._data[idx_local]) - nz = np.nonzero(local_data) - starts = space.starts - global_multi = tuple(nz[d] + int(starts[d]) for d in range(len(nz))) - rows = row_offset + np.ravel_multi_index(global_multi, space.npts) - return rows, local_data[nz] - - def codomain_local_nonzero_rows(vec): - """`local_nonzero_rows`, dispatched over `op.codomain`'s type.""" - if isinstance(op.codomain, BlockVectorSpace): - all_rows, all_vals = [], [] - row_offset = 0 - for b, sp in enumerate(op.codomain.spaces): - r, val = local_nonzero_rows(vec[b], row_offset) - all_rows.append(r) - all_vals.append(val) - row_offset += sp.dimension - return np.concatenate(all_rows), np.concatenate(all_vals) - return local_nonzero_rows(vec, 0) - - if isinstance(op.domain, BlockVectorSpace): - starts = [vi.starts for vi in v] - ends = [vi.ends for vi in v] - npts = [sp.npts for sp in op.domain.spaces] - nsp = len(op.domain.spaces) - ndim = [sp.ndim for sp in op.domain.spaces] - - # Plain NumPy throughout: this is tiny host-side index bookkeeping (rank - # starts/ends, a running column count), never device compute -- `xp.array` - # under the CuPy backend would produce 0-d CuPy scalars that `range()` (and - # plain Python int arithmetic below) cannot consume, the same class of - # NumPy-vs-CuPy scalar-typing trap documented for `AdhocTorus`/`xp.sqrt`. - startsarr = np.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) - allstarts = np.empty(size * len(startsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allstarts = startsarr - else: - comm.Allgather(startsarr, allstarts) - allstarts = allstarts.reshape((size, len(startsarr))) - - endsarr = np.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) - allends = np.empty(size * len(endsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allends = endsarr - else: - comm.Allgather(endsarr, allends) - allends = allends.reshape((size, len(endsarr))) - - for currentrank in range(size): - spoint = 0 - npredim = 0 - for h in range(nsp): - iterables = [ - range(int(allstarts[currentrank][i + npredim]), int(allends[currentrank][i + npredim]) + 1) - for i in range(ndim[h]) - ] - for i in itertools.product(*iterables): - if rank == currentrank: - v[h][i] = 1.0 - v[h].update_ghost_regions() - tmp2 *= 0.0 - op.dot(v, out=tmp2) - c = spoint + int(np.ravel_multi_index(i, npts[h])) - rs, vals = codomain_local_nonzero_rows(tmp2) - row.append(rs) - col.append(np.full(rs.shape, c)) - data.append(vals) - if rank == currentrank: - # No `update_ghost_regions()` here: resetting this rank's own - # entry back to 0 only needs to be visible to neighbors before - # their *own* next `dot()` call, which is exactly what the - # `update_ghost_regions()` at the top of every iteration (run by - # every rank, every iteration, whether or not it owns that - # iteration's unit vector) already provides -- an extra call - # here would just be the same synchronization done twice. - v[h][i] = 0.0 - cumulative = 1 - for i in range(ndim[h]): - cumulative *= npts[h][i] - spoint += cumulative - npredim += ndim[h] - - else: - starts = v.starts - ends = v.ends - npts = op.domain.npts - ndim = op.domain.ndim - - # Plain NumPy, same reasoning as the BlockVectorSpace branch above. - startsarr = np.array([starts[j] for j in range(ndim)], dtype=int) - allstarts = np.empty(size * len(startsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allstarts = startsarr - else: - comm.Allgather(startsarr, allstarts) - allstarts = allstarts.reshape((size, len(startsarr))) - - endsarr = np.array([ends[j] for j in range(ndim)], dtype=int) - allends = np.empty(size * len(endsarr), dtype=int) - if comm is None or isinstance(comm, MockComm): - allends = endsarr - else: - comm.Allgather(endsarr, allends) - allends = allends.reshape((size, len(endsarr))) - - for currentrank in range(size): - iterables = [ - range(int(allstarts[currentrank][i]), int(allends[currentrank][i]) + 1) for i in range(ndim) - ] - for i in itertools.product(*iterables): - if rank == currentrank: - v[i] = 1.0 - v.update_ghost_regions() - op.dot(v, out=tmp2) - c = int(np.ravel_multi_index(i, npts)) - rs, vals = codomain_local_nonzero_rows(tmp2) - row.append(rs) - col.append(np.full(rs.shape, c)) - data.append(vals) - if rank == currentrank: - # See the matching comment in the BlockVectorSpace branch above -- - # no `update_ghost_regions()` needed here, the one at the top of the - # next iteration already covers it. - v[i] = 0.0 - - if comm is None or isinstance(comm, MockComm): - all_rows, all_cols, all_data = row, col, data - else: - gathered_rows = comm.gather(row, root=0) - gathered_cols = comm.gather(col, root=0) - gathered_data = comm.gather(data, root=0) - if rank == 0: - all_rows = [item for sublist in gathered_rows for item in sublist] - all_cols = [item for sublist in gathered_cols for item in sublist] - all_data = [item for sublist in gathered_data for item in sublist] - comm.bcast(all_rows, root=0) - comm.bcast(all_cols, root=0) - comm.bcast(all_data, root=0) - else: - all_rows = comm.bcast(None, root=0) - all_cols = comm.bcast(None, root=0) - all_data = comm.bcast(None, root=0) - - # `row`/`col`/`data` (and therefore `all_rows`/`all_cols`/`all_data`) are lists of - # small per-iteration arrays -- one nonzero-entries batch per unit vector, from - # `codomain_local_nonzero_rows` -- not lists of scalars, so concatenate before - # handing them to `coo_matrix`, which expects flat 1D array-likes. - if all_rows: - all_rows = np.concatenate(all_rows) - all_cols = np.concatenate(all_cols) - all_data = np.concatenate(all_data) - else: - all_rows = all_cols = all_data = np.empty(0, dtype=int) - - mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) - return mat.asformat(format) #============================================================================== def petsc_to_psydac(x, Xh, out=None): """ - Convert a PETSc.Vec object to a StencilVector or BlockVector. It assumes that PETSc was installed with the configuration for complex numbers. - Uses the index conversion functions in feectools.linalg.topetsc.py. - - Parameters - ---------- - x : PETSc.Vec - PETSc vector + Convert a PETSc.Vec object to a StencilVector or BlockVector. - Xh : feectools.linalg.stencil.StencilVectorSpace | feectools.linalg.block.BlockVectorSpace - Space of the coefficients of the Psydac vector. - - out : feectools.linalg.stencil.StencilVector | feectools.linalg.block.BlockVector, optional - The Psydac vector where to store the result. - - Returns - ------- - u : feectools.linalg.stencil.StencilVector | feectools.linalg.block.BlockVector - Psydac vector. In the case of a BlockVector, the blocks must be StencilVector. The general case is not yet implemented. + In the case of a BlockVector, the blocks must be StencilVector. The general case + is not yet implemented. """ - + if isinstance(Xh, BlockVectorSpace): - if any([isinstance(Xh.spaces[b], BlockVectorSpace) for b in range(len(Xh.spaces))]): + if any(isinstance(Xh.spaces[b], BlockVectorSpace) for b in range(len(Xh.spaces))): raise NotImplementedError('Block of blocks not implemented.') - + if out is not None: assert isinstance(out, BlockVector) assert out.space is Xh @@ -1042,28 +83,22 @@ def petsc_to_psydac(x, Xh, out=None): else: u = BlockVector(Xh) - comm = x.comm - dtype = Xh._dtype + comm = x.comm + dtype = Xh._dtype localsize, globalsize = x.getSizes() assert globalsize == u.shape[0], 'Sizes of global vectors do not match' - # Find shift for process k: - # ..get number of points for each block, each process and each dimension: - npts_local_per_block_per_process = xp.array(get_npts_per_block(Xh)) #indexed [b,k,d] for block b and process k and dimension d - # ..get local sizes for each block and each process: - local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) #indexed [b,k] for block b and process k - # ..sum the sizes over all the blocks and the previous processes: - index_shift = 0 + xp.sum(local_sizes_per_block_per_process[:,:comm.Get_rank()], dtype=int) #global variable + npts_local_per_block_per_process = xp.array(get_npts_per_block(Xh)) + local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) + index_shift = 0 + xp.sum(local_sizes_per_block_per_process[:, :comm.Get_rank()], dtype=int) for local_petsc_index in range(localsize): block_index, psydac_index = petsc_local_to_psydac(Xh, local_petsc_index) - # Get value of local PETSc vector passing the global PETSc index - value = x.getValue(local_petsc_index + index_shift) + value = x.getValue(local_petsc_index + index_shift) if value != 0: - u[block_index[0]]._data[psydac_index] = value if dtype is complex else value.real # PETSc always handles dtype specified in the installation configuration - - elif isinstance(Xh, StencilVectorSpace): + u[block_index[0]]._data[psydac_index] = value if dtype is complex else value.real + elif isinstance(Xh, StencilVectorSpace): if out is not None: assert isinstance(out, StencilVector) assert out.space is Xh @@ -1071,25 +106,20 @@ def petsc_to_psydac(x, Xh, out=None): else: u = StencilVector(Xh) - comm = x.comm - dtype = Xh.dtype + comm = x.comm + dtype = Xh.dtype localsize, globalsize = x.getSizes() assert globalsize == u.shape[0], 'Sizes of global vectors do not match' - # Find shift for process k: - # ..get number of points for each process and each dimension: - npts_local_per_block_per_process = xp.array(get_npts_per_block(Xh))[0] #indexed [k,d] for process k and dimension d - # ..get local sizes for each process: - local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) #indexed [k] for process k - # ..sum the sizes over all the previous processes: - index_shift = 0 + xp.sum(local_sizes_per_block_per_process[:comm.Get_rank()], dtype=int) #global variable + npts_local_per_block_per_process = xp.array(get_npts_per_block(Xh))[0] + local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) + index_shift = 0 + xp.sum(local_sizes_per_block_per_process[:comm.Get_rank()], dtype=int) for local_petsc_index in range(localsize): - block_index, psydac_index = petsc_local_to_psydac(Xh, local_petsc_index) - # Get value of local PETSc vector passing the global PETSc index + block_index, psydac_index = petsc_local_to_psydac(Xh, local_petsc_index) value = x.getValue(local_petsc_index + index_shift) if value != 0: - u._data[psydac_index] = value if dtype is complex else value.real # PETSc always handles dtype specified in the installation configuration + u._data[psydac_index] = value if dtype is complex else value.real else: raise ValueError('Xh must be a StencilVectorSpace or a BlockVectorSpace') @@ -1098,26 +128,16 @@ def petsc_to_psydac(x, Xh, out=None): return u + #============================================================================== def _sym_ortho(a, b): """ Stable implementation of Givens rotation. - This function was taken from the scipy repository - https://github.com/scipy/scipy/blob/master/scipy/sparse/linalg/isolve/lsqr.py - Notes - ----- - The routine 'SymOrtho' was added for numerical stability. This is - recommended by S.-C. Choi in [1]_. It removes the unpleasant potential of - ``1/eps`` in some important places (see, for example text following - "Compute the next plane rotation Qk" in minres.py). - - References - ---------- - .. [1] S.-C. Choi, "Iterative Methods for Singular Linear Equations - and Least-Squares Problems", Dissertation, - http://www.stanford.edu/group/SOL/dissertations/sou-cheng-choi-thesis.pdf + This function was taken from the scipy repository: + https://github.com/scipy/scipy/blob/master/scipy/sparse/linalg/isolve/lsqr.py """ + if b == 0: return _scalar_sign(a), 0, abs(a) elif a == 0: @@ -1129,17 +149,19 @@ def _sym_ortho(a, b): r = b / s else: tau = b / a - c = _scalar_sign(a) / sqrt(1+tau*tau) + c = _scalar_sign(a) / sqrt(1 + tau * tau) s = c * tau r = a / c return c, s, r + #============================================================================== def _scalar_sign(x): """ Sign of a real Python scalar. `xp.sign` (array_api_compat) requires its argument to expose a `.dtype` attribute, which plain Python floats don't have. """ + if x > 0: return 1.0 elif x < 0: From bb2c218cd62e508c46495d0125b9b5ff5eb2d9af Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 21 Aug 2026 13:56:01 +0200 Subject: [PATCH 22/23] Resolve all unit test failures with cupy --- conftest.py | 15 +++++ 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 | 10 ++-- feectools/fem/tests/analytical_profiles_1d.py | 2 + feectools/fem/tests/utilities.py | 4 ++ feectools/linalg/direct_solvers.py | 52 ++++++++++++----- feectools/linalg/fft.py | 9 ++- feectools/linalg/kron.py | 11 ++-- feectools/linalg/solvers.py | 12 ++-- feectools/linalg/sparse.py | 13 ++++- feectools/linalg/stencil.py | 8 ++- 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 +++-- 17 files changed, 170 insertions(+), 79 deletions(-) diff --git a/conftest.py b/conftest.py index 6efc25b37..3cf9a8d58 100644 --- a/conftest.py +++ b/conftest.py @@ -1,6 +1,8 @@ """Root-level pytest configuration.""" import pytest import sys +import importlib.util +import os from pathlib import Path @@ -32,6 +34,11 @@ 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 + cupy_without_mpi = ( + os.environ.get("ARRAY_BACKEND", "").lower() == "cupy" + and not os.environ.get("FEECTOOLS_ENABLE_MPI") + ) for item in items: # Skip if module is in skip list @@ -39,6 +46,14 @@ 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")) + + # CuPy deliberately disables MPI unless explicitly enabled; cart + # exchanger tests require a real Cartesian MPI communicator. + if cupy_without_mpi and item.fspath.basename in {"test_cart_2d.py", "test_cart_3d.py"}: + item.add_marker(pytest.mark.skip(reason="CuPy MPI is disabled; set FEECTOOLS_ENABLE_MPI=1")) + # 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/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 af4c4013c..0dd751d31 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -517,6 +517,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 @@ -557,11 +562,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/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 8d02a6b21..d05a4e176 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -139,9 +139,22 @@ 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. Keep the public solver backend-agnostic by + # staging a device right-hand side on the host and returning the + # solution on the caller's backend. + if xp.is_gpu(rhs): + rhs_cpu = xp.to_numpy(rhs) + preout, self._sinfo = self._solver_function( + self._bmat, self._l, self._u, rhs_cpu.T, self._ipiv, + trans=transposed, + ) + out = xp.asarray(preout.T) + else: + preout, self._sinfo = self._solver_function( + self._bmat, self._l, self._u, rhs.T, self._ipiv, + trans=transposed, + ) + out = preout.T else: assert out.shape == rhs.shape @@ -153,17 +166,21 @@ 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, - trans=transposed) + # We want FORTRAN-contiguous data (default is assumed to be C + # contiguous). The LAPACK factorization is host-side regardless + # of the globally selected backend. + 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, - trans=transposed) - out.set(out_cpu) + _, self._sinfo = self._solver_function( + self._bmat, self._l, self._u, out.T, self._ipiv, + overwrite_b=True, trans=transposed, + ) return out @@ -224,7 +241,14 @@ 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 + if xp.is_gpu(rhs): + rhs_cpu = xp.to_numpy(rhs) + result_cpu = self._splu.solve( + rhs_cpu.T, trans='T' if transposed else 'N' + ).T + out = xp.asarray(result_cpu) + else: + out = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T else: assert out.shape == rhs.shape 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 833e3fcbe..cf413ad75 100644 --- a/feectools/linalg/kron.py +++ b/feectools/linalg/kron.py @@ -114,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 @@ -152,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): @@ -190,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() @@ -213,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): diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 56ca82743..21047b09f 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -1902,7 +1902,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] @@ -1910,9 +1912,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 b19d08ea3..132081f42 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -1406,8 +1406,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) @@ -2411,7 +2412,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() 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 From d793c1693edb586f8553e4d64d9e14f9218103f5 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 21 Aug 2026 14:08:09 +0200 Subject: [PATCH 23/23] Run feectools with MPI and cuda --- conftest.py | 10 ---------- feectools/ddm/mpi.py | 24 ++++++++---------------- feectools/ddm/tests/test_cart_2d.py | 7 ++++--- feectools/ddm/tests/test_cart_3d.py | 10 ++++++---- 4 files changed, 18 insertions(+), 33 deletions(-) diff --git a/conftest.py b/conftest.py index 3cf9a8d58..c7c2f8d42 100644 --- a/conftest.py +++ b/conftest.py @@ -2,7 +2,6 @@ import pytest import sys import importlib.util -import os from pathlib import Path @@ -35,10 +34,6 @@ 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 - cupy_without_mpi = ( - os.environ.get("ARRAY_BACKEND", "").lower() == "cupy" - and not os.environ.get("FEECTOOLS_ENABLE_MPI") - ) for item in items: # Skip if module is in skip list @@ -49,11 +44,6 @@ def pytest_collection_modifyitems(config, items): if item.get_closest_marker("petsc") and not petsc_available: item.add_marker(pytest.mark.skip(reason="petsc4py is not installed")) - # CuPy deliberately disables MPI unless explicitly enabled; cart - # exchanger tests require a real Cartesian MPI communicator. - if cupy_without_mpi and item.fspath.basename in {"test_cart_2d.py", "test_cart_3d.py"}: - item.add_marker(pytest.mark.skip(reason="CuPy MPI is disabled; set FEECTOOLS_ENABLE_MPI=1")) - # 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/ddm/mpi.py b/feectools/ddm/mpi.py index fb3a13ee9..de9301c53 100644 --- a/feectools/ddm/mpi.py +++ b/feectools/ddm/mpi.py @@ -90,26 +90,18 @@ def _enabled(name, default=False): try: - # MPI is off by default on the CuPy backend, and on by default otherwise. + # CuPy arrays implement ``__cuda_array_interface__``, which mpi4py can + # pass to CUDA-aware MPI implementations. The DDM exchangers synchronize + # the current CUDA stream before every MPI operation, and ranks bind to a + # node-local GPU, so MPI is supported on the CuPy backend just as it is on + # NumPy. An MPI implementation must of course have CUDA support for + # device-buffer communication. # - # It is no longer *incorrect* to combine the two -- the reductions in - # feectools.linalg stage their (tiny) buffers through the host, the ghost - # exchangers synchronize the device before handing it a buffer, and each - # rank binds to its own GPU. It is, however, still slow: a ghost exchange - # of device memory through MPI derived datatypes costs milliseconds, so a - # single-GPU run pays several times over for communication it does not - # need. Until that is addressed, opt in explicitly: - # - # FEECTOOLS_ENABLE_MPI=1 use MPI on the CuPy backend - # FEECTOOLS_DISABLE_MPI=1 force the serial path on any backend + # FEECTOOLS_DISABLE_MPI=1 remains available to force the serial path on + # any backend. if _enabled('FEECTOOLS_DISABLE_MPI'): raise ImportError('MPI disabled by FEECTOOLS_DISABLE_MPI') - if os.environ.get('ARRAY_BACKEND', '').lower() == 'cupy' \ - and not _enabled('FEECTOOLS_ENABLE_MPI'): - raise ImportError('MPI off by default on the CuPy backend; ' - 'set FEECTOOLS_ENABLE_MPI=1 to use it') - from mpi4py import MPI _comm = MPI.COMM_WORLD 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<=i2