Reduce thread divergence in covariance transport

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