Skip to content

Commit

Permalink
Add SMatrixGPU compatibility to trackParCov (AliceO2Group#13602)
Browse files Browse the repository at this point in the history 
* Add SMatrixGPU compatibility to trackParCov

* Hide TrackParCov<double> from GPU code

* Remove SMatrix.h as not needed anymore
  • Loading branch information
@mconcas
mconcas authored Oct 30, 2024
1 parent b797c12 commit 6a2cc7e
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 22 deletions.
37 changes: 37 additions & 0 deletions Common/MathUtils/include/MathUtils/SMatrixGPU.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -446,6 +446,8 @@ class SMatrixGPU
GPUdi() SMatrixGPU(SMatrixNoInit) {}
GPUd() SMatrixGPU(SMatrixIdentity);
GPUd() SMatrixGPU(const SMatrixGPU<T, D1, D2, R>& rhs);
template <class R2>
GPUd() SMatrixGPU(const SMatrixGPU<T, D1, D2, R2>& rhs);
template <class A, class R2>
GPUd() SMatrixGPU(const Expr<A, T, D1, D2, R2>& rhs);
template <class M>
Expand Down Expand Up @@ -497,6 +499,11 @@ class SMatrixGPU
GPUd() SMatrixRowGPU operator[](unsigned int i) { return SMatrixRowGPU(*this, i); }
template <class R2>
GPUd() SMatrixGPU<T, D1, D2, R>& operator+=(const SMatrixGPU<T, D1, D2, R2>& rhs);
GPUd() SMatrixGPU<T, D1, D2, R>& operator*=(const T& rhs);
template <class R2>
GPUd() SMatrixGPU<T, D1, D2, R>& operator*=(const SMatrixGPU<T, D1, D2, R2>& rhs);
template <class A, class R2>
GPUd() SMatrixGPU<T, D1, D2, R>& operator*=(const Expr<A, T, D1, D2, R2>& rhs);

GPUd() bool Invert();
GPUd() bool IsInUse(const T* p) const;
Expand Down Expand Up @@ -528,6 +535,13 @@ GPUdi() SMatrixGPU<T, D1, D2, R>::SMatrixGPU(const SMatrixGPU<T, D1, D2, R>& rhs
mRep = rhs.mRep;
}

template <class T, unsigned int D1, unsigned int D2, class R>
template <class R2>
GPUd() SMatrixGPU<T, D1, D2, R>::SMatrixGPU(const SMatrixGPU<T, D1, D2, R2>& rhs)
{
operator=(rhs);
}

template <class T, unsigned int D1, unsigned int D2, class R>
GPUdi() T* SMatrixGPU<T, D1, D2, R>::begin()
{
Expand Down Expand Up @@ -1387,6 +1401,29 @@ GPUdi() SMatrixGPU<T, D1, D2, R>& SMatrixGPU<T, D1, D2, R>::operator+=(const SMa
return *this;
}

template <class T, unsigned int D1, unsigned int D2, class R>
GPUdi() SMatrixGPU<T, D1, D2, R>& SMatrixGPU<T, D1, D2, R>::operator*=(const T & rhs)
{
for (unsigned int i = 0; i < R::kSize; ++i) {
mRep.Array()[i] *= rhs;
}
return *this;
}

template <class T, unsigned int D1, unsigned int D2, class R>
template <class R2>
GPUdi() SMatrixGPU<T, D1, D2, R>& SMatrixGPU<T, D1, D2, R>::operator*=(const SMatrixGPU<T, D1, D2, R2>& rhs)
{
return operator=(*this* rhs);
}

template <class T, unsigned int D1, unsigned int D2, class R>
template <class A, class R2>
GPUdi() SMatrixGPU<T, D1, D2, R>& SMatrixGPU<T, D1, D2, R>::operator*=(const Expr<A, T, D1, D2, R2>& rhs)
{
return operator=(*this* rhs);
}

template <class T, unsigned int D1, unsigned int D2, class R>
struct TranspPolicyGPU {
enum {
Expand Down
13 changes: 7 additions & 6 deletions DataFormats/Reconstruction/include/ReconstructionDataFormats/TrackParametrizationWithError.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
#define INCLUDE_RECONSTRUCTIONDATAFORMATS_TRACKPARAMETRIZATIONWITHERROR_H_

#include "ReconstructionDataFormats/TrackParametrization.h"
#include <MathUtils/Cartesian.h>

namespace o2
{
Expand All @@ -38,8 +39,8 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
#endif

using covMat_t = gpu::gpustd::array<value_t, kCovMatSize>;
using MatrixDSym5 = ROOT::Math::SMatrix<double, kNParams, kNParams, ROOT::Math::MatRepSym<double, kNParams>>;
using MatrixD5 = ROOT::Math::SMatrix<double, kNParams, kNParams, ROOT::Math::MatRepStd<double, kNParams, kNParams>>;
using MatrixDSym5 = o2::math_utils::SMatrix<double, kNParams, kNParams, o2::math_utils::MatRepSym<double, kNParams>>;
using MatrixD5 = o2::math_utils::SMatrix<double, kNParams, kNParams, o2::math_utils::MatRepStd<double, kNParams, kNParams>>;

GPUd() TrackParametrizationWithError();
GPUd() TrackParametrizationWithError(value_t x, value_t alpha, const params_t& par, const covMat_t& cov, int charge = 1, const PID pid = PID::Pion);
Expand Down Expand Up @@ -100,12 +101,12 @@ class TrackParametrizationWithError : public TrackParametrization<value_T>
template <typename T>
GPUd() value_t getPredictedChi2(const BaseCluster<T>& p) const;

void buildCombinedCovMatrix(const TrackParametrizationWithError& rhs, MatrixDSym5& cov) const;
value_t getPredictedChi2(const TrackParametrizationWithError& rhs, MatrixDSym5& covToSet) const;
GPUd() void buildCombinedCovMatrix(const TrackParametrizationWithError& rhs, MatrixDSym5& cov) const;
GPUd() value_t getPredictedChi2(const TrackParametrizationWithError& rhs, MatrixDSym5& covToSet) const;
GPUd() value_t getPredictedChi2(const TrackParametrizationWithError& rhs) const;
GPUd() value_t getPredictedChi2Quiet(const TrackParametrizationWithError& rhs) const;
bool update(const TrackParametrizationWithError& rhs, const MatrixDSym5& covInv);
bool update(const TrackParametrizationWithError& rhs);
GPUd() bool update(const TrackParametrizationWithError& rhs, const MatrixDSym5& covInv);
GPUd() bool update(const TrackParametrizationWithError& rhs);

GPUd() bool update(const dim2_t& p, const dim3_t& cov);
GPUd() bool update(const value_t* p, const value_t* cov);
Expand Down
31 changes: 15 additions & 16 deletions DataFormats/Reconstruction/src/TrackParametrizationWithError.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,14 @@
#include "ReconstructionDataFormats/DCA.h"
#include <GPUCommonLogger.h>

#ifndef __OPENCL__
#include <cfloat>
#else
#include <float.h>
#endif

#ifndef GPUCA_GPUCODE_DEVICE
#include <iostream>
#ifndef GPUCA_STANDALONE
#include "Math/SMatrix.h"
#endif
#endif

#ifndef GPUCA_ALIGPUCODE
Expand Down Expand Up @@ -754,11 +757,17 @@ GPUd() auto TrackParametrizationWithError<value_T>::getPredictedChi2Quiet(const
return (d * (szz * d - sdz * z) + z * (sdd * z - d * sdz)) / det;
}

#if !defined(GPUCA_GPUCODE) && !defined(GPUCA_STANDALONE) // Disable function relying on ROOT SMatrix on GPU
//______________________________________________
template <typename value_T>
GPUd() auto TrackParametrizationWithError<value_T>::getPredictedChi2(const TrackParametrizationWithError<value_T>& rhs) const -> value_t
{
MatrixDSym5 cov; // perform matrix operations in double!
return getPredictedChi2(rhs, cov);
}

//______________________________________________
template <typename value_T>
void TrackParametrizationWithError<value_T>::buildCombinedCovMatrix(const TrackParametrizationWithError<value_T>& rhs, MatrixDSym5& cov) const
GPUd() void TrackParametrizationWithError<value_T>::buildCombinedCovMatrix(const TrackParametrizationWithError<value_T>& rhs, MatrixDSym5& cov) const
{
// fill combined cov.matrix (NOT inverted)
cov(kY, kY) = static_cast<double>(getSigmaY2()) + static_cast<double>(rhs.getSigmaY2());
Expand All @@ -778,14 +787,6 @@ void TrackParametrizationWithError<value_T>::buildCombinedCovMatrix(const TrackP
cov(kQ2Pt, kQ2Pt) = static_cast<double>(getSigma1Pt2()) + static_cast<double>(rhs.getSigma1Pt2());
}

//______________________________________________
template <typename value_T>
GPUd() auto TrackParametrizationWithError<value_T>::getPredictedChi2(const TrackParametrizationWithError<value_T>& rhs) const -> value_t
{
MatrixDSym5 cov; // perform matrix operations in double!
return getPredictedChi2(rhs, cov);
}

//______________________________________________
template <typename value_T>
GPUd() auto TrackParametrizationWithError<value_T>::getPredictedChi2(const TrackParametrizationWithError<value_T>& rhs, MatrixDSym5& covToSet) const -> value_t
Expand Down Expand Up @@ -867,7 +868,7 @@ GPUd() bool TrackParametrizationWithError<value_T>::update(const TrackParametriz
}

// updated covariance: Cov0 = Cov0 - K*Cov0
matK *= ROOT::Math::SMatrix<double, kNParams, kNParams, ROOT::Math::MatRepStd<double, kNParams>>(matC0);
matK *= o2::math_utils::SMatrix<double, kNParams, kNParams, o2::math_utils::MatRepStd<double, kNParams>>(matC0);
mC[kSigY2] -= matK(kY, kY);
mC[kSigZY] -= matK(kZ, kY);
mC[kSigZ2] -= matK(kZ, kZ);
Expand Down Expand Up @@ -901,8 +902,6 @@ GPUd() bool TrackParametrizationWithError<value_T>::update(const TrackParametriz
return update(rhs, covI);
}

#endif

//______________________________________________
template <typename value_T>
GPUd() bool TrackParametrizationWithError<value_T>::update(const value_t* p, const value_t* cov)
Expand Down

0 comments on commit 6a2cc7e

Please sign in to comment.