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

Updates for 3.1 (#932)

Browse Source
This commit is contained in:
ANIKET SHIVAM
2023-04-29 09:34:27 -04:00
committed by GitHub
parent 6f8596ce3f
commit 7c04f95415
51 changed files with 1796 additions and 328 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/epilogue/collective/builders/sm90_builder.inl
+35 -13
View File
@@ -290,7 +290,7 @@ struct CollectiveBuilder<
AlignmentD,
Schedule,
cute::enable_if_t<cute::is_same_v<Schedule, TmaWarpSpecialized> ||
cute::is_same_v<Schedule, TmaWarpSpecializedCooperative> >> {
cute::is_same_v<Schedule, TmaWarpSpecializedCooperative> >> {
public:
// Passing void C disables source load
using ElementC = cute::conditional_t<cute::is_void_v<ElementC_>,
@@ -302,16 +302,33 @@ public:
using ThreadOp = thread::LinearCombination<
ElementD, AlignmentD, ElementAccumulator, ElementCompute,
thread::ScaleType::Default, FloatRoundStyle::round_to_nearest, ElementC>;
ScaleType, FloatRoundStyle::round_to_nearest, ElementC>;
private:
using Impl = detail::TmaBuilderImpl<
TileShape_MNK, ClusterShape_MNK, EpilogueTileType, ElementAccumulator, ElementCompute,
ElementC, GmemLayoutTagC, AlignmentC, ElementD, GmemLayoutTagD, AlignmentD,
Schedule, ThreadOp, cutlass::epilogue::Sm90TmaWarpSpecialized<1,2,true>>;
using GmemStrideTypeC = gemm::TagToStrideC_t<GmemLayoutTagC>;
using GmemStrideTypeD = gemm::TagToStrideC_t<GmemLayoutTagD>;
public:
using CollectiveOp = typename Impl::CollectiveOp;
using EpilogueTile_MN = decltype(detail::sm90_compute_tile_shape_or_override<
ElementD, EpilogueTileType, Schedule>());
static constexpr int StagesC = 1;
static constexpr int StagesD = 2;
static constexpr bool DisableReuseSmemC = true;
using CollectiveOp = cutlass::epilogue::collective::CollectiveEpilogue<
cutlass::epilogue::Sm90TmaWarpSpecialized<StagesC,StagesD,DisableReuseSmemC>,
TileShape_MNK,
EpilogueTile_MN,
ElementC_, // need to pass void to expose via GemmUniversal
GmemStrideTypeC,
ElementD,
GmemStrideTypeD,
ThreadOp,
SM90_TMA_LOAD,
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<GmemStrideTypeC, ElementC, TileShape_MNK>()),
decltype(detail::sm90_get_smem_load_op_for_source<GmemStrideTypeC, ElementC>()),
SM90_TMA_STORE,
decltype(detail::sm90_get_epilogue_smem_swizzle_layout_atom<GmemStrideTypeD, ElementD, EpilogueTile_MN>()),
decltype(detail::sm90_get_smem_store_op_for_accumulator<GmemStrideTypeD, ElementD>())
>;
};
// Auto builder
@@ -409,7 +426,7 @@ struct CollectiveBuilder<
AlignmentD,
Schedule,
cute::enable_if_t<cute::is_base_of_v<TmaWarpSpecializedElementwiseBase, Schedule> ||
cute::is_base_of_v<TmaWarpSpecializedCooperativeElementwiseBase, Schedule> >> {
cute::is_base_of_v<TmaWarpSpecializedCooperativeElementwiseBase, Schedule> >> {
public:
using ThreadOp = thread::LinearCombinationGeneric<
@@ -419,10 +436,13 @@ public:
Schedule::Round>;
private:
static constexpr int StagesC = 1;
static constexpr int StagesD = 2;
static constexpr bool DisableReuseSmemC = true;
using Impl = detail::TmaBuilderImpl<
TileShape_MNK, ClusterShape_MNK, EpilogueTileType, ElementAccumulator, ElementCompute,
ElementC, GmemLayoutTagC, AlignmentC, ElementD, GmemLayoutTagD, AlignmentD,
Schedule, ThreadOp, cutlass::epilogue::Sm90TmaWarpSpecialized<1,2,true>>;
Schedule, ThreadOp, cutlass::epilogue::Sm90TmaWarpSpecialized<StagesC,StagesD,DisableReuseSmemC>>;
public:
using CollectiveOp = typename Impl::CollectiveOp;
@@ -459,7 +479,7 @@ struct CollectiveBuilder<
AlignmentD,
Schedule,
cute::enable_if_t<cute::is_base_of_v<TmaWarpSpecializedBiasElementwiseBase, Schedule> ||
cute::is_base_of_v<TmaWarpSpecializedCooperativeBiasElementwiseBase, Schedule> >> {
cute::is_base_of_v<TmaWarpSpecializedCooperativeBiasElementwiseBase, Schedule> >> {
public:
using ThreadOp = thread::LinearCombinationBiasElementwise<
@@ -468,10 +488,12 @@ public:
Schedule::StoreT, typename Schedule::ElementBias>;
private:
static constexpr int StagesC = 1;
static constexpr int StagesD = 2;
using Impl = detail::TmaBuilderImpl<
TileShape_MNK, ClusterShape_MNK, EpilogueTileType, ElementAccumulator, ElementCompute,
ElementC, GmemLayoutTagC, AlignmentC, ElementD, GmemLayoutTagD, AlignmentD,
Schedule, ThreadOp, cutlass::epilogue::Sm90TmaWarpSpecializedBiasElementwise<1,2>>;
Schedule, ThreadOp, cutlass::epilogue::Sm90TmaWarpSpecializedBiasElementwise<StagesC,StagesD>>;
public:
using CollectiveOp = typename Impl::CollectiveOp;
include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp
+11 -11
View File
@@ -82,7 +82,6 @@ public:
//
// Type Aliases
//
// derived types of output thread level operator
using DispatchPolicy = Sm90TmaWarpSpecialized<StagesC_,StagesD_,DisableSmemReuseC_>;
using BlockTileShape = BlockTileShape_;
using EpilogueTile = EpilogueTile_;
@@ -108,7 +107,6 @@ public:
constexpr static bool iskThreadEpilogueOpWithBias = detail::IsThreadEpilogueOpWithBias<ThreadEpilogueOp>::value;
using AlignmentType = typename uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
static_assert(sizeof(ElementC) == 2, "Only 16b source supported for now");
static_assert(sizeof(ElementD) == 2, "Only 16b output supported for now");
static_assert(!is_layout<EpilogueTile>::value && is_tuple<EpilogueTile>::value, "EpilogueTile must be a cute::Tile or cute::Shape");
static_assert(rank(BlockTileShape{}) == 3, "BlockTileShape must be rank-3: [BLK_M,BLK_N,BLK_K]");
@@ -117,17 +115,19 @@ public:
static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
private:
using InternalElementC = std::conditional_t<std::is_void_v<ElementC>,ElementD,ElementC>; // prevents void ref breakages
constexpr static int StagesC = StagesC_;
constexpr static int StagesD = StagesD_;
constexpr static bool is_source_supported = ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::Default ||
ThreadEpilogueOp::kScale == cutlass::epilogue::thread::ScaleType::NoBetaScaling;
static_assert((std::is_void_v<ElementC> && not is_source_supported) || (not std::is_void_v<ElementC> && is_source_supported));
// internal optimization to reuse C shared memory for storing D
using SmemLayoutAtomBitsC = decltype(downcast<sizeof_bits<ElementC>::value>(SmemLayoutAtomC{}));
using SmemLayoutAtomBitsC = decltype(downcast<sizeof_bits<InternalElementC>::value>(SmemLayoutAtomC{}));
using SmemLayoutAtomBitsD = decltype(downcast<sizeof_bits<ElementD>::value>(SmemLayoutAtomD{}));
constexpr static bool ReuseSmemC = not DispatchPolicy::DisableSmemReuseC &&
is_source_supported &&
sizeof(ElementC) == sizeof(ElementD) &&
sizeof(InternalElementC) == sizeof(ElementD) &&
StrideC{} == StrideD{} &&
cute::is_same_v<SmemLayoutAtomBitsC,SmemLayoutAtomBitsD>;
@@ -152,7 +152,7 @@ public:
using LoadPipeline = cutlass::PipelineTransactionAsync<is_source_supported ? StagesC : 0>;
using LoadPipelineState = cutlass::PipelineState<is_source_supported ? StagesC : 0>;
constexpr static uint32_t TmaTransactionBytes =
size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof(ElementC));
size(take<0,2>(SmemLayoutC{})) * static_cast<uint32_t>(sizeof(InternalElementC));
// TMA pipeline for storing D
using StorePipeline = cutlass::PipelineTmaStore<ReuseSmemC ? StagesC : StagesD>;
@@ -161,8 +161,8 @@ public:
struct SharedStorage {
struct TensorStorage : aligned_struct<128> {
cute::conditional_t<not is_source_supported,
detail::EmptyStorage<ElementC>,
array_aligned<ElementC, size(SmemLayoutC{})>> smem_C;
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;
@@ -187,7 +187,7 @@ public:
struct Params {
using TMA_C = decltype(make_tma_copy(
CopyOpG2S{},
make_tensor(static_cast<ElementC const*>(nullptr),
make_tensor(static_cast<InternalElementC const*>(nullptr),
repeat_like(StrideC{}, int32_t(0)), StrideC{}),
SmemLayoutC{}(_,_,0)));
using TMA_D = decltype(make_tma_copy(
@@ -217,7 +217,7 @@ public:
auto M = get<0>(problem_shape_MNKL);
auto N = get<1>(problem_shape_MNKL);
auto L = get<3>(problem_shape_MNKL);
Tensor tensor_c = make_tensor(args.ptr_C, make_layout(make_shape(M,N,L), args.dC));
Tensor tensor_c = make_tensor(static_cast<InternalElementC const*>(args.ptr_C), make_layout(make_shape(M,N,L), args.dC));
Tensor tensor_d = make_tensor(args.ptr_D, make_layout(make_shape(M,N,L), args.dD));
typename Params::TMA_C tma_load_c = make_tma_copy(
CopyOpG2S{},
@@ -409,7 +409,7 @@ public:
// Allocate register tensors
auto tRS_rD_shape = take<0,3>(shape(thread_r2s.partition_S(bEsD))); // (R2S,R2S_M,R2S_N)
Tensor tRS_rC = make_tensor<ElementC>(tRS_rD_shape); // (R2S,R2S_M,R2S_N)
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)
// Vectorized fragment view for thread epilogue op
@@ -418,7 +418,7 @@ public:
Tensor tRS_rD_frg = recast<typename ThreadEpilogueOp::FragmentOutput>(tRS_rD);
// Partition for smem to register copy (tSR_)
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R,ElementC>{}, tiled_r2s);
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R,InternalElementC>{}, tiled_r2s);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
Tensor tSR_sC = thread_s2r.partition_S(bEsC); // (S2R,S2R_M,S2R_N,EPI_M,EPI_N)
Tensor tSR_rC = thread_s2r.retile_D(tRS_rC); // (S2R,S2R_M,S2R_N)
include/cutlass/epilogue/thread/linear_combination_tensor_broadcast.hpp
+1
View File
@@ -130,6 +130,7 @@ public:
using ActivationFunctor = ActivationFunctor_<ElementCompute>;
static constexpr int kCount = 1;
static constexpr ScaleType::Kind kScale = Scale;
using FragmentOutput = Array<ElementOutput, kCount>;
using FragmentAccumulator = Array<ElementAccumulator, kCount>;
include/cutlass/epilogue/threadblock/epilogue.h
+1 -1
View File
@@ -323,7 +323,7 @@ public:
OutputTileIterator destination_iterator, ///< Tile iterator for destination
OutputTileIterator source_iterator) ///< Threadblock tile coordinate in GEMM (in units of threadblock tiles)
{
// Redcuce peer accumulator fragments into one fragment
// 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);