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 d9eca59
Show file tree
Hide file tree
Showing 10 changed files with 201 additions and 222 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
88 changes: 44 additions & 44 deletions src/particles/wakefields/ChargeBinning.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,26 +11,24 @@
#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
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;
}
}
});
}
}
81 changes: 43 additions & 38 deletions src/particles/wakefields/ExecuteWakefield.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@

namespace impactx::particles::wakefields
{

/** Function to handle CSR bending process including charge deposition,
* mean transverse position calculation, wakefield generation, and convolution.
*
Expand All @@ -31,17 +30,19 @@ namespace impactx::particles::wakefields
* @param[in] csr boolean indicating if CSR effects are included
* @param[in] csr_bins number of bins (resolution) for CSR calculations
* @param[in] slice_ds slice spacing along s
* @param[in] print_wakefield for debugging: print the wakefield to convoluted_wakefield.txt
*/
template <typename T_Element>
void HandleWakefield (
impactx::ImpactXParticleContainer* particle_container,
impactx::ImpactXParticleContainer& particle_container,
T_Element const& element_variant,
bool csr,
int csr_bins,
amrex::Real slice_ds
amrex::Real slice_ds,
bool print_wakefield = false
)
{
BL_PROFILE("impactx::particles::wakefields::HandleWakefield");
BL_PROFILE("impactx::particles::wakefields::HandleWakefield")

// Call the CSR bend function
auto const [element_has_csr, R] = impactx::particles::wakefields::CSRBendElement(element_variant);
Expand All @@ -51,62 +52,66 @@ namespace impactx::particles::wakefields
{
// Measure beam size, extract the min, max of particle positions
[[maybe_unused]] auto const [x_min, y_min, t_min, x_max, y_max, t_max] =

Check notice

Code scanning / CodeQL

Unused local variable Note

Variable x_min is not used.

Check notice

Code scanning / CodeQL

Unused local variable Note

Variable y_min is not used.

Check notice

Code scanning / CodeQL

Unused local variable Note

Variable x_max is not used.

Check notice

Code scanning / CodeQL

Unused local variable Note

Variable y_max is not used.
particle_container->MinAndMaxPositions();

using amrex::Real;
particle_container.MinAndMaxPositions();

// Set parameters for charge deposition
bool is_unity_particle_weight = false; // Only true if w = 1
bool GetNumberDensity = true;
bool const is_unity_particle_weight = false; // Only true if w = 1
bool const GetNumberDensity = true;

int padding_factor = 1; // Set amount of zero-padding
int num_bins = csr_bins; // Set resolution
Real bin_min = t_min;
Real bin_max = t_max;
Real bin_size = (bin_max - bin_min) / num_bins;
int const padding_factor = 1; // Set amount of zero-padding
int const num_bins = csr_bins; // Set resolution
amrex::Real const bin_min = t_min;
amrex::Real const bin_max = t_max;
amrex::Real const bin_size = (bin_max - bin_min) / num_bins;

// Allocate memory for the charge profile
Real* dptr_data = new Real[num_bins]();
Real* mean_x = new Real[num_bins]();
Real* mean_y = new Real[num_bins]();
amrex::Gpu::DeviceVector<amrex::Real> charge_distribution(num_bins);
amrex::Gpu::DeviceVector<amrex::Real> mean_x(num_bins);
amrex::Gpu::DeviceVector<amrex::Real> mean_y(num_bins);

// Call charge deposition function
impactx::particles::wakefields::DepositCharge1D(*particle_container, dptr_data, num_bins, bin_min, bin_size, is_unity_particle_weight);
impactx::particles::wakefields::DepositCharge1D(particle_container, charge_distribution, bin_min, bin_size, is_unity_particle_weight);

// FIXME: MPI reduce to rank zero, then the rest on IO processor, then BCast result

Check notice

Code scanning / CodeQL

FIXME comment Note

FIXME comment: MPI reduce to rank zero, then the rest on IO processor, then BCast result

// Call the mean transverse position function
impactx::particles::wakefields::MeanTransversePosition(*particle_container, mean_x, mean_y, num_bins, bin_min, bin_size, is_unity_particle_weight);
impactx::particles::wakefields::MeanTransversePosition(particle_container, mean_x, mean_y, bin_min, bin_size, is_unity_particle_weight);

// Call charge density derivative function
std::vector<amrex::Real> charge_distribution(dptr_data, dptr_data + num_bins);
std::vector<amrex::Real> slopes(num_bins - 1);
impactx::particles::wakefields::DerivativeCharge1D(charge_distribution.data(), slopes.data(), num_bins, bin_size, GetNumberDensity); // Use number derivatives for convolution with CSR
amrex::Gpu::DeviceVector<amrex::Real> slopes(num_bins - 1);
impactx::particles::wakefields::DerivativeCharge1D(charge_distribution, slopes, bin_size, GetNumberDensity); // Use number derivatives for convolution with CSR

// Call CSR wake function
std::vector<amrex::Real> wake_function(num_bins);
for (int i = 0; i < num_bins; ++i)
amrex::Gpu::DeviceVector<amrex::Real> wake_function(num_bins);
amrex::Real * const dptr_wake_function = wake_function.data();
auto const dR = R; // for NVCC capture
amrex::ParallelFor(num_bins, [=] AMREX_GPU_DEVICE(int i)
{
amrex::Real s = bin_min + i * bin_size;
wake_function[i] = impactx::particles::wakefields::w_l_csr(s, R);
}
amrex::Real const s = bin_min + i * bin_size;
dptr_wake_function[i] = impactx::particles::wakefields::w_l_csr(s, dR);
});

// Call convolution function
std::vector<amrex::Real> convoluted_wakefield(padding_factor * (2 * num_bins - 1));
impactx::particles::wakefields::convolve_fft(slopes.data(), wake_function.data(), slopes.size(), wake_function.size(), bin_size, convoluted_wakefield.data(), padding_factor);
amrex::Gpu::DeviceVector<amrex::Real> convoluted_wakefield(padding_factor * (2 * num_bins - 1));
impactx::particles::wakefields::convolve_fft(slopes, wake_function, bin_size, convoluted_wakefield, padding_factor);

// FIXME: MPI BCast convoluted_wakefield to every MPI rank

Check notice

Code scanning / CodeQL

FIXME comment Note

FIXME comment: MPI BCast convoluted_wakefield to every MPI rank

// Check convolution output
std::cout << "Convoluted wakefield: ";
std::ofstream outfile("convoluted_wakefield.txt");
for (int i = 0; i < int(convoluted_wakefield.size()); ++i)
if (print_wakefield && amrex::ParallelDescriptor::IOProcessor())
{
std::cout << convoluted_wakefield[i] << " ";
outfile << convoluted_wakefield[i] << std::endl;
std::cout << "Convoluted wakefield: ";
std::ofstream outfile("convoluted_wakefield.txt");
for (double i : convoluted_wakefield) {
std::cout << i << " ";
outfile << i << std::endl;
}
std::cout << std::endl;
outfile.close();
}
std::cout << std::endl;
outfile.close();
delete[] dptr_data;

// Call function to kick particles with wake
impactx::particles::wakefields::WakePush(*particle_container, convoluted_wakefield, slice_ds, bin_size, t_min, padding_factor);
impactx::particles::wakefields::WakePush(particle_container, convoluted_wakefield, slice_ds, bin_size, t_min, padding_factor);
}
}

Expand Down
Loading

0 comments on commit d9eca59

Please sign in to comment.