Skip to content
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.

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

Implement first in-place matrix multiplications #901

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
35 changes: 18 additions & 17 deletions core/include/detray/propagator/actors/parameter_transporter.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,16 +30,16 @@ struct parameter_transporter : actor {
using bound_to_free_matrix_t = bound_to_free_matrix<algebra_t>;
/// @}

struct get_full_jacobian_kernel {
struct update_full_jacobian_kernel {

template <typename mask_group_t, typename index_t,
typename stepper_state_t>
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<scalar_type>* 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<algebra_t>;
Expand All @@ -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
Expand All @@ -63,8 +63,12 @@ struct parameter_transporter : actor {
const free_matrix_t correction_term =
matrix::identity<free_matrix_t>() + path_correction;

return free_to_bound_jacobian * correction_term *
stepping.transport_jacobian() * bound_to_free_jacobian;
algebra::generic::math::set_inplace_product_right(
beomki-yeo marked this conversation as resolved.
Show resolved Hide resolved
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));
}
};

Expand Down Expand Up @@ -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<get_full_jacobian_kernel>(
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<update_full_jacobian_kernel>(
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
Expand Down
10 changes: 10 additions & 0 deletions core/include/detray/propagator/base_stepper.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
4 changes: 2 additions & 2 deletions core/include/detray/propagator/line_stepper.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 2 additions & 1 deletion core/include/detray/propagator/rk_stepper.ipp
Original file line number Diff line number Diff line change
Expand Up @@ -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 <typename magnetic_field_t, typename algebra_t, typename constraint_t,
Expand Down
Loading