From 8b4321f910c5a4fdaccf9db94686abbda3b82e3d Mon Sep 17 00:00:00 2001 From: Stephen Nicholas Swatman Date: Sun, 8 Dec 2024 21:49:48 +0100 Subject: [PATCH] Implement first in-place matrix multiplications This is the first round of implementations of the new in-place matrix operations introduced in algebra-plugins. While it doesn't necessarily make the code more readable, it does increase performance quite significantly. More to come later. --- .../actors/parameter_transporter.hpp | 35 ++++++++++--------- .../detray/propagator/base_stepper.hpp | 10 ++++++ .../detray/propagator/line_stepper.hpp | 4 +-- core/include/detray/propagator/rk_stepper.ipp | 3 +- 4 files changed, 32 insertions(+), 20 deletions(-) diff --git a/core/include/detray/propagator/actors/parameter_transporter.hpp b/core/include/detray/propagator/actors/parameter_transporter.hpp index de1340334..9e9726136 100644 --- a/core/include/detray/propagator/actors/parameter_transporter.hpp +++ b/core/include/detray/propagator/actors/parameter_transporter.hpp @@ -30,16 +30,16 @@ struct parameter_transporter : actor { using bound_to_free_matrix_t = bound_to_free_matrix; /// @} - struct get_full_jacobian_kernel { + struct update_full_jacobian_kernel { template - DETRAY_HOST_DEVICE inline bound_matrix_t operator()( + DETRAY_HOST_DEVICE inline void operator()( const mask_group_t& /*mask_group*/, const index_t& /*index*/, const transform3_type& trf3, const bound_to_free_matrix_t& bound_to_free_jacobian, const material* vol_mat_ptr, - const stepper_state_t& stepping) const { + stepper_state_t& stepping) const { using frame_t = typename mask_group_t::value_type::shape:: template local_frame_type; @@ -51,7 +51,7 @@ struct parameter_transporter : actor { typename jacobian_engine_t::free_to_bound_matrix_type; // Free to bound jacobian at the destination surface - const free_to_bound_matrix_t free_to_bound_jacobian = + free_to_bound_matrix_t free_to_bound_jacobian = jacobian_engine_t::free_to_bound_jacobian(trf3, stepping()); // Path correction factor @@ -63,8 +63,12 @@ struct parameter_transporter : actor { const free_matrix_t correction_term = matrix::identity() + path_correction; - return free_to_bound_jacobian * correction_term * - stepping.transport_jacobian() * bound_to_free_jacobian; + algebra::generic::math::set_inplace_product_right( + free_to_bound_jacobian, correction_term); + + algebra::generic::math::set_product( + stepping.full_jacobian(), free_to_bound_jacobian, + (stepping.transport_jacobian() * bound_to_free_jacobian)); } }; @@ -104,17 +108,14 @@ struct parameter_transporter : actor { const auto vol_mat_ptr = vol.has_material() ? vol.material_parameters(stepping().pos()) : nullptr; - stepping.set_full_jacobian( - sf.template visit_mask( - sf.transform(gctx), bound_to_free_jacobian, vol_mat_ptr, - propagation._stepping)); - - // Calculate surface-to-surface covariance transport - const bound_matrix_t new_cov = - stepping.full_jacobian() * bound_params.covariance() * - matrix::transpose(stepping.full_jacobian()); - - stepping.bound_params().set_covariance(new_cov); + sf.template visit_mask( + sf.transform(gctx), bound_to_free_jacobian, vol_mat_ptr, + propagation._stepping); + + algebra::generic::math::set_inplace_product_left( + bound_params.covariance(), stepping.full_jacobian()); + algebra::generic::math::set_inplace_product_right_transpose( + bound_params.covariance(), stepping.full_jacobian()); } // Convert free to bound vector diff --git a/core/include/detray/propagator/base_stepper.hpp b/core/include/detray/propagator/base_stepper.hpp index 54d2526b5..5b9f13d82 100644 --- a/core/include/detray/propagator/base_stepper.hpp +++ b/core/include/detray/propagator/base_stepper.hpp @@ -180,12 +180,22 @@ class base_stepper { return m_jac_transport; } + /// @returns the current transport Jacbian. + DETRAY_HOST_DEVICE + inline free_matrix_type &transport_jacobian() { + return m_jac_transport; + } + /// @returns the current full Jacbian. DETRAY_HOST_DEVICE inline const bound_matrix_type &full_jacobian() const { return m_full_jacobian; } + /// @returns the current full Jacbian. + DETRAY_HOST_DEVICE + inline bound_matrix_type &full_jacobian() { return m_full_jacobian; } + /// Set new full Jacbian. DETRAY_HOST_DEVICE inline void set_full_jacobian(const bound_matrix_type &jac) { diff --git a/core/include/detray/propagator/line_stepper.hpp b/core/include/detray/propagator/line_stepper.hpp index 2219ba2bf..ffa3f0d4b 100644 --- a/core/include/detray/propagator/line_stepper.hpp +++ b/core/include/detray/propagator/line_stepper.hpp @@ -68,8 +68,8 @@ class line_stepper final /// NOTE: Let's skip the element for d(time)/d(qoverp) for the /// moment.. - - this->set_transport_jacobian(D * this->transport_jacobian()); + algebra::generic::math::set_inplace_product_left( + this->transport_jacobian(), D); } DETRAY_HOST_DEVICE diff --git a/core/include/detray/propagator/rk_stepper.ipp b/core/include/detray/propagator/rk_stepper.ipp index 1759fd2ab..bc36cf4d1 100644 --- a/core/include/detray/propagator/rk_stepper.ipp +++ b/core/include/detray/propagator/rk_stepper.ipp @@ -331,7 +331,8 @@ DETRAY_HOST_DEVICE inline void detray::rk_stepper< getter::set_block(D, dFdqop, 0u, 7u); getter::set_block(D, dGdqop, 4u, 7u); - this->set_transport_jacobian(D * this->transport_jacobian()); + algebra::generic::math::set_inplace_product_left(this->transport_jacobian(), + D); } template