mirror of
https://github.com/paboyle/Grid.git
synced 2026-08-24 11:29:35 +01:00
Claude 2D block cyclic inverse
This commit is contained in:
@@ -0,0 +1,215 @@
|
||||
/*************************************************************************************
|
||||
|
||||
Grid physics library, www.github.com/paboyle/Grid
|
||||
|
||||
Source file: ./Grid/algorithms/multigrid/BlockCyclic.h
|
||||
|
||||
Copyright (C) 2026
|
||||
|
||||
Author: Peter Boyle <pboyle@bnl.gov>
|
||||
|
||||
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.
|
||||
|
||||
See the full license in the file "LICENSE" in the top level distribution
|
||||
directory
|
||||
*************************************************************************************/
|
||||
/* END LEGAL */
|
||||
#pragma once
|
||||
|
||||
NAMESPACE_BEGIN(Grid);
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
// BlockCyclicLayout: the index arithmetic of a 2D block-cyclic distribution
|
||||
// of an N x N matrix over a Pr x Pc logical process grid with block size nb.
|
||||
//
|
||||
// This is stage 1 of the 2D distributed dense inverse
|
||||
// (documentation/DistributedDenseInverse2D.tex). It is deliberately
|
||||
// COMMUNICATOR-FREE: every mapping is a static pure function of
|
||||
// (N, nb, Pr, Pc), so the whole layout is exhaustively unit-testable on one
|
||||
// rank with no MPI in the loop (Test_blockcyclic). A thin instance layer
|
||||
// binds a world rank to a grid coordinate and caches local extents.
|
||||
//
|
||||
// Conventions (fixed here, relied on by every later stage):
|
||||
//
|
||||
// * Global block b of a dimension with Pg processes is owned by process
|
||||
// coordinate b % Pg (ScaLAPACK csrc=0), and is that process's local
|
||||
// block b / Pg.
|
||||
// * Rank <-> grid coordinate is ROW MAJOR over the process grid:
|
||||
// rank = p*Pc + q , p = rank/Pc , q = rank%Pc .
|
||||
// The eventual ring transport must construct its neighbour tables with
|
||||
// the same convention.
|
||||
// * Local storage is COLUMN MAJOR with ld = mloc, matching BlockRows:
|
||||
// local element (i,j) lives at data[i + j*mloc].
|
||||
// * The trailing partial block (N % nb != 0) belongs to the owner of the
|
||||
// last full-size block position; only that one block is short.
|
||||
//
|
||||
// Element (gi,gj) therefore lives on grid coordinate
|
||||
// ( (gi/nb) % Pr , (gj/nb) % Pc )
|
||||
// at local coordinate
|
||||
// ( ((gi/nb)/Pr)*nb + gi%nb , ((gj/nb)/Pc)*nb + gj%nb ).
|
||||
//
|
||||
// Everything here is host-side integer arithmetic; nothing allocates.
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
class BlockCyclicLayout
|
||||
{
|
||||
public:
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Closest-to-square factorisation Pr*Pc == P with Pr <= Pc.
|
||||
// For fixed P the per-rank SUMMA volume N^2 (1/Pr + 1/Pc) is minimised
|
||||
// at the most square grid. P=288 -> 16 x 18.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
static void ChooseProcessGrid(int P, int &Pr, int &Pc)
|
||||
{
|
||||
GRID_ASSERT(P >= 1);
|
||||
Pr = 1;
|
||||
for(int r=1; (int64_t)r*r <= (int64_t)P; r++)
|
||||
if ( P % r == 0 ) Pr = r;
|
||||
Pc = P / Pr;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Number of rows (or columns) of a dimension of global extent N, block nb,
|
||||
// owned by process coordinate p of Pg. ScaLAPACK "numroc", csrc=0.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
static int64_t NumLocal(int64_t N, int64_t nb, int p, int Pg)
|
||||
{
|
||||
GRID_ASSERT(N >= 0);
|
||||
GRID_ASSERT(nb >= 1);
|
||||
GRID_ASSERT(p >= 0);
|
||||
GRID_ASSERT(p < Pg);
|
||||
int64_t nblocks = N / nb; // full blocks
|
||||
int64_t extra = N % nb; // trailing partial block
|
||||
int64_t full = nblocks / Pg; // full blocks everyone owns
|
||||
int64_t rem = nblocks % Pg; // coords [0,rem) own one more
|
||||
int64_t n = full*nb;
|
||||
if ( p < (int)rem ) n += nb; // an extra full block
|
||||
if ( p == (int)rem ) n += extra; // the partial block, if any
|
||||
return n;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Global index -> (owner coordinate, local index) in one dimension.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
static void GlobalToLocal(int64_t g, int64_t nb, int Pg,
|
||||
int &owner, int64_t &loc)
|
||||
{
|
||||
GRID_ASSERT(g >= 0);
|
||||
int64_t b = g / nb; // global block
|
||||
owner = (int)(b % Pg);
|
||||
loc = (b / Pg)*nb + (g % nb);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// (process coordinate, local index) -> global index in one dimension.
|
||||
// Inverse of GlobalToLocal on the owned set.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
static int64_t LocalToGlobal(int64_t l, int64_t nb, int p, int Pg)
|
||||
{
|
||||
GRID_ASSERT(l >= 0);
|
||||
int64_t lb = l / nb; // local block
|
||||
int64_t b = lb*Pg + p; // global block
|
||||
return b*nb + (l % nb);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Instance layer: bind a rank of a Pr x Pc grid.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
int64_t N; // global matrix dimension (square)
|
||||
int64_t nb; // block size
|
||||
int Pr, Pc; // process grid
|
||||
int me; // world rank within the grid, row major
|
||||
int prow, pcol;// my grid coordinate
|
||||
int64_t mloc,nloc; // my local extents; storage column major, ld = mloc
|
||||
|
||||
BlockCyclicLayout(int64_t N_, int64_t nb_, int Pr_, int Pc_, int me_)
|
||||
{
|
||||
N = N_;
|
||||
nb = nb_;
|
||||
Pr = Pr_;
|
||||
Pc = Pc_;
|
||||
me = me_;
|
||||
GRID_ASSERT( N >= 0 );
|
||||
GRID_ASSERT( nb >= 1 );
|
||||
GRID_ASSERT( Pr >= 1 );
|
||||
GRID_ASSERT( Pc >= 1 );
|
||||
GRID_ASSERT( me >= 0 );
|
||||
GRID_ASSERT( me < Pr*Pc );
|
||||
prow = me / Pc; // ROW MAJOR rank convention
|
||||
pcol = me % Pc;
|
||||
mloc = NumLocal(N, nb, prow, Pr);
|
||||
nloc = NumLocal(N, nb, pcol, Pc);
|
||||
}
|
||||
|
||||
// Owning rank of global element (gi,gj), row-major rank convention.
|
||||
int OwnerRank(int64_t gi, int64_t gj) const
|
||||
{
|
||||
int pr,pc; int64_t li,lj;
|
||||
GlobalToLocal(gi, nb, Pr, pr, li);
|
||||
GlobalToLocal(gj, nb, Pc, pc, lj);
|
||||
return pr*Pc + pc;
|
||||
}
|
||||
|
||||
// My local storage offset of global element (gi,gj).
|
||||
// The caller must know I own it; asserted, not assumed.
|
||||
int64_t LocalOffset(int64_t gi, int64_t gj) const
|
||||
{
|
||||
int pr,pc; int64_t li,lj;
|
||||
GlobalToLocal(gi, nb, Pr, pr, li);
|
||||
GlobalToLocal(gj, nb, Pc, pc, lj);
|
||||
GRID_ASSERT( pr == prow );
|
||||
GRID_ASSERT( pc == pcol );
|
||||
return li + lj*mloc; // column major, ld = mloc
|
||||
}
|
||||
|
||||
// Do I own global element (gi,gj)?
|
||||
int Owns(int64_t gi, int64_t gj) const
|
||||
{
|
||||
return OwnerRank(gi,gj) == me;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Block-aligned global range [g0,g1) -> my contiguous local range [l0,l1).
|
||||
//
|
||||
// For fixed owner p the local index is monotone in the global index, so a
|
||||
// coordinate's owned elements of ANY global range are contiguous in local
|
||||
// storage; and for a BLOCK-ALIGNED range the bounds are exactly
|
||||
// NumLocal(g0) and NumLocal(g1), because NumLocal(g,...) counts the owned
|
||||
// elements below g. This is what lets a windowed product view the local
|
||||
// sub-matrix of a global window as &data[l0 + c0*mloc] with the SAME ld --
|
||||
// no gather, no copy. Verified exhaustively in Test_blockcyclic T7.
|
||||
//
|
||||
// g0 must be a block multiple; g1 a block multiple or N itself.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
static void RangeToLocal(int64_t g0, int64_t g1,
|
||||
int64_t N, int64_t nb, int p, int Pg,
|
||||
int64_t &l0, int64_t &l1)
|
||||
{
|
||||
GRID_ASSERT( 0 <= g0 );
|
||||
GRID_ASSERT( g0 <= g1 );
|
||||
GRID_ASSERT( g1 <= N );
|
||||
GRID_ASSERT( g0 % nb == 0 );
|
||||
GRID_ASSERT( (g1 % nb == 0) || (g1 == N) );
|
||||
l0 = NumLocal(g0, nb, p, Pg);
|
||||
l1 = NumLocal(g1, nb, p, Pg);
|
||||
}
|
||||
|
||||
// Instance forms, rows and columns of my own coordinate.
|
||||
void RowRange(int64_t g0, int64_t g1, int64_t &l0, int64_t &l1) const
|
||||
{ RangeToLocal(g0,g1,N,nb,prow,Pr,l0,l1); }
|
||||
void ColRange(int64_t g0, int64_t g1, int64_t &l0, int64_t &l1) const
|
||||
{ RangeToLocal(g0,g1,N,nb,pcol,Pc,l0,l1); }
|
||||
|
||||
// Size of global block b (the trailing block may be short).
|
||||
int64_t BlockSize(int64_t b) const
|
||||
{
|
||||
int64_t lo = b*nb;
|
||||
GRID_ASSERT( lo < N );
|
||||
return std::min(N, lo+nb) - lo;
|
||||
}
|
||||
};
|
||||
|
||||
NAMESPACE_END(Grid);
|
||||
@@ -0,0 +1,247 @@
|
||||
/*************************************************************************************
|
||||
|
||||
Grid physics library, www.github.com/paboyle/Grid
|
||||
|
||||
Source file: ./Grid/algorithms/multigrid/BlockCyclicRedistribute.h
|
||||
|
||||
Copyright (C) 2026
|
||||
|
||||
Author: Peter Boyle <pboyle@bnl.gov>
|
||||
|
||||
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.
|
||||
|
||||
See the full license in the file "LICENSE" in the top level distribution
|
||||
directory
|
||||
*************************************************************************************/
|
||||
/* END LEGAL */
|
||||
#pragma once
|
||||
|
||||
#include <Grid/algorithms/multigrid/BlockCyclicSumma.h>
|
||||
|
||||
NAMESPACE_BEGIN(Grid);
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
// Stage 4 of the 2D distributed dense inverse: redistribution between the
|
||||
// 1D rank-major row layout (BlockRows: rank r owns contiguous global rows
|
||||
// [rowStart[r], rowStart[r+1]) of an N x N matrix, stored rows x N column
|
||||
// major with ld = rows) and the 2D block-cyclic layout.
|
||||
//
|
||||
// This is what lets the EXISTING stencil->dense import, its certificate,
|
||||
// the fp32 slab conversion and the apply path all remain byte-for-byte
|
||||
// untouched: the 2D inverse slots between them as
|
||||
//
|
||||
// RowsToCyclic -> BlockCyclicSchurInverse::Invert -> CyclicToRows
|
||||
//
|
||||
// Volume is one matrix pass each way -- N^2/P elements per rank (~1 GB at
|
||||
// production), trivial against the inversion itself.
|
||||
//
|
||||
// Transport: PURE POINT-TO-POINT, like everything else in this stack.
|
||||
// Ranks exchange in a round-robin TOURNAMENT (the circle method, on an odd
|
||||
// modulus M so it covers every pair exactly once for any P, with byes):
|
||||
// at round r, ranks x and y are partners iff x+y == r (mod M). Each
|
||||
// meeting handles both directed edges of the pair in ONE SendToRecvFrom,
|
||||
// padded to the larger of the two edge sizes -- SendToRecvFrom carries a
|
||||
// single byte count for both directions, and both endpoints compute the
|
||||
// same max from the shared descriptors, so there is no asymmetric-size
|
||||
// case and no zero-count shape. Pairs with nothing to exchange skip the
|
||||
// round, decided identically at both ends.
|
||||
//
|
||||
// Element enumeration within an edge is canonical -- ascending global
|
||||
// column outer, ascending global row inner -- and each endpoint builds its
|
||||
// OWN local offset tables from the shared descriptors, so no index data is
|
||||
// ever transmitted. The round trip is BITWISE exact (pure data movement,
|
||||
// no arithmetic): Test_schur2d_redist proves it.
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
class BlockCyclicRedistribute
|
||||
{
|
||||
public:
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// The directed edge (1D rank r1, 2D rank r2): global rows of r1's range
|
||||
// whose row-block coordinate is r2's prow; ALL global columns whose
|
||||
// column-block coordinate is r2's pcol. Every rank can enumerate any
|
||||
// edge from (rowStart, layout) alone.
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
static void EdgeRows(const std::vector<int64_t> &rowStart, int r1,
|
||||
const BlockCyclicLayout &L, int r2,
|
||||
std::vector<int64_t> &rows)
|
||||
{
|
||||
rows.clear();
|
||||
int p = r2 / L.Pc; // row-major rank convention
|
||||
for(int64_t i=rowStart[r1]; i<rowStart[r1+1]; i++)
|
||||
if ( (int)((i/L.nb) % L.Pr) == p ) rows.push_back(i);
|
||||
}
|
||||
static void EdgeCols(const BlockCyclicLayout &L, int r2,
|
||||
std::vector<int64_t> &cols)
|
||||
{
|
||||
cols.clear();
|
||||
int q = r2 % L.Pc;
|
||||
for(int64_t b=0; b*L.nb<L.N; b++){
|
||||
if ( (int)(b % L.Pc) != q ) continue;
|
||||
int64_t g0=b*L.nb, g1=std::min(L.N,(b+1)*L.nb);
|
||||
for(int64_t j=g0;j<g1;j++) cols.push_back(j);
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// Gather/scatter one edge between a matrix (device, column major, ld)
|
||||
// and a dense edge buffer, through device offset tables.
|
||||
// buffer(a,b) = elem(rows[a], cols[b]), a fastest.
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
static void MoveEdge(int toBuffer,
|
||||
ComplexD *mat, int64_t ld,
|
||||
const std::vector<int64_t> &roff, // per-row offset in mat
|
||||
const std::vector<int64_t> &coff, // per-col offset in mat
|
||||
ComplexD *buf)
|
||||
{
|
||||
int64_t nr = roff.size();
|
||||
int64_t nc = coff.size();
|
||||
if ( !(nr && nc) ) return;
|
||||
deviceVector<int64_t> dro(nr), dco(nc);
|
||||
acceleratorCopyToDevice((void *)&roff[0], (void *)&dro[0], nr*sizeof(int64_t));
|
||||
acceleratorCopyToDevice((void *)&coff[0], (void *)&dco[0], nc*sizeof(int64_t));
|
||||
int64_t *ro = &dro[0];
|
||||
int64_t *co = &dco[0];
|
||||
if ( toBuffer ) {
|
||||
accelerator_for(idx, (uint64_t)(nr*nc), 1, {
|
||||
int64_t b = idx / nr;
|
||||
int64_t a = idx - b*nr;
|
||||
buf[a + b*nr] = mat[ ro[a] + co[b]*ld ];
|
||||
});
|
||||
} else {
|
||||
accelerator_for(idx, (uint64_t)(nr*nc), 1, {
|
||||
int64_t b = idx / nr;
|
||||
int64_t a = idx - b*nr;
|
||||
mat[ ro[a] + co[b]*ld ] = buf[a + b*nr];
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// My offset tables for an edge, on whichever side I am.
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
static void Offsets1D(const std::vector<int64_t> &rows,
|
||||
const std::vector<int64_t> &cols,
|
||||
int64_t row0,
|
||||
std::vector<int64_t> &roff, std::vector<int64_t> &coff)
|
||||
{
|
||||
roff.resize(rows.size()); coff.resize(cols.size());
|
||||
for(uint64_t a=0;a<rows.size();a++) roff[a] = rows[a]-row0; // local row
|
||||
for(uint64_t b=0;b<cols.size();b++) coff[b] = cols[b]; // global col
|
||||
}
|
||||
static void Offsets2D(const BlockCyclicLayout &L,
|
||||
const std::vector<int64_t> &rows,
|
||||
const std::vector<int64_t> &cols,
|
||||
std::vector<int64_t> &roff, std::vector<int64_t> &coff)
|
||||
{
|
||||
roff.resize(rows.size()); coff.resize(cols.size());
|
||||
for(uint64_t a=0;a<rows.size();a++){
|
||||
int p; int64_t l;
|
||||
BlockCyclicLayout::GlobalToLocal(rows[a], L.nb, L.Pr, p, l);
|
||||
GRID_ASSERT( p == L.prow );
|
||||
roff[a] = l;
|
||||
}
|
||||
for(uint64_t b=0;b<cols.size();b++){
|
||||
int q; int64_t l;
|
||||
BlockCyclicLayout::GlobalToLocal(cols[b], L.nb, L.Pc, q, l);
|
||||
GRID_ASSERT( q == L.pcol );
|
||||
coff[b] = l;
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// The worker. dir=+1 : 1D rows -> block cyclic ; dir=-1 : back.
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
static void Redistribute(int dir, GridBase *grid,
|
||||
const std::vector<int64_t> &rowStart,
|
||||
ComplexD *rows1d, int64_t myrows,
|
||||
BlockCyclicMatrix &A)
|
||||
{
|
||||
BlockCyclicLayout &L = A.layout;
|
||||
int P = grid->ProcessorCount();
|
||||
int me = grid->ThisRank();
|
||||
GRID_ASSERT( (int)rowStart.size() == P+1 );
|
||||
GRID_ASSERT( rowStart[P] == L.N );
|
||||
GRID_ASSERT( rowStart[me+1]-rowStart[me] == myrows );
|
||||
int64_t row0 = rowStart[me];
|
||||
int64_t ld1 = myrows ? myrows : 1;
|
||||
|
||||
std::vector<int64_t> rows, cols, roff, coff;
|
||||
deviceVector<ComplexD> sbuf(1), rbuf(1);
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
// Self edge: purely local, via a bounce buffer (shares all the code).
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
EdgeRows(rowStart, me, L, me, rows);
|
||||
EdgeCols(L, me, cols);
|
||||
if ( rows.size() && cols.size() ){
|
||||
uint64_t ne = rows.size()*cols.size();
|
||||
if ( sbuf.size() < ne ) sbuf.resize(ne);
|
||||
std::vector<int64_t> roff2, coff2;
|
||||
Offsets1D(rows, cols, row0, roff, coff);
|
||||
Offsets2D(L, rows, cols, roff2, coff2);
|
||||
if ( dir > 0 ) {
|
||||
MoveEdge(1, rows1d, ld1, roff, coff, &sbuf[0]);
|
||||
MoveEdge(0, &A.data[0], L.mloc, roff2, coff2, &sbuf[0]);
|
||||
} else {
|
||||
MoveEdge(1, &A.data[0], L.mloc, roff2, coff2, &sbuf[0]);
|
||||
MoveEdge(0, rows1d, ld1, roff, coff, &sbuf[0]);
|
||||
}
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
// Tournament over all pairs: odd modulus M, partner = (r - me) mod M.
|
||||
// Every unordered pair meets exactly once; partner==me or >=P is a bye.
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
int M = (P%2) ? P : P+1;
|
||||
for(int r=0;r<M;r++){
|
||||
int partner = (int)(((int64_t)r - me + 2L*M) % M);
|
||||
if ( partner == me || partner >= P ) continue;
|
||||
|
||||
// outbound edge: my (dir>0 ? 1D rows : 2D data) -> partner
|
||||
// inbound edge: partner -> my (dir>0 ? 2D data : 1D rows)
|
||||
std::vector<int64_t> orow, ocol, irow, icol;
|
||||
if ( dir > 0 ) { EdgeRows(rowStart, me, L, partner, orow); EdgeCols(L, partner, ocol);
|
||||
EdgeRows(rowStart, partner, L, me, irow); EdgeCols(L, me, icol); }
|
||||
else { EdgeRows(rowStart, partner, L, me, orow); EdgeCols(L, me, ocol);
|
||||
EdgeRows(rowStart, me, L, partner, irow); EdgeCols(L, partner, icol); }
|
||||
|
||||
uint64_t nout = orow.size()*ocol.size();
|
||||
uint64_t nin = irow.size()*icol.size();
|
||||
if ( !(nout || nin) ) continue; // both ends compute this identically
|
||||
|
||||
uint64_t nmax = std::max(nout,nin); // symmetric padded transfer
|
||||
if ( sbuf.size() < nmax ) sbuf.resize(nmax);
|
||||
if ( rbuf.size() < nmax ) rbuf.resize(nmax);
|
||||
|
||||
if ( nout ){
|
||||
if ( dir > 0 ) { Offsets1D(orow, ocol, row0, roff, coff);
|
||||
MoveEdge(1, rows1d, ld1, roff, coff, &sbuf[0]); }
|
||||
else { Offsets2D(L, orow, ocol, roff, coff);
|
||||
MoveEdge(1, &A.data[0], L.mloc, roff, coff, &sbuf[0]); }
|
||||
}
|
||||
grid->SendToRecvFrom((void *)&sbuf[0], partner,
|
||||
(void *)&rbuf[0], partner,
|
||||
nmax*sizeof(ComplexD));
|
||||
if ( nin ){
|
||||
if ( dir > 0 ) { Offsets2D(L, irow, icol, roff, coff);
|
||||
MoveEdge(0, &A.data[0], L.mloc, roff, coff, &rbuf[0]); }
|
||||
else { Offsets1D(irow, icol, row0, roff, coff);
|
||||
MoveEdge(0, rows1d, ld1, roff, coff, &rbuf[0]); }
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void RowsToCyclic(GridBase *grid, const std::vector<int64_t> &rowStart,
|
||||
ComplexD *rows1d, int64_t myrows, BlockCyclicMatrix &A)
|
||||
{ Redistribute(+1, grid, rowStart, rows1d, myrows, A); }
|
||||
|
||||
static void CyclicToRows(GridBase *grid, const std::vector<int64_t> &rowStart,
|
||||
BlockCyclicMatrix &A, ComplexD *rows1d, int64_t myrows)
|
||||
{ Redistribute(-1, grid, rowStart, rows1d, myrows, A); }
|
||||
};
|
||||
|
||||
NAMESPACE_END(Grid);
|
||||
@@ -0,0 +1,291 @@
|
||||
/*************************************************************************************
|
||||
|
||||
Grid physics library, www.github.com/paboyle/Grid
|
||||
|
||||
Source file: ./Grid/algorithms/multigrid/BlockCyclicSchurInverse.h
|
||||
|
||||
Copyright (C) 2026
|
||||
|
||||
Author: Peter Boyle <pboyle@bnl.gov>
|
||||
|
||||
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.
|
||||
|
||||
See the full license in the file "LICENSE" in the top level distribution
|
||||
directory
|
||||
*************************************************************************************/
|
||||
/* END LEGAL */
|
||||
#pragma once
|
||||
|
||||
#include <Grid/algorithms/blas/BatchedInverse.h>
|
||||
#include <Grid/algorithms/multigrid/BlockCyclicSumma.h>
|
||||
|
||||
NAMESPACE_BEGIN(Grid);
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
// Stage 3 of the 2D distributed dense inverse: the recursive Schur
|
||||
// complement on a block-cyclic matrix, in place.
|
||||
//
|
||||
// The nine-step algebra is IDENTICAL to RecursiveSchurInverse (1D); what
|
||||
// changes is the decomposition. The recursion splits the GLOBAL INDEX
|
||||
// RANGE at the block boundary nearest the midpoint -- not the rank range --
|
||||
// so every rank owns part of every sub-block at every depth, and the
|
||||
// ownership gating (inI/inJ, dummy operands, zero-width rank ranges) of the
|
||||
// 1D scheme has no analogue here: it is simply gone.
|
||||
//
|
||||
// I = [c0,m) J = [m,c1) (block-aligned, m the mid block boundary)
|
||||
// 1. recurse I : A11 -> A11inv (in place)
|
||||
// 2. Bt = A11inv . A12 (scratch, I x J)
|
||||
// 3. Ct = A21 . A11inv (scratch, J x I)
|
||||
// 4. A22 -= A21 . Bt == S (in place)
|
||||
// 5. recurse J : S -> Sinv (in place)
|
||||
// 6. Tt = Sinv . Ct (scratch, J x I)
|
||||
// 7. Ut = Bt . Sinv (scratch, I x J)
|
||||
// 8. A11 += Ut . Ct == X11 (in place)
|
||||
// 9. A12 = -Ut , A21 = -Tt (window copies)
|
||||
//
|
||||
// SCRATCH SHARING. Four full-size block-cyclic scratch matrices (Bt, Ct,
|
||||
// Tt, Ut) serve the ENTIRE tree, used through windows. This is safe at
|
||||
// every depth because of a window-disjointness invariant:
|
||||
//
|
||||
// * every temporary of a node has its row range in one half of the
|
||||
// node's window and its column range in the other (I x J or J x I);
|
||||
// * everything any DESCENDANT touches -- its A windows and its own
|
||||
// temporaries -- has BOTH ranges inside a single half (I x I during
|
||||
// step 1, J x J during step 5).
|
||||
//
|
||||
// Hence a descendant window and a live ancestor temporary always differ in
|
||||
// at least one dimension by disjoint ranges. Only Bt and Ct are live
|
||||
// across the step-5 recursion (Tt, Ut are written after it), and both are
|
||||
// covered by the invariant.
|
||||
//
|
||||
// LEAF. A leaf is a single diagonal block, and block (b,b) of a
|
||||
// block-cyclic layout lives ENTIRELY on rank (b%Pr, b%Pc). The leaf
|
||||
// inversion is therefore purely local -- pack the strided block dense,
|
||||
// GridBLASInverse, unpack -- with NO communication and no assembly. The
|
||||
// leaf-assembly transport question of the 1D scheme does not arise.
|
||||
// Successive leaves cycle over ranks, so leaf work is naturally spread.
|
||||
//
|
||||
// COMMUNICATION. Every transfer in the whole inversion is a
|
||||
// SendToRecvFrom inside BlockCyclicSumma's rings: pure point-to-point, no
|
||||
// collectives on the critical path, deterministic summation order (so
|
||||
// repeated inversions are bitwise identical). ReportTelemetry() is the
|
||||
// one optional exception: it performs reductions, and is only ever called
|
||||
// explicitly by a caller who wants the numbers.
|
||||
//
|
||||
// NUMERICS. No pivoting, exactly as the 1D scheme: every A11 and every
|
||||
// Schur complement met on the way down must be non-singular. The growth
|
||||
// telemetry stands in for pivoting; note the recursion splits differently
|
||||
// from the 1D rank-range tree, so DIFFERENT sub-blocks are inverted and
|
||||
// telemetry values are NOT comparable with the 1D implementation's --
|
||||
// re-baseline, do not compare.
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
class BlockCyclicSchurInverse
|
||||
{
|
||||
public:
|
||||
BlockCyclicSumma SUMMA;
|
||||
GridBLASInverse INV;
|
||||
|
||||
// Telemetry: accumulated LOCALLY, no comms unless ReportTelemetry().
|
||||
double telLeafMaxInv;
|
||||
uint64_t nLeaf;
|
||||
uint64_t nNode;
|
||||
double tLeaf;
|
||||
double tGemm; // wall in Multiply calls (comms+gemm)
|
||||
double tCopy;
|
||||
|
||||
BlockCyclicSchurInverse()
|
||||
{
|
||||
telLeafMaxInv = 0.0;
|
||||
nLeaf = nNode = 0;
|
||||
tLeaf = tGemm = tCopy = 0.0;
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Window copy-scale: Dst[i0:i1, j0:j1] = alpha * Src[same window].
|
||||
// Both share one layout, so the local bands coincide; pure local kernel.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
void WindowCopyScale(ComplexD alpha,
|
||||
BlockCyclicMatrix &Src, BlockCyclicMatrix &Dst,
|
||||
int64_t i0, int64_t i1, int64_t j0, int64_t j1)
|
||||
{
|
||||
BlockCyclicLayout &L = Dst.layout;
|
||||
GRID_ASSERT( Src.layout.N==L.N && Src.layout.nb==L.nb );
|
||||
GRID_ASSERT( Src.layout.Pr==L.Pr && Src.layout.Pc==L.Pc );
|
||||
int64_t li0,li1, lj0,lj1;
|
||||
L.RowRange(i0,i1, li0,li1);
|
||||
L.ColRange(j0,j1, lj0,lj1);
|
||||
int64_t m = li1-li0, n = lj1-lj0;
|
||||
if ( !(m && n) ) return;
|
||||
ComplexD *src = Src.LocalWindow(li0,lj0);
|
||||
ComplexD *dst = Dst.LocalWindow(li0,lj0);
|
||||
int64_t ldS = Src.layout.mloc;
|
||||
int64_t ldD = L.mloc;
|
||||
tCopy -= usecond();
|
||||
accelerator_for(idx, (uint64_t)(m*n), 1, {
|
||||
int64_t jj = idx / m;
|
||||
int64_t ii = idx - jj*m;
|
||||
dst[ii + jj*ldD] = alpha*src[ii + jj*ldS];
|
||||
});
|
||||
tCopy += usecond();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Leaf: single diagonal block (b,b), entirely on rank (b%Pr, b%Pc).
|
||||
// Local pack -> dense inverse -> unpack; every other rank does nothing
|
||||
// and needs no synchronisation: the next SUMMA's rings pair them up.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
void Leaf(BlockCyclicMatrix &A, int64_t b)
|
||||
{
|
||||
BlockCyclicLayout &L = A.layout;
|
||||
nLeaf++;
|
||||
if ( (int)(b % L.Pr) != L.prow ) return;
|
||||
if ( (int)(b % L.Pc) != L.pcol ) return;
|
||||
|
||||
tLeaf -= usecond();
|
||||
int64_t g0 = b*L.nb;
|
||||
int64_t g1 = std::min(L.N, g0+L.nb);
|
||||
int64_t w = g1-g0;
|
||||
int64_t lr0,lr1, lc0,lc1;
|
||||
L.RowRange(g0,g1, lr0,lr1);
|
||||
L.ColRange(g0,g1, lc0,lc1);
|
||||
GRID_ASSERT( lr1-lr0 == w );
|
||||
GRID_ASSERT( lc1-lc0 == w );
|
||||
|
||||
// Pack the strided block dense (inverseBatched assumes lda == w).
|
||||
deviceVector<ComplexD> dense((uint64_t)w*w);
|
||||
{
|
||||
ComplexD *src = A.LocalWindow(lr0,lc0);
|
||||
ComplexD *dst = &dense[0];
|
||||
int64_t ld = L.mloc;
|
||||
accelerator_for(idx, (uint64_t)(w*w), 1, {
|
||||
int64_t jj = idx / w;
|
||||
int64_t ii = idx - jj*w;
|
||||
dst[ii + jj*w] = src[ii + jj*ld];
|
||||
});
|
||||
}
|
||||
{
|
||||
deviceVector<ComplexD*> bp(1);
|
||||
std::vector<ComplexD*> ptr(1);
|
||||
ptr[0] = &dense[0];
|
||||
acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*));
|
||||
INV.inverseBatched(w, bp);
|
||||
}
|
||||
{
|
||||
ComplexD *src = &dense[0];
|
||||
ComplexD *dst = A.LocalWindow(lr0,lc0);
|
||||
int64_t ld = L.mloc;
|
||||
accelerator_for(idx, (uint64_t)(w*w), 1, {
|
||||
int64_t jj = idx / w;
|
||||
int64_t ii = idx - jj*w;
|
||||
dst[ii + jj*ld] = src[ii + jj*w];
|
||||
});
|
||||
}
|
||||
// Growth telemetry, local only.
|
||||
{
|
||||
std::vector<ComplexD> h((uint64_t)w*w);
|
||||
acceleratorCopyFromDevice(&dense[0], &h[0], h.size()*sizeof(ComplexD));
|
||||
double mx = 0.0;
|
||||
for(auto &z : h){
|
||||
double re=z.real(), im=z.imag();
|
||||
mx = std::max(mx, re*re+im*im);
|
||||
}
|
||||
telLeafMaxInv = std::max(telLeafMaxInv, std::sqrt(mx));
|
||||
}
|
||||
tLeaf += usecond();
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// The recursion, on global BLOCK range [b0,b1). SPMD: every rank calls
|
||||
// with identical arguments; there is no ownership gating to get wrong.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
void SchurNode(BlockCyclicMatrix &A,
|
||||
BlockCyclicMatrix &Bt, BlockCyclicMatrix &Ct,
|
||||
BlockCyclicMatrix &Tt, BlockCyclicMatrix &Ut,
|
||||
int64_t b0, int64_t b1)
|
||||
{
|
||||
BlockCyclicLayout &L = A.layout;
|
||||
int64_t span = b1-b0;
|
||||
GRID_ASSERT( span >= 1 );
|
||||
if ( span == 1 ) { Leaf(A, b0); return; }
|
||||
nNode++;
|
||||
|
||||
int64_t bm = b0 + span/2;
|
||||
int64_t c0 = b0*L.nb;
|
||||
int64_t m = bm*L.nb;
|
||||
int64_t c1 = std::min(L.N, b1*L.nb);
|
||||
ComplexD one (1.0,0.0), mone(-1.0,0.0), zero(0.0,0.0);
|
||||
|
||||
// 1. A11 -> A11inv
|
||||
SchurNode(A,Bt,Ct,Tt,Ut, b0,bm);
|
||||
|
||||
tGemm -= usecond();
|
||||
// 2. Bt[I,J] = A11inv . A12
|
||||
SUMMA.Multiply(one, A, A, zero, Bt, c0,m, m,c1, c0,m );
|
||||
// 3. Ct[J,I] = A21 . A11inv
|
||||
SUMMA.Multiply(one, A, A, zero, Ct, m,c1, c0,m, c0,m );
|
||||
// 4. A22 -= A21 . Bt (the Schur complement, in place)
|
||||
SUMMA.Multiply(mone, A, Bt, one, A, m,c1, m,c1, c0,m );
|
||||
tGemm += usecond();
|
||||
|
||||
// 5. S -> Sinv (Bt, Ct live across this call: see invariant)
|
||||
SchurNode(A,Bt,Ct,Tt,Ut, bm,b1);
|
||||
|
||||
tGemm -= usecond();
|
||||
// 6. Tt[J,I] = Sinv . Ct
|
||||
SUMMA.Multiply(one, A, Ct, zero, Tt, m,c1, c0,m, m,c1 );
|
||||
// 7. Ut[I,J] = Bt . Sinv
|
||||
SUMMA.Multiply(one, Bt, A, zero, Ut, c0,m, m,c1, m,c1 );
|
||||
// 8. A11 += Ut . Ct
|
||||
SUMMA.Multiply(one, Ut, Ct, one, A, c0,m, c0,m, m,c1 );
|
||||
tGemm += usecond();
|
||||
|
||||
// 9. Off-diagonal signs
|
||||
WindowCopyScale(mone, Ut, A, c0,m, m,c1);
|
||||
WindowCopyScale(mone, Tt, A, m,c1, c0,m);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// PUBLIC ENTRY. In-place inverse of the whole matrix. Scratch (4x the
|
||||
// matrix footprint) is allocated here and released on return.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
void Invert(BlockCyclicMatrix &A)
|
||||
{
|
||||
BlockCyclicLayout &L = A.layout;
|
||||
GRID_ASSERT( L.N >= 1 );
|
||||
int64_t nblocks = (L.N + L.nb - 1)/L.nb;
|
||||
|
||||
BlockCyclicMatrix Bt(A.grid, L.N, L.nb, L.Pr, L.Pc);
|
||||
BlockCyclicMatrix Ct(A.grid, L.N, L.nb, L.Pr, L.Pc);
|
||||
BlockCyclicMatrix Tt(A.grid, L.N, L.nb, L.Pr, L.Pc);
|
||||
BlockCyclicMatrix Ut(A.grid, L.N, L.nb, L.Pr, L.Pc);
|
||||
|
||||
telLeafMaxInv = 0.0;
|
||||
nLeaf = nNode = 0;
|
||||
tLeaf = tGemm = tCopy = 0.0;
|
||||
|
||||
SchurNode(A, Bt,Ct,Tt,Ut, 0, nblocks);
|
||||
}
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// Optional, and the ONLY place any reduction happens: call it if you
|
||||
// want the numbers, never from Invert.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
void ReportTelemetry(GridBase *grid)
|
||||
{
|
||||
RealD mx = telLeafMaxInv;
|
||||
grid->GlobalMax(mx);
|
||||
std::cout << GridLogMessage << "BlockCyclicSchurInverse:"
|
||||
<< " nodes " << nNode << " leaves " << nLeaf
|
||||
<< " max|leafinv| " << mx
|
||||
<< " (boss secs: gemm+comms " << tGemm/1.0e6
|
||||
<< " leaf " << tLeaf/1.0e6
|
||||
<< " copy " << tCopy/1.0e6 << ")"
|
||||
<< std::endl;
|
||||
}
|
||||
};
|
||||
|
||||
NAMESPACE_END(Grid);
|
||||
@@ -0,0 +1,299 @@
|
||||
/*************************************************************************************
|
||||
|
||||
Grid physics library, www.github.com/paboyle/Grid
|
||||
|
||||
Source file: ./Grid/algorithms/multigrid/BlockCyclicSumma.h
|
||||
|
||||
Copyright (C) 2026
|
||||
|
||||
Author: Peter Boyle <pboyle@bnl.gov>
|
||||
|
||||
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.
|
||||
|
||||
See the full license in the file "LICENSE" in the top level distribution
|
||||
directory
|
||||
*************************************************************************************/
|
||||
/* END LEGAL */
|
||||
#pragma once
|
||||
|
||||
#include <Grid/algorithms/multigrid/BlockCyclic.h>
|
||||
|
||||
NAMESPACE_BEGIN(Grid);
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
// Stage 2 of the 2D distributed dense inverse: the windowed SUMMA product
|
||||
//
|
||||
// C[i0:i1, j0:j1] <- beta C[i0:i1, j0:j1]
|
||||
// + alpha A[i0:i1, k0:k1] . B[k0:k1, j0:j1]
|
||||
//
|
||||
// on block-cyclic matrices sharing one BlockCyclicLayout. All ranges are
|
||||
// BLOCK ALIGNED (multiples of nb, or N itself at the top end): stage 3's
|
||||
// recursion splits on block boundaries, so nothing else is ever needed, and
|
||||
// alignment makes every local window a contiguous band of local storage
|
||||
// (BlockCyclic.h RangeToLocal, Test_blockcyclic T7).
|
||||
//
|
||||
// Transport is PURE POINT-TO-POINT: SendToRecvFrom on explicit world ranks
|
||||
// computed from the row-major rank convention. No collectives of any kind
|
||||
// -- no sub-communicators, no broadcast, no allgather -- by design: the
|
||||
// collective pathologies measured on this machine (MPI_Allgatherv at
|
||||
// ~0.18 MB/s with skewed counts, mpir_request.h:508 aborts, allreduce at
|
||||
// 39% of the P2P rate) motivated this implementation. SendToRecvFrom is
|
||||
// the most exercised device-buffer path in Grid and the only one never
|
||||
// implicated.
|
||||
//
|
||||
// Algorithm: round-based ring allgather SUMMA. The k range is processed in
|
||||
// rounds of Pc consecutive global blocks. Within a round
|
||||
//
|
||||
// * process column c owns at most one A panel (blocks s with s%Pc == c);
|
||||
// the Pc panels circulate around each process-ROW ring in Pc-1 steps;
|
||||
// * process row r owns up to ceil(Pc/Pr) B panels (blocks s%Pr == r);
|
||||
// they circulate around each process-COLUMN ring in Pr-1 steps;
|
||||
// * every rank then accumulates Cloc += alpha * Apanel_s . Bpanel_s
|
||||
// for each block s of the round, in ascending s: a fixed summation
|
||||
// order, so REPEATED RUNS ARE BITWISE IDENTICAL (no reduction, no
|
||||
// order ambiguity -- the property the P2P doctrine buys).
|
||||
//
|
||||
// Ring chunks are PADDED to a fixed size (full nb panels, fixed
|
||||
// panels-per-origin): SendToRecvFrom carries one byte count for both
|
||||
// directions, so symmetric transfers eliminate every variable-size edge
|
||||
// case at a worst-case ~1/Pc extra volume on ragged rounds. Padding is
|
||||
// never read: GEMMs address only the leading nb_s x width of each slot.
|
||||
//
|
||||
// Per-rank received volume: (k-extent) * (mloc_i + nloc_j) elements --
|
||||
// the N^2 (1/Pr + 1/Pc) SUMMA optimum, ~sqrt(P)/2 below the 1D scheme.
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
class BlockCyclicMatrix
|
||||
{
|
||||
public:
|
||||
GridBase *grid; // borrowed, never owned
|
||||
BlockCyclicLayout layout;
|
||||
deviceVector<ComplexD> data; // column major, ld = layout.mloc
|
||||
|
||||
BlockCyclicMatrix(GridBase *g, int64_t N, int64_t nb, int Pr, int Pc)
|
||||
: grid(g),
|
||||
layout(N, nb, Pr, Pc, g->ThisRank())
|
||||
{
|
||||
GRID_ASSERT( Pr*Pc == g->ProcessorCount() );
|
||||
uint64_t sz = (uint64_t)layout.mloc*layout.nloc;
|
||||
data.resize( sz ? sz : 1 );
|
||||
}
|
||||
|
||||
ComplexD *LocalWindow(int64_t li, int64_t lj)
|
||||
{
|
||||
return &data[0] + li + lj*layout.mloc;
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
// TEST-SCALE import/export of a replicated global matrix (host, O(N^2)
|
||||
// loops, one collective in Export). For unit tests and the stage-3
|
||||
// oracle only; production data enters through the direct block-cyclic
|
||||
// import, never through these.
|
||||
/////////////////////////////////////////////////////////////////////////
|
||||
void ImportGlobal(const std::vector<ComplexD> &G)
|
||||
{
|
||||
int64_t N = layout.N;
|
||||
GRID_ASSERT( (int64_t)G.size() == N*N );
|
||||
std::vector<ComplexD> h((uint64_t)layout.mloc*layout.nloc, ComplexD(0.0,0.0));
|
||||
for(int64_t j=0;j<N;j++){
|
||||
for(int64_t i=0;i<N;i++){
|
||||
if ( layout.Owns(i,j) ) h[layout.LocalOffset(i,j)] = G[i + j*N];
|
||||
}
|
||||
}
|
||||
if ( h.size() )
|
||||
acceleratorCopyToDevice(&h[0], &data[0], h.size()*sizeof(ComplexD));
|
||||
}
|
||||
void ExportGlobal(std::vector<ComplexD> &G)
|
||||
{
|
||||
int64_t N = layout.N;
|
||||
G.assign((uint64_t)N*N, ComplexD(0.0,0.0));
|
||||
std::vector<ComplexD> h((uint64_t)layout.mloc*layout.nloc);
|
||||
if ( h.size() )
|
||||
acceleratorCopyFromDevice(&data[0], &h[0], h.size()*sizeof(ComplexD));
|
||||
for(int64_t j=0;j<N;j++){
|
||||
for(int64_t i=0;i<N;i++){
|
||||
if ( layout.Owns(i,j) ) G[i + j*N] = h[layout.LocalOffset(i,j)];
|
||||
}
|
||||
}
|
||||
if ( N ) grid->GlobalSumVector((ComplexD *)&G[0], (int)(N*N)); // zero-fill: exact
|
||||
}
|
||||
};
|
||||
|
||||
class BlockCyclicSumma
|
||||
{
|
||||
public:
|
||||
GridBLAS BLAS;
|
||||
|
||||
static int Overlap(int64_t a0,int64_t a1,int64_t b0,int64_t b1)
|
||||
{ return (a0 < b1) && (b0 < a1); }
|
||||
|
||||
void Multiply(ComplexD alpha,
|
||||
BlockCyclicMatrix &A,
|
||||
BlockCyclicMatrix &B,
|
||||
ComplexD beta,
|
||||
BlockCyclicMatrix &C,
|
||||
int64_t i0, int64_t i1,
|
||||
int64_t j0, int64_t j1,
|
||||
int64_t k0, int64_t k1)
|
||||
{
|
||||
BlockCyclicLayout &L = C.layout;
|
||||
GridBase *grid = C.grid;
|
||||
const int64_t N = L.N;
|
||||
const int64_t nb = L.nb;
|
||||
const int Pr = L.Pr, Pc = L.Pc;
|
||||
const int prow = L.prow, pcol = L.pcol;
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
// Conformability: one layout, one communicator, aligned ranges.
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
auto same = [&](BlockCyclicLayout &X){
|
||||
GRID_ASSERT( X.N==N ); GRID_ASSERT( X.nb==nb );
|
||||
GRID_ASSERT( X.Pr==Pr ); GRID_ASSERT( X.Pc==Pc );
|
||||
GRID_ASSERT( X.me==L.me );
|
||||
};
|
||||
same(A.layout); same(B.layout);
|
||||
GRID_ASSERT( A.grid==grid ); GRID_ASSERT( B.grid==grid );
|
||||
auto aligned = [&](int64_t g0,int64_t g1){
|
||||
GRID_ASSERT( 0<=g0 ); GRID_ASSERT( g0<=g1 ); GRID_ASSERT( g1<=N );
|
||||
GRID_ASSERT( g0%nb==0 ); GRID_ASSERT( (g1%nb==0)||(g1==N) );
|
||||
};
|
||||
aligned(i0,i1); aligned(j0,j1); aligned(k0,k1);
|
||||
GRID_ASSERT( k1 > k0 ); // pure scaling not supported here
|
||||
// In-place windows are legal only if the written window is disjoint
|
||||
// from anything read (stage 3 uses this; make violation loud).
|
||||
if ( &C==&A ) GRID_ASSERT( !Overlap(j0,j1,k0,k1) );
|
||||
if ( &C==&B ) GRID_ASSERT( !Overlap(i0,i1,k0,k1) );
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
// My local bands of the three windows (contiguous: T7).
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
int64_t li0,li1, lj0,lj1;
|
||||
L.RowRange(i0,i1, li0,li1);
|
||||
L.ColRange(j0,j1, lj0,lj1);
|
||||
const int64_t mloc_i = li1-li0; // my rows of the i window
|
||||
const int64_t nloc_j = lj1-lj0; // my cols of the j window
|
||||
|
||||
const int64_t kb0 = k0/nb;
|
||||
const int64_t kb1 = (k1+nb-1)/nb; // block-aligned or ==N: exact
|
||||
const int64_t S = (Pc + Pr - 1)/Pr; // max B panels per origin row
|
||||
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
// Round buffers, padded to fixed slot sizes (see header comment).
|
||||
// A: Pc slots of mloc_i x nb (slot c = panel of the block owned by c)
|
||||
// B: Pr slots of S x (nb x nloc_j)
|
||||
///////////////////////////////////////////////////////////////////////
|
||||
const uint64_t slotA = (uint64_t)mloc_i*nb;
|
||||
const uint64_t slotB1 = (uint64_t)nb*nloc_j; // one panel
|
||||
const uint64_t slotB = (uint64_t)S*slotB1;
|
||||
deviceVector<ComplexD> Abuf( slotA*Pc ? slotA*Pc : 1 );
|
||||
deviceVector<ComplexD> Bbuf( slotB*Pr ? slotB*Pr : 1 );
|
||||
|
||||
deviceVector<ComplexD *> ap(1), bp(1), cp(1);
|
||||
std::vector<ComplexD *> ptr(1);
|
||||
|
||||
int firstblock = 1;
|
||||
for(int64_t r0=kb0; r0<kb1; r0+=Pc){ // rounds of Pc k-blocks
|
||||
int64_t r1 = std::min(kb1, r0+Pc);
|
||||
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
// Pack MY panels of this round into my origin slots.
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
for(int64_t s=r0; s<r1; s++){
|
||||
int64_t nb_s = L.BlockSize(s);
|
||||
if ( (int)(s%Pc) == pcol && mloc_i ){ // my A panel: block-column s
|
||||
int64_t lc0, lc1;
|
||||
L.ColRange(s*nb, std::min(N,(s+1)*nb), lc0, lc1);
|
||||
GRID_ASSERT( lc1-lc0 == nb_s );
|
||||
ComplexD *src = A.LocalWindow(li0, lc0);
|
||||
ComplexD *dst = &Abuf[0] + slotA*pcol;
|
||||
int64_t ld = A.layout.mloc;
|
||||
int64_t m = mloc_i;
|
||||
accelerator_for(idx, (uint64_t)(m*nb_s), 1, {
|
||||
int64_t jj = idx / m;
|
||||
int64_t ii = idx - jj*m;
|
||||
dst[ii + jj*m] = src[ii + jj*ld];
|
||||
});
|
||||
}
|
||||
if ( (int)(s%Pr) == prow && nloc_j ){ // my B panel: block-row s
|
||||
int64_t lr0, lr1;
|
||||
L.RowRange(s*nb, std::min(N,(s+1)*nb), lr0, lr1);
|
||||
GRID_ASSERT( lr1-lr0 == nb_s );
|
||||
int64_t idxs = (s - r0 - ((prow - r0%Pr + Pr) % Pr)) / Pr; // my panel # in round
|
||||
GRID_ASSERT( idxs >= 0 ); GRID_ASSERT( idxs < S );
|
||||
ComplexD *src = B.LocalWindow(lr0, lj0);
|
||||
ComplexD *dst = &Bbuf[0] + slotB*prow + slotB1*idxs;
|
||||
int64_t ld = B.layout.mloc;
|
||||
int64_t nn = nloc_j;
|
||||
accelerator_for(idx, (uint64_t)(nb_s*nn), 1, {
|
||||
int64_t jj = idx / nb_s;
|
||||
int64_t ii = idx - jj*nb_s;
|
||||
dst[ii + jj*nb_s] = src[ii + jj*ld];
|
||||
});
|
||||
}
|
||||
}
|
||||
accelerator_barrier();
|
||||
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
// Ring allgather along my process ROW: Pc-1 symmetric steps. At
|
||||
// step t send the slot of origin (pcol-t+1), receive origin (pcol-t).
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
if ( Pc > 1 && slotA ){
|
||||
int dest = prow*Pc + (pcol+1)%Pc;
|
||||
int src = prow*Pc + (pcol-1+Pc)%Pc;
|
||||
for(int t=1;t<Pc;t++){
|
||||
int cs = (pcol - t + 1 + Pc*Pc) % Pc;
|
||||
int cr = (pcol - t + Pc*Pc) % Pc;
|
||||
grid->SendToRecvFrom((void *)(&Abuf[0]+slotA*cs), dest,
|
||||
(void *)(&Abuf[0]+slotA*cr), src,
|
||||
slotA*sizeof(ComplexD));
|
||||
}
|
||||
}
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
// Ring allgather along my process COLUMN: Pr-1 symmetric steps.
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
if ( Pr > 1 && slotB ){
|
||||
int dest = ((prow+1)%Pr)*Pc + pcol;
|
||||
int src = ((prow-1+Pr)%Pr)*Pc + pcol;
|
||||
for(int t=1;t<Pr;t++){
|
||||
int rs = (prow - t + 1 + Pr*Pr) % Pr;
|
||||
int rr = (prow - t + Pr*Pr) % Pr;
|
||||
grid->SendToRecvFrom((void *)(&Bbuf[0]+slotB*rs), dest,
|
||||
(void *)(&Bbuf[0]+slotB*rr), src,
|
||||
slotB*sizeof(ComplexD));
|
||||
}
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
// Local update, ascending s: fixed order, bitwise-reproducible.
|
||||
/////////////////////////////////////////////////////////////////////
|
||||
for(int64_t s=r0; s<r1; s++){
|
||||
int64_t nb_s = L.BlockSize(s);
|
||||
if ( !(mloc_i && nloc_j && nb_s) ) { firstblock = 0; continue; }
|
||||
int cA = (int)(s%Pc);
|
||||
int rB = (int)(s%Pr);
|
||||
int64_t idxs = (s - r0 - ((rB - r0%Pr + Pr) % Pr)) / Pr;
|
||||
ComplexD beta_use = firstblock ? beta : ComplexD(1.0,0.0);
|
||||
firstblock = 0;
|
||||
|
||||
ptr[0] = &Abuf[0] + slotA*cA;
|
||||
acceleratorCopyToDevice(&ptr[0], &ap[0], sizeof(ComplexD *));
|
||||
ptr[0] = &Bbuf[0] + slotB*rB + slotB1*idxs;
|
||||
acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD *));
|
||||
ptr[0] = C.LocalWindow(li0, lj0);
|
||||
acceleratorCopyToDevice(&ptr[0], &cp[0], sizeof(ComplexD *));
|
||||
|
||||
BLAS.gemmBatched(GridBLAS_OP_N, GridBLAS_OP_N,
|
||||
(int)mloc_i, (int)nloc_j, (int)nb_s,
|
||||
alpha, ap, (int)mloc_i,
|
||||
bp, (int)nb_s,
|
||||
beta_use, cp, (int)C.layout.mloc);
|
||||
BLAS.synchronise();
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
NAMESPACE_END(Grid);
|
||||
@@ -30,6 +30,8 @@ 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>
|
||||
|
||||
#include <unordered_map>
|
||||
|
||||
@@ -874,13 +876,48 @@ public:
|
||||
BlockRows S;
|
||||
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).
|
||||
////////////////////////////////////////////////////////////////
|
||||
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;
|
||||
RecursiveSchurInverse RSI(grid, N, rowStart, panelBytes);
|
||||
double t2 = usecond();
|
||||
RSI.Invert(S);
|
||||
double t3 = usecond();
|
||||
RSI.ReportTelemetry();
|
||||
: (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;
|
||||
GRID_ASSERT( nb >= 1 );
|
||||
std::cout << GridLogMessage << "DenseCoarseMatrix: 2D SCHUR invert, process grid "
|
||||
<< Pr << " x " << Pc << " nb " << nb
|
||||
<< " (pure P2P: redistribute + SUMMA rings + local leaves)" << std::endl;
|
||||
BlockCyclicMatrix A2(grid, N, nb, Pr, Pc);
|
||||
BlockCyclicSchurInverse RSI2;
|
||||
t2 = usecond();
|
||||
BlockCyclicRedistribute::RowsToCyclic(grid, rowStart, &S.data[0], nrows, A2);
|
||||
RSI2.Invert(A2);
|
||||
BlockCyclicRedistribute::CyclicToRows(grid, rowStart, A2, &S.data[0], nrows);
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user