Project CMSSW displayed by LXR CMSSW_15_0_X_2024-12-05-2300/​HeterogeneousCore/​AlpakaInterface/​interface/​radixSort.h	Source navigation Diff markup Identifier search General search
	Version: CMSSW_15_0_X_2024-12-05-2300 CMSSW_15_0_X_2024-12-04-2300 CMSSW_15_0_X_2024-12-03-2300 CMSSW_15_0_X_2024-12-02-2300 CMSSW_15_0_X_2024-12-01-2300 Revert   Change

File indexing completed on 2024-04-09 02:22:21

 #ifndef HeterogeneousCore_AlpakaInterface_interface_radixSort_h
 #define HeterogeneousCore_AlpakaInterface_interface_radixSort_h
 
 #include <algorithm>
 #include <cstdint>
 #include <numeric>
 #include <type_traits>
 
 #include <alpaka/alpaka.hpp>
 
 #include "HeterogeneousCore/AlpakaInterface/interface/workdivision.h"
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
 
 namespace cms::alpakatools {
 
   template <typename TAcc, typename T>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void dummyReorder(
       const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {}
 
   template <typename TAcc, typename T>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void reorderSigned(
       const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {
     //move negative first...
 
     auto& firstNeg = alpaka::declareSharedVar<uint32_t, __COUNTER__>(acc);
     firstNeg = a[ind[0]] < 0 ? 0 : size;
     alpaka::syncBlockThreads(acc);
 
     // find first negative
     for (auto idx : independent_group_elements(acc, size - 1)) {
       if ((a[ind[idx]] ^ a[ind[idx + 1]]) < 0) {
         firstNeg = idx + 1;
       }
     }
 
     alpaka::syncBlockThreads(acc);
 
     for (auto idx : independent_group_elements(acc, firstNeg, size)) {
       ind2[idx - firstNeg] = ind[idx];
     }
     alpaka::syncBlockThreads(acc);
 
     for (auto idx : independent_group_elements(acc, firstNeg)) {
       ind2[idx + size - firstNeg] = ind[idx];
     }
     alpaka::syncBlockThreads(acc);
 
     for (auto idx : independent_group_elements(acc, size)) {
       ind[idx] = ind2[idx];
     }
   }
 
   template <typename TAcc, typename T>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void reorderFloat(
       const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {
     //move negative first...
 
     auto& firstNeg = alpaka::declareSharedVar<uint32_t, __COUNTER__>(acc);
     firstNeg = a[ind[0]] < 0 ? 0 : size;
     alpaka::syncBlockThreads(acc);
 
     // find first negative
     for (auto idx : independent_group_elements(acc, size - 1)) {
       if ((a[ind[idx]] ^ a[ind[idx + 1]]) < 0)
         firstNeg = idx + 1;
     }
     alpaka::syncBlockThreads(acc);
 
     for (auto idx : independent_group_elements(acc, firstNeg, size)) {
       ind2[size - idx - 1] = ind[idx];
     }
     alpaka::syncBlockThreads(acc);
 
     for (auto idx : independent_group_elements(acc, firstNeg)) {
       ind2[idx + size - firstNeg] = ind[idx];
     }
     alpaka::syncBlockThreads(acc);
 
     for (auto idx : independent_group_elements(acc, size)) {
       ind[idx] = ind2[idx];
     }
   }
 
   // Radix sort implements a bytewise lexicographic order on the input data.
   // Data is reordered into bins indexed by the byte considered. But considering least significant bytes first
   // and respecting the existing order when binning the values, we achieve the lexicographic ordering.
   // The number of bytes actually considered is a parameter template parameter.
   // The post processing reorder
   // function fixes the order when bitwise ordering is not the order for the underlying type (only based on
   // most significant bit for signed types, integer or floating point).
   // The floating point numbers are reinterpret_cast into integers in the calling wrapper
   // This algorithm requires to run in a single block
   template <typename TAcc,
             typename T,   // shall be integer, signed or not does not matter here
             int NS,       // number of significant bytes to use in sorting.
             typename RF>  // The post processing reorder function.
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void radixSortImpl(
       const TAcc& acc, T const* __restrict__ a, uint16_t* ind, uint16_t* ind2, uint32_t size, RF reorder) {
     if constexpr (!requires_single_thread_per_block_v<TAcc>) {
       const auto warpSize = alpaka::warp::getSize(acc);
       const uint32_t threadIdxLocal(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
       [[maybe_unused]] const uint32_t blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Elems>(acc)[0u]);
       // we expect a power of 2 here
       assert(warpSize && (0 == (warpSize & (warpSize - 1))));
       const std::size_t warpMask = warpSize - 1;
 
       // Define the bin size (d=8 => 1 byte bin).
       constexpr int binBits = 8, dataBits = 8 * sizeof(T), totalSortingPassses = dataBits / binBits;
       // Make sure the slices are data aligned
       static_assert(0 == dataBits % binBits);
       // Make sure the NS parameter makes sense
       static_assert(NS > 0 && NS <= sizeof(T));
       constexpr int binsNumber = 1 << binBits;
       constexpr int binsMask = binsNumber - 1;
       // Prefix scan iterations. NS is counted in full bytes and not slices.
       constexpr int initialSortingPass = int(sizeof(T)) - NS;
 
       // Count/index for the prefix scan
       // TODO: rename
       auto& c = alpaka::declareSharedVar<int32_t[binsNumber], __COUNTER__>(acc);
       // Temporary storage for prefix scan. Only really needed for first-of-warp keeping
       // Then used for thread to bin mapping TODO: change type to byte and remap to
       auto& ct = alpaka::declareSharedVar<int32_t[binsNumber], __COUNTER__>(acc);
       // Bin to thread index mapping (used to store the highest thread index within a bin number
       // batch of threads.
       // TODO: currently initialized to an invalid value, but could also be initialized to the
       // lowest possible value (change to bytes?)
       auto& cu = alpaka::declareSharedVar<int32_t[binsNumber], __COUNTER__>(acc);
       // TODO we could also have an explicit caching of the current index for each thread.
 
       // TODO: do those have to be shared?
       auto& ibs = alpaka::declareSharedVar<int, __COUNTER__>(acc);
       auto& currentSortingPass = alpaka::declareSharedVar<int, __COUNTER__>(acc);
 
       ALPAKA_ASSERT_ACC(size > 0);
       // TODO: is this a hard requirement?
       ALPAKA_ASSERT_ACC(blockDimension >= binsNumber);
 
       currentSortingPass = initialSortingPass;
 
       auto j = ind;
       auto k = ind2;
 
       // Initializer index order to trivial increment.
       for (auto idx : independent_group_elements(acc, size)) {
         j[idx] = idx;
       }
       alpaka::syncBlockThreads(acc);
 
       // Iterate on the slices of the data.
       while (alpaka::syncBlockThreadsPredicate<alpaka::BlockAnd>(acc, (currentSortingPass < totalSortingPassses))) {
         for (auto idx : independent_group_elements(acc, binsNumber)) {
           c[idx] = 0;
         }
         alpaka::syncBlockThreads(acc);
         const auto sortingPassShift = binBits * currentSortingPass;
 
         // fill bins (count elements in each bin)
         for (auto idx : independent_group_elements(acc, size)) {
           auto bin = (a[j[idx]] >> sortingPassShift) & binsMask;
           alpaka::atomicAdd(acc, &c[bin], 1, alpaka::hierarchy::Threads{});
         }
         alpaka::syncBlockThreads(acc);
 
         if (!threadIdxLocal && 1 == alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0]) {
           //          printf("Pass=%d, Block=%d, ", currentSortingPass - 1, alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0]);
           size_t total = 0;
           for (int i = 0; i < (int)binsNumber; i++) {
             //            printf("count[%d]=%d ", i, c[i] );
             total += c[i];
           }
           //          printf("total=%zu\n", total);
           assert(total == size);
         }
         // prefix scan "optimized"???...
         // TODO: we might be able to reuse the warpPrefixScan function
         // Warp level prefix scan
         for (auto idx : independent_group_elements(acc, binsNumber)) {
           auto x = c[idx];
           auto laneId = idx & warpMask;
 
           for (int offset = 1; offset < warpSize; offset <<= 1) {
             auto y = alpaka::warp::shfl(acc, x, laneId - offset);
             if (laneId >= (uint32_t)offset)
               x += y;
           }
           ct[idx] = x;
         }
         alpaka::syncBlockThreads(acc);
 
         // Block level completion of prefix scan (add last sum of each preceding warp)
         for (auto idx : independent_group_elements(acc, binsNumber)) {
           auto ss = (idx / warpSize) * warpSize - 1;
           c[idx] = ct[idx];
           for (int i = ss; i > 0; i -= warpSize)
             c[idx] += ct[i];
         }
         // Post prefix scan, c[bin] contains the offsets in index counts to the last index +1 for each bin
 
         /*
         //prefix scan for the nulls  (for documentation)
         if (threadIdxLocal==0)
           for (int i = 1; i < sb; ++i) c[i] += c[i-1];
         */
 
         // broadcast: we will fill the new index array downward, from offset c[bin], with one thread per
         // bin, working on one set of bin size elements at a time.
         // This will reorder the indices by the currently considered slice, otherwise preserving the previous order.
         ibs = size - 1;
         alpaka::syncBlockThreads(acc);
 
         // Iterate on bin-sized slices to (size - 1) / binSize + 1 iterations
         while (alpaka::syncBlockThreadsPredicate<alpaka::BlockAnd>(acc, ibs >= 0)) {
           // Init
           for (auto idx : independent_group_elements(acc, binsNumber)) {
             cu[idx] = -1;
             ct[idx] = -1;
           }
           alpaka::syncBlockThreads(acc);
 
           // Find the highest index for all the threads dealing with a given bin (in cu[])
           // Also record the bin for each thread (in ct[])
           for (auto idx : independent_group_elements(acc, binsNumber)) {
             int i = ibs - idx;
             int32_t bin = -1;
             if (i >= 0) {
               bin = (a[j[i]] >> sortingPassShift) & binsMask;
               ct[idx] = bin;
               alpaka::atomicMax(acc, &cu[bin], int(i), alpaka::hierarchy::Threads{});
             }
           }
           alpaka::syncBlockThreads(acc);
 
           // FIXME: we can slash a memory access.
           for (auto idx : independent_group_elements(acc, binsNumber)) {
             int i = ibs - idx;
             // Are we still in inside the data?
             if (i >= 0) {
               int32_t bin = ct[idx];
               // Are we the thread with the highest index (from previous pass)?
               if (cu[bin] == i) {
                 // With the highest index, we are actually the lowest thread number. We will
                 // work "on behalf of" the higher thread numbers (including ourselves)
                 // No way around scanning and testing for bin in ct[otherThread] number to find the other threads
                 for (int peerThreadIdx = idx; peerThreadIdx < binsNumber; peerThreadIdx++) {
                   if (ct[peerThreadIdx] == bin) {
                     k[--c[bin]] = j[ibs - peerThreadIdx];
                   }
                 }
               }
             }
             /*
             int32_t bin = (i >= 0 ? ((a[j[i]] >> sortingPassShift) & binsMask) : -1);
             if (i >= 0 && i == cu[bin])  // ensure to keep them in order: only one thread per bin is active, rest is idle.
               // 
               for (int ii = idx; ii < sb; ++ii)
                 if (ct[ii] == bin) {
                   auto oi = ii - idx;
                   // assert(i>=oi);if(i>=oi)
                   k[--c[bin]] = j[i - oi]; // i = ibs - idx, oi = ii - idx => i - oi = ibs - ii;
                 }
             */
           }
           alpaka::syncBlockThreads(acc);
 
           if (threadIdxLocal == 0) {
             ibs -= binsNumber;
             // https://github.com/cms-patatrack/pixeltrack-standalone/pull/210
             // TODO: is this really needed?
             alpaka::mem_fence(acc, alpaka::memory_scope::Grid{});
           }
           alpaka::syncBlockThreads(acc);
         }
 
         /*
         // broadcast for the nulls  (for documentation)
         if (threadIdxLocal==0)
         for (int i=size-first-1; i>=0; i--) { // =blockDim.x) {
           auto bin = (a[j[i]] >> d*p)&(sb-1);
           auto ik = atomicSub(&c[bin],1);
           k[ik-1] = j[i];
         }
         */
 
         alpaka::syncBlockThreads(acc);
         ALPAKA_ASSERT_ACC(c[0] == 0);
 
         // swap (local, ok)
         auto t = j;
         j = k;
         k = t;
 
         const uint32_t threadIdxLocal(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
         if (threadIdxLocal == 0)
           ++currentSortingPass;
         alpaka::syncBlockThreads(acc);
       }
 
       if ((dataBits != 8) && (0 == (NS & 1)))
         ALPAKA_ASSERT_ACC(j ==
                           ind);  // dataBits/binBits is even so ind is correct (the result is in the right location)
 
       // TODO this copy is (doubly?) redundant with the reorder
       if (j != ind)  // odd number of sorting passes, we need to move the result to the right array (ind[])
         for (auto idx : independent_group_elements(acc, size)) {
           ind[idx] = ind2[idx];
         };
 
       alpaka::syncBlockThreads(acc);
 
       // now move negative first... (if signed)
       // TODO: the ind2 => ind copy should have beed deferred. We should pass (j != ind) as an extra parameter
       reorder(acc, a, ind, ind2, size);
     } else {
       //static_assert(false);
     }
   }
 
   template <typename TAcc,
             typename T,
             int NS = sizeof(T),  // number of significant bytes to use in sorting
             typename std::enable_if<std::is_unsigned<T>::value && !requires_single_thread_per_block_v<TAcc>, T>::type* =
                 nullptr>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void radixSort(
       const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {
     radixSortImpl<TAcc, T, NS>(acc, a, ind, ind2, size, dummyReorder<TAcc, T>);
   }
 
   template <typename TAcc,
             typename T,
             int NS = sizeof(T),  // number of significant bytes to use in sorting
             typename std::enable_if<std::is_integral<T>::value && std::is_signed<T>::value &&
                                         !requires_single_thread_per_block_v<TAcc>,
                                     T>::type* = nullptr>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void radixSort(
       const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {
     radixSortImpl<TAcc, T, NS>(acc, a, ind, ind2, size, reorderSigned<TAcc, T>);
   }
 
   template <typename TAcc,
             typename T,
             int NS = sizeof(T),  // number of significant bytes to use in sorting
             typename std::enable_if<std::is_floating_point<T>::value && !requires_single_thread_per_block_v<TAcc>,
                                     T>::type* = nullptr>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void radixSort(
       const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {
     static_assert(sizeof(T) == sizeof(int), "radixSort with the wrong type size");
     using I = int;
     radixSortImpl<TAcc, I, NS>(acc, (I const*)(a), ind, ind2, size, reorderFloat<TAcc, I>);
   }
 
   template <typename TAcc,
             typename T,
             int NS = sizeof(T),  // number of significant bytes to use in sorting
             typename std::enable_if<requires_single_thread_per_block_v<TAcc>, T>::type* = nullptr>
   /* not ALPAKA_FN_ACC to avoid trying to compile it for the CUDA or ROCm back-ends */
   ALPAKA_FN_INLINE void radixSort(const TAcc& acc, T const* a, uint16_t* ind, uint16_t* ind2, uint32_t size) {
     static_assert(requires_single_thread_per_block_v<TAcc>, "CPU sort (not a radixSort) called wtth wrong accelerator");
     // Initialize the index array
     std::iota(ind, ind + size, 0);
     /*
     printf("std::stable_sort(a=%p, ind=%p, indmax=%p, size=%d)\n", a, ind, ind + size, size);
     for (uint32_t i=0; i<10 && i<size; i++) {
       printf ("a[%d]=%ld ", i, (long int)a[i]);
     }
     printf("\n");
     for (uint32_t i=0; i<10 && i<size; i++) {
       printf ("ind[%d]=%d ", i, ind[i]);
     }
     printf("\n");
     */
     std::stable_sort(ind, ind + size, [a](uint16_t i0, uint16_t i1) { return a[i0] < a[i1]; });
     /*
     for (uint32_t i=0; i<10 && i<size; i++) {
       printf ("ind[%d]=%d ", i, ind[i]);
     }
     printf("\n");
     */
   }
 
   template <typename TAcc, typename T, int NS = sizeof(T)>
   ALPAKA_FN_ACC ALPAKA_FN_INLINE void radixSortMulti(
       const TAcc& acc, T const* v, uint16_t* index, uint32_t const* offsets, uint16_t* workspace) {
     // TODO: check
     // Sort multiple blocks of data in v[] separated by in chunks located at offsets[]
     // extern __shared__ uint16_t ws[];
     uint16_t* ws = alpaka::getDynSharedMem<uint16_t>(acc);
 
     const uint32_t blockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
     auto a = v + offsets[blockIdx];
     auto ind = index + offsets[blockIdx];
     auto ind2 = nullptr == workspace ? ws : workspace + offsets[blockIdx];
     auto size = offsets[blockIdx + 1] - offsets[blockIdx];
     assert(offsets[blockIdx + 1] >= offsets[blockIdx]);
     if (size > 0)
       radixSort<TAcc, T, NS>(acc, a, ind, ind2, size);
   }
 
   template <typename T, int NS = sizeof(T)>
   struct radixSortMultiWrapper {
     /* We cannot set launch_bounds in alpaka, so both kernel wrappers are identical (keeping CUDA/HIP code for reference for the moment)
 #if defined(__CUDACC__) || defined(__HIPCC__)
     //__global__ void __launch_bounds__(256, 4)
 #endif
 */
     template <typename TAcc>
     ALPAKA_FN_ACC void operator()(const TAcc& acc,
                                   T const* v,
                                   uint16_t* index,
                                   uint32_t const* offsets,
                                   uint16_t* workspace,
                                   size_t sharedMemBytes = 0) const {
       radixSortMulti<TAcc, T, NS>(acc, v, index, offsets, workspace);
     }
   };
 
   template <typename T, int NS = sizeof(T)>
   struct radixSortMultiWrapper2 {
     template <typename TAcc>
     ALPAKA_FN_ACC void operator()(const TAcc& acc,
                                   T const* v,
                                   uint16_t* index,
                                   uint32_t const* offsets,
                                   uint16_t* workspace,
                                   size_t sharedMemBytes = 0) const {
       radixSortMulti<TAcc, T, NS>(acc, v, index, offsets, workspace);
     }
   };
 }  // namespace cms::alpakatools
 
 namespace alpaka::trait {
   // specialize the BlockSharedMemDynSizeBytes trait to specify the amount of
   // block shared dynamic memory for the radixSortMultiWrapper kernel
   template <typename TAcc, typename T, int NS>
   struct BlockSharedMemDynSizeBytes<cms::alpakatools::radixSortMultiWrapper<T, NS>, TAcc> {
     // the size in bytes of the shared memory allocated for a block
     ALPAKA_FN_HOST_ACC static std::size_t getBlockSharedMemDynSizeBytes(
         cms::alpakatools::radixSortMultiWrapper<T, NS> const& /* kernel */,
         alpaka_common::Vec1D /* threads */,
         alpaka_common::Vec1D /* elements */,
         T const* /* v */,
         uint16_t* /* index */,
         uint32_t const* /* offsets */,
         uint16_t* workspace,
         size_t sharedMemBytes) {
       if (workspace != nullptr)
         return 0;
       /* The shared memory workspace is 'blockspace * 2' in CUDA *but that's a value coincidence... TODO: check */
       //printf ("in BlockSharedMemDynSizeBytes<cms::alpakatools::radixSortMultiWrapper<T, NS>, TAcc>: shared mem size = %d\n", (int)sharedMemBytes);
       return sharedMemBytes;
     }
   };
 }  // namespace alpaka::trait
 
 #endif  // HeterogeneousCore_AlpakaInterface_interface_radixSort_h

This page was automatically generated by the 2.2.1 LXR engine. The LXR team