Merge pull request #295 from krasznaa/SoAGetSize-main-20240918

get_size(...) for SoA Containers, main branch (2024.09.20.)

core/include/vecmem/utils/copy.hpp

@@ -1,7 +1,7 @@
 /*
  * VecMem project, part of the ACTS project (R&D line)
  *
- * (c) 2021-2023 CERN for the benefit of the ACTS project
+ * (c) 2021-2024 CERN for the benefit of the ACTS project
  *
  * Mozilla Public License Version 2.0
  */
@@ -179,6 +179,11 @@ class VECMEM_CORE_EXPORT copy {
         edm::host<edm::schema<VARTYPES...>>& to,
         type::copy_type cptype = type::unknown) const;
 
+    /// Get the (outer) size of a resizable SoA container
+    template <typename... VARTYPES>
+    typename edm::view<edm::schema<VARTYPES...>>::size_type get_size(
+        const edm::view<edm::schema<VARTYPES...>>& data) const;
+
     /// @}
 
 protected:

core/include/vecmem/utils/impl/copy.ipp

@@ -444,6 +444,45 @@ copy::event_type copy::operator()(
     return operator()(from_view, vecmem::get_data(to_vec), cptype);
 }
 
+template <typename... VARTYPES>
+typename edm::view<edm::schema<VARTYPES...>>::size_type copy::get_size(
+    const edm::view<edm::schema<VARTYPES...>>& data) const {
+
+    // Start by taking the capacity of the container.
+    typename edm::view<edm::schema<VARTYPES...>>::size_type size =
+        data.capacity();
+
+    // If there are jagged vectors in the container and/or the container is not
+    // resizable, we're done. It is done in two separate if statements to avoid
+    // MSVC trying to be too smart, and giving a warning...
+    if constexpr (std::disjunction_v<
+                      edm::type::details::is_jagged_vector<VARTYPES>...>) {
+        return size;
+    } else {
+        // All the rest is put into an else block, to avoid MSVC trying to be
+        // too smart, and giving a warning about unreachable code...
+        if (data.size().ptr() == nullptr) {
+            return size;
+        }
+
+        // A small security check.
+        assert(data.size().size() ==
+               sizeof(typename edm::view<edm::schema<VARTYPES...>>::size_type));
+
+        // Get the exact size of the container.
+        do_copy(sizeof(typename edm::view<edm::schema<VARTYPES...>>::size_type),
+                data.size().ptr(), &size, type::unknown);
+        // We have to wait for this to finish, since the "size" variable is
+        // not going to be available outside of this function. And
+        // asynchronous SYCL memory copies can happen from variables on the
+        // stack as well...
+        create_event()->wait();
+
+        // Return what we got.
+        return size;
+    }
+}
+
 template <typename TYPE>
 bool copy::copy_view_impl(
     const data::vector_view<std::add_const_t<TYPE>>& from_view,

tests/common/soa_copy_tests.ipp

@@ -52,6 +52,9 @@ void soa_copy_tests_base<CONTAINER>::host_to_fixed_device_to_host_direct() {
                 vecmem::copy::type::host_to_device)
         ->wait();
 
+    // Check the size of the device buffer.
+    EXPECT_EQ(input.size(), main_copy().get_size(device_buffer));
+
     // Create the target host container.
     typename CONTAINER::host target{host_mr()};
 
@@ -134,6 +137,9 @@ void soa_copy_tests_base<CONTAINER>::host_to_resizable_device_to_host() {
                 vecmem::copy::type::host_to_device)
         ->wait();
 
+    // Check the size of the device buffer.
+    EXPECT_EQ(input.size(), main_copy().get_size(device_buffer));
+
     // Create the target host container.
     typename CONTAINER::host target{host_mr()};
 
@@ -166,6 +172,9 @@ void soa_copy_tests_base<
                 vecmem::copy::type::host_to_device)
         ->wait();
 
+    // Check the size of the device buffer.
+    EXPECT_EQ(input.size(), main_copy().get_size(device_buffer1));
+
     // Create the (resizable) device buffer.
     typename CONTAINER::buffer device_buffer2;
     vecmem::testing::make_buffer(device_buffer2, main_mr(), host_mr(),
@@ -176,6 +185,9 @@ void soa_copy_tests_base<
                 vecmem::copy::type::device_to_device)
         ->wait();
 
+    // Check the size of the device buffer.
+    EXPECT_EQ(input.size(), main_copy().get_size(device_buffer2));
+
     // Create the target host container.
     typename CONTAINER::host target{host_mr()};