src/sanity/contiguous_on.hpp

0001 /**
0002  * traccc library, part of the ACTS project (R&D line)
0003  *
0004  * (c) 2024-2026 CERN for the benefit of the ACTS project
0005  *
0006  * Mozilla Public License Version 2.0
0007  */
0008
0009 #pragma once
0010
0011 // VecMem include(s).
0012 #include <vecmem/memory/device_atomic_ref.hpp>
0013 #include <vecmem/memory/memory_resource.hpp>
0014 #include <vecmem/memory/unique_ptr.hpp>
0015 #include <vecmem/utils/copy.hpp>
0016
0017 // SYCL include
0018 #include <sycl/sycl.hpp>
0019
0020 // System include
0021 #include <concepts>
0022 #include <utility>
0023
0024 namespace traccc::sycl {
0025
0026 namespace kernels {
0027
0028 /// Kernel used in implementing @c traccc::sycl::is_contiguous_on
0029 ///
0030 /// It has to be implemented as a functor, as oneAPI 2024.2.1 gets confused when
0031 /// trying to set up this type of code as a lambda.
0032 ///
0033 template <typename CONTAINER, typename P, typename VIEW,
0034           std::equality_comparable S>
0035 struct is_contiguous_on_compress_adjacent {
0036   /// Constructor
0037   is_contiguous_on_compress_adjacent(P projection, VIEW view,
0038                                      vecmem::data::vector_view<S> out_view)
0039       : m_projection(projection), m_view(view), m_out_view(out_view) {}
0040
0041   /// Execution operator for the kernel
0042   void operator()(::sycl::nd_item<1> item) const {
0043     std::size_t tid = item.get_global_linear_id();
0044
0045     const CONTAINER in(m_view);
0046     vecmem::device_vector<S> out(m_out_view);
0047
0048     if (tid > 0 && tid < in.size()) {
0049       S v1 = m_projection(in.at(static_cast<CONTAINER::size_type>(tid - 1)));
0050       S v2 = m_projection(in.at(static_cast<CONTAINER::size_type>(tid)));
0051
0052       if (v1 != v2) {
0053         out.push_back(v2);
0054       }
0055     } else if (tid == 0) {
0056       out.push_back(
0057           m_projection(in.at(static_cast<CONTAINER::size_type>(tid))));
0058     }
0059   }
0060
0061   /// The projection object to use
0062   P m_projection;
0063   /// View to the input container
0064   VIEW m_view;
0065   /// Output array
0066   vecmem::data::vector_view<S> m_out_view;
0067 };
0068
0069 /// Kernel used in implementing @c traccc::sycl::is_contiguous_on
0070 template <typename T>
0071 class is_contiguous_on_all_unique {};
0072
0073 }  // namespace kernels
0074
0075 /**
0076  * @brief Sanity check that a given container is contiguous on a given
0077  *        projection.
0078  *
0079  * For a container $v$ to be contiguous on a projection $\pi$, it must be the
0080  * case that for all indices $i$ and $j$, if $v_i = v_j$, then all indices $k$
0081  * between $i$ and $j$, $v_i = v_j = v_k$.
0082  *
0083  * @note This function runs in O(n^2) time.
0084  *
0085  * @tparam CONTAINER The type of the (device) container.
0086  * @tparam P The type of projection $\pi$, a callable which returns some
0087  * comparable type.
0088  * @tparam VIEW The type of the view for the container.
0089  * @param projection A projection object of type `P`.
0090  * @param mr A memory resource used for allocating intermediate memory.
0091  * @param view The container which to check for contiguity.
0092  * @return true If the container is contiguous on `P`.
0093  * @return false Otherwise.
0094  */
0095 template <typename CONTAINER, typename P, typename VIEW>
0096   requires std::regular_invocable<P,
0097                                   decltype(std::declval<CONTAINER>().at(0))> &&
0098            std::semiregular<P>
0099 bool is_contiguous_on(P&& projection, vecmem::memory_resource& mr,
0100                       const vecmem::copy& copy, ::sycl::queue& queue,
0101                       const VIEW& view) {
0102   // This should never be a performance-critical step, so we can keep the
0103   // block size fixed.
0104   constexpr int local_size = 512;
0105   constexpr int local_size_2d = 16;
0106
0107   // Grab the number of elements in our vector.
0108   const typename VIEW::size_type n = copy.get_size(view);
0109
0110   // Exit early for empty containers.
0111   if (n == 0) {
0112     return true;
0113   }
0114
0115   // Get the output type of the projection.
0116   using projection_t =
0117       std::invoke_result_t<P, decltype(std::declval<CONTAINER>().at(0))>;
0118
0119   // Allocate memory for intermediate values and outputs, then set them up.
0120   vecmem::data::vector_buffer<projection_t> iout(
0121       n, mr, vecmem::data::buffer_type::resizable);
0122   copy.setup(iout)->wait();
0123   vecmem::unique_alloc_ptr<bool> out = vecmem::make_unique_alloc<bool>(mr);
0124
0125   bool initial_out = true;
0126
0127   ::sycl::event kernel2_memcpy_evt = queue.copy(&initial_out, out.get(), 1);
0128
0129   ::sycl::nd_range<1> compress_adjacent_range{
0130       ::sycl::range<1>(((n + local_size - 1) / local_size) * local_size),
0131       ::sycl::range<1>(local_size)};
0132
0133   // Launch the first kernel, which will squash consecutive equal elements
0134   // into one element.
0135   queue
0136       .submit([&](::sycl::handler& h) {
0137         h.parallel_for<kernels::is_contiguous_on_compress_adjacent<
0138             CONTAINER, P, VIEW, projection_t>>(
0139             compress_adjacent_range,
0140             kernels::is_contiguous_on_compress_adjacent<CONTAINER, P, VIEW,
0141                                                         projection_t>(
0142                 projection, view, iout));
0143       })
0144       .wait_and_throw();
0145
0146   typename vecmem::data::vector_view<projection_t>::size_type host_iout_size =
0147       copy.get_size(iout);
0148   uint32_t grid_size_rd = (host_iout_size + local_size_2d - 1) / local_size_2d;
0149   ::sycl::nd_range<2> all_unique_range{
0150       ::sycl::range<2>(grid_size_rd * local_size_2d,
0151                        grid_size_rd * local_size_2d),
0152       ::sycl::range<2>(local_size_2d, local_size_2d)};
0153
0154   // Launch the second kernel, which will check if the values are unique.
0155   ::sycl::event kernel2_evt = queue.submit([&](::sycl::handler& h) {
0156     h.depends_on(kernel2_memcpy_evt);
0157     h.parallel_for<kernels::is_contiguous_on_all_unique<projection_t>>(
0158         all_unique_range, [in_view = vecmem::get_data(iout),
0159                            out = out.get()](::sycl::nd_item<2> item) {
0160           std::size_t tid_x = item.get_global_id(0);
0161           std::size_t tid_y = item.get_global_id(1);
0162
0163           const vecmem::device_vector<projection_t> in(in_view);
0164
0165           if (tid_x < in.size() && tid_y < in.size() && tid_x != tid_y &&
0166               in.at(static_cast<CONTAINER::size_type>(tid_x)) ==
0167                   in.at(static_cast<CONTAINER::size_type>(tid_y))) {
0168             *out = false;
0169           }
0170         });
0171   });
0172
0173   // Get the result from the device and return it.
0174   bool host_out;
0175
0176   queue.memcpy(&host_out, out.get(), sizeof(bool), {kernel2_evt})
0177       .wait_and_throw();
0178
0179   return host_out;
0180 }
0181 }  // namespace traccc::sycl