Separated EndcapGeometry into host-only and buffers struct

SegmentLinking · Jun 24, 2024 · a3f65da · a3f65da
1 parent e0acfa5
commit a3f65da
Show file tree

Hide file tree

Showing 8 changed files with 117 additions and 102 deletions.
diff --git a/RecoTracker/LSTCore/interface/alpaka/EndcapGeometryBuffer.h b/RecoTracker/LSTCore/interface/alpaka/EndcapGeometryBuffer.h
@@ -0,0 +1,56 @@
+#ifndef EndcapGeometryBuffers_h
+#define EndcapGeometryBuffers_h
+
+#include <map>
+#include <iostream>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <vector>
+#include <stdexcept>
+
+#include "HeterogeneousCore/AlpakaInterface/interface/host.h"
+
+#include "RecoTracker/LSTCore/interface/alpaka/Constants.h"
+
+namespace SDL {
+
+  struct endcapGeom {
+    const unsigned int* geoMapDetId;
+    const float* geoMapPhi;
+
+    template <typename TBuff>
+    void setData(const TBuff& endcapgeombuf) {
+      geoMapDetId = alpaka::getPtrNative(endcapgeombuf.geoMapDetId_buf);
+      geoMapPhi = alpaka::getPtrNative(endcapgeombuf.geoMapPhi_buf);
+    }
+  };
+
+  template <typename TDev>
+  struct endcapGeometryBuffer : endcapGeom {
+    Buf<TDev, unsigned int> geoMapDetId_buf;
+    Buf<TDev, float> geoMapPhi_buf;
+
+    endcapGeometryBuffer(TDev const& dev, unsigned int nEndCapMap)
+        : geoMapDetId_buf(allocBufWrapper<unsigned int>(dev, nEndCapMap)),
+          geoMapPhi_buf(allocBufWrapper<float>(dev, nEndCapMap)) {
+      setData(*this);
+    }
+
+    template <typename TQueue, typename TDevSrc>
+    inline void copyFromSrc(TQueue queue, const endcapGeometryBuffer<TDevSrc>& src) {
+      alpaka::memcpy(queue, geoMapDetId_buf, src.geoMapDetId_buf);
+      alpaka::memcpy(queue, geoMapPhi_buf, src.geoMapPhi_buf);
+    }
+
+    template <typename TQueue>
+    endcapGeometryBuffer(TQueue queue, const endcapGeometryBuffer<alpaka::DevCpu>& src, unsigned int nEndCapMap) {
+      copyFromSrc(queue, src);
+    }
+
+    inline SDL::endcapGeom const* data() const { return this; }
+  };
+
+}  // namespace SDL
+
+#endif
diff --git a/RecoTracker/LSTCore/interface/alpaka/LSTESData.h b/RecoTracker/LSTCore/interface/alpaka/LSTESData.h
@@ -1,11 +1,9 @@
 #ifndef LSTESData_H
 #define LSTESData_H
 
-#ifdef LST_IS_CMSSW_PACKAGE
 #include "RecoTracker/LSTCore/interface/alpaka/Constants.h"
-#else
-#include "Constants.h"
-#endif
+#include "RecoTracker/LSTCore/interface/alpaka/EndcapGeometryBuffer.h"
+#include "RecoTracker/LSTCore/interface/alpaka/Module.h"
 
 #include "HeterogeneousCore/AlpakaInterface/interface/CopyToDevice.h"
 
@@ -16,32 +14,29 @@ namespace SDL {
 
   struct pixelMap;
 
-  template <typename>
-  struct modulesBuffer;
-
-  template <typename>
-  class EndcapGeometry;
-
   template <typename TDev>
   struct LSTESData {
     uint16_t nModules;
     uint16_t nLowerModules;
     unsigned int nPixels;
+    unsigned int nEndCapMap;
     std::shared_ptr<const modulesBuffer<TDev>> modulesBuffers;
-    std::shared_ptr<const EndcapGeometry<TDev>> endcapGeometry;
+    std::shared_ptr<const endcapGeometryBuffer<TDev>> endcapGeometryBuffers;
     std::shared_ptr<const pixelMap> pixelMapping;
 
     LSTESData(uint16_t const& nModulesIn,
               uint16_t const& nLowerModulesIn,
               unsigned int const& nPixelsIn,
+              unsigned int const& nEndCapMapIn,
               std::shared_ptr<const modulesBuffer<TDev>> const& modulesBuffersIn,
-              std::shared_ptr<const EndcapGeometry<TDev>> const& endcapGeometryIn,
+              std::shared_ptr<const endcapGeometryBuffer<TDev>> const& endcapGeometryBuffersIn,
               std::shared_ptr<const pixelMap> const& pixelMappingIn)
         : nModules(nModulesIn),
           nLowerModules(nLowerModulesIn),
           nPixels(nPixelsIn),
+          nEndCapMap(nEndCapMapIn),
           modulesBuffers(modulesBuffersIn),
-          endcapGeometry(endcapGeometryIn),
+          endcapGeometryBuffers(endcapGeometryBuffersIn),
           pixelMapping(pixelMappingIn) {}
   };
 
@@ -53,7 +48,23 @@ namespace cms::alpakatools {
   template <>
   struct CopyToDevice<SDL::LSTESData<SDL::DevHost>> {
     template <typename TQueue>
-    static SDL::LSTESData<alpaka::Dev<TQueue>> copyAsync(TQueue& queue, SDL::LSTESData<SDL::DevHost> const& srcData);
+    static SDL::LSTESData<alpaka::Dev<TQueue>> copyAsync(
+        TQueue& queue, SDL::LSTESData<SDL::DevHost> const& srcData) {
+      auto deviceModulesBuffers = std::make_shared<SDL::modulesBuffer<alpaka::Dev<TQueue>>>(
+          alpaka::getDev(queue), srcData.nModules, srcData.nPixels);
+      deviceModulesBuffers->copyFromSrc(queue, *srcData.modulesBuffers);
+      auto deviceEndcapGeometryBuffers =
+          std::make_shared<SDL::endcapGeometryBuffer<alpaka::Dev<TQueue>>>(alpaka::getDev(queue), srcData.nEndCapMap);
+      deviceEndcapGeometryBuffers->copyFromSrc(queue, *srcData.endcapGeometryBuffers);
+
+      return SDL::LSTESData<alpaka::Dev<TQueue>>(srcData.nModules,
+                                                srcData.nLowerModules,
+                                                srcData.nPixels,
+                                                srcData.nEndCapMap,
+                                                deviceModulesBuffers,
+                                                deviceEndcapGeometryBuffers,
+                                                srcData.pixelMapping);
+    }
   };
 }  // namespace cms::alpakatools
 

diff --git a/RecoTracker/LSTCore/src/alpaka/EndcapGeometry.dev.cc b/RecoTracker/LSTCore/src/alpaka/EndcapGeometry.dev.cc
@@ -1,23 +1,8 @@
 #include "EndcapGeometry.h"
 
-template <typename TDev>
-SDL::EndcapGeometry<TDev>::EndcapGeometry(TDev const& devAccIn)
-    : geoMapDetId_buf(allocBufWrapper<unsigned int>(devAccIn, 0)), geoMapPhi_buf(allocBufWrapper<float>(devAccIn, 0)) {}
+SDL::EndcapGeometry<SDL::Dev>::EndcapGeometry(std::string filename) { load(filename); }
 
-template <typename TDev>
-template <typename TQueue>
-SDL::EndcapGeometry<TDev>::EndcapGeometry(TQueue& queue, SDL::EndcapGeometry<SDL::DevHost> const& endcapGeometrySrc)
-    : dxdy_slope_(endcapGeometrySrc.dxdy_slope_),
-      centroid_phis_(endcapGeometrySrc.centroid_phis_),
-      geoMapDetId_buf(allocBufWrapper<unsigned int>(alpaka::getDev(queue), endcapGeometrySrc.nEndCapMap)),
-      geoMapPhi_buf(allocBufWrapper<float>(alpaka::getDev(queue), endcapGeometrySrc.nEndCapMap)),
-      nEndCapMap(endcapGeometrySrc.nEndCapMap) {
-  alpaka::memcpy(queue, geoMapPhi_buf, endcapGeometrySrc.geoMapPhi_buf);
-  alpaka::memcpy(queue, geoMapDetId_buf, endcapGeometrySrc.geoMapDetId_buf);
-}
-
-template <typename TDev>
-void SDL::EndcapGeometry<TDev>::load(std::string filename) {
+void SDL::EndcapGeometry<SDL::Dev>::load(std::string filename) {
   dxdy_slope_.clear();
   centroid_phis_.clear();
 
@@ -46,44 +31,27 @@ void SDL::EndcapGeometry<TDev>::load(std::string filename) {
     }
   }
 
-  fillGeoMapArraysExplicitHost();
+  fillGeoMapArraysExplicit();
 }
 
-template <typename TDev>
-void SDL::EndcapGeometry<TDev>::fillGeoMapArraysExplicitHost() {
-  unsigned int phi_size = centroid_phis_.size();
-
-  // Allocate buffers on host
-  SDL::DevHost const& devHost = cms::alpakatools::host();
-  geoMapPhi_buf = allocBufWrapper<float>(devHost, phi_size);
-  geoMapDetId_buf = allocBufWrapper<unsigned int>(devHost, phi_size);
+void SDL::EndcapGeometry<SDL::Dev>::fillGeoMapArraysExplicit() {
+  nEndCapMap = centroid_phis_.size();
 
-  // Access the raw pointers of the buffers
-  float* mapPhi = alpaka::getPtrNative(geoMapPhi_buf);
-  unsigned int* mapDetId = alpaka::getPtrNative(geoMapDetId_buf);
+  geoMapDetId_buf.reserve(nEndCapMap);
+  geoMapPhi_buf.reserve(nEndCapMap);
 
-  unsigned int counter = 0;
   for (auto it = centroid_phis_.begin(); it != centroid_phis_.end(); ++it) {
     unsigned int detId = it->first;
     float Phi = it->second;
-    mapPhi[counter] = Phi;
-    mapDetId[counter] = detId;
-    counter++;
+    geoMapPhi_buf.push_back(Phi);
+    geoMapDetId_buf.push_back(detId);
   }
-
-  nEndCapMap = counter;
 }
 
-template <typename TDev>
-float SDL::EndcapGeometry<TDev>::getdxdy_slope(unsigned int detid) const {
+float SDL::EndcapGeometry<SDL::Dev>::getdxdy_slope(unsigned int detid) const {
   if (dxdy_slope_.find(detid) != dxdy_slope_.end()) {
     return dxdy_slope_.at(detid);
   } else {
     return 0;
   }
 }
-
-// We need to instantiate these to ensure that the compiler generates them
-template class SDL::EndcapGeometry<SDL::DevHost>;
-template SDL::EndcapGeometry<SDL::Dev>::EndcapGeometry(SDL::QueueAcc& queue,
-                                                       SDL::EndcapGeometry<SDL::DevHost> const& endcapGeometrySrc);
diff --git a/RecoTracker/LSTCore/src/alpaka/EndcapGeometry.h b/RecoTracker/LSTCore/src/alpaka/EndcapGeometry.h
@@ -15,33 +15,27 @@
 #include "Constants.h"
 #endif
 
-#include "HeterogeneousCore/AlpakaInterface/interface/host.h"
-
 namespace SDL {
-
-  template <typename TDev>
-  class EndcapGeometry {
+  template <typename>
+  class EndcapGeometry;
+  template <>
+  class EndcapGeometry<SDL::Dev> {
   private:
     std::map<unsigned int, float> dxdy_slope_;     // dx/dy slope
     std::map<unsigned int, float> centroid_phis_;  // centroid phi
 
-    // Friend all other instantiations of this template
-    template <typename OtherTDev>
-    friend class EndcapGeometry;
-
   public:
-    Buf<TDev, unsigned int> geoMapDetId_buf;
-    Buf<TDev, float> geoMapPhi_buf;
+    std::vector<unsigned int> geoMapDetId_buf;
+    std::vector<float> geoMapPhi_buf;
 
     unsigned int nEndCapMap;
 
-    EndcapGeometry(TDev const& devAccIn);
-    template <typename TQueue>
-    EndcapGeometry(TQueue& queue, EndcapGeometry<DevHost> const& endcapGeometrySrc);
+    EndcapGeometry() = default;
+    EndcapGeometry(std::string filename);
     ~EndcapGeometry() = default;
 
     void load(std::string);
-    void fillGeoMapArraysExplicitHost();
+    void fillGeoMapArraysExplicit();
     float getdxdy_slope(unsigned int detid) const;
   };
 }  // namespace SDL

diff --git a/RecoTracker/LSTCore/src/alpaka/Event.dev.cc b/RecoTracker/LSTCore/src/alpaka/Event.dev.cc
@@ -191,9 +191,9 @@ void SDL::Event<SDL::Acc>::addHitToEvent(std::vector<float> x,
                                                          Endcap,
                                                          TwoS,
                                                          nModules_,
-                                                         endcapGeometry_->nEndCapMap,
-                                                         alpaka::getPtrNative(endcapGeometry_->geoMapDetId_buf),
-                                                         alpaka::getPtrNative(endcapGeometry_->geoMapPhi_buf),
+                                                         nEndCapMap_,
+                                                         alpaka::getPtrNative(endcapGeometryBuffers_->geoMapDetId_buf),
+                                                         alpaka::getPtrNative(endcapGeometryBuffers_->geoMapPhi_buf),
                                                          *modulesBuffers_->data(),
                                                          *hitsInGPU,
                                                          nHits));

diff --git a/RecoTracker/LSTCore/src/alpaka/Event.h b/RecoTracker/LSTCore/src/alpaka/Event.h
@@ -90,9 +90,10 @@ namespace SDL {
     const uint16_t nModules_;
     const uint16_t nLowerModules_;
     const unsigned int nPixels_;
+    const unsigned int nEndCapMap_;
     const std::shared_ptr<const modulesBuffer<Dev>> modulesBuffers_;
     const std::shared_ptr<const pixelMap> pixelMapping_;
-    const std::shared_ptr<const EndcapGeometry<Dev>> endcapGeometry_;
+    const std::shared_ptr<const endcapGeometryBuffer<Dev>> endcapGeometryBuffers_;
 
   public:
     // Constructor used for CMSSW integration. Uses an external queue.
@@ -104,9 +105,10 @@ namespace SDL {
           nModules_(deviceESData->nModules),
           nLowerModules_(deviceESData->nLowerModules),
           nPixels_(deviceESData->nPixels),
+          nEndCapMap_(deviceESData->nEndCapMap),
           modulesBuffers_(deviceESData->modulesBuffers),
           pixelMapping_(deviceESData->pixelMapping),
-          endcapGeometry_(deviceESData->endcapGeometry) {
+          endcapGeometryBuffers_(deviceESData->endcapGeometryBuffers) {
       init(verbose);
     }
     void resetEvent();

diff --git a/RecoTracker/LSTCore/src/alpaka/LSTESData.dev.cc b/RecoTracker/LSTCore/src/alpaka/LSTESData.dev.cc
@@ -44,7 +44,7 @@ namespace {
   }
 
   void loadMapsHost(SDL::MapPLStoLayer& pLStoLayer,
-                    std::shared_ptr<SDL::EndcapGeometry<SDL::DevHost>> endcapGeometry,
+                    std::shared_ptr<SDL::EndcapGeometry<SDL::Dev>> endcapGeometry,
                     std::shared_ptr<SDL::TiltedGeometry<SDL::Dev>> tiltedGeometry,
                     std::shared_ptr<SDL::ModuleConnectionMap<SDL::Dev>> moduleConnectionMap) {
     // Module orientation information (DrDz or phi angles)
@@ -87,12 +87,18 @@ std::unique_ptr<SDL::LSTESData<SDL::DevHost>> SDL::loadAndFillESHost() {
   unsigned int nPixels;
   std::shared_ptr<SDL::modulesBuffer<SDL::DevHost>> modulesBuffers = nullptr;
   auto pLStoLayer = std::make_shared<SDL::MapPLStoLayer>();
-  auto endcapGeometry = std::make_shared<SDL::EndcapGeometry<SDL::DevHost>>(cms::alpakatools::host());
+  auto endcapGeometry = std::make_shared<SDL::EndcapGeometry<SDL::Dev>>();
   auto tiltedGeometry = std::make_shared<SDL::TiltedGeometry<SDL::Dev>>();
   auto pixelMapping = std::make_shared<SDL::pixelMap>();
   auto moduleConnectionMap = std::make_shared<SDL::ModuleConnectionMap<SDL::Dev>>();
   ::loadMapsHost(*pLStoLayer, endcapGeometry, tiltedGeometry, moduleConnectionMap);
 
+  auto endcapGeometryBuffers = std::make_shared<SDL::endcapGeometryBuffer<SDL::DevHost>>(cms::alpakatools::host(),
+                                                                                     endcapGeometry->nEndCapMap);
+  alpaka::QueueCpuBlocking queue(cms::alpakatools::host());
+  alpaka::memcpy(queue, endcapGeometryBuffers->geoMapDetId_buf, endcapGeometry->geoMapDetId_buf, endcapGeometry->nEndCapMap);
+  alpaka::memcpy(queue, endcapGeometryBuffers->geoMapPhi_buf, endcapGeometry->geoMapPhi_buf, endcapGeometry->nEndCapMap);
+
   auto path =
       get_absolute_path_after_check_file_exists(trackLooperDir() + "/data/OT800_IT615_pt0.8/sensor_centroids.bin");
   SDL::loadModulesFromFile(pLStoLayer.get(),
@@ -106,27 +112,5 @@ std::unique_ptr<SDL::LSTESData<SDL::DevHost>> SDL::loadAndFillESHost() {
                            tiltedGeometry.get(),
                            moduleConnectionMap.get());
   return std::make_unique<LSTESData<SDL::DevHost>>(
-      nModules, nLowerModules, nPixels, modulesBuffers, endcapGeometry, pixelMapping);
+      nModules, nLowerModules, nPixels, endcapGeometry->nEndCapMap, modulesBuffers, endcapGeometryBuffers, pixelMapping);
 }
-
-template <typename TQueue>
-SDL::LSTESData<alpaka::Dev<TQueue>> cms::alpakatools::CopyToDevice<SDL::LSTESData<SDL::DevHost>>::copyAsync(
-    TQueue& queue, SDL::LSTESData<SDL::DevHost> const& srcData) {
-  auto deviceModulesBuffers = std::make_shared<SDL::modulesBuffer<alpaka::Dev<TQueue>>>(
-      alpaka::getDev(queue), srcData.nModules, srcData.nPixels);
-  deviceModulesBuffers->copyFromSrc(queue, *srcData.modulesBuffers);
-  auto deviceEndcapGeometry =
-      std::make_shared<SDL::EndcapGeometry<alpaka::Dev<TQueue>>>(queue, *srcData.endcapGeometry);
-
-  return SDL::LSTESData<alpaka::Dev<TQueue>>(srcData.nModules,
-                                             srcData.nLowerModules,
-                                             srcData.nPixels,
-                                             deviceModulesBuffers,
-                                             deviceEndcapGeometry,
-                                             srcData.pixelMapping);
-}
-
-// Make sure it is compiled
-template struct cms::alpakatools::CopyToDevice<SDL::LSTESData<SDL::DevHost>>;
-template SDL::LSTESData<alpaka::Dev<SDL::QueueAcc>> cms::alpakatools::CopyToDevice<
-    SDL::LSTESData<SDL::DevHost>>::copyAsync<SDL::QueueAcc>(SDL::QueueAcc&, SDL::LSTESData<SDL::DevHost> const&);
diff --git a/RecoTracker/LSTCore/src/alpaka/ModuleMethods.h b/RecoTracker/LSTCore/src/alpaka/ModuleMethods.h
@@ -231,7 +231,7 @@ namespace SDL {
                                   unsigned int& nPixels,
                                   std::shared_ptr<modulesBuffer<DevHost>>& modulesBuf,
                                   pixelMap* pixelMapping,
-                                  const EndcapGeometry<DevHost>* endcapGeometry,
+                                  const EndcapGeometry<Dev>* endcapGeometry,
                                   const TiltedGeometry<Dev>* tiltedGeometry,
                                   const ModuleConnectionMap<Dev>* moduleConnectionMap) {
     ModuleMetaData mmd;