clusterTracksByDensity.h

Go to the documentation of this file.

#ifndef RecoTracker_PixelVertexFinding_plugins_alpaka_clusterTracksByDensity_h
#define RecoTracker_PixelVertexFinding_plugins_alpaka_clusterTracksByDensity_h

#include <algorithm>
#include <cmath>
#include <cstdint>

#include <alpaka/alpaka.hpp>

#include "DataFormats/VertexSoA/interface/ZVertexSoA.h"
#include "HeterogeneousCore/AlpakaInterface/interface/HistoContainer.h"
#include "HeterogeneousCore/AlpakaInterface/interface/config.h"
#include "HeterogeneousCore/AlpakaInterface/interface/workdivision.h"
#include "RecoTracker/PixelVertexFinding/interface/PixelVertexWorkSpaceLayout.h"

#include "vertexFinder.h"

namespace ALPAKA_ACCELERATOR_NAMESPACE::vertexFinder {

  using VtxSoAView = ::reco::ZVertexSoAView;
  using WsSoAView = ::vertexFinder::PixelVertexWorkSpaceSoAView;
  // this algo does not really scale as it works in a single block...
  // enough for <10K tracks we have
  //
  // based on Rodrighez&Laio algo
  //
  template <typename TAcc>
  ALPAKA_FN_ACC ALPAKA_FN_INLINE void __attribute__((always_inline))
  clusterTracksByDensity(const TAcc& acc,
                         VtxSoAView& pdata,
                         WsSoAView& pws,
                         int minT,      // min number of neighbours to be "seed"
                         float eps,     // max absolute distance to cluster
                         float errmax,  // max error to be "seed"
                         float chi2max  // max normalized distance to cluster
  ) {
    using namespace vertexFinder;
    constexpr bool verbose = false;  // in principle the compiler should optmize out if false
    const uint32_t threadIdxLocal(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);

    if constexpr (verbose) {
      if (cms::alpakatools::once_per_block(acc))
        printf("params %d %f %f %f\n", minT, eps, errmax, chi2max);
    }
    auto er2mx = errmax * errmax;

    auto& __restrict__ data = pdata;
    auto& __restrict__ ws = pws;
    auto nt = ws.ntrks();
    float const* __restrict__ zt = ws.zt();
    float const* __restrict__ ezt2 = ws.ezt2();

    uint32_t& nvFinal = data.nvFinal();
    uint32_t& nvIntermediate = ws.nvIntermediate();

    uint8_t* __restrict__ izt = ws.izt();
    int32_t* __restrict__ nn = data.ndof();
    int32_t* __restrict__ iv = ws.iv();

    ALPAKA_ASSERT_ACC(zt);
    ALPAKA_ASSERT_ACC(ezt2);
    ALPAKA_ASSERT_ACC(izt);
    ALPAKA_ASSERT_ACC(nn);
    ALPAKA_ASSERT_ACC(iv);

    using Hist = cms::alpakatools::HistoContainer<uint8_t, 256, 16000, 8, uint16_t>;
    auto& hist = alpaka::declareSharedVar<Hist, __COUNTER__>(acc);
    auto& hws = alpaka::declareSharedVar<Hist::Counter[32], __COUNTER__>(acc);

    for (auto j : cms::alpakatools::uniform_elements(acc, Hist::totbins())) {
      hist.off[j] = 0;
    }
    alpaka::syncBlockThreads(acc);

    if constexpr (verbose) {
      if (cms::alpakatools::once_per_block(acc))
        printf("booked hist with %d bins, size %d for %d tracks\n", hist.totbins(), hist.capacity(), nt);
    }
    ALPAKA_ASSERT_ACC(static_cast<int>(nt) <= hist.capacity());

    // fill hist  (bin shall be wider than "eps")
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      ALPAKA_ASSERT_ACC(i < ::zVertex::MAXTRACKS);
      int iz = int(zt[i] * 10.);  // valid if eps<=0.1
      // iz = std::clamp(iz, INT8_MIN, INT8_MAX);  // sorry c++17 only
      iz = std::min(std::max(iz, INT8_MIN), INT8_MAX);
      izt[i] = iz - INT8_MIN;
      ALPAKA_ASSERT_ACC(iz - INT8_MIN >= 0);
      ALPAKA_ASSERT_ACC(iz - INT8_MIN < 256);
      hist.count(acc, izt[i]);
      iv[i] = i;
      nn[i] = 0;
    }
    alpaka::syncBlockThreads(acc);
    if (threadIdxLocal < 32)
      hws[threadIdxLocal] = 0;  // used by prefix scan...
    alpaka::syncBlockThreads(acc);
    hist.finalize(acc, hws);
    alpaka::syncBlockThreads(acc);
    ALPAKA_ASSERT_ACC(hist.size() == nt);
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      hist.fill(acc, izt[i], uint16_t(i));
    }
    alpaka::syncBlockThreads(acc);
    // count neighbours
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      if (ezt2[i] > er2mx)
        continue;
      auto loop = [&](uint32_t j) {
        if (i == j)
          return;
        auto dist = std::abs(zt[i] - zt[j]);
        if (dist > eps)
          return;
        if (dist * dist > chi2max * (ezt2[i] + ezt2[j]))
          return;
        nn[i]++;
      };

      cms::alpakatools::forEachInBins(hist, izt[i], 1, loop);
    }
    alpaka::syncBlockThreads(acc);

    // find closest above me .... (we ignore the possibility of two j at same distance from i)
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      float mdist = eps;
      auto loop = [&](uint32_t j) {
        if (nn[j] < nn[i])
          return;
        if (nn[j] == nn[i] && zt[j] >= zt[i])
          return;  // if equal use natural order...
        auto dist = std::abs(zt[i] - zt[j]);
        if (dist > mdist)
          return;
        if (dist * dist > chi2max * (ezt2[i] + ezt2[j]))
          return;  // (break natural order???)
        mdist = dist;
        iv[i] = j;  // assign to cluster (better be unique??)
      };
      cms::alpakatools::forEachInBins(hist, izt[i], 1, loop);
    }
    alpaka::syncBlockThreads(acc);

#ifdef GPU_DEBUG
    //  mini verification
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      if (iv[i] != int(i))
        ALPAKA_ASSERT_ACC(iv[iv[i]] != int(i));
    }
    alpaka::syncBlockThreads(acc);
#endif

    // consolidate graph (percolate index of seed)
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      auto m = iv[i];
      while (m != iv[m])
        m = iv[m];
      iv[i] = m;
    }

#ifdef GPU_DEBUG
    alpaka::syncBlockThreads(acc);
    //  mini verification
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      if (iv[i] != int(i))
        ALPAKA_ASSERT_ACC(iv[iv[i]] != int(i));
    }
#endif

#ifdef GPU_DEBUG
    // and verify that we did not spit any cluster...
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      auto minJ = i;
      auto mdist = eps;
      auto loop = [&](uint32_t j) {
        if (nn[j] < nn[i])
          return;
        if (nn[j] == nn[i] && zt[j] >= zt[i])
          return;  // if equal use natural order...
        auto dist = std::abs(zt[i] - zt[j]);
        if (dist > mdist)
          return;
        if (dist * dist > chi2max * (ezt2[i] + ezt2[j]))
          return;
        mdist = dist;
        minJ = j;
      };
      cms::alpakatools::forEachInBins(hist, izt[i], 1, loop);
      // should belong to the same cluster...
      ALPAKA_ASSERT_ACC(iv[i] == iv[minJ]);
      ALPAKA_ASSERT_ACC(nn[i] <= nn[iv[i]]);
    }
    alpaka::syncBlockThreads(acc);
#endif

    auto& foundClusters = alpaka::declareSharedVar<unsigned int, __COUNTER__>(acc);
    foundClusters = 0;
    alpaka::syncBlockThreads(acc);

    // find the number of different clusters, identified by a tracks with clus[i] == i and density larger than threshold;
    // mark these tracks with a negative id.
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      if (iv[i] == int(i)) {
        if (nn[i] >= minT) {
          auto old = alpaka::atomicInc(acc, &foundClusters, 0xffffffff, alpaka::hierarchy::Threads{});
          iv[i] = -(old + 1);
        } else {  // noise
          iv[i] = -9998;
        }
      }
    }
    alpaka::syncBlockThreads(acc);

    ALPAKA_ASSERT_ACC(foundClusters < ::zVertex::MAXVTX);

    // propagate the negative id to all the tracks in the cluster.
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      if (iv[i] >= 0) {
        // mark each track in a cluster with the same id as the first one
        iv[i] = iv[iv[i]];
      }
    }
    alpaka::syncBlockThreads(acc);

    // adjust the cluster id to be a positive value starting from 0
    for (auto i : cms::alpakatools::uniform_elements(acc, nt)) {
      iv[i] = -iv[i] - 1;
    }

    nvIntermediate = nvFinal = foundClusters;
    if constexpr (verbose) {
      if (cms::alpakatools::once_per_block(acc))
        printf("found %d proto vertices\n", foundClusters);
    }
  }
  class ClusterTracksByDensityKernel {
  public:
    template <typename TAcc>
    ALPAKA_FN_ACC void operator()(const TAcc& acc,
                                  VtxSoAView pdata,
                                  WsSoAView pws,
                                  int minT,      // min number of neighbours to be "seed"
                                  float eps,     // max absolute distance to cluster
                                  float errmax,  // max error to be "seed"
                                  float chi2max  // max normalized distance to cluster
    ) const {
      clusterTracksByDensity(acc, pdata, pws, minT, eps, errmax, chi2max);
    }
  };

}  // namespace ALPAKA_ACCELERATOR_NAMESPACE::vertexFinder

#endif  // RecoTracker_PixelVertexFinding_plugins_alpaka_clusterTracksByDensity_h