Skip to content

Commit

Permalink
[alpaka] Refactor prefixScan implementation
Browse files Browse the repository at this point in the history
  • Loading branch information
@antoniopetre @fwyzard
antoniopetre authored and fwyzard committed Oct 12, 2021
1 parent 0733b90 commit fb7bd6f
Show file tree
Hide file tree
Showing 3 changed files with 39 additions and 80 deletions.
15 changes: 7 additions & 8 deletions src/alpaka/AlpakaCore/HistoContainer.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ namespace cms {
T const *__restrict__ v,
uint32_t const *__restrict__ offsets) const {
const uint32_t nt = offsets[nh];

cms::alpakatools::for_each_element_in_grid_strided(acc, nt, [&](uint32_t i) {
auto off = alpaka_std::upper_bound(offsets, offsets + nh + 1, i);
ALPAKA_ASSERT_OFFLOAD((*off) > 0);
Expand Down Expand Up @@ -74,17 +75,15 @@ namespace cms {
const unsigned int nblocks = (num_items + nthreads - 1) / nthreads;
const Vec1 blocksPerGrid(nblocks);

auto d_pc = cms::alpakatools::allocDeviceBuf<int32_t>(1u);
int32_t *pc = alpaka::getPtrNative(d_pc);
alpaka::memset(queue, d_pc, 0, 1u);

const WorkDiv1 &workDiv = cms::alpakatools::make_workdiv(blocksPerGrid, threadsPerBlockOrElementsPerThread);
alpaka::enqueue(queue,
alpaka::createTaskKernel<ALPAKA_ACCELERATOR_NAMESPACE::Acc1>(
workDiv, multiBlockPrefixScanFirstStep<uint32_t>(), poff, poff, num_items));

const WorkDiv1 &workDivWith1Block =
cms::alpakatools::make_workdiv(Vec1::all(1), threadsPerBlockOrElementsPerThread);
alpaka::enqueue(
queue,
alpaka::createTaskKernel<ALPAKA_ACCELERATOR_NAMESPACE::Acc1>(
workDivWith1Block, multiBlockPrefixScanSecondStep<uint32_t>(), poff, poff, num_items, nblocks));
workDiv, multiBlockPrefixScan<uint32_t>(), poff, poff, num_items, pc));
alpaka::wait(queue);
}

template <typename Histo, typename T>
Expand Down
79 changes: 24 additions & 55 deletions src/alpaka/AlpakaCore/prefixScan.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
#include <cstdint>
#include "CUDACore/CMSUnrollLoop.h"
#include "AlpakaCore/alpakaConfig.h"
#include "AlpakaCore/threadfence.h"

#ifdef ALPAKA_ACC_GPU_CUDA_ENABLED

Expand Down Expand Up @@ -51,7 +52,7 @@ namespace cms {
uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
uint32_t const gridBlockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
uint32_t const blockThreadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
assert(ws);
ALPAKA_ASSERT_OFFLOAD(ws);
ALPAKA_ASSERT_OFFLOAD(size <= 1024);
ALPAKA_ASSERT_OFFLOAD(0 == blockDimension % 32);
auto first = blockThreadIdx;
Expand Down Expand Up @@ -99,7 +100,7 @@ namespace cms {
uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
uint32_t const gridBlockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
uint32_t const blockThreadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
assert(ws);
ALPAKA_ASSERT_OFFLOAD(ws);
ALPAKA_ASSERT_OFFLOAD(size <= 1024);
ALPAKA_ASSERT_OFFLOAD(0 == blockDimension % 32);
auto first = blockThreadIdx;
Expand Down Expand Up @@ -134,43 +135,46 @@ namespace cms {

// limited to 1024*1024 elements....
template <typename T>
struct multiBlockPrefixScanFirstStep {
struct multiBlockPrefixScan {
template <typename T_Acc>
ALPAKA_FN_ACC void operator()(const T_Acc& acc, T const* ci, T* co, int32_t size) const {
ALPAKA_FN_ACC void operator()(const T_Acc& acc, T const* ci, T* co, int32_t size, int32_t* pc) const {
uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
uint32_t const threadDimension(alpaka::getWorkDiv<alpaka::Thread, alpaka::Elems>(acc)[0u]);
uint32_t const blockIdx(alpaka::getIdx<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
uint32_t const threadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);

auto& ws = alpaka::declareSharedVar<T[32], __COUNTER__>(acc);
// first each block does a scan of size 1024; (better be enough blocks....)
#ifndef NDEBUG
uint32_t const gridDimension(alpaka::getWorkDiv<alpaka::Grid, alpaka::Blocks>(acc)[0u]);
ALPAKA_ASSERT_OFFLOAD(gridDimension / threadDimension <= 1024);
#endif
int off = blockDimension * blockIdx * threadDimension;
auto& ws = alpaka::declareSharedVar<T[32], __COUNTER__>(acc);
if (size - off > 0)
blockPrefixScan(acc, ci + off, co + off, std::min(int(blockDimension * threadDimension), size - off), ws);
}
};

// limited to 1024*1024 elements....
template <typename T>
struct multiBlockPrefixScanSecondStep {
template <typename T_Acc>
ALPAKA_FN_ACC void operator()(const T_Acc& acc, T const* ci, T* co, int32_t size, int32_t numBlocks) const {
uint32_t const blockDimension(alpaka::getWorkDiv<alpaka::Block, alpaka::Threads>(acc)[0u]);
uint32_t const threadDimension(alpaka::getWorkDiv<alpaka::Thread, alpaka::Elems>(acc)[0u]);
auto& isLastBlockDone = alpaka::declareSharedVar<bool, __COUNTER__>(acc);
if (0 == threadIdx) {
cms::alpakatools::threadfence(acc);
auto value = alpaka::atomicAdd(acc, pc, 1, alpaka::hierarchy::Blocks{}); // block counter
isLastBlockDone = (value == (int(gridDimension) - 1));
}

uint32_t const threadIdx(alpaka::getIdx<alpaka::Block, alpaka::Threads>(acc)[0u]);
alpaka::syncBlockThreads(acc);

auto* const psum(alpaka::getDynSharedMem<T>(acc));
if (!isLastBlockDone)
return;

ALPAKA_ASSERT_OFFLOAD(int(gridDimension) == *pc);

auto& psum = alpaka::declareSharedVar<T[1024], __COUNTER__>(acc);

ALPAKA_ASSERT_OFFLOAD(static_cast<int32_t>(blockDimension * threadDimension) >= gridDimension);

// first each block does a scan of size 1024; (better be enough blocks....)
ALPAKA_ASSERT_OFFLOAD(static_cast<int32_t>(blockDimension * threadDimension) >= numBlocks);
for (int elemId = 0; elemId < static_cast<int>(threadDimension); ++elemId) {
int index = +threadIdx * threadDimension + elemId;

if (index < numBlocks) {
if (index < gridDimension) {
int lastElementOfPreviousBlockId = index * blockDimension * threadDimension - 1;
psum[index] = (lastElementOfPreviousBlockId < size and lastElementOfPreviousBlockId >= 0)
? co[lastElementOfPreviousBlockId]
Expand All @@ -179,9 +183,7 @@ namespace cms {
}

alpaka::syncBlockThreads(acc);

auto& ws = alpaka::declareSharedVar<T[32], __COUNTER__>(acc);
blockPrefixScan(acc, psum, psum, numBlocks, ws);
blockPrefixScan(acc, psum, psum, gridDimension, ws);

for (int elemId = 0; elemId < static_cast<int>(threadDimension); ++elemId) {
int first = threadIdx * threadDimension + elemId;
Expand All @@ -192,40 +194,7 @@ namespace cms {
}
}
};

} // namespace alpakatools
} // namespace cms

namespace alpaka {
namespace traits {

//#############################################################################
//! The trait for getting the size of the block shared dynamic memory for a kernel.
template <typename T, typename TAcc>
struct BlockSharedMemDynSizeBytes<cms::alpakatools::multiBlockPrefixScanSecondStep<T>, TAcc> {
//-----------------------------------------------------------------------------
//! \return The size of the shared memory allocated for a block.
template <typename TVec>
ALPAKA_FN_HOST_ACC static auto getBlockSharedMemDynSizeBytes(
cms::alpakatools::multiBlockPrefixScanSecondStep<T> const& myKernel,
TVec const& blockThreadExtent,
TVec const& threadElemExtent,
T const* ci,
T* co,
int32_t size,
int32_t numBlocks) -> T {
alpaka::ignore_unused(myKernel);
alpaka::ignore_unused(blockThreadExtent);
alpaka::ignore_unused(threadElemExtent);
alpaka::ignore_unused(ci);
alpaka::ignore_unused(co);
alpaka::ignore_unused(size);

return static_cast<size_t>(numBlocks) * sizeof(T);
}
};

} // namespace traits
} // namespace alpaka

#endif // HeterogeneousCore_AlpakaUtilities_interface_prefixScan_h
25 changes: 8 additions & 17 deletions src/alpaka/test/alpaka/prefixScan_t.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -176,23 +176,14 @@ int main() {
cms::alpakatools::make_workdiv(blocksPerGrid4, threadsPerBlockOrElementsPerThread4);

std::cout << "launch multiBlockPrefixScan " << num_items << ' ' << nBlocks << std::endl;
alpaka::enqueue(queue,
alpaka::createTaskKernel<Acc1>(workDivMultiBlock,
cms::alpakatools::multiBlockPrefixScanFirstStep<uint32_t>(),
input_d,
output1_d,
num_items));

const Vec1 blocksPerGridSecondStep(Vec1::all(1));
const WorkDiv1& workDivMultiBlockSecondStep =
cms::alpakatools::make_workdiv(blocksPerGridSecondStep, threadsPerBlockOrElementsPerThread4);
alpaka::enqueue(queue,
alpaka::createTaskKernel<Acc1>(workDivMultiBlockSecondStep,
cms::alpakatools::multiBlockPrefixScanSecondStep<uint32_t>(),
input_d,
output1_d,
num_items,
nBlocks));
auto d_pc(alpaka::allocBuf<int32_t, Idx>(device, size));
int32_t* pc = alpaka::getPtrNative(d_pc);

alpaka::memset(queue, d_pc, 0, size);
alpaka::enqueue(
queue,
alpaka::createTaskKernel<Acc1>(
workDivMultiBlock, cms::alpakatools::multiBlockPrefixScan<uint32_t>(), input_d, output1_d, num_items, pc));

alpaka::enqueue(queue, alpaka::createTaskKernel<Acc1>(workDivMultiBlock, verify(), output1_d, num_items));

Expand Down

0 comments on commit fb7bd6f

Please sign in to comment.