portelli/Grid
1
0
Fork 0
mirror of https://github.com/paboyle/Grid.git synced 2026-09-14 05:23:17 +01:00
Code Releases Activity

Preparing for multigrid parameter consolidation and clean up of code, rationalise the different variants.

Browse Source
This commit is contained in:
Peter Boyle
2026-09-06 14:18:00 -04:00
parent 482f3cbaa2
commit a07545adc4
21 changed files with 262 additions and 2435 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Grid/algorithms/multigrid/DenseCoarseMatrix.h
+62 -424
View File
@@ -29,7 +29,6 @@ Author: Peter Boyle <pboyle@bnl.gov>
#include <Grid/algorithms/blas/BatchedBlas.h>
#include <Grid/algorithms/blas/BatchedInverse.h>
#include <Grid/algorithms/multigrid/RecursiveSchurInverse.h>
#include <Grid/algorithms/multigrid/BlockCyclicSchurInverse.h>
#include <Grid/algorithms/multigrid/BlockCyclicRedistribute.h>
@@ -39,46 +38,33 @@ 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).
// row-distributed A^{-1} of a GeneralCoarsenedMatrix.
//
// 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.
// assembled from purely LOCAL _A[p] data: no operator applies, no comms.
// 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).
//
// - Inversion is END-TO-END fp64 through the 2D block-cyclic recursive Schur
// complement (BlockCyclicSchurInverse): fp64 rank-major import ->
// RowsToCyclic -> in-place recursion (pure point-to-point SUMMA rings and
// local leaves; bitwise reproducible) -> CyclicToRows -> ONE terminal
// rounding into the fp32 apply slab. Distributed at every N and P; banked
// 3.87x against its retired 1D predecessor and >10x against SLATE at
// N=138240 on 288 GCDs.
//
// - 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).
// SPLITK chunk-GEMMs by pointer offset into the resident slab (lda = N),
// partials reduced in one accelerator_for. The source vector is assembled by
// a cartesian ring ALLGATHER of device buffers (pure P2P; ~8x fewer bytes
// than a padded allreduce, and no collective size cliffs). Platform-agnostic:
// deviceVector + GridBLAS run the SAME code on HIP/CUDA/SYCL and CPU(Eigen).
//
// Setup: SLAB_FILE=<stem> loads per-rank <stem>.<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
// DENSE_SCHUR (0/absent: single-GCD gather-invert; 1: distributed
// recursive Schur; 2: AUDIT -- run BOTH on the same imported A, report
// the slab difference, keep the Schur result) DENSE_PANEL_BYTES
//
// The DENSE_SCHUR=1 path is the RecursiveSchurInverse distributed
// factorisation: it lifts the fp32 N ~ 90k boss-HBM ceiling (the CC-grid
// 256-rank SIMD cap remains -- separate issue). Internal only: slab layout,
// apply path, SLAB_FILE format and VERIFY are identical in every mode.
// VERIFY ||A Ainv x - x||/||x|| certifies the DEVICE slab + split-K path at the
// end of Import, since the single-RHS apply routes through the same core.
//
// Tensor-depth agnostic: site scalar objects treated as contiguous ComplexD
// (iScalar wrappers add no data), so any MG level's coarse operator imports.
@@ -111,15 +97,15 @@ public:
std::vector<int64_t> myGsite; // global lex site index of my site ss
std::vector<ComplexF> 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;
static const int SPLITK = 32; // requested split-K chunk count, snapped DOWN to a divisor of N
// Apply machinery: resident slab + persistent buffers + AOT split-K pointers.
GridBLAS BLAS;
deviceVector<ComplexF> dSlab;
deviceVector<ComplexF> dX; // N x MRHS_MAX
deviceVector<ComplexF> dY; // nrows x MRHS_MAX
deviceVector<ComplexF> dG; // N x MRHS_MAX lex-major staging for the allgather (devSum==4)
deviceVector<ComplexF> dG; // N x MRHS_MAX lex-major staging for the allgather
deviceVector<int> dLex2Rank;// lex index of a process coordinate -> its rank (allgather block order -> row-block order)
deviceVector<int64_t> dRm2G; // rank-major index (rank*nrows + ss*nbasis + b) -> global column (gsite*nbasis + b) of x / the slab
int myLex;
@@ -130,8 +116,6 @@ public:
std::vector<ComplexF> hX;
std::vector<ComplexF> hY;
int NK; // split-K chunk count (divides N)
int devSum;
double schurAuditRel; // DENSE_SCHUR=2: rel slab diff single-vs-schur (-1 = not run)
DenseCoarseMatrix(GridBase *g)
: grid(g)
@@ -142,7 +126,6 @@ public:
N = grid->gSites() * nbasis;
lsites = grid->lSites();
nrows = (int64_t)lsites * nbasis;
schurAuditRel = -1.0;
std::cout << GridLogMessage << "DenseCoarseMatrix: N = " << N
<< " (" << grid->gSites() << " sites x " << nbasis << ")"
@@ -180,47 +163,11 @@ public:
void Import(CoarseOp &Op)
{
double t0 = usecond();
////////////////////////////////////////////////////////////////////
// 0. Slab cache: SLAB_FILE=<stem> -> per-rank raw file <stem>.<rank>.
// 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(Op); // slab <- my rows of A (LOCAL, no comms)
ImportCertificate(Op); // dense apply == Op.M, before inversion
InvertDense(Op); // 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;
}
@@ -233,11 +180,9 @@ public:
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;
// Split-K chunk count, snapped DOWN to a divisor of N.
NK = 1;
for(int j=1;j<=req;j++) if ( (N % j) == 0 ) NK = j;
for(int j=1;j<=SPLITK;j++) if ( (N % j) == 0 ) NK = j;
int64_t Kc = N / NK;
dX.resize((uint64_t)N*MRHS_MAX);
@@ -255,12 +200,7 @@ public:
for(int j=0;j<NK;j++) h[j] = &dPartial[0] + (uint64_t)j*nrows*MRHS_MAX; // compact, ldc=nrows
acceleratorCopyToDevice(&h[0],&cptrs[0],NK*sizeof(ComplexF*));
devSum = getenv("DENSE_DEVICE_SUM") ? atoi(getenv("DENSE_DEVICE_SUM")) : 0;
const char *sumName[5] = {"host allreduce","DEVICE-buffer allreduce (GPU-aware MPI)",
"DEVICE cartesian ring allreduce (P2P)","DEVICE flat ring allreduce (P2P)",
"DEVICE cartesian ring ALLGATHER (P2P, ~8x fewer bytes than the padded allreduce)"};
GRID_ASSERT(devSum>=0 && devSum<=4);
if ( devSum==4 ) {
{
dG.resize((uint64_t)N*MRHS_MAX);
// allgather delivers blocks in lexicographic-coordinate order; the row
// blocks of x are in RANK order. Same table as BuildRankMajorMap.
@@ -282,8 +222,8 @@ public:
acceleratorCopyToDevice(&rm2g[0], &dRm2G[0], N*sizeof(int64_t));
}
std::cout << GridLogMessage << "DenseCoarseMatrix: slab resident on device ("
<< sbytes/1024./1024. << " MB/rank), split-K NK=" << NK << " (Kc=" << Kc << "); "
<< sumName[devSum] << std::endl;
<< sbytes/1024./1024. << " MB/rank), split-K NK=" << NK << " (Kc=" << Kc
<< "); DEVICE cartesian ring ALLGATHER (P2P)" << std::endl;
}
////////////////////////////////////////////////////////////////////
@@ -355,8 +295,6 @@ public:
{
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); });
@@ -367,28 +305,11 @@ public:
// extract the unpadded field before peeking with unpadded coordinates
// (exactly as MultiGeneralCoarsenedMatrix::CopyMatrix does).
CoarseMatrix Aun(grid); Op.ExtractMatrix(p,Aun);
if ( getenv("DENSE_IMPORT_DEBUG") ) {
// Peek-path vs field-norm audit: sum |peekLocalSite|^2 must match norm2
double pk = 0.0;
autoView(Adbg, Aun, CpuRead);
for(int ss=0; ss<lsites; ss++){
Msobj m;
peekLocalSite(m, Adbg, myLcoor[ss]);
ComplexD *md = (ComplexD *)&m;
for(int i=0; i<nbasis*nbasis; i++) pk += md[i].real()*md[i].real() + md[i].imag()*md[i].imag();
}
RealD gpk = pk;
grid->GlobalSumVector(&gpk, 1);
std::cout << GridLogMessage << "DenseCoarseMatrix: DEBUG p=" << p
<< " norm2(_A[p]) " << norm2(Aun)
<< " norm2(Extract) " << norm2(Aun)
<< " sum|peek|^2 " << gpk << std::endl;
}
autoView(Av, Aun, CpuRead);
thread_for(ss, lsites, {
Coordinate ncoor(nd);
for(int d=0; d<nd; d++){
int64_t g = grid->_lstart[d] + myLcoor[ss][d] + sign*shift[d];
int64_t g = grid->_lstart[d] + myLcoor[ss][d] + shift[d];
ncoor[d] = (int)((g % gdims[d] + gdims[d]) % gdims[d]);
}
int64_t nsite;
@@ -402,8 +323,8 @@ public:
// column (nbr,a). BUG LEDGER 2026-08-14: the original mapping
// wrote (s,a),(nbr,b) -- caught by the IMPORT CERTIFICATE on its
// FIRST fresh-import exercise (Test_schur_dense_coarse); every
// production slab predates this path (probe-import SLAB_FILEs),
// so no production output is suspect.
// production slab predated this path (probe-import era), so no
// production output is suspect.
for(int b=0; b<nbasis; b++){
ComplexF *row = &slab[(uint64_t)(ss*nbasis+b)*N + nsite*nbasis];
for(int a=0; a<nbasis; a++)
@@ -499,250 +420,15 @@ public:
<< rel << std::endl;
if ( rel >= 1.0e-3 ) {
std::cout << GridLogMessage << "DenseCoarseMatrix: IMPORT CERTIFICATE FAILED. If O(1), the "
<< "stencil shift-sign convention is opposite: rerun with DENSE_IMPORT_SIGN=-1"
<< "stencil shift-sign convention of the coarse operator has changed: "
<< "the import in ImportDense/ImportDenseFP64 must change with it"
<< std::endl;
}
GRID_ASSERT(rel < 1.0e-3);
}
////////////////////////////////////////////////////////////////////
// 3. Invert dispatcher. slab holds my rows of A on entry, my rows
// of A^{-1} on exit, in every mode.
// DENSE_SCHUR absent/0 : single-GCD gather-invert (the oracle)
// DENSE_SCHUR=1 : distributed recursive Schur
// DENSE_SCHUR=2 : AUDIT -- both on the same A; report the
// slab difference; keep the Schur result
// (so VERIFY certifies the new path).
////////////////////////////////////////////////////////////////////
template<class CoarseOp>
void InvertDense(CoarseOp &Op)
{
char *sc = getenv("DENSE_SCHUR");
int mode = sc ? atoi(sc) : 0;
if ( mode == 0 )
{
InvertDenseSingle();
return;
}
if ( mode == 1 )
{
InvertDenseSchur(Op);
return;
}
GRID_ASSERT( mode == 2 );
std::vector<ComplexF> Aimp(slab); // imported A
InvertDenseSingle();
std::vector<ComplexF> ref(slab); // Ainv, single path
slab = Aimp;
InvertDenseSchur(Op); // slab = Ainv, Schur path
// NaN-PROOF comparison: max() masks NaN, so count non-finite
// entries in each result explicitly.
double mx = 0.0;
double mr = 0.0;
int64_t badschur = 0;
int64_t badsingle = 0;
for(uint64_t i=0; i<(uint64_t)nrows*N; i++)
{
double as = abs(ComplexD(slab[i]));
double ar = abs(ComplexD(ref[i]));
if ( !std::isfinite(as) ) badschur++;
if ( !std::isfinite(ar) ) badsingle++;
if ( std::isfinite(as) && std::isfinite(ar) )
{
mx = std::max(mx, (double)abs(ComplexD(slab[i]) - ComplexD(ref[i])));
mr = std::max(mr, ar);
}
}
RealD gmx = mx;
RealD gmr = mr;
RealD gbs = (RealD)badschur;
RealD gbr = (RealD)badsingle;
grid->GlobalMax(gmx);
grid->GlobalMax(gmr);
grid->GlobalSumVector(&gbs, 1);
grid->GlobalSumVector(&gbr, 1);
schurAuditRel = gmx/gmr;
std::cout << GridLogMessage << "DenseCoarseMatrix: DENSE_SCHUR=2 AUDIT "
<< "max|Ainv_schur - Ainv_single| = " << gmx
<< " relative " << schurAuditRel
<< " non-finite: schur " << (int64_t)gbs << " single " << (int64_t)gbr
<< " (two fp32 roundings of the same inverse; expect ~ growth * eps32)"
<< std::endl;
GRID_ASSERT( gbs == 0 );
GRID_ASSERT( gbr == 0 );
}
////////////////////////////////////////////////////////////////////
// 3a. Single-GCD invert: chunked zero-fill+GlobalSum gather of A
// streamed to the boss GCD, cgetrf_64 (ILP64), rows of A^{-1} by
// blocked identity cgetrs_64 + broadcast; each rank keeps its own
// rows (in `slab`, overwriting A). Proven path; the SCHUR oracle.
////////////////////////////////////////////////////////////////////
void InvertDenseSingle(void)
{
double t1 = usecond();
int boss = grid->IsBoss();
std::vector<ComplexF> 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<int64_t,int> rowmap; // global row -> my slab row
for(int ss=0; ss<lsites; ss++)
for(int a=0; a<nbasis; a++)
rowmap[ myGsite[ss]*nbasis + a ] = ss*nbasis + a;
std::vector<ComplexF> chunk((uint64_t)CHUNKROWS * N);
for(int64_t row0=0; row0<N; row0+=CHUNKROWS){
int64_t nrow = std::min(CHUNKROWS, N-row0);
uint64_t nelem = (uint64_t)nrow * N;
for(uint64_t i=0;i<nelem;i++) chunk[i]=ComplexF(0.0,0.0);
for(int64_t r=row0; r<row0+nrow; r++){
auto it = rowmap.find(r);
if (it != rowmap.end()) {
uint64_t src = (uint64_t)(it->second) * N;
uint64_t dst = (uint64_t)(r-row0) * N;
for(int64_t j=0;j<N;j++) chunk[dst+j] = slab[src+j];
}
}
grid->GlobalSumVector(&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<nelem;i++) Afull[dst+i] = chunk[i];
#endif
}
}
}
double t2 = usecond();
std::cout << GridLogMessage << "DenseCoarseMatrix: gather to boss took "
<< (t2-t1)/1.0e6 << " s" << std::endl;
////////////////////////////////////////////////////////////////////
// Factor in place on the boss (fp32, ILP64). Row-major buffer handed
// to column-major LAPACK => 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);
GRID_ASSERT( hipDeviceSynchronize() == hipSuccess );
int64_t info_h = -1;
GRID_ASSERT( hipMemcpy(&info_h, dInfo, sizeof(int64_t),
hipMemcpyDeviceToHost) == hipSuccess );
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);
GRID_ASSERT( hipFree(dInfo) == hipSuccess );
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<std::complex<float>,Eigen::Dynamic,Eigen::Dynamic,Eigen::RowMajor> MatF;
Eigen::Map<MatF> A(reinterpret_cast<std::complex<float>*>(&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<int64_t,int> rowmap;
for(int ss=0; ss<lsites; ss++)
for(int a=0; a<nbasis; a++)
rowmap[ myGsite[ss]*nbasis + a ] = ss*nbasis + a;
std::vector<ComplexF> chunk((uint64_t)CHUNKROWS * N);
for(int64_t row0=0; row0<N; row0+=CHUNKROWS){
int64_t nrow = std::min(CHUNKROWS, N-row0);
uint64_t nelem = (uint64_t)nrow * N;
if (boss) {
#ifdef GRID_HIP
// Identity block E: column j = e_{row0+j}; solve A^T X = E so X's
// columns are rows [row0,row0+nrow) of A^{-1}.
for(uint64_t i=0;i<nelem;i++) chunk[i] = ComplexF(0.0,0.0);
for(int64_t j=0;j<nrow;j++) chunk[(uint64_t)j*N + (uint64_t)(row0+j)] = ComplexF(1.0,0.0);
GRID_ASSERT( hipMemcpy(dB, &chunk[0], nelem*sizeof(ComplexF), hipMemcpyHostToDevice) == hipSuccess );
auto strs = rocsolver_cgetrs_64(GridBLASInverse::Handle(), rocblas_operation_none,
(int64_t)N, (int64_t)nrow,
dA, (int64_t)N, dIpiv, dB, (int64_t)N);
GRID_ASSERT(strs == rocblas_status_success);
GRID_ASSERT( hipDeviceSynchronize() == hipSuccess );
GRID_ASSERT( hipMemcpy(&chunk[0], dB, nelem*sizeof(ComplexF), hipMemcpyDeviceToHost) == hipSuccess );
#else
uint64_t src = (uint64_t)row0 * N;
for(uint64_t i=0;i<nelem;i++) chunk[i] = Afull[src+i];
#endif
}
grid->Broadcast(0, &chunk[0], nelem*sizeof(ComplexF));
for(int64_t r=row0; r<row0+nrow; r++){
auto it = rowmap.find(r);
if (it != rowmap.end()) {
uint64_t dst = (uint64_t)(it->second) * N;
uint64_t src = (uint64_t)(r-row0) * N;
for(int64_t j=0;j<N;j++) slab[dst+j] = chunk[src+j];
}
}
}
}
#ifdef GRID_HIP
if (boss) {
if (dA) GRID_ASSERT( hipFree(dA) == hipSuccess );
if (dB) GRID_ASSERT( hipFree(dB) == hipSuccess );
if (dIpiv) GRID_ASSERT( hipFree(dIpiv) == hipSuccess );
}
#endif
double t4 = usecond();
std::cout << GridLogMessage << "DenseCoarseMatrix: blocked getrs solve+scatter took "
<< (t4-t3)/1.0e6 << " s" << std::endl;
}
////////////////////////////////////////////////////////////////////
// 3b. Global column -> rank-major column map, computed LOCALLY.
// 3a. Global column -> rank-major column map, computed LOCALLY.
// Rank-major ordering: rank q's rows/columns are the contiguous
// block [q*nrows, (q+1)*nrows), ordered by q's local site index
// (uniform local volumes make ownership arithmetic exact).
@@ -789,7 +475,7 @@ public:
}
////////////////////////////////////////////////////////////////////
// 3c. Direct stencil -> fp64 rank-major import of MY ROWS of A (the
// 3b. Direct stencil -> fp64 rank-major import of MY ROWS of A (the
// end-to-end fp64 path: the stencil source IS ComplexD; nothing is
// rounded through fp32 on the way into the inversion). Same
// loop/sign/accumulate/transposed-contraction discipline as
@@ -805,8 +491,6 @@ public:
void ImportDenseFP64(CoarseOp &Op, BlockRows &S, std::vector<int64_t> &g2rm)
{
Coordinate gdims = grid->GlobalDimensions();
int sign = getenv("DENSE_IMPORT_SIGN") ? atoi(getenv("DENSE_IMPORT_SIGN")) : 1;
GRID_ASSERT( sign==1 || sign==-1 );
std::vector<ComplexD> h((uint64_t)nrows*N, ComplexD(0.0,0.0));
for(int p=0; p<Op.Geometry().npoint; p++)
@@ -818,7 +502,7 @@ public:
Coordinate ncoor(nd);
for(int d=0; d<nd; d++)
{
int64_t g = grid->_lstart[d] + myLcoor[ss][d] + sign*shift[d];
int64_t g = grid->_lstart[d] + myLcoor[ss][d] + shift[d];
ncoor[d] = (int)((g % gdims[d] + gdims[d]) % gdims[d]);
}
int64_t nsite;
@@ -868,14 +552,14 @@ public:
}
////////////////////////////////////////////////////////////////////
// 3d. Distributed recursive Schur invert, END-TO-END fp64 (decision
// 3c. The inverse: distributed recursive Schur, END-TO-END fp64 (decision
// 2026-08-14): stencil (ComplexD) -> fp64 rank-major import ->
// fp64 recursion -> ONE terminal rounding into the fp32 apply
// slab. Everything downstream (device residency, split-K apply,
// VERIFY, SLAB_FILE) is untouched.
// VERIFY) is untouched.
////////////////////////////////////////////////////////////////////
template<class CoarseOp>
void InvertDenseSchur(CoarseOp &Op)
void InvertDense(CoarseOp &Op)
{
double t1 = usecond();
int P = grid->ProcessorCount();
@@ -905,26 +589,17 @@ public:
ImportDenseFP64(Op, S, g2rm);
////////////////////////////////////////////////////////////////
// DENSE_SCHUR2D=1 : invert via the 2D block-cyclic recursion
// (BlockCyclicSchurInverse) instead of the 1D rank-range one.
// The SAME imported rank-major rows S go in and come back, so the
// import certificate above and the slab rounding / VERIFY below are
// identical for both paths: a clean A/B on one imported operator.
//
// Everything in the 2D path -- redistribution, SUMMA rings, leaf --
// is point-to-point SendToRecvFrom; no collectives at all.
// DENSE_NB overrides the block size (default: rows-per-rank, which
// makes the redistribution edges maximally regular).
// The 2D block-cyclic recursion (BlockCyclicSchurInverse).
// Everything in it -- redistribution, SUMMA rings, leaves -- is
// point-to-point SendToRecvFrom; no collectives at all. Block size
// nb = rows-per-rank makes the redistribution edges maximally
// regular.
////////////////////////////////////////////////////////////////
int use2d = getenv("DENSE_SCHUR2D") ? atoi(getenv("DENSE_SCHUR2D")) : 0;
int64_t panelBytes = getenv("DENSE_PANEL_BYTES") ? atol(getenv("DENSE_PANEL_BYTES"))
: (int64_t)1024*1024*1024; // 1D path only
double t2, t3;
if ( use2d )
{
int Pr,Pc;
BlockCyclicLayout::ChooseProcessGrid(P, Pr, Pc);
int64_t nb = getenv("DENSE_NB") ? atol(getenv("DENSE_NB")) : nrows;
int64_t nb = nrows;
GRID_ASSERT( nb >= 1 );
std::cout << GridLogMessage << "DenseCoarseMatrix: 2D SCHUR invert, process grid "
<< Pr << " x " << Pc << " nb " << nb
@@ -938,14 +613,6 @@ public:
t3 = usecond();
RSI2.ReportTelemetry(grid);
}
else
{
RecursiveSchurInverse RSI(grid, N, rowStart, panelBytes);
t2 = usecond();
RSI.Invert(S);
t3 = usecond();
RSI.ReportTelemetry();
}
// The single terminal rounding: fp64 inverse -> fp32 apply slab
// (row-major, global columns)
@@ -962,28 +629,26 @@ public:
}
double t4 = usecond();
std::cout << GridLogMessage << "DenseCoarseMatrix: SCHUR fp64 distributed invert took "
<< (t4-t1)/1.0e6 << " s (recursion " << (t3-t2)/1.0e6 << " s), panelBytes "
<< panelBytes << std::endl;
<< (t4-t1)/1.0e6 << " s (recursion " << (t3-t2)/1.0e6 << " s)" << std::endl;
}
////////////////////////////////////////////////////////////////////
// CORE apply on packed data: hX[N x nr] (zero-filled, my sites only)
// -> allreduce -> split-K GEMM against the resident slab -> reduce
// partials -> hY[nrows x nr] (column major). Platform-agnostic:
// -> ring allgather -> 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.
// tprof (optional): per-phase microseconds {allgather, H2D, gemm+reduce,
// D2H}, printed by the caller on GridLogPerformance.
////////////////////////////////////////////////////////////////////
void SlabApplyPacked(int nr, double *tprof)
{
GRID_TRACE("DenseSlabApply");
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==4) {
{
// ALLGATHER: x is not a reduction -- every rank owns the rows of x at
// global columns myGsite[ss]*nbasis+b (scattered by site coordinate, NOT
// a contiguous block) and needs all of it. Only MY rows go host->device
@@ -1011,31 +676,6 @@ public:
});
}
t3 = usecond();
} else if (devSum) {
{ GRID_TRACE("DenseH2D");
acceleratorCopyToDevice(&hX[0],&dX[0],nX*sizeof(ComplexF));
}
t2 = usecond();
{ GRID_TRACE("DenseAllreduce");
// DENSE_DEVICE_SUM=1 : device-buffer MPI_Allreduce (Cray MPICH aborts
// above ~8 MB: 12 RHS at N=138240 is 13.3 MB)
// DENSE_DEVICE_SUM=2 : CartesianRingAllReduce, P2P only, no size cliff
// DENSE_DEVICE_SUM=3 : flat RingAllReduce, P2P only
// DENSE_DEVICE_SUM=4 : CartesianRingAllGather (branch above)
if (devSum==2) CartesianRingAllReduce(grid,(ComplexF *)&dX[0],nX);
else if (devSum==3) RingAllReduce(grid,(ComplexF *)&dX[0],nX);
else grid->GlobalSumVector((ComplexF *)&dX[0], (int)nX);
}
t3 = usecond();
} else {
{ GRID_TRACE("DenseAllreduce");
grid->GlobalSumVector(&hX[0], (int)nX);
}
t2 = usecond();
{ GRID_TRACE("DenseH2D");
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.
@@ -1065,8 +705,8 @@ public:
}
double t5 = usecond();
if (tprof) {
tprof[0] = devSum ? (t3-t2) : (t2-t1); // allreduce
tprof[1] = devSum ? (t2-t1) : (t3-t2); // H2D
tprof[0] = t3-t2; // allgather
tprof[1] = t2-t1; // H2D
tprof[2] = t4-t3; // gemm+reduce
tprof[3] = t5-t4; // D2H
}
@@ -1208,16 +848,14 @@ public:
double t6 = usecond();
std::cout << GridLogMessage << "DenseCoarseMatrix: apply6D " << nr << " rhs took "
<< (t6-t0)/1000.0 << " ms" << std::endl;
if ( getenv("DENSE_APPLY_PROFILE") ) {
std::cout << GridLogMessage << "DenseCoarseMatrix: apply6D profile:"
<< " pack " << (t1-t0)/1000.0
<< (devSum==4 ? " allgather " : " allreduce ") << tprof[0]/1000.0
<< " H2D " << tprof[1]/1000.0
<< " gemm+reduce "<< tprof[2]/1000.0
<< " D2H " << tprof[3]/1000.0
<< " unpack " << (t6-t5)/1000.0
<< " ms" << std::endl;
}
std::cout << GridLogPerformance << "DenseCoarseMatrix: apply6D profile:"
<< " pack " << (t1-t0)/1000.0
<< " allgather " << tprof[0]/1000.0
<< " H2D " << tprof[1]/1000.0
<< " gemm+reduce "<< tprof[2]/1000.0
<< " D2H " << tprof[3]/1000.0
<< " unpack " << (t6-t5)/1000.0
<< " ms" << std::endl;
}
};