Global (currently TPC only) refit on GPU using either GPU or TrackParCov track model

Detectors/Base/include/DetectorsBase/Propagator.h

@@ -111,6 +111,7 @@ class Propagator
   GPUd() const o2::base::MatLayerCylSet* getMatLUT() const { return mMatLUT; }
   GPUd() void setGPUField(const o2::gpu::GPUTPCGMPolynomialField* field) { mGPUField = field; }
   GPUd() const o2::gpu::GPUTPCGMPolynomialField* getGPUField() const { return mGPUField; }
+  GPUd() void setBz(float bz) { mBz = bz; }
 
 #ifndef GPUCA_GPUCODE
   static Propagator* Instance()

Detectors/Base/src/Propagator.cxx

@@ -98,7 +98,10 @@ int Propagator::initFieldFromGRP(const o2::parameters::GRPObject* grp, bool verb
   }
   return 0;
 }
-
+#elif !defined(GPUCA_GPUCODE)
+Propagator::Propagator()
+{
+} // empty dummy constructor for standalone benchmark
 #endif
 
 //_______________________________________________________________________

Detectors/TPC/workflow/src/CATrackerSpec.cxx

@@ -39,6 +39,9 @@
 #include "TPCdEdxCalibrationSplines.h"
 #include "DPLUtils/DPLRawParser.h"
 #include "DetectorsBase/MatLayerCylSet.h"
+#include "DetectorsBase/Propagator.h"
+#include "DetectorsBase/GeometryManager.h"
+#include "DetectorsCommonDataFormats/NameConf.h"
 #include "DetectorsRaw/HBFUtils.h"
 #include "TPCBase/RDHUtils.h"
 #include "GPUO2InterfaceConfiguration.h"
@@ -127,7 +130,9 @@ DataProcessorSpec getCATrackerSpec(CompletionPolicyData* policyData, ca::Config
       tracker = std::make_unique<GPUCATracking>();
 
       // Create configuration object and fill settings
-      const auto grp = o2::parameters::GRPObject::loadFrom("o2sim_grp.root");
+      const auto grp = o2::parameters::GRPObject::loadFrom(o2::base::NameConf::getGRPFileName());
+      o2::base::GeometryManager::loadGeometry();
+      o2::base::Propagator::initFieldFromGRP(o2::base::NameConf::getGRPFileName());
       if (grp) {
         config.configEvent.solenoidBz = 5.00668f * grp->getL3Current() / 30000.;
         config.configEvent.continuousMaxTimeBin = grp->isDetContinuousReadOut(o2::detectors::DetID::TPC) ? -1 : 0; // Number of timebins in timeframe if continuous, 0 otherwise
@@ -216,15 +221,16 @@ DataProcessorSpec getCATrackerSpec(CompletionPolicyData* policyData, ca::Config
       if (config.configCalib.fastTransform == nullptr) {
         throw std::invalid_argument("GPUCATracking: initialization of the TPC transformation failed");
       }
+
       if (confParam.matLUTFile.size()) {
         config.configCalib.matLUT = o2::base::MatLayerCylSet::loadFromFile(confParam.matLUTFile.c_str(), "MatBud");
       }
+
       if (confParam.dEdxFile.size()) {
         processAttributes->dEdxSplines.reset(new TPCdEdxCalibrationSplines(confParam.dEdxFile.c_str()));
       } else {
         processAttributes->dEdxSplines.reset(new TPCdEdxCalibrationSplines);
       }
-
       config.configCalib.dEdxSplines = processAttributes->dEdxSplines.get();
 
       if (boost::filesystem::exists(confParam.gainCalibFile)) {
@@ -239,9 +245,10 @@ DataProcessorSpec getCATrackerSpec(CompletionPolicyData* policyData, ca::Config
       }
       config.configCalib.tpcPadGain = processAttributes->tpcPadGainCalib.get();
 
+      config.configCalib.o2Propagator = Propagator::Instance();
+
       // Sample code what needs to be done for the TRD Geometry, when we extend this to TRD tracking.
-      /*o2::base::GeometryManager::loadGeometry();
-      o2::trd::Geometry gm;
+      /* o2::trd::Geometry gm;
       gm.createPadPlaneArray();
       gm.createClusterMatrixArray();
       std::unique_ptr<o2::trd::GeometryFlat> gf(gm);

GPU/Common/GPUDefGPUParameters.h

@@ -345,6 +345,12 @@
   #ifndef GPUCA_LB_GPUTPCStartHitsSorter
     #define GPUCA_LB_GPUTPCStartHitsSorter 256
   #endif
+  #ifndef GPUCA_LB_GPUTrackingRefitKernel_mode0asGPU
+    #define GPUCA_LB_GPUTrackingRefitKernel_mode0asGPU 256
+  #endif
+  #ifndef GPUCA_LB_GPUTrackingRefitKernel_mode1asTrackParCov
+    #define GPUCA_LB_GPUTrackingRefitKernel_mode1asTrackParCov 256
+  #endif
   #define GPUCA_GET_THREAD_COUNT(...) GPUCA_M_FIRST(__VA_ARGS__)
 #else
   // The following defaults are needed to compile the host code

GPU/GPUTracking/Base/GPUConstantMem.h

@@ -47,6 +47,7 @@ class GPUTRDTracker_t
 #include "GPUTPCCompression.h"
 #include "GPUITSFitter.h"
 #include "GPUTPCClusterFinder.h"
+#include "GPUTrackingRefit.h"
 #else
 #include "GPUO2FakeClasses.h"
 #endif
@@ -71,6 +72,7 @@ struct GPUConstantMem {
   GPUTRDTrackerGPU trdTracker;
   GPUTPCClusterFinder tpcClusterer[GPUCA_NSLICES];
   GPUITSFitter itsFitter;
+  GPUTrackingRefitProcessor trackingRefit;
   GPUTrackingInOutPointers ioPtrs;
   GPUCalibObjectsConst calibObjects;
   GPUErrors errorCodes;
@@ -133,13 +135,6 @@ GPUdi() void GPUProcessor::raiseError(unsigned int code, unsigned int param1, un
   GetConstantMem()->errorCodes.raiseError(code, param1, param2, param3);
 }
 
-#if defined(GPUCA_NOCOMPAT_ALLCINT) && (!defined(GPUCA_GPULIBRARY) || !defined(GPUCA_ALIROOT_LIB)) && defined(HAVE_O2HEADERS)
-GPUd() float GPUTPCClusterFinder::getGainCorrection(tpccf::Row row, tpccf::Pad pad) const
-{
-  return GetConstantMem()->calibObjects.tpcPadGain->getGainCorrection(mISlice, row, pad);
-}
-#endif
-
 } // namespace gpu
 } // namespace GPUCA_NAMESPACE
 

GPU/GPUTracking/Base/GPUDataTypes.h

@@ -48,6 +48,7 @@ namespace o2
 class MCCompLabel;
 namespace base
 {
+class Propagator;
 class MatLayerCylSet;
 } // namespace base
 namespace trd
@@ -136,7 +137,7 @@ class GPUDataTypes
                               TPCRaw = 64 };
 
 #ifdef GPUCA_NOCOMPAT_ALLOPENCL
-  static constexpr const char* const RECO_STEP_NAMES[] = {"TPC Transformation", "TPC Sector Tracking", "TPC Track Merging and Fit", "TPC Compression", "TRD Tracking", "ITS Tracking", "TPC dEdx Computation", "TPC Cluster Finding", "TPC Decompression"};
+  static constexpr const char* const RECO_STEP_NAMES[] = {"TPC Transformation", "TPC Sector Tracking", "TPC Track Merging and Fit", "TPC Compression", "TRD Tracking", "ITS Tracking", "TPC dEdx Computation", "TPC Cluster Finding", "TPC Decompression", "Global Refit"};
   static constexpr const char* const GENERAL_STEP_NAMES[] = {"Prepare", "QA"};
   typedef bitfield<RecoStep, unsigned int> RecoStepField;
   typedef bitfield<InOutType, unsigned int> InOutTypeField;
@@ -176,6 +177,7 @@ struct GPUCalibObjectsTemplate {
   typename S<o2::trd::GeometryFlat>::type* trdGeometry = nullptr;
   typename S<TPCdEdxCalibrationSplines>::type* dEdxSplines = nullptr;
   typename S<TPCPadGainCalib>::type* tpcPadGain = nullptr;
+  typename S<o2::base::Propagator>::type* o2Propagator = nullptr;
 };
 typedef GPUCalibObjectsTemplate<DefaultPtr> GPUCalibObjects;
 typedef GPUCalibObjectsTemplate<ConstPtr> GPUCalibObjectsConst;

GPU/GPUTracking/Base/GPUO2DataTypes.h

@@ -20,6 +20,7 @@
 #include "DataFormatsTPC/ClusterNative.h"
 #include "DataFormatsTPC/Digit.h"
 #include "DetectorsBase/MatLayerCylSet.h"
+#include "DetectorsBase/Propagator.h"
 #include "TRDBase/GeometryFlat.h"
 #else
 #include "GPUO2FakeClasses.h"

GPU/GPUTracking/Base/GPUO2FakeClasses.h

@@ -102,6 +102,9 @@ class GPUTPCCompression
 class GPUTPCClusterFinder
 {
 };
+class GPUTrackingRefitProcessor
+{
+};
 struct GPUTPCCFChainContext {
 };
 #ifndef __OPENCL__

GPU/GPUTracking/Base/GPUParam.cxx

@@ -27,6 +27,9 @@ using namespace GPUCA_NAMESPACE::gpu;
 #endif
 #include <cstring>
 #include <tuple>
+#ifdef HAVE_O2HEADERS
+#include "DetectorsBase/Propagator.h"
+#endif
 
 #include "utils/qconfigrtc.h"
 
@@ -244,3 +247,21 @@ std::string GPUParamRTC::generateRTCCode(const GPUParam& param, bool useConstexp
 static_assert(alignof(GPUCA_NAMESPACE::gpu::GPUParam) == alignof(GPUCA_NAMESPACE::gpu::GPUSettingsRec));
 static_assert(alignof(GPUCA_NAMESPACE::gpu::GPUParam) == alignof(GPUCA_NAMESPACE::gpu::GPUSettingsParam));
 static_assert(sizeof(GPUCA_NAMESPACE::gpu::GPUParam) - sizeof(GPUCA_NAMESPACE::gpu::GPUParamRTC) == sizeof(GPUCA_NAMESPACE::gpu::GPUSettingsRec) + sizeof(GPUCA_NAMESPACE::gpu::GPUSettingsParam) - sizeof(GPUCA_NAMESPACE::gpu::gpu_rtc::GPUSettingsRec) - sizeof(GPUCA_NAMESPACE::gpu::gpu_rtc::GPUSettingsParam));
+
+o2::base::Propagator* GPUParam::GetDefaultO2Propagator(bool useGPUField) const
+{
+  o2::base::Propagator* prop = nullptr;
+#ifdef HAVE_O2HEADERS
+  if (useGPUField == false) {
+    throw std::runtime_error("o2 propagator withouzt gpu field unsupported");
+  }
+  prop = o2::base::Propagator::Instance();
+  if (useGPUField) {
+    prop->setGPUField(&polynomialField);
+    prop->setBz(polynomialField.GetNominalBz());
+  }
+#else
+  throw std::runtime_error("o2 propagator unsupported");
+#endif
+  return prop;
+}

GPU/GPUTracking/Base/GPUParam.h

@@ -21,6 +21,14 @@
 #include "GPUTPCGeometry.h"
 #include "GPUTPCGMPolynomialField.h"
 
+namespace o2
+{
+namespace base
+{
+class Propagator;
+} // namespace base
+} // namespace o2
+
 namespace GPUCA_NAMESPACE
 {
 namespace gpu
@@ -63,6 +71,7 @@ struct GPUParam : public internal::GPUParam_t<GPUSettingsRec, GPUSettingsParam>
   void SetDefaults(const GPUSettingsEvent* e, const GPUSettingsRec* r = nullptr, const GPUSettingsProcessing* p = nullptr, const GPURecoStepConfiguration* w = nullptr);
   void UpdateEventSettings(const GPUSettingsEvent* e, const GPUSettingsProcessing* p = nullptr);
   void LoadClusterErrors(bool Print = 0);
+  o2::base::Propagator* GetDefaultO2Propagator(bool useGPUField = false) const;
 #endif
 
   GPUd() float Alpha(int iSlice) const

GPU/GPUTracking/Base/GPUReconstructionCPU.h

@@ -39,6 +39,7 @@
 #include "GPUTPCConvertKernel.h"
 #include "GPUTPCCompressionKernels.h"
 #include "GPUTPCClusterFinderKernels.h"
+#include "GPUTrackingRefitKernel.h"
 #endif
 
 namespace GPUCA_NAMESPACE

GPU/GPUTracking/Base/GPUReconstructionIncludesDevice.h

@@ -97,6 +97,10 @@ using namespace GPUCA_NAMESPACE::gpu;
 // Files for ITS Track Fit
 #include "GPUITSFitterKernels.cxx"
 
+// Files for Refit
+#include "GPUTrackingRefit.cxx"
+#include "GPUTrackingRefitKernel.cxx"
+
 #if !defined(GPUCA_O2_LIB) && defined(__HIPCC__) && !defined(GPUCA_NO_ITS_TRAITS) && !defined(GPUCA_GPUCODE_GENRTC)
 #include "VertexerTraitsHIP.hip.cxx"
 #include "ContextHIP.hip.cxx"

GPU/GPUTracking/Base/GPUReconstructionKernels.h

@@ -69,25 +69,26 @@ GPUCA_KRNL_LB((GPUTPCCompressionGatherKernels, buffered32   ), (simple), (), ())
 GPUCA_KRNL_LB((GPUTPCCompressionGatherKernels, buffered64   ), (simple), (), ())
 GPUCA_KRNL_LB((GPUTPCCompressionGatherKernels, buffered128  ), (simple), (), ())
 GPUCA_KRNL_LB((GPUTPCCompressionGatherKernels, multiBlock   ), (simple), (), ())
-
 GPUCA_KRNL_LB((GPUTPCCFCheckPadBaseline                     ), (single), (), ())
-GPUCA_KRNL_LB((GPUTPCCFChargeMapFiller,  fillIndexMap       ), (single), (), ())
-GPUCA_KRNL_LB((GPUTPCCFChargeMapFiller,  fillFromDigits     ), (single), (), ())
-GPUCA_KRNL_LB((GPUTPCCFChargeMapFiller,  findFragmentStart  ), (single), (, char setPositions), (, setPositions))
+GPUCA_KRNL_LB((GPUTPCCFChargeMapFiller, fillIndexMap        ), (single), (), ())
+GPUCA_KRNL_LB((GPUTPCCFChargeMapFiller, fillFromDigits      ), (single), (), ())
+GPUCA_KRNL_LB((GPUTPCCFChargeMapFiller, findFragmentStart   ), (single), (, char setPositions), (, setPositions))
 GPUCA_KRNL_LB((GPUTPCCFPeakFinder                           ), (single), (), ())
-GPUCA_KRNL_LB((GPUTPCCFNoiseSuppression,   noiseSuppression ), (single), (), ())
-GPUCA_KRNL_LB((GPUTPCCFNoiseSuppression,   updatePeaks      ), (single), (), ())
+GPUCA_KRNL_LB((GPUTPCCFNoiseSuppression, noiseSuppression   ), (single), (), ())
+GPUCA_KRNL_LB((GPUTPCCFNoiseSuppression, updatePeaks        ), (single), (), ())
 GPUCA_KRNL_LB((GPUTPCCFDeconvolution                        ), (single), (), ())
 GPUCA_KRNL_LB((GPUTPCCFClusterizer                          ), (single), (, char onlyMC), (, onlyMC))
-GPUCA_KRNL((   GPUTPCCFMCLabelFlattener,   setRowOffsets    ), (single), (), ())
-GPUCA_KRNL((   GPUTPCCFMCLabelFlattener,   flatten          ), (single), (, GPUPtr1(GPUTPCLinearLabels*, out)), (, GPUPtr2(GPUTPCLinearLabels*, out)))
-GPUCA_KRNL_LB((GPUTPCCFStreamCompaction,   scanStart        ), (single), (, int iBuf, int stage), (, iBuf, stage))
-GPUCA_KRNL_LB((GPUTPCCFStreamCompaction,   scanUp           ), (single), (, int iBuf, int nElems), (, iBuf, nElems))
-GPUCA_KRNL_LB((GPUTPCCFStreamCompaction,   scanTop          ), (single), (, int iBuf, int nElems), (, iBuf, nElems))
-GPUCA_KRNL_LB((GPUTPCCFStreamCompaction,   scanDown         ), (single), (, int iBuf, unsigned int offset, int nElems), (, iBuf, offset, nElems))
-GPUCA_KRNL_LB((GPUTPCCFStreamCompaction,   compactDigits    ), (single), (, int iBuf, int stage, GPUPtr1(ChargePos*, in), GPUPtr1(ChargePos*, out)), (, iBuf, stage, GPUPtr2(ChargePos*, in), GPUPtr2(ChargePos*, out)))
+GPUCA_KRNL((   GPUTPCCFMCLabelFlattener, setRowOffsets      ), (single), (), ())
+GPUCA_KRNL((   GPUTPCCFMCLabelFlattener, flatten            ), (single), (, GPUPtr1(GPUTPCLinearLabels*, out)), (, GPUPtr2(GPUTPCLinearLabels*, out)))
+GPUCA_KRNL_LB((GPUTPCCFStreamCompaction, scanStart          ), (single), (, int iBuf, int stage), (, iBuf, stage))
+GPUCA_KRNL_LB((GPUTPCCFStreamCompaction, scanUp             ), (single), (, int iBuf, int nElems), (, iBuf, nElems))
+GPUCA_KRNL_LB((GPUTPCCFStreamCompaction, scanTop            ), (single), (, int iBuf, int nElems), (, iBuf, nElems))
+GPUCA_KRNL_LB((GPUTPCCFStreamCompaction, scanDown           ), (single), (, int iBuf, unsigned int offset, int nElems), (, iBuf, offset, nElems))
+GPUCA_KRNL_LB((GPUTPCCFStreamCompaction, compactDigits      ), (single), (, int iBuf, int stage, GPUPtr1(ChargePos*, in), GPUPtr1(ChargePos*, out)), (, iBuf, stage, GPUPtr2(ChargePos*, in), GPUPtr2(ChargePos*, out)))
 GPUCA_KRNL_LB((GPUTPCCFDecodeZS                             ), (single), (, int firstHBF), (, firstHBF))
 GPUCA_KRNL_LB((GPUTPCCFGather                               ), (single), (, GPUPtr1(o2::tpc::ClusterNative*, dest)), (, GPUPtr2(o2::tpc::ClusterNative*, dest)))
+GPUCA_KRNL_LB((GPUTrackingRefitKernel, mode0asGPU           ), (simple), (), ())
+GPUCA_KRNL_LB((GPUTrackingRefitKernel, mode1asTrackParCov   ), (simple), (), ())
 #endif
 #endif
 // clang-format on

GPU/GPUTracking/Base/GPUSettingsList.h

@@ -85,6 +85,7 @@ AddOptionRTC(loopInterpolationInExtraPass, char, -1, "", 0, "Perform loop interp
 AddOptionRTC(mergerReadFromTrackerDirectly, char, 1, "", 0, "Forward data directly from tracker to merger on GPU")
 AddOptionRTC(useMatLUT, char, 0, "", 0, "Use material lookup table for TPC refit")
 AddOptionRTC(trdStopTrkAfterNMissLy, unsigned char, 6, "", 0, "Abandon track following after N layers without a TRD match")
+AddOptionRTC(trackingRefitGPUModel, char, 1, "", 0, "Use GPU track model for the Global Track Refit")
 AddCustomCPP(void SetMinTrackPt(float v) { MaxTrackQPt = v > 0.001 ? (1. / v) : (1. / 0.001); })
 AddVariable(dummyRTC, float, 0.f) // Ensure non empty struct and proper alignment even if all normal members are constexpr
 AddHelp("help", 'h')

GPU/GPUTracking/CMakeLists.txt

@@ -175,7 +175,8 @@ if(ALIGPU_BUILD_TYPE STREQUAL "O2" OR CONFIG_O2_EXTENSIONS)
         TPCClusterFinder/GPUTPCCFGather.cxx
         TPCClusterFinder/TPCPadGainCalib.cxx
         dEdx/TPCdEdxCalibrationSplines.cxx
-        Refit/GPUTrackingRefit.cxx)
+        Refit/GPUTrackingRefit.cxx
+        Refit/GPUTrackingRefitKernel.cxx)
 
     set(SRCS_NO_H ${SRCS_NO_H}
         TPCClusterFinder/GPUTPCClusterFinderDump.cxx)

GPU/GPUTracking/Global/GPUChainTracking.cxx

@@ -125,6 +125,9 @@ void GPUChainTracking::RegisterPermanentMemoryAndProcessors()
       mRec->RegisterGPUProcessor(&processors()->tpcClusterer[i], GetRecoStepsGPU() & RecoStep::TPCClusterFinding);
     }
   }
+  if (GetRecoSteps() & RecoStep::Refit) {
+    mRec->RegisterGPUProcessor(&processors()->trackingRefit, GetRecoStepsGPU() & RecoStep::Refit);
+  }
 #endif
 #ifdef GPUCA_KERNEL_DEBUGGER_OUTPUT
   mRec->RegisterGPUProcessor(&processors()->debugOutput, true);
@@ -162,6 +165,9 @@ void GPUChainTracking::RegisterGPUProcessors()
       mRec->RegisterGPUDeviceProcessor(&processorsShadow()->tpcClusterer[i], &processors()->tpcClusterer[i]);
     }
   }
+  if (GetRecoStepsGPU() & RecoStep::Refit) {
+    mRec->RegisterGPUDeviceProcessor(&processorsShadow()->trackingRefit, &processors()->trackingRefit);
+  }
 #endif
 #ifdef GPUCA_KERNEL_DEBUGGER_OUTPUT
   mRec->RegisterGPUDeviceProcessor(&processorsShadow()->debugOutput, &processors()->debugOutput);
@@ -259,6 +265,10 @@ bool GPUChainTracking::ValidateSteps()
     GPUError("Cannot run gain calibration without calibration object");
     return false;
   }
+  if ((GetRecoSteps() & GPUDataTypes::RecoStep::Refit) && !param().rec.trackingRefitGPUModel && (processors()->calibObjects.o2Propagator == nullptr || processors()->calibObjects.matLUT == nullptr)) {
+    GPUError("Cannot run refit with o2 track model without o2 propagator");
+    return false;
+  }
   return true;
 }
 
@@ -386,8 +396,15 @@ int GPUChainTracking::Init()
     if (processors()->calibObjects.tpcPadGain) {
       memcpy((void*)mFlatObjectsShadow.mCalibObjects.tpcPadGain, (const void*)processors()->calibObjects.tpcPadGain, sizeof(*processors()->calibObjects.tpcPadGain));
     }
+    if (processors()->calibObjects.o2Propagator) {
+      memcpy((void*)mFlatObjectsShadow.mCalibObjects.o2Propagator, (const void*)processors()->calibObjects.o2Propagator, sizeof(*processors()->calibObjects.o2Propagator));
+      mFlatObjectsShadow.mCalibObjects.o2Propagator->setGPUField(&processorsDevice()->param.polynomialField);
+      mFlatObjectsShadow.mCalibObjects.o2Propagator->setBz(param().polynomialField.GetNominalBz());
+      mFlatObjectsShadow.mCalibObjects.o2Propagator->setMatLUT(mFlatObjectsShadow.mCalibObjects.matLUT);
+    }
 #endif
     TransferMemoryResourceLinkToGPU(RecoStep::NoRecoStep, mFlatObjectsShadow.mMemoryResFlat);
+    memcpy((void*)&processorsShadow()->calibObjects, (void*)&mFlatObjectsDevice.mCalibObjects, sizeof(mFlatObjectsDevice.mCalibObjects));
     WriteToConstantMemory(RecoStep::NoRecoStep, (char*)&processors()->calibObjects - (char*)processors(), &mFlatObjectsDevice.mCalibObjects, sizeof(mFlatObjectsDevice.mCalibObjects), -1); // First initialization, for users not using RunChain
     processorsShadow()->errorCodes.setMemory(mInputsShadow->mErrorCodes);
     WriteToConstantMemory(RecoStep::NoRecoStep, (char*)&processors()->errorCodes - (char*)processors(), &processorsShadow()->errorCodes, sizeof(processorsShadow()->errorCodes), -1);
@@ -454,7 +471,9 @@ void* GPUChainTracking::GPUTrackingFlatObjects::SetPointersFlatObjects(void* mem
   if (mChainTracking->GetTRDGeometry()) {
     computePointerWithAlignment(mem, mCalibObjects.trdGeometry, 1);
   }
-
+  if (mChainTracking->GetO2Propagator()) {
+    computePointerWithAlignment(mem, mCalibObjects.o2Propagator, 1);
+  }
 #endif
   return mem;
 }
@@ -2393,22 +2412,24 @@ int GPUChainTracking::DoTRDGPUTracking()
 int GPUChainTracking::RunRefit()
 {
 #ifdef HAVE_O2HEADERS
-  GPUTrackingRefit re;
-  re.SetPtrsFromGPUConstantMem(processorsShadow());
-  re.SetPropagatorDefault();
-  for (unsigned int i = 0; i < mIOPtrs.nMergedTracks; i++) {
-    if (mIOPtrs.mergedTracks[i].OK()) {
-      printf("\nRefitting track %d\n", i);
-      GPUTPCGMMergedTrack t = mIOPtrs.mergedTracks[i];
-      int retval = re.RefitTrackAsGPU(t, false, true);
-      printf("Refit error code: %d\n", retval);
+  bool doGPU = GetRecoStepsGPU() & RecoStep::Refit;
+  GPUTrackingRefitProcessor& Refit = processors()->trackingRefit;
+  GPUTrackingRefitProcessor& RefitShadow = doGPU ? processorsShadow()->trackingRefit : Refit;
 
-      printf("\nRefitting track TrackParCov %d\n", i);
-      t = mIOPtrs.mergedTracks[i];
-      retval = re.RefitTrackAsTrackParCov(t, false, true);
-      printf("Refit error code: %d\n", retval);
-    }
+  const auto& threadContext = GetThreadContext();
+  SetupGPUProcessor(&Refit, false);
+  RefitShadow.SetPtrsFromGPUConstantMem(processorsShadow(), doGPU ? &processorsDevice()->param : nullptr);
+  RefitShadow.SetPropagator(doGPU ? processorsShadow()->calibObjects.o2Propagator : GetO2Propagator());
+  RefitShadow.mPTracks = (doGPU ? processorsShadow() : processors())->tpcMerger.OutputTracks();
+  WriteToConstantMemory(RecoStep::Refit, (char*)&processors()->trackingRefit - (char*)processors(), &RefitShadow, sizeof(RefitShadow), 0);
+  //TransferMemoryResourcesToGPU(RecoStep::Refit, &Refit, 0);
+  if (param().rec.trackingRefitGPUModel) {
+    runKernel<GPUTrackingRefitKernel, GPUTrackingRefitKernel::mode0asGPU>(GetGrid(mIOPtrs.nMergedTracks, 0), krnlRunRangeNone);
+  } else {
+    runKernel<GPUTrackingRefitKernel, GPUTrackingRefitKernel::mode1asTrackParCov>(GetGrid(mIOPtrs.nMergedTracks, 0), krnlRunRangeNone);
   }
+  //TransferMemoryResourcesToHost(RecoStep::Refit, &Refit, 0);
+  SynchronizeStream(0);
 #endif
   return 0;
 }
@@ -2657,3 +2678,10 @@ void GPUChainTracking::ClearErrorCodes()
   }
   TransferMemoryResourceLinkToGPU(RecoStep::NoRecoStep, mInputsHost->mResourceErrorCodes, 0);
 }
+
+void GPUChainTracking::SetDefaultO2PropagatorForGPU()
+{
+  o2::base::Propagator* prop = param().GetDefaultO2Propagator(true);
+  prop->setMatLUT(processors()->calibObjects.matLUT);
+  SetO2Propagator(prop);
+}