perf(numba): the mapper x mapper block and the MGE operated matrix (#507)

autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py

@@ -492,6 +492,44 @@ def curvature_matrix_with_added_to_diag_from(
     return curvature_matrix
 
 
+# ---------------------------------------------------------------------------
+# Two-stage mapper x mapper threshold (PyAutoArray#507 step 2)
+# ---------------------------------------------------------------------------
+#
+# `curvature_matrix_via_sparse_operator_from` has two implementations with
+# identical contracts (see their docstrings): the direct quadruple loop, whose
+# cost is `sum_pairs u0 * u1` irregular read-modify-writes and does not depend
+# on `pix_pixels`; and the two-stage source-space accumulator, whose cost grows
+# linearly in `pix_pixels` but is dominated by contiguous, vectorisable AXPYs.
+# The two-stage form therefore must lose eventually, and this is the measured
+# point at which it is switched off.
+#
+# Measured by sweeping `pix_pixels` on the two production HST geometries -- the
+# real 15361-pixel sparse operator and real mapper mappings, with only the
+# source-space extent varied (autolens_profiling, OMP_NUM_THREADS=1). Speed-up
+# of two-stage over direct:
+#
+#   pix_pixels          128    784   1250   2048   4096   6144   8192
+#   rectangular u0=4.0  3.10   2.93   2.47   2.10   1.36   1.12   1.04
+#   delaunay    u0=1.55 1.79   1.72   1.53   1.33   1.04   1.01   0.98
+#
+# So the crossover is geometry-dependent -- the narrow-mapping Delaunay
+# geometry reaches parity near 8192 source pixels while the wide-mapping
+# bilinear one is still ahead there -- and there is *no* crossover anywhere
+# inside the range any PyAuto pixelization is run at. The threshold below is
+# deliberately the conservative end of the measured envelope rather than a
+# fitted crossover: at 4096 both geometries still favour the two-stage kernel,
+# and above it the two forms are within a few per cent of each other anyway, so
+# picking the direct one costs nothing. It is *not* a silent heuristic: it is a
+# single explicit constant, and it can be overridden per call through the
+# `two_stage_max_pix_pixels` argument (which is how both branches are tested).
+#
+# For scale: a 4096-pixel source makes `F` a 134 MB dense matrix whose Cholesky
+# factorisation alone is ~2e10 flops, so the inversion is long past being
+# dominated by this kernel by the time the threshold is in play.
+CURVATURE_TWO_STAGE_MAX_PIX_PIXELS = 4096
+
+
 @numba_util.jit()
 def curvature_matrix_via_sparse_operator_from(
     psf_precision_operator: np.ndarray,
@@ -501,6 +539,95 @@ def curvature_matrix_via_sparse_operator_from(
     data_weights: np.ndarray,
     pix_lengths: np.ndarray,
     pix_pixels: int,
+    two_stage_max_pix_pixels: int = CURVATURE_TWO_STAGE_MAX_PIX_PIXELS,
+) -> np.ndarray:
+    """
+    Returns the mapper x mapper block of the curvature matrix `F` (see Warren & Dye
+    2003) from the sparse PSF precision operator.
+
+    This is the entry point the inversion calls. It selects between two
+    implementations with identical inputs, outputs and contracts:
+
+    - `curvature_matrix_via_sparse_operator_two_stage_from`, which accumulates a
+      dense source-space vector per data pixel and then writes whole rows of `F`
+      with contiguous AXPYs. Faster for every source-space size any PyAuto
+      pixelization uses (2.9x at the HST rectangular fiducial), because it
+      replaces `sum_pairs u0 * u1` irregular read-modify-writes into a multi-MB
+      matrix with `sum_pairs u1` L1-resident scatters plus vectorisable dense
+      adds.
+    - `curvature_matrix_via_sparse_operator_direct_from`, the quadruple loop,
+      whose cost does not grow with `pix_pixels` and which therefore takes over
+      above `two_stage_max_pix_pixels`.
+
+    The two agree to floating-point reassociation, not bit-identically: the
+    two-stage form sums the same products in a different order. The measured
+    maximum relative difference on the production HST geometries is ~4e-13,
+    three orders inside the `rtol=1e-6` the likelihood is pinned at.
+
+    Parameters
+    ----------
+    psf_precision_operator
+        A matrix that precomputes the values for fast computation of the curvature matrix in a memory efficient way.
+    psf_precision_indexes
+        The image-pixel indexes of the values stored in the w tilde preload matrix, which are used to compute
+        the weights of the data values when computing the curvature matrix.
+    psf_precision_lengths
+        The number of image pixels in every row of `psf_precision_operator`, which is iterated over when computing the
+        curvature matrix.
+    data_to_pix_unique
+        An array that maps every data pixel index (e.g. the masked image pixel indexes in 1D) to its unique set of
+        pixelization pixel indexes (see `data_slim_to_pixelization_unique_from`).
+    data_weights
+        For every unique mapping between a set of data sub-pixels and a pixelization pixel, the weight of these mapping
+        based on the number of sub-pixels that map to pixelization pixel.
+    pix_lengths
+        A 1D array describing how many unique pixels each data pixel maps too, which is used to iterate over
+        `data_to_pix_unique` and `data_weights`.
+    pix_pixels
+        The total number of pixels in the pixelization that reconstructs the data.
+    two_stage_max_pix_pixels
+        The largest `pix_pixels` for which the two-stage implementation is used;
+        above it the direct loop is. Defaults to the measured
+        `CURVATURE_TWO_STAGE_MAX_PIX_PIXELS`; exposed so that both branches can
+        be exercised without building a source space of that size.
+
+    Returns
+    -------
+    ndarray
+        The curvature matrix `F` (see Warren & Dye 2003).
+    """
+
+    if pix_pixels <= two_stage_max_pix_pixels:
+        return curvature_matrix_via_sparse_operator_two_stage_from(
+            psf_precision_operator=psf_precision_operator,
+            psf_precision_indexes=psf_precision_indexes,
+            psf_precision_lengths=psf_precision_lengths,
+            data_to_pix_unique=data_to_pix_unique,
+            data_weights=data_weights,
+            pix_lengths=pix_lengths,
+            pix_pixels=pix_pixels,
+        )
+
+    return curvature_matrix_via_sparse_operator_direct_from(
+        psf_precision_operator=psf_precision_operator,
+        psf_precision_indexes=psf_precision_indexes,
+        psf_precision_lengths=psf_precision_lengths,
+        data_to_pix_unique=data_to_pix_unique,
+        data_weights=data_weights,
+        pix_lengths=pix_lengths,
+        pix_pixels=pix_pixels,
+    )
+
+
+@numba_util.jit()
+def curvature_matrix_via_sparse_operator_direct_from(
+    psf_precision_operator: np.ndarray,
+    psf_precision_indexes: np.ndarray,
+    psf_precision_lengths: np.ndarray,
+    data_to_pix_unique: np.ndarray,
+    data_weights: np.ndarray,
+    pix_lengths: np.ndarray,
+    pix_pixels: int,
 ) -> np.ndarray:
     """
     Returns the curvature matrix `F` (see Warren & Dye 2003) by computing it using `psf_precision_operator`
@@ -554,6 +681,203 @@ def curvature_matrix_via_sparse_operator_from(
 
     curvature_index = 0
 
+    for data_0 in range(data_pixels):
+
+        # Hoisted once per data pixel and once per stored pair (PyAutoArray#507
+        # step 1). The unhoisted form re-read `data_weights[data_1, pix_1_index]`
+        # and `data_to_pix_unique[data_1, pix_1_index]` inside the innermost
+        # loop, so every one of `data_1`'s u1 mappings was gathered u0 times from
+        # a wide-stride 2-D array; and `curvature_matrix[pix_0, pix_1]` re-did the
+        # row offset multiply on every accumulation. Taking 1-D row views and the
+        # `curvature_matrix[pix_0]` row view leaves the arithmetic untouched --
+        # the accumulated expression below is operand-for-operand the one it
+        # replaces, so the result is bit-identical, which
+        # `curvature_matrix_via_sparse_operator_reference_from` and its test
+        # assert directly.
+        pix_lengths_0 = pix_lengths[data_0]
+        pix_row_0 = data_to_pix_unique[data_0]
+        weight_row_0 = data_weights[data_0]
+
+        for data_1_index in range(psf_precision_lengths[data_0]):
+            data_1 = psf_precision_indexes[curvature_index]
+            psf_precision_value = psf_precision_operator[curvature_index]
+
+            curvature_index += 1
+
+            pix_lengths_1 = pix_lengths[data_1]
+            pix_row_1 = data_to_pix_unique[data_1]
+            weight_row_1 = data_weights[data_1]
+
+            for pix_0_index in range(pix_lengths_0):
+                data_0_weight = weight_row_0[pix_0_index]
+                curvature_row = curvature_matrix[pix_row_0[pix_0_index]]
+
+                for pix_1_index in range(pix_lengths_1):
+                    curvature_row[pix_row_1[pix_1_index]] += (
+                        data_0_weight
+                        * weight_row_1[pix_1_index]
+                        * psf_precision_value
+                    )
+
+    for i in range(pix_pixels):
+        for j in range(i, pix_pixels):
+            curvature_matrix[i, j] += curvature_matrix[j, i]
+
+    for i in range(pix_pixels):
+        for j in range(i, pix_pixels):
+            curvature_matrix[j, i] = curvature_matrix[i, j]
+
+    return curvature_matrix
+
+
+@numba_util.jit()
+def curvature_matrix_via_sparse_operator_two_stage_from(
+    psf_precision_operator: np.ndarray,
+    psf_precision_indexes: np.ndarray,
+    psf_precision_lengths: np.ndarray,
+    data_to_pix_unique: np.ndarray,
+    data_weights: np.ndarray,
+    pix_lengths: np.ndarray,
+    pix_pixels: int,
+) -> np.ndarray:
+    """
+    The mapper x mapper block of `F`, computed in two stages via a dense
+    source-space accumulator (PyAutoArray#507 step 2). Same inputs, same outputs
+    and same halved-diagonal / `A + A.T` contract as
+    `curvature_matrix_via_sparse_operator_direct_from`; only the order of the
+    summation differs, so the two agree to floating-point reassociation
+    (`rtol=1e-6` is the pinned tolerance; the observed difference is far smaller)
+    rather than bit-identically.
+
+    The direct kernel's innermost work is
+
+        for each stored pair (data_0, data_1):
+            for each of data_0's u0 mappings (pix_0, w0):
+                for each of data_1's u1 mappings (pix_1, w1):
+                    F[pix_0, pix_1] += w0 * w1 * W(data_0, data_1)
+
+    -- `u0 * u1` irregular read-modify-writes per stored pair, scattered across a
+    `pix_pixels ** 2` matrix that at HST resolution is 4.9 MB, far outside L2.
+    Measured at the HST rectangular fiducial that is 1.77e8 accumulations at
+    ~1.4 ns each.
+
+    The same sum factorises, because `w0` does not depend on `data_1`:
+
+        a[p] = sum over data_0's stored pairs of W(data_0, data_1) * w1(data_1, p)
+        F[pix_0, :] += w0 * a[:]     for each of data_0's u0 mappings
+
+    Stage 1 accumulates `a` -- one scatter per (stored pair, mapping of data_1),
+    i.e. `u1` rather than `u0 * u1` -- into an L1-resident `pix_pixels` vector.
+    Stage 2 is `u0` contiguous, vectorisable AXPYs of length `pix_pixels` into
+    whole rows of `F`.
+
+    The trade is `sum_pairs u0 * u1` scattered RMWs for
+    `sum_pairs u1` L1 scatters plus `(sum_data u0 + data_pixels) * pix_pixels`
+    dense ops. It is therefore *not* unconditionally faster: it wins when the
+    source space is small relative to the PSF overlap and the mappings are wide
+    (bilinear, u0 = 4), and loses when the source space is large and the mappings
+    narrow (barycentric Delaunay, u0 ~ 1.55). `curvature_matrix_via_sparse_operator_from`
+    chooses between the two from the measured geometry -- see
+    `CURVATURE_TWO_STAGE_COST_RATIO_THRESHOLD`.
+
+    Parameters and return value are exactly those of
+    `curvature_matrix_via_sparse_operator_direct_from`.
+    """
+
+    data_pixels = psf_precision_lengths.shape[0]
+
+    curvature_matrix = np.zeros((pix_pixels, pix_pixels))
+
+    # The per-data-pixel source-space accumulator, allocated once and zeroed
+    # after every data pixel (it is left clean by the loop below, so the
+    # allocation's own zeros are correct for the first data pixel).
+    source_accumulator = np.zeros(pix_pixels)
+
+    curvature_index = 0
+
+    for data_0 in range(data_pixels):
+
+        pair_length_0 = psf_precision_lengths[data_0]
+        pix_lengths_0 = pix_lengths[data_0]
+
+        if pix_lengths_0 == 0:
+            # No mappings for this data pixel: it contributes nothing, but its
+            # stored pairs must still be stepped over.
+            curvature_index += pair_length_0
+            continue
+
+        # -- stage 1: accumulate the dense source-space vector ---------------
+        for data_1_index in range(pair_length_0):
+            data_1 = psf_precision_indexes[curvature_index]
+            psf_precision_value = psf_precision_operator[curvature_index]
+
+            curvature_index += 1
+
+            pix_row_1 = data_to_pix_unique[data_1]
+            weight_row_1 = data_weights[data_1]
+
+            for pix_1_index in range(pix_lengths[data_1]):
+                source_accumulator[pix_row_1[pix_1_index]] += (
+                    weight_row_1[pix_1_index] * psf_precision_value
+                )
+
+        # -- stage 2: one contiguous AXPY per mapping of data_0 --------------
+        pix_row_0 = data_to_pix_unique[data_0]
+        weight_row_0 = data_weights[data_0]
+
+        for pix_0_index in range(pix_lengths_0):
+            data_0_weight = weight_row_0[pix_0_index]
+            curvature_row = curvature_matrix[pix_row_0[pix_0_index]]
+
+            for pix_1 in range(pix_pixels):
+                curvature_row[pix_1] += data_0_weight * source_accumulator[pix_1]
+
+        for pix_1 in range(pix_pixels):
+            source_accumulator[pix_1] = 0.0
+
+    for i in range(pix_pixels):
+        for j in range(i, pix_pixels):
+            curvature_matrix[i, j] += curvature_matrix[j, i]
+
+    for i in range(pix_pixels):
+        for j in range(i, pix_pixels):
+            curvature_matrix[j, i] = curvature_matrix[i, j]
+
+    return curvature_matrix
+
+
+@numba_util.jit()
+def curvature_matrix_via_sparse_operator_reference_from(
+    psf_precision_operator: np.ndarray,
+    psf_precision_indexes: np.ndarray,
+    psf_precision_lengths: np.ndarray,
+    data_to_pix_unique: np.ndarray,
+    data_weights: np.ndarray,
+    pix_lengths: np.ndarray,
+    pix_pixels: int,
+) -> np.ndarray:
+    """
+    The straightforward quadruple loop that `curvature_matrix_via_sparse_operator_from`
+    is an optimisation of, retained verbatim as the reference the optimised kernel is
+    tested against.
+
+    It is not called by the inversion. It exists so that the restructurings applied to
+    the production kernel -- the hoisted row views of PyAutoArray#507 step 1, and the
+    two-stage source-space accumulator of step 2 -- have a fixed, unrestructured
+    definition to be asserted against, rather than only the dense `M.T W M` oracle
+    (which is a far smaller problem than the kernel is ever run on, since it
+    materializes the [image_pixels, image_pixels] precision operator).
+
+    Parameters and return value are exactly those of
+    `curvature_matrix_via_sparse_operator_from`.
+    """
+
+    data_pixels = psf_precision_lengths.shape[0]
+
+    curvature_matrix = np.zeros((pix_pixels, pix_pixels))
+
+    curvature_index = 0
+
     for data_0 in range(data_pixels):
         for data_1_index in range(psf_precision_lengths[data_0]):
             data_1 = psf_precision_indexes[curvature_index]