diff --git a/Grid/algorithms/blas/BatchedInverse.h b/Grid/algorithms/blas/BatchedInverse.h new file mode 100644 index 000000000..6549e1770 --- /dev/null +++ b/Grid/algorithms/blas/BatchedInverse.h @@ -0,0 +1,282 @@ +/************************************************************************************* + + Grid physics library, www.github.com/paboyle/Grid + + Source file: BatchedInverse.h + + Copyright (C) 2026 + +Author: Peter Boyle + + 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. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License along + with this program; if not, write to the Free Software Foundation, Inc., + 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA. + + See the full license in the file "LICENSE" in the top level distribution directory +*************************************************************************************/ +/* END LEGAL */ +#pragma once + +#include + +#ifdef GRID_HIP +#include +#endif +// GRID_CUDA: batched LU inversion lives in cuBLAS (getrfBatched/getriBatched); +// cublas_v2.h already included via BatchedBlas.h. +// GRID_SYCL: oneapi/mkl.hpp already included via BatchedBlas.h (lapack::getrf/getri). + +NAMESPACE_BEGIN(Grid); + +/////////////////////////////////////////////////////////////////////////////// +// GridBLASInverse: cross-platform batched dense matrix inversion. +// +// HIGH LEVEL contract (deliberately NOT a getrf/getrs interface): invert a +// batch of dense N x N matrices IN PLACE, +// +// A[i] <- A[i]^{-1} i = 0 .. batchCount-1 +// +// Layout: column major, lda = N, contiguous per batch element; pointer list +// exactly as GridBLAS::gemmBatched (deviceVector of device pointers). +// Each backend chooses HOW: +// HIP : rocSOLVER getrf_batched + getri_batched +// CUDA : cuBLAS getrfBatched + getriBatched (out-of-place getri; workspace +// hidden here, result copied back so the surface stays in-place) +// SYCL : oneMKL LAPACK getrf + getri per batch element (USM, in-order queue) +// CPU : Eigen PartialPivLU (the correctness oracle for all of the above) +// +// The int32 vendor-batched entry points bound N < 2^31 (asserted); the huge +// single-matrix ILP64 path (getrf_64 + blocked identity-getrs harvest, proven +// in the dense coarse-coarse setup at N=69120) migrates here as a batch==1 +// large-N dispatch in a follow-up -- the recursive Schur leaves are the +// batched consumers this surface is shaped for. +// +// NB GPU-backend call signatures are written to vendor documentation but the +// air-gapped development loop compiles only the CPU/Eigen path; verify the +// rocSOLVER/cuBLAS/oneMKL calls against headers on first device compile. +// Semantics are locked by the CPU unit test (Test_batched_blas). +/////////////////////////////////////////////////////////////////////////////// +class GridBLASInverse { +public: + +#ifdef GRID_HIP + // rocSOLVER runs on a rocblas_handle (distinct type from hipblasHandle_t) + static rocblas_handle & Handle(void) { + static rocblas_handle h; + static int init = 0; + if ( !init ) { + auto st = rocblas_create_handle(&h); + GRID_ASSERT(st == rocblas_status_success); + init = 1; + } + return h; + } +#endif +#ifdef GRID_CUDA + // cuBLAS batched LU shares the GridBLAS handle + static cublasHandle_t & Handle(void) { + GridBLAS::Init(); + return GridBLAS::gridblasHandle; + } +#endif +#ifdef GRID_SYCL + static sycl::queue * & Handle(void) { + GridBLAS::Init(); + return GridBLAS::gridblasHandle; + } +#endif + + GridBLASInverse() {}; + ~GridBLASInverse() {}; + + void inverseBatched(int64_t N, deviceVector &Amat) + { + int32_t batchCount = Amat.size(); + GRID_ASSERT(batchCount > 0); + +#ifdef GRID_HIP + GRID_ASSERT( N < 2147483647L ); + rocblas_int n = (rocblas_int)N; + rocblas_int lda = (rocblas_int)N; + + deviceVector ipiv((uint64_t)batchCount*N); + deviceVector info(batchCount); + + auto st1 = rocsolver_cgetrf_batched(Handle(), n, n, + (rocblas_float_complex *const *)&Amat[0], lda, + &ipiv[0], (rocblas_stride)N, + &info[0], batchCount); + GRID_ASSERT(st1 == rocblas_status_success); + auto st2 = rocsolver_cgetri_batched(Handle(), n, + (rocblas_float_complex *const *)&Amat[0], lda, + &ipiv[0], (rocblas_stride)N, + &info[0], batchCount); + GRID_ASSERT(st2 == rocblas_status_success); + accelerator_barrier(); + std::vector info_h(batchCount); + acceleratorCopyFromDevice(&info[0],&info_h[0],batchCount*sizeof(rocblas_int)); + for(int i=0;i abort loudly +#endif +#ifdef GRID_CUDA + GRID_ASSERT( N < 2147483647L ); + int n = (int)N; + + deviceVector ipiv((uint64_t)batchCount*N); + deviceVector info(batchCount); + + auto st1 = cublasCgetrfBatched(Handle(), n, + (cuComplex **)&Amat[0], n, + &ipiv[0], &info[0], batchCount); + GRID_ASSERT(st1 == CUBLAS_STATUS_SUCCESS); + + // getri is OUT of place: hidden workspace keeps the surface in-place + deviceVector work((uint64_t)batchCount*N*N); + deviceVector Cptr(batchCount); + std::vector Cptr_h(batchCount); + std::vector Aptr_h(batchCount); + for(int i=0;i info_h(batchCount); + acceleratorCopyFromDevice(&info[0],&info_h[0],batchCount*sizeof(int)); + for(int i=0;i Aptr_h(batchCount); + acceleratorCopyFromDevice(&Amat[0],&Aptr_h[0],batchCount*sizeof(ComplexF*)); + + int64_t lwf = oneapi::mkl::lapack::getrf_scratchpad_size >(*q,N,N,N); + int64_t lwi = oneapi::mkl::lapack::getri_scratchpad_size >(*q,N,N); + deviceVector scratchf(lwf); + deviceVector scratchi(lwi); + deviceVector ipiv(N); + for(int i=0;i*)Aptr_h[i],N,&ipiv[0], + (std::complex*)&scratchf[0],lwf); + oneapi::mkl::lapack::getri(*q,N, (std::complex*)Aptr_h[i],N,&ipiv[0], + (std::complex*)&scratchi[0],lwi); + } + q->wait(); +#endif +#if !defined(GRID_SYCL) && !defined(GRID_CUDA) && !defined(GRID_HIP) + // Reference implementation; the oracle the unit test locks semantics with. + thread_for (p, batchCount, { + Eigen::Map eA(Amat[p],N,N); + Eigen::PartialPivLU lu(eA); + eA = lu.inverse(); + }); +#endif + } + + void inverseBatched(int64_t N, deviceVector &Amat) + { + int32_t batchCount = Amat.size(); + GRID_ASSERT(batchCount > 0); + +#ifdef GRID_HIP + GRID_ASSERT( N < 2147483647L ); + rocblas_int n = (rocblas_int)N; + rocblas_int lda = (rocblas_int)N; + + deviceVector ipiv((uint64_t)batchCount*N); + deviceVector info(batchCount); + + auto st1 = rocsolver_zgetrf_batched(Handle(), n, n, + (rocblas_double_complex *const *)&Amat[0], lda, + &ipiv[0], (rocblas_stride)N, + &info[0], batchCount); + GRID_ASSERT(st1 == rocblas_status_success); + auto st2 = rocsolver_zgetri_batched(Handle(), n, + (rocblas_double_complex *const *)&Amat[0], lda, + &ipiv[0], (rocblas_stride)N, + &info[0], batchCount); + GRID_ASSERT(st2 == rocblas_status_success); + accelerator_barrier(); + std::vector info_h(batchCount); + acceleratorCopyFromDevice(&info[0],&info_h[0],batchCount*sizeof(rocblas_int)); + for(int i=0;i ipiv((uint64_t)batchCount*N); + deviceVector info(batchCount); + + auto st1 = cublasZgetrfBatched(Handle(), n, + (cuDoubleComplex **)&Amat[0], n, + &ipiv[0], &info[0], batchCount); + GRID_ASSERT(st1 == CUBLAS_STATUS_SUCCESS); + + deviceVector work((uint64_t)batchCount*N*N); + deviceVector Cptr(batchCount); + std::vector Cptr_h(batchCount); + std::vector Aptr_h(batchCount); + for(int i=0;i info_h(batchCount); + acceleratorCopyFromDevice(&info[0],&info_h[0],batchCount*sizeof(int)); + for(int i=0;i Aptr_h(batchCount); + acceleratorCopyFromDevice(&Amat[0],&Aptr_h[0],batchCount*sizeof(ComplexD*)); + + int64_t lwf = oneapi::mkl::lapack::getrf_scratchpad_size >(*q,N,N,N); + int64_t lwi = oneapi::mkl::lapack::getri_scratchpad_size >(*q,N,N); + deviceVector scratchf(lwf); + deviceVector scratchi(lwi); + deviceVector ipiv(N); + for(int i=0;i*)Aptr_h[i],N,&ipiv[0], + (std::complex*)&scratchf[0],lwf); + oneapi::mkl::lapack::getri(*q,N, (std::complex*)Aptr_h[i],N,&ipiv[0], + (std::complex*)&scratchi[0],lwi); + } + q->wait(); +#endif +#if !defined(GRID_SYCL) && !defined(GRID_CUDA) && !defined(GRID_HIP) + thread_for (p, batchCount, { + Eigen::Map eA(Amat[p],N,N); + Eigen::PartialPivLU lu(eA); + eA = lu.inverse(); + }); +#endif + } +}; + +NAMESPACE_END(Grid);