diff --git a/feectools/core/bsplines.py b/feectools/core/bsplines.py index 189d95f0e..22fcfcbf5 100644 --- a/feectools/core/bsplines.py +++ b/feectools/core/bsplines.py @@ -16,6 +16,7 @@ """ import cunumpy as xp +from cunumpy import PyccelKernel from cunumpy.xp import array_backend import numpy as np @@ -38,6 +39,27 @@ cell_index_p, basis_ders_on_irregular_grid_p) +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +find_span_p = PyccelKernel(find_span_p) +find_spans_p = PyccelKernel(find_spans_p) +basis_funs_p = PyccelKernel(basis_funs_p) +basis_funs_array_p = PyccelKernel(basis_funs_array_p) +basis_funs_1st_der_p = PyccelKernel(basis_funs_1st_der_p) +basis_funs_all_ders_p = PyccelKernel(basis_funs_all_ders_p) +collocation_matrix_p = PyccelKernel(collocation_matrix_p) +histopolation_matrix_p = PyccelKernel(histopolation_matrix_p) +greville_p = PyccelKernel(greville_p) +breakpoints_p = PyccelKernel(breakpoints_p) +elements_spans_p = PyccelKernel(elements_spans_p) +make_knots_p = PyccelKernel(make_knots_p) +elevate_knots_p = PyccelKernel(elevate_knots_p) +quadrature_grid_p = PyccelKernel(quadrature_grid_p) +basis_ders_on_quad_grid_p = PyccelKernel(basis_ders_on_quad_grid_p) +basis_integrals_p = PyccelKernel(basis_integrals_p) +cell_index_p = PyccelKernel(cell_index_p) +basis_ders_on_irregular_grid_p = PyccelKernel(basis_ders_on_irregular_grid_p) + __all__ = ('find_span', 'find_spans', 'basis_funs', @@ -84,7 +106,7 @@ def find_span(knots, degree, x): Knot span index. """ x = float(x) - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) return find_span_p(knots, degree, x) #============================================================================== @@ -116,8 +138,8 @@ def find_spans(knots, degree, x, out=None): spans : array of ints Knots span indexes. """ - knots = xp.ascontiguousarray(knots, dtype=float) - x = xp.ascontiguousarray(x, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + x = xp.ascontiguousarray(xp.asarray(x), dtype=float) if out is None: out = xp.zeros_like(x, dtype=int) else: @@ -155,7 +177,7 @@ def basis_funs(knots, degree, x, span, out=None): 1D array containing the values of ``degree + 1`` non-zero Bsplines at location ``x``. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) # Get native float x = float(x) if out is None: @@ -193,8 +215,8 @@ def basis_funs_array(knots, degree, span, x, out=None): 2D array of shape ``(len(x), degree + 1)`` containing the values of ``degree + 1`` non-zero Bsplines at each location in ``x``. """ - knots = xp.ascontiguousarray(knots, dtype=float) - x = xp.ascontiguousarray(x, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + x = xp.ascontiguousarray(xp.asarray(x), dtype=float) if out is None: out = xp.zeros(x.shape + (degree + 1,), dtype=float) else: @@ -240,7 +262,7 @@ def basis_funs_1st_der(knots, degree, x, span, out=None): ---------- .. [2] SELALIB, Semi-Lagrangian Library. http://selalib.gforge.inria.fr """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) # Get native float to work on windows x = float(x) if out is None: @@ -291,7 +313,7 @@ def basis_funs_all_ders(knots, degree, x, span, n, normalization='B', out=None): ders[i,j] = (d/dx)^i B_k(x) with k=(span-degree+j), for 0 <= i <= n and 0 <= j <= degree+1. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) # Get native float to work on windows x = float(x) if out is None: @@ -346,8 +368,8 @@ def collocation_matrix(knots, degree, periodic, normalization, xgrid, out=None, if xgrid.size == 1: return xp.ones((1, 1), dtype=float) - knots = xp.ascontiguousarray(knots, dtype=float) - xgrid = xp.ascontiguousarray(xgrid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + xgrid = xp.ascontiguousarray(xp.asarray(xgrid), dtype=float) if out is None: nb = len(knots) - degree - 1 if periodic: @@ -430,8 +452,8 @@ def histopolation_matrix(knots, degree, periodic, normalization, xgrid, multipli if not xp.all(xp.diff(xgrid) > 0): raise ValueError("Grid points must be ordered, with no repetitions: {}".format(xgrid)) - knots = xp.ascontiguousarray(knots, dtype=float) - xgrid = xp.ascontiguousarray(xgrid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + xgrid = xp.ascontiguousarray(xp.asarray(xgrid), dtype=float) elevated_knots = elevate_knots(knots, degree, periodic, multiplicity=multiplicity) normalization = normalization == "M" @@ -477,7 +499,7 @@ def breakpoints(knots, degree, tol=1e-15, out=None): breaks : numpy.ndarray (1D) Abscissas of all breakpoints. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = xp.zeros(len(knots), dtype=float) else: @@ -572,7 +594,7 @@ def elements_spans(knots, degree, out=None): spans = xp.searchsorted( knots, breaks[:-1], side='right' ) - 1 """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = np.zeros(len(knots), dtype=xp.int64) else: @@ -848,8 +870,8 @@ def basis_ders_on_quad_grid(knots, degree, quad_grid, nders, normalization, offs """ offset = int(offset) ne, nq = quad_grid.shape - knots = xp.ascontiguousarray(knots, dtype=float) - quad_grid = xp.ascontiguousarray(quad_grid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + quad_grid = xp.ascontiguousarray(xp.asarray(quad_grid), dtype=float) if out is None: out = xp.zeros((ne, degree + 1, nders + 1, nq), dtype=float) else: @@ -892,7 +914,7 @@ def basis_integrals(knots, degree, out=None): to (len(knots)-degree-1). In the periodic case the last (degree) values in the array are redundant, as they are a copy of the first (degree) values. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = xp.zeros(len(knots) - degree - 1, dtype=float) else: @@ -934,8 +956,8 @@ def cell_index(breaks, i_grid, tol=1e-15, out=None): ``cell_index[i]`` is the index of the cell in which ``i_grid[i]`` belong. """ - breaks = xp.ascontiguousarray(breaks, dtype=float) - i_grid = xp.ascontiguousarray(i_grid, dtype=float) + breaks = xp.ascontiguousarray(xp.asarray(breaks), dtype=float) + i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float) if out is None: out = np.zeros_like(i_grid, dtype=xp.int64) else: @@ -990,8 +1012,8 @@ def basis_ders_on_irregular_grid(knots, degree, i_grid, cell_index, nders, norma . il: local basis function (0 <= il <= degree) . id: derivative (0 <= id <= nders ) """ - knots = xp.ascontiguousarray(knots, dtype=float) - i_grid = xp.ascontiguousarray(i_grid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float) if out is None: nx = i_grid.shape[0] out = xp.zeros((nx, degree + 1, nders + 1), dtype=float) 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 2b2b58b41..837220c4e 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -2,18 +2,17 @@ import os import numpy as np -import cunumpy as xp -from cunumpy.xp import array_backend +import numpy as xp # this module is host-only MPI/index bookkeeping, never device data from itertools import product -# Initialize CUDA context before MPI if using CuPy backend -if array_backend.backend == "cupy": - try: - import cupy as cp - cp.cuda.Device(0).use() - cp.cuda.Stream.null.synchronize() - except Exception: - pass +from cunumpy.xp import array_backend, to_numpy + +# Initialize the CUDA context before MPI if using CuPy backend, binding this +# rank to its own GPU. Must stay above the feectools.ddm.mpi import, which +# initialises MPI as a side effect. +from feectools.ddm.device import bind_local_device + +bind_local_device() from feectools.ddm.mpi import mpi as MPI from feectools.ddm.mpi import MockMPI @@ -482,6 +481,12 @@ class CartDecomposition(): """ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads, shifts ): + # global_starts/global_ends are host-side decomposition metadata; callers + # may hand them in as CuPy arrays (e.g. built with cunumpy under the CuPy + # backend), so coerce them to NumPy up front. + global_starts = [ to_numpy(gs) for gs in global_starts ] + global_ends = [ to_numpy(ge) for ge in global_ends ] + # Check input arguments # TODO: check that arguments are identical across all processes assert len( npts ) == len( global_starts ) == len( global_ends ) == len( pads ) == len(shifts) 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/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..159e2d8e2 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 @@ -169,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 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/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 513635c52..f5b9ed2eb 100644 --- a/feectools/fem/tensor.py +++ b/feectools/fem/tensor.py @@ -44,6 +44,23 @@ eval_fields_3d_weighted, eval_fields_3d_irregular_weighted) +from cunumpy import PyccelKernel + +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +eval_fields_1d_no_weights = PyccelKernel(eval_fields_1d_no_weights) +eval_fields_1d_irregular_no_weights = PyccelKernel(eval_fields_1d_irregular_no_weights) +eval_fields_1d_weighted = PyccelKernel(eval_fields_1d_weighted) +eval_fields_1d_irregular_weighted = PyccelKernel(eval_fields_1d_irregular_weighted) +eval_fields_2d_no_weights = PyccelKernel(eval_fields_2d_no_weights) +eval_fields_2d_irregular_no_weights = PyccelKernel(eval_fields_2d_irregular_no_weights) +eval_fields_2d_weighted = PyccelKernel(eval_fields_2d_weighted) +eval_fields_2d_irregular_weighted = PyccelKernel(eval_fields_2d_irregular_weighted) +eval_fields_3d_no_weights = PyccelKernel(eval_fields_3d_no_weights) +eval_fields_3d_irregular_no_weights = PyccelKernel(eval_fields_3d_irregular_no_weights) +eval_fields_3d_weighted = PyccelKernel(eval_fields_3d_weighted) +eval_fields_3d_irregular_weighted = PyccelKernel(eval_fields_3d_irregular_weighted) + __all__ = ('TensorFemSpace',) #=============================================================================== @@ -500,7 +517,7 @@ def eval_fields(self, grid, *fields, weights=None, npts_per_cell=None, overlap=0 # -> grid is tensor-product, but npts_per_cell is not the same in each cell elif grid[0].ndim == 1 and npts_per_cell is None: out_fields = self.eval_fields_irregular_tensor_grid(grid, *fields, weights=weights, overlap=overlap) - return [xp.ascontiguousarray(out_fields[..., i]) for i in range(len(fields))] + return [xp.ascontiguousarray(xp.asarray(out_fields[..., i])) for i in range(len(fields))] # Case 3. 1D arrays of coordinates and npts_per_cell is a tuple or an integer # -> grid is tensor-product, and each cell has the same number of evaluation points @@ -512,7 +529,7 @@ def eval_fields(self, grid, *fields, weights=None, npts_per_cell=None, overlap=0 grid[i] = xp.reshape(grid[i], (ncells_i, npts_per_cell[i])) out_fields = self.eval_fields_regular_tensor_grid(grid, *fields, weights=weights, overlap=overlap) # return a list - return [xp.ascontiguousarray(out_fields[..., i]) for i in range(len(fields))] + return [xp.ascontiguousarray(xp.asarray(out_fields[..., i])) for i in range(len(fields))] # Case 4. (self.ldim)D arrays of coordinates and no npts_per_cell # -> unstructured grid diff --git a/feectools/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/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/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 4848595c3..8e76376a9 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -3,21 +3,27 @@ # 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 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 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 @@ -50,6 +56,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), @@ -62,6 +72,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,16 +117,18 @@ 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): +class StencilVectorSpace(ReductionWorkspace, VectorSpace): """ Vector space for n-dimensional stencil format. Two different initializations are possible: @@ -189,6 +219,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 @@ -214,7 +256,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 @@ -266,23 +308,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._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 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 ) - return x._dot_recv_data[0] - else: - return inner_func(*inner_args) + 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): @@ -315,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 @@ -477,8 +616,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 @@ -1099,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) @@ -1114,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 @@ -1126,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): """ @@ -1687,7 +1876,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 ) @@ -1764,8 +1953,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 ) @@ -1854,7 +2045,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 +2199,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 +2223,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 +2358,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 +3118,7 @@ def _tocoo_no_pads(self): M = coo_matrix( (data,(rows,cols)), - shape = [xp.prod(nr),xp.prod(nc)], + shape = [math.prod(nr),math.prod(nc)], dtype = self.domain.dtype) return M diff --git a/feectools/linalg/tests/test_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'])) 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'])) 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']))