solver_fused.hpp

Go to the documentation of this file.

 
#ifndef DEM_SOLVER_FUSED_HPP
#define DEM_SOLVER_FUSED_HPP
 
#include <algorithm>
#include <cstdlib>
#include <vector>
 
#include <Kokkos_Core.hpp>
 
#include "contact_preprocessing.hpp"  // CpExec/CpMem
 
namespace peclet::dem {
 
struct FusedSweepCtx {
  Kokkos::View<const int*, CpMem> offsDev;
  Kokkos::View<unsigned*, CpMem> bar;
  int maxBucket = 0;
};
 
struct FusedLoopSpec {
  int maxIters;
  float tol;
  bool strictLess;  // position loop breaks on res < tol; velocity loops on res <= tol
};
 
inline bool demFusedWanted(bool graphReplayAvailable) {
#ifdef KOKKOS_ENABLE_CUDA
  static const int mode = [] {
    if (std::getenv("PECLET_DEM_NO_FUSED"))
      return 0;
    if (std::getenv("PECLET_DEM_FUSED"))
      return 1;
    return -1;  // auto
  }();
  return mode < 0 ? !graphReplayAvailable : mode == 1;
#else
  (void)graphReplayAvailable;
  return false;
#endif
}
 
#ifdef KOKKOS_ENABLE_CUDA
 
__device__ inline void demGridBarrier(unsigned* count, unsigned k) {
  __syncthreads();  // this block's work items are done
  if (threadIdx.x == 0) {
    __threadfence();  // publish this block's writes before declaring arrival
    atomicAdd(count, 1u);
    const unsigned target = (k + 1u) * gridDim.x;
    while (*reinterpret_cast<volatile unsigned*>(count) < target) {
    }
    __threadfence();  // acquire the other blocks' writes
  }
  __syncthreads();  // release the block into the next colour
}
 
inline constexpr int kFusedBarWords = 32769;  // 4096 blocks * 8 + 1
 
template <class Sweep>
__global__ void demFusedColorSweepK(Sweep f, Kokkos::View<const int*, CpMem> perm,
                                    Kokkos::View<const int*, CpMem> offs, int numColors,
                                    unsigned* bar) {
  const int stride = gridDim.x * blockDim.x;
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  unsigned k = 0;
  for (int c = 0; c < numColors; ++c) {
    const int b = offs(c), e = offs(c + 1);
    if (b == e)
      continue;
    for (int i = b + tid; i < e; i += stride)
      f.solveOne(perm(i));
    if (c + 1 < numColors)
      demGridBarrier(bar, k++);
  }
}
 
template <class Sweep>
__global__ void demFusedSweepLoopK(Sweep f, Kokkos::View<const int*, CpMem> perm,
                                   Kokkos::View<const int*, CpMem> offs, int numColors,
                                   int maxIters, float tol, bool strictLess, float* res,
                                   unsigned* bar) {
  const int stride = gridDim.x * blockDim.x;
  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
  unsigned k = 0;
  for (int it = 0; it < maxIters; ++it) {
    if (tid == 0)
      *res = 0.0f;
    demGridBarrier(bar, k++);
    for (int c = 0; c < numColors; ++c) {
      const int b = offs(c), e = offs(c + 1);
      if (b == e)
        continue;
      for (int i = b + tid; i < e; i += stride)
        f.solveOne(perm(i));
      demGridBarrier(bar, k++);
    }
    const float r = *reinterpret_cast<volatile float*>(res);
    if (strictLess ? (r < tol) : (r <= tol))
      break;
    demGridBarrier(bar, k++);  // every block has read res before the next zero
  }
}
 
inline constexpr int kFusedBlock = 256;
 
inline int demFusedGridCap() {
  static const int cap = [] {
    const char* e = std::getenv("PECLET_DEM_FUSED_GRID");
    return e ? std::atoi(e) : 1 << 30;
  }();
  return cap;
}
 
template <class Kernel>
inline int demFusedMaxGrid(Kernel kernel) {
  static const int maxGrid = [kernel] {
    int perSm = 0;
    if (cudaOccupancyMaxActiveBlocksPerMultiprocessor(&perSm, kernel, kFusedBlock, 0) !=
        cudaSuccess) {
      (void)cudaGetLastError();
      return 0;
    }
    int dev = 0, sm = 0;
    if (cudaGetDevice(&dev) != cudaSuccess ||
        cudaDeviceGetAttribute(&sm, cudaDevAttrMultiProcessorCount, dev) != cudaSuccess) {
      (void)cudaGetLastError();
      return 0;
    }
    return perSm * sm;
  }();
  return maxGrid;
}
 
template <class Sweep>
inline bool demLaunchFusedColorSweep(CpExec& space, const Sweep& f,
                                     Kokkos::View<const int*, CpMem> perm,
                                     const FusedSweepCtx& ctx, int numColors) {
  const int maxGrid = std::min(demFusedMaxGrid(demFusedColorSweepK<Sweep>),
                               (static_cast<int>(ctx.bar.extent(0)) - 1) / 8);
  if (maxGrid <= 0 || numColors <= 0 || ctx.maxBucket <= 0)
    return false;
  const int want = std::min((ctx.maxBucket + kFusedBlock - 1) / kFusedBlock,
                            std::max(1, demFusedGridCap()));
  const int grid = want < maxGrid ? want : maxGrid;
  cudaStream_t str = space.cuda_stream();
  cudaMemsetAsync(ctx.bar.data(), 0, (static_cast<std::size_t>(grid) * 8 + 1) * sizeof(unsigned),
                  str);
  demFusedColorSweepK<Sweep>
      <<<grid, kFusedBlock, 0, str>>>(f, perm, ctx.offsDev, numColors, ctx.bar.data());
  return true;
}
 
template <class Sweep>
inline bool demLaunchFusedSweepLoop(CpExec& space, const Sweep& f,
                                    Kokkos::View<const int*, CpMem> perm,
                                    const FusedSweepCtx& ctx, int numColors,
                                    const FusedLoopSpec& spec, float* res) {
  const int maxGrid = std::min(demFusedMaxGrid(demFusedSweepLoopK<Sweep>),
                               (static_cast<int>(ctx.bar.extent(0)) - 1) / 8);
  if (maxGrid <= 0 || numColors <= 0 || ctx.maxBucket <= 0 || res == nullptr)
    return false;
  const int want = std::min((ctx.maxBucket + kFusedBlock - 1) / kFusedBlock,
                            std::max(1, demFusedGridCap()));
  const int grid = want < maxGrid ? want : maxGrid;
  cudaStream_t str = space.cuda_stream();
  cudaMemsetAsync(ctx.bar.data(), 0, (static_cast<std::size_t>(grid) * 8 + 1) * sizeof(unsigned),
                  str);
  demFusedSweepLoopK<Sweep><<<grid, kFusedBlock, 0, str>>>(f, perm, ctx.offsDev, numColors,
                                                           spec.maxIters, spec.tol,
                                                           spec.strictLess, res, ctx.bar.data());
  return true;
}
 
#endif  // KOKKOS_ENABLE_CUDA
 
inline FusedSweepCtx demMakeFusedCtx(CpExec& space, const std::vector<int>& offs,
                                     Kokkos::View<int*, CpMem> offsDev,
                                     Kokkos::View<unsigned*, CpMem> bar) {
  FusedSweepCtx ctx;
#ifdef KOKKOS_ENABLE_CUDA
  if (offs.size() < 2 || offs.size() > offsDev.extent(0) + static_cast<std::size_t>(0))
    return ctx;
  int maxBucket = 0;
  for (std::size_t c = 0; c + 1 < offs.size(); ++c)
    maxBucket = std::max(maxBucket, offs[c + 1] - offs[c]);
  if (maxBucket <= 0)
    return ctx;
  const Kokkos::View<const int*, Kokkos::HostSpace, Kokkos::MemoryUnmanaged> h(
      offs.data(), offs.size());
  auto d = Kokkos::subview(offsDev, Kokkos::pair<std::size_t, std::size_t>(0, offs.size()));
  Kokkos::deep_copy(space, d, h);
  ctx.offsDev = Kokkos::subview(Kokkos::View<const int*, CpMem>(offsDev),
                                Kokkos::pair<std::size_t, std::size_t>(0, offs.size()));
  ctx.bar = bar;
  ctx.maxBucket = maxBucket;
#else
  (void)space;
  (void)offs;
  (void)offsDev;
  (void)bar;
#endif
  return ctx;
}
 
}  // namespace peclet::dem
 
#endif  // DEM_SOLVER_FUSED_HPP