From 5e3fe2441229683829d33b7e06f9b21e282fd4d5 Mon Sep 17 00:00:00 2001 From: Peter Boyle Date: Wed, 19 Aug 2026 19:31:46 -0400 Subject: [PATCH] Recursive Schur complement based distributed matrix inverse by Claude after discussing plan and instructions from me --- .../multigrid/RecursiveSchurInverse.h | 82 ++++++------------- 1 file changed, 26 insertions(+), 56 deletions(-) diff --git a/Grid/algorithms/multigrid/RecursiveSchurInverse.h b/Grid/algorithms/multigrid/RecursiveSchurInverse.h index a46e6a964..be62822aa 100644 --- a/Grid/algorithms/multigrid/RecursiveSchurInverse.h +++ b/Grid/algorithms/multigrid/RecursiveSchurInverse.h @@ -36,34 +36,20 @@ NAMESPACE_BEGIN(Grid); // RecursiveSchurInverse: distributed dense inversion by recursive Schur // complement over a binary rank-range tree. // -// CONTRACT: the caller presents an N x N matrix in RANK-MAJOR row ordering, -// distributed by rows -- rank r owns global rows [rowStart[r], rowStart[r+1]) -// -- and receives its rows of the INVERSE in the same layout. This class -// knows nothing of lattices or coarse operators; it consumes a GridBase for -// world collectives, GridBLAS for GEMMs and GridBLASInverse for the leaf -// inversions (all of which have Eigen reference backends, so the whole -// algorithm unit-tests on a CPU-only laptop build under mpirun). +// Contract: rank r owns global rows [rowStart[r], rowStart[r+1]) of an +// N x N matrix in rank-major ordering, and receives its rows of the +// inverse in the same layout. Consumes GridBase collectives, GridBLAS +// and GridBLASInverse only; Eigen reference backends permit CPU unit +// testing under mpirun (Test_schur_inverse). // -// PRECISION (decision 2026-08-14, superseding the fp32-merge design): the -// ENTIRE inversion runs in fp64 (ComplexD). The apply-side fp32 gain is -// taken where it matters -- inside the iterative process -- by rounding the -// finished inverse ONCE when the caller stores it in the fp32 apply slab. -// Consequences: merge-growth error accumulates in eps64 and the terminal -// rounding gives representation-only ~eps32 accuracy independent of growth; -// the Newton-Schulz refinement and the fp32 escalation ladder are DELETED -// (resurrectable from git history if a future scale forces reduced-precision -// merges). Setup cost: ~2x panel-gather bytes and ~2x transient memory, -// once per setup; fp64 GEMM runs at fp32 rate on CDNA2/PVC. +// Arithmetic is fp64 throughout; the caller rounds once into fp32 storage. // -// EXECUTION MODEL: SPMD full-tree walk. Every rank executes the identical -// recursion call sequence; participation in DATA is ownership-gated, and -// every collective is a world-communicator zero-fill GlobalSumVector. No -// sub-communicators exist, so no deadlock surface exists. +// Execution: SPMD full-tree walk. Every rank makes the identical call +// sequence; data participation is ownership-gated; all collectives are +// world-wide, so no deadlock surface exists. // -// STORAGE CONVENTION (pinned by unit test T1b, Test_schur_inverse.cc): -// BlockRows is COLUMN-MAJOR with ld = rows, matching the BLAS world: -// element (i,j) lives at data[ i + j*ld ]; a column window [col0, col0+w) -// is the contiguous slice starting at data[ col0*ld ]. +// Storage: BlockRows is column-major, ld = rows; element (i,j) at +// data[i + j*ld]; a column window is the contiguous slice at data[col0*ld]. /////////////////////////////////////////////////////////////////////////////// /////////////////////////////////////////////////////////////////////////////// @@ -113,14 +99,12 @@ public: GridBLAS BLAS; GridBLASInverse INV; - // Growth telemetry (diagnostic, not load-bearing at fp64): one entry per - // merge node, walk order + // Growth telemetry: one entry per merge node, walk order std::vector telNormB; // ||B||_F = ||A11inv A12||_F std::vector telSratio; // ||S||_F / ||A22||_F double telLeafMaxInv; // max |(leaf inverse)_ij| over leaves - // Phase timers/counters (this rank), accumulated across the whole walk: - // where does the setup wall actually go? Reported by ReportTelemetry. + // Phase timers/counters, reported by ReportTelemetry double tMemset; // device panel zero-fill double tDeposit; // owner rows -> panel (device kernel) double tAllreduce; // GlobalSumVector on device panels @@ -132,17 +116,8 @@ public: uint64_t nAllreduce; // panel collectives uint64_t nGatherGemm; // GatherGemm calls - // PERSISTENT DEVICE panel, grow-only across the whole walk. - // Persistent: fresh per-call allocations defeat the MPI registration - // cache (measured 13% at N=138k). DEVICE-RESIDENT: the collective runs - // on the device pointer via GPU-aware MPI, and panel assembly (zero-fill - // + owner deposit) is device-side -- the entire host round-trip (stage - // D2H, 92s memset, deposit, H2D at N=138k) is deleted. Measured - // motivation (MPI_benchmark/gather_mpi, Frontier, 16-128 ranks): - // host-buffer allreduce 0.8-1.3 GB/s vs device-buffer 5.9-6.8 GB/s; - // device allgatherv/ring 20-24 GB/s is the planned stage-2 pattern. - // DEVICE BUILDS REQUIRE GPU-aware MPI (MPICH_GPU_SUPPORT_ENABLED=1, - // standard in production); CPU builds are unaffected (device == host). + // Persistent grow-only device panel; assembly and collectives are + // device-resident. Device builds require GPU-aware MPI. deviceVector dPanelBuf; /////////////////////////////////////////////////////////////////////////// @@ -207,7 +182,7 @@ public: } /////////////////////////////////////////////////////////////////////////// - // THE communication primitive (plan 3.4 / 4B.3). + // The communication primitive. // // C(:, colC : colC+widthB) <- beta * C(:, colC : colC+widthB) // + alpha * A(:, colA : colA+widthA) * Bsub @@ -215,15 +190,12 @@ public: // Bsub is the widthA x widthB sub-block of a row-distributed operand // owned by ranks [rB0, rB1): owner r contributes its rows of // B(:, colB : colB+widthB) at sub-block row offset - // rowStart[r] - rowStart[rB0]. The sub-block is gathered in panelBytes - // row-chunks, DEVICE-NATIVE: device zero-fill + device deposit kernel + - // world GlobalSumVector on the device pointer (GPU-aware MPI on device - // builds; see the dPanelBuf comment for the measured motivation). + // rowStart[r] - rowStart[rB0], gathered in panelBytes row-chunks by + // device zero-fill + deposit kernel + GlobalSumVector. // - // SPMD rules: EVERY rank calls (the collectives are world-wide); - // non-owners of B add zeros; ranks with A.rows == 0 skip all local - // compute but still make every collective call. Column offsets are - // LOCAL buffer offsets -- non-participants pass 0. + // Every rank calls; non-owners of B add zeros; ranks with A.rows == 0 + // skip local compute but make every collective call. Column offsets + // are local buffer offsets -- non-participants pass 0. /////////////////////////////////////////////////////////////////////////// void GatherGemm(ComplexD alpha, BlockRows &A, int64_t colA, int64_t widthA, @@ -280,6 +252,8 @@ public: { int64_t kchunk = std::min(kc, k-k0); + // Zeroing must complete before MPI reads the panel; it is the sole + // producer on non-owner ranks. tMemset -= usecond(); acceleratorMemSet(&dPanel[0], 0, (uint64_t)kchunk*n*sizeof(ComplexD)); tMemset += usecond(); @@ -289,8 +263,7 @@ public: int64_t i1 = std::min(k0+kchunk, myOff+B.rows); if ( i1 > i0 ) { - // Device-side deposit straight from the B window: strided - // block copy as a flat kernel (len rows of n columns). + // Deposit my rows: strided block copy, len rows x n columns int64_t len = i1-i0; int64_t brows = B.rows; int64_t dof = i0-k0; @@ -380,11 +353,8 @@ public: } /////////////////////////////////////////////////////////////////////////// - // Leaf inversion. Purely LOCAL -- the calling rank owns the whole - // width x width leaf (width == my row count); no collectives, so the - // SPMD walk stays uniform with other ranks doing nothing. The window - // is contiguous (ld == rows == width): invert IN PLACE via - // GridBLASInverse. Everything is already fp64; no promote/demote. + // Leaf inversion: local, in place on the contiguous diagonal window. + // No collectives. /////////////////////////////////////////////////////////////////////////// void LeafInvert(int64_t col0, int64_t width, BlockRows &Arows) {