GPU stack 5/8: Poisson/diffusion solve with fewer passes over the vectors

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

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<double> pick(int c, Array3D<double> a0, Array3D<double> a1, Array3D<double> 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<double> x, Array3D<double> m_in, Array3D<double> out0,
+                                                   Array3D<double> out1, Array3D<double> out2, Array3D<double> m0,
+                                                   Array3D<double> m1, Array3D<double> 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<double> out = pick(c, out0, out1, out2);
+    Array3D<double> 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<double> y0, Array3D<double> y1, Array3D<double> y2,
+                                                     Array3D<double> m0, Array3D<double> m1, Array3D<double> m2,
+                                                     Array3D<double> out, Array3D<double> 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<double> y = pick(c, y0, y1, y2);
+        Array3D<double> 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);
+}