leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 Packages Projects Releases Wiki Activity

CUTLASS 3.2.1 (#1113)

Browse Source 
* Updates for 3.2.1 release.

* Minor fix in gemm op profiler for raster order.

* Add scheduler mapping for raster order in the kernels.
This commit is contained in:
ANIKET SHIVAM
2023-09-26 17:24:26 -04:00
committed by GitHub
parent e0aaa3c3b3
commit 90d3b0fb18
428 changed files with 22252 additions and 21761 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+91 -1
View File
@@ -259,6 +259,95 @@ struct Sm90TmaBuilderImpl {
>;
};
///////////////////////////////////////////////////////////////////////////////
// Descriptor classes for defining EVT nodes
// Some of the epilogue visitor nodes require non-intuitive template arguments
// such as CopyOpS2R for AuxLoad node. Traditionaly, these are resolved by the
// builder classes. Here we provide a set of descriptor classes that resolve
// these template arguments from more intuitive types such as Stride, Layout
// Get TileShape, EpilogueTile, Dispatch Policy, StagesC, and STagesD
template<
typename TileShape_MNK,
typename EpilogueTileType,
typename ElementC,
typename ElementD,
typename Schedule
>
struct EpilogueDescriptor {
using TileShape = TileShape_MNK;
using EpilogueTile =
decltype(
detail::sm90_compute_tile_shape_or_override<
ElementD, EpilogueTileType, Schedule
>()
);
using DispatchPolicy =
decltype(
detail::sm90_get_tma_dispatch_policy<
TileShape_MNK, EpilogueTile,
ElementC, ElementD, Schedule
>()
);
constexpr static int StagesC = DispatchPolicy::StagesC;
constexpr static int StagesD = DispatchPolicy::StagesD;
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxLoad node
template<
typename EpilogueDescriptor,
typename StrideOrLayoutTag,
typename ElementAux
>
struct AuxLoadDescriptor {
constexpr static int Stages = EpilogueDescriptor::StagesC;
using EpilogueTile = typename EpilogueDescriptor::EpilogueTile;
using Element = ElementAux;
using Stride = cutlass::detail::TagToStrideC_t<StrideOrLayoutTag>;
using SmemLayoutAtom =
decltype(
detail::sm90_get_epilogue_smem_swizzle_layout_atom<
Stride, ElementAux, typename EpilogueDescriptor::EpilogueTile
>()
);
using CopyOpS2R =
decltype(detail::sm90_get_smem_load_op_for_source<Stride, ElementAux>());
};
// Get Stride, SmemLayout, and CopyOpS2R for AuxStore node
template<
typename EpilogueDescriptor,
typename StrideOrLayoutTag,
typename ElementAux
>
struct AuxStoreDescriptor {
constexpr static int Stages = EpilogueDescriptor::StagesD;
using EpilogueTile = typename EpilogueDescriptor::EpilogueTile;
using Element = ElementAux;
using Stride = cutlass::detail::TagToStrideC_t<StrideOrLayoutTag>;
using SmemLayoutAtom =
decltype(
detail::sm90_get_epilogue_smem_swizzle_layout_atom<
Stride, ElementAux, typename EpilogueDescriptor::EpilogueTile
>()
);
using CopyOpR2S =
decltype(detail::sm90_get_smem_store_op_for_accumulator<Stride, ElementAux>());
};
template<
typename EpilogueDescriptor,
typename ElementVector
>
struct RowBroadcastDescriptor {
constexpr static int Stages = ceil_div(
EpilogueDescriptor::StagesC,
size(shape_div(take<0, 2>(typename EpilogueDescriptor::TileShape{}), typename EpilogueDescriptor::EpilogueTile{}))
) + 1;
using Element = ElementVector;
};
} // namespace detail
///////////////////////////////////////////////////////////////////////////////
@@ -426,7 +515,8 @@ private:
ElementD,
GmemLayoutTagD,
AlignmentD,
EpilogueSchedule
EpilogueSchedule,
FusionOperation
>;
public:
include/cutlass/epilogue/collective/collective_builder.hpp
+1
View File
@@ -45,6 +45,7 @@ struct EpilogueTileAuto {};
// Used to let the builder pick the epilogue schedule automatically.
// Can be overridden with kernel schedule tags in cutlass/gemm/dispatch_policy.hpp
struct EpilogueScheduleAuto {};
struct EpilogueIm2ColScheduleAuto {};
template <
class ArchTag,
include/cutlass/epilogue/collective/detail.hpp
+2 -2
View File
@@ -126,14 +126,14 @@ public:
CUTLASS_HOST_DEVICE
static constexpr int
get_load_pipe_increment([[maybe_unused]] TileShapeMNK) {
return 0;
return 1;
}
template<class TileShapeMNK>
CUTLASS_HOST_DEVICE
static constexpr int
get_store_pipe_increment([[maybe_unused]] TileShapeMNK) {
return 0;
return 1;
}
CUTLASS_DEVICE
include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp
+73 -56
View File
@@ -40,6 +40,7 @@
#include "cutlass/epilogue/collective/detail.hpp"
#include "cutlass/epilogue/thread/scale_type.h"
#include "cutlass/epilogue/fusion/callbacks.hpp"
#include "cutlass/detail/layout.hpp"
#include "cutlass/trace.h"
#include "cute/tensor.hpp"
@@ -119,40 +120,52 @@ public:
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]");
private:
using InternalElementC = cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
using SmemElementC = cute::conditional_t<cute::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
constexpr static int StagesC = StagesC_;
constexpr static int StagesD = StagesD_;
constexpr static bool ReuseSmemC = ReuseSmemC_;
constexpr static bool is_source_supported = not cute::is_void_v<ElementC>;
// internal optimization to reuse C shared memory for storing D
using SmemLayoutAtomBitsC = decltype(downcast<sizeof_bits<InternalElementC>::value>(SmemLayoutAtomC{}));
using SmemLayoutAtomBitsD = decltype(downcast<sizeof_bits<ElementD>::value>(SmemLayoutAtomD{}));
constexpr static bool support_smem_reuse = is_source_supported &&
sizeof(InternalElementC) == sizeof(ElementD) &&
StrideC{} == StrideD{} &&
StagesD <= StagesC &&
cute::is_same_v<SmemLayoutAtomBitsC,SmemLayoutAtomBitsD>;
constexpr static bool ReuseSmemC = DispatchPolicy::ReuseSmemC;
static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met");
constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
public:
using SmemLayoutC = decltype(tile_to_shape(
SmemLayoutAtomC{},
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesC>{}),
cute::conditional_t<is_m_major_C, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
using SmemLayoutD = decltype(tile_to_shape(
SmemLayoutAtomD{},
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesD>{}),
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<ReuseSmemC ? StagesC : StagesD>{}),
cute::conditional_t<is_m_major_D, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
constexpr static bool support_smem_reuse = is_source_supported && StagesD <= StagesC
&& cosize(take<0,2>(SmemLayoutC{})) == cosize(take<0,2>(SmemLayoutD{}));
static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met");
constexpr static size_t SmemAlignmentD = cutlass::detail::alignment_for_swizzle(SmemLayoutD{});
constexpr static size_t SmemAlignmentC = cutlass::detail::alignment_for_swizzle(SmemLayoutC{});
struct TensorStorageWithC {
alignas(SmemAlignmentC) array_aligned<SmemElementC, size(SmemLayoutC{})> smem_C;
alignas(SmemAlignmentD) array_aligned<ElementD, size(SmemLayoutD{})> smem_D;
using FusionStorage = typename FusionCallbacks::SharedStorage;
FusionStorage thread;
};
struct TensorStorageWithoutC {
alignas(SmemAlignmentD) array_aligned<ElementD, size(SmemLayoutD{})> smem_D;
using FusionStorage = typename FusionCallbacks::SharedStorage;
FusionStorage thread;
};
public:
// TMA pipeline for loading C
using LoadPipeline = cutlass::PipelineTransactionAsync<StagesC>;
using LoadPipelineState = cutlass::PipelineState<StagesC>;
constexpr static uint32_t TmaTransactionBytes =
size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof(InternalElementC));
size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof(SmemElementC));
// TMA pipeline for storing D
using StorePipeline = cute::conditional_t<ReuseSmemC,
@@ -161,17 +174,9 @@ public:
using StorePipelineState = cutlass::PipelineState<ReuseSmemC ? StagesC : StagesD>;
struct SharedStorage {
struct TensorStorage : aligned_struct<128> {
cute::conditional_t<not is_source_supported,
detail::EmptyStorage<InternalElementC>,
array_aligned<InternalElementC, size(SmemLayoutC{})>> smem_C;
alignas(128) cute::conditional_t<ReuseSmemC,
detail::EmptyStorage<ElementD>,
array_aligned<ElementD, size(SmemLayoutD{})>> smem_D;
using FusionStorage = typename FusionCallbacks::SharedStorage;
alignas(128) FusionStorage thread;
} tensors;
using TensorStorage =
cute::conditional_t<not is_source_supported or ReuseSmemC, TensorStorageWithoutC, TensorStorageWithC>;
TensorStorage tensors;
using PipelineStorage = typename LoadPipeline::SharedStorage;
PipelineStorage pipeline;
@@ -192,7 +197,7 @@ public:
struct Params {
using TMA_C = decltype(make_tma_copy(
CopyOpG2S{},
make_tensor(static_cast<InternalElementC const*>(nullptr),
make_tensor(static_cast<SmemElementC const*>(nullptr),
repeat_like(StrideC{}, int32_t(0)), StrideC{}),
SmemLayoutC{}(_,_,0)));
using TMA_D = decltype(make_tma_copy(
@@ -316,21 +321,22 @@ public:
int thread_idx,
TensorStorage& shared_tensors) {
using namespace cute;
using _X = Underscore;
// Indexing variables
auto [M, N, K, L] = problem_shape_mnkl;
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
// Represent the full source tensor, slice to get the tile this CTA is currently responsible for
Tensor mC_mnl = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gC_mnl = local_tile(mC_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1,_X>{}); // (CTA_M,CTA_N,m,n,l)
Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord); // (CTA_M,CTA_N)
Tensor mC = params.tma_load_c.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gC = local_tile(mC, take<0,2>(CtaTileMNK{}), make_coord(m_coord,n_coord,l_coord)); // (CTA_M,CTA_N)
// Apply epilogue subtile, get matching smem tensor
auto ptr_sC = make_smem_ptr(shared_tensors.smem_C.data());
SmemElementC* ptr_sC = reinterpret_cast<SmemElementC*>(shared_tensors.smem_D.data());
if constexpr (not ReuseSmemC and is_source_supported) {
ptr_sC = shared_tensors.smem_C.data();
}
Tensor gC_epi = local_tile(gC, EpilogueTile{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sC_epi = make_tensor(ptr_sC, SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sC_epi = make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
// Prepare the thread(b)lock's (G)mem to (S)mem TMA tiled copy (bGS_)
ThrCopy thrblk_g2s = params.tma_load_c.get_slice(Int<0>{});
@@ -420,8 +426,9 @@ public:
int thread_idx,
TensorStorage& shared_tensors) {
using namespace cute;
using _X = Underscore;
using ElementAccumulator = typename AccEngine::value_type;
using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::ElementCompute;
using ElementCompute = cute::conditional_t<cute::is_void_v<ElementCompute_>,ElementAccumulator,ElementCompute_>;
static_assert(is_rmem<AccEngine>::value, "Accumulator must be RF resident.");
static_assert(rank(AccLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)");
@@ -439,16 +446,22 @@ public:
auto epi_tile_n = size<1>(EpilogueTile{});
// Represent the full output tensor, slice to get the tile this CTA is responsible for
Tensor mD_mnl = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gD_mnl = local_tile(mD_mnl, tile_shape_MNK, make_coord(_,_,_), Step<_1,_1,_X>{}); // (CTA_M,CTA_N,m,n,l)
Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord); // (CTA_M,CTA_N)
// Apply epilogue subtiling, construct corresponding pipelined smem tensors
auto ptr_sC = make_smem_ptr(shared_tensors.smem_C.data());
auto ptr_sD = make_smem_ptr(shared_tensors.smem_D.data());
Tensor mD = params.tma_store_d.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gD = local_tile(mD, take<0,2>(CtaTileMNK{}), make_coord(m_coord,n_coord,l_coord)); // (CTA_M,CTA_N)
// Apply epilogue subtiling
Tensor gD_epi = local_tile(gD, EpilogueTile{}, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sC_epi = make_tensor(ptr_sC, SmemLayoutC{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sD_epi = make_tensor(ptr_sD, SmemLayoutD{}); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
// Construct the corresponding pipelined smem tensors
SmemElementC* ptr_sC = reinterpret_cast<SmemElementC*>(shared_tensors.smem_D.data());
if constexpr (not ReuseSmemC and is_source_supported) {
ptr_sC = shared_tensors.smem_C.data();
}
ElementD* ptr_sD = shared_tensors.smem_D.data();
Tensor sC_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sC), SmemLayoutC{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
Tensor sD_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sD), SmemLayoutD{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
// Get the smallest tiled copy we can use to retile the accumulators
using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, cutlass::half_t>;
@@ -458,14 +471,11 @@ public:
TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom<CopyOpR2S,ElementD>{}, tiled_copy_C_atom);
ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx);
Tensor tRS_rAcc = thread_r2s.retile_S(accumulators); // ((R2S,R2S_V),MMA_M,MMA_N)
Tensor tRS_sD = conditional_return<ReuseSmemC>(
thread_r2s.partition_D(recast<ElementD>(sC_epi)), // (R2S,R2S_M,R2S_N,PIPE_C)
thread_r2s.partition_D(sD_epi) ); // (R2S,R2S_M,R2S_N,PIPE_D)
Tensor tRS_sD = thread_r2s.partition_D(sD_epi); // (R2S,R2S_M,R2S_N,PIPE_D)
// Allocate register tensors
auto tRS_rD_shape = take<0,3>(shape(thread_r2s.partition_S(sD_epi)));
Tensor tRS_rC = make_tensor<InternalElementC>(tRS_rD_shape); // (R2S,R2S_M,R2S_N)
Tensor tRS_rD = make_tensor<ElementD >(tRS_rD_shape); // (R2S,R2S_M,R2S_N)
// Allocate D registers
Layout tRS_rD_layout = make_layout(take<0,3>(shape(thread_r2s.partition_S(sD_epi))));
Tensor tRS_rD = make_tensor<ElementD>(tRS_rD_layout); // (R2S,R2S_M,R2S_N)
// Vectorized fragment view
constexpr int FragmentSize = DispatchPolicy::FragmentSize;
@@ -474,16 +484,23 @@ public:
CUTE_STATIC_ASSERT(size<0>(tRS_rAcc) % FragmentSize == 0, "Fragment size does not vectorize properly");
// (t)hread-partition for (s)mem to (r)egister copy (tSR_)
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R,InternalElementC>{}, tiled_copy_C_atom);
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R, SmemElementC>{}, tiled_copy_C_atom);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C)
Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N)
Tensor tSR_sC = thread_s2r.partition_S(sC_epi); // (S2R,S2R_M,S2R_N,PIPE_C)
Layout tSR_rC_layout = thread_s2r.retile_D(tRS_rD).layout(); // (S2R,S2R_M,S2R_N)
// Allocate C registers
// If C smem load is a non-vectorized dst(i) = src(i) then we can allocate C registers directly in the compute type
// to eliminate some redundant pack+unpack instruction sequences for sub-word types
constexpr bool IsDirectS2R = cute::is_same_v<CopyOpS2R,DefaultCopy>
&& decltype(max_common_vector(tSR_rC_layout, tSR_sC.layout()))::value <= 1;
using RegisterElementC = cute::conditional_t<IsDirectS2R, ElementCompute, SmemElementC>;
Tensor tRS_rC = make_tensor<RegisterElementC>(tRS_rD_layout); // (R2S,R2S_M,R2S_N)
Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N)
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
ThrCopy thrblk_s2g = params.tma_store_d.get_slice(Int<0>{});
Tensor bSG_sD = conditional_return<ReuseSmemC>(
thrblk_s2g.partition_S(recast<ElementD>(sC_epi)), // (S2G,S2G_M,S2G_N,PIPE_C)
thrblk_s2g.partition_S(sD_epi) ); // (S2G,S2G_M,S2G_N,PIPE_D)
Tensor bSG_sD = thrblk_s2g.partition_S(sD_epi); // (S2G,S2G_M,S2G_N,PIPE_D)
Tensor bSG_gD = thrblk_s2g.partition_D(gD_epi); // (S2G,S2G_M,S2G_N,EPI_M,EPI_N)
CUTE_STATIC_ASSERT(mma_tile_m == epi_tile_m, "EPI_TILE_M must equal MMA_TILE_M");
include/cutlass/epilogue/fusion/callbacks.hpp
+2
View File
@@ -62,6 +62,7 @@ struct FusionCallbacksTraits {
using Operation = T;
using CtaTile_MNK = void;
using EpilogueTile_MN = void;
using ElementCompute = void;
};
template <
@@ -78,6 +79,7 @@ struct FusionCallbacksTraits<
using Operation = Operation_;
using CtaTile_MNK = CtaTile_MNK_;
using EpilogueTile_MN = EpilogueTile_MN_;
using ElementCompute = typename Operation::ElementCompute;
};
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/fusion/operations.hpp
+3 -5
View File
@@ -60,7 +60,7 @@ struct FusionOperation {
using ElementBias = void;
static constexpr int AlignmentBias = 0;
static constexpr bool IsPerRowBiasSupported = false;
template <class> using ActivationFn = void;
using ActivationFn = void;
static constexpr bool IsEltActSupported = false;
using ElementAux = void;
@@ -108,8 +108,7 @@ template<
>
struct LinCombEltAct
: LinearCombination<ElementOutput_, ElementCompute_, ElementScalar_, RoundStyle_> {
template <class T>
using ActivationFn = ActivationFn_<T>;
using ActivationFn = ActivationFn_<ElementCompute_>;
static constexpr bool IsEltActSupported = true;
};
@@ -142,8 +141,7 @@ template<
struct LinCombPerRowBiasEltAct
: LinCombPerRowBias<ElementOutput_, ElementCompute_,
ElementBias_, ElementScalar_, AlignmentBias_, RoundStyle_> {
template <class T>
using ActivationFn = ActivationFn_<T>;
using ActivationFn = ActivationFn_<ElementCompute_>;
static constexpr bool IsEltActSupported = true;
};
include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp
+63 -25
View File
@@ -217,6 +217,9 @@ struct FusionCallbacks<
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
{ // unary op: activation(beta * C + (alpha * acc))
@@ -230,7 +233,7 @@ struct FusionCallbacks<
}, // end binary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args: activation
activation // unary args: activation
}; // end unary op
}
};
@@ -258,7 +261,7 @@ using Sm90LinCombPerRowBias =
Sm90EVT<Sm90Compute<multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ScalarBroadcast<ElementScalar>, // alpha
Sm90AccFetch, // acc
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,_0>, AlignmentBias> // bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias
>
>;
@@ -293,7 +296,10 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
operator typename Impl::Arguments() const {
return
@@ -303,7 +309,7 @@ struct FusionCallbacks<
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
@@ -373,7 +379,13 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
@@ -384,12 +396,12 @@ struct FusionCallbacks<
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args : activation
activation // unary args : activation
}; // end unary op
}
};
@@ -461,10 +473,9 @@ struct FusionCallbacks<
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
> {
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
using Impl =
Sm90LinCombPerRowBiasEltActAux<
CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>, SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
using Operation =
@@ -478,7 +489,15 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
ElementAux* aux_ptr = nullptr;
StrideAux dAux = {};
@@ -492,14 +511,14 @@ struct FusionCallbacks<
{ // ternary op : alpha * acc + bias
{{alpha}, {alpha_ptr}}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{aux_ptr, dAux} // unary args : store
}, // end unary op
{} // unary args : activation
activation // unary args : activation
}; // end unary op
}
};
@@ -528,7 +547,7 @@ using Sm90PerRowLinCombPerRowBias =
Sm90EVT<Sm90Compute<multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,_0>, AlignmentScalar>, // alpha
Sm90AccFetch, // acc
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,_0>, AlignmentBias> // bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias
>
>;
@@ -591,7 +610,13 @@ struct FusionCallbacks<
ElementScalar beta = ElementScalar(0);
ElementScalar const* alpha_ptr = nullptr;
ElementScalar const* beta_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
@@ -600,14 +625,14 @@ struct FusionCallbacks<
{beta_ptr, beta}, // leaf args : beta
{}, // leaf args : C
{ // ternary op : alpha * acc + bias
{alpha_ptr, alpha}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
{alpha_ptr, alpha}, // leaf args : alpha
{}, // leaf args : acc
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args : activation
activation // unary args : activation
}; // end unary op
}
};
@@ -650,7 +675,7 @@ using Sm90ScaledLinCombPerRowBias =
Sm90EVT<Sm90Compute<multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,_0>, 3>, // scale_a * scale_b * alpha
Sm90AccFetch, // acc
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,_0>, AlignmentBias> // bias
Sm90ColBroadcast<0, CtaTileShapeMNK, ElementBias, Stride<_1,_0,int>, AlignmentBias> // bias
>
>;
@@ -728,7 +753,12 @@ struct FusionCallbacks<
ElementScalar const* scale_c_ptr = nullptr;
ElementScalar const* scale_d_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
operator typename Impl::Arguments() const {
return
@@ -742,14 +772,14 @@ struct FusionCallbacks<
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{scale_a, scale_b, alpha},
{scale_a_ptr, scale_b_ptr, alpha_ptr}
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
}, // end ternary op
{} // unary args : activation
activation // unary args : activation
}, // end unary op
{{scale_d},
{scale_d_ptr}
@@ -855,10 +885,10 @@ struct FusionCallbacks<
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
> {
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
using Impl =
Sm90ScaledLinCombPerRowBiasEltActAmaxAux<
CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>,
SmemLayoutAtom, CopyOpR2S, ActivationFn,
ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
>;
using Operation =
@@ -885,9 +915,17 @@ struct FusionCallbacks<
ElementScalar scale_aux = ElementScalar(1);
ElementScalar const* scale_aux_ptr = nullptr;
using StrideBias = Stride<_1,_0,int>;
ElementBias const* bias_ptr = nullptr;
StrideBias dBias = {};
using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
ActivationArguments activation = ActivationArguments();
ElementAmax* amax_D_ptr = nullptr;
ElementAmax* amax_aux_ptr = nullptr;
using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
ElementAux* aux_ptr = nullptr;
StrideAux dAux = {};
@@ -905,9 +943,9 @@ struct FusionCallbacks<
{ // ternary op : (scale_a * scale_b * alpha) * acc + bias
{{scale_a, scale_b, alpha},
{scale_a_ptr, scale_b_ptr, alpha_ptr}
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr}, // leaf args : bias
}, // leaf args : (scale_a * scale_b * alpha)
{}, // leaf args : acc
{bias_ptr, ElementBias(0), dBias}, // leaf args : bias
{} // ternary args : multiply_add
}, // end ternary op
{} // ternary args : multiply_add
@@ -924,7 +962,7 @@ struct FusionCallbacks<
{ // unary op : reduce(activation(Z))
{ // unary op : activation(Z)
{}, // leaf args : Z
{} // unary args : activation
activation // unary args : activation
}, // end unary op
{amax_D_ptr_} // unary args : reduce
}, // end unary op
include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp
+121
View File
@@ -99,7 +99,9 @@ struct Sm90Compute : Sm90VisitorImpl<> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -118,6 +120,123 @@ struct Sm90Compute : Sm90VisitorImpl<> {
};
// partial specialization for compute fns that define an Arguments member, e.g. activation hyperparameters
template<
template <class> class ComputeFn,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle
>
struct Sm90Compute<
ComputeFn,
ElementOutput,
ElementCompute,
RoundStyle,
cute::void_t<typename ComputeFn<ElementCompute>::Arguments>
> {
struct SharedStorage { };
using Arguments = typename ComputeFn<ElementCompute>::Arguments;
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
CUTLASS_DEVICE bool
is_producer_load_needed() const {
return false;
}
CUTLASS_DEVICE bool
is_C_load_needed() const {
return false;
}
CUTLASS_HOST_DEVICE
Sm90Compute() { }
CUTLASS_HOST_DEVICE
Sm90Compute(Params const& params, SharedStorage const& shared_storage)
: params(params) {}
Params const params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
get_producer_load_callbacks(
ProblemShapeMNKL problem_shape_mnkl,
TileShapeMNK tile_shape_mnk,
TileCoordMNKL tile_coord_mnkl,
EpilogueTile epi_tile,
int thread_idx) {
return EmptyProducerLoadCallbacks{};
}
struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
CUTLASS_DEVICE
ConsumerStoreCallbacks(Params const& params)
: params(params) {}
Params const& params;
template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
CUTLASS_DEVICE Array<ElementOutput, FragmentSize>
visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n,
Array<ElementInputs, FragmentSize> const&... frg_inputs) {
return transform_apply(cute::make_tuple(frg_inputs...),
[&] (auto&& frg_input) {
using ElementInput = typename cute::remove_cvref_t<decltype(frg_input)>::Element;
using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
return convert_input(frg_input);
},
[&] (auto&&... cvt_frg_inputs) {
using ComputeOutput = ComputeFn<Array<ElementCompute, FragmentSize>>;
using ConvertOutput = NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize, RoundStyle>;
ComputeOutput compute_output{};
ConvertOutput convert_output{};
return convert_output(compute_output(cvt_frg_inputs..., params));
}
);
}
};
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
>
CUTLASS_DEVICE auto
get_consumer_store_callbacks(
ProblemShapeMNKL problem_shape_mnkl,
TileShapeMNK tile_shape_mnk,
TileCoordMNKL tile_coord_mnkl,
EpilogueTile epi_tile,
TiledCopy tiled_copy,
int thread_idx,
SrcTensor const& tCrC) {
return ConsumerStoreCallbacks(params);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Performance Optimized Specializations
@@ -215,7 +334,9 @@ struct Sm90TreeVisitor<
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
+45 -15
View File
@@ -71,7 +71,9 @@ struct Sm90AccFetch : Sm90VisitorImpl<> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -129,7 +131,9 @@ struct Sm90SrcFetch : Sm90VisitorImpl<> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -181,7 +185,8 @@ struct Sm90AuxLoad {
cute::conditional_t<is_m_major, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
struct SharedStorage {
alignas(128) array_aligned<Element, size(SmemLayout{})> smem_aux;
alignas(cutlass::detail::alignment_for_swizzle(SmemLayout{}))
array_aligned<Element, size(SmemLayout{})> smem_aux;
};
struct Arguments {
@@ -222,9 +227,9 @@ struct Sm90AuxLoad {
Sm90AuxLoad() { }
CUTLASS_HOST_DEVICE
Sm90AuxLoad(Params const& params, SharedStorage& shared_storage)
Sm90AuxLoad(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params),
smem_aux(shared_storage.smem_aux.data()) { }
smem_aux(const_cast<Element*>(shared_storage.smem_aux.data())) { }
Params const* params_ptr;
Element* smem_aux;
@@ -273,7 +278,9 @@ struct Sm90AuxLoad {
};
template <
class TileShapeMNK
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL
>
CUTLASS_DEVICE auto
get_producer_load_callbacks(
@@ -284,8 +291,9 @@ struct Sm90AuxLoad {
int thread_idx) {
auto [M, N, K, L] = problem_shape_mnkl;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mAux = params_ptr->tma_load_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gAux = sm90_tensor_to_cta_tile(mAux, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
Tensor gAux = local_tile(mAux, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor gAux_epi = local_tile(gAux, epi_tile, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE)
@@ -339,7 +347,9 @@ struct Sm90AuxLoad {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class TiledCopy,
class SrcTensor
>
@@ -363,7 +373,8 @@ struct Sm90AuxLoad {
make_tiled_copy_S(Copy_Atom<CopyOpS2R,Element>{}, tiled_copy),
make_tiled_copy_D(Copy_Atom<CopyOpS2R,Element>{}, tiled_copy)
);
Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE)
Tensor sAux_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(smem_aux), SmemLayout{})); // (EPI_TILE_M,EPI_TILE_N,PIPE)
auto tSR_sAux = tiled_s2r.get_slice(thread_idx).partition_S(sAux_epi); // (S2R,S2R_M,S2R_N,PIPE)
@@ -378,6 +389,7 @@ struct Sm90AuxLoad {
/////////////////////////////////////////////////////////////////////////////////////////////////
// Scalar broadcast
// Supports reduction over multiple broadcasts to support fusions such as fp8 scaling factors
template<
class Element,
class StrideMNL = Stride<_0,_0,_0>,
@@ -387,7 +399,8 @@ template<
struct Sm90ScalarBroadcast {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0, _0>>) || // scalar broadcast, e.g. alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>)); // batched scalar broadcast, e.g. per-batch alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0, _1>>) || // batched scalar broadcast, e.g. per-batch alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>));
struct SharedStorage { };
@@ -419,7 +432,7 @@ struct Sm90ScalarBroadcast {
Sm90ScalarBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90ScalarBroadcast(Params const& params, SharedStorage& shared_storage)
Sm90ScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) {
// Get the scalar for non-batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,_0>>) {
@@ -431,7 +444,9 @@ struct Sm90ScalarBroadcast {
Params const* params_ptr;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -442,7 +457,9 @@ struct Sm90ScalarBroadcast {
EpilogueTile epi_tile,
int thread_idx) {
// Get the scalar for batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
if constexpr (
cute::is_same_v<StrideMNL, Stride<_0,_0,_1>> ||
cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
update_scalar(l_coord);
}
@@ -470,7 +487,9 @@ struct Sm90ScalarBroadcast {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -486,7 +505,9 @@ struct Sm90ScalarBroadcast {
SrcTensor const& tCrC) {
// Get the scalar for batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
if constexpr (
cute::is_same_v<StrideMNL, Stride<_0,_0,_1>> ||
cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
update_scalar(l_coord);
}
@@ -541,7 +562,7 @@ struct Sm90RowBroadcast {
// Accumulator doesn't distribute row elements evenly amongst threads so we must buffer in smem
struct SharedStorage {
array_aligned<Element, size<1>(CtaTileShapeMNK{}) * Stages> smem_row;
alignas(16) array_aligned<Element, size<1>(CtaTileShapeMNK{}) * Stages> smem_row;
};
struct Arguments {
@@ -562,9 +583,9 @@ struct Sm90RowBroadcast {
Sm90RowBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90RowBroadcast(Params const& params, SharedStorage& shared_storage)
Sm90RowBroadcast(Params const& params, SharedStorage const& shared_storage)
: params(params),
smem_row(shared_storage.smem_row.data()) { }
smem_row(const_cast<Element*>(shared_storage.smem_row.data())) { }
Params params;
Element* smem_row;
@@ -613,7 +634,9 @@ struct Sm90RowBroadcast {
};
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -625,8 +648,9 @@ struct Sm90RowBroadcast {
int thread_idx) {
auto [M, N, K, L] = problem_shape_mnkl;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
Tensor gRow = sm90_tensor_to_cta_tile(mRow, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
Tensor gRow = local_tile(mRow, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor sRow = make_tensor(make_smem_ptr(smem_row), // (CTA_M,CTA_N,PIPE)
make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{}), Stages),
make_stride(_0{},_1{},size<1>(CtaTileShapeMNK{})));
@@ -680,7 +704,9 @@ struct Sm90RowBroadcast {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -757,13 +783,15 @@ struct Sm90ColBroadcast {
Sm90ColBroadcast() { }
CUTLASS_HOST_DEVICE
Sm90ColBroadcast(Params const& params, SharedStorage& shared_storage)
Sm90ColBroadcast(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -819,7 +847,9 @@ struct Sm90ColBroadcast {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp
+32 -12
View File
@@ -83,7 +83,8 @@ struct Sm90AuxStore {
cute::conditional_t<is_m_major, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
struct SharedStorage {
alignas(128) array_aligned<Element, size(SmemLayout{})> smem_aux;
alignas(cutlass::detail::alignment_for_swizzle(SmemLayout{}))
array_aligned<Element, size(SmemLayout{})> smem_aux;
};
struct Arguments {
@@ -125,9 +126,9 @@ struct Sm90AuxStore {
Sm90AuxStore() { }
CUTLASS_HOST_DEVICE
Sm90AuxStore(Params const& params, SharedStorage& shared_storage)
Sm90AuxStore(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params),
smem_aux(shared_storage.smem_aux.data()) { }
smem_aux(const_cast<Element*>(shared_storage.smem_aux.data())) { }
Params const* params_ptr;
Element* smem_aux;
@@ -143,7 +144,9 @@ struct Sm90AuxStore {
}
template <
class TileShapeMNK
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL
>
CUTLASS_DEVICE auto
get_producer_load_callbacks(
@@ -233,7 +236,9 @@ struct Sm90AuxStore {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class TiledCopy,
class SrcTensor
>
@@ -248,14 +253,16 @@ struct Sm90AuxStore {
SrcTensor const& tCrC) {
auto [M, N, K, L] = problem_shape_mnkl;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mAux = params_ptr->tma_store_aux.get_tma_tensor(make_shape(M,N,L)); // (M,N,L)
Tensor gAux = sm90_tensor_to_cta_tile(mAux, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
Tensor gAux = local_tile(mAux, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor tC_gAux = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
gAux, epi_tile, tiled_copy, thread_idx);
Tensor tC_rAux = make_tensor<Element>(take<0,3>(shape(tC_gAux))); // (CPY,CPY_M,CPY_N)
Tensor sAux_epi = make_tensor(make_smem_ptr(smem_aux), SmemLayout{}); // (EPI_TILE_M,EPI_TILE_N,PIPE)
Tensor sAux_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(smem_aux), SmemLayout{})); // (EPI_TILE_M,EPI_TILE_N,PIPE)
Tensor gAux_epi = local_tile(gAux, epi_tile, _); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
auto tiled_r2s = conditional_return<ReferenceSrc>(
@@ -297,8 +304,8 @@ template <
struct Sm90ScalarReduction {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0, _0>>) || // scalar reduction, e.g. tensor max element
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>)); // batched scalar reduction, e.g. per-batch max element
(cute::is_same_v<StrideMNL, Stride<_0,_0, _1>>) || // batched scalar reduction, e.g. per-batch max element
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>));
struct SharedStorage { };
struct Arguments {
@@ -329,13 +336,15 @@ struct Sm90ScalarReduction {
Sm90ScalarReduction() { }
CUTLASS_HOST_DEVICE
Sm90ScalarReduction(Params const& params, SharedStorage& shared_storage)
Sm90ScalarReduction(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params const params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -417,7 +426,9 @@ struct Sm90ScalarReduction {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -502,13 +513,15 @@ struct Sm90RowReduction {
Sm90RowReduction() { }
CUTLASS_HOST_DEVICE
Sm90RowReduction(Params const& params, SharedStorage& shared_storage)
Sm90RowReduction(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -619,7 +632,9 @@ struct Sm90RowReduction {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -707,13 +722,15 @@ struct Sm90ColReduction {
Sm90ColReduction() { }
CUTLASS_HOST_DEVICE
Sm90ColReduction(Params const& params, SharedStorage& shared_storage)
Sm90ColReduction(Params const& params, SharedStorage const& shared_storage)
: params(params) { }
Params params;
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -765,10 +782,11 @@ struct Sm90ColReduction {
Array frg_I = convert_input(frg_input);
Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
Tensor tCcCol_mn = tCcCol(_,_,_,epi_m,epi_n);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
if (elem_less(tCcCol(i), residue_mn)) {
if (elem_less(tCcCol_mn(i), residue_mn)) {
ElementCompute& tCrCol_vmn = tCrCol_mn(epi_v * FragmentSize + i);
tCrCol_vmn = reduce_input(tCrCol_vmn, frg_I[i]);
}
@@ -808,7 +826,9 @@ struct Sm90ColReduction {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp
+48 -56
View File
@@ -51,34 +51,12 @@ using cute::tuple;
namespace detail {
// Convenience aliases
using ProblemShapeMNKL = tuple<int,int,int,int>;
using TileCoordMNKL = tuple<int,int,int,int>;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Partitioning Helpers
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template <
class Engine, class LayoutMNL,
class TileShapeMNK
>
CUTLASS_HOST_DEVICE
constexpr auto
sm90_tensor_to_cta_tile(
Tensor<Engine, LayoutMNL> mT, // (M,N,L)
TileShapeMNK tile_shape_mnk, // (CTA_M,CTA_N,CTA_K)
TileCoordMNKL tile_coord_mnkl) {
using _X = Underscore;
auto [m, n, k, l] = tile_coord_mnkl;
Tensor mT_mnl = local_tile(mT, tile_shape_mnk, make_coord(_,_,_), Step<_1,_1,_X>{}); // (CTA_M,CTA_N)
return mT_mnl(_,_,m,n,l);
}
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class CtaTileMN,
@@ -106,6 +84,7 @@ template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class Engine, class LayoutMNL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy
>
@@ -118,7 +97,8 @@ sm90_partition_for_epilogue(
EpilogueTile epi_tile, // (EPI_TILE_M,EPI_TILE_N)
TiledCopy tiled_copy,
int thread_idx) {
Tensor cT = sm90_tensor_to_cta_tile(mT, tile_shape_mnk, tile_coord_mnkl); // (CTA_M,CTA_N)
auto [m, n, k, l] = tile_coord_mnkl;
Tensor cT = local_tile(mT, take<0,2>(tile_shape_mnk), make_coord(m,n,l)); // (CTA_M,CTA_N)
Tensor tCcT =
sm90_partition_for_epilogue<ReferenceSrc>(cT, epi_tile, tiled_copy, thread_idx); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
@@ -156,7 +136,7 @@ struct Sm90VisitorImplBase {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops(transform_apply(tuple<Ops...>{}, params, shared_storage,
[] (auto&& op, auto const& op_params, auto&& op_storage) {
using Op = cute::remove_cvref_t<decltype(op)>;
@@ -262,7 +242,9 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase<Ops...> {
// Producer load callbacks factory
// All operations must redefine this, but most can just dispatch to the base impl
template <
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile
>
CUTLASS_DEVICE auto
@@ -363,7 +345,9 @@ struct Sm90VisitorImpl : Sm90VisitorImplBase<Ops...> {
// All operations must redefine this
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -446,7 +430,9 @@ struct Sm90TreeVisitor : Sm90VisitorImpl<ChildOps..., NodeOp> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -516,7 +502,9 @@ struct Sm90SplitTreeVisitor : Sm90VisitorImpl<InputTree, AuxOutTrees..., OutputT
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -613,7 +601,9 @@ struct Sm90TopologicalVisitor : Sm90VisitorImpl<Ops...> {
template <
bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
class ProblemShapeMNKL,
class TileShapeMNK,
class TileCoordMNKL,
class EpilogueTile,
class TiledCopy,
class SrcTensor
@@ -651,9 +641,11 @@ namespace detail {
template <class Op0>
struct Sm90VisitorImplBase<Op0> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
};
// Retain tuple for SharedStorage because empty structs have 1B alignment
// tuples use multiple inheritance, avoids this problem
using SharedStorage = tuple<
typename Op0::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -675,9 +667,9 @@ struct Sm90VisitorImplBase<Op0> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0)
Op0(params.op_0, get<0>(shared_storage))
}) {}
tuple<Op0> ops;
@@ -686,10 +678,10 @@ struct Sm90VisitorImplBase<Op0> {
template <class Op0, class Op1>
struct Sm90VisitorImplBase<Op0, Op1> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
typename Op1::SharedStorage op_1;
};
using SharedStorage = tuple<
typename Op0::SharedStorage,
typename Op1::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -714,10 +706,10 @@ struct Sm90VisitorImplBase<Op0, Op1> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0),
Op1(params.op_1, shared_storage.op_1)
Op0(params.op_0, get<0>(shared_storage)),
Op1(params.op_1, get<1>(shared_storage))
}) {}
tuple<Op0, Op1> ops;
@@ -726,11 +718,11 @@ struct Sm90VisitorImplBase<Op0, Op1> {
template <class Op0, class Op1, class Op2>
struct Sm90VisitorImplBase<Op0, Op1, Op2> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
typename Op1::SharedStorage op_1;
typename Op2::SharedStorage op_2;
};
using SharedStorage = tuple<
typename Op0::SharedStorage,
typename Op1::SharedStorage,
typename Op2::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -758,11 +750,11 @@ struct Sm90VisitorImplBase<Op0, Op1, Op2> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0),
Op1(params.op_1, shared_storage.op_1),
Op2(params.op_2, shared_storage.op_2)
Op0(params.op_0, get<0>(shared_storage)),
Op1(params.op_1, get<1>(shared_storage)),
Op2(params.op_2, get<2>(shared_storage))
}) {}
tuple<Op0, Op1, Op2> ops;
@@ -771,12 +763,12 @@ struct Sm90VisitorImplBase<Op0, Op1, Op2> {
template <class Op0, class Op1, class Op2, class Op3>
struct Sm90VisitorImplBase<Op0, Op1, Op2, Op3> {
struct SharedStorage {
typename Op0::SharedStorage op_0;
typename Op1::SharedStorage op_1;
typename Op2::SharedStorage op_2;
typename Op3::SharedStorage op_3;
};
using SharedStorage = tuple<
typename Op0::SharedStorage,
typename Op1::SharedStorage,
typename Op2::SharedStorage,
typename Op3::SharedStorage
>;
struct Arguments {
typename Op0::Arguments op_0;
@@ -807,12 +799,12 @@ struct Sm90VisitorImplBase<Op0, Op1, Op2, Op3> {
Sm90VisitorImplBase() {}
CUTLASS_HOST_DEVICE
Sm90VisitorImplBase(Params const& params, SharedStorage& shared_storage)
Sm90VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
: ops({
Op0(params.op_0, shared_storage.op_0),
Op1(params.op_1, shared_storage.op_1),
Op2(params.op_2, shared_storage.op_2),
Op3(params.op_3, shared_storage.op_3)
Op0(params.op_0, get<0>(shared_storage)),
Op1(params.op_1, get<1>(shared_storage)),
Op2(params.op_2, get<2>(shared_storage)),
Op3(params.op_3, get<3>(shared_storage))
}) {}
tuple<Op0, Op1, Op2, Op3> ops;
include/cutlass/epilogue/thread/activation.h
+119 -250
View File
@@ -49,38 +49,6 @@ namespace cutlass {
namespace epilogue {
namespace thread {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <typename T>
struct LinearCombinationGenericParams {
T alpha; ///< scales accumulators
T beta; ///< scales source tensor
T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams():
alpha(T(1)),
beta(T(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T alpha,
T beta = T(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T const *alpha_ptr,
T const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Identity operator
@@ -92,13 +60,6 @@ struct Identity {
T operator()(T value) const {
return value;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T, int N>
@@ -107,12 +68,59 @@ struct Identity<Array<T, N> > {
Array<T, N> operator()(Array<T, N> const &value) const {
return value;
}
};
using Params = LinearCombinationGenericParams<T>;
/// Scale operator
template <typename T>
struct Scale {
struct Arguments {
T scale = T(1);
};
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
T operator()(T const& value, T const& scale) const {
multiplies<T> mul;
return mul(scale, value);
}
CUTLASS_HOST_DEVICE
T operator()(T const& value, Arguments const& args = Arguments()) const {
return this->operator()(value, args.scale);
}
};
template <typename T, int N>
struct Scale<Array<T, N>> {
using Arguments = typename Scale<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& values, T const& scale) const {
multiplies<Array<T, N>> mul;
return mul(scale, values);
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const& values, Arguments const& args = Arguments()) const {
return this->operator()(values, args.scale);
}
};
/// Specialization to compose other activations with a defined unary operator
/// e.g. Scale<Identity<T>>
template <template <class> class Activation, typename T>
struct Scale<Activation<T>> {
using Arguments = typename Scale<T>::Arguments;
CUTLASS_HOST_DEVICE
T operator()(T const &value, decltype(Arguments{}.scale) const& scale) const {
multiplies<T> mul;
Activation<T> act;
return mul(scale, act(value));
}
CUTLASS_HOST_DEVICE
T operator()(T const& value, Arguments const& args = Arguments()) const {
return this->operator()(value, args.scale);
}
};
@@ -134,14 +142,6 @@ struct ReLu {
return mx(value, T(0));
}
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T>
@@ -162,90 +162,87 @@ struct ReLu<Array<T, N>> {
maximum<Array<T, N>> mx;
return mx(frag, T(0));
}
};
/// Host-constructable parameters structure
using Params = LinearCombinationGenericParams<T>;
// Generic clamp
template <typename T>
struct Clamp {
struct Arguments {
T lower_bound = cutlass::platform::numeric_limits<T>::min();
T upper_bound = cutlass::platform::numeric_limits<T>::max();
};
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &frag, Params const &params_) const {
return this->operator()(frag);
T operator()(T const& value, T const& lower_bound, T const& upper_bound) const {
maximum<T> mx;
minimum<T> mn;
return mn(mx(value, lower_bound), upper_bound);
}
CUTLASS_HOST_DEVICE
T operator()(T const& value, Arguments const& args = Arguments()) const {
return this->operator()(value, args.lower_bound, args.upper_bound);
}
};
template <typename T, int N>
struct Clamp<Array<T,N>> {
using Arguments = typename Clamp<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T,N> operator()(Array<T,N> const& values, T const& lower_bound, T const& upper_bound) const {
maximum<Array<T,N>> mx;
minimum<Array<T,N>> mn;
return mn(mx(values, lower_bound), upper_bound);
}
CUTLASS_HOST_DEVICE
Array<T,N> operator()(Array<T,N> const& values, Arguments const& args = Arguments()) const {
return this->operator()(values, args.lower_bound, args.upper_bound);
}
};
// Leaky Relu operator
template <typename T>
struct LeakyReLU {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
// Methods
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
struct Arguments {
T leaky_alpha = T(0);
};
CUTLASS_HOST_DEVICE
T operator()(T const &value, T const & alpha_recip) const {
T res = value > T(0) ? value : value * alpha_recip;
T operator()(T const& value, T const& leaky_alpha) const {
T res = value > T(0) ? value : value * leaky_alpha;
return res;
}
CUTLASS_HOST_DEVICE
T operator()(T const &value, Params const &params_) const {
this->operator()(value, params_.leaky_alpha);
T operator()(T const& value, Arguments const& args = Arguments()) const {
this->operator()(value, args.leaky_alpha);
}
};
template <typename T, int N>
struct LeakyReLU<Array<T, N> > {
struct Params: LinearCombinationGenericParams<T> {
T leaky_alpha; ///< leaky_alpha
using LinearCombinationGenericParams<T>::LinearCombinationGenericParams;
// Methods
CUTLASS_HOST_DEVICE
Params():
LinearCombinationGenericParams<T>(),
leaky_alpha(T(1)) {}
CUTLASS_HOST_DEVICE
Params(
T alpha,
T beta,
T leaky_alpha = T(1)
): LinearCombinationGenericParams<T>(alpha, beta), leaky_alpha(leaky_alpha) {}
};
using Arguments = typename LeakyReLU<T>::Arguments;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, T const & alpha_recip) const {
Array<T, N> operator()(Array<T, N> const& values, T const& leaky_alpha) const {
Array<T, N> y;
LeakyReLU<T> leaky_op;
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < int(value.size()); ++i) {
y[i] = leaky_op(value[i], alpha_recip);
for (int i = 0; i < int(values.size()); ++i) {
y[i] = leaky_op(values[i], leaky_alpha);
}
return y;
}
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value, params_.leaky_alpha);
Array<T, N> operator()(Array<T, N> const& values, Arguments const& args = Arguments()) const {
return this->operator()(values, args.leaky_alpha);
}
};
@@ -253,15 +250,8 @@ struct LeakyReLU<Array<T, N> > {
template <typename T>
struct Tanh {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
return fast_tanh(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
T operator()(T const &value) const {
return fast_tanh(value);
}
};
@@ -279,13 +269,6 @@ struct Tanh<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <int N>
@@ -296,14 +279,6 @@ struct Tanh<Array<half_t, N>> {
Array<T, N> operator()(Array<T, N> const& z) const {
fast_tanh_op<Array<T, N>> tanh;
return tanh(z);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
@@ -311,15 +286,8 @@ struct Tanh<Array<half_t, N>> {
template <typename T>
struct Sigmoid {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
return T(1) / (T(1) + fast_exp(-scalar));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
T operator()(T const &value) const {
return T(1) / (T(1) + fast_exp(-value));
}
};
@@ -337,13 +305,6 @@ struct Sigmoid<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <int N>
@@ -368,13 +329,6 @@ struct Sigmoid<Array<half_t, N>> {
fast_exp(neg(z))));
#endif
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &z, Params const &params_) const {
return this->operator()(z);
}
};
// SiLu (swish) operator introduced by Elfwing et al. in the following paper
@@ -385,16 +339,9 @@ struct Sigmoid<Array<half_t, N>> {
template <typename T>
struct SiLu {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
T operator()(T const &value) const {
Sigmoid<T> sigmoid;
return scalar * sigmoid(scalar);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
return value * sigmoid(value);
}
};
@@ -406,13 +353,6 @@ struct SiLu<Array<T, N>> {
multiplies<Array<T, N>> mul;
return mul(value, sigmoid_op(value));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
// Hardswish operator introduced by Howard et al. in the following paper
@@ -429,13 +369,6 @@ struct HardSwish {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 / T(6);
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <>
@@ -449,13 +382,6 @@ struct HardSwish<float> {
T relu6 = mn(mx(x + T(3), T(0)), T(6));
return x * relu6 * 0.16666667f;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <typename T, int N>
@@ -472,13 +398,6 @@ struct HardSwish<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
template <int N>
@@ -494,13 +413,6 @@ struct HardSwish<Array<half_t, N> > {
return mul(mul(mn(mx(add(value, T(3)), T(0)), T(6)), value), T(0.16666667f));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &x, Params const &params_) const {
return this->operator()(x);
}
};
//
@@ -516,48 +428,27 @@ struct HardSwish<Array<half_t, N> > {
template <typename T>
struct GELU {
CUTLASS_HOST_DEVICE
T operator()(T const &scalar) const {
return T(cutlass::constants::half<T>() * scalar *
(cutlass::constants::one<T>() + (T)erff((float)(scalar * cutlass::constants::half_root_two<T>()))));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
T operator()(T const &value) const {
return T(cutlass::constants::half<T>() * value *
(cutlass::constants::one<T>() + (T)erff((float)(value * cutlass::constants::half_root_two<T>()))));
}
};
template <>
struct GELU<float> {
CUTLASS_HOST_DEVICE
float operator()(float const &scalar) const {
return cutlass::constants::half<float>() * scalar *
(cutlass::constants::one<float>() + erff(scalar * cutlass::constants::half_root_two<float>() ));
}
using Params = LinearCombinationGenericParams<float>;
CUTLASS_HOST_DEVICE
float operator()(float const &scalar, Params const &params_) const {
return this->operator()(scalar);
float operator()(float const &value) const {
return cutlass::constants::half<float>() * value *
(cutlass::constants::one<float>() + erff(value * cutlass::constants::half_root_two<float>() ));
}
};
template <>
struct GELU<double> {
CUTLASS_HOST_DEVICE
double operator()(double const &scalar) const {
return cutlass::constants::half<double>() * scalar *
(cutlass::constants::one<double>() + erf( scalar * cutlass::constants::half_root_two<double>() ));
}
using Params = LinearCombinationGenericParams<double>;
CUTLASS_HOST_DEVICE
double operator()(double const &scalar, Params const &params_) const {
return this->operator()(scalar);
double operator()(double const &value) const {
return cutlass::constants::half<double>() * value *
(cutlass::constants::one<double>() + erf( value * cutlass::constants::half_root_two<double>() ));
}
};
@@ -575,15 +466,11 @@ struct GELU<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T>
using ScaledGELU = Scale<GELU<T>>;
// GELU operator implemented using the Taylor series approximation
template <typename T>
struct GELU_taylor {
@@ -597,13 +484,6 @@ struct GELU_taylor {
return T(cutlass::constants::half<T>() * z *
(cutlass::constants::one<T>() + fast_tanh(k0 * z * (cutlass::constants::one<T>() + k1 * z * z))));
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
T operator()(T const &scalar, Params const &params_) const {
return this->operator()(scalar);
}
};
template <int N>
@@ -630,13 +510,6 @@ struct GELU_taylor<Array<half_t, N> > {
return y;
}
using Params = LinearCombinationGenericParams<half_t>;
CUTLASS_HOST_DEVICE
Array<half_t, N> operator()(Array<half_t, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T, int N>
@@ -654,15 +527,11 @@ struct GELU_taylor<Array<T, N> > {
return y;
}
using Params = LinearCombinationGenericParams<T>;
CUTLASS_HOST_DEVICE
Array<T, N> operator()(Array<T, N> const &value, Params const &params_) const {
return this->operator()(value);
}
};
template <typename T>
using ScaledGELU_taylor = Scale<GELU_taylor<T>>;
/// Computes backwards pass for GELU operator assuming d_t is the layer gradient and
/// z is computed from the forward pass.
template <typename T>
include/cutlass/epilogue/thread/linear_combination_generic.h
+60 -3
View File
@@ -49,6 +49,51 @@ namespace thread {
/////////////////////////////////////////////////////////////////////////////////////////////////
template <class Activation, class = void>
struct GenericActivationTraits {
static constexpr bool IsArgumentsNeeded = false;
struct Arguments {};
};
template <class Activation>
struct GenericActivationTraits<Activation, decltype(typename Activation::Arguments(), void())> {
static constexpr bool IsArgumentsNeeded = true;
using Arguments = typename Activation::Arguments;
};
template <typename T>
struct LinearCombinationGenericParams {
T alpha; ///< scales accumulators
T beta; ///< scales source tensor
T const *alpha_ptr; ///< pointer to accumulator scalar - if not null, loads it from memory
T const *beta_ptr; ///< pointer to source scalar - if not null, loads it from memory
//
// Methods
//
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams():
alpha(T(1)),
beta(T(0)),
alpha_ptr(nullptr),
beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T alpha,
T beta = T(0)
): alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
CUTLASS_HOST_DEVICE
LinearCombinationGenericParams(
T const *alpha_ptr,
T const *beta_ptr = nullptr
): alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
};
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Applies a linear combination operator followed by an activation function to an array of elements.
///
/// D = activation(alpha * accumulator + beta * source + uniform)
@@ -84,7 +129,11 @@ public:
static FloatRoundStyle const kRound = Round;
/// Host-constructable parameters structure
using Params = typename ActivationFunctor<FragmentCompute>::Params;
struct Params
: LinearCombinationGenericParams<ElementCompute>,
GenericActivationTraits<ActivationFunctor<ElementCompute>>::Arguments {
using LinearCombinationGenericParams<ElementCompute>::LinearCombinationGenericParams;
};
private:
@@ -161,7 +210,11 @@ public:
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator, intermediate); // D = alpha * Accum + X
}
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
if constexpr (GenericActivationTraits<ActivationFunctor<ElementCompute>>::IsArgumentsNeeded) {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
} else {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate);
}
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
@@ -192,7 +245,11 @@ public:
intermediate = mul_add_accumulator(params_.alpha, converted_accumulator); // D = alpha * Accum
}
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
if constexpr (GenericActivationTraits<ActivationFunctor<FragmentCompute>>::IsArgumentsNeeded) {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate, params_);
} else {
intermediate = skip_elementwise_ ? intermediate : activation(intermediate);
}
// Convert to destination numeric type
NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
include/cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h
+495
View File
@@ -0,0 +1,495 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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 Functor performing elementwise operations used by epilogues.
*/
#pragma once
/////////////////////////////////////////////////////////////////////////////////////////////////
#include "cutlass/epilogue/threadblock/epilogue_base.h"
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
namespace epilogue {
namespace threadblock {
////////////////////////////////////////////////////////////////////////////////////////////////////
/// Epilogue operator
template <
typename DefaultEpilogue, ///< Default Epilogue Descriptor
typename FusionCallbacks_, ///< The called fusion callbacks
int Stages = 2, ///< Software pipeline stages for epilogue
int IterationsUnroll = true ///< Used to reduce binary size when epilogue op is large
>
class EpilogueWithVisitorCallbacks :
public EpilogueBase<
typename DefaultEpilogue::Shape,
typename DefaultEpilogue::WarpMmaOperator::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::AccumulatorFragmentIterator,
typename DefaultEpilogue::WarpTileIterator,
typename DefaultEpilogue::Padding,
DefaultEpilogue::kFragmentsPerIteration>,
public EpilogueBaseStreamK<
typename DefaultEpilogue::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::WarpMmaOperator,
typename DefaultEpilogue::AccumulatorFragmentIterator>
{
public:
static_assert(Stages <= 2, "Sm80 EVT only support upto 2 Stages.");
// Whether the epilogue is pipelined
static bool constexpr Pipelined = Stages > 1;
using FusionCallbacks = FusionCallbacks_;
using OutputTileIterator = typename DefaultEpilogue::OutputTileIterator;
// Number of epilogue iterations.
// Each iteration processes a 8xThreadblockTile::kN output tile
static const int kIterations = OutputTileIterator::kIterations;
using Base = EpilogueBase<
typename DefaultEpilogue::Shape,
typename DefaultEpilogue::WarpMmaOperator::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::AccumulatorFragmentIterator,
typename DefaultEpilogue::WarpTileIterator,
typename DefaultEpilogue::Padding,
DefaultEpilogue::kFragmentsPerIteration>;
using BaseStreamK = EpilogueBaseStreamK<
typename DefaultEpilogue::Shape,
DefaultEpilogue::kPartitionsK,
typename DefaultEpilogue::WarpMmaOperator,
typename DefaultEpilogue::AccumulatorFragmentIterator>;
static int const kPartitionsK = DefaultEpilogue::kPartitionsK;
using AccumulatorFragmentIterator = typename DefaultEpilogue::AccumulatorFragmentIterator;
using WarpTileIterator = typename DefaultEpilogue::WarpTileIterator;
using SharedLoadIterator = typename DefaultEpilogue::SharedLoadIterator;
/// The complete warp-level accumulator tile
using AccumulatorTile = typename Base::AccumulatorTile;
/// Accumulator element
using ElementAccumulator = typename WarpTileIterator::Element;
struct OutputOp{
using ElementAccumulator = ElementAccumulator;
using Params = typename FusionCallbacks::Arguments;
};
/// Fragment type used by the accumulator tile's fragment iterator
using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment;
// Output access size
static int const kElementsPerAccess = DefaultEpilogue::kElementsPerAccess;
/// Array type used by output functor
using AccumulatorAccessType = Array<
typename WarpTileIterator::Element, kElementsPerAccess>;
static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1 ? Base::kFragmentsPerIteration : kPartitionsK;
static int constexpr kSmemPointerOffset = Base::SharedStorage::StorageShape::kCount / kSmemTiles;
using Params = typename FusionCallbacks::Params;
static size_t constexpr kSmemStageOffset = sizeof(Base::SharedStorage) / sizeof(ElementAccumulator);
static int constexpr kAccumulatorFragmentCount = AccumulatorTile::kElements / (kIterations * AccumulatorAccessType::kElements) / kPartitionsK;
struct SharedStorage {
typename Base::SharedStorage acc_smem[Stages];
typename FusionCallbacks::SharedStorage callback_smem;
};
private:
/// Loads fragment from shared memory aligned with output tensor
SharedLoadIterator shared_load_iterator_;
FusionCallbacks fusion_callbacks;
public:
/// Constructor
CUTLASS_DEVICE
EpilogueWithVisitorCallbacks(
const Params &params_callbacks, ///< Epilogue Visitor params
SharedStorage &shared_storage, ///< Shared storage object
int thread_idx, ///< ID of a thread within the threadblock
int warp_idx, ///< ID of warp within threadblock
int lane_idx ///< Id of thread within warp
):
Base(shared_storage.acc_smem[0], thread_idx, warp_idx, lane_idx),
BaseStreamK(thread_idx),
shared_load_iterator_(shared_storage.acc_smem[0].reference(), thread_idx),
fusion_callbacks(params_callbacks, shared_storage.callback_smem)
{ }
/// Aggregates the accumulator sets shared by peer blocks in the global workspace,
/// performing epilogue computations, writing to output
template <class ProblemShape>
CUTLASS_DEVICE
void reduce(
int peer_idx_begin,
int peer_idx_end,
int reduce_fragment_idx,
void *element_workspace,
cutlass::gemm::GemmCoord threadblock_tile_offset,
ProblemShape problem_shape,
int thread_idx)
{
auto callbacks = fusion_callbacks.get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
);
callbacks.begin_epilogue();
// Reduce peer accumulator fragments into one fragment
AccumulatorFragment accum_fragment;
BaseStreamK::reduce(accum_fragment, peer_idx_begin, peer_idx_end, reduce_fragment_idx, element_workspace);
// Store fragment to shared memory
this->warp_tile_iterator_.store(accum_fragment);
__syncthreads();
callbacks.begin_step(reduce_fragment_idx);
// Load fragment from shared memory
typename SharedLoadIterator::Fragment aligned_accum_fragment;
shared_load_iterator_.load(aligned_accum_fragment);
// Add fragments shared by other k partitions
if (kPartitionsK > 1)
{
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
typename SharedLoadIterator::Fragment aligned_addend_fragment;
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_addend_fragment);
aligned_accum_fragment = add_fragments(aligned_accum_fragment, aligned_addend_fragment);
}
}
//
// Iterate over output fragment
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const*>(&aligned_accum_fragment);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
callbacks.begin_row(row_idx);
}
callbacks.visit(
reduce_fragment_idx,
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
callbacks.end_row(row_idx);
}
}
callbacks.end_step(reduce_fragment_idx);
callbacks.end_epilogue();
}
/// Streams the result to global memory
template <class ProblemShape>
CUTLASS_DEVICE
void operator()(
AccumulatorTile const &accumulators,
cutlass::gemm::GemmCoord threadblock_tile_offset,
ProblemShape problem_shape,
int thread_idx
) { ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
auto callbacks = fusion_callbacks.get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
);
callbacks.begin_epilogue();
//
// Iterator over warp-level accumulator fragment
//
AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
//
// Iterate over accumulator tile
//
if constexpr(Pipelined){
__syncthreads();
//
// Pipeline Prologue
//
size_t warp_iterator_offset = kSmemStageOffset;
size_t smem_iterator_offset = kSmemStageOffset;
callbacks.begin_step(0);
acc2smem_source_needed<cutlass::make_index_sequence<kIterations>>::push(
0, accum_fragment_iterator, this->warp_tile_iterator_);
this->warp_tile_iterator_.add_pointer_offset(warp_iterator_offset);
warp_iterator_offset = -warp_iterator_offset;
//
// Pipeline Loop
//
#pragma unroll(IterationsUnroll ? kIterations : 1)
for (int iter_idx = 1; iter_idx < kIterations + 1; ++iter_idx) {
__syncthreads();
// Skip the load for epilogue
if (iter_idx < kIterations) {
callbacks.begin_step(iter_idx);
acc2smem_source_needed<cutlass::make_index_sequence<kIterations>>::push(
iter_idx, accum_fragment_iterator, this->warp_tile_iterator_);
this->warp_tile_iterator_.add_pointer_offset(warp_iterator_offset);
warp_iterator_offset = -warp_iterator_offset;
}
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment[0]);
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
if (kPartitionsK > 1) {
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_accum_fragment[i]);
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
}
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
shared_load_iterator_.add_pointer_offset(smem_iterator_offset);
smem_iterator_offset = -smem_iterator_offset;
//
// Iterate over output fragments
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
callbacks.begin_row(row_idx);
}
callbacks.visit(
iter_idx-1,
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
callbacks.end_row(row_idx);
}
}
//
// Conclude the step
//
callbacks.end_step(iter_idx-1);
}
} else {
#pragma unroll(IterationsUnroll ? kIterations : 1)
for (int iter_idx = 0; iter_idx < kIterations; ++iter_idx) {
//
// Load the source
//
callbacks.begin_step(iter_idx);
//
// Convert and store fragment
//
__syncthreads();
acc2smem_source_needed<cutlass::make_index_sequence<kIterations>>::push(
iter_idx, accum_fragment_iterator, this->warp_tile_iterator_);
__syncthreads();
//
// Load fragments from shared memory
//
typename SharedLoadIterator::Fragment aligned_accum_fragment[kPartitionsK];
shared_load_iterator_.load(aligned_accum_fragment[0]);
// If the number of k-slices is > 1 - perform a reduction amongst the k-slices
if (kPartitionsK > 1) {
plus <typename SharedLoadIterator::Fragment> add_fragments;
CUTLASS_PRAGMA_UNROLL
for ( int i = 1; i < kPartitionsK; ++i) {
shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
shared_load_iterator_.load(aligned_accum_fragment[i]);
aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0], aligned_accum_fragment[i]);
}
shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) * kSmemPointerOffset);
}
//
// Iterate over output fragments
//
AccumulatorAccessType const *accum_frag_ptr =
reinterpret_cast<AccumulatorAccessType const *>(&aligned_accum_fragment[0]);
CUTLASS_PRAGMA_UNROLL
for (int idx = 0; idx < kAccumulatorFragmentCount; ++idx) {
int row_idx = idx / SharedLoadIterator::ThreadMap::Iterations::kColumn;
int col_idx = idx % SharedLoadIterator::ThreadMap::Iterations::kColumn;
// Start a new row of the output fragment
if (!col_idx) {
callbacks.begin_row(row_idx);
}
callbacks.visit(
iter_idx,
row_idx,
col_idx,
idx,
accum_frag_ptr[idx]
);
// End the row of the output fragment
if (col_idx + 1 == SharedLoadIterator::ThreadMap::Iterations::kColumn) {
callbacks.end_row(row_idx);
}
}
//
// Conclude the step
//
callbacks.end_step(iter_idx);
}
}
callbacks.end_epilogue();
}
private:
template<class Seq>
struct acc2smem_source_needed;
template <size_t... Seq>
struct acc2smem_source_needed<cutlass::index_sequence<Seq...>> {
template<int Advance>
CUTLASS_DEVICE
static void helper(AccumulatorFragmentIterator accum_fragment_iterator,
WarpTileIterator &warp_tile_iterator) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < Advance; i++) {
++accum_fragment_iterator;
}
typename AccumulatorFragmentIterator::Fragment accum_fragment;
accum_fragment_iterator.load(accum_fragment);
warp_tile_iterator.store(accum_fragment);
}
CUTLASS_DEVICE
static void push(size_t pos,
AccumulatorFragmentIterator const &iterator_begin,
WarpTileIterator &warp_tile_iterator) {
int dummy[] = {(pos == Seq) && (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
}
};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace threadblock
} // namespace epilogue
} // namespace cutlass
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_2x.hpp
+433
View File
@@ -0,0 +1,433 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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 Visitor tree operation base implementation to enable composable fusions
for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using cute::tuple;
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace detail {
template <class... Ops>
struct VisitorImpl2x: fusion::detail::Sm90VisitorImplBase<Ops...> {
using fusion::detail::Sm90VisitorImplBase<Ops...>::Sm90VisitorImplBase;
using fusion::detail::Sm90VisitorImplBase<Ops...>::ops;
template <class CallbacksTuple>
struct Callbacks {
// Callbacks can store non-persistent variables (e.g. tensors) or copies of persistent variables
CallbacksTuple callbacks_tuple;
/// Called at the start of the epilogue just before iterating over accumulator slices
CUTLASS_DEVICE void
begin_epilogue() {
for_each(callbacks_tuple,
[] (auto& callbacks) {
callbacks.begin_epilogue();
}
);
}
/// Called at the start of one step before starting accumulator exchange
CUTLASS_DEVICE void
begin_step(int step_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.begin_step(step_idx);
}
);
}
/// Called at the start of a row
CUTLASS_DEVICE void
begin_row(int row_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.begin_row(row_idx);
}
);
}
/// Called after accumulators have been exchanged for each accumulator vector
template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInputs, FragmentSize> const&... frg_inputs) // depends on the N-naryness of the op
= delete; // Must be implemented for each operation
/// Called at the start of a row
CUTLASS_DEVICE void
end_row(int row_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.end_row(row_idx);
}
);
}
/// Called after all accumulator elements have been visited
CUTLASS_DEVICE void
end_step(int step_idx) {
for_each(callbacks_tuple,
[&] (auto& callbacks) {
callbacks.end_step(step_idx);
}
);
}
/// Called after all steps have been completed
CUTLASS_DEVICE void
end_epilogue() {
for_each(callbacks_tuple,
[] (auto& callbacks) {
callbacks.end_epilogue();
}
);
}
};
// Callbacks factory
// All operations must redefine this
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return transform_apply(ops,
[&] (auto& op) {
return op.get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape);
},
[] (auto&&... callbacks) {
auto callbacks_tuple = cute::make_tuple(callbacks...);
return Callbacks<decltype(callbacks_tuple)>{callbacks_tuple};
}
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Convenience aliases
using EmptyCallbacks = VisitorImpl2x<>::Callbacks<cute::tuple<>>;
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace detail
using namespace detail;
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Tree visitor
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template <class NodeOp, class... ChildOps>
struct TreeVisitor2x : VisitorImpl2x<ChildOps..., NodeOp> {
using VisitorImpl2x<ChildOps..., NodeOp>::VisitorImpl2x;
template<class CallbacksImpl>
struct Callbacks : CallbacksImpl {
CUTLASS_DEVICE
Callbacks(CallbacksImpl&& impl)
: CallbacksImpl(cute::forward<CallbacksImpl>(impl)) {}
using CallbacksImpl::callbacks_tuple;
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
constexpr int Rm1 = sizeof...(ChildOps);
return cute::detail::tapply(callbacks_tuple,
[&] (auto& child_callbacks) {
return child_callbacks.visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc);
},
[&] (auto&&... frg_inputs) {
return get<Rm1>(callbacks_tuple).visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc, frg_inputs...);
},
make_seq<Rm1>{}
);
}
};
// Callbacks factory
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks<
decltype(VisitorImpl2x<ChildOps..., NodeOp>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
))>(
VisitorImpl2x<ChildOps..., NodeOp>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
)
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class ElementCompute,
class EdgeTuple,
class... Ops
>
struct TopologicalVisitor2x : VisitorImpl2x<Ops...> {
static_assert(is_static_v<EdgeTuple>);
static_assert(rank(EdgeTuple{}) == sizeof...(Ops));
static_assert(sizeof...(Ops) > 1);
using VisitorImpl2x<Ops...>::VisitorImpl2x;
template<class CallbacksImpl>
struct Callbacks : CallbacksImpl {
CUTLASS_DEVICE
Callbacks(CallbacksImpl&& impl)
: CallbacksImpl(cute::forward<CallbacksImpl>(impl)) {}
using CallbacksImpl::callbacks_tuple;
template <typename ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
constexpr int Rm1 = sizeof...(Ops) - 1;
auto frg_compute_tuple = cute::repeat<Rm1>(Array<ElementCompute, FragmentSize>{});
return cute::detail::tapply(EdgeTuple{}, callbacks_tuple, frg_compute_tuple,
// Visit the first R-1 ops in topological order
[&] (auto&& edge_seq, auto& callbacks, auto& frg_compute) {
frg_compute = cute::detail::apply(frg_compute_tuple,
// Compute the current op with children inputs
[&] (auto const&... frg_inputs) {
auto frg_output = callbacks.visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc, frg_inputs...);
using ElementOutput = typename decltype(frg_output)::Element;
using ConvertOutput = NumericArrayConverter<ElementCompute, ElementOutput, FragmentSize>;
ConvertOutput convert_output{};
return convert_output(frg_output);
},
// Get inputs in the sequence given by the children indices of the current op
edge_seq
);
return frg_compute;
},
// Visit the last op
[&] (auto const&...) {
return cute::detail::apply(frg_compute_tuple,
// Compute the last op with children inputs
[&] (auto const&... frg_inputs) {
return get<Rm1>(callbacks_tuple).visit(iter_idx, row_idx, column_idx, frg_idx, frg_acc, frg_inputs...);
},
// Get inputs in the sequence given by the children indices of the last op
get<Rm1>(EdgeTuple{})
);
},
// Transform to visit R-1 ops, apply to visit last op
make_seq<Rm1>{}
);
}
};
// Callbacks factory
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks<decltype(
VisitorImpl2x<Ops...>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
))>(
VisitorImpl2x<Ops...>::
get_callbacks(
threadblock_tile_offset,
thread_idx,
problem_shape
)
);
}
};
template <class NodeOp, class... ChildOps>
using Sm80EVT = TreeVisitor2x<NodeOp, ChildOps...>;
template<
class ElementCompute,
class EdgeTuple,
class... Ops
>
using Sm80TopologicalVisitor = TopologicalVisitor2x<ElementCompute, EdgeTuple, Ops...>;
using X = Underscore;
/////////////////////////////////////////////////////////////////////////////////////////////////
// OutputTileThreadLayout translate the CUTLASS 2.X OutputTileOptimalThreadMap into cute layout
// used by CUTLASS 3.X Epilogue
template <
typename ThreadblockShape_,
typename WarpShape_,
typename Element_,
int ElementsPerAccess,
int Stages_=1
>
struct OutputTileThreadLayout: DefaultThreadMapTensorOp<
ThreadblockShape_,
WarpShape_,
ThreadblockShape_::kK/WarpShape_::kK,
Element_,
ElementsPerAccess>::Type {
using Base = typename DefaultThreadMapTensorOp<
ThreadblockShape_,
WarpShape_,
ThreadblockShape_::kK/WarpShape_::kK,
Element_,
ElementsPerAccess>::Type;
using Base::Base;
// Software pipeline stages in epilogue
static_assert(Stages_ <= 2, "Sm80 EVT only support upto 2 Stages.");
static const int Stages = Stages_;
using ThreadShape = cute::Shape<
cute::Int<Base::Detail::kAccessWidth>, // lane col idx
cute::Int<Base::Detail::kAccessRows>, // lane row idx
cute::Int<Base::Detail::kWarpsRemainingForRows>, // warp row idx
cute::Int<Base::Shape::kGroup>, // group idx
cute::Int<Base::Shape::kCluster> // cluster idx
>;
using Shape = typename Base::Shape;
using Count = typename Base::Count;
using ThreadMapShape = cute::Shape<
// Column
Int<Base::kElementsPerAccess>, // vector
Int<Base::Detail::kAccessWidth>, // lane_col_coord
Int<Base::Iterations::kColumn>, // iteration::column
// Row
Int<Base::Detail::kAccessRows>, // lane_row_coord
Int<Base::Iterations::kRow>, // iterations in row
Int<Base::Detail::kWarpsRemainingForRows>, // warp_row_coord
Int<Count::kRow>, // iteration::row
Int<Count::kGroup>, // iteration::group
Int<Shape::kGroup>, // group_coord
Int<Count::kCluster>, // iteration::cluster
Int<Shape::kCluster> // cluster_coord
>;
// The shape of CTA Tile
using CtaShapeMNL = cute::Shape<
Int<
Shape::kRow * Count::kRow *
Shape::kGroup * Count::kGroup *
Shape::kCluster * Count::kCluster
>,
Int<Shape::kColumn * Count::kColumn>,
_1
>;
static const int kElementsPerAccess = ElementsPerAccess;
//
// Methods
//
CUTLASS_DEVICE
static auto tid2coord(int thread_idx) {
return make_layout(ThreadShape{})[thread_idx];
}
template <class TensorInput>
CUTLASS_DEVICE
static auto partition(TensorInput &&xT, int thread_idx, gemm::GemmCoord threadblock_tile_offset) {
// (BLK_M,BLK_N)
Tensor bCxT = local_tile(
xT, CtaShapeMNL{}, make_coord(_,_,_), Step<_1,_1, X>{}
)(_,_,threadblock_tile_offset.m(),threadblock_tile_offset.n(),threadblock_tile_offset.k());
auto [lane_col_coord, lane_row_coord, warp_row_coord, group_coord, cluster_coord] = tid2coord(thread_idx);
// transform to column-major
Tensor bCxT_nm = make_tensor(
std::forward<decltype(bCxT)>(bCxT).data(), make_layout(get<1>(bCxT.layout()), get<0>(bCxT.layout()))
).compose(make_layout(ThreadMapShape{}));
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
return bCxT_nm(_,lane_col_coord,_,lane_row_coord,_,warp_row_coord,_,_,group_coord,_,cluster_coord);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_compute.hpp
+109
View File
@@ -0,0 +1,109 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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 Visitor tree compute operations for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using namespace detail;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// N-nary Elementwise Compute Operation
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
template <class> class ComputeFn,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class = void
>
struct VisitorCompute : VisitorImpl2x<> {
using VisitorImpl2x<>::VisitorImpl2x;
struct Callbacks : EmptyCallbacks {
template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
CUTLASS_DEVICE Array<ElementOutput, FragmentSize>
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInputs, FragmentSize> const&... frg_inputs) {
return transform_apply(cute::make_tuple(frg_inputs...),
[&] (auto&& frg_input) {
using ElementInput = typename cute::remove_cvref_t<decltype(frg_input)>::Element;
using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
return convert_input(frg_input);
},
[&] (auto&&... cvt_frg_inputs) {
using ComputeOutput = ComputeFn<Array<ElementCompute, FragmentSize>>;
using ConvertOutput = NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize, RoundStyle>;
ComputeOutput compute_output{};
ConvertOutput convert_output{};
return convert_output(compute_output(cvt_frg_inputs...));
}
);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks();
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp
+559
View File
@@ -0,0 +1,559 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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 Visitor tree load operations for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
#include "cute/tensor.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using namespace detail;
using X = Underscore;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Elementwise Fetch Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// returns accumulator
struct VisitorAccFetch : VisitorImpl2x<> {
using VisitorImpl2x<>::VisitorImpl2x;
struct Callbacks : EmptyCallbacks {
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE Array<ElementAccumulator, FragmentSize>
visit(int iter_idx, int row_idx, int column_idx, int frg_idx, Array<ElementAccumulator, FragmentSize> const& frg_acc) {
return frg_acc;
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
return Callbacks{};
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Broadcast Load Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// Scalar broadcast
template<
class Element,
class StrideMNL = Stride<_0,_0,_0>,
int BroadcastCount = 1,
template <class> class ReductionFn = multiplies
>
struct VisitorScalarBroadcast {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0,_0>>) || // scalar broadcast, e.g. alpha
(cute::is_same_v<StrideMNL, Stride<_0,_0,_1>>) ||
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>)); // batched scalar broadcast, e.g. per-batch alpha
struct SharedStorage { };
struct Arguments {
Element scalars[BroadcastCount] = {};
Element const* scalar_ptrs[BroadcastCount] = {};
StrideMNL dScalar = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
CUTLASS_HOST_DEVICE
VisitorScalarBroadcast() { }
CUTLASS_HOST_DEVICE
VisitorScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) {
// Get the scalar for non-batched broadcast
if constexpr (cute::is_same_v<StrideMNL, Stride<_0,_0,_0>>) {
update_scalar();
}
}
Element scalar;
Params const* params_ptr;
struct Callbacks: EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(Element scalar)
: scalar(scalar) {}
Element scalar;
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Array<Element, FragmentSize> frg_scalar;
frg_scalar.fill(scalar);
return frg_scalar;
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
// Get the scalar for batched broadcast
if constexpr (
cute::is_same_v<StrideMNL, Stride<_0,_0,_1>> ||
cute::is_same_v<StrideMNL, Stride<_0,_0,int>>) {
update_scalar(threadblock_tile_offset.k());
}
return Callbacks(scalar);
}
private:
CUTLASS_DEVICE void
update_scalar(int l_coord = 0) {
int l_offset = l_coord * size<2>(params_ptr->dScalar);
if (params_ptr->scalar_ptrs[0] != nullptr) {
scalar = params_ptr->scalar_ptrs[0][l_offset];
} else {
// batch stride is ignored for nullptr fallback
scalar = params_ptr->scalars[0];
}
// Do reduction over multiple broadcasts if necessary
ReductionFn<Element> reduction_fn;
CUTLASS_PRAGMA_UNROLL
for (int i = 1; i < BroadcastCount; ++i) {
if (params_ptr->scalar_ptrs[i] != nullptr) {
scalar = reduction_fn(scalar, params_ptr->scalar_ptrs[i][l_offset]);
} else {
// batch stride is ignored for nullptr fallback
scalar = reduction_fn(scalar, params_ptr->scalars[i]);
}
}
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Elementwise Load Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class ThreadMap,
class Element,
class StrideMNL
>
struct VisitorAuxLoad{
struct Arguments {
Element* ptr_aux = nullptr;
Element null_default = Element(0);
StrideMNL dAux = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
// Software pipeline stages
static const int Stages = ThreadMap::Stages;
struct SharedStorage {};
// Global load type
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
static int constexpr VecLength = sizeof(VecType) / sizeof(Element);
CUTLASS_HOST_DEVICE
VisitorAuxLoad() { }
CUTLASS_HOST_DEVICE
VisitorAuxLoad(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gAux,
RTensor&& tC_rAux,
CTensor&& tC_cAux,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gAux(cute::forward<GTensor>(tC_gAux)),
tC_rAux(cute::forward<RTensor>(tC_rAux)),
tC_cAux(cute::forward<CTensor>(tC_cAux)),
problem_shape(problem_shape),
params_ptr(params_ptr) { }
GTensor tC_gAux;
RTensor tC_rAux;
CTensor tC_cAux;
Params const* params_ptr;
ProblemShape problem_shape;
CUTLASS_DEVICE void
begin_step(int step_idx) {
clear(tC_rAux(_,_,_,step_idx%Stages));
auto src_v = filter(tC_gAux(_,_,_,step_idx));
auto coord_v = filter(tC_cAux(_,_,_,step_idx));
auto dst_v = filter(tC_rAux(_,_,_,step_idx%Stages));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(src_v); ++i) {
bool guard = elem_less(coord_v(i), problem_shape);
cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const*)&src_v(i), guard);
}
}
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Tensor tC_rAux_frg = recast<Array<Element, FragmentSize>>(coalesce(tC_rAux(_,_,_,iter_idx%Stages)));
return tC_rAux_frg(frg_idx);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mAux = make_tensor(
make_gmem_ptr(params_ptr->ptr_aux),
problem_shape,
params_ptr->dAux); // (M,N,L)
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
Tensor tC_gAux = recast<VecType>(
group_modes<3,6>(ThreadMap::partition(mAux, thread_idx, threadblock_tile_offset)));
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, Stages
Tensor tC_rAux = make_tensor<VecType>(
make_layout(flatten(make_shape(take<0,3>(tC_gAux.shape()), Int<Stages>{}))));
// Generate the pred tensor
Tensor cAux = make_identity_tensor(mAux.shape());
Tensor tC_cAux = local_partition(
group_modes<3,6>(ThreadMap::partition(cAux, thread_idx, threadblock_tile_offset)),
Shape<Int<VecLength>>{},
(_0{})
);
return Callbacks<
decltype(tC_gAux), decltype(tC_rAux),
decltype(tC_cAux), ProblemShape>(
cute::move(tC_gAux),
cute::move(tC_rAux),
cute::move(tC_cAux),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Row vector broadcast
template<
class ThreadMap,
class Element,
class StrideMNL
>
struct VisitorRowBroadcast {
struct Arguments {
Element const* ptr_row = nullptr;
Element null_default = Element(0);
StrideMNL dRow = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage {};
// Global load type
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
static int constexpr VecLength = sizeof(VecType) / sizeof(Element);
CUTLASS_HOST_DEVICE
VisitorRowBroadcast() { }
CUTLASS_HOST_DEVICE
VisitorRowBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gRow,
RTensor&& tC_rRow,
CTensor&& tC_cRow,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gRow(cute::forward<GTensor>(tC_gRow)),
tC_rRow(cute::forward<RTensor>(tC_rRow)),
tC_cRow(cute::forward<CTensor>(tC_cRow)),
n(get<1>(problem_shape)),
params_ptr(params_ptr) { }
GTensor tC_gRow;
RTensor tC_rRow;
CTensor tC_cRow;
Params const* params_ptr;
int n;
CUTLASS_DEVICE void
begin_epilogue() {
clear(tC_rRow);
auto src_v = filter(tC_gRow);
auto coord_v = filter(tC_cRow);
auto dst_v = filter(tC_rRow);
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(src_v); ++i) {
bool guard = get<1>(coord_v(i)) < n;
cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const*)&src_v(i), guard);
}
}
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Tensor rRow_frg = recast<Array<Element, FragmentSize>>(coalesce(tC_rRow));
return rRow_frg(column_idx);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mRow = make_tensor(
make_gmem_ptr(params_ptr->ptr_row),
problem_shape,
params_ptr->dRow);
// VECTOR, FRAGMENT_COLUMN
Tensor tC_gRow = recast<VecType>(
ThreadMap::partition(mRow, thread_idx, threadblock_tile_offset)
)(_,_,_0{},_0{},_0{},_0{});
Tensor tC_rRow = make_tensor_like(tC_gRow);
// Generate the pred tensor
Tensor cRow = make_identity_tensor(mRow.shape());
Tensor tC_cRow = local_partition(
ThreadMap::partition(cRow, thread_idx, threadblock_tile_offset)(_,_,_0{},_0{},_0{},_0{}),
Shape<Int<VecLength>>{},
(_0{})
);
return Callbacks<
decltype(tC_gRow), decltype(tC_rRow),
decltype(tC_cRow), ProblemShape>(
cute::move(tC_gRow),
cute::move(tC_rRow),
cute::move(tC_cRow),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Column vector broadcast
template<
class ThreadMap,
class Element,
class StrideMNL = Stride<_1,_0,_0>
>
struct VisitorColBroadcast {
struct Arguments {
Element const* ptr_col = nullptr;
Element null_default = Element(0);
StrideMNL dCol = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage { };
CUTLASS_HOST_DEVICE
VisitorColBroadcast() { }
CUTLASS_HOST_DEVICE
VisitorColBroadcast(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gCol,
RTensor&& tC_rCol,
CTensor&& tC_cCol,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gCol(cute::forward<GTensor>(tC_gCol)),
tC_rCol(cute::forward<RTensor>(tC_rCol)),
tC_cCol(cute::forward<CTensor>(tC_cCol)),
m(get<0>(problem_shape)),
params_ptr(params_ptr) { }
GTensor tC_gCol;
RTensor tC_rCol;
CTensor tC_cCol;
Params const* params_ptr;
int m;
CUTLASS_DEVICE void
begin_epilogue() {
clear(tC_rCol);
Tensor pred = make_tensor<bool>(shape(tC_gCol));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(pred); ++i) {
pred(i) = get<0>(tC_cCol(i)) < m;
}
copy_if(pred, tC_gCol, tC_rCol);
}
template <class ElementAccumulator, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc) {
Array<Element, FragmentSize> frg_col;
frg_col.fill(tC_rCol(row_idx,iter_idx));
return frg_col;
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mCol = make_tensor(
make_gmem_ptr(params_ptr->ptr_col),
problem_shape,
params_ptr->dCol);
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
Tensor tC_gCol = group_modes<1,4>(
ThreadMap::partition(mCol, thread_idx, threadblock_tile_offset)(_0{},_0{},_,_,_,_));
Tensor tC_rCol = make_tensor_like(tC_gCol);
// Generate the pred tensor
Tensor cCol = make_identity_tensor(mCol.shape());
Tensor tC_cCol = group_modes<1,4>(
ThreadMap::partition(cCol, thread_idx, threadblock_tile_offset)(_0{},_0{},_,_,_,_));
return Callbacks<
decltype(tC_gCol), decltype(tC_rCol),
decltype(tC_cCol), ProblemShape>(
cute::move(tC_gCol),
cute::move(tC_rCol),
cute::move(tC_cCol),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp
+781
View File
@@ -0,0 +1,781 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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 Visitor tree store operations for the CUTLASS 2x epilogue
*/
#pragma once
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::epilogue::threadblock {
using namespace cute;
using namespace detail;
using X = Underscore;
/////////////////////////////////////////////////////////////////////////////////////////////////
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Elementwise Store Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
class ThreadMap,
class Element,
FloatRoundStyle RoundStyle,
class StrideMNL
>
struct VisitorAuxStore{
struct Arguments {
Element* ptr_aux = nullptr;
StrideMNL dAux = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage {};
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<Element>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
static int constexpr VecLength = sizeof(VecType) / sizeof(Element);
CUTLASS_HOST_DEVICE
VisitorAuxStore() { }
CUTLASS_HOST_DEVICE
VisitorAuxStore(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class RTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gAux,
RTensor&& tC_rAux,
CTensor&& tC_cAux,
ProblemShape problem_shape,
Params const* params_ptr
):
tC_gAux(cute::forward<GTensor>(tC_gAux)),
tC_rAux(cute::forward<RTensor>(tC_rAux)),
tC_cAux(cute::forward<CTensor>(tC_cAux)),
problem_shape(problem_shape),
params_ptr(params_ptr) { }
GTensor tC_gAux;
RTensor tC_rAux;
CTensor tC_cAux;
Params const* params_ptr;
ProblemShape problem_shape;
CUTLASS_DEVICE void
begin_step(int step_idx) {
clear(tC_rAux);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
using ConvertInput = NumericArrayConverter<Element, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
Tensor tC_rAux_frg = recast<Array<Element, FragmentSize>>(coalesce(tC_rAux));
tC_rAux_frg(frg_idx) = convert_input(frg_input);
return frg_input;
}
CUTLASS_DEVICE void
end_step(int step_idx) {
auto src_v = filter(tC_rAux);
auto coord_v = filter(tC_cAux(_,_,_,step_idx));
auto dst_v = filter(tC_gAux(_,_,_,step_idx));
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(src_v); ++i) {
bool guard = elem_less(coord_v(i), problem_shape);
cutlass::arch::global_store<VecType, sizeof(VecType)>(src_v(i), (void*)&dst_v(i), guard);
}
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mAux = make_tensor(
make_gmem_ptr(params_ptr->ptr_aux),
problem_shape,
params_ptr->dAux); // (M,N,L)
// VECTOR, FRAGMENT_COLUMN, FRAGMENT_ROW, ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER
Tensor tC_gAux = recast<VecType>(group_modes<3,6>(ThreadMap::partition(mAux, thread_idx, threadblock_tile_offset)));
Tensor tC_rAux = make_tensor_like(take<0,3>(tC_gAux));
// Generate the pred tensor
Tensor cAux = make_identity_tensor(mAux.shape());
Tensor tC_cAux = local_partition(
group_modes<3,6>(ThreadMap::partition(cAux, thread_idx, threadblock_tile_offset)),
Shape<Int<VecLength>>{},
(_0{})
);
return Callbacks<
decltype(tC_gAux), decltype(tC_rAux),
decltype(tC_cAux), ProblemShape>(
cute::move(tC_gAux),
cute::move(tC_rAux),
cute::move(tC_cAux),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Reduction Store Operations
//
/////////////////////////////////////////////////////////////////////////////////////////////////
// Helper functions
template <
template <class> class ReduceFn,
int kThreads, class T>
CUTLASS_DEVICE
void intra_warp_row_reduce(T& value) {
using ReduceInput = ReduceFn<T>;
ReduceInput reduce_input{};
constexpr int kHalfThreads = kThreads >> 1;
CUTLASS_PRAGMA_UNROLL
for (int i = kHalfThreads; i > 0; i >>= 1) {
value = reduce_input(value, __shfl_xor_sync(0xFFFFFFFF, value, i));
}
}
template <
template <class> class ReduceFn,
FloatRoundStyle RoundStyle,
class ElementCompute,
class ElementFragment, int FragmentSize>
CUTLASS_DEVICE
void fragment_reduce(ElementCompute& value, Array<ElementFragment, FragmentSize> const& frg) {
using ReduceInput = ReduceFn<ElementCompute>;
ReduceInput reduce_input{};
using ConvertInput = NumericConverter<ElementCompute, ElementFragment, RoundStyle>;
ConvertInput convert_input{};
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
value = reduce_input(value, convert_input(frg[i]));
}
}
template<
template <class> class AtomicReduceFn,
FloatRoundStyle RoundStyle,
class ElementCompute,
class ElementOutput>
CUTLASS_DEVICE
void atomic_reduce(ElementOutput* ptr, ElementCompute const& value) {
using ReduceOutput = AtomicReduceFn<ElementOutput>;
using ConvertOutput = NumericConverter<ElementOutput, ElementCompute, RoundStyle>;
ReduceOutput reduce_output{};
ConvertOutput convert_output{};
reduce_output(ptr, convert_output(value));
}
// Col vector reduction
template <
template <class> class RegReduceFn,
template <class> class AtomicReduceFn,
class ThreadMap,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class StrideMNL = Stride<_1,_0,_0>
>
struct VisitorColReduction {
struct Arguments {
ElementOutput* ptr_col = nullptr;
ElementCompute reduction_identity = 0;
StrideMNL dCol = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage { };
CUTLASS_HOST_DEVICE
VisitorColReduction() { }
CUTLASS_HOST_DEVICE
VisitorColReduction(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class GTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
GTensor&& tC_gCol,
CTensor&& tC_cCol,
ProblemShape problem_shape,
Params const* params_ptr,
int thread_idx
):
tC_gCol(cute::forward<GTensor>(tC_gCol)),
tC_cCol(cute::forward<CTensor>(tC_cCol)),
m(get<0>(problem_shape)),
n(get<1>(problem_shape)),
params_ptr(params_ptr) {
// The partial reduction results of each warp are further
// reduced to the first thread in each row.
// Only the first thread in each row is the writing thread
is_writing_thread = thread_idx % ThreadMap::Detail::kAccessWidth == 0;
}
GTensor tC_gCol;
CTensor tC_cCol;
Params const* params_ptr;
int m;
int n;
int curr_iter_idx;
bool is_writing_thread;
ElementCompute reduction_accum;
CUTLASS_DEVICE void
begin_row(int row_idx) {
reduction_accum = ElementCompute(params_ptr->reduction_identity);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
curr_iter_idx = iter_idx;
int coord_n = get<1>(tC_cCol(column_idx, row_idx, iter_idx));
if (coord_n < n) {
fragment_reduce<RegReduceFn, RoundStyle>(reduction_accum, frg_input);
}
// Intra-warp reduction
if (column_idx + 1 == ThreadMap::Iterations::kColumn) {
intra_warp_row_reduce<RegReduceFn, ThreadMap::Detail::kAccessWidth>(reduction_accum);
}
return frg_input;
}
CUTLASS_DEVICE auto
end_row(int row_idx) {
bool guard = get<0>(tC_cCol(_0{}, row_idx,curr_iter_idx)) < m;
if (guard && is_writing_thread) {
atomic_reduce<AtomicReduceFn, RoundStyle>(&tC_gCol(row_idx,curr_iter_idx), reduction_accum);
}
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mCol = make_tensor(
make_gmem_ptr(params_ptr->ptr_col),
problem_shape,
params_ptr->dCol);
// FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tC_gCol = group_modes<1,4>(
ThreadMap::partition(mCol, thread_idx, threadblock_tile_offset)(_0{},_0{},_,_,_,_));
// Generate the pred tensor
Tensor cCol = make_identity_tensor(mCol.shape());
// FRAGMENT_COL, FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tC_cCol = group_modes<2,5>(
ThreadMap::partition(cCol, thread_idx, threadblock_tile_offset)(_0{},_,_,_,_,_));
return Callbacks<
decltype(tC_gCol), decltype(tC_cCol),
ProblemShape>(
cute::move(tC_gCol),
cute::move(tC_cCol),
problem_shape,
params_ptr,
thread_idx
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Row vector reduction
template <
template <class> class RegReduceFn,
template <class> class AtomicReduceFn,
class ThreadMap,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class StrideMNL = Stride<_0,_1,_0>
>
struct VisitorRowReduction {
struct Arguments {
ElementOutput* ptr_row = nullptr;
ElementCompute reduction_identity = 0;
StrideMNL dRow = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
using SharedStorageShape = decltype(select<0,1,2,3,5,8,10>(typename ThreadMap::ThreadMapShape{}));
struct SharedStorage {
AlignedArray<ElementCompute, size(SharedStorageShape{}), 16> reduction;
};
static int constexpr vec_bits = ThreadMap::kElementsPerAccess * sizeof_bits<ElementOutput>::value;
using VecType = uint_bit_t<cute::min(128, vec_bits)>;
CUTLASS_HOST_DEVICE
VisitorRowReduction() { }
CUTLASS_HOST_DEVICE
VisitorRowReduction(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params),
smem_reduce(const_cast<ElementCompute*>(shared_storage.reduction.data())) { }
Params const* params_ptr;
ElementCompute* smem_reduce;
template <
class RTensorR2S, class STensorR2S, class CTensorR2S,
class STensorS2R, class RTensorS2R, class CTensorS2R,
class GTensor, class CTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
// R->S
RTensorR2S&& tRS_rSrc,
STensorR2S&& tRS_sRows,
CTensorR2S&& tRS_cSrc,
// S->R
STensorS2R&& tSR_sRows,
RTensorS2R&& tSR_rRows,
CTensorS2R&& tSR_cRows,
// R->G
GTensor&& tC_gRow,
CTensor&& tC_cRow,
ProblemShape problem_shape,
Params const* params_ptr
):
// R->S
tRS_rSrc(cute::forward<RTensorR2S>(tRS_rSrc)),
tRS_sRows(cute::forward<STensorR2S>(tRS_sRows)),
tRS_cSrc(cute::forward<CTensorR2S>(tRS_cSrc)),
// S->R
tSR_sRows(cute::forward<STensorS2R>(tSR_sRows)),
tSR_rRows(cute::forward<RTensorS2R>(tSR_rRows)),
tSR_cRows(cute::forward<CTensorS2R>(tSR_cRows)),
// R->G
tC_gRow(cute::forward<GTensor>(tC_gRow)),
tC_cRow(cute::forward<CTensor>(tC_cRow)),
m(get<0>(problem_shape)),
n(get<1>(problem_shape)),
params_ptr(params_ptr) { }
// R->S
RTensorR2S tRS_rSrc;
STensorR2S tRS_sRows;
CTensorR2S tRS_cSrc;
// S->R
STensorS2R tSR_sRows;
RTensorS2R tSR_rRows;
CTensorS2R tSR_cRows;
// R->G
GTensor tC_gRow;
CTensor tC_cRow;
Params const* params_ptr;
int n;
int m;
CUTLASS_DEVICE void
begin_epilogue() {
fill(tRS_rSrc, params_ptr->reduction_identity);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto // returns an Array
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
ConvertInput convert_input{};
Tensor tRS_rRow_frg = recast<Array<ElementCompute, FragmentSize>>(coalesce(tRS_rSrc));
int coord_m = get<0>(tRS_cSrc(column_idx,row_idx,iter_idx));
if (coord_m < m)
reduction(tRS_rRow_frg[column_idx], convert_input(frg_input));
return frg_input;
}
CUTLASS_DEVICE void
end_epilogue() {
//
// Store the partially reduced value to SMEM
//
// Guard against uses of the existing SMEM tile
__syncthreads();
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tRS_rSrc); ++i) {
copy_vec<VecType>(filter(tRS_rSrc), filter(tRS_sRows));
}
__syncthreads();
//
// Now, threads are assigned several columns of the output. They fetch over all rows from
// the compacted SMEM tile and perform a reduction.
//
fill(tSR_rRows, params_ptr->reduction_identity);
using ReduceInputReg = RegReduceFn<ElementCompute>;
ReduceInputReg reduce_input_reg{};
CUTLASS_PRAGMA_UNROLL
for (int j = 0; j < size(tSR_rRows); ++j) {
if (get<0>(tSR_cRows(j)) < get<1>(typename ThreadMap::CtaShapeMNL{}) && get<1>(tC_cRow(j)) < n) {
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tSR_sRows) / size(tSR_rRows); ++i) {
tSR_rRows(j) = reduce_input_reg(tSR_rRows(j), tSR_sRows(i + j * size(tSR_sRows) / size(tSR_rRows)));
}
atomic_reduce<AtomicReduceFn, RoundStyle>(&tC_gRow(j), tSR_rRows(j));
}
}
}
private:
template <int FragmentSize>
CUTLASS_DEVICE ElementCompute
reduction(Array<ElementCompute, FragmentSize>& reduce_buffer, Array<ElementCompute, FragmentSize> const& result) {
using ReduceInput = RegReduceFn<ElementCompute>;
ReduceInput reduce_input{};
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < FragmentSize; ++i) {
reduce_buffer[i] = reduce_input(reduce_buffer[i], result[i]);
}
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor mRow = make_tensor(
make_gmem_ptr(params_ptr->ptr_row),
problem_shape,
params_ptr->dRow);
//
// Step 1: reduce fragment input (Src) into tRS_rSrc
//
// VECTOR,FRAGMENT_COL
Tensor tRS_rSrc = make_tensor<ElementCompute>(select<0,2>(typename ThreadMap::ThreadMapShape{}));
Tensor cSrc = make_identity_tensor(mRow.shape());
// FRAGMENT_COLUMN, FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tRS_cSrc = group_modes<2,5>(ThreadMap::partition(cSrc, thread_idx, threadblock_tile_offset)(_0{},_,_,_,_,_));
//
// Step 2: copy the partial results in tRS_rSrc to sRows in shared memory
//
// VECTOR,ACCESS_WIDTH,FRAGMENT_COL,ACCESS_ROWS,WARPS_PER_ROW,GROUPS,CLUSTERS
Tensor sRows = make_tensor(
make_smem_ptr(smem_reduce), SharedStorageShape{}
);
auto [lane_col_coord, lane_row_coord, warp_row_coord, group_coord, cluster_coord] = ThreadMap::tid2coord(thread_idx);
Tensor tRS_sRows = sRows(_,lane_col_coord,_,lane_row_coord,warp_row_coord,group_coord,cluster_coord);
//
// Step 3: copy the partial results in sRows to tSR_sRow for reduction
//
// VECTOR*ACCESS_WIDTH*FRAGMENT_COL,ACCESS_ROWS*WARPS_PER_ROW*GROUPS*CLUSTERS
Tensor sRows_nm = coalesce(group_modes<1,5>(group_modes<0,3>(sRows)), Shape<_1,_1>{});
// SMEM_ROW/THREADS,ACCESS_ROWS*WARPS_PER_ROW*GROUPS*CLUSTERS
Tensor tSR_sRows = local_partition(sRows_nm, Shape<Int<ThreadMap::kThreads>,_1>{}, thread_idx);
// SMEM_ROW/THREADS
Tensor tSR_rRows = make_tensor_like(tSR_sRows(_,_0{}));
// Coord
Tensor cRows_nm = make_identity_tensor(sRows_nm.shape());
Tensor tSR_cRows = local_partition(cRows_nm, Shape<Int<ThreadMap::kThreads>,_1>{}, thread_idx)(_,_0{});
//
// Step 4: atomically reduce the results to global memory
//
Tensor tC_gRow = local_partition(
// Cta tile
local_tile(
mRow, typename ThreadMap::CtaShapeMNL{}, make_coord(_,_,_),Step<_1,_1, X>{}
)(_,_,threadblock_tile_offset.m(),threadblock_tile_offset.n(),threadblock_tile_offset.k()),
// Partition to threads
Shape<_1,Int<ThreadMap::kThreads>>{}, thread_idx
)(_0{},_);
Tensor cRow = make_identity_tensor(mRow.shape());
Tensor tC_cRow = local_partition(
// Cta tile
local_tile(
cRow, typename ThreadMap::CtaShapeMNL{}, make_coord(_,_,_), Step<_1,_1, X>{}
)(_,_,threadblock_tile_offset.m(),threadblock_tile_offset.n(),threadblock_tile_offset.k()),
// Partition to threads
Shape<_1,Int<ThreadMap::kThreads>>{}, thread_idx
)(_0{},_);
return Callbacks<
decltype(tRS_rSrc), decltype(tRS_sRows),
decltype(tRS_cSrc), decltype(tSR_sRows),
decltype(tSR_rRows), decltype(tSR_cRows),
decltype(tC_gRow), decltype(tC_cRow),
ProblemShape>(
// R->S
cute::move(tRS_rSrc),
cute::move(tRS_sRows),
cute::move(tRS_cSrc),
// S->R
cute::move(tSR_sRows),
cute::move(tSR_rRows),
cute::move(tSR_cRows),
// R->G
cute::move(tC_gRow),
cute::move(tC_cRow),
problem_shape,
params_ptr
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
// Scalar reduction
template <
template <class> class RegReduceFn,
template <class> class AtomicReduceFn,
class ThreadMap,
class ElementOutput,
class ElementCompute,
FloatRoundStyle RoundStyle,
class StrideMNL = Stride<_0,_0,_0>
>
struct VisitorScalarReduction {
static_assert(
(cute::is_same_v<StrideMNL, Stride<_0,_0, _0>>) || // scalar reduction, e.g. tensor max element
(cute::is_same_v<StrideMNL, Stride<_0,_0, _1>>) || // batched scalar reduction, e.g. per-batch max element
(cute::is_same_v<StrideMNL, Stride<_0,_0,int>>));
struct Arguments {
ElementOutput* ptr_scalar = nullptr;
ElementCompute reduction_identity = 0;
StrideMNL dScalar = {};
};
using Params = Arguments;
template <class ProblemShape>
static constexpr Params
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
return args;
}
struct SharedStorage { };
CUTLASS_HOST_DEVICE
VisitorScalarReduction(){ };
CUTLASS_HOST_DEVICE
VisitorScalarReduction(Params const& params, SharedStorage const& shared_storage)
: params_ptr(&params) { }
Params const* params_ptr;
template <class CTensor, class GTensor, class ProblemShape>
struct Callbacks : EmptyCallbacks {
CUTLASS_DEVICE
Callbacks(
CTensor&& tC_cSrc,
GTensor&& tC_gScalar,
ProblemShape problem_shape,
Params const* params_ptr,
int thread_idx
):
tC_cSrc(cute::forward<CTensor>(tC_cSrc)),
tC_gScalar(cute::forward<GTensor>(tC_gScalar)),
problem_shape(problem_shape),
params_ptr(params_ptr) {
// The partial reduction results of each warp are further
// reduced to this first thread.
// Only the first thread of each warp is the writing thread
is_writing_thread = thread_idx % ThreadMap::kWarpSize == 0;
}
GTensor tC_gScalar;
CTensor tC_cSrc;
Params const* params_ptr;
ProblemShape problem_shape;
bool is_writing_thread;
ElementCompute reduction_accum;
CUTLASS_DEVICE void
begin_epilogue() {
reduction_accum = ElementCompute(params_ptr->reduction_identity);
}
template <class ElementAccumulator, class ElementInput, int FragmentSize>
CUTLASS_DEVICE auto
visit(int iter_idx, int row_idx, int column_idx, int frg_idx,
Array<ElementAccumulator, FragmentSize> const& frg_acc,
Array<ElementInput, FragmentSize> const& frg_input) {
auto coord = tC_cSrc(column_idx, row_idx, iter_idx);
if (elem_less(coord, problem_shape)) {
fragment_reduce<RegReduceFn, RoundStyle>(reduction_accum, frg_input);
}
return frg_input;
}
CUTLASS_DEVICE auto
end_epilogue() {
// Intra-warp reduction
intra_warp_row_reduce<RegReduceFn, ThreadMap::kWarpSize>(reduction_accum);
// Atomically reduce to global memory
atomic_reduce<AtomicReduceFn, RoundStyle>(&tC_gScalar(_0{},_0{}), reduction_accum);
}
};
template <class ProblemShape>
CUTLASS_DEVICE auto
get_callbacks(
gemm::GemmCoord threadblock_tile_offset,
int thread_idx,
ProblemShape problem_shape
) {
Tensor cSrc = make_identity_tensor(problem_shape);
// FRAGMENT_COL, FRAGMENT_ROW, (ITERATION_ROW, ITERATION_GROUP, ITERATION_CLUSTER)
Tensor tC_cSrc = group_modes<2,5>(
ThreadMap::partition(cSrc, thread_idx, threadblock_tile_offset)(_0{},_,_,_,_,_)
);
Tensor mScalar = make_tensor(
make_gmem_ptr(params_ptr->ptr_scalar),
problem_shape,
params_ptr->dScalar
);
Tensor tC_gScalar = mScalar(_,_,threadblock_tile_offset.k());
return Callbacks<
decltype(tC_cSrc), decltype(tC_gScalar),
ProblemShape>(
cute::move(tC_cSrc),
cute::move(tC_gScalar),
problem_shape,
params_ptr,
thread_idx
);
}
};
/////////////////////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::epilogue::threadblock
/////////////////////////////////////////////////////////////////////////////////////////////////
include/cutlass/epilogue/threadblock/fusion/visitors.hpp
+38
View File
@@ -0,0 +1,38 @@
/***************************************************************************************************
* Copyright (c) 2023 - 2023 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 Higher-level header file includes all the CUTLASS 2x visitors
*/
#include "cutlass/epilogue/threadblock/fusion/visitor_2x.hpp"
#include "cutlass/epilogue/threadblock/fusion/visitor_load.hpp"
#include "cutlass/epilogue/threadblock/fusion/visitor_store.hpp"
#include "cutlass/epilogue/threadblock/fusion/visitor_compute.hpp"
include/cutlass/epilogue/threadblock/predicated_tile_iterator_params.h
+10
View File
@@ -32,6 +32,16 @@
\brief
*/
/*
Note: CUTLASS 3x increases the host compiler requirements to C++17. However, certain
existing integrations of CUTLASS require C++11 host compilers.
Until this requirement can be lifted, certain headers with this annotation are required
to be remain consistent with C++11 syntax.
C++11 compatibility is enforced by this unit test: `cutlass_test_unit_core_cpp11`.
*/
#pragma once
#include "cutlass/cutlass.h"