diff --git a/conftest.py b/conftest.py index 6efc25b37..c7c2f8d42 100644 --- a/conftest.py +++ b/conftest.py @@ -1,6 +1,7 @@ """Root-level pytest configuration.""" import pytest import sys +import importlib.util from pathlib import Path @@ -32,6 +33,7 @@ def pytest_collection_modifyitems(config, items): items_to_remove = [] skip = pytest.mark.skip(reason="Requires optional dependency (sympde)") + petsc_available = importlib.util.find_spec("petsc4py") is not None for item in items: # Skip if module is in skip list @@ -39,6 +41,9 @@ def pytest_collection_modifyitems(config, items): items_to_remove.append(item) continue + if item.get_closest_marker("petsc") and not petsc_available: + item.add_marker(pytest.mark.skip(reason="petsc4py is not installed")) + # If running with xdist, automatically skip mpi and petsc tests if config.pluginmanager.has_plugin("xdist"): if item.get_closest_marker("mpi") or item.get_closest_marker("petsc"): diff --git a/feectools/core/bsplines.py b/feectools/core/bsplines.py index 189d95f0e..2ec9d7cd5 100644 --- a/feectools/core/bsplines.py +++ b/feectools/core/bsplines.py @@ -16,6 +16,7 @@ """ import cunumpy as xp +from cunumpy import PyccelKernel from cunumpy.xp import array_backend import numpy as np @@ -38,6 +39,27 @@ cell_index_p, basis_ders_on_irregular_grid_p) +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +find_span_p = PyccelKernel(find_span_p) +find_spans_p = PyccelKernel(find_spans_p) +basis_funs_p = PyccelKernel(basis_funs_p) +basis_funs_array_p = PyccelKernel(basis_funs_array_p) +basis_funs_1st_der_p = PyccelKernel(basis_funs_1st_der_p) +basis_funs_all_ders_p = PyccelKernel(basis_funs_all_ders_p) +collocation_matrix_p = PyccelKernel(collocation_matrix_p) +histopolation_matrix_p = PyccelKernel(histopolation_matrix_p) +greville_p = PyccelKernel(greville_p) +breakpoints_p = PyccelKernel(breakpoints_p) +elements_spans_p = PyccelKernel(elements_spans_p) +make_knots_p = PyccelKernel(make_knots_p) +elevate_knots_p = PyccelKernel(elevate_knots_p) +quadrature_grid_p = PyccelKernel(quadrature_grid_p) +basis_ders_on_quad_grid_p = PyccelKernel(basis_ders_on_quad_grid_p) +basis_integrals_p = PyccelKernel(basis_integrals_p) +cell_index_p = PyccelKernel(cell_index_p) +basis_ders_on_irregular_grid_p = PyccelKernel(basis_ders_on_irregular_grid_p) + __all__ = ('find_span', 'find_spans', 'basis_funs', @@ -84,7 +106,7 @@ def find_span(knots, degree, x): Knot span index. """ x = float(x) - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) return find_span_p(knots, degree, x) #============================================================================== @@ -116,8 +138,8 @@ def find_spans(knots, degree, x, out=None): spans : array of ints Knots span indexes. """ - knots = xp.ascontiguousarray(knots, dtype=float) - x = xp.ascontiguousarray(x, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + x = xp.ascontiguousarray(xp.asarray(x), dtype=float) if out is None: out = xp.zeros_like(x, dtype=int) else: @@ -155,7 +177,7 @@ def basis_funs(knots, degree, x, span, out=None): 1D array containing the values of ``degree + 1`` non-zero Bsplines at location ``x``. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) # Get native float x = float(x) if out is None: @@ -193,8 +215,8 @@ def basis_funs_array(knots, degree, span, x, out=None): 2D array of shape ``(len(x), degree + 1)`` containing the values of ``degree + 1`` non-zero Bsplines at each location in ``x``. """ - knots = xp.ascontiguousarray(knots, dtype=float) - x = xp.ascontiguousarray(x, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + x = xp.ascontiguousarray(xp.asarray(x), dtype=float) if out is None: out = xp.zeros(x.shape + (degree + 1,), dtype=float) else: @@ -240,7 +262,7 @@ def basis_funs_1st_der(knots, degree, x, span, out=None): ---------- .. [2] SELALIB, Semi-Lagrangian Library. http://selalib.gforge.inria.fr """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) # Get native float to work on windows x = float(x) if out is None: @@ -291,7 +313,7 @@ def basis_funs_all_ders(knots, degree, x, span, n, normalization='B', out=None): ders[i,j] = (d/dx)^i B_k(x) with k=(span-degree+j), for 0 <= i <= n and 0 <= j <= degree+1. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) # Get native float to work on windows x = float(x) if out is None: @@ -346,8 +368,8 @@ def collocation_matrix(knots, degree, periodic, normalization, xgrid, out=None, if xgrid.size == 1: return xp.ones((1, 1), dtype=float) - knots = xp.ascontiguousarray(knots, dtype=float) - xgrid = xp.ascontiguousarray(xgrid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + xgrid = xp.ascontiguousarray(xp.asarray(xgrid), dtype=float) if out is None: nb = len(knots) - degree - 1 if periodic: @@ -430,8 +452,8 @@ def histopolation_matrix(knots, degree, periodic, normalization, xgrid, multipli if not xp.all(xp.diff(xgrid) > 0): raise ValueError("Grid points must be ordered, with no repetitions: {}".format(xgrid)) - knots = xp.ascontiguousarray(knots, dtype=float) - xgrid = xp.ascontiguousarray(xgrid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + xgrid = xp.ascontiguousarray(xp.asarray(xgrid), dtype=float) elevated_knots = elevate_knots(knots, degree, periodic, multiplicity=multiplicity) normalization = normalization == "M" @@ -477,7 +499,7 @@ def breakpoints(knots, degree, tol=1e-15, out=None): breaks : numpy.ndarray (1D) Abscissas of all breakpoints. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = xp.zeros(len(knots), dtype=float) else: @@ -518,8 +540,7 @@ def greville(knots, degree, periodic, out=None, multiplicity=1): # Greville points are index arrays, keep on NumPy if isinstance(knots, (list, tuple)): knots = np.asarray(knots, dtype=float) - if hasattr(knots, 'get'): - knots = knots.get() # Convert CuPy to NumPy + knots = xp.to_numpy(knots) knots = np.ascontiguousarray(knots, dtype=float) if out is None: n = len(knots) - 2 * degree - 2 + multiplicity if periodic else len(knots) - degree - 1 @@ -572,7 +593,7 @@ def elements_spans(knots, degree, out=None): spans = xp.searchsorted( knots, breaks[:-1], side='right' ) - 1 """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = np.zeros(len(knots), dtype=xp.int64) else: @@ -624,7 +645,7 @@ def make_knots(breaks, degree, periodic, multiplicity=1, out=None): # Consistency checks assert len(breaks) > 1 # Convert to numpy for comparison since assertion needs Python bool - breaks_np = breaks.get() if hasattr(breaks, 'get') else breaks + breaks_np = xp.to_numpy(breaks) if isinstance(breaks_np, (list, tuple)): breaks_np = np.asarray(breaks_np) assert all( np.diff(breaks_np) > 0 ) @@ -638,8 +659,7 @@ def make_knots(breaks, degree, periodic, multiplicity=1, out=None): # Keep breaks on NumPy for initialization - knots are index arrays needed for CPU operations breaks = np.asarray(breaks, dtype=float) if isinstance(breaks, (list, tuple)) else breaks - if hasattr(breaks, 'get'): - breaks = breaks.get() # Convert CuPy to NumPy + breaks = xp.to_numpy(breaks) breaks = np.ascontiguousarray(breaks, dtype=float) if out is None: # Knots are index arrays, keep them on NumPy @@ -693,8 +713,7 @@ def elevate_knots(knots, degree, periodic, multiplicity=1, tol=1e-15, out=None): multiplicity = int(multiplicity) if isinstance(knots, (list, tuple)): knots = np.asarray(knots, dtype=float) - if hasattr(knots, 'get'): - knots = knots.get() # Convert CuPy to NumPy + knots = xp.to_numpy(knots) knots = np.ascontiguousarray(knots, dtype=float) if out is None: if periodic: @@ -771,14 +790,13 @@ def quadrature_grid(breaks, quad_rule_x, quad_rule_w): assert max(quad_rule_x) <= +1 # Convert breaks to numpy if CuPy (breaks/grids should stay on CPU) - if hasattr(breaks, 'get'): - breaks = breaks.get() + breaks = xp.to_numpy(breaks) breaks = np.ascontiguousarray(breaks, dtype=float) if array_backend.backend == "cupy": # Convert CuPy arrays to NumPy - quad_rule_x = quad_rule_x.get() if hasattr(quad_rule_x, 'get') else quad_rule_x - quad_rule_w = quad_rule_w.get() if hasattr(quad_rule_w, 'get') else quad_rule_w + quad_rule_x = xp.to_numpy(quad_rule_x) + quad_rule_w = xp.to_numpy(quad_rule_w) quad_rule_x = np.ascontiguousarray(quad_rule_x, dtype=float) quad_rule_w = np.ascontiguousarray(quad_rule_w, dtype=float) @@ -848,8 +866,8 @@ def basis_ders_on_quad_grid(knots, degree, quad_grid, nders, normalization, offs """ offset = int(offset) ne, nq = quad_grid.shape - knots = xp.ascontiguousarray(knots, dtype=float) - quad_grid = xp.ascontiguousarray(quad_grid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + quad_grid = xp.ascontiguousarray(xp.asarray(quad_grid), dtype=float) if out is None: out = xp.zeros((ne, degree + 1, nders + 1, nq), dtype=float) else: @@ -892,7 +910,7 @@ def basis_integrals(knots, degree, out=None): to (len(knots)-degree-1). In the periodic case the last (degree) values in the array are redundant, as they are a copy of the first (degree) values. """ - knots = xp.ascontiguousarray(knots, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) if out is None: out = xp.zeros(len(knots) - degree - 1, dtype=float) else: @@ -934,8 +952,8 @@ def cell_index(breaks, i_grid, tol=1e-15, out=None): ``cell_index[i]`` is the index of the cell in which ``i_grid[i]`` belong. """ - breaks = xp.ascontiguousarray(breaks, dtype=float) - i_grid = xp.ascontiguousarray(i_grid, dtype=float) + breaks = xp.ascontiguousarray(xp.asarray(breaks), dtype=float) + i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float) if out is None: out = np.zeros_like(i_grid, dtype=xp.int64) else: @@ -990,8 +1008,8 @@ def basis_ders_on_irregular_grid(knots, degree, i_grid, cell_index, nders, norma . il: local basis function (0 <= il <= degree) . id: derivative (0 <= id <= nders ) """ - knots = xp.ascontiguousarray(knots, dtype=float) - i_grid = xp.ascontiguousarray(i_grid, dtype=float) + knots = xp.ascontiguousarray(xp.asarray(knots), dtype=float) + i_grid = xp.ascontiguousarray(xp.asarray(i_grid), dtype=float) if out is None: nx = i_grid.shape[0] out = xp.zeros((nx, degree + 1, nders + 1), dtype=float) diff --git a/feectools/core/tests/test_bsplines.py b/feectools/core/tests/test_bsplines.py index 5ec11e777..92a1431e2 100644 --- a/feectools/core/tests/test_bsplines.py +++ b/feectools/core/tests/test_bsplines.py @@ -159,7 +159,7 @@ def test_histopolation_matrix(lims, nc, p, periodic, tol=1e-13): def test_cell_index(i_grid, expected): breaks = xp.array([0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.]) out = cell_index(breaks, xp.asarray(i_grid)) - assert xp.array_equal(expected, out) + assert xp.array_equal(xp.asarray(expected), out) #============================================================================== # SCRIPT FUNCTIONALITY: PLOT BASIS FUNCTIONS diff --git a/feectools/core/tests/test_bsplines_kernel.py b/feectools/core/tests/test_bsplines_kernel.py index 441057eb3..ef6e148a4 100644 --- a/feectools/core/tests/test_bsplines_kernel.py +++ b/feectools/core/tests/test_bsplines_kernel.py @@ -2,14 +2,16 @@ import pytest import cunumpy as xp +import numpy as np from feectools.core.bsplines_kernels import cell_index_p def test_cell_index_p(): - breaks = xp.array([0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.]) - breaks = xp.ascontiguousarray(breaks, dtype=float) - out = xp.zeros_like(breaks, dtype=xp.int64) + # This directly tests the raw Pyccel kernel, which intentionally accepts + # NumPy host arrays only; CuPy coverage belongs to the public wrapper. + breaks = np.ascontiguousarray([0. , 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.], dtype=float) + out = np.zeros_like(breaks, dtype=np.int64) tol = 1e-15 # limit case: code should decide wether point is in or out, not fall in infinite loop @@ -26,15 +28,14 @@ def test_cell_index_p(): assert status == expected_status # checking that the values match those of searchsorted (-1) for arbitrary grid points - i_grid = xp.array([0.14320482, 0.86569833, 0.77775327, 0.00895956, 0.074629 , + i_grid = np.array([0.14320482, 0.86569833, 0.77775327, 0.00895956, 0.074629 , 0.45682646, 0.5384352 , 0.20915311, 0.73121977, 0.01057414, 0.33756086, 0.17839759, 0.14023414, 0.09846206, 0.79970392, 0.65330406, 0.82716552, 0.24185731, 0.24054685, 0.72466651, 0.69125033, 0.3136558 , 0.64794089, 0.47975527, 0.99802844, 0.64402598, 0.41263526, 0.28178414, 0.57274384, 0.73218562]) - out = xp.zeros_like(i_grid, dtype=xp.int64) + out = np.zeros_like(i_grid, dtype=np.int64) status = cell_index_p(breaks, i_grid, tol, out) assert status == 0 - nps = xp.searchsorted(breaks, i_grid)-1 - assert xp.allclose(out, nps) - + nps = np.searchsorted(breaks, i_grid)-1 + assert np.allclose(out, nps) diff --git a/feectools/core/tests/test_bsplines_pyccel.py b/feectools/core/tests/test_bsplines_pyccel.py index e39bd370e..4be129e6d 100644 --- a/feectools/core/tests/test_bsplines_pyccel.py +++ b/feectools/core/tests/test_bsplines_pyccel.py @@ -187,8 +187,9 @@ def basis_funs_all_ders_true(knots, degree, x, span, n, normalization='B'): # Normalization to get M-Splines if normalization == 'M': - ders *= [(degree + 1) / (knots[i + degree + 1] - knots[i]) \ - for i in range(span - degree, span + 1)] + scaling = xp.asarray([(degree + 1) / (knots[i + degree + 1] - knots[i]) + for i in range(span - degree, span + 1)]) + ders *= scaling return ders #============================================================================== @@ -221,7 +222,15 @@ def collocation_matrix_true(knots, degree, periodic, normalization, xgrid): for i,x in enumerate( xgrid ): span = find_span_true( knots, degree, x ) basis = basis_funs_true( knots, degree, x, span ) - mat[i,js(span)] = normalize(basis, span) + values = normalize(basis, span) + if periodic: + # NumPy and CuPy differ for indexed assignment with repeated + # indices (which occurs when nb <= degree). The production + # kernel assigns in loop order, so make the reference explicit. + for j, value in zip(js(span), values): + mat[i, j] = value + else: + mat[i, js(span)] = values # Mitigate round-off errors mat[abs(mat) < 1e-14] = 0.0 @@ -293,7 +302,7 @@ def histopolation_matrix_true(knots, degree, periodic, normalization, xgrid): # Compute span for each row (index of last non-zero basis function) # TODO: would be better to have this ready beforehand # TODO: use tolerance instead of comparing against zero - spans = [(row != 0).argmax() + (degree+1) for row in C] + spans = [int((row != 0).argmax()) + (degree+1) for row in C] # Compute histopolation matrix from collocation matrix of higher degree m = C.shape[0] - 1 diff --git a/feectools/ddm/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..0dd751d31 100644 --- a/feectools/ddm/cart.py +++ b/feectools/ddm/cart.py @@ -1,19 +1,20 @@ # coding: utf-8 import os +import cunumpy # only for its backend-agnostic to_numpy(), see below -- not aliased to + # xp here, since that alias is reserved for plain NumPy in this module. import numpy as np -import cunumpy as xp -from cunumpy.xp import array_backend +import numpy as xp # this module is host-only MPI/index bookkeeping, never device data from itertools import product -# 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 +483,12 @@ class CartDecomposition(): """ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads, shifts ): + # global_starts/global_ends are host-side decomposition metadata; callers + # may hand them in as CuPy arrays (e.g. built with cunumpy under the CuPy + # backend), so coerce them to NumPy up front. + global_starts = [ to_numpy(gs) for gs in global_starts ] + global_ends = [ to_numpy(ge) for ge in global_ends ] + # Check input arguments # TODO: check that arguments are identical across all processes assert len( npts ) == len( global_starts ) == len( global_ends ) == len( pads ) == len(shifts) @@ -494,8 +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 @@ -510,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 @@ -522,7 +534,11 @@ def __init__( self, domain_decomposition, npts, global_starts, global_ends, pads # Know my coordinates in the topology self._coords = domain_decomposition.coords # Convert coords to NumPy for indexing (MPI coords should be on CPU) - coords_np = [c.get() if hasattr(c, 'get') else c for c in self._coords] + # cunumpy.to_numpy, not used here: self._coords may hold plain Python ints + # (mpi4py's Get_coords returns a plain list), and to_numpy would wrap those + # into 0-d NumPy arrays via np.asarray -- the wrong type to index a tuple of + # global_starts/ends with below. is_gpu leaves non-CuPy values untouched. + coords_np = [c.get() if cunumpy.is_gpu(c) else c for c in self._coords] # Start/end values of global indices (without ghost regions) self._starts = tuple( self._global_starts[axis][c] for axis,c in zip(range(self._ndims), coords_np) ) @@ -546,11 +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/ddm/device.py b/feectools/ddm/device.py new file mode 100644 index 000000000..b5cc715bc --- /dev/null +++ b/feectools/ddm/device.py @@ -0,0 +1,106 @@ +""" +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 + +import cunumpy as xp +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(xp.is_gpu(a) 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..de9301c53 100644 --- a/feectools/ddm/mpi.py +++ b/feectools/ddm/mpi.py @@ -80,12 +80,28 @@ 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") - + # 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. + # + # 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') + from mpi4py import MPI _comm = MPI.COMM_WORLD @@ -93,7 +109,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/ddm/tests/test_cart_2d.py b/feectools/ddm/tests/test_cart_2d.py index 66ec3fba1..d9ae40853 100644 --- a/feectools/ddm/tests/test_cart_2d.py +++ b/feectools/ddm/tests/test_cart_2d.py @@ -92,7 +92,9 @@ def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis= #--------------------------------------------------------------------------- # Fill in true domain with u[i1_loc,i2_loc,:]=[i1_glob,i2_glob] - u[p1:-p1,p2:-p2,:] = [[(i1,i2) for i2 in range(s2,e2+1)] for i1 in range(s1,e1+1)] + u[p1:-p1,p2:-p2,:] = xp.asarray( + [[(i1, i2) for i2 in range(s2, e2 + 1)] for i1 in range(s1, e1 + 1)] + ) request = synchronizer.prepare_communications(u) @@ -109,7 +111,7 @@ def run_cart_2d( data_exchanger_type, verbose=False , nprocs=None, reverse_axis= val = lambda i1,i2: (i1%n1,i2) if 0<=i2cupy conversion, + # cupy->numpy needs an explicit .get()/xp.to_numpy()) -- these + # diagonals are tiny and only ever feed this one-time host-side + # sparse assembly, never a device computation. + import numpy as np + + maindiag = np.ones(domain_local) * (-sign) + adddiag = np.ones(domain_local) * sign # handle special case with not self.domain.parallel and not with_pads and periodic if self.domain.periods[d] and not self.domain.parallel and not with_pads: # then: add element to other side of the array - adddiagcirc = xp.array([sign]) + adddiagcirc = np.array([sign]) offsets = (-codomain_local+1, 0, 1) diags = (adddiagcirc, maindiag, adddiag) else: diff --git a/feectools/feec/global_geometric_projectors.py b/feectools/feec/global_geometric_projectors.py index 3dc7752be..8ddf73b14 100644 --- a/feectools/feec/global_geometric_projectors.py +++ b/feectools/feec/global_geometric_projectors.py @@ -12,6 +12,20 @@ from feectools.fem.basic import FemField from feectools.feec import dof_kernels +from cunumpy import PyccelKernel + +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +for _name in ( + 'evaluate_dofs_1d_0form', 'evaluate_dofs_1d_1form', + 'evaluate_dofs_2d_0form', 'evaluate_dofs_2d_1form_hcurl', 'evaluate_dofs_2d_1form_hdiv', + 'evaluate_dofs_2d_2form', 'evaluate_dofs_2d_vec', + 'evaluate_dofs_3d_0form', 'evaluate_dofs_3d_1form', 'evaluate_dofs_3d_2form', + 'evaluate_dofs_3d_3form', 'evaluate_dofs_3d_vec', +): + setattr(dof_kernels, _name, PyccelKernel(getattr(dof_kernels, _name))) +del _name + from feectools.fem.tensor import TensorFemSpace from feectools.fem.vector import VectorFemSpace, MultipatchFemSpace @@ -29,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] @@ -169,7 +178,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 @@ -816,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 @@ -838,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 @@ -864,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 @@ -894,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 @@ -927,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 @@ -952,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 @@ -978,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 @@ -1008,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 @@ -1043,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 @@ -1084,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 @@ -1112,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 @@ -1141,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 4df8bae0b..51987884b 100644 --- a/feectools/fem/partitioning.py +++ b/feectools/fem/partitioning.py @@ -1,8 +1,10 @@ # -*- coding: UTF-8 -*- import os +import cunumpy # only for its backend-agnostic to_numpy(), see below -- not aliased to + # xp here, since that alias is reserved for plain NumPy in this module. import numpy as np -import cunumpy as xp +import numpy as xp # this module is host-only MPI/index bookkeeping, never device data from feectools.ddm.cart import CartDecomposition, InterfaceCartDecomposition, create_interfaces_cart from feectools.core.bsplines import elements_spans @@ -58,9 +60,8 @@ def partition_coefficients(domain_decomposition, spaces, min_blocks=None): m = multiplicity[axis] # Convert to numpy if CuPy (needed for MPI operations later) - if hasattr(ee, 'get'): - ee = ee.get() - + ee = cunumpy.to_numpy(ee) + global_ends [axis] = m*(ee+1)-1 global_ends [axis][-1] = npts[axis]-1 global_starts[axis] = xp.array([0] + (global_ends[axis][:-1]+1).tolist()) @@ -69,8 +70,8 @@ def partition_coefficients(domain_decomposition, spaces, min_blocks=None): min_blocks = [None] * ndims for s, e, V, mb in zip(global_starts, global_ends, spaces, min_blocks): - s_host = s.get() if hasattr(s, 'get') else np.asarray(s) - e_host = e.get() if hasattr(e, 'get') else np.asarray(e) + s_host = cunumpy.to_numpy(s) + e_host = cunumpy.to_numpy(e) local_sizes = e_host - s_host + 1 if V.periodic or mb is None: diff --git a/feectools/fem/splines.py b/feectools/fem/splines.py index ac62f9fa7..b292cfa2c 100644 --- a/feectools/fem/splines.py +++ b/feectools/fem/splines.py @@ -189,19 +189,13 @@ def init_interpolation( self, dtype=float ): # Convert to CSC format and compute sparse LU decomposition # Convert to LAPACK banded format (see DGBTRF function) - if hasattr(imat, 'get'): - imat = imat.get() - else: - imat = _np.asanyarray(imat) + imat = xp.to_numpy(imat) self._interpolator = SparseSolver( csc_matrix( imat ) ) else: # Convert to LAPACK banded format (see DGBTRF function) - if array_backend.backend == "cupy": - imat = imat.get() - else: - imat = _np.asanyarray(imat) + imat = xp.to_numpy(imat) dmat = dia_matrix( imat ) l = abs( dmat.offsets.min() ) u = dmat.offsets.max() @@ -231,10 +225,7 @@ def init_histopolation( self, dtype=float): xgrid = self.ext_greville, multiplicity = self._multiplicity ) - if hasattr(imat, 'get'): - imat = imat.get() - else: - imat = _np.asanyarray(imat) + imat = xp.to_numpy(imat) self.hmat= imat if self.periodic: diff --git a/feectools/fem/tensor.py b/feectools/fem/tensor.py index 513635c52..f5b9ed2eb 100644 --- a/feectools/fem/tensor.py +++ b/feectools/fem/tensor.py @@ -44,6 +44,23 @@ eval_fields_3d_weighted, eval_fields_3d_irregular_weighted) +from cunumpy import PyccelKernel + +# Kernels generated by Pyccel only understand NumPy arrays; wrap them so they +# can also be called with CuPy arrays (see cunumpy.kernel.PyccelKernel). +eval_fields_1d_no_weights = PyccelKernel(eval_fields_1d_no_weights) +eval_fields_1d_irregular_no_weights = PyccelKernel(eval_fields_1d_irregular_no_weights) +eval_fields_1d_weighted = PyccelKernel(eval_fields_1d_weighted) +eval_fields_1d_irregular_weighted = PyccelKernel(eval_fields_1d_irregular_weighted) +eval_fields_2d_no_weights = PyccelKernel(eval_fields_2d_no_weights) +eval_fields_2d_irregular_no_weights = PyccelKernel(eval_fields_2d_irregular_no_weights) +eval_fields_2d_weighted = PyccelKernel(eval_fields_2d_weighted) +eval_fields_2d_irregular_weighted = PyccelKernel(eval_fields_2d_irregular_weighted) +eval_fields_3d_no_weights = PyccelKernel(eval_fields_3d_no_weights) +eval_fields_3d_irregular_no_weights = PyccelKernel(eval_fields_3d_irregular_no_weights) +eval_fields_3d_weighted = PyccelKernel(eval_fields_3d_weighted) +eval_fields_3d_irregular_weighted = PyccelKernel(eval_fields_3d_irregular_weighted) + __all__ = ('TensorFemSpace',) #=============================================================================== @@ -500,7 +517,7 @@ def eval_fields(self, grid, *fields, weights=None, npts_per_cell=None, overlap=0 # -> grid is tensor-product, but npts_per_cell is not the same in each cell elif grid[0].ndim == 1 and npts_per_cell is None: out_fields = self.eval_fields_irregular_tensor_grid(grid, *fields, weights=weights, overlap=overlap) - return [xp.ascontiguousarray(out_fields[..., i]) for i in range(len(fields))] + return [xp.ascontiguousarray(xp.asarray(out_fields[..., i])) for i in range(len(fields))] # Case 3. 1D arrays of coordinates and npts_per_cell is a tuple or an integer # -> grid is tensor-product, and each cell has the same number of evaluation points @@ -512,7 +529,7 @@ def eval_fields(self, grid, *fields, weights=None, npts_per_cell=None, overlap=0 grid[i] = xp.reshape(grid[i], (ncells_i, npts_per_cell[i])) out_fields = self.eval_fields_regular_tensor_grid(grid, *fields, weights=weights, overlap=overlap) # return a list - return [xp.ascontiguousarray(out_fields[..., i]) for i in range(len(fields))] + return [xp.ascontiguousarray(xp.asarray(out_fields[..., i])) for i in range(len(fields))] # Case 4. (self.ldim)D arrays of coordinates and no npts_per_cell # -> unstructured grid diff --git a/feectools/fem/tests/analytical_profiles_1d.py b/feectools/fem/tests/analytical_profiles_1d.py index 9966399da..7921260d3 100644 --- a/feectools/fem/tests/analytical_profiles_1d.py +++ b/feectools/fem/tests/analytical_profiles_1d.py @@ -32,6 +32,7 @@ def poly_order( self ): return -1 def eval( self, x, diff=0 ): + x = xp.asarray(x) return self._k**diff * xp.cos( 0.5*math.pi*diff + self._k*x + self._phi ) def max_norm( self, diff=0 ): @@ -58,6 +59,7 @@ def poly_order( self ): return -1 def eval( self, x, diff=0 ): + x = xp.asarray(x) return self._k**diff * xp.sin( 0.5*math.pi*diff + self._k*x + self._phi ) def max_norm( self, diff=0 ): diff --git a/feectools/fem/tests/utilities.py b/feectools/fem/tests/utilities.py index fa1cd0d26..46596f3e2 100644 --- a/feectools/fem/tests/utilities.py +++ b/feectools/fem/tests/utilities.py @@ -8,6 +8,10 @@ def horner( x, *poly_coeffs ): """ Use Horner's Scheme to evaluate a polynomial of coefficients *poly_coeffs at location x. """ + # Spline metadata (notably Greville abscissas) is intentionally host-side. + # Convert it at the numerical API boundary so CuPy coefficients and NumPy + # coordinates can be combined just as they can under NumPy. + x = xp.asarray(x) p = 0 for c in poly_coeffs[::-1]: p = p*x + c diff --git a/feectools/linalg/basic.py b/feectools/linalg/basic.py index 266172d00..b5990d297 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 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: + 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/direct_solvers.py b/feectools/linalg/direct_solvers.py index 1b6096d12..d05a4e176 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 @@ -143,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 @@ -157,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 @@ -228,18 +241,29 @@ 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 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. + 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/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/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/kron.py b/feectools/linalg/kron.py index 80b013f14..cf413ad75 100644 --- a/feectools/linalg/kron.py +++ b/feectools/linalg/kron.py @@ -1,6 +1,7 @@ #coding = utf-8 from functools import reduce +import numpy as np import cunumpy as xp from scipy.sparse import kron from scipy.sparse import coo_matrix @@ -113,7 +114,10 @@ def dot(self, x, out=None): for jj in xp.ndindex(*pnrows): i_mats = [mat._data[s, j] for s,j,mat in zip(xx, jj, mats)] ii_jj = tuple(i+j+(s-1)*p for i,j,p,s in zip(ii, jj, pads, shifts)) - v += x._data[ii_jj] * xp.prod(i_mats) + # ``array_api_compat.cupy`` does not accept a Python list in + # ``prod``; multiplying the scalar factors also avoids a + # temporary device array in this innermost loop. + v += x._data[ii_jj] * reduce(lambda a, b: a * b, i_mats, 1) out._data[xx] = v @@ -151,7 +155,7 @@ def __getitem__(self, key): cols = key[self.ndim:] mats = self.mats elements = [A[i,j] for A,i,j in zip(mats, rows, cols)] - return xp.prod(elements) + return reduce(lambda a, b: a * b, elements, 1) def tostencil(self): @@ -189,7 +193,7 @@ def _tostencil(M, mats, nrows, nrows_extra, pads, xpads): for kk in xp.ndindex( *ndiags ): values = [mat[i,k] for mat,i,k in zip(mats, ii, kk)] - M[(*ii, *kk)] = xp.prod(values) + M[(*ii, *kk)] = reduce(lambda a, b: a * b, values, 1) # handle partly-multiplied rows new_nrows = nrows.copy() @@ -212,7 +216,7 @@ def _tostencil(M, mats, nrows, nrows_extra, pads, xpads): for kk in xp.ndindex( *ndiags ): values = [mat[i,k] for mat,i,k in zip(mats, ii, kk)] - M[(*ii, *kk)] = xp.prod(values) + M[(*ii, *kk)] = reduce(lambda a, b: a * b, values, 1) new_nrows[d] += er def tosparse(self): @@ -440,15 +444,19 @@ def _setup_solvers(self): Computes the distribution of elements and sets up the solvers (which potentially utilize MPI). """ - # slice sizes - starts = xp.array(self._domain.starts) - ends = xp.array(self._domain.ends) + 1 + # slice sizes -- domain-decomposition bookkeeping (sizes/starts/ends), + # always host-resident regardless of the active array backend, matching + # self._domain.starts/ends (which are already plain host ints, see + # CartDecomposition in feectools.ddm.cart) and the MPI sizes/displacements + # computed from them below in KroneckerSolverParallelPass. + starts = np.array(self._domain.starts) + ends = np.array(self._domain.ends) + 1 self._slice = tuple([slice(s, e) for s,e in zip(starts, ends)]) # local and global sizes nglobals = self._domain.npts nlocals = ends - starts - self._localsize = xp.prod(nlocals) + self._localsize = np.prod(nlocals) mglobals = self._localsize // nlocals self._nlocals = nlocals @@ -491,7 +499,9 @@ def _setup_permutations(self): # we use a single permutation for all steps # it is: (n, 1, 2, ..., n-1) - self._perm = xp.arange(self._ndim) + # host bookkeeping (ndim-sized), like self._shapes/self._nlocals it + # indexes -- see the note on _setup_solvers above. + self._perm = np.arange(self._ndim) self._perm[1:] = self._perm[:-1] self._perm[0] = self._ndim - 1 @@ -793,20 +803,20 @@ def __init__(self, solver, mpi_type, i, cart, mglobal, nglobal, nlocal, localsiz cartstart = cart.global_starts[i] cartsize = cartend - cartstart - # source MPI sizes and disps - # distribute the data like - # (N+1, N+1, ..., N+1, N, N, ...) - # where N = floor(mglobaldata / comm.size) + # source MPI sizes and disps -- these are passed straight to + # mpi4py's Alltoallv as counts/displacements, which (like + # cart.global_starts/global_ends above) must be host int arrays + # regardless of backend; keep this whole computation on numpy. mlocal_pre = mglobal // comm.size mlocal_add = mglobal % comm.size - sourcesizes = xp.full((comm.size,), mlocal_pre, dtype=int) + sourcesizes = np.full((comm.size,), mlocal_pre, dtype=int) sourcesizes[:mlocal_add] += 1 mlocal = sourcesizes[comm.rank] sourcesizes *= nlocal # disps, created from the sizes - sourcedisps = xp.zeros((comm.size+1,), dtype=int) - xp.cumsum(sourcesizes, out=sourcedisps[1:]) + sourcedisps = np.zeros((comm.size+1,), dtype=int) + np.cumsum(sourcesizes, out=sourcedisps[1:]) sourcedisps = sourcedisps[:-1] # target MPI sizes and disps diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index d2e673a5f..21047b09f 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) @@ -1897,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] @@ -1905,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] @@ -2131,4 +2140,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) diff --git a/feectools/linalg/sparse.py b/feectools/linalg/sparse.py index d14a9cb55..3af174dd5 100644 --- a/feectools/linalg/sparse.py +++ b/feectools/linalg/sparse.py @@ -2,6 +2,7 @@ from scipy.sparse import sparray, csr_array, bsr_array from scipy.sparse import spmatrix, csr_matrix, bsr_matrix +import cunumpy as xp from feectools.linalg.basic import LinearOperator from feectools.linalg.basic import VectorSpace, Vector, LinearOperator @@ -97,7 +98,17 @@ def _dot_recursive(self, v, out, ind_V=0, ind_W=0): dim_W = W.dimension dim_V = V.dimension - out[index_global_W].flat += self._matrix[ind_W:ind_W+dim_W, ind_V:ind_V+dim_V] @ v[index_global_V].flat + matrix = self._matrix[ind_W:ind_W+dim_W, ind_V:ind_V+dim_V] + values = v[index_global_V].flat + # SciPy sparse matrices are host-only. Stage just this local + # vector slice when the active backend is CuPy, then assign the + # result back through the backend-neutral vector interface. + if xp.is_gpu(v._data): + product = matrix @ xp.to_numpy(v[index_global_V]).ravel() + target = out[index_global_W] + target[...] += xp.asarray(product).reshape(target.shape) + else: + out[index_global_W].flat += matrix @ values elif isinstance(v, BlockVector): diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 4848595c3..132081f42 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 @@ -38,17 +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 xp.is_gpu(val) #========================================================================# Dictionary used to select correct kernel functions based on dimensionality kernels = { @@ -62,6 +66,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 +111,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 +213,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 +250,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 +302,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 +446,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 +610,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 +1235,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,9 +1267,9 @@ 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'): + if xp.is_gpu(out._data): import cupy as cp out._data[:] = cp.asarray(out_data_np) else: @@ -1126,6 +1279,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): """ @@ -1175,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 @@ -1223,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) @@ -1687,7 +1871,7 @@ def tocoo_local(self, order='C'): M = coo_matrix( (data,(rows,cols)), - shape = [xp.prod(nr),xp.prod(nc)], + shape = [math.prod(nr),math.prod(nc)], dtype = self._domain.dtype ) @@ -1749,23 +1933,12 @@ def _tocoo_no_pads(self , order='C'): data[:ind] = cp.asarray(data_np[:ind]) rows[:ind] = cp.asarray(rows_np[:ind]) cols[:ind] = cp.asarray(cols_np[:ind]) - nrl = [_np.int64(e-s+1) for s,e in zip(self.codomain.starts, self.codomain.ends)] - ncl = [_np.int64(i) for i in self._data.shape[nd:]] - ss = [_np.int64(i) for i in ss] - nr = [_np.int64(i) for i in nr] - nc = [_np.int64(i) for i in nc] - dm = [_np.int64(i) for i in dm] - cm = [_np.int64(i) for i in cm] - cpads = [_np.int64(i) for i in cpads] - pp = [_np.int64(i) for i in pp] - stencil2coo = kernels['stencil2coo'][order][nd] - ind = stencil2coo(self._data, data, rows, cols, *nrl, *ncl, *ss, *nr, *nc, *dm, *cm, *cpads, *pp) - - if array_backend.backend == "cupy": + def _host(a): + return xp.to_numpy(a) M = coo_matrix( - (data[:ind].get(), (rows[:ind].get(), cols[:ind].get())), + (_host(data[:ind]), (_host(rows[:ind]), _host(cols[:ind]))), shape=[int(_np.prod(nr)), int(_np.prod(nc))], dtype=self.dtype ) @@ -1854,7 +2027,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 +2181,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 +2205,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 +2340,7 @@ def _get_diagonal_indices(self): nrows = [e - s + 1 for s, e in zip(self.codomain.starts, self.codomain.ends)] ndim = self.domain.ndim - indices = [xp.zeros(xp.prod(nrows), dtype=int) for _ in range(2 * ndim)] + indices = [xp.zeros(math.prod(nrows), dtype=int) for _ in range(2 * ndim)] for l, xx in enumerate(xp.ndindex(*nrows)): ii = [m * p + x for m, p, x in zip(dm, dp, xx)] @@ -2239,7 +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() @@ -2927,7 +3101,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..01a293e54 --- /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 xp.to_numpy(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_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_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_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_mpi_device.py b/feectools/linalg/tests/test_mpi_device.py new file mode 100644 index 000000000..aa55737f2 --- /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 = xp.to_numpy(data) + s0, s1 = int(V.starts[0]), int(V.starts[1]) + p0, p1 = int(V.pads[0]), int(V.pads[1]) + + expected = np.empty_like(data) + for k0 in range(data.shape[0]): + for k1 in range(data.shape[1]): + expected[k0, k1] = glob[(s0 - p0 + k0) % NPTS[0], + (s1 - p1 + k1) % NPTS[1]] + + assert np.allclose(data, expected, rtol=0.0, atol=1e-14) + + +# =============================================================================== +def test_matvec_matches_global_reference(): + """A @ v must equal the stencil applied to the global field, whatever the + decomposition. This is the check that catches an unsynchronized ghost + exchange: a self-consistency check between two distributed runs does not, + because both would be wrong identically.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + v = scatter(V, glob) + A = laplacian(V) + + w = A.dot(v) + ref = reference_apply(glob) + + # Compare through a global reduction, so the check is decomposition-free. + got = float(w.inner(v)) + expected = float((ref * glob).sum()) + assert abs(got - expected) <= 1e-9 * abs(expected) + + # And entry by entry on the rows this rank owns + data = w._data + data = xp.to_numpy(data) + p0, p1 = int(V.pads[0]), int(V.pads[1]) + for i1 in range(int(V.starts[0]), int(V.ends[0]) + 1): + for i2 in range(int(V.starts[1]), int(V.ends[1]) + 1): + k0 = p0 + i1 - int(V.starts[0]) + k1 = p1 + i2 - int(V.starts[1]) + assert abs(data[k0, k1] - ref[i1, i2]) <= 1e-12 + + +# =============================================================================== +def test_matvec_after_device_kernels_without_explicit_sync(): + """The exchange must be safe when the vector was just written by kernels + and the ghost update happens implicitly inside `A.dot` -- the ordering the + PCG loop produces. + + A missing synchronization here is a data race, so the test has to force it + rather than hope for it: a long chain of asynchronous work on the vector is + queued and the exchange is triggered immediately afterwards, leaving the + stream busy while MPI reads the buffer. + """ + comm = MPI.COMM_WORLD + npts, pads = (512, 512), (1, 2) + V = make_space(npts, pads, comm=comm) + glob = global_field(npts) + A = laplacian(V) + + r = scatter(V, glob) + + # Queue work that writes r, mathematically the identity so the expected + # result is unchanged. On a device the arrays are large and the chain long + # enough that kernels are still queued when the exchange starts -- which is + # what makes the race reproducible rather than occasional. There is no race + # on the host, so one pass is enough there. + passes = 200 if xp.is_gpu(r._data) else 1 + for _ in range(passes): + r._data *= 2.0 + r._data *= 0.5 + + r.ghost_regions_in_sync = False + w = A.dot(r) # triggers the implicit ghost update + + ref = reference_apply(glob) + got = float(w.inner(r)) + expected = float((ref * glob).sum()) + assert abs(got - expected) <= 1e-9 * abs(expected) + + +# =============================================================================== +def test_ghost_exchange_synchronizes_before_mpi(monkeypatch): + """ + The exchangers must call `synchronize_for_mpi` before giving a buffer to + MPI. + + This is checked structurally rather than by observing corrupted data, + because the underlying race is not deterministic: whether MPI actually + reads a half-written buffer depends on which internal protocol it picks for + the message, and some of those happen to synchronize with the CuPy stream + by accident. Relying on that accident is exactly the bug, so the contract + is what gets tested. + """ + import feectools.ddm.blocking_data_exchanger as blocking + import feectools.ddm.nonblocking_data_exchanger as nonblocking + + calls = [] + for module in (blocking, nonblocking): + monkeypatch.setattr(module, 'synchronize_for_mpi', + lambda *args: calls.append(args)) + + V = make_space(NPTS, PADS, comm=MPI.COMM_WORLD) + v = StencilVector(V) + v.ghost_regions_in_sync = False + v.update_ghost_regions() + + assert calls, 'ghost exchange handed a buffer to MPI without synchronizing' + assert any(v._data is arg for args in calls for arg in args), \ + 'the synchronized buffer was not the one being exchanged' + + +# =============================================================================== +def test_axpy_then_matvec_is_correct(): + """`mul_iadd` writes on the device; the following exchange must see it.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + A = laplacian(V) + + x = scatter(V, glob) + y = scatter(V, glob) + x.mul_iadd(2.0, y) # x = 3 * glob + w = A.dot(x) + + ref = reference_apply(3.0 * glob) + got = float(w.inner(x)) + expected = float((ref * (3.0 * glob)).sum()) + assert abs(got - expected) <= 1e-9 * abs(expected) + + +# =============================================================================== +def test_inner_matches_global_reference(): + """Reductions must equal the value computed from the global field.""" + comm = MPI.COMM_WORLD + V = make_space(NPTS, PADS, comm=comm) + glob = global_field(NPTS) + other = np.flipud(glob).copy() + + x = scatter(V, glob) + y = scatter(V, other) + + assert abs(float(x.inner(y)) - float((glob * other).sum())) <= 1e-9 + 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'])) diff --git a/feectools/linalg/tests/test_stencil_interface_matrix.py b/feectools/linalg/tests/test_stencil_interface_matrix.py index cfb755f2a..8406fd350 100644 --- a/feectools/linalg/tests/test_stencil_interface_matrix.py +++ b/feectools/linalg/tests/test_stencil_interface_matrix.py @@ -5,6 +5,7 @@ #---------------------------------------------------------------------------# import pytest import cunumpy as xp +import numpy as np from random import random from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix, StencilInterfaceMatrix @@ -22,7 +23,8 @@ def compute_global_starts_ends(domain_decomposition, npts, pads): global_ends [axis] = ee.copy() global_ends [axis][-1] = npts[axis]-1 - global_starts[axis] = xp.array([0] + (global_ends[axis][:-1]+1).tolist()) + # Cartesian partition data is host metadata, including on CuPy. + global_starts[axis] = np.array([0] + (global_ends[axis][:-1]+1).tolist()) for s, e, p in zip(global_starts, global_ends, pads): assert all(e - s + 1 >= p) @@ -94,7 +96,7 @@ def test_stencil_interface_matrix_1d_serial_init(dtype, n1, p1, s1, axis, ext, P assert M.domain_start == (0,) * M.dim assert M.codomain_start == (0,) * M.dim assert M.flip == (1,) * M.dim - assert xp.array_equal(M.permutation, [0]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([0])) assert M.pads == (p1,) assert M.backend == None assert M._data.shape == (p1 + 1 + 2 * p1 * s1, 1 + 2 * p1) @@ -149,9 +151,9 @@ def test_stencil_interface_matrix_2d_serial_init(dtype, n1, n2, p1, p2, s1, s2, elif axis2 == 1: assert M._data.shape == (n1 + 2 * p1 * s1, p2 + 1 + 2 * p2 * s2, 1 + 2 * p1, 1 + 2 * p2) if axis1 == axis2: - assert xp.array_equal(M.permutation, [0, 1]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([0, 1])) else: - assert xp.array_equal(M.permutation, [1, 0]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([1, 0])) assert M.shape == (n1 * n2, n1 * n2) # =============================================================================== @@ -211,11 +213,11 @@ def test_stencil_interface_matrix_3d_serial_init(dtype, n1, n2, n3, p1, p2, p3, assert M._data.shape == ( n1 + 2 * p1 * s1, n2 + 2 * p2 * s2, p3 + 1 + 2 * p3 * s3, 1 + 2 * p1, 1 + 2 * p2, 1 + 2 * p3) if axis1 == axis2: - assert xp.array_equal(M.permutation, [0, 1, 2]) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray([0, 1, 2])) else: permutation = [0, 1, 2] permutation[axis1], permutation[axis2] = permutation[axis2], permutation[axis1] - assert xp.array_equal(M.permutation, permutation) + assert xp.array_equal(xp.asarray(M.permutation), xp.asarray(permutation)) assert M.shape == (n1 * n2 * n3, n1 * n2 * n3) #=============================================================================== # Parallel TESTS diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 42d5636b7..fa64a0802 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -1,12 +1,13 @@ # coding: utf-8 -import cunumpy as xp from math import sqrt -from feectools.linalg.basic import Vector +import cunumpy as xp + +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', @@ -14,30 +15,19 @@ '_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() @@ -48,58 +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 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. + Convert a PETSc.Vec object to a StencilVector or BlockVector. - Parameters - ---------- - x : PETSc.Vec - PETSc vector - - 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 @@ -107,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 @@ -136,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') @@ -163,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: @@ -194,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: 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) 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"