pw_gatherscatter.h

Go to the documentation of this file.

#include "pw_basis.h"
#include "source_base/global_function.h"
#include "source_base/timer.h"
#include "typeinfo"
namespace ModulePW
{
template <typename T>
void PW_Basis::gatherp_scatters(std::complex<T>* in, std::complex<T>* out) const
{
    //ModuleBase::timer::tick(this->classname, "gatherp_scatters");
    
    if(this->poolnproc == 1) //In this case nst=nstot, nz = nplane, 
    {
#ifdef _OPENMP
#pragma omp parallel for
#endif
        for(int is = 0 ; is < this->nst ; ++is)
        {
            int ixy = this->istot2ixy[is];
            //int ixy = (ixy / fftny)*ny + ixy % fftny;
            std::complex<T> *outp = &out[is*nz];
            std::complex<T> *inp = &in[ixy*nz];
            for(int iz = 0 ; iz < this->nz ; ++iz)
            {
                outp[iz] = inp[iz];
            }
        }
        //ModuleBase::timer::tick(this->classname, "gatherp_scatters");
        return;
    }
#ifdef __MPI
    //change (nplane fftnxy) to (nplane,nstot)
    // Hence, we can send them at one time.
#ifdef _OPENMP
#pragma omp parallel for
#endif
    for (int istot = 0;istot < nstot; ++istot)
    {
        int ixy = this->istot2ixy[istot];
        //int ixy = (ixy / fftny)*ny + ixy % fftny;
        std::complex<T> *outp = &out[istot*nplane];
        std::complex<T> *inp = &in[ixy*nplane];
        for (int iz = 0; iz < nplane; ++iz)
        {
            outp[iz] = inp[iz];
        }
    }
 
    //exchange data
    //(nplane,nstot) to (numz[ip],ns, poolnproc)
    if(typeid(T) == typeid(double))
    {
        MPI_Alltoallv(out, numr, startr, MPI_DOUBLE_COMPLEX, in, numg, startg, MPI_DOUBLE_COMPLEX, this->pool_world);
    }
    else if(typeid(T) == typeid(float))
    {
        MPI_Alltoallv(out, numr, startr, MPI_COMPLEX, in, numg, startg, MPI_COMPLEX, this->pool_world);
    }
 
    // change (nz,ns) to (numz[ip],ns, poolnproc)
#ifdef _OPENMP
#pragma omp parallel for collapse(2)
#endif
    for (int ip = 0; ip < this->poolnproc ;++ip)
    {
        for (int is = 0; is < this->nst; ++is)
        {
            int nzip = this->numz[ip];
            std::complex<T> *outp0 = &out[startz[ip]];
            std::complex<T> *inp0 = &in[startg[ip]];
            std::complex<T> *outp = &outp0[is * nz];
            std::complex<T> *inp = &inp0[is * nzip ];
            for (int izip = 0; izip < nzip; ++izip)
            {
                outp[izip] = inp[izip];
            }
        }
    }
#endif
    //ModuleBase::timer::tick(this->classname, "gatherp_scatters");
    return;
}
 
template <typename T>
void PW_Basis::gathers_scatterp(std::complex<T>* in, std::complex<T>* out) const
{
    // ModuleBase::timer::tick(this->classname, "gathers_scatterp");
    if(this->poolnproc == 1) //In this case nrxx=fftnx*fftny*nz, nst = nstot, 
    {
#ifdef _OPENMP
#pragma omp parallel for schedule(static, 4096/sizeof(T))
#endif
        for(int i = 0; i < this->nrxx; ++i)
        {
            out[i] = std::complex<T>(0, 0);
        }
 
#ifdef _OPENMP
#pragma omp parallel for
#endif
        for(int is = 0 ; is < this->nst ; ++is)
        {
            int ixy = istot2ixy[is];
            //int ixy = (ixy / fftny)*ny + ixy % fftny;
            std::complex<T> *outp = &out[ixy*nz];
            std::complex<T> *inp = &in[is*nz];
            for(int iz = 0 ; iz < this->nz ; ++iz)
            {
                outp[iz] = inp[iz];
            }
        }
        // ModuleBase::timer::tick(this->classname, "gathers_scatterp");
        return;
    }
#ifdef __MPI
    // change (nz,ns) to (numz[ip],ns, poolnproc)
    // Hence, we can send them at one time. 
#ifdef _OPENMP
#pragma omp parallel for collapse(2)
#endif
    for (int ip = 0; ip < this->poolnproc ;++ip)
    {
        for (int is = 0; is < this->nst; ++is)
        {
            int nzip = this->numz[ip];
            std::complex<T> *outp0 = &out[startg[ip]];
            std::complex<T> *inp0 = &in[startz[ip]];
            std::complex<T> *outp = &outp0[is * nzip];
            std::complex<T> *inp = &inp0[is * nz ];
            for (int izip = 0; izip < nzip; ++izip)
            {
                outp[izip] = inp[izip];
            }
        }
    }
 
    //exchange data
    //(numz[ip],ns, poolnproc) to (nplane,nstot)
    if(typeid(T) == typeid(double))
    {
        MPI_Alltoallv(out, numg, startg, MPI_DOUBLE_COMPLEX, in, numr, startr, MPI_DOUBLE_COMPLEX, this->pool_world);
    }
    else if(typeid(T) == typeid(float))
    {
        MPI_Alltoallv(out, numg, startg, MPI_COMPLEX, in, numr, startr, MPI_COMPLEX, this->pool_world);
    }
 
#ifdef _OPENMP
#pragma omp parallel for schedule(static, 4096/sizeof(T))
#endif
    for(int i = 0; i < this->nrxx; ++i)
    {
        out[i] = std::complex<T>(0, 0);
    }
    //change (nplane,nstot) to (nplane fftnxy)
#ifdef _OPENMP
#pragma omp parallel for
#endif
    for (int istot = 0;istot < nstot; ++istot)
    {
        int ixy = this->istot2ixy[istot];
        //int ixy = (ixy / fftny)*ny + ixy % fftny;
        std::complex<T> *outp = &out[ixy * nplane];
        std::complex<T> *inp = &in[istot * nplane];
        for (int iz = 0; iz < nplane; ++iz)
        {
            outp[iz] = inp[iz];
        }
    }
#endif
    // ModuleBase::timer::tick(this->classname, "gathers_scatterp");
    return;
}
 
 
 
}