mirror of
https://github.com/paboyle/Grid.git
synced 2026-10-07 00:08:06 +01:00
301 lines
10 KiB
C++
301 lines
10 KiB
C++
/*************************************************************************************
|
|
|
|
Grid physics library, www.github.com/paboyle/Grid
|
|
|
|
Source file: MultiRHSDeflation.h
|
|
|
|
Copyright (C) 2023
|
|
|
|
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.
|
|
|
|
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
|
|
|
|
NAMESPACE_BEGIN(Grid);
|
|
|
|
|
|
/* Need helper object for BLAS accelerated mrhs projection
|
|
|
|
i) MultiRHS Deflation
|
|
|
|
Import Evecs -> nev x vol x internal
|
|
Import vector of Lattice objects -> nrhs x vol x internal
|
|
=> Cij (nrhs x Nev) via GEMM.
|
|
=> Guess (nrhs x vol x internal) = C x evecs (via GEMM)
|
|
Export
|
|
|
|
|
|
ii) MultiRHS block projection
|
|
|
|
Import basis -> nblock x nbasis x (block x internal)
|
|
Import vector of fine lattice objects -> nblock x nrhs x (block x internal)
|
|
|
|
=> coarse_(nrhs x nbasis )^block = via batched GEMM
|
|
|
|
iii) Alternate interface:
|
|
Import higher dim Lattice object-> vol x nrhs layout
|
|
|
|
*/
|
|
template<class Field>
|
|
class MultiRHSDeflation
|
|
{
|
|
public:
|
|
|
|
typedef typename Field::scalar_type scalar;
|
|
typedef typename Field::scalar_object scalar_object;
|
|
typedef typename Field::vector_object vobj;
|
|
|
|
int nev;
|
|
std::vector<RealD> eval;
|
|
GridBase *grid;
|
|
uint64_t vol;
|
|
uint64_t words;
|
|
|
|
deviceVector<scalar> BLAS_E; // nev x vol -- the eigenbasis (up to a 1/sqrt(lambda))
|
|
deviceVector<scalar> BLAS_R; // nrhs x vol -- the sources
|
|
deviceVector<scalar> BLAS_G; // nrhs x vol -- the guess
|
|
deviceVector<scalar> BLAS_C; // nrhs x nev -- the coefficients
|
|
|
|
MultiRHSDeflation(){};
|
|
~MultiRHSDeflation(){ Deallocate(); };
|
|
|
|
void Deallocate(void)
|
|
{
|
|
nev=0;
|
|
grid=nullptr;
|
|
vol=0;
|
|
words=0;
|
|
// deviceVector is a std::vector with a device allocator: resize(0) drops
|
|
// the size and keeps the capacity, returning no device memory. Swapping
|
|
// with an empty vector destroys the buffer.
|
|
deviceVector<scalar>().swap(BLAS_E);
|
|
deviceVector<scalar>().swap(BLAS_R);
|
|
deviceVector<scalar>().swap(BLAS_C);
|
|
deviceVector<scalar>().swap(BLAS_G);
|
|
}
|
|
// Resident (non-evictable) device memory held by this deflator.
|
|
uint64_t DeviceBytes(void)
|
|
{
|
|
return (BLAS_E.capacity()+BLAS_R.capacity()+BLAS_C.capacity()+BLAS_G.capacity())*sizeof(scalar);
|
|
}
|
|
void Allocate(int _nev,GridBase *_grid)
|
|
{
|
|
nev=_nev;
|
|
grid=_grid;
|
|
vol = grid->lSites();
|
|
words = sizeof(scalar_object)/sizeof(scalar);
|
|
eval.resize(nev);
|
|
BLAS_E.resize (vol * words * nev );
|
|
std::cout << GridLogMessage << " Allocate for "<<nev<<" eigenvectors and volume "<<vol<<std::endl;
|
|
}
|
|
void ImportEigenVector(Field &evec,RealD &_eval, int ev)
|
|
{
|
|
// std::cout << " ev " <<ev<<" eval "<<_eval<< std::endl;
|
|
GRID_ASSERT(ev<eval.size());
|
|
eval[ev] = _eval;
|
|
|
|
int64_t offset = ev*vol*words;
|
|
autoView(v,evec,AcceleratorRead);
|
|
acceleratorCopyDeviceToDevice(&v[0],&BLAS_E[offset],sizeof(scalar_object)*vol);
|
|
|
|
}
|
|
void ImportEigenBasis(std::vector<Field> &evec,std::vector<RealD> &_eval)
|
|
{
|
|
ImportEigenBasis(evec,_eval,0,evec.size());
|
|
}
|
|
// Could use to import a batch of eigenvectors
|
|
void ImportEigenBasis(std::vector<Field> &evec,std::vector<RealD> &_eval, int _ev0, int _nev)
|
|
{
|
|
GRID_ASSERT(_ev0+_nev<=evec.size());
|
|
|
|
Allocate(_nev,evec[0].Grid());
|
|
|
|
// Imports a sub-batch of eigenvectors, _ev0, ..., _ev0+_nev-1
|
|
for(int e=0;e<nev;e++){
|
|
std::cout << "Importing eigenvector "<<e<<" evalue "<<_eval[_ev0+e]<<std::endl;
|
|
ImportEigenVector(evec[_ev0+e],_eval[_ev0+e],e);
|
|
}
|
|
}
|
|
/////////////////////////////////////////////////////////////////////////
|
|
// Sources as a vector of D-dimensional fields: each is one contiguous
|
|
// column of BLAS_R already, so import/export are straight copies.
|
|
// The eigenvectors may live on a D grid or on the Nrhs=1 D+1 grid of the
|
|
// same local volume; the two are the same bytes, so only the site count
|
|
// is checked, not grid identity.
|
|
/////////////////////////////////////////////////////////////////////////
|
|
void DeflateSources(std::vector<Field> &source,std::vector<Field> & guess)
|
|
{
|
|
int nrhs = source.size();
|
|
GRID_ASSERT(source.size()==guess.size());
|
|
GRID_ASSERT(grid->lSites() == guess[0].Grid()->lSites());
|
|
conformable(guess[0],source[0]);
|
|
|
|
int64_t vw = vol * words;
|
|
BLAS_R.resize(nrhs * vw); // cost free if size doesn't change
|
|
BLAS_G.resize(nrhs * vw); // cost free if size doesn't change
|
|
BLAS_C.resize(nev * nrhs);// cost free if size doesn't change
|
|
|
|
for(int r=0;r<nrhs;r++){
|
|
int64_t offset = r*vw;
|
|
autoView(v,source[r],AcceleratorRead);
|
|
acceleratorCopyDeviceToDevice(&v[0],&BLAS_R[offset],sizeof(scalar_object)*vol);
|
|
}
|
|
DeflateBLAS(nrhs);
|
|
for(int r=0;r<nrhs;r++){
|
|
int64_t offset = r*vw;
|
|
autoView(v,guess[r],AcceleratorWrite);
|
|
acceleratorCopyDeviceToDevice(&BLAS_G[offset],&v[0],sizeof(scalar_object)*vol);
|
|
}
|
|
}
|
|
/////////////////////////////////////////////////////////////////////////
|
|
// Sources as ONE D+1 dimensional multiRHS field, rhs innermost
|
|
// (dimension 0, undistributed) on an UNVECTORISED coarse space, so the
|
|
// site index is rhs + nrhs*x with x the D-dimensional site. The GEMM
|
|
// wants x contiguous per rhs, so the import is a permutation, one fused
|
|
// pass over the field; likewise the export. No slices.
|
|
/////////////////////////////////////////////////////////////////////////
|
|
void DeflateSources(const Field &source_mrhs, Field &guess_mrhs)
|
|
{
|
|
conformable(source_mrhs,guess_mrhs);
|
|
GridBase *hi = source_mrhs.Grid();
|
|
GRID_ASSERT(hi->_ndimension == grid->_ndimension+1);
|
|
GRID_ASSERT(hi->_processors[0] == 1);
|
|
GRID_ASSERT(hi->_simd_layout[0]== 1);
|
|
GRID_ASSERT(vobj::Nsimd() == 1);
|
|
int nrhs = hi->_fdimensions[0];
|
|
GRID_ASSERT(hi->lSites() == nrhs*vol);
|
|
|
|
int64_t vw = vol * words;
|
|
BLAS_R.resize(nrhs * vw);
|
|
BLAS_G.resize(nrhs * vw);
|
|
BLAS_C.resize(nev * nrhs);
|
|
|
|
ImportSourcesMrhs(source_mrhs,nrhs);
|
|
DeflateBLAS(nrhs);
|
|
ExportGuessMrhs(guess_mrhs,nrhs);
|
|
}
|
|
void ImportSourcesMrhs(const Field &source_mrhs,int nrhs)
|
|
{
|
|
const int64_t lwords = words, lnrhs = nrhs, vw = vol*words;
|
|
autoView(v,source_mrhs,AcceleratorRead);
|
|
auto vp = &v[0];
|
|
scalar *R = &BLAS_R[0];
|
|
accelerator_for(scr, lnrhs*vol, 1, {
|
|
int64_t r = scr % lnrhs;
|
|
int64_t x = scr / lnrhs;
|
|
const scalar *s = (const scalar *)&vp[scr];
|
|
for(int64_t w=0;w<lwords;w++) R[r*vw + x*lwords + w] = s[w];
|
|
});
|
|
}
|
|
void ExportGuessMrhs(Field &guess_mrhs,int nrhs)
|
|
{
|
|
const int64_t lwords = words, lnrhs = nrhs, vw = vol*words;
|
|
autoView(v,guess_mrhs,AcceleratorWrite);
|
|
auto vp = &v[0];
|
|
scalar *G = &BLAS_G[0];
|
|
accelerator_for(scr, lnrhs*vol, 1, {
|
|
int64_t r = scr % lnrhs;
|
|
int64_t x = scr / lnrhs;
|
|
scalar *g = (scalar *)&vp[scr];
|
|
for(int64_t w=0;w<lwords;w++) g[w] = G[r*vw + x*lwords + w];
|
|
});
|
|
}
|
|
/////////////////////////////////////////////////////////////////////////
|
|
// BLAS_R (nrhs columns of vol*words) -> BLAS_G, through the imported
|
|
// eigenbasis: C = E^dag R / lambda, G = E C.
|
|
/////////////////////////////////////////////////////////////////////////
|
|
void DeflateBLAS(int nrhs)
|
|
{
|
|
int64_t vw = vol * words;
|
|
RealD t0 = usecond();
|
|
|
|
/*
|
|
* in Fortran column major notation (cuBlas order)
|
|
*
|
|
* Exe = [e1(x)][..][en(x)]
|
|
*
|
|
* Rxr = [r1(x)][..][rm(x)]
|
|
*
|
|
* C_er = E^dag R
|
|
* C_er = C_er / lambda_e
|
|
* G_xr = Exe Cer
|
|
*/
|
|
deviceVector<scalar *> Ed(1);
|
|
deviceVector<scalar *> Rd(1);
|
|
deviceVector<scalar *> Cd(1);
|
|
deviceVector<scalar *> Gd(1);
|
|
|
|
scalar * Eh = & BLAS_E[0];
|
|
scalar * Rh = & BLAS_R[0];
|
|
scalar * Ch = & BLAS_C[0];
|
|
scalar * Gh = & BLAS_G[0];
|
|
|
|
acceleratorPut(Ed[0],Eh);
|
|
acceleratorPut(Rd[0],Rh);
|
|
acceleratorPut(Cd[0],Ch);
|
|
acceleratorPut(Gd[0],Gh);
|
|
|
|
GridBLAS BLAS;
|
|
|
|
/////////////////////////////////////////
|
|
// C_er = E^dag R
|
|
/////////////////////////////////////////
|
|
BLAS.gemmBatched(GridBLAS_OP_C,GridBLAS_OP_N,
|
|
nev,nrhs,vw,
|
|
scalar(1.0),
|
|
Ed,
|
|
Rd,
|
|
scalar(0.0), // wipe out C
|
|
Cd);
|
|
BLAS.synchronise();
|
|
|
|
GRID_ASSERT(BLAS_C.size()==nev*nrhs);
|
|
|
|
std::vector<scalar> HOST_C(BLAS_C.size()); // nrhs . nev -- the coefficients
|
|
acceleratorCopyFromDevice(&BLAS_C[0],&HOST_C[0],BLAS_C.size()*sizeof(scalar));
|
|
grid->GlobalSumVector(&HOST_C[0],nev*nrhs);
|
|
for(int e=0;e<nev;e++){
|
|
scalar lam(1.0/eval[e]); // in the coarse scalar: fp32 coarse spaces too
|
|
for(int r=0;r<nrhs;r++){
|
|
int off = e+nev*r;
|
|
HOST_C[off]=HOST_C[off] * lam;
|
|
// std::cout << "C["<<e<<"]["<<r<<"] ="<<HOST_C[off]<< " eval[e] "<<eval[e] <<std::endl;
|
|
}
|
|
}
|
|
acceleratorCopyToDevice(&HOST_C[0],&BLAS_C[0],BLAS_C.size()*sizeof(scalar));
|
|
|
|
|
|
/////////////////////////////////////////
|
|
// Guess G_xr = Exe Cer
|
|
/////////////////////////////////////////
|
|
BLAS.gemmBatched(GridBLAS_OP_N,GridBLAS_OP_N,
|
|
vw,nrhs,nev,
|
|
scalar(1.0),
|
|
Ed, // x . nev
|
|
Cd, // nev . nrhs
|
|
scalar(0.0),
|
|
Gd);
|
|
BLAS.synchronise();
|
|
RealD t1 = usecond();
|
|
std::cout << GridLogMessage << "MultiRHSDeflation for "<<nrhs<<" sources with "<<nev<<" eigenvectors took " << (t1-t0)/1e3 <<" ms"<<std::endl;
|
|
}
|
|
};
|
|
|
|
NAMESPACE_END(Grid);
|