Merge pull request #2 from cericeci/CMSSW_12_6_0_pre5_dev_forJosh

GPU->RecoVertex conversion, basis for the new fitter algo
cericeci · Feb 28, 2023 · a098464 · a098464
2 parents f0b75f9 + 655ad21
commit a098464
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 6 deletions.
diff --git a/CUDADataFormats/Track/interface/TrackForPVSoAT.h b/CUDADataFormats/Track/interface/TrackForPVSoAT.h
@@ -54,6 +54,8 @@ class TrackForPVSoAHeterogeneousT {
   eigenSoA::ScalarSoA<double, S> px;
   eigenSoA::ScalarSoA<double, S> py;
   eigenSoA::ScalarSoA<double, S> pz;
+  eigenSoA::ScalarSoA<double, S> x;
+  eigenSoA::ScalarSoA<double, S> y;
   eigenSoA::ScalarSoA<double, S> dxError;
   eigenSoA::ScalarSoA<double, S> dyError;
   eigenSoA::ScalarSoA<double, S> dzError;
@@ -91,7 +93,16 @@ class VertexForPVSoAHeterogeneousT {
   eigenSoA::ScalarSoA<double, S> z;
   eigenSoA::ScalarSoA<double, S> x;
   eigenSoA::ScalarSoA<double, S> y;
+  eigenSoA::ScalarSoA<double, S> t; // Without time for the moment
+  eigenSoA::ScalarSoA<double, S> chi2;
+  eigenSoA::ScalarSoA<unsigned int, S> ndof;
+  eigenSoA::ScalarSoA<double, S> errz;
+  eigenSoA::ScalarSoA<double, S> errx;
+  eigenSoA::ScalarSoA<double, S> erry;  
   eigenSoA::ScalarSoA<unsigned int, S> ntracks;
+  // The track-vertex association matrices
+  eigenSoA::MatrixSoA<Vector512d, S> track_id;
+  eigenSoA::MatrixSoA<Vector512d, S> track_weight;
   eigenSoA::ScalarSoA<double, S> rho;
   // Auxiliar vectors
   eigenSoA::ScalarSoA<double, S> aux1;

diff --git a/RecoVertex/PrimaryVertexProducer/plugins/PrimaryVertexProducerCUDA.cc b/RecoVertex/PrimaryVertexProducer/plugins/PrimaryVertexProducerCUDA.cc
@@ -182,13 +182,13 @@ void PrimaryVertexProducerCUDA::produce(edm::Event& iEvent, const edm::EventSetu
     edm::LogError("UnusableBeamSpot") << "No beam spot available from EventSetup";
   }
 
-  bool validBS = true;
+  /*bool validBS = true;
   VertexState beamVertexState(beamSpot);
   if ((beamVertexState.error().cxx() <= 0.) || (beamVertexState.error().cyy() <= 0.) ||
       (beamVertexState.error().czz() <= 0.)) {
     validBS = false;
     edm::LogError("UnusableBeamSpot") << "Beamspot with invalid errors " << beamVertexState.error().matrix();
-  }
+  }*/
 
   //if this is a recovery iteration, check if we already have a valid PV
   if (fRecoveryIteration) {
@@ -435,12 +435,39 @@ void PrimaryVertexProducerCUDA::produce(edm::Event& iEvent, const edm::EventSetu
 
     //should already be sorted
     //if (GPUverticesObject->nTrueVertex(0) > 1) sort(GPUverticesObject->order(0), GPUverticesObject->order(GPUverticesObject->nTrueVertex(0)-1), VertexHigherPtSquared());
-    for (unsigned int i = 0; i < GPUverticesObject->nTrueVertex(0); i++) {
-      auto iv = GPUverticesObject->order(i);
-    //for (TrackForPV::VertexForPvSoA::const_iterator iv = GPUverticesObject->order(0); iv != GPUverticesObject->order(GPUverticesObject->nTrueVertex(0)-1); iv++) {
-      vColl.push_back(*iv);
+    //   for (unsigned int i = 0; i < GPUverticesObject->nTrueVertex(0); i++) {
+    //   auto iv = GPUverticesObject->order(i);
+    //   for (TrackForPV::VertexForPvSoA::const_iterator iv = GPUverticesObject->order(0); iv != GPUverticesObject->order(GPUverticesObject->nTrueVertex(0)-1); iv++) {
+    //  vColl.push_back(*iv);
+    //  }
+    // We have to copy the vertex back to CPU first
+    cudaCheck(cudaMemcpy(CPUverticesObject, GPUverticesObject, sizeof(TrackForPV::VertexForPVSoA), cudaMemcpyDeviceToHost));
+    // Then we iterate over them and apply the conversion
+    for (unsigned int ivertex = 0; ivertex < CPUverticesObject->nTrueVertex(blockIdx.x) ; ivertex++){
+      if (CPUverticesObject->isGood(ivertex)){
+	// I.e. the vertex is correct, so we fill a new one, first we get the error matrix
+        AlgebraicSymMatrix33 newErr;
+        newErr(0, 0) = CPUverticesObject->errx(ivertex);
+        newErr(1, 1) = CPUverticesObject->erry(ivertex);
+        newErr(2, 2) = CPUverticesObject->errz(ivertex);
+        // Then we build the new vertex
+	reco::Vertex newVertex = reco::Vertex(reco::Vertex::Point(CPUverticesObject->x(ivertex), CPUverticesObject->y(ivertex), CPUverticesObject->z(ivertex)),
+                GlobalError(newErr).matrix4D(),
+                CPUverticesObject->t(ivertex), // Without time, for the moment
+                CPUverticesObject->chi2(ivertex),
+                CPUverticesObject->ndof(ivertex),
+                CPUverticesObject->ntracks(ivertex));
+        // And we fill up the track information
+	for (unsigned int itrack = 0; itrack < CPUverticesObject->ntracks(ivertex) ; itrack++){
+          newVertex.add(t_tks.at(CPUverticesObject->track_id(ivertex)(itrack)).trackBaseRef(), CPUverticesObject->track_weight(ivertex)(itrack)); // They are never refitted tracks so this is ok
+        }
+	// We push the new vertex into the collection then
+	vColl.push_back(newVertex);
+      }
     }
 
+
+    // This we can keep as is, if we found no vertex, fill a dummy one
     if (vColl.empty()) {
       GlobalError bse(beamSpot.rotatedCovariance3D());
       if ((bse.cxx() <= 0.) || (bse.cyy() <= 0.) || (bse.czz() <= 0.)) {