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

@@ -9,6 +9,7 @@
 
 // Project include(s)
 #include "detray/definitions/detail/qualifiers.hpp"
+#include "detray/definitions/track_parametrization.hpp"
 #include "detray/geometry/detail/surface_kernels.hpp"
 #include "detray/propagator/detail/jacobian_engine.hpp"
 #include "detray/tracks/detail/transform_track_parameters.hpp"
@@ -28,6 +29,7 @@ struct tracking_surface_kernels : public surface_kernels<algebra_t> {
     using bound_param_vector_type = bound_parameters_vector<algebra_t>;
     using free_param_vector_type = free_parameters_vector<algebra_t>;
     using free_matrix_type = free_matrix<algebra_t>;
+    using free_to_path_matrix_type = free_to_path_matrix<algebra_t>;
 
     /// A functor to get from a free to a bound vector
     struct free_to_bound_vector {
@@ -71,8 +73,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 +90,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,10 +109,27 @@ 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);
         }
     };
-};
 
+    /// A function object to get the free to path derivative
+    struct free_to_path_derivative {
+
+        template <typename mask_group_t, typename index_t>
+        DETRAY_HOST_DEVICE inline free_to_path_matrix_type operator()(
+            const mask_group_t& /*mask_group*/, const index_t& /*index*/,
+            const transform3_type& trf3, const vector3_type& pos,
+            const vector3_type& dir, const vector3_type& dtds) const {
+
+            using frame_t = typename mask_group_t::value_type::local_frame;
+
+            return detail::jacobian_engine<
+                algebra_t>::template free_to_path_derivative<frame_t>(pos, dir,
+                                                                      dtds,
+                                                                      trf3);
+        }
+    };
+};
 }  // namespace detray::detail

core/include/detray/geometry/tracking_surface.hpp

@@ -100,6 +100,17 @@ class tracking_surface : public geometry::surface<detector_t> {
         return this->template visit_mask<typename kernels::path_correction>(
             this->transform(ctx), pos, dir, dtds, dqopds);
     }
+
+    /// @returns the free to path derivative
+    DETRAY_HOST_DEVICE
+    constexpr auto free_to_path_derivative(const context &ctx,
+                                           const vector3_type &pos,
+                                           const vector3_type &dir,
+                                           const vector3_type &dtds) const {
+        return this
+            ->template visit_mask<typename kernels::free_to_path_derivative>(
+                this->transform(ctx), pos, dir, dtds);
+    }
 };
 
 template <typename detector_t, typename descr_t>

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

@@ -29,44 +29,29 @@ 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);
-
-            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));
+            return jac_to_local;
         }
     };
 
@@ -143,9 +128,30 @@ 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.free_to_path_derivative(
+            gctx, stepping().pos(), stepping().dir(), stepping.dtds());
+
+        const auto path_correction =
+            path_to_free_derivative * free_to_path_derivative;
+
+        const auto correction_term =
+            matrix::identity<free_matrix<algebra_t>>() + path_correction;
+
+        return free_to_bound_jacobian *
+               (correction_term *
+                (stepping.transport_jacobian() * bound_to_free_jacobian));
     }
 
 };  // namespace detray