Merge pull request #40285 from ericcano/multiCollections

Add support to multi layout portable collections (any number)

DataFormats/Portable/README.md

@@ -126,3 +126,33 @@ should explicitly use the `PortableHostObject<T>` and `PortableHostCollection<T>
 Modules that implement portable interfaces (_e.g._ producers) should use the generic types based on
 `ALPAKA_ACCELERATOR_NAMESPACE::PortableObject<T>` or `PortableObject<T, TDev>`, and
 `ALPAKA_ACCELERATOR_NAMESPACE::PortableCollection<T>` or `PortableCollection<T, TDev>`.
+
+## Multi layout collections
+
+Some use cases require multiple sets of columns of different sizes. This is can be achieved in a single
+`PortableCollection` using `PortableCollection2<T1, T2>`, `PortableCollection3<T1, T2, T3>` and so on up to
+`PortableCollection5<...>`. The numbered, fixed size wrappers are needed in order to be added to the ROOT dictionary.
+Behind the scenes recursive `PortableHostMultiCollection<T0, ...>` and
+`ALPAKA_ACCELERATOR_NAMESPACE::PortableDeviceMultiCollection<TDev, T0, ...>` (note the reversed parameter order) provide
+the actual class definitions.
+
+## ROOT dictionary declaration helper scripts
+
+In order to be serialized by ROOT, the products need to be added to its dictionary. This happens during `scram build`
+as instructed in `<module>/src/classes_dev.xml` and `<module>/src/alpaka/classes_cuda_def.xml` and
+`<module>/src/alpaka/classes_rocm_def.xml`. Two scripts generate the code to be added to the xml files.
+Both scripts expect the collections to be aliased as in:
+```
+using TestDeviceMultiCollection3 = PortableCollection3<TestSoA, TestSoA2, TestSoA3>;
+```
+
+For the host xml, SoA layouts have to be listed and duplicates should be removed manually is multiple
+collections share a same layout. The scripts are called as follows:
+```
+./DataFormats/Portable/scripts/portableHostCollectionHints portabletest::TestHostMultiCollection3  \
+            portabletest::TestSoALayout portabletest::TestSoALayout2 portabletest::TestSoALayout3
+
+./DataFormats/Portable/scripts/portableDeviceCollectionHints portabletest::TestHostMultiCollection3
+```
+The layouts should not be added as parameters for the device collection. Those script can be use equally with the
+single layout collections or multi layout collections.

DataFormats/Portable/interface/PortableCollection.h

@@ -22,12 +22,18 @@ namespace traits {
     using CollectionType = PortableHostCollection<T>;
   };
 
+  template <typename TDev, typename T0, typename... Args>
+  class PortableMultiCollectionTrait;
 }  // namespace traits
 
 // type alias for a generic SoA-based product
 template <typename T, typename TDev, typename = std::enable_if_t<alpaka::isDevice<TDev>>>
 using PortableCollection = typename traits::PortableCollectionTrait<T, TDev>::CollectionType;
 
+// type alias for a generic SoA-based product
+template <typename TDev, typename T0, typename... Args>
+using PortableMultiCollection = typename traits::PortableMultiCollectionTrait<TDev, T0, Args...>::CollectionType;
+
 // define how to copy PortableCollection between host and device
 namespace cms::alpakatools {
   template <typename TLayout, typename TDevice>
@@ -40,6 +46,16 @@ namespace cms::alpakatools {
     }
   };
 
+  template <typename TDev, typename T0, typename... Args>
+  struct CopyToHost<PortableDeviceMultiCollection<TDev, T0, Args...>> {
+    template <typename TQueue>
+    static auto copyAsync(TQueue& queue, PortableDeviceMultiCollection<TDev, T0, Args...> const& srcData) {
+      PortableHostMultiCollection<T0, Args...> dstData(srcData.sizes(), queue);
+      alpaka::memcpy(queue, dstData.buffer(), srcData.buffer());
+      return dstData;
+    }
+  };
+
   template <typename TLayout>
   struct CopyToDevice<PortableHostCollection<TLayout>> {
     template <typename TQueue>
@@ -50,6 +66,17 @@ namespace cms::alpakatools {
       return dstData;
     }
   };
+
+  template <typename TDev, typename T0, typename... Args>
+  struct CopyToDevice<PortableHostMultiCollection<TDev, T0, Args...>> {
+    template <typename TQueue>
+    static auto copyAsync(TQueue& queue, PortableHostMultiCollection<TDev, T0, Args...> const& srcData) {
+      using TDevice = typename alpaka::trait::DevType<TQueue>::type;
+      PortableDeviceMultiCollection<TDev, T0, Args...> dstData(srcData.sizes(), queue);
+      alpaka::memcpy(queue, dstData.buffer(), srcData.buffer());
+      return dstData;
+    }
+  };
 }  // namespace cms::alpakatools
 
 #endif  // DataFormats_Portable_interface_PortableCollection_h

DataFormats/Portable/interface/PortableCollectionCommon.h

@@ -0,0 +1,103 @@
+#ifndef DataFormats_Portable_interface_PortableCollectionCommon_h
+#define DataFormats_Portable_interface_PortableCollectionCommon_h
+
+#include <cstddef>
+#include <type_traits>
+#include <array>
+
+namespace portablecollection {
+
+  // Note: if there are other uses for this, it could be moved to a central place
+  template <std::size_t Start, std::size_t End, std::size_t Inc = 1, typename F>
+  constexpr void constexpr_for(F&& f) {
+    if constexpr (Start < End) {
+      f(std::integral_constant<std::size_t, Start>());
+      constexpr_for<Start + Inc, End, Inc>(std::forward<F>(f));
+    }
+  }
+
+  template <std::size_t Idx, typename T>
+  struct CollectionLeaf {
+    CollectionLeaf() = default;
+    CollectionLeaf(std::byte* buffer, int32_t elements) : layout_(buffer, elements), view_(layout_) {}
+    template <std::size_t N>
+    CollectionLeaf(std::byte* buffer, std::array<int32_t, N> const& sizes)
+        : layout_(buffer, sizes[Idx]), view_(layout_) {
+      static_assert(N >= Idx);
+    }
+    using Layout = T;
+    using View = typename Layout::View;
+    using ConstView = typename Layout::ConstView;
+    Layout layout_;  //
+    View view_;      //!
+    // Make sure types are not void.
+    static_assert(not std::is_same<T, void>::value);
+  };
+
+  template <std::size_t Idx, typename T, typename... Args>
+  struct CollectionImpl : public CollectionLeaf<Idx, T>, public CollectionImpl<Idx + 1, Args...> {
+    CollectionImpl() = default;
+    CollectionImpl(std::byte* buffer, int32_t elements) : CollectionLeaf<Idx, T>(buffer, elements) {}
+
+    template <std::size_t N>
+    CollectionImpl(std::byte* buffer, std::array<int32_t, N> const& sizes)
+        : CollectionLeaf<Idx, T>(buffer, sizes),
+          CollectionImpl<Idx + 1, Args...>(CollectionLeaf<Idx, T>::layout_.metadata().nextByte(), sizes) {}
+  };
+
+  template <std::size_t Idx, typename T>
+  struct CollectionImpl<Idx, T> : public CollectionLeaf<Idx, T> {
+    CollectionImpl() = default;
+    CollectionImpl(std::byte* buffer, int32_t elements) : CollectionLeaf<Idx, T>(buffer, elements) {}
+
+    template <std::size_t N>
+    CollectionImpl(std::byte* buffer, std::array<int32_t, N> const& sizes) : CollectionLeaf<Idx, T>(buffer, sizes) {
+      static_assert(N == Idx + 1);
+    }
+  };
+
+  template <typename... Args>
+  struct Collections : public CollectionImpl<0, Args...> {};
+
+  // return the type at the Idx position in Args...
+  template <std::size_t Idx, typename... Args>
+  using TypeResolver = typename std::tuple_element<Idx, std::tuple<Args...>>::type;
+
+  // count how many times the type T occurs in Args...
+  template <typename T, typename... Args>
+  inline constexpr std::size_t typeCount = ((std::is_same<T, Args>::value ? 1 : 0) + ... + 0);
+
+  // count the non-void elements of Args...
+  template <typename... Args>
+  inline constexpr std::size_t membersCount = sizeof...(Args);
+
+  // if the type T occurs in Tuple, TupleTypeIndex has a static member value with the corresponding index;
+  // otherwise there is no such data  member.
+  template <typename T, typename Tuple>
+  struct TupleTypeIndex {};
+
+  template <typename T, typename... Args>
+  struct TupleTypeIndex<T, std::tuple<T, Args...>> {
+    static_assert(typeCount<T, Args...> == 0, "the requested type appears more than once among the arguments");
+    static constexpr std::size_t value = 0;
+  };
+
+  template <typename T, typename U, typename... Args>
+  struct TupleTypeIndex<T, std::tuple<U, Args...>> {
+    static_assert(not std::is_same_v<T, U>);
+    static_assert(typeCount<T, Args...> == 1, "the requested type does not appear among the arguments");
+    static constexpr std::size_t value = 1 + TupleTypeIndex<T, std::tuple<Args...>>::value;
+  };
+
+  // if the type T occurs in Args..., TypeIndex has a static member value with the corresponding index;
+  // otherwise there is no such data  member.
+  template <typename T, typename... Args>
+  using TypeIndex = TupleTypeIndex<T, std::tuple<Args...>>;
+
+  // return the index where the type T occurs in Args...
+  template <typename T, typename... Args>
+  inline constexpr std::size_t typeIndex = TypeIndex<T, Args...>::value;
+
+}  // namespace portablecollection
+
+#endif  // DataFormats_Portable_interface_PortableCollectionCommon_h