Pixel Alpaka Migration: Tracking [IV]

DataFormats/TrackSoA/BuildFile.xml

@@ -0,0 +1,12 @@
+<use name="alpaka"/>
+<use name="rootcore"/>
+<use name="eigen"/>
+<use name="DataFormats/Common"/>
+<use name="DataFormats/Portable"/>
+<use name="DataFormats/SoATemplate" source_only="1"/>
+<use name="DataFormats/TrackerCommon" source_only="1"/>
+<use name="HeterogeneousCore/AlpakaInterface"/>
+<flags ALPAKA_BACKENDS="!serial"/>
+<export>
+    <lib name="1"/>
+</export>

DataFormats/TrackSoA/README.md

@@ -0,0 +1,60 @@
+# TrackSoA Data Formats
+
+`DataFormat`s meant to be used on Host (CPU) or Device (GPU) for
+storing information about `TrackSoA`s created during the Pixel-local Reconstruction
+chain. It stores data in an SoA manner.  
+
+The host format is inheriting from `DataFormats/Portable/interface/PortableHostCollection.h`,
+while the device format is inheriting from `DataFormats/Portable/interface/PortableDeviceCollection.h`
+
+Both formats use the same SoA Layout (`TrackSoA::Layout`) which is generated
+via the `GENERATE_SOA_LAYOUT` macro in the `TrackDefinitions.h` file.
+
+## Notes
+
+-`hitIndices` and `detIndices`, instances of `HitContainer`, have been added into the
+layout as `SOA_SCALAR`s, meaning that they manage their own data independently from the SoA
+`Layout`. This could be improved in the future, if `HitContainer` (aka a `OneToManyAssoc` of fixed size)
+is replaced, but there don't seem to be any conflicts in including it in the `Layout` like this.
+- Host and Device classes should **not** be created via inheritance, as they're done here,
+but via composition. See [this discussion](https://github.com/cms-sw/cmssw/pull/40465#discussion_r1066039309).
+
+## TracksHost
+
+The version of the data format to be used for storing `TrackSoA` data on the CPU. 
+Instances of this class are to be used for:
+
+- Having a place to copy data to host from device, via `Memcpy`, or
+- Running host-side algorithms using data stored in an SoA manner.
+
+## TracksDevice
+
+The version of the data format to be used for storing `TrackSoA` data on the GPU.
+
+Instances of `TracksDevice` are to be created on host and be
+used on device only. To do so, the instance's `view()` method is to be called
+to pass a `View` to any kernel launched. Accessing data from the `view()` is not
+possible on the host side.
+
+## TracksSoACollection
+
+Depending on the Alpaka accelerator back-end enabled, `TrackSoACollection` is an alias to either the Host or Device SoA:
+
+```cpp
+template <typename TrackerTraits>
+  using TrackSoACollection = std::conditional_t<std::is_same_v<Device, alpaka::DevCpu>,
+                                                          TrackSoAHost<TrackerTraits>,
+                                                          TrackSoADevice<TrackerTraits, Device>>;
+```
+
+## Utilities
+
+`alpaka/TrackUtilities.h` contains a collection of methods which were originally
+defined as class methods inside either `TrackSoAHeterogeneousT` and `TrajectoryStateSoAT`
+which have been adapted to operate on `View` instances, so that they are callable
+from within `__global__` kernels, on both CPU and CPU. 
+
+## Use case
+
+See `test/TrackSoAHeterogeneous_test.cpp` for a simple example of instantiation,
+processing and copying from device to host.

DataFormats/TrackSoA/interface/TrackDefinitions.h

@@ -0,0 +1,32 @@
+#ifndef DataFormats_Track_interface_TrackDefinitions_h
+#define DataFormats_Track_interface_TrackDefinitions_h
+#include <string>
+#include <algorithm>
+#include <stdexcept>
+
+namespace pixelTrack {
+
+  enum class Quality : uint8_t { bad = 0, edup, dup, loose, strict, tight, highPurity, notQuality };
+  constexpr uint32_t qualitySize{uint8_t(Quality::notQuality)};
+  constexpr std::string_view qualityName[qualitySize]{"bad", "edup", "dup", "loose", "strict", "tight", "highPurity"};
+  inline Quality qualityByName(std::string_view name) {
+    auto qp = std::find(qualityName, qualityName + qualitySize, name) - qualityName;
+    auto ret = static_cast<Quality>(qp);
+
+    if (ret == pixelTrack::Quality::notQuality)
+      throw std::invalid_argument(std::string(name) + " is not a pixelTrack::Quality!");
+
+    return ret;
+  }
+
+#ifdef GPU_SMALL_EVENTS
+  // kept for testing and debugging
+  constexpr uint32_t maxNumber() { return 2 * 1024; }
+#else
+  // tested on MC events with 55-75 pileup events
+  constexpr uint32_t maxNumber() { return 32 * 1024; }
+#endif
+
+}  // namespace pixelTrack
+
+#endif

DataFormats/TrackSoA/interface/TracksDevice.h

@@ -0,0 +1,38 @@
+#ifndef DataFormats_Track_interface_TracksDevice_h
+#define DataFormats_Track_interface_TracksDevice_h
+
+#include <cstdint>
+#include <alpaka/alpaka.hpp>
+#include "DataFormats/TrackSoA/interface/TracksSoA.h"
+#include "DataFormats/TrackSoA/interface/TrackDefinitions.h"
+#include "DataFormats/Portable/interface/PortableDeviceCollection.h"
+
+// TODO: The class is created via inheritance of the PortableCollection.
+// This is generally discouraged, and should be done via composition.
+// See: https://github.com/cms-sw/cmssw/pull/40465#discussion_r1067364306
+template <typename TrackerTraits, typename TDev>
+class TracksDevice : public PortableDeviceCollection<reco::TrackLayout<TrackerTraits>, TDev> {
+public:
+  static constexpr int32_t S = TrackerTraits::maxNumberOfTuples;  //TODO: this could be made configurable at runtime
+  TracksDevice() = default;                                       // necessary for ROOT dictionaries
+
+  using PortableDeviceCollection<reco::TrackLayout<TrackerTraits>, TDev>::view;
+  using PortableDeviceCollection<reco::TrackLayout<TrackerTraits>, TDev>::const_view;
+  using PortableDeviceCollection<reco::TrackLayout<TrackerTraits>, TDev>::buffer;
+
+  // Constructor which specifies the SoA size
+  template <typename TQueue>
+  explicit TracksDevice<TrackerTraits, TDev>(TQueue& queue)
+      : PortableDeviceCollection<reco::TrackLayout<TrackerTraits>, TDev>(S, queue) {}
+};
+
+namespace pixelTrack {
+
+  template <typename TDev>
+  using TracksDevicePhase1 = TracksDevice<pixelTopology::Phase1, TDev>;
+  template <typename TDev>
+  using TracksDevicePhase2 = TracksDevice<pixelTopology::Phase2, TDev>;
+
+}  // namespace pixelTrack
+
+#endif  // DataFormats_Track_TracksDevice_H

DataFormats/TrackSoA/interface/TracksHost.h

@@ -0,0 +1,42 @@
+#ifndef DataFormats_Track_TracksHost_H
+#define DataFormats_Track_TracksHost_H
+
+#include <cstdint>
+#include <alpaka/alpaka.hpp>
+#include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
+#include "DataFormats/TrackSoA/interface/TracksSoA.h"
+#include "DataFormats/TrackSoA/interface/TrackDefinitions.h"
+#include "DataFormats/Portable/interface/PortableHostCollection.h"
+
+// TODO: The class is created via inheritance of the PortableHostCollection.
+// This is generally discouraged, and should be done via composition.
+// See: https://github.com/cms-sw/cmssw/pull/40465#discussion_r1067364306
+template <typename TrackerTraits>
+class TracksHost : public PortableHostCollection<reco::TrackLayout<TrackerTraits>> {
+public:
+  static constexpr int32_t S = TrackerTraits::maxNumberOfTuples;  //TODO: this could be made configurable at runtime
+  TracksHost() = default;  // Needed for the dictionary; not sure if line above is needed anymore
+
+  using PortableHostCollection<reco::TrackLayout<TrackerTraits>>::view;
+  using PortableHostCollection<reco::TrackLayout<TrackerTraits>>::const_view;
+  using PortableHostCollection<reco::TrackLayout<TrackerTraits>>::buffer;
+
+  // Constructor which specifies the SoA size
+  template <typename TQueue>
+  explicit TracksHost<TrackerTraits>(TQueue& queue)
+      : PortableHostCollection<reco::TrackLayout<TrackerTraits>>(S, queue) {}
+
+  // Constructor which specifies the DevHost
+  explicit TracksHost(alpaka_common::DevHost const& host)
+      : PortableHostCollection<reco::TrackLayout<TrackerTraits>>(S, host) {}
+};
+
+namespace pixelTrack {
+
+  using TracksHostPhase1 = TracksHost<pixelTopology::Phase1>;
+  using TracksHostPhase2 = TracksHost<pixelTopology::Phase2>;
+  using TracksHostHIonPhase1 = TracksHost<pixelTopology::HIonPhase1>;
+
+}  // namespace pixelTrack
+
+#endif  // DataFormats_Track_TracksHost_H

DataFormats/TrackSoA/interface/TracksSoA.h

@@ -0,0 +1,56 @@
+#ifndef DataFormats_Track_interface_TrackLayout_h
+#define DataFormats_Track_interface_TrackLayout_h
+
+#include <Eigen/Core>
+#include "HeterogeneousCore/AlpakaInterface/interface/OneToManyAssoc.h"
+#include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
+#include "DataFormats/SoATemplate/interface/SoALayout.h"
+#include "DataFormats/TrackSoA/interface/TrackDefinitions.h"
+
+namespace reco {
+
+  template <typename TrackerTraits>
+  struct TrackSoA {
+    static constexpr int32_t S = TrackerTraits::maxNumberOfTuples;
+    static constexpr int32_t H = TrackerTraits::avgHitsPerTrack;
+    // Aliases in order to not confuse the GENERATE_SOA_LAYOUT
+    // macro with weird colons and angled brackets.
+    using Vector5f = Eigen::Matrix<float, 5, 1>;
+    using Vector15f = Eigen::Matrix<float, 15, 1>;
+    using Quality = pixelTrack::Quality;
+
+    using hindex_type = uint32_t;
+
+    using HitContainer = cms::alpakatools::OneToManyAssocSequential<hindex_type, S + 1, H * S>;
+
+    GENERATE_SOA_LAYOUT(Layout,
+                        SOA_COLUMN(Quality, quality),
+                        SOA_COLUMN(float, chi2),
+                        SOA_COLUMN(int8_t, nLayers),
+                        SOA_COLUMN(float, eta),
+                        SOA_COLUMN(float, pt),
+                        SOA_EIGEN_COLUMN(Vector5f, state),
+                        SOA_EIGEN_COLUMN(Vector15f, covariance),
+                        SOA_SCALAR(int, nTracks),
+                        SOA_SCALAR(HitContainer, hitIndices),
+                        SOA_SCALAR(HitContainer, detIndices))
+  };
+
+  template <typename TrackerTraits>
+  using TrackLayout = typename reco::TrackSoA<TrackerTraits>::template Layout<>;
+  template <typename TrackerTraits>
+  using TrackSoAView = typename reco::TrackSoA<TrackerTraits>::template Layout<>::View;
+  template <typename TrackerTraits>
+  using TrackSoAConstView = typename reco::TrackSoA<TrackerTraits>::template Layout<>::ConstView;
+
+  template <typename TrackerTraits>
+  ALPAKA_FN_HOST_ACC ALPAKA_FN_INLINE static constexpr float charge(const TrackSoAConstView<TrackerTraits> &tracks,
+                                                                    int32_t i) {
+    //was: std::copysign(1.f, tracks[i].state()(2)). Will be constexpr with C++23
+    float v = tracks[i].state()(2);
+    return float((0.0f < v) - (v < 0.0f));
+  }
+
+}  // namespace reco
+
+#endif