mirror of
https://github.com/paboyle/Grid.git
synced 2026-08-30 14:29:36 +01:00
Faster leaf inverse
This commit is contained in:
Peter Boyle
@@ -99,6 +99,49 @@ public:
|
||||
GridBLASInverse() {};
|
||||
~GridBLASInverse() {};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
// SINGLE large matrix, in place: blocked LU + identity solve.
|
||||
// HIP : rocsolver_zgetrf_64 (blocked, GEMM-based) then ONE zgetrs_64 with
|
||||
// the N x N identity as RHS (trsm-based) into a scratch N x N, copied
|
||||
// back. Extra device memory: N*N*16 B (1.2 GB at N=8640).
|
||||
// other backends: falls through to inverseBatched(batch 1).
|
||||
// Why: rocSOLVER getri_batched is a small-matrix routine -- measured 0.35 s
|
||||
// at N=1920, 0.53 s at 4320, 1.84 s at 8640 (n^1.2-1.8, i.e. overhead
|
||||
// bound) as the big-leaf inverse of the 2D Schur recursion, where 287
|
||||
// ranks wait on it. The _64 getrf/getrs pair is the path proven in the
|
||||
// 1-rank dense coarse-coarse setup at N=69120 (DenseCoarseMatrix.h).
|
||||
// Same in-place, column-major, lda=N contract as inverseBatched.
|
||||
// NB: written to rocSOLVER's documented z*_64 signatures; not compiled on
|
||||
// HIP in the air-gapped loop -- verify on first hipcc build.
|
||||
///////////////////////////////////////////////////////////////////////////
|
||||
void inverseLU(int64_t N, ComplexD *A)
|
||||
{
|
||||
#ifdef GRID_HIP
|
||||
rocblas_handle handle = Handle();
|
||||
deviceVector<int64_t> ipiv((uint64_t)N);
|
||||
deviceVector<int64_t> info(1);
|
||||
auto st1 = rocsolver_zgetrf_64(handle, N, N, (rocblas_double_complex *)A, N, &ipiv[0], &info[0]);
|
||||
GRID_ASSERT(st1 == rocblas_status_success);
|
||||
accelerator_barrier();
|
||||
int64_t info_h = -1; acceleratorCopyFromDevice(&info[0], &info_h, sizeof(int64_t));
|
||||
GRID_ASSERT(info_h == 0);
|
||||
deviceVector<ComplexD> X((uint64_t)N*N);
|
||||
{ ComplexD *x = &X[0]; const int64_t NN = N;
|
||||
accelerator_for(idx, (uint64_t)N*N, 1, { int64_t j = idx/NN, i = idx - j*NN; x[idx] = (i==j) ? ComplexD(1.0,0.0) : ComplexD(0.0,0.0); });
|
||||
accelerator_barrier(); }
|
||||
auto st2 = rocsolver_zgetrs_64(handle, rocblas_operation_none, N, N,
|
||||
(rocblas_double_complex *)A, N, &ipiv[0],
|
||||
(rocblas_double_complex *)&X[0], N);
|
||||
GRID_ASSERT(st2 == rocblas_status_success);
|
||||
accelerator_barrier();
|
||||
acceleratorCopyDeviceToDevice((void *)&X[0], (void *)A, (uint64_t)N*N*sizeof(ComplexD));
|
||||
#else
|
||||
deviceVector<ComplexD*> bp(1); std::vector<ComplexD*> ptr(1); ptr[0] = A;
|
||||
acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*));
|
||||
inverseBatched(N, bp);
|
||||
#endif
|
||||
}
|
||||
|
||||
void inverseBatched(int64_t N, deviceVector<ComplexF*> &Amat)
|
||||
{
|
||||
int32_t batchCount = Amat.size();
|
||||
|
||||
@@ -99,6 +99,7 @@ public:
|
||||
// single-block panels of exactly these levels). Instead: gather the
|
||||
// (s*nb)^2 sub-block to one rank, invert locally, scatter back.
|
||||
int leafSpan = -1;
|
||||
int leafLU = -1; // SCHUR2D_LEAF_LU=1: big-leaf inverse via GridBLASInverse::inverseLU (blocked getrf_64 + identity getrs_64) instead of getri_batched
|
||||
uint64_t nBigLeaf = 0; int64_t maxBigW = 0;
|
||||
double tBigGather = 0, tBigInv = 0, tBigScatter = 0;
|
||||
uint64_t nNode;
|
||||
@@ -285,11 +286,16 @@ public:
|
||||
|
||||
// ---- invert on root ----
|
||||
tBigInv -= usecond();
|
||||
if ( leafLU < 0 ) leafLU = getenv("SCHUR2D_LEAF_LU") ? atoi(getenv("SCHUR2D_LEAF_LU")) : 0;
|
||||
if ( me == root ) {
|
||||
if ( leafLU ) {
|
||||
INV.inverseLU(W, &dense[0]);
|
||||
} else {
|
||||
deviceVector<ComplexD*> bp(1); std::vector<ComplexD*> ptr(1); ptr[0] = &dense[0];
|
||||
acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*));
|
||||
INV.inverseBatched(W, bp);
|
||||
}
|
||||
}
|
||||
tBigInv += usecond();
|
||||
|
||||
// ---- scatter ----
|
||||
@@ -500,7 +506,8 @@ public:
|
||||
<< std::endl;
|
||||
if ( nBigLeaf ) {
|
||||
RealD ti = tBigInv/1.0e6; grid->GlobalMax(ti); // inverse runs on the root of each leaf: report the max over ranks
|
||||
std::cout << GridLogMessage << "BlockCyclicSchurInverse: BIG LEAVES (SCHUR2D_LEAF_SPAN=" << leafSpan << "): " << nBigLeaf
|
||||
std::cout << GridLogMessage << "BlockCyclicSchurInverse: BIG LEAVES (SCHUR2D_LEAF_SPAN=" << leafSpan
|
||||
<< (leafLU>0 ? ", SCHUR2D_LEAF_LU=1: getrf_64+getrs_64" : ", getri_batched") << "): " << nBigLeaf
|
||||
<< " leaves, max W " << maxBigW
|
||||
<< " boss secs: gather " << tBigGather/1.0e6 << " scatter " << tBigScatter/1.0e6
|
||||
<< " inverse (max over ranks) " << ti << std::endl;
|
||||
|
||||
@@ -158,11 +158,21 @@ run_cell () {
|
||||
# histogram ">= 2.0 MB" row -- should shrink to ~nothing at s>=9; handshake column = partner wait
|
||||
# Solver settings are the banked adaptive point (solve time is the same in every cell; only setup changes).
|
||||
# name Fso Fss Csn fine coarse
|
||||
export SCHUR2D_LEAF_SPAN=1; run_cell I1_span01 6 0.1 2 gcr gcr # reference: today's recursion to single-block leaves
|
||||
export SCHUR2D_LEAF_SPAN=9; run_cell I2_span09 6 0.1 2 gcr gcr # 32 leaves of W=4320 (300 MB)
|
||||
export SCHUR2D_LEAF_SPAN=18; run_cell I3_span18 6 0.1 2 gcr gcr # 16 leaves of W=8640 (1.2 GB): first span where every process row owns a piece
|
||||
export SCHUR2D_LEAF_SPAN=4; run_cell I4_span04 6 0.1 2 gcr gcr # 72 leaves of W=1920: if rocSOLVER is slow at 8640, the other end
|
||||
unset SCHUR2D_LEAF_SPAN
|
||||
# Sweep result (job 2026-08-27): span 1 132.5 s | 4 37.0 s | 9 27.6 s | 18 30.6 s; VERIFY 3.3939e-5 in all.
|
||||
# Span 9 optimal: above it the SERIAL leaf chain (32 x 0.53 s at W=4320, 16 x 1.84 s at 8640,
|
||||
# 287 ranks waiting) dominates; below it the partial rings return.
|
||||
#export SCHUR2D_LEAF_SPAN=1; run_cell I1_span01 6 0.1 2 gcr gcr # 132.5 s
|
||||
#export SCHUR2D_LEAF_SPAN=9; run_cell I2_span09 6 0.1 2 gcr gcr # 27.6 s
|
||||
#export SCHUR2D_LEAF_SPAN=18; run_cell I3_span18 6 0.1 2 gcr gcr # 30.6 s
|
||||
#export SCHUR2D_LEAF_SPAN=4; run_cell I4_span04 6 0.1 2 gcr gcr # 37.0 s
|
||||
# Leaf inverse: rocSOLVER getri_batched scales n^1.2-1.8 here (overhead bound). SCHUR2D_LEAF_LU=1
|
||||
# replaces it with blocked zgetrf_64 + one identity zgetrs_64 (the N=69120 1-rank path).
|
||||
# Readout: "BIG LEAVES ... inverse (max over ranks)" vs 0.53 s (span 9) / 1.84 s (span 18);
|
||||
# if the leaf gets ~3x cheaper, span 18's clean rings (10.3 GB/s) may win overall.
|
||||
export SCHUR2D_LEAF_LU=1
|
||||
export SCHUR2D_LEAF_SPAN=9; run_cell I5_span09_LU 6 0.1 2 gcr gcr # vs I2 27.6 s
|
||||
export SCHUR2D_LEAF_SPAN=18; run_cell I6_span18_LU 6 0.1 2 gcr gcr # vs I3 30.6 s
|
||||
unset SCHUR2D_LEAF_SPAN SCHUR2D_LEAF_LU
|
||||
|
||||
echo "========================================================="
|
||||
echo "summary"
|
||||
|
||||
Reference in New Issue
Block a user