From 6f7a2ad7c7a5bb0db11e7f3164145e2745eb09d8 Mon Sep 17 00:00:00 2001 From: Peter Boyle Date: Wed, 12 Aug 2026 12:47:40 -0400 Subject: [PATCH] Dense coarse inverse wrapper class, bundling up a bunch of useful work --- Grid/algorithms/multigrid/DenseCoarseMatrix.h | 702 ++++++++++++++++++ 1 file changed, 702 insertions(+) create mode 100644 Grid/algorithms/multigrid/DenseCoarseMatrix.h diff --git a/Grid/algorithms/multigrid/DenseCoarseMatrix.h b/Grid/algorithms/multigrid/DenseCoarseMatrix.h new file mode 100644 index 000000000..5e1703822 --- /dev/null +++ b/Grid/algorithms/multigrid/DenseCoarseMatrix.h @@ -0,0 +1,702 @@ +/************************************************************************************* + + Grid physics library, www.github.com/paboyle/Grid + + Source file: ./lib/algorithms/multigrid/DenseCoarseMatrix.h + + Copyright (C) 2026 + +Author: Peter Boyle + + This program is free software; you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation; either version 2 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License along + with this program; if not, write to the Free Software Foundation, Inc., + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + + See the full license in the file "LICENSE" in the top level distribution directory +*************************************************************************************/ +/* END LEGAL */ +#pragma once + +#include +#include + +#include + +NAMESPACE_BEGIN(Grid); + +////////////////////////////////////////////////////////////////////////////////////// +// DenseCoarseMatrix: a coarsened operator treated as a DENSE matrix -- explicit, +// row-distributed A^{-1} of a GeneralCoarsenedMatrix. Library-grade successor of +// the example-local DistributedDenseInverse (Example_pvdagm_mrhs_3level_dense.cc, +// FROZEN as the regression baseline). +// +// What is new versus the example class: +// - Stencil -> dense DIRECT IMPORT. The coarse operator IS the dense matrix +// unrolled: Dense[(s,a),(s+shift_p,b)] += A[p][s]_{a,b}. Rows of my sites are +// assembled from purely LOCAL _A[p] data: no operator applies, no comms -- the +// O(N) probe assembly (93 s at N=69120) is retired. ACCUMULATE (+=) because on +// short axes distinct shifts wrap to the same neighbour. An IMPORT CERTIFICATE +// compares the dense apply against Op.M on a NON-CONSTANT vector (a constant one +// cannot see a shift-sign error); DENSE_IMPORT_SIGN=-1 flips the convention +// without recompiling. +// - Split-K apply through GridBLAS.gemmBatched with EXPLICIT leading dimensions +// (arXiv:2409.03904 fig 11): the tiny-output/huge-K GEMM Y = slab^T X becomes +// DENSE_SPLITK chunk-GEMMs by pointer offset into the resident slab (lda = N), +// partials reduced in one accelerator_for. Platform-agnostic: deviceVector + +// GridBLAS run the SAME code on HIP/CUDA/SYCL and CPU(Eigen). +// - deviceVector everywhere in the apply path; the ONE surviving naked-HIP block +// is the boss inversion buffer (quarantined below, documented). +// +// Setup: SLAB_FILE= loads per-rank . (header-guarded N/nrows/ +// nbasis -- the interchange format shared with the frozen example; the STEM must +// encode cfg/mass/blocking/nbasis, only the header is guarded). Absent: direct +// import -> import certificate -> chunked zero-fill+GlobalSum gather streamed to +// the boss GCD -> cgetrf_64 (ILP64) -> rows of A^{-1} via blocked identity +// cgetrs_64 + broadcast, each rank keeping the rows of its own sites -> save. +// VERIFY ||A Ainv x - x||/||x|| runs in BOTH paths (and now certifies the DEVICE +// slab + split-K path, since the single-RHS apply routes through the same core). +// +// Env: SLAB_FILE DENSE_SPLITK (default 32, snapped to a divisor of N) +// DENSE_DEVICE_SUM DENSE_IMPORT_SIGN DENSE_APPLY_PROFILE DENSE_CC_CHECK +// +// Eventual internal upgrade (unchanged surface): RecursiveSchur distributed +// factorisation replacing the single-GCD gather/invert, lifting BOTH the fp32 +// N ~ 90k boss-HBM ceiling AND the CC-grid 256-rank SIMD cap; leaves land on +// GridBLASInverse::inverseBatched. +// +// Tensor-depth agnostic: site scalar objects treated as contiguous ComplexD +// (iScalar wrappers add no data), so any MG level's coarse operator imports. +////////////////////////////////////////////////////////////////////////////////////// +template +class DenseCoarseMatrix : public LinearFunction::CoarseVector> { +public: + typedef GeneralCoarsenedMatrix GeneralCoarseOp; + typedef typename GeneralCoarseOp::CoarseVector Field; + typedef typename GeneralCoarseOp::CoarseMatrix CoarseMatrix; + using LinearFunction::operator(); + typedef typename Field::vector_object vobj; + typedef typename vobj::scalar_object sobj; + typedef typename CoarseMatrix::vector_object Mvobj; + typedef typename Mvobj::scalar_object Msobj; + + GeneralCoarseOp &_Op; // the coarse operator: stencil source + certificate oracle + GridBase *grid; + int nd; + int64_t N; // dense rank = gSites * nbasis + int lsites; // my local sites + int64_t nrows; // my rows = lsites * nbasis + std::vector myLcoor; // local coordinate of my site ss + std::vector myGsite; // global lex site index of my site ss + std::vector slab; // nrows x N row-major: A during setup, rows of A^{-1} after + + static const int64_t CHUNKROWS = 1024; // getrs harvest block (trsm efficiency + fewer broadcasts) + static const int MRHS_MAX = 32; + + // Apply machinery: resident slab + persistent buffers + AOT split-K pointers. + GridBLAS BLAS; + deviceVector dSlab; + deviceVector dX; // N x MRHS_MAX + deviceVector dY; // nrows x MRHS_MAX + deviceVector dPartial; // NK x (nrows x MRHS_MAX) + deviceVector aptrs; // slab K-chunk pointers (lda = N) + deviceVector xptrs; // X K-chunk pointers (ldb = N) + deviceVector cptrs; // partial buffers (ldc = nrows) + std::vector hX; + std::vector hY; + int NK; // split-K chunk count (divides N) + int devSum; + + DenseCoarseMatrix(GeneralCoarseOp &Op, GridBase *g) + : _Op(Op), grid(g) + { + GRID_ASSERT( sizeof(sobj) == nbasis*sizeof(ComplexD) ); + GRID_ASSERT( sizeof(Msobj) == nbasis*nbasis*sizeof(ComplexD) ); + GRID_ASSERT( grid == Op.Grid() ); + nd = grid->_ndimension; + N = grid->gSites() * nbasis; + lsites = grid->lSites(); + nrows = (int64_t)lsites * nbasis; + + std::cout << GridLogMessage << "DenseCoarseMatrix: N = " << N + << " (" << grid->gSites() << " sites x " << nbasis << ")" + << " rows/rank = " << nrows + << " slab = " << (double)nrows*N*sizeof(ComplexF)/1024./1024. << " MB/rank" + << std::endl; + + //////////////////////////////////////////////////////////////////// + // Enumerate my sites: local coords and global lexicographic indices + //////////////////////////////////////////////////////////////////// + Coordinate ldims = grid->LocalDimensions(); + Coordinate gdims = grid->GlobalDimensions(); + myLcoor.resize(lsites); + myGsite.resize(lsites); + for(int ss=0; ss_lstart[d] + lcoor[d]; + int64_t gsite; + Lexicographic::IndexFromCoor(gcoor, gsite, gdims); + myLcoor[ss] = lcoor; + myGsite[ss] = gsite; + } + + slab.resize((uint64_t)nrows * N); + + double t0 = usecond(); + //////////////////////////////////////////////////////////////////// + // 0. Slab cache: SLAB_FILE= -> per-rank raw file .. + // SAME format as the frozen example (interchange compatible). + //////////////////////////////////////////////////////////////////// + bool loaded = false; + char *sfile = getenv("SLAB_FILE"); + std::string slabfile; + if (sfile) { + slabfile = std::string(sfile) + "." + std::to_string(grid->ThisRank()); + FILE *f = fopen(slabfile.c_str(),"rb"); + if (f) { + int64_t hdr[4] = {0,0,0,0}; + GRID_ASSERT( fread(hdr,sizeof(int64_t),4,f) == 4 ); + GRID_ASSERT( hdr[0] == (int64_t)0x44454E5345 ); // magic "DENSE" + GRID_ASSERT( hdr[1] == N && hdr[2] == (int64_t)nrows && hdr[3] == (int64_t)nbasis ); + uint64_t nelem = (uint64_t)nrows * N; + GRID_ASSERT( fread(&slab[0], sizeof(ComplexF), nelem, f) == nelem ); + fclose(f); + loaded = true; + std::cout << GridLogMessage << "DenseCoarseMatrix: slab loaded from " + << slabfile << " -- skipping import/factor/solve" << std::endl; + } else { + std::cout << GridLogMessage << "DenseCoarseMatrix: slab cache " << slabfile + << " absent -- full setup, will write it" << std::endl; + } + } + if (!loaded) { + ImportDense(); // slab <- my rows of A (LOCAL, no comms) + ImportCertificate(); // dense apply == Op.M on a non-constant vector + InvertDense(); // slab <- my rows of A^{-1} + double t1 = usecond(); + if (sfile) { + FILE *f = fopen(slabfile.c_str(),"wb"); + GRID_ASSERT(f != nullptr); + int64_t hdr[4] = { (int64_t)0x44454E5345, N, (int64_t)nrows, (int64_t)nbasis }; + GRID_ASSERT( fwrite(hdr,sizeof(int64_t),4,f) == 4 ); + uint64_t nelem = (uint64_t)nrows * N; + GRID_ASSERT( fwrite(&slab[0], sizeof(ComplexF), nelem, f) == nelem ); + fclose(f); + std::cout << GridLogMessage << "DenseCoarseMatrix: slab written to " << slabfile << std::endl; + } + std::cout << GridLogMessage << "DenseCoarseMatrix: import+invert took " + << (t1-t0)/1.0e6 << " s" << std::endl; + } + + //////////////////////////////////////////////////////////////////// + // Device residency + persistent apply buffers + AOT split-K pointers + //////////////////////////////////////////////////////////////////// + { + uint64_t sbytes = (uint64_t)nrows * N * sizeof(ComplexF); + dSlab.resize((uint64_t)nrows*N); + acceleratorCopyToDevice(&slab[0],&dSlab[0],sbytes); + + // DENSE_SPLITK: requested chunk count, snapped DOWN to a divisor of N. + int req = getenv("DENSE_SPLITK") ? atoi(getenv("DENSE_SPLITK")) : 32; + if (req < 1) req = 1; + NK = 1; + for(int j=1;j<=req;j++) if ( (N % j) == 0 ) NK = j; + int64_t Kc = N / NK; + + dX.resize((uint64_t)N*MRHS_MAX); + dY.resize((uint64_t)nrows*MRHS_MAX); + dPartial.resize((uint64_t)NK*nrows*MRHS_MAX); + hX.resize((uint64_t)N*MRHS_MAX); + hY.resize((uint64_t)nrows*MRHS_MAX); + + aptrs.resize(NK); xptrs.resize(NK); cptrs.resize(NK); + std::vector h(NK); + for(int j=0;j dense import of MY ROWS of A (no comms): + // Dense[(s,a),(wrap(s+shift_p),b)] += A[p][s]_{a,b} + //////////////////////////////////////////////////////////////////// + void ImportDense(void) + { + double t = -usecond(); + Coordinate gdims = grid->GlobalDimensions(); + int sign = getenv("DENSE_IMPORT_SIGN") ? atoi(getenv("DENSE_IMPORT_SIGN")) : 1; + GRID_ASSERT( sign==1 || sign==-1 ); + + uint64_t nelem = (uint64_t)nrows * N; + thread_for(i, nelem, { slab[i] = ComplexF(0.0,0.0); }); + + for(int p=0; p<_Op.geom.npoint; p++){ + Coordinate shift = _Op.geom.shifts[p]; + autoView(Av, _Op._A[p], CpuRead); + thread_for(ss, lsites, { + Coordinate ncoor(nd); + for(int d=0; d_lstart[d] + myLcoor[ss][d] + sign*shift[d]; + ncoor[d] = (int)((g % gdims[d] + gdims[d]) % gdims[d]); + } + int64_t nsite; + Lexicographic::IndexFromCoor(ncoor, nsite, gdims); + Msobj m; + peekLocalSite(m, Av, myLcoor[ss]); + ComplexD *md = (ComplexD *)&m; + for(int a=0; adense import took " + << t/1.0e6 << " s (" << _Op.geom.npoint << " points, local, no comms)" << std::endl; + } + + //////////////////////////////////////////////////////////////////// + // 2. IMPORT CERTIFICATE: dense rows vs Op.M on a NON-CONSTANT vector. + // (Constant x has x[s+d]==x[s-d]: blind to a shift-sign error.) + //////////////////////////////////////////////////////////////////// + void ImportCertificate(void) + { + Field x(grid); Field Ax(grid); Field Dx(grid); + for(int ss=0; ss xh((uint64_t)N, ComplexD(0.0,0.0)); + for(int ss=0; ssGlobalSumVector(&xh[0], (int)N); + std::vector yh(nrows); + thread_for(r, nrows, { + ComplexD acc(0.0,0.0); + const ComplexF *row = &slab[(uint64_t)r * N]; + for(int64_t j=0; jIsBoss(); + std::vector Afull; +#ifdef GRID_HIP + // QUARANTINED naked HIP: the boss-only N^2 inversion buffer (34GB at + // N=65536) must come from raw HBM; EvictAll flushes the device-copy + // layer to make the window. (FreePool of the allocator free-list + // awaits the type-dispatched fix.) Confined to setup; the apply path + // is pure Grid primitives. + rocblas_float_complex *dA = nullptr; + rocblas_float_complex *dB = nullptr; + int64_t *dIpiv = nullptr; + uint64_t Abytes = (uint64_t)N * N * sizeof(ComplexF); + MemoryManager::EvictAll(); + if (boss) { + auto aerr = hipMalloc((void **)&dA, Abytes); + if (aerr != hipSuccess) { + std::cout << GridLogMessage << "DenseCoarseMatrix: hipMalloc of " + << Abytes/1024./1024./1024. << " GB FAILED -- reduce --device-mem" << std::endl; + GRID_ASSERT(aerr == hipSuccess); + } + std::cout << GridLogMessage << "DenseCoarseMatrix: device inversion buffer allocated (" + << Abytes/1024./1024./1024. << " GB)" << std::endl; + } +#else + if (boss) Afull.resize((uint64_t)N * N); +#endif + { + std::unordered_map rowmap; // global row -> my slab row + for(int ss=0; ss chunk((uint64_t)CHUNKROWS * N); + for(int64_t row0=0; row0second) * N; + uint64_t dst = (uint64_t)(r-row0) * N; + for(int64_t j=0;jGlobalSumVector(&chunk[0], (int)nelem); + if (boss) { +#ifdef GRID_HIP + GRID_ASSERT( hipMemcpy((char *)dA + (uint64_t)row0*N*sizeof(ComplexF), + &chunk[0], nelem*sizeof(ComplexF), + hipMemcpyHostToDevice) == hipSuccess ); +#else + uint64_t dst = (uint64_t)row0 * N; + for(uint64_t i=0;i LU of A^T. + //////////////////////////////////////////////////////////////////// + if (boss) { +#ifdef GRID_HIP + std::cout << GridLogMessage << "DenseCoarseMatrix: rocSOLVER cgetrf_64 (ILP64 LU) N=" << N + << " in place on resident device buffer" << std::endl; + rocblas_handle handle = GridBLASInverse::Handle(); + int64_t *dInfo; + GRID_ASSERT( hipMalloc((void **)&dIpiv, N*sizeof(int64_t)) == hipSuccess ); + GRID_ASSERT( hipMalloc((void **)&dInfo, sizeof(int64_t)) == hipSuccess ); + auto st1 = rocsolver_cgetrf_64(handle, (int64_t)N, (int64_t)N, dA, (int64_t)N, dIpiv, dInfo); + hipDeviceSynchronize(); + int64_t info_h = -1; + hipMemcpy(&info_h, dInfo, sizeof(int64_t), hipMemcpyDeviceToHost); + std::cout << GridLogMessage << "DenseCoarseMatrix: cgetrf_64 status " << (int)st1 + << " info = " << (int)info_h << std::endl; + GRID_ASSERT(st1 == rocblas_status_success); + GRID_ASSERT(info_h == 0); + hipFree(dInfo); + GRID_ASSERT( hipMalloc((void **)&dB, (uint64_t)CHUNKROWS*N*sizeof(ComplexF)) == hipSuccess ); + // dA holds the LU of A^T; rows of A^{-1} are produced blockwise below via + // cgetrs_64 on identity-column blocks: A^T X = E => X columns = rows of + // A^{-1}, in exactly the linear layout the harvest expects. +#else + // Eigen fallback: small local CPU tests only. + std::cout << GridLogMessage << "DenseCoarseMatrix: Eigen fallback inversion N=" << N + << (N > 10000 ? " (WARNING: SLOW; use the HIP/rocSOLVER path)" : "") + << std::endl; + typedef Eigen::Matrix,Eigen::Dynamic,Eigen::Dynamic,Eigen::RowMajor> MatF; + Eigen::Map A(reinterpret_cast*>(&Afull[0]), N, N); + MatF Ainv = A.inverse(); + A = Ainv; +#endif + } + double t3 = usecond(); + std::cout << GridLogMessage << "DenseCoarseMatrix: factorisation took " + << (t3-t2)/1.0e6 << " s" << std::endl; + + //////////////////////////////////////////////////////////////////// + // Blocked solve + broadcast: rows of A^{-1} chunk by chunk; each + // rank keeps the rows of its own sites (ownership-aligned). + //////////////////////////////////////////////////////////////////// + { + std::unordered_map rowmap; + for(int ss=0; ss chunk((uint64_t)CHUNKROWS * N); + for(int64_t row0=0; row0Broadcast(0, &chunk[0], nelem*sizeof(ComplexF)); + for(int64_t r=row0; rsecond) * N; + uint64_t src = (uint64_t)(r-row0) * N; + for(int64_t j=0;j allreduce -> split-K GEMM against the resident slab -> reduce + // partials -> hY[nrows x nr] (column major). Platform-agnostic: + // deviceVector + GridBLAS (Eigen fallback on CPU builds). + // fp32 allreduce is EXACT: zero-fill assembly gives every element + // exactly one contributing rank. + //////////////////////////////////////////////////////////////////// + void SlabApplyPacked(int nr, double *tprof) + { + GRID_ASSERT(nr <= MRHS_MAX); + uint64_t nX = (uint64_t)N * nr; + uint64_t nY = (uint64_t)nrows * nr; + int64_t Kc = N / NK; + double t1 = usecond(); + double t2, t3; + if (devSum) { + acceleratorCopyToDevice(&hX[0],&dX[0],nX*sizeof(ComplexF)); + t2 = usecond(); + grid->GlobalSumVector((ComplexF *)&dX[0], (int)nX); + t3 = usecond(); + } else { + grid->GlobalSumVector(&hX[0], (int)nX); + t2 = usecond(); + acceleratorCopyToDevice(&hX[0],&dX[0],nX*sizeof(ComplexF)); + t3 = usecond(); + } + // Y = op(slab,T) . X : row-major slab (nrows x N) == col-major A^T + // (N x nrows, lda=N) => transpose gives the nrows x N operator. + // Split-K: NK chunk-GEMMs by pointer offset (AOT lists), then reduce. + ComplexF one (1.0,0.0); + ComplexF zero(0.0,0.0); + BLAS.gemmBatched(GridBLAS_OP_T, GridBLAS_OP_N, + (int)nrows, nr, (int)Kc, + one, aptrs, (int)N, + xptrs, (int)N, + zero, cptrs, (int)nrows); + BLAS.synchronise(); + { + ComplexF *pp = &dPartial[0]; + ComplexF *py = &dY[0]; + uint64_t stride = (uint64_t)nrows*MRHS_MAX; + int nk = NK; + accelerator_for(i, nY, 1, { + ComplexF acc(0.0,0.0); + for(int j=0;j certifies device slab). + //////////////////////////////////////////////////////////////////// + virtual void operator()(const Field &src, Field &psi) + { + uint64_t nX = (uint64_t)N; + thread_for(i, nX, { hX[i]=ComplexF(0.0,0.0); }); + for(int ss=0; ss