378 lines
12 KiB
C++
378 lines
12 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
* SPDX-License-Identifier: BSD-3-Clause
|
|
*
|
|
* Redistribution and use in source and binary forms, with or without
|
|
* modification, are permitted provided that the following conditions are met:
|
|
*
|
|
* 1. Redistributions of source code must retain the above copyright notice, this
|
|
* list of conditions and the following disclaimer.
|
|
*
|
|
* 2. Redistributions in binary form must reproduce the above copyright notice,
|
|
* this list of conditions and the following disclaimer in the documentation
|
|
* and/or other materials provided with the distribution.
|
|
*
|
|
* 3. Neither the name of the copyright holder nor the names of its
|
|
* contributors may be used to endorse or promote products derived from
|
|
* this software without specific prior written permission.
|
|
*
|
|
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
|
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
|
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
|
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
|
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
|
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
|
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
|
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
|
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
|
*
|
|
**************************************************************************************************/
|
|
/*! \file
|
|
\brief Implementation of a CTA-wide barrier for inter-CTA synchronization.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "cutlass/cutlass.h"
|
|
#include "cutlass/arch/barrier.h"
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
namespace cutlass {
|
|
|
|
namespace detail {
|
|
|
|
//
|
|
// Utilities for abstracting synchronization methods for barriers
|
|
//
|
|
|
|
struct SyncthreadsSync {
|
|
CUTLASS_DEVICE
|
|
static void sync() {
|
|
__syncthreads();
|
|
}
|
|
};
|
|
|
|
struct SyncwarpSync {
|
|
CUTLASS_DEVICE
|
|
static void sync() {
|
|
__syncwarp();
|
|
}
|
|
};
|
|
|
|
template <
|
|
int ThreadCount,
|
|
int BarrierId
|
|
>
|
|
struct NamedBarrierSync {
|
|
CUTLASS_DEVICE
|
|
static void sync() {
|
|
cutlass::arch::NamedBarrier::sync(ThreadCount, static_cast<arch::ReservedNamedBarriers>(BarrierId));
|
|
}
|
|
};
|
|
|
|
} // namepspace detail
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/// Group or CTA-wide semaphore for inter-CTA synchronization.
|
|
template <class Sync>
|
|
struct GenericBarrier {
|
|
|
|
public:
|
|
|
|
/// Flag type
|
|
using T = int;
|
|
|
|
/// Initial flag value
|
|
static const T INIT = 0;
|
|
|
|
|
|
protected:
|
|
|
|
/// Load flag, as a strong acquire operation (int specialization)
|
|
CUTLASS_DEVICE
|
|
static int ld_acquire(int *ptr)
|
|
{
|
|
int state = 0;
|
|
|
|
#if (__CUDA_ARCH__ >= 700)
|
|
/// SM70 and newer use memory consistency qualifiers
|
|
|
|
// Acquire pattern using acquire modifier
|
|
asm volatile ("ld.global.acquire.gpu.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr));
|
|
|
|
#else
|
|
asm volatile ("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr));
|
|
#endif // (__CUDA_ARCH__ >= 700)
|
|
|
|
return state;
|
|
}
|
|
|
|
|
|
/// Reduce into flag, with release pattern (int specialization)
|
|
CUTLASS_DEVICE
|
|
static void red_release(int *ptr, int val)
|
|
{
|
|
#if (__CUDA_ARCH__ >= 700)
|
|
/// SM70 and newer use memory consistency qualifiers
|
|
|
|
// Release pattern using acq_rel fence + relaxed modifier. (The fence also releases data
|
|
// that was weakly-written by other threads prior to the last syncthreads)
|
|
asm volatile ("fence.acq_rel.gpu;\n");
|
|
asm volatile ("red.relaxed.gpu.global.add.s32 [%0], %1;\n" : : "l"(ptr), "r"(val));
|
|
|
|
#else
|
|
__threadfence();
|
|
atomicAdd(ptr, val);
|
|
#endif // (__CUDA_ARCH__ >= 700)
|
|
}
|
|
|
|
|
|
public:
|
|
|
|
/// Uses thread[0] to wait for at least the specified count of signals on the given flag counter
|
|
CUTLASS_DEVICE
|
|
static void wait_lt(void *lock_ptr, int thread_idx, int flag_idx, int count)
|
|
{
|
|
T *flag_ptr = reinterpret_cast<T*>(lock_ptr) + flag_idx;
|
|
|
|
if (thread_idx == 0)
|
|
{
|
|
// Spin-loop
|
|
#pragma unroll 1
|
|
while(ld_acquire(flag_ptr) < count) {}
|
|
}
|
|
|
|
Sync::sync();
|
|
}
|
|
|
|
/// Uses thread[0] to wait for at least the specified count of signals on the given flag counter
|
|
CUTLASS_DEVICE
|
|
static void wait_eq(void *lock_ptr, int thread_idx, int flag_idx, T val = 1)
|
|
{
|
|
T *flag_ptr = reinterpret_cast<T*>(lock_ptr) + flag_idx;
|
|
|
|
if (thread_idx == 0)
|
|
{
|
|
// Spin-loop
|
|
#pragma unroll 1
|
|
while(ld_acquire(flag_ptr) != val) {}
|
|
}
|
|
Sync::sync();
|
|
}
|
|
|
|
/// Uses thread[0] to wait for the specified count of signals on the given flag counter
|
|
CUTLASS_DEVICE
|
|
static void wait_eq_reset(void *lock_ptr, int thread_idx, int flag_idx, T val = 1) {
|
|
T *flag_ptr = reinterpret_cast<T*>(lock_ptr) + flag_idx;
|
|
|
|
if (thread_idx == 0)
|
|
{
|
|
// Spin-loop
|
|
#pragma unroll 1
|
|
while(atomicCAS(flag_ptr, val, 0) != val) {}
|
|
}
|
|
|
|
Sync::sync();
|
|
}
|
|
|
|
/// Increment the arrival count for a flag
|
|
CUTLASS_DEVICE
|
|
static void arrive_inc(void *lock_ptr, int thread_idx, int flag_idx, int val = 1)
|
|
{
|
|
T* flag_ptr = reinterpret_cast<T*>(lock_ptr) + flag_idx;
|
|
|
|
Sync::sync();
|
|
|
|
if (thread_idx == 0)
|
|
{
|
|
red_release(flag_ptr, val);
|
|
}
|
|
}
|
|
|
|
|
|
/// Increment the arrival counts for a range of flags
|
|
CUTLASS_DEVICE
|
|
static void arrive_range_inc(void *lock_ptr, int thread_idx, int first_flag_idx, int count = 1, int val = 1)
|
|
{
|
|
int flag_idx = first_flag_idx + thread_idx;
|
|
T* flag_ptr = reinterpret_cast<T*>(lock_ptr) + flag_idx;
|
|
|
|
// Barrier to make sure all other threads in group have written their data
|
|
Sync::sync();
|
|
|
|
// Select threads increment their flags
|
|
if (thread_idx < count) {
|
|
red_release(flag_ptr, val);
|
|
}
|
|
}
|
|
};
|
|
|
|
using Barrier = GenericBarrier<detail::SyncthreadsSync>;
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/** Structure for managing multiple NamedBarriers to be used by different warp groups, allowing
|
|
* runtime index values to be used to call into named barriers with compile-time-constant IDs.
|
|
*
|
|
* @param ThreadCount_ Number of threads that will wait on a NamedBarrier with a given ID
|
|
* @param Offset Value added to the ID passed in by the user to determine the NamedBarrier ID to call into
|
|
* @param MaxNumNamedBarriers The maximum number of unique barrier IDs that will be requested on this type
|
|
**/
|
|
template <
|
|
uint32_t ThreadCount_,
|
|
uint32_t Offset = 0,
|
|
uint32_t MaxNumNamedBarriers = 16
|
|
>
|
|
struct NamedBarrierManager {
|
|
|
|
static_assert(MaxNumNamedBarriers <= arch::NamedBarrier::HardwareMaxNumNamedBarriers);
|
|
static_assert(MaxNumNamedBarriers + Offset <= arch::NamedBarrier::HardwareMaxNumNamedBarriers, "Barrier IDs cannot exceed 15");
|
|
|
|
// Number of threads participating in the barrier
|
|
static constexpr uint32_t ThreadCount = ThreadCount_;
|
|
|
|
template <uint32_t BarrierId>
|
|
using BarrierSync = cutlass::GenericBarrier<cutlass::detail::NamedBarrierSync<ThreadCount, BarrierId>>;
|
|
|
|
// Underlying type used by all barriers for synchronization. Does not depend on
|
|
// template parameter BarrierId, so passing in 0 suffices.
|
|
using T = typename BarrierSync<0>::T;
|
|
|
|
using IntegerSequence = cute::make_integer_sequence<uint32_t, MaxNumNamedBarriers>;
|
|
|
|
CUTLASS_DEVICE
|
|
static
|
|
void wait_lt(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, int count) {
|
|
wait_lt_helper(idx, lock_ptr, thread_idx, flag_idx, count, IntegerSequence{});
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
wait_eq(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, T val = 1) {
|
|
wait_eq_helper<false>(idx, lock_ptr, thread_idx, flag_idx, val, IntegerSequence{});
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
wait_eq_reset(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, T val = 1) {
|
|
wait_eq_helper<true>(idx, lock_ptr, thread_idx, flag_idx, val, IntegerSequence{});
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
arrive_inc(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, int val = 1) {
|
|
arrive_inc_helper(idx, lock_ptr, thread_idx, flag_idx, val, IntegerSequence{});
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
arrive_range_inc(uint32_t idx, void *lock_ptr, int thread_idx, int first_flag_idx, int count = 1, int val = 1) {
|
|
arrive_range_inc_helper(idx, lock_ptr, thread_idx, first_flag_idx, count, val, IntegerSequence{});
|
|
}
|
|
|
|
private:
|
|
CUTLASS_DEVICE
|
|
static void
|
|
check_barrier_in_range([[maybe_unused]] uint32_t idx) {
|
|
assert((idx >= MaxNumNamedBarriers) && "Index exceeds barrier count");
|
|
}
|
|
|
|
template <uint32_t... Idx>
|
|
CUTLASS_DEVICE
|
|
static void
|
|
wait_lt_helper(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, int count, cute::integer_sequence<uint32_t, Idx...>) {
|
|
check_barrier_in_range(idx);
|
|
((Idx == idx && (BarrierSync<Idx + Offset>::wait_lt(lock_ptr, thread_idx, flag_idx, count), true)) || ...);
|
|
}
|
|
|
|
template <bool Reset, uint32_t... Idx>
|
|
CUTLASS_DEVICE
|
|
static void
|
|
wait_eq_helper(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, T val, cute::integer_sequence<uint32_t, Idx...>) {
|
|
check_barrier_in_range(idx);
|
|
if constexpr (Reset) {
|
|
((Idx == idx && (BarrierSync<Idx + Offset>::wait_eq_reset(lock_ptr, thread_idx, flag_idx, val), true)) || ...);
|
|
}
|
|
else {
|
|
((Idx == idx && (BarrierSync<Idx + Offset>::wait_eq(lock_ptr, thread_idx, flag_idx, val), true)) || ...);
|
|
}
|
|
}
|
|
|
|
template <uint32_t... Idx>
|
|
CUTLASS_DEVICE
|
|
static void
|
|
arrive_inc_helper(uint32_t idx, void *lock_ptr, int thread_idx, int flag_idx, int val, cute::integer_sequence<uint32_t, Idx...>) {
|
|
check_barrier_in_range(idx);
|
|
((Idx == idx && (BarrierSync<Idx + Offset>::arrive_inc(lock_ptr, thread_idx, flag_idx, val), true)) || ...);
|
|
}
|
|
|
|
template <uint32_t... Idx>
|
|
CUTLASS_DEVICE
|
|
static void
|
|
arrive_range_inc_helper(uint32_t idx, void *lock_ptr, int thread_idx, int first_flag_idx, int count, int val, cute::integer_sequence<uint32_t, Idx...>) {
|
|
check_barrier_in_range(idx);
|
|
((Idx == idx && (BarrierSync<Idx + Offset>::arrive_range_inc(lock_ptr, thread_idx, first_flag_idx, count, val), true)) || ...);
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
/** Structure for synchronizing via contiguous barriers (e.g., __syncwarp, __syncthreads)
|
|
* via an API that mirrors that of NamedBarrierManager
|
|
*
|
|
* @param Synchronizer Synchronization helper exposing a `sync()` method to perform synchronization
|
|
**/
|
|
template <
|
|
class Synchronizer,
|
|
uint32_t ThreadCount_
|
|
>
|
|
struct SyncManager {
|
|
|
|
// Number of threads participating in the barrier
|
|
static constexpr uint32_t ThreadCount = ThreadCount_;
|
|
|
|
using BarrierSync = cutlass::GenericBarrier<Synchronizer>;
|
|
|
|
// Underlying type used by all barriers for synchronization.
|
|
using T = typename BarrierSync::T;
|
|
|
|
CUTLASS_DEVICE
|
|
static
|
|
void wait_lt(uint32_t, void *lock_ptr, int thread_idx, int flag_idx, int count) {
|
|
BarrierSync::wait_lt(lock_ptr, thread_idx, flag_idx, count);
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
wait_eq(uint32_t, void *lock_ptr, int thread_idx, int flag_idx, T val = 1) {
|
|
BarrierSync::wait_eq(lock_ptr, thread_idx, flag_idx, val);
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
wait_eq_reset(uint32_t, void *lock_ptr, int thread_idx, int flag_idx, T val = 1) {
|
|
BarrierSync::wait_eq_reset(lock_ptr, thread_idx, flag_idx, val);
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
arrive_inc(uint32_t, void *lock_ptr, int thread_idx, int flag_idx, int val = 1) {
|
|
BarrierSync::arrive_inc(lock_ptr, thread_idx, flag_idx, val);
|
|
}
|
|
|
|
CUTLASS_DEVICE
|
|
static void
|
|
arrive_range_inc(uint32_t idx, void *lock_ptr, int thread_idx, int first_flag_idx, int count = 1, int val = 1) {
|
|
BarrierSync::arrive_range_inc(lock_ptr, thread_idx, first_flag_idx, count, val);
|
|
}
|
|
};
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|
|
|
|
} // namespace cutlass
|
|
|
|
/////////////////////////////////////////////////////////////////////////////////////////////////
|