diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 69f8866b5..c2c09d408 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -843,6 +843,11 @@ the [porting order](#porting-order); part of #688). with device permutations and batched sparse triangular solves. Small 1d factors and the pole Schur complement use cached dense inverses and device matrix products. The full poloidal inverse is never made dense. Projectors support `schur`, `lu`, `ilu`, boundary conditions and transposed solves on both backends. +- **Shared factorizations.** Polar spaces belonging to one Derham share identical 1d factors and poloidal + solvers, including their cached device copies. Keys use exact canonical sparse matrix contents, dtype, + solver kind/settings, Schur pole mask and outer factors, and backend/device. Changed weights and boundary + matrices get different factors; no approximate matching is used. Weak entries disappear when no live solver + uses them, so repeated weight updates do not retain obsolete factorizations. Separate Derham objects are isolated. - **Mass setup.** Callable weights are evaluated with coordinates on the active backend. Sampled weights and spline quadrature data are copied explicitly to the host for the 2d reduction/assembly and factorization. These are setup transfers; repeated operator/preconditioner applications have no coefficient transfers on @@ -917,5 +922,4 @@ fallback for that kernel. They take only arrays and scalars, no argument objects runs with regular MPI too, without involving feectools halos. - **Separate follow-ups.** `SplineFunction.__call__` already dispatches CUDA kernels for marker and dense/sparse grid evaluation. Remaining evaluation work is the host-only `eval_tp_fixed_loc` helper and distributed - field-evaluation MPI handling/coverage. FFT eta3 solves and sharing factorizations across different projector - objects are additional optimizations. + field-evaluation MPI handling/coverage. FFT eta3 solves are an additional optimization. diff --git a/doc/sections/subsections/numerics-geomFE.rst b/doc/sections/subsections/numerics-geomFE.rst index 59a90211f..fdce1b521 100644 --- a/doc/sections/subsections/numerics-geomFE.rst +++ b/doc/sections/subsections/numerics-geomFE.rst @@ -291,7 +291,9 @@ it is exact for axisymmetric weights and mappings. On the CuPy backend, extraction, derivatives, Kronecker operators and these solves keep their coefficients on the device. Factorization and mass-preconditioner assembly run once on the host; -cached device factors apply subsequent solves. Polar MPI collectives support CUDA-aware device +cached device factors apply subsequent solves. Identical factors are shared across polar spaces +and projector/preconditioner instances of the same Derham. Matrix contents and solver settings +are matched exactly; updated weights use new factors, and unused entries are released. Polar MPI collectives support CUDA-aware device buffers or host staging through ``cunumpy.mpi.mpi_buffer``. The feectools halo exchanges require CUDA-aware MPI on real GPUs. ``SplineFunction.__call__`` dispatches CUDA kernels for marker and dense/sparse grid evaluation. The fixed local-grid helper ``eval_tp_fixed_loc`` remains diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index d11f4aa5d..5f41e4696 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import logging +from weakref import WeakValueDictionary import cunumpy as xp import feectools.core.bsplines as bsp @@ -809,6 +810,8 @@ def __init__( assert self.domain is not None self._ck_blocks = PolarExtractionBlocksC1(self.domain, self) + # Polar spaces share live factors; obsolete mass-weight factors are not retained. + self._polar_factorization_cache = WeakValueDictionary() V0pol, V1pol, V2pol, V3pol, Vvpol = self._assemble_polar_extraction_operators() self._V0pol = V0pol diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index b6fb1fb6a..eeb3d1e94 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -35,6 +35,7 @@ def __init__(self, derham: "Derham", space_id: str): self._dtype = float self._comm = derham.comm self._space_id = space_id + self._factorization_cache = derham._polar_factorization_cache # dimensions of 1d spaces self._n = derham.spline_attributes["H1"].nbasis[0] diff --git a/src/struphy/polar/kron_2dx1d.py b/src/struphy/polar/kron_2dx1d.py index 6c89c80a6..987e228d4 100644 --- a/src/struphy/polar/kron_2dx1d.py +++ b/src/struphy/polar/kron_2dx1d.py @@ -215,6 +215,41 @@ def poloidal_extraction_matrix( # -------------------------------------------------------------------------------------- # Poloidal solvers # -------------------------------------------------------------------------------------- +def _matrix_key(matrix): + """Exact numerical contents, independent of sparse storage order and explicit zeros.""" + matrix = spa.csc_matrix(matrix, copy=True) + matrix.sum_duplicates() + matrix.eliminate_zeros() + matrix.sort_indices() + return ( + matrix.shape, + matrix.dtype.str, + matrix.indptr.astype(np.int64, copy=False).tobytes(), + matrix.indices.astype(np.int64, copy=False).tobytes(), + matrix.data.tobytes(), + ) + + +def _cache_context(): + """Keep device copies separate when solvers are built on different backends/devices.""" + backend = xp.get_backend() + if backend == "cupy": + import cupy + + return backend, cupy.cuda.runtime.getDevice() + return backend, None + + +def _factorized_1d(matrix, cache=None): + key = ("1d", _cache_context(), _matrix_key(matrix)) + factor = None if cache is None else cache.get(key) + if factor is None: + factor = _FactorizedSolve(splu(spa.csc_matrix(matrix)), dense=True) + if cache is not None: + cache[key] = factor + return factor + + class _FactorizedSolve: """Host factorization with a cached device solve; dense inverses are only used for small 1d factors.""" @@ -260,15 +295,15 @@ def solve(self, Y: np.ndarray) -> np.ndarray: class _KroneckerBlockDiagonal: """Exact solver for a block diagonal matrix with Kronecker blocks A1 x A2 (C-ordering of the unknowns).""" - def __init__(self, factors: list[tuple[np.ndarray, np.ndarray]]): + def __init__(self, factors: list[tuple[np.ndarray, np.ndarray]], cache=None): self._sizes = [] self._lus = [] for A1, A2 in factors: self._sizes.append((A1.shape[0], A2.shape[0])) self._lus.append( ( - _FactorizedSolve(splu(spa.csc_matrix(A1)), dense=True), - _FactorizedSolve(splu(spa.csc_matrix(A2)), dense=True), + _factorized_1d(A1, cache), + _factorized_1d(A2, cache), ) ) self.exact = True @@ -356,6 +391,7 @@ def make_poloidal_solver( schur_outer: str = "lu", drop_tol: float = ILU_DROP_TOL, fill_factor: float = ILU_FILL_FACTOR, + _cache=None, ): """ Solver for the poloidal matrix A. @@ -382,23 +418,48 @@ def make_poloidal_solver( Parameters of the incomplete LU factorization (``scipy.sparse.linalg.spilu``). """ assert kind in POLOIDAL_SOLVERS, f"Unknown poloidal solver {kind!r}, must be one of {POLOIDAL_SOLVERS}." - if kind in ("lu", "ilu"): - return _SparseLU(A, kind, drop_tol=drop_tol, fill_factor=fill_factor) + if kind == "schur": + assert pole is not None + settings = ( + np.asarray(pole, dtype=bool).tobytes(), + None if outer_kronecker is None else tuple((_matrix_key(a), _matrix_key(b)) for a, b in outer_kronecker), + schur_outer, + float(drop_tol), + float(fill_factor), + ) + else: + settings = (float(drop_tol), float(fill_factor)) if kind == "ilu" else () + key = ("poloidal", _cache_context(), kind, _matrix_key(A), settings) + if _cache is not None: + cached = _cache.get(key) + if cached is not None: + return cached - assert pole is not None - if outer_kronecker is not None: - outer = _KroneckerBlockDiagonal(outer_kronecker) - if __debug__: - # the outer block must match the Kronecker factors - o = np.flatnonzero(~pole) - Aoo = spa.csr_matrix(A)[o][:, o] - K = spa.block_diag([spa.kron(spa.csr_matrix(A1), spa.csr_matrix(A2)) for A1, A2 in outer_kronecker]) - diff = abs(Aoo - K).max() if Aoo.nnz + K.nnz > 0 else 0.0 - assert diff <= 1e-12 * max(abs(Aoo).max(), 1.0), f"Outer block is not Kronecker ({diff = })." + if kind in ("lu", "ilu"): + solver = _SparseLU(A, kind, drop_tol=drop_tol, fill_factor=fill_factor) else: - o = np.flatnonzero(~pole) - outer = _SparseLU(spa.csr_matrix(A)[o][:, o], schur_outer, drop_tol=drop_tol, fill_factor=fill_factor) - return _SchurSolver(A, pole, outer) + if outer_kronecker is not None: + outer = _KroneckerBlockDiagonal(outer_kronecker, cache=_cache) + if __debug__: + # the outer block must match the Kronecker factors + o = np.flatnonzero(~pole) + Aoo = spa.csr_matrix(A)[o][:, o] + K = spa.block_diag([spa.kron(spa.csr_matrix(A1), spa.csr_matrix(A2)) for A1, A2 in outer_kronecker]) + diff = abs(Aoo - K).max() if Aoo.nnz + K.nnz > 0 else 0.0 + assert diff <= 1e-12 * max(abs(Aoo).max(), 1.0), f"Outer block is not Kronecker ({diff = })." + else: + o = np.flatnonzero(~pole) + outer = make_poloidal_solver( + spa.csr_matrix(A)[o][:, o], + schur_outer, + drop_tol=drop_tol, + fill_factor=fill_factor, + _cache=_cache, + ) + solver = _SchurSolver(A, pole, outer) + if _cache is not None: + _cache[key] = solver + return solver def pole_mask(space: PolarDerhamSpace, comps: tuple[int, ...]) -> np.ndarray: @@ -732,13 +793,14 @@ def __init__( super().__init__(space, groups) self._kind = poloidal_solver self._solver_kwargs = solver_kwargs - self._lu_3 = [_FactorizedSolve(splu(spa.csc_matrix(g.A_3)), dense=True) for g in groups] + self._lu_3 = [_factorized_1d(g.A_3, space._factorization_cache) for g in groups] self._pol_solvers = [ make_poloidal_solver( g.A_pol, poloidal_solver, pole=g.pole if g.pole is not None else pole_mask(space, g.comps), outer_kronecker=g.outer_kronecker, + _cache=space._factorization_cache, **solver_kwargs, ) for g in groups diff --git a/src/struphy/polar/tests/test_factorization_cache.py b/src/struphy/polar/tests/test_factorization_cache.py new file mode 100644 index 000000000..31db5b8b2 --- /dev/null +++ b/src/struphy/polar/tests/test_factorization_cache.py @@ -0,0 +1,176 @@ +"""Reuse polar factors without retaining obsolete matrices or confusing solver configurations.""" + +import gc +import subprocess +import sys +from weakref import WeakValueDictionary + +import cunumpy as xp +import numpy as np +import pytest +import scipy.sparse as sparse + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.polar import kron_2dx1d as kron +from struphy.polar.tests.test_polar_cupy import DenseDeviceSparse, fill, make_polar_derham, to_host + + +@pytest.mark.parametrize("kind", ("schur", "lu", "ilu")) +def test_poloidal_factors_match_contents_and_settings(kind): + cache = WeakValueDictionary() + A = sparse.csr_matrix([[4.0, 1.0, 0.0], [2.0, 5.0, 1.0], [0.0, 1.0, 3.0]]) + pole = np.array([True, False, False]) + first = kron.make_poloidal_solver(A, kind, pole=pole, _cache=cache) + same = kron.make_poloidal_solver(A.tocsc(), kind, pole=pole.copy(), _cache=cache) + assert first is same + unsorted = sparse.coo_matrix( + ([4.0, 0.5, 0.5, 2.0, 5.0, 1.0, 1.0, 3.0, 0.0], ([0, 0, 0, 1, 1, 1, 2, 2, 0], [0, 1, 1, 0, 1, 2, 1, 2, 2])), + shape=(3, 3), + ) + assert kron.make_poloidal_solver(unsorted, kind, pole=pole, _cache=cache) is first + # Mutating the original matrix must never retrieve stale factors. + A.data *= 2.0 + changed = kron.make_poloidal_solver(A, kind, pole=pole, _cache=cache) + assert changed is not first + other_kind = "ilu" if kind != "ilu" else "lu" + assert kron.make_poloidal_solver(A, other_kind, pole=pole, _cache=cache) is not changed + rhs = np.arange(6.0).reshape(3, 2) + np.testing.assert_allclose(changed.solve(rhs), first.solve(rhs) / 2, rtol=1e-12, atol=1e-12) + transposed = kron.make_poloidal_solver(A.T, kind, pole=pole, _cache=cache) + assert transposed is not changed + np.testing.assert_allclose(A.T @ transposed.solve(rhs), rhs, atol=1e-12) + if kind == "ilu": + different_options = kron.make_poloidal_solver(A, kind, drop_tol=1e-2, pole=pole, _cache=cache) + assert different_options is not changed + elif kind == "schur": + other_pole = np.array([False, True, False]) + different_options = kron.make_poloidal_solver(A, kind, pole=other_pole, _cache=cache) + assert different_options is not changed + approximate = kron.make_poloidal_solver(A, kind, pole=pole, schur_outer="ilu", _cache=cache) + assert approximate is not changed and not approximate.exact + outer = [(np.eye(1), A.toarray()[1:, 1:])] + separable = kron.make_poloidal_solver(A, kind, pole=pole, outer_kronecker=outer, _cache=cache) + assert separable is not changed + np.testing.assert_allclose(A @ separable.solve(rhs), rhs, atol=1e-12) + + +def test_1d_factors_reuse_exact_values_and_release_old_entries(monkeypatch): + cache = WeakValueDictionary() + calls = [] + original = kron.splu + + def factorize(matrix): + calls.append(matrix.shape) + return original(matrix) + + monkeypatch.setattr(kron, "splu", factorize) + A = np.array([[3.0, 1.0], [0.0, 2.0]]) + first = kron._factorized_1d(A, cache) + same = kron._factorized_1d(A.copy(), cache) + assert same is first and len(calls) == 1 + A[0, 0] += 1e-14 # no tolerance-based matching of distinct operators + changed = kron._factorized_1d(A, cache) + assert changed is not first and len(calls) == 2 + single_precision = kron._factorized_1d(A.astype(np.float32), cache) + assert single_precision is not changed and len(calls) == 3 + assert len(cache) == 3 + del first, same, changed, single_precision + gc.collect() + assert len(cache) == 0 + + +def test_projectors_share_factors_across_spaces_and_solver_instances(monkeypatch): + _, derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) + projectors = [getattr(derham, "P" + form) for form in ("0", "1", "2", "3")] + solvers = [P.pc for P in projectors] + assert all(s.domain._factorization_cache is solvers[0].domain._factorization_cache for s in solvers) + seen = {} + reused = 0 + for solver in solvers: + for group, factor in zip(solver.groups, solver._lu_3): + key = kron._matrix_key(group.A_3) + if key in seen: + assert factor is seen[key] + reused += 1 + seen[key] = factor + assert reused > 0 + + def unexpected_factorization(*args, **kwargs): + pytest.fail("an identical solver repeated the factorization") + + monkeypatch.setattr(kron, "splu", unexpected_factorization) + monkeypatch.setattr(kron, "lu_factor", unexpected_factorization) + duplicate = kron.PolarKron2Dx1DSolver(solvers[1].domain, solvers[1].groups) + assert duplicate._lu_3[0] is solvers[1]._lu_3[0] + assert duplicate._pol_solvers[0] is solvers[1]._pol_solvers[0] + _, other_derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) + assert other_derham.P0.I.domain._factorization_cache is not solvers[0].domain._factorization_cache + + +def test_mass_weight_update_reuses_1d_factors_and_rebuilds_poloidal_factors(): + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + _, _, mass = make_polar_derham("IGAPolarCylinder") + pc = PolarMassMatrixPreconditioner(mass.M0) + previous = pc.solver + updated = mass.create_weighted_mass( + "H1", "H1", weights=[[lambda e1, e2, e3: 1 + e1 + 0 * e2 + 0 * e3]], assemble=True + ) + pc.update_mass_operator(updated) + assert pc.solver._lu_3[0] is previous._lu_3[0] + assert pc.solver._pol_solvers[0] is not previous._pol_solvers[0] + x = updated._V_boundary_op.dot(fill(updated.domain.zeros(), seed=96)) + actual = pc.dot(updated.dot(x)) + for a, b in zip(to_host(actual, interior=True), to_host(x, interior=True)): + np.testing.assert_allclose(a, b, rtol=1e-10, atol=1e-11) + same_weights = pc.solver + pc.update_mass_operator(updated) + assert pc.solver._pol_solvers[0] is same_weights._pol_solvers[0] + + +def check_device_factor_reuse(fake=False): + from contextlib import nullcontext + + from struphy.pic.tests.cuda_emulation import emulated_launches + from struphy.polar.linear_operators import set_device_sparse_module + + previous = set_device_sparse_module(DenseDeviceSparse) if fake else None + cache = WeakValueDictionary() + matrix = np.array([[3.0, 1.0], [0.0, 2.0]]) + with xp.use_backend("numpy"): + host_factor = kron._factorized_1d(matrix, cache) + try: + with xp.use_backend("cupy"), emulated_launches() if fake else nullcontext(): + device_factor = kron._factorized_1d(matrix, cache) + assert device_factor is not host_factor + assert kron._factorized_1d(matrix.copy(), cache) is device_factor + np.testing.assert_allclose(xp.to_numpy(device_factor.solve(xp.asarray(matrix))), np.eye(2), atol=1e-12) + _, derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) + first = derham.P1.pc + x = fill(first.domain.zeros(), seed=91) + expected = first.dot(x) + duplicate = kron.PolarKron2Dx1DSolver(first.domain, first.groups) + assert duplicate._lu_3[0] is first._lu_3[0] + assert duplicate._pol_solvers[0] is first._pol_solvers[0] + with xp.profiling.count_transfers() as transfers: + actual = duplicate.dot(x) + assert transfers.total == 0, transfers.report() + for a, b in zip(to_host(actual), to_host(expected)): + np.testing.assert_allclose(a, b, rtol=1e-12, atol=1e-12) + finally: + if fake: + set_device_sparse_module(previous) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake cannot replace installed CuPy") +def test_device_factors_are_shared_fake_cupy(): + code = "from struphy.polar.tests.test_factorization_cache import check_device_factor_reuse; check_device_factor_reuse(True)" + result = subprocess.run( + [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr[-4000:] + + +@pytest.mark.skipif(not xp.cupy_available(), reason="requires a GPU") +def test_device_factors_are_shared_on_cupy(): + check_device_factor_reuse()