Skip to content

Commit

Permalink
Update Particle Container to Pure SoA
Browse files Browse the repository at this point in the history 
Transition particle containers to pure SoA layouts.
  • Loading branch information
@ax3l
ax3l committed Feb 8, 2024
1 parent 5127908 commit 6968ba3
Show file tree
Hide file tree
Showing 49 changed files with 974 additions and 778 deletions.
2 changes: 1 addition & 1 deletion cmake/dependencies/ABLASTR.cmake
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,7 +178,7 @@ set(ImpactX_openpmd_src ""
set(ImpactX_ablastr_repo "https://github.com/ECP-WarpX/WarpX.git"
CACHE STRING
"Repository URI to pull and build ABLASTR from if(ImpactX_ablastr_internal)")
set(ImpactX_ablastr_branch "24.02"
set(ImpactX_ablastr_branch "6e332e9479baa2769ff0ac22adb51c25c67627da"
CACHE STRING
"Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")

Expand Down
2 changes: 1 addition & 1 deletion examples/epac2004_benchmarks/input_fodo_rf_SC.in
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,4 +125,4 @@ geometry.prob_relative = 4.0
###############################################################################
# Diagnostics
###############################################################################
diag.slice_step_diagnostics = true
diag.slice_step_diagnostics = false
22 changes: 11 additions & 11 deletions examples/fodo/run_fodo_programmable.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,16 +77,16 @@ def my_drift(pge, pti, refpart):

else:
array = np.array
# access AoS data such as positions and cpu/id
aos = pti.aos()
aos_arr = array(aos, copy=False)

# access SoA data such as momentum
# access particle attributes
soa = pti.soa()
real_arrays = soa.GetRealData()
px = array(real_arrays[0], copy=False)
py = array(real_arrays[1], copy=False)
pt = array(real_arrays[2], copy=False)
real_arrays = soa.get_real_data()
x = array(real_arrays[0], copy=False)
y = array(real_arrays[1], copy=False)
t = array(real_arrays[2], copy=False)
px = array(real_arrays[3], copy=False)
py = array(real_arrays[4], copy=False)
pt = array(real_arrays[5], copy=False)

# length of the current slice
slice_ds = pge.ds / pge.nslice
Expand All @@ -96,9 +96,9 @@ def my_drift(pge, pti, refpart):
betgam2 = pt_ref**2 - 1.0

# advance position and momentum (drift)
aos_arr[:]["x"] += slice_ds * px[:]
aos_arr[:]["y"] += slice_ds * py[:]
aos_arr[:]["z"] += (slice_ds / betgam2) * pt[:]
x[:] += slice_ds * px[:]
y[:] += slice_ds * py[:]
t[:] += (slice_ds / betgam2) * pt[:]


def my_ref_drift(pge, refpart):
Expand Down
72 changes: 50 additions & 22 deletions examples/pytorch_surrogate_model/run_ml_surrogate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,14 @@
from urllib import request

import numpy as np

try:
import cupy as cp

cupy_available = True
except ImportError:
cupy_available = False

from surrogate_model_definitions import surrogate_model

try:
Expand All @@ -20,9 +28,10 @@
sys.exit(0)

from impactx import (
Config,
CoordSystem,
ImpactX,
ImpactXParIter,
TransformationDirection,
coordinate_transformation,
distribution,
elements,
Expand Down Expand Up @@ -78,7 +87,22 @@ def __init__(self, stage_i, surrogate_model, surrogate_length, stage_start):
self.ds = surrogate_length

def surrogate_push(self, pc, step):
array = np.array
# CPU/GPU logic
if Config.have_gpu:
if cupy_available:
array = cp.array
stack = cp.stack
device = torch.device("cuda")
else:
print("Warning: GPU found but cupy not available! Try managed...")
array = np.array
stack = np.stack
device = torch.device("cpu")
if Config.gpu_backend == "SYCL":
print("Warning: SYCL GPU backend not yet implemented for Python")

else:
array = np.array

ref_part = pc.ref_particle()
ref_z_i = ref_part.z
Expand All @@ -99,26 +123,30 @@ def surrogate_push(self, pc, step):
ref_part_final = torch.tensor([0, 0, ref_z_f, 0, 0, ref_uz_f])

# transform
coordinate_transformation(pc, TransformationDirection.to_fixed_t)
coordinate_transformation(pc, direction=CoordSystem.t)

for lvl in range(pc.finest_level + 1):
for pti in ImpactXParIter(pc, level=lvl):
aos = pti.aos()
aos_arr = array(aos, copy=False)

soa = pti.soa()
real_arrays = soa.GetRealData()
px = array(real_arrays[0], copy=False)
py = array(real_arrays[1], copy=False)
pt = array(real_arrays[2], copy=False)
data_arr = (
torch.tensor(
np.vstack(
[aos_arr["x"], aos_arr["y"], aos_arr["z"], real_arrays[:3]]
)
)
.float()
.T
real_arrays = soa.get_real_data()
x = array(real_arrays[0], copy=False)
y = array(real_arrays[1], copy=False)
t = array(real_arrays[2], copy=False)
px = array(real_arrays[3], copy=False)
py = array(real_arrays[4], copy=False)
pt = array(real_arrays[5], copy=False)
data_arr = torch.as_tensor(
stack(

Check failure

Code scanning / CodeQL

Potentially uninitialized local variable Error

Local variable 'stack' may be used before it is initialized.
[
x,
y,
t,
px,
py,
py,
]
),
device=device,

Check failure

Code scanning / CodeQL

Potentially uninitialized local variable Error

Local variable 'device' may be used before it is initialized.
)

data_arr[:, 0] += ref_part.x
Expand Down Expand Up @@ -146,9 +174,9 @@ def surrogate_push(self, pc, step):
data_arr_post_model[:, 3 + ii] -= ref_part_final[3 + ii]
data_arr_post_model[:, 3 + ii] /= ref_beta_gamma_final

aos_arr["x"] = data_arr_post_model[:, 0]
aos_arr["y"] = data_arr_post_model[:, 1]
aos_arr["z"] = data_arr_post_model[:, 2]
x[:] = data_arr_post_model[:, 0]
y[:] = data_arr_post_model[:, 1]
t[:] = data_arr_post_model[:, 2]
px[:] = data_arr_post_model[:, 3]
py[:] = data_arr_post_model[:, 4]
pt[:] = data_arr_post_model[:, 5]
Expand All @@ -173,7 +201,7 @@ def surrogate_push(self, pc, step):
# ref_part.s += pge1.ds
# ref_part.t += pge1.ds / ref_beta

coordinate_transformation(pc, TransformationDirection.to_fixed_s)
coordinate_transformation(pc, direction=CoordSystem.s)
## Done!


Expand Down
18 changes: 10 additions & 8 deletions src/particles/CollectLost.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
#include <AMReX_GpuLaunch.H>
#include <AMReX_GpuQualifiers.H>
#include <AMReX_Math.H>
#include <AMReX_Particle.H>
#include <AMReX_ParticleTransformation.H>
#include <AMReX_RandomEngine.H>

Expand All @@ -27,9 +28,9 @@ namespace impactx
using DstData = ImpactXParticleContainer::ParticleTileType::ParticleTileDataType;

AMREX_GPU_HOST_DEVICE
void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept {
dst.m_aos[dst_ip] = src.m_aos[src_ip];

void operator() (DstData const &dst, SrcData const &src, int src_ip, int dst_ip) const noexcept
{
dst.m_idcpu[dst_ip] = src.m_idcpu[src_ip];
for (int j = 0; j < SrcData::NAR; ++j)
dst.m_rdata[j][dst_ip] = src.m_rdata[j][src_ip];
for (int j = 0; j < src.m_num_runtime_real; ++j)
Expand All @@ -42,7 +43,7 @@ namespace impactx
// dst.m_runtime_idata[j][dst_ip] = src.m_runtime_idata[j][src_ip];

// flip id to positive in destination
dst.id(dst_ip) = amrex::Math::abs(dst.id(dst_ip));
amrex::ParticleIDWrapper{dst.m_idcpu[dst_ip]}.make_valid();

// remember the current s of the ref particle when lost
dst.m_runtime_rdata[s_index][dst_ip] = s_lost;
Expand Down Expand Up @@ -85,7 +86,7 @@ namespace impactx
auto const predicate = [] AMREX_GPU_HOST_DEVICE (const SrcData& src, int ip)
/* NVCC 11.3.109 chokes in C++17 on this: noexcept */
{
return src.id(ip) < 0;
return !amrex::ConstParticleIDWrapper{src.m_idcpu[ip]}.is_valid();
};

auto& ptile_dest = dest.DefineAndReturnParticleTile(
Expand Down Expand Up @@ -130,9 +131,11 @@ namespace impactx
{
int n_removed = 0;
auto ptile_src_data = ptile_source.getParticleTileData();
auto const ptile_soa = ptile_source.GetStructOfArrays();
auto const ptile_idcpu = ptile_soa.GetIdCPUData().dataPtr();
for (int ip = 0; ip < np; ++ip)
{
if (ptile_source.id(ip) < 0)
if (!amrex::ConstParticleIDWrapper{ptile_idcpu[ip]}.is_valid())
n_removed++;
else
{
Expand All @@ -141,8 +144,7 @@ namespace impactx
// move down
int const new_index = ip - n_removed;

ptile_src_data.m_aos[new_index] = ptile_src_data.m_aos[ip];

ptile_src_data.m_idcpu[new_index] = ptile_src_data.m_idcpu[ip];
for (int j = 0; j < SrcData::NAR; ++j)
ptile_src_data.m_rdata[j][new_index] = ptile_src_data.m_rdata[j][ip];
for (int j = 0; j < ptile_src_data.m_num_runtime_real; ++j)
Expand Down
Loading

0 comments on commit 6968ba3

Please sign in to comment.