diff options
Diffstat (limited to 'Source/FieldSolver/SpectralSolver')
5 files changed, 200 insertions, 59 deletions
diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H index 0487e5226..34743525e 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.H @@ -1,6 +1,8 @@ #ifndef WARPX_PSATD_ALGORITHM_H_ #define WARPX_PSATD_ALGORITHM_H_ +#include <SpectralKSpace.H> +#include <SpectralFieldData.H> #include <SpectralBaseAlgorithm.H> /* \brief Class that updates the field in spectral space @@ -8,16 +10,26 @@ */ class PsatdAlgorithm : public SpectralBaseAlgorithm { + using SpectralCoefficients = amrex::FabArray< amrex::BaseFab <amrex::Real> >; + public: PsatdAlgorithm(const SpectralKSpace& spectral_kspace, const amrex::DistributionMapping& dm, const int norder_x, const int norder_y, - const int norder_z, const bool nodal, - const amrex::Real dt); - // Redefine update equation from base class - virtual void pushSpectralFields(SpectralFieldData& f) const override final; + const int norder_z, const bool nodal, const amrex::Real dt); + PsatdAlgorithm() = default; // Default constructor + PsatdAlgorithm& operator=(PsatdAlgorithm&& algorithm) = default; + void pushSpectralFields(SpectralFieldData& f) const; + void InitializeCoefficience(const SpectralKSpace& spectral_kspace, + const amrex::DistributionMapping& dm, + const amrex::Real dt); private: + // Modified finite-order vectors + KVectorComponent modified_kx_vec, modified_kz_vec; +#if (AMREX_SPACEDIM==3) + KVectorComponent modified_ky_vec; +#endif SpectralCoefficients C_coef, S_ck_coef, X1_coef, X2_coef, X3_coef; }; diff --git a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp index 37892d35a..8dd2a830f 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralAlgorithms/PsatdAlgorithm.cpp @@ -22,58 +22,60 @@ PsatdAlgorithm::PsatdAlgorithm(const SpectralKSpace& spectral_kspace, X2_coef = SpectralCoefficients(ba, dm, 1, 0); X3_coef = SpectralCoefficients(ba, dm, 1, 0); - // Fill them with the right values: - // Loop over boxes and allocate the corresponding coefficients - // for each box owned by the local MPI proc - for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi){ - - const Box& bx = ba[mfi]; - - // Extract pointers for the k vectors - const Real* modified_kx = modified_kx_vec[mfi].dataPtr(); -#if (AMREX_SPACEDIM==3) - const Real* modified_ky = modified_ky_vec[mfi].dataPtr(); -#endif - const Real* modified_kz = modified_kz_vec[mfi].dataPtr(); - // Extract arrays for the coefficients - Array4<Real> C = C_coef[mfi].array(); - Array4<Real> S_ck = S_ck_coef[mfi].array(); - Array4<Real> X1 = X1_coef[mfi].array(); - Array4<Real> X2 = X2_coef[mfi].array(); - Array4<Real> X3 = X3_coef[mfi].array(); - - // Loop over indices within one box - ParallelFor(bx, - [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept - { - // Calculate norm of vector - const Real k_norm = std::sqrt( - std::pow(modified_kx[i], 2) + -#if (AMREX_SPACEDIM==3) - std::pow(modified_ky[j], 2) + - std::pow(modified_kz[k], 2)); -#else - std::pow(modified_kz[j], 2)); -#endif - - // Calculate coefficients - constexpr Real c = PhysConst::c; - constexpr Real ep0 = PhysConst::ep0; - if (k_norm != 0){ - C(i,j,k) = std::cos(c*k_norm*dt); - S_ck(i,j,k) = std::sin(c*k_norm*dt)/(c*k_norm); - X1(i,j,k) = (1. - C(i,j,k))/(ep0 * c*c * k_norm*k_norm); - X2(i,j,k) = (1. - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm); - X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm); - } else { // Handle k_norm = 0, by using the analytical limit - C(i,j,k) = 1.; - S_ck(i,j,k) = dt; - X1(i,j,k) = 0.5 * dt*dt / ep0; - X2(i,j,k) = c*c * dt*dt / (6.*ep0); - X3(i,j,k) = - c*c * dt*dt / (3.*ep0); - } - }); - } + InitializeCoefficience(spectral_kspace, dm, dt); +// // Fill them with the right values: +// // Loop over boxes and allocate the corresponding coefficients +// // for each box owned by the local MPI proc +// for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi){ +// +// //const Box& bx = ba[mfi]; +// const Box bx = ba[mfi]; +// +// // Extract pointers for the k vectors +// const Real* modified_kx = modified_kx_vec[mfi].dataPtr(); +//#if (AMREX_SPACEDIM==3) +// const Real* modified_ky = modified_ky_vec[mfi].dataPtr(); +//#endif +// const Real* modified_kz = modified_kz_vec[mfi].dataPtr(); +// // Extract arrays for the coefficients +// Array4<Real> C = C_coef[mfi].array(); +// Array4<Real> S_ck = S_ck_coef[mfi].array(); +// Array4<Real> X1 = X1_coef[mfi].array(); +// Array4<Real> X2 = X2_coef[mfi].array(); +// Array4<Real> X3 = X3_coef[mfi].array(); +// +// // Loop over indices within one box +// ParallelFor(bx, +// [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept +// { +// // Calculate norm of vector +// const Real k_norm = std::sqrt( +// std::pow(modified_kx[i], 2) + +//#if (AMREX_SPACEDIM==3) +// std::pow(modified_ky[j], 2) + +// std::pow(modified_kz[k], 2)); +//#else +// std::pow(modified_kz[j], 2)); +//#endif +// +// // Calculate coefficients +// constexpr Real c = PhysConst::c; +// constexpr Real ep0 = PhysConst::ep0; +// if (k_norm != 0){ +// C(i,j,k) = std::cos(c*k_norm*dt); +// S_ck(i,j,k) = std::sin(c*k_norm*dt)/(c*k_norm); +// X1(i,j,k) = (1. - C(i,j,k))/(ep0 * c*c * k_norm*k_norm); +// X2(i,j,k) = (1. - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm); +// X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm); +// } else { // Handle k_norm = 0, by using the analytical limit +// C(i,j,k) = 1.; +// S_ck(i,j,k) = dt; +// X1(i,j,k) = 0.5 * dt*dt / ep0; +// X2(i,j,k) = c*c * dt*dt / (6.*ep0); +// X3(i,j,k) = - c*c * dt*dt / (3.*ep0); +// } +// }); +// } }; /* Advance the E and B field in spectral space (stored in `f`) @@ -130,7 +132,7 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{ #endif constexpr Real c2 = PhysConst::c*PhysConst::c; constexpr Real inv_ep0 = 1./PhysConst::ep0; - constexpr Complex I = Complex{0,1}; + const Complex I = Complex{0,1}; const Real C = C_arr(i,j,k); const Real S_ck = S_ck_arr(i,j,k); const Real X1 = X1_arr(i,j,k); @@ -160,3 +162,64 @@ PsatdAlgorithm::pushSpectralFields(SpectralFieldData& f) const{ }); } }; + +void PsatdAlgorithm::InitializeCoefficience(const SpectralKSpace& spectral_kspace, + const amrex::DistributionMapping& dm, + const amrex::Real dt) +{ + const BoxArray& ba = spectral_kspace.spectralspace_ba; + // Fill them with the right values: + // Loop over boxes and allocate the corresponding coefficients + // for each box owned by the local MPI proc + for (MFIter mfi(ba, dm); mfi.isValid(); ++mfi){ + + //const Box& bx = ba[mfi]; + const Box bx = ba[mfi]; + + // Extract pointers for the k vectors + const Real* modified_kx = modified_kx_vec[mfi].dataPtr(); +#if (AMREX_SPACEDIM==3) + const Real* modified_ky = modified_ky_vec[mfi].dataPtr(); +#endif + const Real* modified_kz = modified_kz_vec[mfi].dataPtr(); + // Extract arrays for the coefficients + Array4<Real> C = C_coef[mfi].array(); + Array4<Real> S_ck = S_ck_coef[mfi].array(); + Array4<Real> X1 = X1_coef[mfi].array(); + Array4<Real> X2 = X2_coef[mfi].array(); + Array4<Real> X3 = X3_coef[mfi].array(); + + // Loop over indices within one box + ParallelFor(bx, + [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept + { + // Calculate norm of vector + const Real k_norm = std::sqrt( + std::pow(modified_kx[i], 2) + +#if (AMREX_SPACEDIM==3) + std::pow(modified_ky[j], 2) + + std::pow(modified_kz[k], 2)); +#else + std::pow(modified_kz[j], 2)); +#endif + + + // Calculate coefficients + constexpr Real c = PhysConst::c; + constexpr Real ep0 = PhysConst::ep0; + if (k_norm != 0){ + C(i,j,k) = std::cos(c*k_norm*dt); + S_ck(i,j,k) = std::sin(c*k_norm*dt)/(c*k_norm); + X1(i,j,k) = (1. - C(i,j,k))/(ep0 * c*c * k_norm*k_norm); + X2(i,j,k) = (1. - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm); + X3(i,j,k) = (C(i,j,k) - S_ck(i,j,k)/dt)/(ep0 * k_norm*k_norm); + } else { // Handle k_norm = 0, by using the analytical limit + C(i,j,k) = 1.; + S_ck(i,j,k) = dt; + X1(i,j,k) = 0.5 * dt*dt / ep0; + X2(i,j,k) = c*c * dt*dt / (6.*ep0); + X3(i,j,k) = - c*c * dt*dt / (3.*ep0); + } + }); + } +} diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H index 8e58aa1d8..1d64817ef 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.H +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.H @@ -26,6 +26,7 @@ class SpectralFieldData // the local MPI rank) #ifdef AMREX_USE_GPU // Add cuFFT-specific code + using FFTplans = amrex::LayoutData<cufftHandle>; #else using FFTplans = amrex::LayoutData<fftw_plan>; #endif diff --git a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp index 02fa2015f..5998bdd2b 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralFieldData.cpp @@ -54,6 +54,25 @@ SpectralFieldData::SpectralFieldData( const BoxArray& realspace_ba, IntVect fft_size = realspace_ba[mfi].length(); #ifdef AMREX_USE_GPU // Add cuFFT-specific code + // Creating 3D plan for real to complex -- double precision + cudaDeviceSynchronize(); + cufftResult result; + result = cufftPlan3d( &forward_plan[mfi], fft_size[2], + fft_size[1],fft_size[0], CUFFT_D2Z); + if ( result != CUFFT_SUCCESS ) { + amrex::Print() << " cufftplan3d forward failed! \n"; + } + // Add 2D cuFFT-spacific code for D2Z + // Note that D2Z is inherently forward plan + + result = cufftPlan3d( &backward_plan[mfi], fft_size[2], + fft_size[1], fft_size[0], CUFFT_Z2D); + // Add 2D cuFFT-specific code for Z2D + if ( result != CUFFT_SUCCESS ) { + amrex::Print() << " cufftplan3d backward failed! \n"; + } + cudaDeviceSynchronize(); + #else // Create FFTW plans forward_plan[mfi] = @@ -87,6 +106,8 @@ SpectralFieldData::~SpectralFieldData() for ( MFIter mfi(tmpRealField); mfi.isValid(); ++mfi ){ #ifdef AMREX_USE_GPU // Add cuFFT-specific code + cufftDestroy( forward_plan[mfi] ); + cufftDestroy( backward_plan[mfi] ); #else // Destroy FFTW plans fftw_destroy_plan( forward_plan[mfi] ); @@ -129,14 +150,30 @@ SpectralFieldData::ForwardTransform( const MultiFab& mf, Array4<Real> tmp_arr = tmpRealField[mfi].array(); ParallelFor( realspace_bx, [=] AMREX_GPU_DEVICE(int i, int j, int k) noexcept { - tmp_arr(i,j,k) = mf_arr(i,j,k,i_comp); + tmp_arr(i,j,k) = mf_arr(i,j,k,i_comp); }); +//#ifdef AMREX_USE_GPU +// cudaDeviceSynchronize(); +//#endif } // Perform Fourier transform from `tmpRealField` to `tmpSpectralField` #ifdef AMREX_USE_GPU // Add cuFFT-specific code ; make sure that this is done on the same // GPU stream as the above copy + cudaDeviceSynchronize(); + cufftResult result; + cudaStream_t stream = amrex::Cuda::Device::cudaStream(); + amrex::Print() << " stream is " << stream << "\n"; + cufftSetStream ( forward_plan[mfi], stream); + result = cufftExecD2Z( forward_plan[mfi], + tmpRealField[mfi].dataPtr(), + reinterpret_cast<cuDoubleComplex*>( + tmpSpectralField[mfi].dataPtr()) ); + if ( result != CUFFT_SUCCESS ) { + amrex::Print() << " cufftplan3d execute failed ! \n"; + } + cudaDeviceSynchronize(); #else fftw_execute( forward_plan[mfi] ); #endif @@ -227,6 +264,22 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, #ifdef AMREX_USE_GPU // Add cuFFT-specific code ; make sure that this is done on the same // GPU stream as the above copy + cudaDeviceSynchronize(); + cufftResult result; + cudaStream_t stream = amrex::Cuda::Device::cudaStream(); + amrex::Print() << " stream is " << stream << "\n"; + cufftSetStream ( backward_plan[mfi], stream); + result = cufftExecZ2D( backward_plan[mfi], + reinterpret_cast<cuDoubleComplex*>( + tmpSpectralField[mfi].dataPtr()), + tmpRealField[mfi].dataPtr() ); + if ( result != CUFFT_SUCCESS ) { + amrex::Print() << " cufftplan3d execute inverse failed ! \n"; + } + if ( result == CUFFT_SUCCESS ) { + amrex::Print() << " created cufft inverse transform\n"; + } + cudaDeviceSynchronize(); #else fftw_execute( backward_plan[mfi] ); #endif @@ -245,6 +298,14 @@ SpectralFieldData::BackwardTransform( MultiFab& mf, // Copy and normalize field mf_arr(i,j,k,i_comp) = inv_N*tmp_arr(i,j,k); }); +//#ifdef AMREX_USE_GPU +// cudaDeviceSynchronize(); +//#endif + +#ifdef AMREX_USE_GPU + cudaDeviceSynchronize(); +#endif + amrex::Print() << " divided by 1/N \n"; } } } diff --git a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp index 2fe78cedd..6a88a52a3 100644 --- a/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp +++ b/Source/FieldSolver/SpectralSolver/SpectralKSpace.cpp @@ -142,9 +142,13 @@ SpectralKSpace::getSpectralShiftFactor( const DistributionMapping& dm, case ShiftType::TransformFromCellCentered: sign = -1.; break; case ShiftType::TransformToCellCentered: sign = 1.; } - constexpr Complex I{0,1}; + const Complex I{0,1}; for (int i=0; i<k.size(); i++ ){ +#ifdef AMREX_USE_GPU + shift[i] = thrust::exp( I*sign*k[i]*0.5*dx[i_dim] ); +#else shift[i] = std::exp( I*sign*k[i]*0.5*dx[i_dim] ); +#endif } } return shift_factor; |