From 68b9fa86d7f7ee54a8e6b7ba8f53d3b09c7b7f73 Mon Sep 17 00:00:00 2001 From: Peter Boyle Date: Thu, 27 Aug 2026 12:53:26 -0400 Subject: [PATCH] Faster leaf inverse --- Grid/algorithms/blas/BatchedInverse.h | 43 +++++++++++++++++++ .../multigrid/BlockCyclicSchurInverse.h | 15 +++++-- systems/Frontier/smoother_modes.job | 20 ++++++--- 3 files changed, 69 insertions(+), 9 deletions(-) diff --git a/Grid/algorithms/blas/BatchedInverse.h b/Grid/algorithms/blas/BatchedInverse.h index 6549e1770..40aaa107d 100644 --- a/Grid/algorithms/blas/BatchedInverse.h +++ b/Grid/algorithms/blas/BatchedInverse.h @@ -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 ipiv((uint64_t)N); + deviceVector 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 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 bp(1); std::vector ptr(1); ptr[0] = A; + acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*)); + inverseBatched(N, bp); +#endif + } + void inverseBatched(int64_t N, deviceVector &Amat) { int32_t batchCount = Amat.size(); diff --git a/Grid/algorithms/multigrid/BlockCyclicSchurInverse.h b/Grid/algorithms/multigrid/BlockCyclicSchurInverse.h index 1e7c73731..fbe2a4fd0 100644 --- a/Grid/algorithms/multigrid/BlockCyclicSchurInverse.h +++ b/Grid/algorithms/multigrid/BlockCyclicSchurInverse.h @@ -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,10 +286,15 @@ public: // ---- invert on root ---- tBigInv -= usecond(); + if ( leafLU < 0 ) leafLU = getenv("SCHUR2D_LEAF_LU") ? atoi(getenv("SCHUR2D_LEAF_LU")) : 0; if ( me == root ) { - deviceVector bp(1); std::vector ptr(1); ptr[0] = &dense[0]; - acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*)); - INV.inverseBatched(W, bp); + if ( leafLU ) { + INV.inverseLU(W, &dense[0]); + } else { + deviceVector bp(1); std::vector ptr(1); ptr[0] = &dense[0]; + acceleratorCopyToDevice(&ptr[0], &bp[0], sizeof(ComplexD*)); + INV.inverseBatched(W, bp); + } } tBigInv += usecond(); @@ -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; diff --git a/systems/Frontier/smoother_modes.job b/systems/Frontier/smoother_modes.job index 790ea6690..655ed5653 100644 --- a/systems/Frontier/smoother_modes.job +++ b/systems/Frontier/smoother_modes.job @@ -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"