Reduce thread divergence in covariance transport

Currently, the covariance transport uses the Jacobian engine which is
templated on the frame type. While this makes the code easier to read
and write, it requires the compiler to duplicate a lot of code for each
of the frame types, and this duplicated code counts as multiple branches
for the sake of GPU execution. Thus, this templating increases the
amount of thread divergence.

This commit refactors the Jacobian engine into smaller parts, some of
which are templated on the frame type and some of which are not. Client
code and then take a more fine-grained approach to branching and improve
divergence.

Author: stephenswat

core/include/detray/geometry/detail/tracking_surface_kernels.hpp

@@ -71,8 +71,9 @@ struct tracking_surface_kernels : public surface_kernels<algebra_t> {
 
             using frame_t = typename mask_group_t::value_type::local_frame;
 
-            return detail::jacobian_engine<frame_t>::free_to_bound_jacobian(
-                trf3, free_vec);
+            return detail::jacobian_engine<
+                algebra_t>::template free_to_bound_jacobian<frame_t>(trf3,
+                                                                     free_vec);
         }
     };
 
@@ -87,8 +88,9 @@ struct tracking_surface_kernels : public surface_kernels<algebra_t> {
 
             using frame_t = typename mask_group_t::value_type::local_frame;
 
-            return detail::jacobian_engine<frame_t>::bound_to_free_jacobian(
-                trf3, mask_group[index], bound_vec);
+            return detail::jacobian_engine<algebra_t>::
+                template bound_to_free_jacobian<frame_t>(
+                    trf3, mask_group[index], bound_vec);
         }
     };
 
@@ -105,8 +107,8 @@ struct tracking_surface_kernels : public surface_kernels<algebra_t> {
 
             using frame_t = typename mask_group_t::value_type::local_frame;
 
-            return detail::jacobian_engine<frame_t>::path_correction(
-                pos, dir, dtds, dqopds, trf3);
+            return detail::jacobian_engine<algebra_t>::template path_correction<
+                frame_t>(pos, dir, dtds, dqopds, trf3);
         }
     };
 };

core/include/detray/propagator/actors/parameter_transporter.hpp

@@ -29,44 +29,45 @@ struct parameter_transporter : actor {
     using bound_matrix_t = bound_matrix<algebra_t>;
     // Matrix type for bound to free jacobian
     using bound_to_free_matrix_t = bound_to_free_matrix<algebra_t>;
+    // Matrix type for free to bound jacobian
+    using free_to_bound_matrix_t = free_to_bound_matrix<algebra_t>;
     /// @}
 
-    struct get_full_jacobian_kernel {
-
+    struct get_free_to_bound_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 free_to_bound_matrix_t 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 {
-
             using frame_t = typename mask_group_t::value_type::shape::
                 template local_frame_type<algebra_t>;
 
-            using jacobian_engine_t = detail::jacobian_engine<frame_t>;
-
-            using free_matrix_t = free_matrix<algebra_t>;
-            using free_to_bound_matrix_t =
-                typename jacobian_engine_t::free_to_bound_matrix_type;
+            // Declare jacobian for bound to free coordinate transform
+            free_to_bound_matrix_t jac_to_local =
+                matrix::zero<free_to_bound_matrix_t>();
 
-            // Free to bound jacobian at the destination surface
-            const free_to_bound_matrix_t free_to_bound_jacobian =
-                jacobian_engine_t::free_to_bound_jacobian(trf3, stepping());
+            detail::jacobian_engine<algebra_t>::
+                template free_to_bound_jacobian_step_1<frame_t>(
+                    jac_to_local, trf3, stepping().pos(), stepping().dir());
 
-            // Path correction factor
-            const free_matrix_t path_correction =
-                jacobian_engine_t::path_correction(
-                    stepping().pos(), stepping().dir(), stepping.dtds(),
-                    stepping.dqopds(vol_mat_ptr), trf3);
+            return jac_to_local;
+        }
+    };
 
-            const free_matrix_t correction_term =
-                matrix::identity<free_matrix_t>() + path_correction;
+    struct get_free_to_path_derivative_kernel {
+        template <typename mask_group_t, typename index_t,
+                  typename stepper_state_t>
+        DETRAY_HOST_DEVICE inline auto operator()(
+            const mask_group_t& /*mask_group*/, const index_t& /*index*/,
+            const transform3_type& trf3,
+            const stepper_state_t& stepping) const {
+            using frame_t = typename mask_group_t::value_type::shape::
+                template local_frame_type<algebra_t>;
 
-            return free_to_bound_jacobian *
-                   (correction_term *
-                    (stepping.transport_jacobian() * bound_to_free_jacobian));
+            return detail::jacobian_engine<algebra_t>::
+                template free_to_path_derivative<frame_t>(
+                    stepping().pos(), stepping().dir(), stepping.dtds(), trf3);
         }
     };
 
@@ -143,9 +144,32 @@ struct parameter_transporter : actor {
                                      ? vol.material_parameters(stepping().pos())
                                      : nullptr;
 
-        return sf.template visit_mask<get_full_jacobian_kernel>(
-            sf.transform(gctx), bound_to_free_jacobian, vol_mat_ptr,
-            propagation._stepping);
+        auto free_to_bound_jacobian =
+            sf.template visit_mask<get_free_to_bound_jacobian_kernel>(
+                sf.transform(gctx), propagation._stepping);
+
+        detail::jacobian_engine<algebra_t>::free_to_bound_jacobian_step_2(
+            free_to_bound_jacobian, stepping().dir());
+
+        const auto path_to_free_derivative =
+            detail::jacobian_engine<algebra_t>::path_to_free_derivative(
+                stepping().dir(), stepping.dtds(),
+                stepping.dqopds(vol_mat_ptr));
+
+        const auto free_to_path_derivative =
+            sf.template visit_mask<get_free_to_path_derivative_kernel>(
+                sf.transform(gctx), propagation._stepping);
+
+        const auto path_correction =
+            path_to_free_derivative * free_to_path_derivative;
+
+        const auto correction_term =
+            matrix::identity<std::decay_t<decltype(path_correction)>>() +
+            path_correction;
+
+        return free_to_bound_jacobian *
+               (correction_term *
+                (stepping.transport_jacobian() * bound_to_free_jacobian));
     }
 
 };  // namespace detray

core/include/detray/propagator/detail/jacobian_engine.hpp

@@ -22,15 +22,12 @@
 namespace detray::detail {
 
 /// @brief Generate Jacobians
-template <typename frame_t>
-requires std::is_object_v<typename frame_t::loc_point> struct jacobian_engine {
+template <typename algebra_t>
+struct jacobian_engine {
 
     /// @name Type definitions for the struct
     /// @{
-    using frame_type = frame_t;
-    using jacobian_t = jacobian<frame_t>;
-
-    using algebra_type = typename frame_t::algebra_type;
+    using algebra_type = algebra_t;
     using transform3_type = dtransform3D<algebra_type>;
     using scalar_type = dscalar<algebra_type>;
     using point3_type = dpoint3D<algebra_type>;
@@ -42,39 +39,33 @@ requires std::is_object_v<typename frame_t::loc_point> struct jacobian_engine {
     using path_to_free_matrix_type = path_to_free_matrix<algebra_type>;
     /// @}
 
-    template <typename mask_t>
-    DETRAY_HOST_DEVICE static inline bound_to_free_matrix_type
-    bound_to_free_jacobian(
-        const transform3_type& trf3, const mask_t& mask,
-        const bound_parameters_vector<algebra_type>& bound_vec) {
+    template <typename frame_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline void bound_to_free_jacobian_step_1(
+            bound_to_free_matrix_type& jac_to_global,
+            const transform3_type& trf3, const vector3_type& pos,
+            const vector3_type& dir) {
+        using jacobian_t = jacobian<frame_t>;
 
-        // Declare jacobian for bound to free coordinate transform
-        bound_to_free_matrix_type jac_to_global =
-            matrix::zero<bound_to_free_matrix_type>();
+        // Set d(x,y,z)/d(loc0, loc1)
+        jacobian_t::set_bound_pos_to_free_pos_derivative(jac_to_global, trf3,
+                                                         pos, dir);
+    }
+
+    DETRAY_HOST_DEVICE static inline void bound_to_free_jacobian_step_2(
+        bound_to_free_matrix_type& jac_to_global,
+        const bound_parameters_vector<algebra_type>& bound_vec) {
+        // Set d(bound time)/d(free time)
+        getter::element(jac_to_global, e_free_time, e_bound_time) = 1.f;
 
         // Get trigonometric values
         const scalar_type theta{bound_vec.theta()};
         const scalar_type phi{bound_vec.phi()};
-
         const scalar_type cos_theta{math::cos(theta)};
         const scalar_type sin_theta{math::sin(theta)};
         const scalar_type cos_phi{math::cos(phi)};
         const scalar_type sin_phi{math::sin(phi)};
 
-        // Global position and direction
-        const free_track_parameters<algebra_type> free_params =
-            bound_to_free_vector(trf3, mask, bound_vec);
-
-        const vector3_type pos = free_params.pos();
-        const vector3_type dir = free_params.dir();
-
-        // Set d(x,y,z)/d(loc0, loc1)
-        jacobian_t::set_bound_pos_to_free_pos_derivative(jac_to_global, trf3,
-                                                         pos, dir);
-
-        // Set d(bound time)/d(free time)
-        getter::element(jac_to_global, e_free_time, e_bound_time) = 1.f;
-
         // Set d(n_x,n_y,n_z)/d(phi, theta)
         getter::element(jac_to_global, e_free_dir0, e_bound_phi) =
             -sin_theta * sin_phi;
@@ -86,42 +77,71 @@ requires std::is_object_v<typename frame_t::loc_point> struct jacobian_engine {
             cos_theta * sin_phi;
         getter::element(jac_to_global, e_free_dir2, e_bound_theta) = -sin_theta;
         getter::element(jac_to_global, e_free_qoverp, e_bound_qoverp) = 1.f;
+    }
+
+    template <typename frame_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline void bound_to_free_jacobian_step_3(
+            bound_to_free_matrix_type& jac_to_global,
+            const transform3_type& trf3, const vector3_type& pos,
+            const vector3_type& dir) {
+        using jacobian_t = jacobian<frame_t>;
 
         // Set d(x,y,z)/d(phi, theta)
         jacobian_t::set_bound_angle_to_free_pos_derivative(jac_to_global, trf3,
                                                            pos, dir);
-
-        return jac_to_global;
     }
 
-    DETRAY_HOST_DEVICE
-    static inline free_to_bound_matrix_type free_to_bound_jacobian(
-        const transform3_type& trf3,
-        const free_track_parameters<algebra_type>& free_params) {
-
-        // Declare jacobian for bound to free coordinate transform
-        free_to_bound_matrix_type jac_to_local =
-            matrix::zero<free_to_bound_matrix_type>();
+    template <typename frame_t, typename mask_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline bound_to_free_matrix_type
+        bound_to_free_jacobian(
+            const transform3_type& trf3, const mask_t& mask,
+            const bound_parameters_vector<algebra_type>& bound_vec) {
+        bound_to_free_matrix_type jac_to_global =
+            matrix::zero<bound_to_free_matrix_type>();
 
         // Global position and direction
+        const free_track_parameters<algebra_type> free_params =
+            bound_to_free_vector(trf3, mask, bound_vec);
+
         const vector3_type pos = free_params.pos();
         const vector3_type dir = free_params.dir();
 
-        const scalar_type theta{vector::theta(dir)};
-        const scalar_type phi{vector::phi(dir)};
+        bound_to_free_jacobian_step_1<frame_t>(jac_to_global, trf3, pos, dir);
+        bound_to_free_jacobian_step_2(jac_to_global, bound_vec);
+        bound_to_free_jacobian_step_3<frame_t>(jac_to_global, trf3, pos, dir);
 
-        const scalar_type cos_theta{math::cos(theta)};
-        const scalar_type sin_theta{math::sin(theta)};
-        const scalar_type cos_phi{math::cos(phi)};
-        const scalar_type sin_phi{math::sin(phi)};
+        return jac_to_global;
+    }
+
+    template <typename frame_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline void free_to_bound_jacobian_step_1(
+            free_to_bound_matrix_type& jac_to_local,
+            const transform3_type& trf3, const vector3_type& pos,
+            const vector3_type& dir) {
+        using jacobian_t = jacobian<frame_t>;
 
         // Set d(loc0, loc1)/d(x,y,z)
         jacobian_t::set_free_pos_to_bound_pos_derivative(jac_to_local, trf3,
                                                          pos, dir);
+    }
+
+    DETRAY_HOST_DEVICE static inline void free_to_bound_jacobian_step_2(
+        free_to_bound_matrix_type& jac_to_local, const vector3_type& dir) {
 
         // Set d(free time)/d(bound time)
         getter::element(jac_to_local, e_bound_time, e_free_time) = 1.f;
 
+        const scalar_type theta{vector::theta(dir)};
+        const scalar_type phi{vector::phi(dir)};
+
+        const scalar_type cos_theta{math::cos(theta)};
+        const scalar_type sin_theta{math::sin(theta)};
+        const scalar_type cos_phi{math::cos(phi)};
+        const scalar_type sin_phi{math::sin(phi)};
+
         // Set d(phi, theta)/d(n_x, n_y, n_z)
         // @note This codes have a serious bug when theta is equal to zero...
         getter::element(jac_to_local, e_bound_phi, e_free_dir0) =
@@ -136,17 +156,42 @@ requires std::is_object_v<typename frame_t::loc_point> struct jacobian_engine {
 
         // Set d(Free Qop)/d(Bound Qop)
         getter::element(jac_to_local, e_bound_qoverp, e_free_qoverp) = 1.f;
+    }
+
+    template <typename frame_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline free_to_bound_matrix_type
+        free_to_bound_jacobian(
+            const transform3_type& trf3,
+            const free_track_parameters<algebra_type>& free_params) {
+
+        // Declare jacobian for bound to free coordinate transform
+        free_to_bound_matrix_type jac_to_local =
+            matrix::zero<free_to_bound_matrix_type>();
+
+        // Global position and direction
+        const vector3_type pos = free_params.pos();
+        const vector3_type dir = free_params.dir();
+
+        free_to_bound_jacobian_step_1<frame_t>(jac_to_local, trf3, pos, dir);
+        free_to_bound_jacobian_step_2(jac_to_local, dir);
 
         return jac_to_local;
     }
 
-    DETRAY_HOST_DEVICE static inline free_matrix<algebra_type> path_correction(
-        const vector3_type& pos, const vector3_type& dir,
-        const vector3_type& dtds, const scalar_type dqopds,
-        const transform3_type& trf3) {
+    template <typename frame_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline auto free_to_path_derivative(
+            const vector3_type& pos, const vector3_type& dir,
+            const vector3_type& dtds, const transform3_type& trf3) {
+        using jacobian_t = jacobian<frame_t>;
 
-        free_to_path_matrix_type path_derivative =
-            jacobian_t::path_derivative(trf3, pos, dir, dtds);
+        return jacobian_t::path_derivative(trf3, pos, dir, dtds);
+    }
+
+    DETRAY_HOST_DEVICE static inline auto path_to_free_derivative(
+        const vector3_type& dir, const vector3_type& dtds,
+        const scalar_type dqopds) {
 
         path_to_free_matrix_type derivative =
             matrix::zero<path_to_free_matrix_type>();
@@ -158,7 +203,18 @@ requires std::is_object_v<typename frame_t::loc_point> struct jacobian_engine {
         getter::element(derivative, e_free_dir2, 0u) = dtds[2];
         getter::element(derivative, e_free_qoverp, 0u) = dqopds;
 
-        return derivative * path_derivative;
+        return derivative;
+    }
+
+    template <typename frame_t>
+    requires std::is_object_v<typename frame_t::loc_point>
+        DETRAY_HOST_DEVICE static inline free_matrix<algebra_type>
+        path_correction(const vector3_type& pos, const vector3_type& dir,
+                        const vector3_type& dtds, const scalar_type dqopds,
+                        const transform3_type& trf3) {
+
+        return path_to_free_derivative(dir, dtds, dqopds) *
+               free_to_path_derivative<frame_t>(pos, dir, dtds, trf3);
     }
 };
 

core/include/detray/utils/curvilinear_frame.hpp

@@ -27,8 +27,7 @@ struct curvilinear_frame {
     using unit_vectors_type = unit_vectors<vector3>;
     using bound_to_free_matrix_type = bound_to_free_matrix<algebra_t>;
     using bound_vector_type = bound_parameters_vector<algebra_t>;
-    using jacobian_engine_type =
-        detail::jacobian_engine<cartesian2D<algebra_t>>;
+    using jacobian_engine_type = detail::jacobian_engine<algebra_t>;
     using free_track_parameters_type = free_track_parameters<algebra_t>;
 
     DETRAY_HOST_DEVICE
@@ -47,8 +46,9 @@ struct curvilinear_frame {
 
     DETRAY_HOST_DEVICE
     bound_to_free_matrix_type bound_to_free_jacobian() const {
-        return jacobian_engine_type().bound_to_free_jacobian(
-            m_trf, mask<rectangle2D, algebra_t>{}, m_bound_vec);
+        return jacobian_engine_type()
+            .template bound_to_free_jacobian<cartesian2D<algebra_t>>(
+                m_trf, mask<rectangle2D, algebra_t>{}, m_bound_vec);
     }
 
     transform3_type m_trf{};