From 55c90ed6b6b9a30420ca2caddbd581f7b56a939f Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Tue, 25 Aug 2026 09:19:31 -0700 Subject: [PATCH] SIMD: Call the vectorized transcendentals The SIMD library evaluates sin, cos, sinh and cosh one lane at a time, so the elements that call them per particle got little out of ImpactX_SIMD=ON. AMReX grew amrex::Math overloads that hand them to a vector math library instead; call those. They have to be named. An unqualified sin(x) on a SIMD argument resolves to the SIMD library's own overload through argument-dependent lookup, and neither a using-declaration nor a using-directive changes that: the library's overload either ties, which is ambiguous, or wins partial ordering, which is silent and slow. So amrex::Math::sin(x), the way amrex::Math::powi and amrex::Math::sincos are already spelled here. 36 call sites across six elements. The others were already fine: most precompute their trigonometry scalar-side in compute_constants, and the ParallelForSIMD and ParticleReduceSIMD kernels contain no transcendentals at all -- beamoptic.H only dispatches to the elements, and the beam-moments reduction is sums of products. Measured on the apochromatic example, 2e6 particles, 4 slices, one thread pinned to a P-core, alternating runs, best of seven each: push::ChrQuad 1.917 s -> 0.810 s 2.37x (sin, cos, sinh, cosh per particle) push::ChrDrift 0.066 s -> 0.065 s 1.02x (no transcendentals, the control) whole run 2.783 s -> 1.648 s 1.69x ChrDrift is the useful row: same ParallelForSIMD machinery, no transcendentals, no change. Run-to-run spread was 3 to 6 percent. The reference particle comes out bit-identical. The beam moments differ by at most 1.1e-10 relative, worst in dispersion_y, which is where cancellation concentrates the few ULP a vector math library costs. Points AMReX and vir-simd at the branches that carry the two halves of this until they land upstream. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/dependencies/gcc.sh | 9 +++++---- .github/workflows/macos.yml | 7 +++---- cmake/dependencies/ABLASTR.cmake | 7 +++++-- src/elements/ChrPlasmaLens.H | 24 ++++++++++++------------ src/elements/ChrQuad.H | 20 ++++++++++---------- src/elements/ChrUniformAcc.H | 4 ++-- src/elements/ExactQuad.H | 8 ++++---- src/elements/ExactSbend.H | 2 +- src/elements/ShortRF.H | 4 ++-- 9 files changed, 44 insertions(+), 41 deletions(-) diff --git a/.github/workflows/dependencies/gcc.sh b/.github/workflows/dependencies/gcc.sh index 351a96d896..9aeb60fc83 100755 --- a/.github/workflows/dependencies/gcc.sh +++ b/.github/workflows/dependencies/gcc.sh @@ -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 diff --git a/.github/workflows/macos.yml b/.github/workflows/macos.yml index f88b825851..5ed3156ca6 100644 --- a/.github/workflows/macos.yml +++ b/.github/workflows/macos.yml @@ -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: | diff --git a/cmake/dependencies/ABLASTR.cmake b/cmake/dependencies/ABLASTR.cmake index dff5d3b2ed..840d15cb5e 100644 --- a/cmake/dependencies/ABLASTR.cmake +++ b/cmake/dependencies/ABLASTR.cmake @@ -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)") diff --git a/src/elements/ChrPlasmaLens.H b/src/elements/ChrPlasmaLens.H index 009c5693a6..bcef8c5964 100644 --- a/src/elements/ChrPlasmaLens.H +++ b/src/elements/ChrPlasmaLens.H @@ -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) @@ -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; @@ -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) @@ -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) { diff --git a/src/elements/ChrQuad.H b/src/elements/ChrQuad.H index 67b624ab71..fc7f437522 100644 --- a/src/elements/ChrQuad.H +++ b/src/elements/ChrQuad.H @@ -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; @@ -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 @@ -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 diff --git a/src/elements/ChrUniformAcc.H b/src/elements/ChrUniformAcc.H index 6f9ebe3687..a9ad9369bd 100644 --- a/src/elements/ChrUniformAcc.H +++ b/src/elements/ChrUniformAcc.H @@ -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 @@ -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); diff --git a/src/elements/ExactQuad.H b/src/elements/ExactQuad.H index db2f2bd659..2b8eb4758b 100644 --- a/src/elements/ExactQuad.H +++ b/src/elements/ExactQuad.H @@ -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) diff --git a/src/elements/ExactSbend.H b/src/elements/ExactSbend.H index 3741d15c40..fcfd2ff6f7 100644 --- a/src/elements/ExactSbend.H +++ b/src/elements/ExactSbend.H @@ -229,7 +229,7 @@ namespace impactx::elements amrex::ParticleIDWrapper{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); diff --git a/src/elements/ShortRF.H b/src/elements/ShortRF.H index 811073a7c4..8058b626b5 100644 --- a/src/elements/ShortRF.H +++ b/src/elements/ShortRF.H @@ -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; @@ -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;