diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..5456b2eeb --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,249 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch + + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 02240a6d0..42a340606 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1666,6 +1667,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2084,6 +2087,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2131,29 +2160,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2240,6 +2294,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2515,6 +2570,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2525,6 +2581,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2535,11 +2592,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2552,11 +2611,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: