Skip to content

[WIP] SIMD: Vectorized Transcendentals in Element Pushes #1632

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ax3l wants to merge 1 commit into
base: development
Choose a base branch
from
Open

[WIP] SIMD: Vectorized Transcendentals in Element Pushes #1632

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions .github/workflows/dependencies/gcc.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,11 @@ sudo apt-get install -y \
wget

# vir-simd
wget https://github.com/mattkretz/vir-simd/archive/refs/tags/v0.4.4.tar.gz
tar -xvf v0.4.4.tar.gz
rm -rf v0.4.4.tar.gz
cmake -S vir-simd-0.4.4 -B vir-simd-build
# TODO: back to the release tarball once vir/simd_vecmath.h is in one. It is
# what makes the SIMD transcendentals call a vector math library instead
# of evaluating them one lane at a time.
git clone --depth 1 --branch topic-vecmath https://github.com/ax3l/vir-simd.git vir-simd-src
cmake -S vir-simd-src -B vir-simd-build
sudo cmake --build vir-simd-build --target install

python3 -m pip install -U pip
Expand Down
7 changes: 3 additions & 4 deletions .github/workflows/macos.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,9 @@ jobs:
set -e
- name: install vir-simd
run: |
wget https://github.com/mattkretz/vir-simd/archive/refs/tags/v0.4.4.tar.gz
tar -xvf v0.4.4.tar.gz
rm -rf v0.4.4.tar.gz
cmake -S vir-simd-0.4.4 -B vir-simd-build
# TODO: back to the release tarball once vir/simd_vecmath.h is in one
git clone --depth 1 --branch topic-vecmath https://github.com/ax3l/vir-simd.git vir-simd-src
cmake -S vir-simd-src -B vir-simd-build
sudo cmake --build vir-simd-build --target install
- name: install pip dependencies
run: |
Expand Down
7 changes: 5 additions & 2 deletions cmake/dependencies/ABLASTR.cmake
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,10 +183,13 @@ set(ImpactX_ablastr_branch "26.08"
"Repository branch for ImpactX_ablastr_repo if(ImpactX_ablastr_internal)")

# AMReX is transitively pulled through ABLASTR
set(ImpactX_amrex_repo "https://github.com/AMReX-Codes/amrex.git"
# TODO: back to AMReX-Codes/amrex once the SIMD transcendentals land there.
# https://github.com/AMReX-Codes/amrex/pull/5644 adds the amrex::Math
# overloads the vectorized elements call.
set(ImpactX_amrex_repo "https://github.com/ax3l/amrex.git"
CACHE STRING
"Repository URI to pull and build AMReX from if(ImpactX_amrex_internal)")
set(ImpactX_amrex_branch "26.08"
set(ImpactX_amrex_branch "topic-simd-vecmath"
CACHE STRING
"Repository branch for ImpactX_amrex_repo if(ImpactX_amrex_internal)")

Expand Down
24 changes: 12 additions & 12 deletions src/elements/ChrPlasmaLens.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -188,10 +188,10 @@ namespace impactx::elements
T_Real const t0 = t - term*m_slice_ds/delta1;

T_Real const w = omega*delta1;
T_Real const term1 = -(powi<2>(p2) - powi<2>(q2) * powi<2>(w)) * sin(2_prt*m_slice_ds*omega);
T_Real const term2 = -(powi<2>(p1) - powi<2>(q1) * powi<2>(w)) * sin(2_prt*m_slice_ds*omega);
T_Real const term3 = -2_prt * q2 * p2 * w * cos(2_prt*m_slice_ds*omega);
T_Real const term4 = -2_prt * q1 * p1 * w * cos(2_prt*m_slice_ds*omega);
T_Real const term1 = -(powi<2>(p2) - powi<2>(q2) * powi<2>(w)) * amrex::Math::sin(2_prt*m_slice_ds*omega);
T_Real const term2 = -(powi<2>(p1) - powi<2>(q1) * powi<2>(w)) * amrex::Math::sin(2_prt*m_slice_ds*omega);
T_Real const term3 = -2_prt * q2 * p2 * w * amrex::Math::cos(2_prt*m_slice_ds*omega);
T_Real const term4 = -2_prt * q1 * p1 * w * amrex::Math::cos(2_prt*m_slice_ds*omega);
T_Real const term5 = 2_prt * omega * (q1*p1*delta1 + q2*p2*delta1
-(powi<2>(p1) + powi<2>(p2))*m_slice_ds - (powi<2>(q1) + powi<2>(q2)) * powi<2>(w)*m_slice_ds);
tout = t0 + (-1_prt+m_beta*pt)
Expand All @@ -201,8 +201,8 @@ namespace impactx::elements
}
else if (m_g < 0_prt)
{
auto const sinh_ods = sinh(omega*m_slice_ds);
auto const cosh_ods = cosh(omega*m_slice_ds);
auto const sinh_ods = amrex::Math::sinh(omega*m_slice_ds);
auto const cosh_ods = amrex::Math::cosh(omega*m_slice_ds);

// advance transverse position and momentum (defocusing)
xout = cosh_ods * x + sinh_ods / (omega * delta1) * px;
Expand All @@ -216,10 +216,10 @@ namespace impactx::elements
T_Real const t0 = t - term*m_slice_ds/delta1;

T_Real const w = omega*delta1;
T_Real const term1 = -(powi<2>(p2) + powi<2>(q2) * powi<2>(w)) * sinh(2_prt*m_slice_ds*omega);
T_Real const term2 = -(powi<2>(p1) + powi<2>(q1) * powi<2>(w)) * sinh(2_prt*m_slice_ds*omega);
T_Real const term3 = -2_prt * q2 * p2 * w * cosh(2_prt*m_slice_ds*omega);
T_Real const term4 = -2_prt * q1 * p1 * w * cosh(2_prt*m_slice_ds*omega);
T_Real const term1 = -(powi<2>(p2) + powi<2>(q2) * powi<2>(w)) * amrex::Math::sinh(2_prt*m_slice_ds*omega);
T_Real const term2 = -(powi<2>(p1) + powi<2>(q1) * powi<2>(w)) * amrex::Math::sinh(2_prt*m_slice_ds*omega);
T_Real const term3 = -2_prt * q2 * p2 * w * amrex::Math::cosh(2_prt*m_slice_ds*omega);
T_Real const term4 = -2_prt * q1 * p1 * w * amrex::Math::cosh(2_prt*m_slice_ds*omega);
T_Real const term5 = 2_prt * omega * (q1*p1*delta1 + q2*p2*delta1
-(powi<2>(p1) + powi<2>(p2))*m_slice_ds - (powi<2>(q1) + powi<2>(q2)) * powi<2>(w)*m_slice_ds);
tout = t0 + (-1_prt+m_beta*pt)
Expand Down Expand Up @@ -344,8 +344,8 @@ namespace impactx::elements

// compute trigonometric quantities
auto const [sin_omega_ds, cos_omega_ds] = amrex::Math::sincos(omega*m_slice_ds);
T_Real const sinh_omega_ds = sinh(omega*m_slice_ds);
T_Real const cosh_omega_ds = cosh(omega*m_slice_ds);
T_Real const sinh_omega_ds = amrex::Math::sinh(omega*m_slice_ds);
T_Real const cosh_omega_ds = amrex::Math::cosh(omega*m_slice_ds);

if (m_g > 0.0_prt)
{
Expand Down
20 changes: 10 additions & 10 deletions src/elements/ChrQuad.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,19 +178,19 @@ namespace impactx::elements
bool const focusing = m_g > 0_prt;
// raw off-diagonal trig: oscillatory (sin) for the focusing plane,
// hyperbolic (sinh) for the defocusing plane (delta1 > 0 always).
T_Real const sx = focusing ? sin(omega * m_slice_ds) : sinh(omega * m_slice_ds);
T_Real const sy = focusing ? sinh(omega * m_slice_ds) : sin(omega * m_slice_ds);
T_Real const sx = focusing ? amrex::Math::sin(omega * m_slice_ds) : amrex::Math::sinh(omega * m_slice_ds);
T_Real const sy = focusing ? amrex::Math::sinh(omega * m_slice_ds) : amrex::Math::sin(omega * m_slice_ds);
// sign on the px,py kick: -1 (x focusing) / +1 (x defocusing)
T_Real const ax = focusing ? -1_prt : 1_prt;

// per-particle 2x2 transfer-map blocks. omega is chromatic (it depends
// on delta1), so unlike our linear Quad these cannot be cached in compute_constants.
T_Real const w = omega * delta1;
T_Real const R11 = focusing ? cos(omega * m_slice_ds) : cosh(omega * m_slice_ds);
T_Real const R11 = focusing ? amrex::Math::cos(omega * m_slice_ds) : amrex::Math::cosh(omega * m_slice_ds);
T_Real const R12 = sx / w;
T_Real const R21 = ax * w * sx;
T_Real const R22 = R11;
T_Real const R33 = focusing ? cosh(omega * m_slice_ds) : cos(omega * m_slice_ds);
T_Real const R33 = focusing ? amrex::Math::cosh(omega * m_slice_ds) : amrex::Math::cos(omega * m_slice_ds);
T_Real const R34 = sy / w;
T_Real const R43 = -ax * w * sy;
T_Real const R44 = R33;
Expand All @@ -212,10 +212,10 @@ namespace impactx::elements
T_Real const term = pt + delta / m_beta;
T_Real const t0 = tout - term * m_slice_ds / delta1;

T_Real const term1 = -(powi<2>(p2) + powi<2>(q2) * powi<2>(w)) * sinh(2_prt * m_slice_ds * omega);
T_Real const term2 = -(powi<2>(p1) - powi<2>(q1) * powi<2>(w)) * sin(2_prt * m_slice_ds * omega);
T_Real const term3 = -2_prt * q2 * p2 * w * cosh(2_prt * m_slice_ds * omega);
T_Real const term4 = -2_prt * q1 * p1 * w * cos(2_prt * m_slice_ds * omega);
T_Real const term1 = -(powi<2>(p2) + powi<2>(q2) * powi<2>(w)) * amrex::Math::sinh(2_prt * m_slice_ds * omega);
T_Real const term2 = -(powi<2>(p1) - powi<2>(q1) * powi<2>(w)) * amrex::Math::sin(2_prt * m_slice_ds * omega);
T_Real const term3 = -2_prt * q2 * p2 * w * amrex::Math::cosh(2_prt * m_slice_ds * omega);
T_Real const term4 = -2_prt * q1 * p1 * w * amrex::Math::cos(2_prt * m_slice_ds * omega);
T_Real const term5 = 2_prt * omega * (
q1 * p1 * delta1 + q2 * p2 * delta1
-(powi<2>(p1) + powi<2>(p2)) * m_slice_ds
Expand Down Expand Up @@ -337,8 +337,8 @@ namespace impactx::elements

// compute trigonometric quantities
auto const [sin_omega_ds, cos_omega_ds] = amrex::Math::sincos(omega_ds);
T_Real const sinh_omega_ds = sinh(omega_ds);
T_Real const cosh_omega_ds = cosh(omega_ds);
T_Real const sinh_omega_ds = amrex::Math::sinh(omega_ds);
T_Real const cosh_omega_ds = amrex::Math::cosh(omega_ds);

// The focusing/defocusing/drift cases differ only by which plane is
// oscillatory vs. hyperbolic and by an overall sign. omega is real
Expand Down
4 changes: 2 additions & 2 deletions src/elements/ChrUniformAcc.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ namespace impactx::elements
// compute focusing constant (1/m) and rotation angle (in rad)
T_Real const theta_ialpha = m_acc == 0_prt
? m_slice_ds / pzi_tot
: log(numer / denom) / m_ez;
: amrex::Math::log(numer / denom) / m_ez;
T_Real const theta = m_alpha * theta_ialpha;
auto const [sin_theta, cos_theta] = amrex::Math::sincos(theta);
T_Real const sin_theta_ialpha = m_alpha == 0_prt
Expand Down Expand Up @@ -355,7 +355,7 @@ namespace impactx::elements
// compute focusing constant (1/m) and rotation angle (in rad)
T_Real const theta_over_alpha = (m_ez == 0_prt) ?
m_slice_ds / pzi_tot :
log(numer / denom) / m_ez;
amrex::Math::log(numer / denom) / m_ez;
T_Real const theta = m_alpha * theta_over_alpha;
auto const [sin2Gth, cos2Gth] = amrex::Math::sincos(2_prt * refpart.gyromagnetic_anomaly * theta);

Expand Down
8 changes: 4 additions & 4 deletions src/elements/ExactQuad.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -231,10 +231,10 @@ namespace impactx::elements
T_Real tout = t;
T_Real ptout = pt;

amrex::ParticleReal const sin_omega_ds = sin(m_omega*tau);
amrex::ParticleReal const cos_omega_ds = cos(m_omega*tau);
amrex::ParticleReal const sinh_omega_ds = sinh(m_omega*tau);
amrex::ParticleReal const cosh_omega_ds = cosh(m_omega*tau);
amrex::ParticleReal const sin_omega_ds = amrex::Math::sin(m_omega*tau);
amrex::ParticleReal const cos_omega_ds = amrex::Math::cos(m_omega*tau);
amrex::ParticleReal const sinh_omega_ds = amrex::Math::sinh(m_omega*tau);
amrex::ParticleReal const cosh_omega_ds = amrex::Math::cosh(m_omega*tau);
amrex::ParticleReal const slice_bg = tau / m_betgam2;

if (m_g > 0.0_prt)
Expand Down
2 changes: 1 addition & 1 deletion src/elements/ExactSbend.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,7 +229,7 @@ namespace impactx::elements
amrex::ParticleIDWrapper<T_IdCpu>{idcpu}.make_invalid(mask2);
{
T_Real const pzf = sqrt(powi<2>(pperp)-powi<2>(pxout));
T_Real const theta = m_slice_phi + asin(px/pperp) - asin(pxout/pperp);
T_Real const theta = m_slice_phi + amrex::Math::asin(px/pperp) - amrex::Math::asin(pxout/pperp);

// update position coordinates
x = -m_rc + rho*m_cos_phi + m_rc*(pzf + px*m_sin_phi - pzi*m_cos_phi);
Expand Down
4 changes: 2 additions & 2 deletions src/elements/ShortRF.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ namespace impactx::elements
// pyout = py;

// tout = t;
ptout = pt - m_V * cos(m_k * t + m_phi) + m_V_cos_phi;
ptout = pt - m_V * amrex::Math::cos(m_k * t + m_phi) + m_V_cos_phi;

// assign updated values
// x = xout;
Expand Down Expand Up @@ -281,7 +281,7 @@ namespace impactx::elements
// Integrated electric field normalized by q/mc^2 (half-step):
T_Real const Ex = 0.0_prt;
T_Real const Ey = 0.0_prt;
T_Real const Ez = 0.5_prt * m_V * cos(m_k * t + m_phi);
T_Real const Ez = 0.5_prt * m_V * amrex::Math::cos(m_k * t + m_phi);

// Quantities required to evaluate the full Thomas-BMT precession vector (before phase space kick).
amrex::ParticleReal beta = m_bgi / gami;
Expand Down
Loading