Faster leaf inverse

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<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();

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,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;

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
+# 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.
-export SCHUR2D_LEAF_SPAN=9;  run_cell   I2_span09   6  0.1  2  gcr  gcr     # 32 leaves of W=4320 (300 MB)
+# Span 9 optimal: above it the SERIAL leaf chain (32 x 0.53 s at W=4320, 16 x 1.84 s at 8640,
-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
+# 287 ranks waiting) dominates; below it the partial rings return.
-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
+#export SCHUR2D_LEAF_SPAN=1;  run_cell   I1_span01   6  0.1  2  gcr  gcr     # 132.5 s
-unset SCHUR2D_LEAF_SPAN
+#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"