Skip to content

Commit

Permalink
GPU Support
Browse files Browse the repository at this point in the history
  • Loading branch information
@ax3l
ax3l committed Aug 17, 2024
1 parent 5d5b75c commit 610689a
Show file tree
Hide file tree
Showing 10 changed files with 209 additions and 230 deletions.
2 changes: 1 addition & 1 deletion src/ImpactX.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,7 @@ namespace impactx {
}

// Wakefield calculation: call wakefield function to apply wake effects
particles::wakefields::HandleWakefield(amr_data->m_particle_container.get(), element_variant, csr, csr_bins, slice_ds);
particles::wakefields::HandleWakefield(*amr_data->m_particle_container, element_variant, csr, csr_bins, slice_ds);

// Space-charge calculation: turn off if there is only 1 particle
if (space_charge &&
Expand Down
5 changes: 1 addition & 4 deletions src/particles/wakefields/CSRBendElement.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,10 +47,7 @@ namespace impactx::particles::wakefields
R = std::abs(element.m_rc);
element_has_csr = true;
}
else
{
element_has_csr = false;
}

return {element_has_csr, R};
}

Expand Down
20 changes: 7 additions & 13 deletions src/particles/wakefields/ChargeBinning.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,15 @@ namespace impactx::particles::wakefields
/** Function to calculate charge density profile
*
* @param[in] myspc the particle species to deposit along s
* @param[out] dptr_data the output array (1D)
* @param[in] num_bins number of bins in dptr_data
* @param[out] charge_distribution the output array (1D)
* @param[in] bin_min lower end of the beam in s
* @param[in] bin_size size of the beam in s divided by num_bins
* @param[in] is_unity_particle_weight ignore the particle weighting per macro particle,
* otherwise treat each particle as one physical particle
*/
void DepositCharge1D (
impactx::ImpactXParticleContainer& myspc,
amrex::Real* dptr_data,
int num_bins,
amrex::Gpu::DeviceVector<amrex::Real> & charge_distribution,
amrex::Real bin_min,
amrex::Real bin_size,
bool is_unity_particle_weight = false
Expand All @@ -36,15 +34,13 @@ namespace impactx::particles::wakefields
/** Function to calculate the slope of the number (or charge) density
*
* @param[in] charge_distribution deposited charge in 1D along s
* @param[out] slopes derivative of charge_distribution along s
* @param[in] num_bins number of bins in charge_distribution and in len(slopes)-1
* @param[out] slopes derivative of charge_distribution along s with len(charge_distribution)-1
* @param[in] bin_size size of the beam in s divided by num_bins
* @param[in] GetNumberDensity number density if true, otherwise charge density
*/
void DerivativeCharge1D (
amrex::Real const * charge_distribution,
amrex::Real* slopes,
int num_bins,
amrex::Gpu::DeviceVector<amrex::Real> const & charge_distribution,
amrex::Gpu::DeviceVector<amrex::Real> & slopes,
amrex::Real bin_size,
bool GetNumberDensity = true
);
Expand All @@ -54,17 +50,15 @@ namespace impactx::particles::wakefields
* @param[in] myspc the particle species to deposit along s
* @param[out] mean_x the output array for mean x positions (1D)
* @param[out] mean_y the output array for mean y positions (1D)
* @param[in] num_bins number of bins in the mean_x and mean_y arrays
* @param[in] bin_min lower end of the beam in s
* @param[in] bin_size size of the beam in s divided by num_bins
* @param[in] is_unity_particle_weight ignore the particle weighting per macro particle,
* otherwise treat each particle as one physical particle
*/
void MeanTransversePosition (
impactx::ImpactXParticleContainer& myspc,
amrex::Real* mean_x,
amrex::Real* mean_y,
int num_bins,
amrex::Gpu::DeviceVector<amrex::Real> & mean_x,
amrex::Gpu::DeviceVector<amrex::Real> & mean_y,
amrex::Real bin_min,
amrex::Real bin_size,
bool is_unity_particle_weight = false
Expand Down
100 changes: 50 additions & 50 deletions src/particles/wakefields/ChargeBinning.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,50 +11,48 @@
#include "particles/ImpactXParticleContainer.H"

#include <cmath>
#include <vector>


namespace impactx::particles::wakefields
{
void DepositCharge1D (
impactx::ImpactXParticleContainer& myspc,
amrex::Real* dptr_data,
int num_bins,
amrex::Gpu::DeviceVector<amrex::Real> & charge_distribution,
amrex::Real bin_min,
amrex::Real bin_size,
bool is_unity_particle_weight
)
{ // Access and change data directly using '&' to pass by reference

// Determine the number of grid levels in the simulation
int const nlevs = std::max(0, myspc.finestLevel() + 1);
{
int const num_bins = charge_distribution.size();
amrex::Real * const dptr_data = charge_distribution.data();

// Loop over each grid level
for (int lev = 0; lev < nlevs; ++lev)
int const nlevs = myspc.finestLevel();
for (int lev = 0; lev <= nlevs; ++lev)
{
// OpenMP parallelization if enabled
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
#endif
{
//Loop over particles at the current grid level
// Loop over particles at the current grid level
for (impactx::ParIterSoA pti(myspc, lev); pti.isValid(); ++pti)
{
auto& soa = pti.GetStructOfArrays(); // Access data directly from StructOfArrays (soa)

//Number of particles
// Number of particles
long const np = pti.numParticles();

//Access particle weights and momenta
// Access particle weights and momenta
amrex::ParticleReal* const AMREX_RESTRICT d_w = soa.GetRealData(impactx::RealSoA::w).dataPtr();

//Access particle positions
// Access particle positions
amrex::ParticleReal* const AMREX_RESTRICT pos_z = soa.GetRealData(impactx::RealSoA::z).dataPtr();

//Parallel loop over particles
// Parallel loop over particles
amrex::ParallelFor(np, [=] AMREX_GPU_DEVICE(int i)
{
//Access particle z-position directly
// Access particle z-position directly
amrex::ParticleReal const z = pos_z[i]; // (Macro)Particle longitudinal position at i
auto const w = amrex::Real(d_w[i]); // (Macro)Particle weight at i

Expand Down Expand Up @@ -86,55 +84,61 @@ namespace impactx::particles::wakefields
}

void DerivativeCharge1D (
amrex::Real const* charge_distribution,
amrex::Real* slopes,
int num_bins,
amrex::Gpu::DeviceVector<amrex::Real> const & charge_distribution,
amrex::Gpu::DeviceVector<amrex::Real> & slopes,
amrex::Real bin_size,
bool GetNumberDensity
)
{
int const num_bins = charge_distribution.size();
amrex::Real const * const dptr_charge_distribution = charge_distribution.data();
amrex::Real * const dptr_slopes = slopes.data();

for (int i = 0; i < num_bins - 1; ++i)
amrex::ParallelFor(num_bins - 1, [=] AMREX_GPU_DEVICE(int i)
{
//Compute the charge density derivative
amrex::Real const charge_derivative = (charge_distribution[i + 1] - charge_distribution[i]) / bin_size;
// Compute the charge density derivative
amrex::Real const charge_derivative = (dptr_charge_distribution[i + 1] - dptr_charge_distribution[i]) / bin_size;

//If GetNumberDensity = True, convert charge density derivative to number density derivative for CSR convolution
// If GetNumberDensity = True, convert charge density derivative to number density derivative for CSR convolution
if (GetNumberDensity)
{
slopes[i] = charge_derivative / ablastr::constant::SI::q_e;
dptr_slopes[i] = charge_derivative / ablastr::constant::SI::q_e;
}
else
{
slopes[i] = charge_derivative;
dptr_slopes[i] = charge_derivative;
}
}
});
}

void MeanTransversePosition (
impactx::ImpactXParticleContainer& myspc,
amrex::Real* mean_x,
amrex::Real* mean_y,
int num_bins,
amrex::Gpu::DeviceVector<amrex::Real> & mean_x,
amrex::Gpu::DeviceVector<amrex::Real> & mean_y,
amrex::Real bin_min,
amrex::Real bin_size,
bool is_unity_particle_weight
)
{
using namespace amrex::literals;

int const nlevs = std::max(0, myspc.finestLevel() + 1);


int const num_bins = mean_x.size();
amrex::Real* dptr_mean_x = mean_x.data();
amrex::Real* dptr_mean_y = mean_y.data();

// Declare arrays for sums of positions and weights
std::vector<amrex::Real> sum_x(num_bins, 0.0);
std::vector<amrex::Real> sum_y(num_bins, 0.0);
std::vector<amrex::Real> sum_w(num_bins, 0.0);
auto sum_x = amrex::Gpu::DeviceVector<amrex::Real>(num_bins, 0.0_rt);
auto sum_y = amrex::Gpu::DeviceVector<amrex::Real>(num_bins, 0.0_rt);
auto sum_w = amrex::Gpu::DeviceVector<amrex::Real>(num_bins, 0.0_rt);

amrex::Real* sum_w_ptr = sum_w.data();
amrex::Real* sum_x_ptr = sum_x.data();
amrex::Real* sum_y_ptr = sum_y.data();
amrex::Real* const sum_w_ptr = sum_w.data();
amrex::Real* const sum_x_ptr = sum_x.data();
amrex::Real* const sum_y_ptr = sum_y.data();

for (int lev = 0; lev < nlevs; ++lev)
int const nlevs = myspc.finestLevel();
for (int lev = 0; lev <= nlevs; ++lev)
{
#ifdef AMREX_USE_OMP
#pragma omp parallel if (amrex::Gpu::notInLaunchRegion())
Expand All @@ -160,32 +164,28 @@ namespace impactx::particles::wakefields
int const bin = int(amrex::Math::floor((z - bin_min) / bin_size));
if (bin < 0 || bin >= num_bins) { return; }

amrex::Real const weight = is_unity_particle_weight ? 1.0_rt : w; // Check is macroparticle made up of 1 or more particles

sum_w_ptr[bin] += weight; // Deposit the number of particles composing macroparticle
sum_x_ptr[bin] += x * weight; // Deposit x position multiplied by number of particles at this position
sum_y_ptr[bin] += y * weight;
amrex::Real const weight = is_unity_particle_weight ? 1.0_rt : w; // Check is macroparticle made up of 1 or more particles

//HostDevice::Atomic::Add(&sum_w[bin], weight); // Deposit the number of particles composing macroparticle
//HostDevice::Atomic::Add(&sum_x[bin], x * weight); // Deposit x position multiplied by number of particles at this position
//HostDevice::Atomic::Add(&sum_y[bin], y * weight);
amrex::HostDevice::Atomic::Add(&sum_w_ptr[bin], weight); // Deposit the number of particles composing macroparticle
amrex::HostDevice::Atomic::Add(&sum_x_ptr[bin], x * weight); // Deposit x position multiplied by number of particles at this position
amrex::HostDevice::Atomic::Add(&sum_y_ptr[bin], y * weight);
});
}
}
}

for (int i = 0; i < num_bins; ++i)
amrex::ParallelFor(num_bins, [=] AMREX_GPU_DEVICE(int i)
{
if (sum_w[i] > 0) // Ensure number of particles in a bin is >= 1 before taking the mean
if (sum_w_ptr[i] > 0) // Ensure number of particles in a bin is >= 1 before taking the mean
{
mean_x[i] = sum_x_ptr[i] / sum_w_ptr[i];
mean_y[i] = sum_y_ptr[i] / sum_w_ptr[i];
dptr_mean_x[i] = sum_x_ptr[i] / sum_w_ptr[i];
dptr_mean_y[i] = sum_y_ptr[i] / sum_w_ptr[i];
}
else
{
mean_x[i] = 0.0;
mean_y[i] = 0.0;
dptr_mean_x[i] = 0.0;
dptr_mean_y[i] = 0.0;
}
}
});
}
}
Loading

0 comments on commit 610689a

Please sign in to comment.