migrate to SoATemplate

RecoTracker/LST/plugins/alpaka/LSTModulesDevESProducer.cc

@@ -1,5 +1,4 @@
 // LST includes
-#include "RecoTracker/LSTCore/interface/Module.h"
 #include "RecoTracker/LSTCore/interface/alpaka/LST.h"
 
 #include "FWCore/ParameterSet/interface/ParameterSet.h"

RecoTracker/LSTCore/interface/Constants.h

@@ -2,6 +2,15 @@
 #define RecoTracker_LSTCore_interface_Constants_h
 
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
+#include "DataFormats/Common/interface/StdArray.h"
+
+#if defined(FP16_Base)
+#if defined ALPAKA_ACC_GPU_CUDA_ENABLED
+#include <cuda_fp16.h>
+#elif defined ALPAKA_ACC_GPU_HIP_ENABLED
+#include <hip/hip_fp16.h>
+#endif
+#endif
 
 #ifdef CACHE_ALLOC
 #include "HeterogeneousCore/AlpakaInterface/interface/CachedBufAlloc.h"
@@ -55,27 +64,68 @@ namespace lst {
 
   constexpr unsigned int size_superbins = 45000;
 
+// Half precision wrapper functions.
+#if defined(FP16_Base)
+#define __F2H __float2half
+#define __H2F __half2float
+  typedef __half float FPX;
+#else
+#define __F2H
+#define __H2F
+  typedef float FPX;
+#endif
+
+// Needed for files that are compiled by g++ to not throw an error.
+// uint4 is defined only for CUDA, so we will have to revisit this soon when running on other backends.
+#if !defined(ALPAKA_ACC_GPU_CUDA_ENABLED) && !defined(ALPAKA_ACC_GPU_HIP_ENABLED)
+  struct uint4 {
+    unsigned int x;
+    unsigned int y;
+    unsigned int z;
+    unsigned int w;
+  };
+#endif
+
   // Defining the constant host device variables right up here
   // Currently pixel tracks treated as LSs with 2 double layers (IT layers 1+2 and 3+4) and 4 hits. To be potentially handled better in the future.
+  struct Params_Modules {
+    using ArrayU16xMaxConnected = edm::StdArray<uint16_t, max_connected_modules>;
+  };
   struct Params_pLS {
     static constexpr int kLayers = 2, kHits = 4;
   };
   struct Params_LS {
     static constexpr int kLayers = 2, kHits = 4;
+    using ArrayUxLayers = edm::StdArray<unsigned int, kLayers>;
   };
   struct Params_T3 {
     static constexpr int kLayers = 3, kHits = 6;
+    using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
+    using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
+    using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
   };
   struct Params_pT3 {
     static constexpr int kLayers = 5, kHits = 10;
+    using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
+    using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
+    using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
   };
   struct Params_T5 {
     static constexpr int kLayers = 5, kHits = 10;
+    using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
+    using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
+    using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
   };
   struct Params_pT5 {
     static constexpr int kLayers = 7, kHits = 14;
+    using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
+    using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
+    using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
   };
 
+  using ArrayIx2 = edm::StdArray<int, 2>;
+  using ArrayUx2 = edm::StdArray<unsigned int, 2>;
+
 }  //namespace lst
 
 #endif

RecoTracker/LSTCore/interface/EndcapGeometryDevHostCollection.h

@@ -0,0 +1,10 @@
+#ifndef RecoTracker_LSTCore_interface_EndcapGeometryDevHostCollection_h
+#define RecoTracker_LSTCore_interface_EndcapGeometryDevHostCollection_h
+
+#include "RecoTracker/LSTCore/interface/EndcapGeometryDevSoA.h"
+#include "DataFormats/Portable/interface/PortableHostCollection.h"
+
+namespace lst {
+  using EndcapGeometryDevHostCollection = PortableHostCollection<EndcapGeometryDevSoA>;
+}  // namespace lst
+#endif

RecoTracker/LSTCore/interface/EndcapGeometryDevSoA.h

@@ -0,0 +1,18 @@
+#ifndef RecoTracker_LSTCore_interface_EndcapGeometryDevSoA_h
+#define RecoTracker_LSTCore_interface_EndcapGeometryDevSoA_h
+
+#include "DataFormats/SoATemplate/interface/SoALayout.h"
+#include "DataFormats/Portable/interface/PortableCollection.h"
+
+namespace lst {
+
+  GENERATE_SOA_LAYOUT(EndcapGeometryDevSoALayout, SOA_COLUMN(unsigned int, geoMapDetId), SOA_COLUMN(float, geoMapPhi))
+
+  using EndcapGeometryDevSoA = EndcapGeometryDevSoALayout<>;
+
+  using EndcapGeometryDev = EndcapGeometryDevSoA::View;
+  using EndcapGeometryDevConst = EndcapGeometryDevSoA::ConstView;
+
+}  // namespace lst
+
+#endif

RecoTracker/LSTCore/interface/HitsHostCollection.h

@@ -0,0 +1,10 @@
+#ifndef RecoTracker_LSTCore_interface_HitsHostCollection_h
+#define RecoTracker_LSTCore_interface_HitsHostCollection_h
+
+#include "RecoTracker/LSTCore/interface/HitsSoA.h"
+#include "DataFormats/Portable/interface/PortableHostCollection.h"
+
+namespace lst {
+  using HitsHostCollection = PortableHostMultiCollection<HitsSoA, HitsRangesSoA>;
+}  // namespace lst
+#endif

RecoTracker/LSTCore/interface/HitsSoA.h

@@ -0,0 +1,43 @@
+#ifndef RecoTracker_LSTCore_interface_HitsSoA_h
+#define RecoTracker_LSTCore_interface_HitsSoA_h
+
+#include "DataFormats/SoATemplate/interface/SoALayout.h"
+#include "DataFormats/Portable/interface/PortableCollection.h"
+
+#include "RecoTracker/LSTCore/interface/Constants.h"
+
+namespace lst {
+
+  GENERATE_SOA_LAYOUT(HitsSoALayout,
+                      SOA_COLUMN(float, xs),
+                      SOA_COLUMN(float, ys),
+                      SOA_COLUMN(float, zs),
+                      SOA_COLUMN(uint16_t, moduleIndices),
+                      SOA_COLUMN(unsigned int, idxs),
+                      SOA_COLUMN(unsigned int, detid),
+                      SOA_COLUMN(float, rts),
+                      SOA_COLUMN(float, phis),
+                      SOA_COLUMN(float, etas),
+                      SOA_COLUMN(float, highEdgeXs),
+                      SOA_COLUMN(float, highEdgeYs),
+                      SOA_COLUMN(float, lowEdgeXs),
+                      SOA_COLUMN(float, lowEdgeYs))
+
+  GENERATE_SOA_LAYOUT(HitsRangesSoALayout,
+                      SOA_COLUMN(ArrayIx2, hitRanges),
+                      SOA_COLUMN(int, hitRangesLower),
+                      SOA_COLUMN(int, hitRangesUpper),
+                      SOA_COLUMN(int8_t, hitRangesnLower),
+                      SOA_COLUMN(int8_t, hitRangesnUpper))
+
+  using HitsSoA = HitsSoALayout<>;
+  using HitsRangesSoA = HitsRangesSoALayout<>;
+
+  using Hits = HitsSoA::View;
+  using HitsConst = HitsSoA::ConstView;
+  using HitsRanges = HitsRangesSoA::View;
+  using HitsRangesConst = HitsRangesSoA::ConstView;
+
+}  // namespace lst
+
+#endif

RecoTracker/LSTCore/interface/LSTESData.h

@@ -2,8 +2,8 @@
 #define RecoTracker_LSTCore_interface_LSTESData_h
 
 #include "RecoTracker/LSTCore/interface/Constants.h"
-#include "RecoTracker/LSTCore/interface/EndcapGeometryBuffer.h"
-#include "RecoTracker/LSTCore/interface/Module.h"
+#include "RecoTracker/LSTCore/interface/EndcapGeometryDevHostCollection.h"
+#include "RecoTracker/LSTCore/interface/ModulesHostCollection.h"
 #include "RecoTracker/LSTCore/interface/PixelMap.h"
 
 #include "HeterogeneousCore/AlpakaInterface/interface/CopyToDevice.h"
@@ -19,23 +19,23 @@ namespace lst {
     uint16_t nLowerModules;
     unsigned int nPixels;
     unsigned int nEndCapMap;
-    ModulesBuffer<TDev> modulesBuffers;
-    EndcapGeometryBuffer<TDev> endcapGeometryBuffers;
+    std::unique_ptr<const PortableMultiCollection<TDev, ModulesSoA, ModulesPixelSoA>> modules;
+    std::unique_ptr<const PortableCollection<EndcapGeometryDevSoA, TDev>> endcapGeometry;
     std::shared_ptr<const PixelMap> pixelMapping;
 
     LSTESData(uint16_t const& nModulesIn,
               uint16_t const& nLowerModulesIn,
               unsigned int const& nPixelsIn,
               unsigned int const& nEndCapMapIn,
-              ModulesBuffer<TDev> const& modulesBuffersIn,
-              EndcapGeometryBuffer<TDev> const& endcapGeometryBuffersIn,
+              std::unique_ptr<const PortableMultiCollection<TDev, ModulesSoA, ModulesPixelSoA>> modulesIn,
+              std::unique_ptr<const PortableCollection<EndcapGeometryDevSoA, TDev>> endcapGeometryIn,
               std::shared_ptr<const PixelMap> const& pixelMappingIn)
         : nModules(nModulesIn),
           nLowerModules(nLowerModulesIn),
           nPixels(nPixelsIn),
           nEndCapMap(nEndCapMapIn),
-          modulesBuffers(modulesBuffersIn),
-          endcapGeometryBuffers(endcapGeometryBuffersIn),
+          modules(std::move(modulesIn)),
+          endcapGeometry(std::move(endcapGeometryIn)),
           pixelMapping(pixelMappingIn) {}
   };
 
@@ -44,24 +44,49 @@ namespace lst {
 }  // namespace lst
 
 namespace cms::alpakatools {
+
+  // The templated definition in CMSSW doesn't work when using CPU as the device
+  template <>
+  struct CopyToDevice<PortableHostCollection<lst::EndcapGeometryDevSoA>> {
+    template <typename TQueue>
+    static auto copyAsync(TQueue& queue, PortableHostCollection<lst::EndcapGeometryDevSoA> const& srcData) {
+      using TDevice = typename alpaka::trait::DevType<TQueue>::type;
+      PortableCollection<lst::EndcapGeometryDevSoA, TDevice> dstData(srcData->metadata().size(), queue);
+      alpaka::memcpy(queue, dstData.buffer(), srcData.buffer());
+      return dstData;
+    }
+  };
+
+  template <>
+  struct CopyToDevice<PortableHostMultiCollection<lst::ModulesSoA, lst::ModulesPixelSoA>> {
+    template <typename TQueue>
+    static auto copyAsync(TQueue& queue,
+                          PortableHostMultiCollection<lst::ModulesSoA, lst::ModulesPixelSoA> const& srcData) {
+      using TDevice = typename alpaka::trait::DevType<TQueue>::type;
+      PortableMultiCollection<TDevice, lst::ModulesSoA, lst::ModulesPixelSoA> dstData(srcData.sizes(), queue);
+      alpaka::memcpy(queue, dstData.buffer(), srcData.buffer());
+      return dstData;
+    }
+  };
+
   template <>
   struct CopyToDevice<lst::LSTESData<alpaka_common::DevHost>> {
     template <typename TQueue>
     static lst::LSTESData<alpaka::Dev<TQueue>> copyAsync(TQueue& queue,
                                                          lst::LSTESData<alpaka_common::DevHost> const& srcData) {
-      auto deviceModulesBuffers =
-          lst::ModulesBuffer<alpaka::Dev<TQueue>>(alpaka::getDev(queue), srcData.nModules, srcData.nPixels);
-      deviceModulesBuffers.copyFromSrc(queue, srcData.modulesBuffers);
-      auto deviceEndcapGeometryBuffers =
-          lst::EndcapGeometryBuffer<alpaka::Dev<TQueue>>(alpaka::getDev(queue), srcData.nEndCapMap);
-      deviceEndcapGeometryBuffers.copyFromSrc(queue, srcData.endcapGeometryBuffers);
+      auto deviceModules =
+          std::make_unique<PortableMultiCollection<alpaka::Dev<TQueue>, lst::ModulesSoA, lst::ModulesPixelSoA>>(
+              CopyToDevice<PortableHostMultiCollection<lst::ModulesSoA, lst::ModulesPixelSoA>>::copyAsync(
+                  queue, *srcData.modules));
+      auto deviceEndcapGeometry = std::make_unique<PortableCollection<lst::EndcapGeometryDevSoA, alpaka::Dev<TQueue>>>(
+          CopyToDevice<PortableHostCollection<lst::EndcapGeometryDevSoA>>::copyAsync(queue, *srcData.endcapGeometry));
 
       return lst::LSTESData<alpaka::Dev<TQueue>>(srcData.nModules,
                                                  srcData.nLowerModules,
                                                  srcData.nPixels,
                                                  srcData.nEndCapMap,
-                                                 std::move(deviceModulesBuffers),
-                                                 std::move(deviceEndcapGeometryBuffers),
+                                                 std::move(deviceModules),
+                                                 std::move(deviceEndcapGeometry),
                                                  srcData.pixelMapping);
     }
   };

RecoTracker/LSTCore/interface/MiniDoubletsHostCollection.h

@@ -0,0 +1,10 @@
+#ifndef RecoTracker_LSTCore_interface_MiniDoubletsHostCollection_h
+#define RecoTracker_LSTCore_interface_MiniDoubletsHostCollection_h
+
+#include "RecoTracker/LSTCore/interface/MiniDoubletsSoA.h"
+#include "DataFormats/Portable/interface/PortableHostCollection.h"
+
+namespace lst {
+  using MiniDoubletsHostCollection = PortableHostMultiCollection<MiniDoubletsSoA, MiniDoubletsOccupancySoA>;
+}  // namespace lst
+#endif