diff --git a/src/struphy/feec/fused_derivatives.py b/src/struphy/feec/fused_derivatives.py new file mode 100644 index 000000000..dd3ef5919 --- /dev/null +++ b/src/struphy/feec/fused_derivatives.py @@ -0,0 +1,217 @@ +"""The discrete gradient with its boundary operators as one CUDA kernel. + +``Derham.grad`` is ``B1 @ G @ B0.T``: the boundary operator of the 0-forms (a 0/1 mask), the derivatives of the +three directions (feectools ``DirectionalDerivativeOperator``, a difference of neighbouring coefficients) and the +boundary operator of the 1-forms. Applied as composed operators, each product makes three passes over the vectors +(and one more over the input for the masked copy), and its transpose four. :func:`fuse_boundary_gradient` returns an +operator that applies the whole chain with one kernel on the GPU (``fused_derivatives_cuda.cu``), with the same +operations in the same order, so the result is bit for bit the same, and the composed operator otherwise. +""" + +from functools import cache +from pathlib import Path + +import cunumpy as xp +import numpy as np +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ComposedLinearOperator, LinearOperator +from feectools.linalg.block import BlockLinearOperator, BlockVector +from feectools.linalg.stencil import StencilVector + +from struphy.feec.linear_operators import BoundaryOperator + +# entries of the kernel metadata per component, see fused_derivatives_cuda.cu +_META = 8 + + +@cache +def _kernels(): + """The forward and transposed CUDA kernels (compiled on first launch).""" + from cunumpy.kernels import CudaKernel + + path = Path(__file__).with_name("fused_derivatives_cuda.cu") + return ( + CudaKernel.from_file(path, "fused_boundary_gradient"), + CudaKernel.from_file(path, "fused_boundary_gradient_t"), + ) + + +def _derivative_blocks(block: BlockLinearOperator): + """The derivatives of a gradient block (one block column) or its transpose (one block row), else None.""" + if not isinstance(block, BlockLinearOperator): + return None + n_rows, n_cols = block.n_block_rows, block.n_block_cols + if min(n_rows, n_cols) != 1 or not 1 <= max(n_rows, n_cols) <= 3: + return None + ops = [block[k, 0] if n_cols == 1 else block[0, k] for k in range(max(n_rows, n_cols))] + if not all(isinstance(op, DirectionalDerivativeOperator) and op.domain.ndim == 3 for op in ops): + return None + transposed = n_cols > 1 or (n_rows == 1 and ops[0].transposed) + if any(op.transposed != transposed for op in ops): + return None + return ops, transposed + + +def fuse_boundary_gradient(op: LinearOperator) -> LinearOperator: + """``op`` as a :class:`FusedBoundaryGradient` if it is a gradient with boundary operators (``B_out @ G @ B_in``, + ``G`` a block of directional derivatives) or the transpose of one, else ``op`` itself. + + Parameters + ---------- + op : LinearOperator + E.g. ``derham.grad`` or ``derham.grad.T``. + + Returns + ------- + LinearOperator + The same operator; on the GPU applied with one kernel. + """ + if isinstance(op, FusedBoundaryGradient) or not isinstance(op, ComposedLinearOperator): + return op + factors = op.multiplicands + if len(factors) != 3: + return op + b_out, block, b_in = factors + if not (isinstance(b_out, BoundaryOperator) and isinstance(b_in, BoundaryOperator)): + return op + found = _derivative_blocks(block) + if found is None: + return op + derivatives, transposed = found + return FusedBoundaryGradient(op, b_out, derivatives, b_in, transposed) + + +class FusedBoundaryGradient(LinearOperator): + """``B_out @ G @ B_in`` (a gradient with its boundary operators, or its transpose) with one kernel on the GPU. + + The kernel computes the masked differences on the owned entries of the output from the input with its ghost + regions, as the composed operators do (see ``fused_derivatives_cuda.cu``). It is used for device data on a GPU; the + composed operator is applied in every other case (host data, cunumpy's fake CuPy, polar or other vectors). + + Parameters + ---------- + op : ComposedLinearOperator + The composed operator, applied when the kernel does not apply. + b_out, b_in : BoundaryOperator + The boundary operators on the output and input side. + derivatives : list[DirectionalDerivativeOperator] + The derivatives per component (of the output, or of the input if transposed). + transposed : bool + Whether ``op`` is the transpose of a gradient (1-forms to 0-forms). + """ + + def __init__(self, op, b_out, derivatives, b_in, transposed: bool): + self._op = op + self._b_out = b_out + self._b_in = b_in + self._derivatives = derivatives + self._transposed = transposed + self._meta = None + self._input_masks = None + + @property + def operator(self) -> LinearOperator: + """The composed operator.""" + return self._op + + @property + def domain(self): + return self._op.domain + + @property + def codomain(self): + return self._op.codomain + + @property + def dtype(self): + return self._op.dtype + + def tosparse(self): + return self._op.tosparse() + + def toarray(self): + return self._op.toarray() + + def transpose(self, conjugate=False): + return fuse_boundary_gradient(self._op.transpose(conjugate=conjugate)) + + def dot(self, v, out=None): + if self._use_kernel(v): + return self._dot_fused(v, out) + return self._op.dot(v, out=out) + + def _use_kernel(self, v) -> bool: + """Whether :meth:`dot` uses the kernel: device data on a real GPU (not cunumpy's fake CuPy).""" + from cunumpy.kernel_testing import fake_cupy_active + + blocks = v.blocks if isinstance(v, BlockVector) else [v] + if not all(isinstance(b, StencilVector) and xp.is_gpu(b._data) for b in blocks): + return False + return not fake_cupy_active() + + def _metadata(self): + """Owned region, direction and sign of each derivative, as the kernels take them (device array, cached).""" + if self._meta is None: + rows = [] + for d in self._derivatives: + # the owned entries of the derivative's codomain, the same padded indices in its domain + idslice = d._idslice + rows.append( + [s.start for s in idslice] + [s.stop - s.start for s in idslice] + [d.diffdir, int(d.negative)] + ) + assert all(len(r) == _META for r in rows) + self._meta = xp.asarray(np.array(rows, dtype=np.int64).ravel()) + return self._meta + + def _input_mask_arrays(self): + """The 0/1 mask of the input boundary operator per block, with its ghost regions exchanged (cached). + + The composed operators mask the input and then exchange the ghost regions of the masked vector (the boundary + operator leaves copies of boundary coefficients in the ghost regions unmasked). The kernels read the input with + its ghost regions exchanged instead, times this mask, whose ghost entries are then those of the owners: the + products are the exchanged values of the masked input, in serial and in parallel. + """ + if self._input_masks is None: + ones = self.domain.zeros() + for b in getattr(ones, "blocks", [ones]): + b._data[...] = 1.0 + ones.ghost_regions_in_sync = True + mask = self._b_in.dot(ones) + mask.update_ghost_regions() + self._input_masks = [b._data for b in getattr(mask, "blocks", [mask])] + return self._input_masks + + def _dot_fused(self, v, out=None): + """The product with the kernels (device data).""" + if out is None: + out = self.codomain.zeros() + masks_out = self._b_out._device_masks(out) + if masks_out is None: + return self._op.dot(v, out=out) + masks_in = self._input_mask_arrays() + + # the kernels read the ghost regions of the input + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + forward, transposed = _kernels() + meta = self._metadata() + n_comp = len(self._derivatives) + + def three(arrays): + return [arrays[min(k, len(arrays) - 1)] for k in range(3)] + + if not self._transposed: + outs = three([b._data for b in out.blocks]) + m_out = three([m for m, _ in masks_out]) + n_threads = sum(int(np.prod([s.stop - s.start for s in d._idslice])) for d in self._derivatives) + forward(v._data, masks_in[0], *outs, *m_out, meta, n_comp, n_threads=n_threads) + else: + ys = three([b._data for b in v.blocks]) + m_in = three(masks_in) + n_threads = int(np.prod([s.stop - s.start for s in self._derivatives[0]._idslice])) + transposed(*ys, *m_in, out._data, masks_out[0][0], meta, n_comp, n_threads=n_threads) + + # as the composed operators: the ghost regions of the output are not written + out.ghost_regions_in_sync = False + return out diff --git a/src/struphy/feec/fused_derivatives_cuda.cu b/src/struphy/feec/fused_derivatives_cuda.cu new file mode 100644 index 000000000..5fb66fbc3 --- /dev/null +++ b/src/struphy/feec/fused_derivatives_cuda.cu @@ -0,0 +1,97 @@ +// The discrete gradient with its boundary operators as one kernel (struphy.feec.fused_derivatives). +#include "cunumpy/array_view.cuh" +#include "cunumpy/index.cuh" + +// Per component c, 8 entries of `meta`: first owned index (3, padded coordinates), owned sizes (3), derivative +// direction, negative (0 or 1). +#define FUSED_META 8 + +__device__ inline Array3D pick(int c, Array3D a0, Array3D a1, Array3D a2) { + return c == 0 ? a0 : (c == 1 ? a1 : a2); +} + +/** + * out_c = B_out,c D_c B_in x for the components c of a gradient, B the 0/1 masks of the boundary operators and D_c the + * derivative in direction dir_c (feectools DirectionalDerivativeOperator): on the owned entries i of out_c (padded + * coordinates, shared by x and out_c as in DirectionalDerivativeOperator), + * + * out_c[i] = ((x[i + e] * m_in[i + e]) - (x[i] * m_in[i])) * m_out_c[i] (operands swapped if negative), + * + * the operations of the composed operators in the same order (bit for bit the same). One thread per owned entry of + * all components (n_threads = sum of the owned sizes), component 0 first; the ghost regions of out are not written. + * + * @param x, m_in Input vector data (ghost regions in sync) and the mask of the input boundary operator. + * @param out0, out1, out2, m0, m1, m2 Data of the output components and the masks of the output boundary operator. + * @param meta FUSED_META entries per component, see above. + * @param n_comp Number of components (1 to 3). + */ +extern "C" __global__ void fused_boundary_gradient(Array3D x, Array3D m_in, Array3D out0, + Array3D out1, Array3D out2, Array3D m0, + Array3D m1, Array3D m2, const long long* meta, + int n_comp) { + long long n_total = 0; + for (int c = 0; c < n_comp; ++c) + n_total += meta[FUSED_META * c + 3] * meta[FUSED_META * c + 4] * meta[FUSED_META * c + 5]; + CUNUMPY_THREAD_1D(t, n_total); + + // CUDA-only: the component and owned entry of this thread + int c = 0; + long long local = t; + while (c < n_comp - 1) { + const long long size = meta[FUSED_META * c + 3] * meta[FUSED_META * c + 4] * meta[FUSED_META * c + 5]; + if (local < size) break; + local -= size; + ++c; + } + const long long* m = meta + FUSED_META * c; + const long long i3 = m[2] + local % m[5]; + const long long i2 = m[1] + (local / m[5]) % m[4]; + const long long i1 = m[0] + local / (m[4] * m[5]); + const int dir = (int)m[6]; + const long long j1 = i1 + (dir == 0), j2 = i2 + (dir == 1), j3 = i3 + (dir == 2); + + const double a = x(j1, j2, j3) * m_in(j1, j2, j3); + const double b = x(i1, i2, i3) * m_in(i1, i2, i3); + const double d = m[7] ? b - a : a - b; + Array3D out = pick(c, out0, out1, out2); + Array3D mask = pick(c, m0, m1, m2); + out(i1, i2, i3) = d * mask(i1, i2, i3); +} + +/** + * The transpose: out = B_out (sum_c D_c^T B_in,c y_c), with D_c^T the transposed derivative; on the owned entries i of + * out, + * + * t_c = (y_c[i - e_c] * m_c[i - e_c]) - (y_c[i] * m_c[i]) (operands swapped if negative), + * out[i] = ((t_0 + t_1) + t_2) * m_out[i], + * + * summed in the order of BlockLinearOperator.dot. One thread per owned entry of out (n_threads = its owned size). + * + * @param y0, y1, y2, m0, m1, m2 Input components (ghost regions in sync) and the masks of the input boundary operator. + * @param out, m_out Output data and the mask of the output boundary operator. + * @param meta FUSED_META entries per component (the owned region of out, the same for every component). + * @param n_comp Number of components (1 to 3). + */ +extern "C" __global__ void fused_boundary_gradient_t(Array3D y0, Array3D y1, Array3D y2, + Array3D m0, Array3D m1, Array3D m2, + Array3D out, Array3D m_out, + const long long* meta, int n_comp) { + CUNUMPY_THREAD_1D(t, meta[3] * meta[4] * meta[5]); + const long long i3 = meta[2] + t % meta[5]; + const long long i2 = meta[1] + (t / meta[5]) % meta[4]; + const long long i1 = meta[0] + t / (meta[4] * meta[5]); + + double acc = 0.0; + for (int c = 0; c < n_comp; ++c) { + const long long* m = meta + FUSED_META * c; + const int dir = (int)m[6]; + const long long j1 = i1 - (dir == 0), j2 = i2 - (dir == 1), j3 = i3 - (dir == 2); + Array3D y = pick(c, y0, y1, y2); + Array3D mask = pick(c, m0, m1, m2); + const double a = y(j1, j2, j3) * mask(j1, j2, j3); + const double b = y(i1, i2, i3) * mask(i1, i2, i3); + const double d = m[7] ? b - a : a - b; + acc = c == 0 ? d : acc + d; + } + out(i1, i2, i3) = acc * m_out(i1, i2, i3); +} diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index b36318a38..056817730 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -1,10 +1,80 @@ +import cunumpy as xp from feectools.linalg.basic import LinearOperator, Vector, VectorSpace -from feectools.linalg.stencil import StencilVectorSpace +from feectools.linalg.block import BlockVector +from feectools.linalg.stencil import StencilVector, StencilVectorSpace from struphy.feec.utilities import apply_essential_bc_to_array from struphy.polar.basic import PolarDerhamSpace +def _blocks(v): + """The StencilVector blocks of a StencilVector or BlockVector (of StencilVectors), else None.""" + if isinstance(v, StencilVector): + return [v] + if isinstance(v, BlockVector) and all(isinstance(b, StencilVector) for b in v.blocks): + return list(v.blocks) + return None + + +_MASKED_COPY_SRC = r""" +extern "C" __global__ void masked_copy(const unsigned long long* __restrict__ ptrs, const long long* __restrict__ offsets, + int n_comp, long long total) +{ + long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + int c = 0; + while (c < n_comp - 1 && i >= offsets[c + 1]) ++c; + long long j = i - offsets[c]; + const double* v = (const double*)ptrs[3 * c]; + double* out = (double*)ptrs[3 * c + 1]; + const double* mask = (const double*)ptrs[3 * c + 2]; + out[j] = v[j] * mask[j]; +} +""" + + +def _masked_copy_fused(op, v_blocks, out_blocks, masks) -> bool: + """``out_blocks[c] = v_blocks[c] * masks[c]`` for all blocks with one kernel (pointer tables cached on ``op``). + + False (nothing done) unless every block is C-contiguous float64 of the shape of its mask, and on cunumpy's fake + CuPy (tests without a GPU), which cannot launch raw kernels.""" + from cunumpy.kernel_testing import fake_cupy_active + + if fake_cupy_active(): + return False + arrays = [(vb._data, ob._data, m) for vb, ob, (m, _) in zip(v_blocks, out_blocks, masks)] + for a, b, m in arrays: + if not ( + a.dtype == b.dtype == m.dtype == xp.float64 + and a.shape == b.shape == m.shape + and a.flags.c_contiguous + and b.flags.c_contiguous + and m.flags.c_contiguous + ): + return False + import cupy as cp + import numpy as np + + cache = op.__dict__.setdefault("_masked_copy_cache", {}) + key = tuple(int(x.data.ptr) for triple in arrays for x in triple) + entry = cache.get(key) + if entry is None: + if "kernel" not in cache: + cache["kernel"] = cp.RawKernel(_MASKED_COPY_SRC, "masked_copy") + sizes = [a.size for a, _, _ in arrays] + offsets = np.concatenate([[0], np.cumsum(sizes)]).astype(np.int64) + entry = (cp.asarray(np.array(key, dtype=np.uint64)), cp.asarray(offsets), int(offsets[-1])) + if len(cache) > 64: + kernel = cache["kernel"] + cache.clear() + cache["kernel"] = kernel + cache[key] = entry + ptrs, offsets, total = entry + if total: + cache["kernel"](((total + 255) // 256,), (256,), (ptrs, offsets, np.int32(len(arrays)), np.int64(total))) + return True + + class BoundaryOperator(LinearOperator): r""" Applies homogeneous Dirichlet boundary conditions to a vector. @@ -196,6 +266,22 @@ def dot(self, v, out=None): else: assert isinstance(out, Vector) assert out.space == self._codomain + masks = self._device_masks(v) + if masks is not None: + # copy and boundary conditions in one kernel per block (see _device_masks) + in_sync = True + v_blocks, out_blocks = _blocks(v), _blocks(out) + if len(v_blocks) > 1 and _masked_copy_fused(self, v_blocks, out_blocks, masks): + pass + else: + for v_blk, out_blk, (mask, _) in zip(v_blocks, out_blocks, masks): + xp.multiply(v_blk._data, mask, out=out_blk._data) + for v_blk, (_, stale) in zip(v_blocks, masks): + in_sync = in_sync and v_blk.ghost_regions_in_sync and not stale + # one flag for all blocks: a BlockVector flagged in sync with stale blocks would have them updated + # one at a time by the operators that check each block, instead of together + out.ghost_regions_in_sync = in_sync + return out v.copy(out=out) # apply boundary conditions to output vector @@ -203,6 +289,28 @@ def dot(self, v, out=None): return out + def _device_masks(self, v): + """Per block of ``v``: (0/1 mask of the entries the boundary conditions keep, whether they stale the ghosts). + + For device (CuPy) data of Stencil/Block vectors, ``dot`` computes ``out = v * mask`` (one kernel per block) + instead of a copy and one fill per boundary and component. The masks are found once, by applying + :func:`apply_essential_bc_to_array` to a vector of ones, so they zero exactly the entries it zeroes (ghost + entries included) and flag the same blocks as having stale ghost regions. None on the host backend and for + other vectors (e.g. polar vectors), which use the copy and :func:`apply_essential_bc_to_array`. + """ + if not hasattr(self, "_masks_cache"): + self._masks_cache = None + blocks = _blocks(v) + if blocks is not None and all(xp.is_gpu(b._data) for b in blocks): + probe = v.space.zeros() + probe_blocks = _blocks(probe) + for b in probe_blocks: + b._data[...] = 1.0 + b.ghost_regions_in_sync = True + apply_essential_bc_to_array(self._space_id, probe, self.bc) + self._masks_cache = [(b._data.copy(), not b.ghost_regions_in_sync) for b in probe_blocks] + return self._masks_cache + def transpose(self, conjugate=False): """ Returns the transposed operator. diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 1020f0ef6..bd13cc954 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2245,6 +2245,15 @@ def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): y = op.dot(y, out=target) return y + def without_boundary_ops(self) -> "MassWithoutBoundaryOps": + """This operator applied without its boundary operators, see :class:`MassWithoutBoundaryOps`.""" + return MassWithoutBoundaryOps(self) + + @property + def boundary_ops(self): + """The boundary operators ``(B_V, B_W)`` applied on the domain and codomain side.""" + return self._V_boundary_op, self._W_boundary_op + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2796,6 +2805,49 @@ def info(self, use_rst=False): return info(self, use_rst=use_rst) +class MassWithoutBoundaryOps(LinearOperator): + """A :class:`WeightedMassOperator` applied without its boundary operators: ``E_W M E_V^T`` instead of + ``B_W E_W M E_V^T B_V^T`` (:meth:`WeightedMassOperator.dot` with ``apply_bc=False``). + + Between operators that apply the same boundary operators on the outside, e.g. in ``grad.T @ M1 @ grad`` with + ``grad = B1 G B0^T``, this is the same operator: boundary operators are 0/1 diagonal masks, hence symmetric and + idempotent (``B1^T B1 = B1``), and the masked products are bit for bit the same, with two fewer passes over the + vectors per product. Get it with :meth:`WeightedMassOperator.without_boundary_ops`. + """ + + def __init__(self, mass: "WeightedMassOperator"): + self._mass = mass + + @property + def mass(self): + """The wrapped mass operator.""" + return self._mass + + @property + def domain(self): + return self._mass.domain + + @property + def codomain(self): + return self._mass.codomain + + @property + def dtype(self): + return self._mass.dtype + + def tosparse(self): + raise NotImplementedError() + + def toarray(self): + raise NotImplementedError() + + def transpose(self, conjugate=False): + return MassWithoutBoundaryOps(self._mass.transpose(conjugate=conjugate)) + + def dot(self, v, out=None): + return self._mass.dot(v, out=out, apply_bc=False) + + class StencilMatrixFreeMassOperator(LinearOperator): r"""Class implementing matrix-free weighted mass operators between StencilVectorSpaces. diff --git a/src/struphy/feec/tests/test_fused_derivatives.py b/src/struphy/feec/tests/test_fused_derivatives.py new file mode 100644 index 000000000..b6e3707d5 --- /dev/null +++ b/src/struphy/feec/tests/test_fused_derivatives.py @@ -0,0 +1,145 @@ +"""The fused gradient with boundary operators (struphy.feec.fused_derivatives) against the composed operators. + +On the host the fused operator applies the composed one. Without a GPU, the CUDA kernels run on cunumpy's fake CuPy with +the launches emulated on the CPU (in a subprocess) and must give bit for bit the products of the composed operators. +""" + +import numpy as np +import pytest + +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child + +# (degree, boundary conditions): Dirichlet at both ends, at one end, none (periodic) +CASES = { + "dirichlet": ((2, 3, 1), (("dirichlet", "dirichlet"), None, None)), + "mixed": ((1, 2, 3), (None, ("dirichlet", "free"), None)), + "periodic": ((3, 1, 2), (None, None, None)), +} + + +def _derham(case): + from maybempi import MPI + + from struphy import domains + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + degree, bcs = CASES[case] + return Derham( + TensorProductGrid(num_elements=(6, 5, 4)), + DerhamOptions(degree=degree, bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domains.Cuboid(), + ) + + +def _random(space, seed): + import cunumpy as xp + + rng = np.random.default_rng(seed) + v = space.zeros() + for b in getattr(v, "blocks", [v]): + b._data[...] = xp.asarray(rng.standard_normal(b._data.shape)) + v.update_ghost_regions() + return v + + +def _owned(v): + import cunumpy as xp + + return xp.to_numpy(v.toarray()) + + +@pytest.mark.parametrize("case", list(CASES)) +def test_fuse_recognizes_gradient(case): + """``derham.grad`` and its transpose are fused; on the host the fused operators give the composed products. + + Without essential boundary conditions the gradient has no boundary operators (only the derivatives), and it is + returned as it is.""" + from struphy.feec.fused_derivatives import FusedBoundaryGradient, fuse_boundary_gradient + + derham = _derham(case) + for op in (derham.grad, derham.grad.T): + fused = fuse_boundary_gradient(op) + if case == "periodic": + assert fused is op + continue + assert isinstance(fused, FusedBoundaryGradient) + assert fused.domain is op.domain and fused.codomain is op.codomain + v = _random(op.domain, 1) + np.testing.assert_array_equal(_owned(fused.dot(v)), _owned(op.dot(v))) + assert isinstance(fused.T, FusedBoundaryGradient) + # other operators are returned as they are + assert fuse_boundary_gradient(derham.curl) is derham.curl + + +def check_fused_gradient_on_device(): + """On the CuPy backend: the kernels give bit for bit the products of the composed operators.""" + import cunumpy as xp + + import struphy.feec.fused_derivatives as fd + from struphy.feec.fused_derivatives import fuse_boundary_gradient + + # count the launches of the two kernels: the checks below must not pass through the composed operators + kernels, launches = fd._kernels, [0, 0] + + def counted(k, i): + def launch(*args, **kwargs): + launches[i] += 1 + return k(*args, **kwargs) + + return launch + + fd._kernels = lambda: tuple(counted(k, i) for i, k in enumerate(kernels())) + try: + _check_fused_gradient(xp, fuse_boundary_gradient) + finally: + fd._kernels = kernels + assert launches[0] > 0 and launches[1] > 0, launches + + +def _check_fused_gradient(xp, fuse_boundary_gradient): + with xp.use_backend("cupy"): + for case in CASES: + derham = _derham(case) + for op in (derham.grad, derham.grad.T): + fused = fuse_boundary_gradient(op) + if fused is op: + continue + v = _random(op.domain, 2) + expected = _owned(op.dot(v)) + out = op.codomain.zeros() + result = fused._dot_fused(v, out=out) + assert result is out + assert not out.ghost_regions_in_sync + assert np.abs(expected).max() > 0, case + np.testing.assert_array_equal(_owned(result), expected, err_msg=case) + # an input whose ghost regions are not in sync + v.ghost_regions_in_sync = False + np.testing.assert_array_equal(_owned(fused._dot_fused(v)), expected, err_msg=case) + + +def check_fused_gradient_on_fake_cupy(): + """:func:`check_fused_gradient_on_device` with the CUDA kernel launches emulated on the CPU.""" + from cunumpy.kernel_testing import emulated_launches, fake_cupy_active + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + assert fake_cupy_active() + with emulated_launches(options=EMULATION_OPTIONS): + check_fused_gradient_on_device() + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +def test_fused_gradient_fake_cupy(): + """Without a GPU: the kernels against the composed operators on cunumpy's fake CuPy (serial subprocess).""" + run_fake_cupy_child( + "from struphy.feec.tests.test_fused_derivatives import check_fused_gradient_on_fake_cupy; " + "check_fused_gradient_on_fake_cupy()" + ) + + +@pytest.mark.skipif(not _cupy_installed(), reason="needs CuPy and a GPU") +def test_fused_gradient_on_gpu(): + check_fused_gradient_on_device() diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 20f5bf2eb..24c5a3da6 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -10,6 +10,7 @@ from scope_profiler import ProfileManager from struphy.feec import preconditioner +from struphy.feec.fused_derivatives import fuse_boundary_gradient from struphy.feec.mass import L2Projector, WeightedAverageProjection, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions, MultiGridPreconditioner @@ -333,10 +334,24 @@ def options(self, new): @profile def allocate(self): # always stabilize + auto_stabilized = False if xp.abs(self.options.sigma_1) < 1e-14: self.options.sigma_1 = 1e-14 + auto_stabilized = True logger.warning(f"Stabilizing Poisson solve with {self.options.sigma_1 =}") + # The automatic stabilization only makes singular problems (e.g. periodic in all directions) solvable. With a + # Dirichlet boundary in some direction, grad.T @ M1 @ grad is already positive definite and 1e-14 * stab_mat + # changes the solution only at round-off, but applying it costs a whole stab_mat product per solver iteration; + # it is then left out of the solver's operator. (M1perp can be degenerate, so it keeps the stabilization.) + self._stabilize_lhs = not ( + auto_stabilized + and self.options.diffusion_mat == "M1" + and any(any(sides) for sides in self.derham.dirichlet_bc) + ) + if not self._stabilize_lhs: + logger.info("Dirichlet boundary present: the automatic stabilization is left out of the solver operator") + # model parameters self._sigma_1 = self.options.sigma_1 self._sigma_2 = self.options.sigma_2 @@ -442,7 +457,21 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: if diffusion_mat is None: self._diffusion_op = None else: - self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad + # grad = B1 G B0^T already applies the boundary operator B1 of the 1-forms on both sides of diffusion_mat; + # if diffusion_mat applies the same one, it is left out there (the same operator, see + # MassWithoutBoundaryOps), except for the multigrid preconditioner, which coarsens the operator by its + # structure + mid = diffusion_mat + grad, grad_t = self.derham.grad, self.derham.grad.T + B1 = self.derham.boundary_ops["1"] + if self.options.precond != "MultiGrid" and all( + op is B1 for op in getattr(diffusion_mat, "boundary_ops", (None,)) + ): + mid = diffusion_mat.without_boundary_ops() + if self.options.precond != "MultiGrid": + # on the GPU, the gradient and its transpose with their boundary operators as one kernel each + grad, grad_t = fuse_boundary_gradient(grad), fuse_boundary_gradient(grad_t) + self._diffusion_op = grad_t @ mid @ grad # preconditioner and solver for Ax=b # the multigrid and stiffness preconditioners depend on sigma_1, they are updated in __call__ @@ -582,8 +611,13 @@ def __call__(self, dt): # compute rhs phin = self.variables.phi.spline.vector - rhs = self._stab_mat.dot(phin, out=self._rhs) - rhs *= sig_2 + if sig_2 == 0.0: + # no stab_mat product only to multiply it by zero (e.g. Poisson: sigma_2 = 0) + rhs = self._rhs + rhs *= 0.0 + else: + rhs = self._stab_mat.dot(phin, out=self._rhs) + rhs *= sig_2 self._rhs2 *= 0.0 for src, coeff in zip(self.sources, self.coeffs): @@ -634,8 +668,10 @@ def __call__(self, dt): # compute lhs if self._diffusion_op is None: self._solver.linop = sig_1 * self._stab_mat - else: + elif self._stabilize_lhs: self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + else: + self._solver.linop = self._diffusion_op if sig_1 != self._pc_sig_1: if self._mg is not None: self._mg.update(self._solver.linop)