//===---- Reduction.cpp - OpenMP device reduction implementation - C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file contains the implementation of reduction with KMPC interface.
//
//===----------------------------------------------------------------------===//

#include "DeviceTypes.h"
#include "Interface.h"
#include "Mapping.h"
#include "State.h"
#include "Synchronization.h"

using namespace ompx;

static constexpr uint32_t kmpc_min(uint32_t a, uint32_t b) {
  return a < b ? a : b;
}

// Round down to the nearest multiple of the warp size. Return 1 if the value is
// less than the warp size.
static uint32_t round_down_to_warpsize(uint32_t s) {
  if (s < mapping::getWarpSize())
    return 1;
  return (s & ~static_cast<uint32_t>(mapping::getWarpSize() - 1u));
}

static void gpu_regular_warp_reduce(void *reduce_data,
                                    ShuffleReductFnTy shflFct) {
  for (uint32_t mask = mapping::getWarpSize() / 2; mask > 0; mask /= 2) {
    shflFct(reduce_data, /*LaneId - not used= */ 0,
            /*Offset = */ mask, /*AlgoVersion=*/0);
  }
}

static void gpu_irregular_warp_reduce(void *reduce_data,
                                      ShuffleReductFnTy shflFct, uint32_t size,
                                      uint32_t tid) {
  uint32_t curr_size;
  uint32_t mask;
  curr_size = size;
  mask = curr_size / 2;
  while (mask > 0) {
    shflFct(reduce_data, /*LaneId = */ tid, /*Offset=*/mask, /*AlgoVersion=*/1);
    curr_size = (curr_size + 1) / 2;
    mask = curr_size / 2;
  }
}

static uint32_t gpu_irregular_simd_reduce(void *reduce_data,
                                          ShuffleReductFnTy shflFct) {
  uint32_t size, remote_id, physical_lane_id;
  physical_lane_id = mapping::getThreadIdInBlock() % mapping::getWarpSize();
  __kmpc_impl_lanemask_t lanemask_lt = mapping::lanemaskLT();
  __kmpc_impl_lanemask_t Liveness = mapping::activemask();
  uint32_t logical_lane_id = utils::popc(Liveness & lanemask_lt) * 2;
  __kmpc_impl_lanemask_t lanemask_gt = mapping::lanemaskGT();
  do {
    Liveness = mapping::activemask();
    remote_id = utils::ctz(Liveness & lanemask_gt);
    size = utils::popc(Liveness);
    logical_lane_id /= 2;
    shflFct(reduce_data, /*LaneId =*/logical_lane_id,
            /*Offset=*/remote_id - physical_lane_id, /*AlgoVersion=*/2);
  } while (logical_lane_id % 2 == 0 && size > 1);
  return (logical_lane_id == 0);
}

// Reduction within a block on the GPU.
//
// Template parameters:
// - checkLiveness: Whether to check the liveness of the lanes. This is
//                  useful if gpu_block_reduce is called in a context where
//                  partial warps or L2 parallel regions are possible.
// Parameters:
// - reduce_data: Pointer to the reduction data
// - shflFct:     Shuffle reduction function
// - cpyFct:      Inter-warp copy function (copies data from each warp's thread
//                0 to the lanes of the zeroth warp)
// - NumThreads:  Number of threads to consider / values to reduce
// - ThreadId:    Thread ID in block (getThreadIdInBlock() in SPMD and 0 in
//                Generic mode)
//
// Returns:
// - 1 if the thread is the zeroth thread of the block
// - 0 otherwise
//
// Note that it is expected that the caller checks for NumThreads <= 1 and acts
// in a way that suits the callers situation. If checkLiveness is false, this
// function performs a regular warp reduce unconditionally.
//
template <bool checkLiveness = true>
static uint32_t gpu_block_reduce(void *reduce_data, ShuffleReductFnTy shflFct,
                                 InterWarpCopyFnTy cpyFct, uint32_t NumThreads,
                                 uint32_t BlockThreadId) {
  if constexpr (checkLiveness) {
    __kmpc_impl_lanemask_t Liveness = mapping::activemask();
    // Check for partial warp with non-contiguous lanes.
    if (Liveness == lanes::All) {
      gpu_regular_warp_reduce(reduce_data, shflFct);
    } else if (!(Liveness & (Liveness + 1))) {
      // Partial warp but contiguous lanes.
      gpu_irregular_warp_reduce(reduce_data, shflFct, utils::popc(Liveness),
                                BlockThreadId % mapping::getWarpSize());
    } else {
      // Dispersed lanes. Only threads in L2 parallel region may enter here.
      return gpu_irregular_simd_reduce(reduce_data, shflFct);
    }
  } else {
    gpu_regular_warp_reduce(reduce_data, shflFct);
  }

  // When we have more than [mapping::getWarpSize()] number of threads
  // a block reduction is performed here.
  //
  // Only L1 parallel region can enter this if condition.

  if (NumThreads > mapping::getWarpSize()) {
    uint32_t WarpsNeeded = utils::roundUp(NumThreads, mapping::getWarpSize());
    // Gather all the reduced values from each warp to the first warp.
    cpyFct(reduce_data, WarpsNeeded);

    uint32_t WarpId = BlockThreadId / mapping::getWarpSize();
    if (WarpId == 0)
      gpu_irregular_warp_reduce(reduce_data, shflFct, WarpsNeeded,
                                BlockThreadId);
  }

  return BlockThreadId == 0;
}

static int32_t nvptx_parallel_reduce_nowait(void *reduce_data,
                                            ShuffleReductFnTy shflFct,
                                            InterWarpCopyFnTy cpyFct) {
  uint32_t BlockThreadId = mapping::getThreadIdInBlock();
  if (mapping::isMainThreadInGenericMode(/*IsSPMD=*/false))
    BlockThreadId = 0;
  uint32_t NumThreads = omp_get_num_threads();
  if (NumThreads == 1)
    return 1;

  //
  // This reduce function handles reduction within a team. It handles
  // parallel regions in both L1 and L2 parallelism levels. It also
  // supports Generic, SPMD, and NoOMP modes.
  //
  // 1. Reduce within a warp.
  // 2. Warp master copies value to warp 0 via shared memory.
  // 3. Warp 0 reduces to a single value.
  // 4. The reduced value is available in the thread that returns 1.
  //

#if __has_builtin(__nvvm_reflect)
  if (__nvvm_reflect("__CUDA_ARCH") >= 700) {
    uint32_t WarpsNeeded = utils::roundUp(NumThreads, mapping::getWarpSize());
    uint32_t WarpId = mapping::getWarpIdInBlock();

    // Volta execution model:
    // For the Generic execution mode a parallel region either has 1 thread and
    // beyond that, always a multiple of 32. For the SPMD execution mode we may
    // have any number of threads.
    if ((NumThreads % mapping::getWarpSize() == 0) ||
        (WarpId < WarpsNeeded - 1))
      gpu_regular_warp_reduce(reduce_data, shflFct);
    else if (NumThreads > 1) // Only SPMD execution mode comes thru this case.
      gpu_irregular_warp_reduce(
          reduce_data, shflFct,
          /*LaneCount=*/NumThreads % mapping::getWarpSize(),
          /*LaneId=*/mapping::getThreadIdInBlock() % mapping::getWarpSize());

    // When we have more than [mapping::getWarpSize()] number of threads
    // a block reduction is performed here.
    //
    // Only L1 parallel region can enter this if condition.
    if (NumThreads > mapping::getWarpSize()) {
      // Gather all the reduced values from each warp
      // to the first warp.
      cpyFct(reduce_data, WarpsNeeded);

      if (WarpId == 0)
        gpu_irregular_warp_reduce(reduce_data, shflFct, WarpsNeeded,
                                  BlockThreadId);
    }
    return BlockThreadId == 0;
  }
#endif

  return gpu_block_reduce(reduce_data, shflFct, cpyFct, NumThreads,
                          BlockThreadId);
}

extern "C" {
[[clang::always_inline]]
int32_t __kmpc_nvptx_parallel_reduce_nowait_v2(IdentTy *Loc,
                                               uint64_t reduce_data_size,
                                               void *reduce_data,
                                               ShuffleReductFnTy shflFct,
                                               InterWarpCopyFnTy cpyFct) {
  return nvptx_parallel_reduce_nowait(reduce_data, shflFct, cpyFct);
}

// Reduction across teams on the GPU.
//
// Parameters:
// - Loc: Location of the reduction
// - reduce_data: Pointer to the reduction data
// - shflFct:  Shuffle reduction function
// - cpyFct:   Inter-warp copy function (copies data from each warp's thread 0
//             to the lanes of the zeroth warp)
// - lgcpyFct: List-global copy function (copies the reduction data from the
//             local thread to the global buffer)
// - glcpyFct: Global copy function (copies the reduction data from the global
//             buffer to the local thread)
// - glredFct: Global reduce function (reduces the reduction data from the
//             global buffer to the local thread)
//
// Returns:
// - 1 if this thread must write the final reduced value back to the shared
//   reduction variable (i.e. thread 0 of the single team when NumTeams == 1,
//   or thread 0 of the last team to finish its partial reduction otherwise).
// - 0 otherwise.
//
[[clang::always_inline]]
int32_t __kmpc_gpu_xteam_reduce_nowait(IdentTy *Loc, void *reduce_data,
                                       ShuffleReductFnTy shflFct,
                                       InterWarpCopyFnTy cpyFct,
                                       ListGlobalFnTy lgcpyFct,
                                       ListGlobalFnTy glcpyFct,
                                       ListGlobalFnTy glredFct) {
  uint32_t ThreadId;
  uint32_t NumThreads;

  if (mapping::isSPMDMode()) {
    // In SPMD mode all workers participate in the teams reduction.
    ThreadId = mapping::getThreadIdInBlock();
    NumThreads = mapping::getNumberOfThreadsInBlock();
  } else {
    // In generic mode, only the team master participates in the teams
    // reduction because the workers are waiting for parallel work.
    if (!mapping::isMainThreadInGenericMode())
      return 0;
    ThreadId = 0;
    NumThreads = 1;
  }

  uint32_t TeamId = omp_get_team_num();
  uint32_t NumTeams = omp_get_num_teams();

  // Fast path for single-team kernels: no cross-team work required,
  // the team-local reduction already produced the final result.
  if (NumTeams <= 1)
    return ThreadId == 0;

  Global<uint32_t> &TeamsDone =
      state::getKernelLaunchEnvironment().ReductionTeamsDone;
  void *GlobalBuffer = state::getKernelLaunchEnvironment().ReductionBuffer;
  [[clang::loader_uninitialized]] static Local<uint32_t> TeamsDoneResult;

  // Save the team's reduced value in the global buffer and atomically
  // increment the teams-done counter.
  if (ThreadId == 0) {
    lgcpyFct(GlobalBuffer, TeamId, reduce_data);
    // We let the atomic inc wrap around if the value gets larger than
    // NumTeams-1, which makes the counter self-reset.
    TeamsDoneResult = atomic::inc(&TeamsDone, NumTeams - 1u, atomic::acq_rel,
                                  atomic::MemScopeTy::device);
  }

  // This sync is needed so that all threads from last team see the shared teams
  // done counter value and know that they are in the last team.
  if (mapping::isSPMDMode())
    synchronize::threadsAligned(atomic::acq_rel);

  // If teams done counter reaches NumTeams-1, this is the last team.
  if (TeamsDoneResult != NumTeams - 1u)
    return 0;

  // The last team performs final reduction across all team values.
  NumThreads = round_down_to_warpsize(kmpc_min(NumThreads, NumTeams));
  if (ThreadId >= NumThreads)
    return 0;

  // Make sure that global buffer is fresh.
  fence::kernel(atomic::acquire);
  // Get the team values from the global buffer.
  glcpyFct(GlobalBuffer, ThreadId, reduce_data);
  // In case we have more teams than threads, we need to iterate over the
  // remaining teams.
  for (uint32_t I = NumThreads + ThreadId; I < NumTeams; I += NumThreads)
    glredFct(GlobalBuffer, I, reduce_data);

  if (NumThreads == 1)
    return 1;

  return gpu_block_reduce<false>(reduce_data, shflFct, cpyFct, NumThreads,
                                 ThreadId);
}
} // extern "C"