CUTLASS 3.3.0 (#1167)
* Release 3.3.0 Adds support for mixed precision GEMMs On Hopper and Ampere Adds support for < 16B aligned GEMMs on Hopper Enhancements to EVT Enhancements to Python interface Enhancements to Sub-byte type handling in CuTe Several other bug-fixes and performance improvements. * minor doc update
This commit is contained in:
Pradeep Ramani
@@ -69,11 +69,11 @@ constexpr int
|
||||
compute_stage_count_or_override(StageCountAutoCarveout<carveout_bytes> stage_count) {
|
||||
// 32 bytes to account for barriers etc.
|
||||
constexpr int stage_barrier_bytes = 32;
|
||||
constexpr int a_bytes = static_cast<int>(sizeof(ElementA));
|
||||
constexpr int b_bytes = static_cast<int>(sizeof(ElementB));
|
||||
constexpr int a_bits = static_cast<int>(sizeof_bits<ElementA>::value);
|
||||
constexpr int b_bits = static_cast<int>(sizeof_bits<ElementB>::value);
|
||||
constexpr int stage_bytes =
|
||||
(a_bytes * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
|
||||
(b_bytes * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
|
||||
(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) / 8 +
|
||||
(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) / 8 +
|
||||
stage_barrier_bytes;
|
||||
|
||||
return (CapacityBytes - carveout_bytes) / stage_bytes;
|
||||
@@ -95,13 +95,39 @@ is_warpspecialized_transpose_B(){
|
||||
constexpr bool IsInputSizeTwoBytes = is_input_size_two_bytes<ElementA, ElementB>();
|
||||
constexpr bool IsLayoutAmnBmn = cutlass::gemm::detail::is_mn_major_A<LayoutA>() &&
|
||||
cutlass::gemm::detail::is_mn_major_B<LayoutB>();
|
||||
constexpr bool IsWarpSpecialized = cute::is_base_of_v<KernelTmaWarpSpecialized, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelTmaWarpSpecializedPingpong, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelTmaWarpSpecializedCooperative, KernelScheduleType>;
|
||||
constexpr bool IsWarpSpecialized = cute::is_base_of_v<KernelTmaWarpSpecialized, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelTmaWarpSpecializedPingpong, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelTmaWarpSpecializedCooperative, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelCpAsyncWarpSpecialized, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelCpAsyncWarpSpecializedPingpong, KernelScheduleType> ||
|
||||
cute::is_base_of_v<KernelCpAsyncWarpSpecializedCooperative, KernelScheduleType>;
|
||||
constexpr bool IsWarpSpecializedTransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn && IsWarpSpecialized;
|
||||
return IsWarpSpecializedTransposeB;
|
||||
}
|
||||
|
||||
template <typename ElementA, typename ElementB>
|
||||
struct Sm90TypeWidths {
|
||||
static constexpr bool IsElementALarger = (cute::sizeof_bits_v<ElementA>) > cute::sizeof_bits_v<ElementB>;
|
||||
using WideType = cute::conditional_t<IsElementALarger, ElementA, ElementB>;
|
||||
using NarrowType = cute::conditional_t<IsElementALarger, ElementB, ElementA>;
|
||||
};
|
||||
|
||||
|
||||
template <class ElementA, class LayoutA, class ElementB, class LayoutB>
|
||||
constexpr bool
|
||||
sm90_is_narrow_type_k_major() {
|
||||
using Widths = Sm90TypeWidths<ElementA, ElementB>;
|
||||
using NarrowType = typename Widths::NarrowType;
|
||||
using WideType = typename Widths::WideType;
|
||||
|
||||
constexpr bool IsANarrow = cute::is_same_v<NarrowType, ElementA>;
|
||||
constexpr cute::GMMA::Major NarrowGmmaMajor = IsANarrow ? detail::gmma_rs_tag_to_major_A<LayoutA>() :
|
||||
detail::gmma_rs_tag_to_major_B<LayoutB>();
|
||||
|
||||
constexpr bool IsNarrowLayoutKMajor = NarrowGmmaMajor == cute::GMMA::Major::K;
|
||||
return IsNarrowLayoutKMajor;
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -297,6 +323,135 @@ struct CollectiveBuilder<
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// GMMA_TMA_WS_RS Mixed GEMM
|
||||
template <
|
||||
class ElementPairA_,
|
||||
class GmemLayoutPairA_,
|
||||
int AlignmentA,
|
||||
class ElementPairB_,
|
||||
class GmemLayoutPairB_,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
arch::OpClassTensorOp,
|
||||
ElementPairA_,
|
||||
GmemLayoutPairA_,
|
||||
AlignmentA,
|
||||
ElementPairB_,
|
||||
GmemLayoutPairB_,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
(cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedMixedInput> ||
|
||||
cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedPingpongMixedInput> ||
|
||||
cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedCooperativeMixedInput>)>
|
||||
> {
|
||||
|
||||
public:
|
||||
static constexpr bool IsATransformed = cute::sizeof_bits_v<ElementPairA_> < cute::sizeof_bits_v<ElementPairB_>;
|
||||
|
||||
// Split out items for processessing, no splitting for now since scales aren't supported.
|
||||
using ElementA = ElementPairA_;
|
||||
using ElementB = ElementPairB_;
|
||||
|
||||
using GmemLayoutA = GmemLayoutPairA_;
|
||||
using GmemLayoutB = GmemLayoutPairB_;
|
||||
|
||||
static_assert(is_static<TileShape_MNK>::value);
|
||||
static_assert(is_static<ClusterShape_MNK>::value);
|
||||
static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>(),
|
||||
"Should meet TMA alignment requirement\n");
|
||||
#ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED
|
||||
static_assert(cutlass::detail::dependent_false<ElementA> == 0, "Unsupported Toolkit for SM90 Collective Builder\n");
|
||||
#endif
|
||||
static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_rs_tag_to_major_A<GmemLayoutA>();
|
||||
static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_rs_tag_to_major_B<GmemLayoutB>();
|
||||
static constexpr bool IsWarpSpecializedTransposeB = detail::is_warpspecialized_transpose_B<
|
||||
ElementA, GmemLayoutA, ElementB, GmemLayoutB, KernelScheduleType>();
|
||||
static_assert(!IsWarpSpecializedTransposeB, "Mixed input GEMM does not support WS transpose B.");
|
||||
|
||||
// If A is scaled, then we don't need to swap. Otherwise, we must ensure B goes to RF and we must swap the operands.
|
||||
static constexpr bool SwapAB = !IsATransformed;
|
||||
static_assert(detail::sm90_is_narrow_type_k_major<ElementA, GmemLayoutA, ElementB, GmemLayoutB>(), "The narrow type must be K-major.");
|
||||
|
||||
static_assert((IsATransformed && (cute::sizeof_bits_v<ElementA> <= 8) && (sizeof(ElementB) == 2)) ||
|
||||
(!IsATransformed && (cute::sizeof_bits_v<ElementB> <= 8) && (sizeof(ElementA) == 2)) ||
|
||||
(GmmaMajorA == cute::GMMA::Major::K && GmmaMajorB == cute::GMMA::Major::K),
|
||||
"The unscaled element must be 2 bytes OR both inputs must be K-major");
|
||||
// When we relax the above assertion, we must handle setting the tile mma GmmaMajorB correctly.
|
||||
static constexpr cute::GMMA::Major TiledMmaGmmaMajorB = SwapAB ? GmmaMajorA : GmmaMajorB;
|
||||
|
||||
using ElementMma = cute::conditional_t<IsATransformed, ElementB, ElementA>;
|
||||
using AtomLayoutMNK = cute::conditional_t<cute::is_same_v<KernelScheduleType, KernelTmaWarpSpecializedCooperativeMixedInput>,
|
||||
Layout<Shape<_2,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
|
||||
|
||||
using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::rs_op_selector<
|
||||
ElementMma, ElementMma, ElementAccumulator, TileShape_MNK, GMMA::Major::K, TiledMmaGmmaMajorB>(), AtomLayoutMNK{}));
|
||||
|
||||
using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{})));
|
||||
using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{})));
|
||||
using SmemLayoutAtomA = decltype(detail::rs_smem_selector<GmmaMajorA, ElementA,
|
||||
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>());
|
||||
using SmemLayoutAtomB = decltype(detail::rs_smem_selector<GmmaMajorB, ElementB,
|
||||
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>());
|
||||
|
||||
using RealElementA = cute::conditional_t<SwapAB, ElementB, ElementA>;
|
||||
using RealElementB = cute::conditional_t<SwapAB, ElementA, ElementB>;
|
||||
static constexpr int PipelineStages = detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes,
|
||||
RealElementA, RealElementB, TileShape_MNK>(StageCountType{});
|
||||
|
||||
using SmemCopyAtomA = cute::conditional_t<SwapAB, void, Copy_Atom<cute::DefaultCopy, ElementA>>;
|
||||
using SmemCopyAtomB = cute::conditional_t<SwapAB, Copy_Atom<cute::DefaultCopy, ElementB>, void>;
|
||||
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<PipelineStages, ClusterShape_MNK, KernelScheduleType>;
|
||||
|
||||
// We pack the scale data with the operand that will be optionally scaled and converted before MMA.
|
||||
using StrideAPair = TagToStrideA_t<GmemLayoutA>;
|
||||
using StrideBPair = TagToStrideB_t<GmemLayoutB>;
|
||||
|
||||
using GmemTiledCopyAPair = GmemTiledCopyA;
|
||||
using SmemLayoutAtomAPair = SmemLayoutAtomA;
|
||||
using SmemCopyAtomAPair = SmemCopyAtomA;
|
||||
|
||||
using GmemTiledCopyBPair = GmemTiledCopyB;
|
||||
using SmemLayoutAtomBPair = SmemLayoutAtomB;
|
||||
using SmemCopyAtomBPair = SmemCopyAtomB;
|
||||
|
||||
|
||||
// If the src type of the converter is the same as ElementA,
|
||||
// interpret this as if the user wanted to apply the scale to the A matrix.
|
||||
using CollectiveOp = CollectiveMma<
|
||||
DispatchPolicy,
|
||||
TileShape_MNK,
|
||||
ElementPairA_,
|
||||
StrideAPair,
|
||||
ElementPairB_,
|
||||
StrideBPair,
|
||||
TiledMma,
|
||||
GmemTiledCopyAPair,
|
||||
SmemLayoutAtomAPair,
|
||||
SmemCopyAtomAPair,
|
||||
cute::identity,
|
||||
GmemTiledCopyBPair,
|
||||
SmemLayoutAtomBPair,
|
||||
SmemCopyAtomBPair,
|
||||
cute::identity
|
||||
>;
|
||||
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// GMMA_TMA_WS_FP8_FAST_ACCUM_SS
|
||||
template <
|
||||
class ElementA,
|
||||
@@ -490,7 +645,8 @@ template <
|
||||
class StageCountType,
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
struct [[deprecated("Use one of KernelCpAsyncWarpSpecialized schedules instead")]]
|
||||
CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
arch::OpClassTensorOp,
|
||||
ElementA,
|
||||
@@ -506,6 +662,61 @@ struct CollectiveBuilder<
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
cute::is_same_v<KernelScheduleType, KernelMultistage>>
|
||||
> {
|
||||
// Map to warp-specialized kernels for better performance
|
||||
using CollectiveOp = typename CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
arch::OpClassTensorOp,
|
||||
ElementA,
|
||||
GmemLayoutA,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
GmemLayoutB,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
KernelCpAsyncWarpSpecialized
|
||||
>::CollectiveOp;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// GMMA_CpAsync_WS_SS
|
||||
template <
|
||||
class ElementA,
|
||||
class GmemLayoutA,
|
||||
int AlignmentA,
|
||||
class ElementB,
|
||||
class GmemLayoutB,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
arch::OpClassTensorOp,
|
||||
ElementA,
|
||||
GmemLayoutA,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
GmemLayoutB,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
(cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecialized> ||
|
||||
cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecializedCooperative> ||
|
||||
cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecializedPingpong>) &&
|
||||
not detail::is_use_rmem_A<ElementA, GmemLayoutA, ElementB, GmemLayoutB>()
|
||||
>
|
||||
> {
|
||||
static_assert(is_static<TileShape_MNK>::value);
|
||||
static_assert(is_static<ClusterShape_MNK>::value);
|
||||
@@ -523,14 +734,19 @@ struct CollectiveBuilder<
|
||||
static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A<ElementA, GmemLayoutA>();
|
||||
static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B<ElementB, GmemLayoutB>();
|
||||
|
||||
using AtomLayoutMNK = cute::conditional_t<cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecializedCooperative>,
|
||||
Layout<Shape<cute::Int<(size<0>(TileShape_MNK{}) < 128) ? 1 : 2>,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
|
||||
|
||||
using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector<
|
||||
MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>()));
|
||||
MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{}));
|
||||
|
||||
static constexpr int NumLoadWarpGroups = cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecialized> ? 2 : 1;
|
||||
|
||||
using GmemTiledCopyA = decltype(detail::make_cp_async_gmem_tiled_copy<
|
||||
128, ElementA, AlignmentA, TagToStrideA_t<GmemLayoutA>,
|
||||
NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementA, AlignmentA, TagToStrideA_t<GmemLayoutA>,
|
||||
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
|
||||
using GmemTiledCopyB = decltype(detail::make_cp_async_gmem_tiled_copy<
|
||||
128, ElementB, AlignmentB, TagToStrideB_t<GmemLayoutB>,
|
||||
NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementB, AlignmentB, TagToStrideB_t<GmemLayoutB>,
|
||||
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
|
||||
|
||||
using SmemLayoutAtomA = decltype(detail::ss_smem_selector<
|
||||
@@ -541,8 +757,11 @@ struct CollectiveBuilder<
|
||||
static constexpr int PipelineStages = detail::compute_stage_count_or_override<
|
||||
detail::sm90_smem_capacity_bytes, MmaElementA, MmaElementB, TileShape_MNK>(StageCountType{});
|
||||
|
||||
using DispatchPolicy = MainloopSm90CpAsyncGmmaWarpSpecialized<
|
||||
PipelineStages, ClusterShape_MNK, KernelScheduleType>;
|
||||
|
||||
using CollectiveOp = CollectiveMma<
|
||||
MainloopSm90CpAsyncGmma<PipelineStages, ClusterShape_MNK>,
|
||||
DispatchPolicy,
|
||||
TileShape_MNK,
|
||||
ElementA,
|
||||
TagToStrideA_t<GmemLayoutA>,
|
||||
@@ -562,6 +781,110 @@ struct CollectiveBuilder<
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// GMMA_CpAsync_WS_RS
|
||||
template <
|
||||
class ElementA,
|
||||
class GmemLayoutA,
|
||||
int AlignmentA,
|
||||
class ElementB,
|
||||
class GmemLayoutB,
|
||||
int AlignmentB,
|
||||
class ElementAccumulator,
|
||||
class TileShape_MNK,
|
||||
class ClusterShape_MNK,
|
||||
class StageCountType,
|
||||
class KernelScheduleType
|
||||
>
|
||||
struct CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
arch::OpClassTensorOp,
|
||||
ElementA,
|
||||
GmemLayoutA,
|
||||
AlignmentA,
|
||||
ElementB,
|
||||
GmemLayoutB,
|
||||
AlignmentB,
|
||||
ElementAccumulator,
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
KernelScheduleType,
|
||||
cute::enable_if_t<
|
||||
(cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecialized> ||
|
||||
cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecializedCooperative> ||
|
||||
cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecializedPingpong>) &&
|
||||
detail::is_use_rmem_A<ElementA, GmemLayoutA, ElementB, GmemLayoutB>()
|
||||
>
|
||||
> {
|
||||
static_assert(is_static<TileShape_MNK>::value);
|
||||
static_assert(is_static<ClusterShape_MNK>::value);
|
||||
#ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED
|
||||
static_assert(cutlass::detail::dependent_false<ElementA> == 0, "Unsupported Toolkit for SM90 Collective Builder\n");
|
||||
#endif
|
||||
|
||||
// For fp32 types, map to tf32 MMA value type
|
||||
using MmaElementA = cute::conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
|
||||
using MmaElementB = cute::conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
|
||||
|
||||
static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::cp_async_min_alignment_bytes>(),
|
||||
"Minimum alignment required for cp.async is 4B.");
|
||||
|
||||
static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_rs_tag_to_major_A<GmemLayoutA>();
|
||||
static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_rs_tag_to_major_B<GmemLayoutB>();
|
||||
static constexpr bool SwapAB = detail::is_swapAB<ElementA, GmemLayoutA, ElementB, GmemLayoutB>();
|
||||
static constexpr bool IsWarpSpecializedTransposeB = detail::is_warpspecialized_transpose_B<
|
||||
ElementA, GmemLayoutA, ElementB, GmemLayoutB, KernelScheduleType>();
|
||||
|
||||
using AtomLayoutMNK = cute::conditional_t<cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecializedCooperative>,
|
||||
Layout<Shape<cute::Int<(size<0>(TileShape_MNK{}) < 128) ? 1 : 2>,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
|
||||
|
||||
using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::rs_op_selector<
|
||||
MmaElementA, MmaElementB, ElementAccumulator, TileShape_MNK, GMMA::Major::K, GMMA::Major::K>(), AtomLayoutMNK{}));
|
||||
|
||||
static constexpr int NumLoadWarpGroups = 1;
|
||||
|
||||
using GmemTiledCopyA = decltype(detail::make_cp_async_gmem_tiled_copy<
|
||||
NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementA, AlignmentA, TagToStrideA_t<GmemLayoutA>,
|
||||
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
|
||||
using GmemTiledCopyB = decltype(detail::make_cp_async_gmem_tiled_copy<
|
||||
NumThreadsPerWarpGroup * NumLoadWarpGroups, ElementB, AlignmentB, TagToStrideB_t<GmemLayoutB>,
|
||||
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
|
||||
|
||||
using SmemLayoutAtomA = decltype(detail::rs_smem_selector<GmmaMajorA, MmaElementA,
|
||||
decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>());
|
||||
using SmemLayoutAtomB = decltype(detail::rs_smem_selector<GmmaMajorB, MmaElementB,
|
||||
decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{})), IsWarpSpecializedTransposeB>());
|
||||
|
||||
static constexpr int PipelineStages = detail::compute_stage_count_or_override<
|
||||
detail::sm90_smem_capacity_bytes, MmaElementA, MmaElementB, TileShape_MNK>(StageCountType{});
|
||||
|
||||
using DispatchPolicy = MainloopSm90CpAsyncGmmaRmemAWarpSpecialized<
|
||||
PipelineStages, ClusterShape_MNK, KernelScheduleType>;
|
||||
|
||||
using SmemCopyAtomA = cute::conditional_t<SwapAB, void, Copy_Atom<cute::DefaultCopy, ElementA>>;
|
||||
using SmemCopyAtomB = cute::conditional_t<SwapAB, Copy_Atom<cute::DefaultCopy, ElementB>, void>;
|
||||
|
||||
using CollectiveOp = CollectiveMma<
|
||||
DispatchPolicy,
|
||||
TileShape_MNK,
|
||||
ElementA,
|
||||
TagToStrideA_t<GmemLayoutA>,
|
||||
ElementB,
|
||||
TagToStrideB_t<GmemLayoutB>,
|
||||
TiledMma,
|
||||
GmemTiledCopyA,
|
||||
SmemLayoutAtomA,
|
||||
SmemCopyAtomA,
|
||||
cute::identity,
|
||||
GmemTiledCopyB,
|
||||
SmemLayoutAtomB,
|
||||
SmemCopyAtomB,
|
||||
cute::identity
|
||||
>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// GMMA auto kernel schedule
|
||||
template <
|
||||
class ElementA,
|
||||
@@ -601,15 +924,26 @@ struct CollectiveBuilder<
|
||||
static constexpr bool IsTmaCompatible = detail::is_aligned<
|
||||
ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>();
|
||||
|
||||
static constexpr bool IsMixedWidthInput = cute::sizeof_bits_v<ElementA> != cute::sizeof_bits_v<ElementB>;
|
||||
|
||||
#if ((__CUDACC_VER_MAJOR__ > 12) || ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 1)))
|
||||
// Persistent schedules perform best for CUDA Toolkits with version >= 12.1
|
||||
// KernelTmaWarpSpecializedCooperative requires TileShape_M to be at least 128
|
||||
using KernelWarpSpecializedSchedule = cute::conditional_t<size<0>(TileShape_MNK{}) == Int<64>{},
|
||||
using KernelTmaWarpSpecializedScheduleSameInput = cute::conditional_t<size<0>(TileShape_MNK{}) == Int<64>{},
|
||||
KernelTmaWarpSpecializedPingpong, KernelTmaWarpSpecializedCooperative>;
|
||||
|
||||
using KernelTmaWarpSpecializedScheduleMixedInput = cute::conditional_t<size<0>(TileShape_MNK{}) == Int<64>{},
|
||||
KernelTmaWarpSpecializedPingpongMixedInput, KernelTmaWarpSpecializedCooperativeMixedInput>;
|
||||
|
||||
using KernelTmaWarpSpecializedSchedule = cute::conditional_t<IsMixedWidthInput, KernelTmaWarpSpecializedScheduleMixedInput, KernelTmaWarpSpecializedScheduleSameInput>;
|
||||
#else
|
||||
using KernelWarpSpecializedSchedule = KernelTmaWarpSpecialized;
|
||||
using KernelTmaWarpSpecializedSchedule = cute::conditional_t<IsMixedWidthInput, KernelTmaWarpSpecializedMixedInput, KernelTmaWarpSpecialized>;
|
||||
#endif
|
||||
|
||||
// Non-persistent schedule is a safer choice for CpAsync kernels due to register pressure
|
||||
using KernelCpAsyncWarpSpecializedSchedule = KernelCpAsyncWarpSpecialized;
|
||||
using KernelSchedule = cute::conditional_t<IsTmaCompatible, KernelTmaWarpSpecializedSchedule, KernelCpAsyncWarpSpecializedSchedule>;
|
||||
static_assert((cute::is_same_v<KernelSchedule, KernelTmaWarpSpecializedSchedule> && IsMixedWidthInput) || !IsMixedWidthInput, "Only TMA warp specialized kernels are supported for mixed width input.");
|
||||
using CollectiveOp = typename CollectiveBuilder<
|
||||
arch::Sm90,
|
||||
arch::OpClassTensorOp,
|
||||
@@ -623,7 +957,7 @@ static constexpr bool IsTmaCompatible = detail::is_aligned<
|
||||
TileShape_MNK,
|
||||
ClusterShape_MNK,
|
||||
StageCountType,
|
||||
cute::conditional_t<IsTmaCompatible, KernelWarpSpecializedSchedule, KernelMultistage>
|
||||
KernelSchedule
|
||||
>::CollectiveOp;
|
||||
};
|
||||
|
||||
|
||||
@@ -67,9 +67,11 @@ struct CollectiveMma {
|
||||
|
||||
#include "cutlass/gemm/collective/sm70_mma_twostage.hpp"
|
||||
#include "cutlass/gemm/collective/sm80_mma_multistage.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_ss.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -0,0 +1,662 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cute/arch/copy_sm90.hpp"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
|
||||
#include "cute/algorithm/functional.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
#include "cute/algorithm/gemm.hpp"
|
||||
#include "cute/tensor_predicate.hpp"
|
||||
#include "cute/numeric/arithmetic_tuple.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp"
|
||||
#include "cutlass/trace.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// WarpSpecialized Mainloop
|
||||
template <
|
||||
int Stages,
|
||||
class ClusterShape_,
|
||||
class TileShape_,
|
||||
class KernelSchedule,
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
class ElementB_,
|
||||
class StrideB_,
|
||||
class TiledMma_,
|
||||
class GmemTiledCopyA_,
|
||||
class SmemLayoutAtomA_,
|
||||
class SmemCopyAtomA_,
|
||||
class TransformA_,
|
||||
class GmemTiledCopyB_,
|
||||
class SmemLayoutAtomB_,
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90CpAsyncGmmaRmemAWarpSpecialized<Stages,ClusterShape_,KernelSchedule>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
ElementB_,
|
||||
StrideB_,
|
||||
TiledMma_,
|
||||
GmemTiledCopyA_,
|
||||
SmemLayoutAtomA_,
|
||||
SmemCopyAtomA_,
|
||||
TransformA_,
|
||||
GmemTiledCopyB_,
|
||||
SmemLayoutAtomB_,
|
||||
SmemCopyAtomB_,
|
||||
TransformB_>
|
||||
{
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90CpAsyncGmmaRmemAWarpSpecialized<Stages,ClusterShape_,KernelSchedule>;
|
||||
using TileShape = TileShape_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ElementA = ElementA_;
|
||||
using StrideA = StrideA_;
|
||||
using ElementB = ElementB_;
|
||||
using StrideB = StrideB_;
|
||||
using TiledMma = TiledMma_;
|
||||
using ElementAccumulator = typename TiledMma::ValTypeC;
|
||||
using GmemTiledCopyA = GmemTiledCopyA_;
|
||||
using GmemTiledCopyB = GmemTiledCopyB_;
|
||||
using SmemLayoutAtomA = SmemLayoutAtomA_;
|
||||
using SmemLayoutAtomB = SmemLayoutAtomB_;
|
||||
using SmemCopyAtomA = SmemCopyAtomA_;
|
||||
using SmemCopyAtomB = SmemCopyAtomB_;
|
||||
|
||||
// Swap and transpose A/B for A k-major layout and B mn-major layout since WGMMA is k-major only (e.g. tf32, Fp32, Int8, Fp8 WGMMA)
|
||||
static constexpr bool IsLayoutAkBmn =
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> &&
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor>;
|
||||
|
||||
static constexpr bool IsInputSizeTwoBytes = sizeof(ElementA) == 2 && sizeof(ElementB) == 2;
|
||||
static constexpr bool SwapAB = !IsInputSizeTwoBytes && IsLayoutAkBmn;
|
||||
using InternalGmemTiledCopyA = cute::conditional_t<!SwapAB, GmemTiledCopyA, GmemTiledCopyB>;
|
||||
using InternalGmemTiledCopyB = cute::conditional_t<!SwapAB, GmemTiledCopyB, GmemTiledCopyA>;
|
||||
using InternalSmemLayoutAtomA = cute::conditional_t<!SwapAB, SmemLayoutAtomA, SmemLayoutAtomB>;
|
||||
using InternalSmemLayoutAtomB = cute::conditional_t<!SwapAB, SmemLayoutAtomB, SmemLayoutAtomA>;
|
||||
using InternalSmemCopyAtomA = cute::conditional_t<!SwapAB, SmemCopyAtomA, SmemCopyAtomB>;
|
||||
using InternalSmemCopyAtomB = cute::conditional_t<!SwapAB, SmemCopyAtomB, SmemCopyAtomA>;
|
||||
// TMA converts f32 input to tf32 when copying from GMEM to SMEM
|
||||
// For all other types, cast to size equivalent uint type to avoid any rounding by TMA.
|
||||
static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>;
|
||||
static constexpr bool ConvertF32toTF32B = cute::is_same_v<float, ElementB>;
|
||||
using ConvertedElementA = cute::conditional_t<ConvertF32toTF32A, tfloat32_t, uint_bit_t<sizeof_bits_v<ElementA>>>;
|
||||
using ConvertedElementB = cute::conditional_t<ConvertF32toTF32B, tfloat32_t, uint_bit_t<sizeof_bits_v<ElementB>>>;
|
||||
using InternalElementA = cute::conditional_t<!SwapAB, ConvertedElementA, ConvertedElementB>;
|
||||
using InternalElementB = cute::conditional_t<!SwapAB, ConvertedElementB, ConvertedElementA>;
|
||||
using InternalStrideA = cute::conditional_t<!SwapAB, StrideA, StrideB>;
|
||||
using InternalStrideB = cute::conditional_t<!SwapAB, StrideB, StrideA>;
|
||||
|
||||
using TransformA = TransformA_;
|
||||
using TransformB = TransformB_;
|
||||
using ArchTag = typename DispatchPolicy::ArchTag;
|
||||
|
||||
using MainloopPipeline = cutlass::PipelineAsync<DispatchPolicy::Stages>;
|
||||
using PipelineState = typename MainloopPipeline::PipelineState;
|
||||
using PipelineParams = typename MainloopPipeline::Params;
|
||||
|
||||
static_assert(rank(InternalSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<0>(TileShape{}) % size<0>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
static_assert(rank(InternalSmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<1>(TileShape{}) % size<0>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
using SmemLayoutA = decltype(tile_to_shape(
|
||||
InternalSmemLayoutAtomA{},
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
using SmemLayoutB = decltype(tile_to_shape(
|
||||
InternalSmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
|
||||
// If A mn-layout and B mn-layout, transposing B matrix since WGMMA is k-major only (e.g. tf32, fp32, fp8, int8).
|
||||
static constexpr bool IsLayoutAmnBmn =
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::ColumnMajor> &&
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor>;
|
||||
static constexpr bool TransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn;
|
||||
using TransposeOperandB = decltype(cutlass::transform::collective::detail::make_transpose_operand_b(
|
||||
0, 0, TiledMma{}, SmemLayoutB{}, InternalSmemLayoutAtomB{},
|
||||
InternalElementB{}, cute::bool_constant<TransposeB>{}));
|
||||
|
||||
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more.");
|
||||
static_assert(not cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
|
||||
cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
|
||||
"MMA atom must source A from rmem and B operand from smem_desc for this mainloop.");
|
||||
|
||||
using GmmaSmemLayoutAtomB = decltype(transform::collective::detail::gmma_smem_transpose_or_passthrough<
|
||||
TransposeB, InternalSmemLayoutAtomB, InternalElementB>());
|
||||
|
||||
// SmemLayoutB for GMMA is different from SmemLayoutB for TMA if TransposeB
|
||||
using GmmaSmemLayoutB = decltype(tile_to_shape(
|
||||
GmmaSmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
|
||||
static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time.");
|
||||
static_assert(TransposeB xor (cute::is_same_v<SmemLayoutB, GmmaSmemLayoutB>),
|
||||
"Should be same layout if not TransposeB.");
|
||||
static_assert(!TransposeB || ((size<1>(SmemLayoutB{}) * sizeof_bits<InternalElementB>::value) / 8) == 128,
|
||||
"SmemLayoutB K must be 128bytes to be transposed.");
|
||||
static_assert(!transform::collective::detail::use_universal_transposition<InternalSmemLayoutAtomB, InternalElementB>(),
|
||||
"Warp specialized ARF kernels have not supported universal B transposition yet.");
|
||||
|
||||
struct SharedStorage
|
||||
{
|
||||
struct TensorStorage : cute::aligned_struct<256> {
|
||||
cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>, 256> smem_A;
|
||||
cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>, 256> smem_B;
|
||||
} tensors;
|
||||
|
||||
using PipelineStorage = typename MainloopPipeline::SharedStorage;
|
||||
PipelineStorage pipeline;
|
||||
};
|
||||
using TensorStorage = typename SharedStorage::TensorStorage;
|
||||
using PipelineStorage = typename SharedStorage::PipelineStorage;
|
||||
|
||||
// Host side kernel arguments
|
||||
struct Arguments {
|
||||
ElementA const* ptr_A = nullptr;
|
||||
StrideA dA{};
|
||||
ElementB const* ptr_B = nullptr;
|
||||
StrideB dB{};
|
||||
uint32_t mma_promotion_interval = 4;
|
||||
};
|
||||
|
||||
// Device side kernel params
|
||||
struct Params {
|
||||
InternalElementA const* ptr_A = nullptr;
|
||||
InternalStrideA dA{};
|
||||
InternalElementB const* ptr_B = nullptr;
|
||||
InternalStrideB dB{};
|
||||
uint32_t mma_promotion_interval = 4;
|
||||
};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
template <class ProblemShape>
|
||||
static constexpr Params
|
||||
to_underlying_arguments(
|
||||
[[maybe_unused]] ProblemShape const& problem_shape,
|
||||
Arguments const& args,
|
||||
[[maybe_unused]] void* workspace) {
|
||||
if constexpr (not SwapAB) {
|
||||
return {
|
||||
reinterpret_cast<InternalElementA const*>(args.ptr_A),
|
||||
args.dA,
|
||||
reinterpret_cast<InternalElementB const*>(args.ptr_B),
|
||||
args.dB
|
||||
};
|
||||
}
|
||||
else {
|
||||
return {
|
||||
reinterpret_cast<InternalElementA const*>(args.ptr_B),
|
||||
args.dB,
|
||||
reinterpret_cast<InternalElementB const*>(args.ptr_A),
|
||||
args.dA
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
template<class ProblemShape>
|
||||
CUTLASS_HOST_DEVICE static bool
|
||||
can_implement(
|
||||
ProblemShape const& problem_shape,
|
||||
[[maybe_unused]] Arguments const& args) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
bool implementable = true;
|
||||
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyA::NumValSrc>(cute::make_shape(M,K,L), StrideA{});
|
||||
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyB::NumValSrc>(cute::make_shape(N,K,L), StrideB{});
|
||||
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
|
||||
}
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
|
||||
static constexpr int K_PIPE_MMAS = 1;
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class TensorA,
|
||||
class TensorB,
|
||||
class KTileIterator,
|
||||
class ResidueMNK
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
load(
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write,
|
||||
TensorA const& gA_in,
|
||||
TensorB const& gB_in,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
ResidueMNK residue_mnk,
|
||||
int thread_idx,
|
||||
TensorStorage& shared_tensors)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
static_assert(is_gmem<TensorA>::value, "A tensor must be gmem resident.");
|
||||
static_assert(is_gmem<TensorB>::value, "B tensor must be gmem resident.");
|
||||
|
||||
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
// Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k)
|
||||
// This aligns the tensor with BLK_K for all but the 0th k_tile
|
||||
Tensor gA = domain_offset(make_coord(0, get<2>(residue_mnk), 0), gA_in);
|
||||
Tensor gB = domain_offset(make_coord(0, get<2>(residue_mnk), 0), gB_in);
|
||||
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
InternalGmemTiledCopyA gmem_tiled_copy_a;
|
||||
InternalGmemTiledCopyB gmem_tiled_copy_b;
|
||||
auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx);
|
||||
auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx);
|
||||
|
||||
Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
|
||||
Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
|
||||
Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k)
|
||||
Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE)
|
||||
|
||||
// Allocate predicate tensors for m and n
|
||||
Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});
|
||||
Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
|
||||
|
||||
// Construct identity layout for sA and sB
|
||||
Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k)
|
||||
Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k)
|
||||
|
||||
// Repeat the partitioning with identity layouts
|
||||
Tensor tAcA = gmem_thr_copy_a.partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k)
|
||||
Tensor tBcB = gmem_thr_copy_b.partition_S(cB); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k)
|
||||
|
||||
// Set predicates for m bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < size<0>(tApA); ++m) {
|
||||
tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk); // blk_m coord < residue_m
|
||||
}
|
||||
// Set predicates for n bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < size<0>(tBpB); ++n) {
|
||||
tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk); // blk_n coord < residue_n
|
||||
}
|
||||
|
||||
// 0-th stage with predication on k to account for residue
|
||||
{
|
||||
// LOCK smem_pipe_write for _writing_
|
||||
pipeline.producer_acquire(smem_pipe_write);
|
||||
int write_stage = smem_pipe_write.index();
|
||||
|
||||
// Copy gmem to smem for *k_tile_iter, predicating for k residue
|
||||
Tensor tAgAk = tAgA(_,_,_,*k_tile_iter);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tAsA); ++k) {
|
||||
if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gA shifted)
|
||||
copy_if(gmem_tiled_copy_a, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,write_stage));
|
||||
}
|
||||
else {
|
||||
clear(tAsA(_,_,k,write_stage));
|
||||
}
|
||||
}
|
||||
Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tBsB); ++k) {
|
||||
if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gB shifted)
|
||||
copy_if(gmem_tiled_copy_b, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,write_stage));
|
||||
}
|
||||
else {
|
||||
clear(tBsB(_,_,k,write_stage));
|
||||
}
|
||||
}
|
||||
|
||||
++k_tile_iter;
|
||||
--k_tile_count;
|
||||
|
||||
// UNLOCK smem_pipe_write
|
||||
pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive);
|
||||
|
||||
// Advance smem_pipe_write
|
||||
++smem_pipe_write;
|
||||
}
|
||||
|
||||
// Mainloop
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > 0; --k_tile_count) {
|
||||
// LOCK smem_pipe_write for _writing_
|
||||
pipeline.producer_acquire(smem_pipe_write);
|
||||
int write_stage = smem_pipe_write.index();
|
||||
|
||||
// Copy gmem to smem for *k_tile_iter
|
||||
copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
++k_tile_iter;
|
||||
|
||||
// UNLOCK smem_pipe_write
|
||||
pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive);
|
||||
|
||||
// Advance smem_pipe_write
|
||||
++smem_pipe_write;
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
|
||||
CUTLASS_DEVICE void
|
||||
load_tail(
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write) {
|
||||
// Issue the epilogue waits
|
||||
/* This helps avoid early exit of blocks in Cluster
|
||||
* Waits for all stages to either be released (all
|
||||
* Consumer UNLOCKs), or if the stage was never used
|
||||
* then would just be acquired since the phase was
|
||||
* still inverted from make_producer_start_state
|
||||
*/
|
||||
pipeline.producer_tail(smem_pipe_write);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Consumer Perspective
|
||||
template <
|
||||
class FrgTensorC
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
mma(MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_read,
|
||||
FrgTensorC& accum,
|
||||
int k_tile_count,
|
||||
int thread_idx,
|
||||
TensorStorage& shared_tensors,
|
||||
Params const& mainloop_params)
|
||||
{
|
||||
using namespace cute;
|
||||
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
|
||||
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(InternalSmemLayoutAtomA{}) == 2, "InternalSmemLayoutAtomA must be rank 2.");
|
||||
static_assert(rank(InternalSmemLayoutAtomB{}) == 2, "InternalSmemLayoutAtomB must be rank 2.");
|
||||
static_assert(!cute::is_void_v<InternalSmemCopyAtomA>,
|
||||
"SM90 GMMA mainloops must specify a non-void copy atom for smem sourced instructions.");
|
||||
static_assert(cute::is_void_v<InternalSmemCopyAtomB>,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
// Obtain warp index
|
||||
int warp_idx = canonical_warp_idx_sync();
|
||||
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
|
||||
|
||||
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE)
|
||||
|
||||
// If TransposeB, GMMA will read from transposed B layout SMEM
|
||||
Tensor gmma_sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
//
|
||||
// Define C accumulators and A/B partitioning
|
||||
//
|
||||
|
||||
TiledMma tiled_mma;
|
||||
auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
|
||||
|
||||
// Allocate fragments and descriptors
|
||||
Tensor tCsA = thread_mma.partition_A(sA);
|
||||
Tensor tCrA = thread_mma.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCsB = thread_mma.partition_B(gmma_sB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
//
|
||||
// Copy Atom A retiling
|
||||
//
|
||||
|
||||
|
||||
auto smem_tiled_copy_A = make_tiled_copy_A(InternalSmemCopyAtomA{}, tiled_mma);
|
||||
|
||||
auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx);
|
||||
|
||||
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); // (CPY,CPY_M,CPY_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
|
||||
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
//
|
||||
static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX),
|
||||
"ERROR : Incorrect number of MMAs in flight");
|
||||
|
||||
// We release buffers to producer warps(dma load) with some mmas in flight
|
||||
PipelineState smem_pipe_release = smem_pipe_read;
|
||||
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
|
||||
TransposeOperandB transpose = cutlass::transform::collective::detail::make_transpose_operand_b(
|
||||
warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{},
|
||||
InternalSmemLayoutAtomB{}, InternalElementB{},
|
||||
cute::bool_constant<TransposeB>{});
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// first k tile
|
||||
{
|
||||
pipeline.consumer_wait(smem_pipe_read);
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
++smem_pipe_read;
|
||||
|
||||
bool skip_wait = (pipeline.consumer_try_wait(smem_pipe_read) == BarrierStatus::WaitDone);
|
||||
|
||||
// copy smem->rmem for A operand
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,read_stage), tCrA_copy_view(_,_,0));
|
||||
// transpose B operand in SMEM
|
||||
transpose(sB, gmma_sB, read_stage, 0);
|
||||
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
if (k_block == 0) {
|
||||
transpose(sB, gmma_sB, read_stage, 1);
|
||||
transpose.synchronize();
|
||||
}
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
}
|
||||
|
||||
warpgroup_wait<2>();
|
||||
|
||||
|
||||
if (k_tile_count - 1 > 0) {
|
||||
if (!skip_wait) {
|
||||
pipeline.consumer_wait(smem_pipe_read);
|
||||
}
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
}
|
||||
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// Mainloop GMMAs
|
||||
--k_tile_count;
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > 1; --k_tile_count) {
|
||||
|
||||
//
|
||||
// Compute on k_tile
|
||||
//
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
++smem_pipe_read;
|
||||
bool skip_wait = (pipeline.consumer_try_wait(smem_pipe_read) == BarrierStatus::WaitDone);
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
|
||||
if (k_block == size<2>(tCrA) - 1) {
|
||||
if (!skip_wait) {
|
||||
pipeline.consumer_wait(smem_pipe_read);
|
||||
}
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
// transpose B operand in SMEM
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
} else {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
// transpose B operand in SMEM
|
||||
if (k_block < 2) {
|
||||
transpose.synchronize(k_block); // make transpose of k_block available
|
||||
}
|
||||
if (k_block == 0) {
|
||||
transpose(sB, gmma_sB, read_stage, 1);
|
||||
}
|
||||
}
|
||||
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
if (k_block == 1) {
|
||||
// release prior barrier
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
if (k_tile_count > 0) {
|
||||
//
|
||||
// Compute on k_tile
|
||||
//
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
if (k_block < 2) {
|
||||
transpose.synchronize(k_block); // make k_block transpose available
|
||||
}
|
||||
if (k_block == 0) {
|
||||
transpose(sB, gmma_sB, read_stage, 1);
|
||||
}
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
if (k_block == 1) {
|
||||
// release prior barrier
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
warpgroup_fence_operand(accum);
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
}
|
||||
|
||||
/// Perform a Consumer Epilogue to release all buffers
|
||||
CUTLASS_DEVICE void
|
||||
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
|
||||
// Prologue GMMAs
|
||||
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
|
||||
k_tile_count -= prologue_mma_count;
|
||||
|
||||
smem_pipe_release.advance(k_tile_count);
|
||||
|
||||
// Wait on all GMMAs to complete
|
||||
warpgroup_wait<0>();
|
||||
|
||||
for (int count = 0; count < prologue_mma_count; ++count) {
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -1,609 +0,0 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cutlass/arch/reg_reconfig.h"
|
||||
|
||||
#include "cute/arch/copy_sm90.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
#include "cute/algorithm/gemm.hpp"
|
||||
|
||||
#include <cuda/std/type_traits>
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
int Stages,
|
||||
class ClusterShape,
|
||||
class TileShape_,
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
class ElementB_,
|
||||
class StrideB_,
|
||||
class TiledMma_,
|
||||
class GmemTiledCopyA_,
|
||||
class SmemLayoutAtomA_,
|
||||
class SmemCopyAtomA_,
|
||||
class TransformA_,
|
||||
class GmemTiledCopyB_,
|
||||
class SmemLayoutAtomB_,
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90CpAsyncGmmaUnpredicated<Stages, ClusterShape>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
ElementB_,
|
||||
StrideB_,
|
||||
TiledMma_,
|
||||
GmemTiledCopyA_,
|
||||
SmemLayoutAtomA_,
|
||||
SmemCopyAtomA_,
|
||||
TransformA_,
|
||||
GmemTiledCopyB_,
|
||||
SmemLayoutAtomB_,
|
||||
SmemCopyAtomB_,
|
||||
TransformB_>
|
||||
{
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90CpAsyncGmmaUnpredicated<Stages, ClusterShape>;
|
||||
using TileShape = TileShape_;
|
||||
using ElementA = ElementA_;
|
||||
using StrideA = StrideA_;
|
||||
using ElementB = ElementB_;
|
||||
using StrideB = StrideB_;
|
||||
using TiledMma = TiledMma_;
|
||||
using ElementAccumulator = typename TiledMma::ValTypeC;
|
||||
using GmemTiledCopyA = GmemTiledCopyA_;
|
||||
using GmemTiledCopyB = GmemTiledCopyB_;
|
||||
using SmemLayoutAtomA = SmemLayoutAtomA_;
|
||||
using SmemLayoutAtomB = SmemLayoutAtomB_;
|
||||
using SmemCopyAtomA = SmemCopyAtomA_;
|
||||
using SmemCopyAtomB = SmemCopyAtomB_;
|
||||
using TransformA = TransformA_;
|
||||
using TransformB = TransformB_;
|
||||
using ArchTag = typename DispatchPolicy::ArchTag;
|
||||
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
using SmemLayoutA = decltype(tile_to_shape(
|
||||
SmemLayoutAtomA{},
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
using SmemLayoutB = decltype(tile_to_shape(
|
||||
SmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
|
||||
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
|
||||
static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
|
||||
cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
|
||||
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
|
||||
|
||||
struct SharedStorage
|
||||
{
|
||||
cute::array_aligned<ElementA, cute::cosize_v<SmemLayoutA>> smem_a;
|
||||
cute::array_aligned<ElementB, cute::cosize_v<SmemLayoutB>> smem_b;
|
||||
};
|
||||
|
||||
struct Arguments {
|
||||
ElementA const* ptr_A;
|
||||
StrideA dA;
|
||||
ElementB const* ptr_B;
|
||||
StrideB dB;
|
||||
};
|
||||
|
||||
using Params = Arguments;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CollectiveMma() = default;
|
||||
|
||||
template <class ProblemShape>
|
||||
static constexpr Params
|
||||
to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) {
|
||||
(void) workspace;
|
||||
return args;
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
template <
|
||||
class TensorA,
|
||||
class TensorB,
|
||||
class FrgTensorC,
|
||||
class KTileIterator,
|
||||
class ResidueMNK
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
operator() (
|
||||
TensorA gA,
|
||||
TensorB gB,
|
||||
FrgTensorC& accum,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
ResidueMNK residue_mnk,
|
||||
int thread_idx,
|
||||
char *smem_buf,
|
||||
Params const& mainloop_params)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
(void) residue_mnk;
|
||||
|
||||
static_assert(is_gmem<TensorA>::value, "A tensor must be gmem resident.");
|
||||
static_assert(is_gmem<TensorB>::value, "B tensor must be gmem resident.");
|
||||
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2.");
|
||||
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2.");
|
||||
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(cute::is_same<TransformA, cute::identity>::value,
|
||||
"SM90 warpgroup MMA must specify transforms through MMA_Atom.");
|
||||
static_assert(cute::is_same<TransformB, cute::identity>::value,
|
||||
"SM90 warpgroup MMA must specify transforms through MMA_Atom.");
|
||||
static_assert(cute::is_same<SmemCopyAtomA, void>::value,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
static_assert(cute::is_same<SmemCopyAtomA, void>::value,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
GmemTiledCopyA gmem_tiled_copy_a;
|
||||
GmemTiledCopyB gmem_tiled_copy_b;
|
||||
auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx);
|
||||
auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx);
|
||||
|
||||
Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
|
||||
Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
|
||||
Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k)
|
||||
Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE)
|
||||
|
||||
// Tile MMA atom and compute thread partitions across A, B and C
|
||||
TiledMma tiled_mma;
|
||||
auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
|
||||
|
||||
// Allocate registers for pipelining
|
||||
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_N,PIPE)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
|
||||
|
||||
//
|
||||
// Prologue
|
||||
//
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_pipe = 0; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) {
|
||||
copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe));
|
||||
copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe));
|
||||
cp_async_fence();
|
||||
++k_tile_iter;
|
||||
--k_tile_count;
|
||||
}
|
||||
|
||||
// Current pipe index in smem to read from
|
||||
int smem_pipe_read = 0;
|
||||
// Current pipe index in smem to write to
|
||||
int smem_pipe_write = DispatchPolicy::Stages-1;
|
||||
|
||||
//
|
||||
// Pipelined Main Loop
|
||||
//
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > -(DispatchPolicy::Stages-1); --k_tile_count)
|
||||
{
|
||||
// Copy gmem to smem before computing gemm on each k-pipe
|
||||
// pipe index in smem where the next gmem tile will be read into
|
||||
copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write));
|
||||
copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write));
|
||||
cp_async_fence();
|
||||
if (k_tile_count > 0) { ++k_tile_iter; }
|
||||
|
||||
//
|
||||
// Compute on k_tile
|
||||
//
|
||||
warpgroup_fence_operand(accum);
|
||||
warpgroup_arrive();
|
||||
|
||||
cp_async_wait<DispatchPolicy::Stages-2>();
|
||||
cute::gemm(tiled_mma, tCrA(_,_,_,smem_pipe_read), tCrB(_,_,_,smem_pipe_read), accum);
|
||||
warpgroup_commit_batch();
|
||||
|
||||
//
|
||||
// Advance the pipe
|
||||
//
|
||||
++smem_pipe_read;
|
||||
smem_pipe_read = (smem_pipe_read == DispatchPolicy::Stages) ? smem_pipe_read = 0 : smem_pipe_read;
|
||||
|
||||
++smem_pipe_write;
|
||||
smem_pipe_write = (smem_pipe_write == DispatchPolicy::Stages) ? smem_pipe_write = 0 : smem_pipe_write;
|
||||
|
||||
// Wait for the pipeline MMAs to drain
|
||||
warpgroup_wait<0>();
|
||||
warpgroup_fence_operand(accum);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
int Stages,
|
||||
class ClusterShape,
|
||||
class TileShape_,
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
class ElementB_,
|
||||
class StrideB_,
|
||||
class TiledMma_,
|
||||
class GmemTiledCopyA_,
|
||||
class SmemLayoutAtomA_,
|
||||
class SmemCopyAtomA_,
|
||||
class TransformA_,
|
||||
class GmemTiledCopyB_,
|
||||
class SmemLayoutAtomB_,
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90CpAsyncGmma<Stages, ClusterShape>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
ElementB_,
|
||||
StrideB_,
|
||||
TiledMma_,
|
||||
GmemTiledCopyA_,
|
||||
SmemLayoutAtomA_,
|
||||
SmemCopyAtomA_,
|
||||
TransformA_,
|
||||
GmemTiledCopyB_,
|
||||
SmemLayoutAtomB_,
|
||||
SmemCopyAtomB_,
|
||||
TransformB_>
|
||||
{
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90CpAsyncGmma<Stages, ClusterShape>;
|
||||
using TileShape = TileShape_;
|
||||
using ElementA = ElementA_;
|
||||
using StrideA = StrideA_;
|
||||
using ElementB = ElementB_;
|
||||
using StrideB = StrideB_;
|
||||
using TiledMma = TiledMma_;
|
||||
using ElementAccumulator = typename TiledMma::ValTypeC;
|
||||
using GmemTiledCopyA = GmemTiledCopyA_;
|
||||
using GmemTiledCopyB = GmemTiledCopyB_;
|
||||
using SmemLayoutAtomA = SmemLayoutAtomA_;
|
||||
using SmemLayoutAtomB = SmemLayoutAtomB_;
|
||||
using SmemCopyAtomA = SmemCopyAtomA_;
|
||||
using SmemCopyAtomB = SmemCopyAtomB_;
|
||||
using TransformA = TransformA_;
|
||||
using TransformB = TransformB_;
|
||||
using ArchTag = typename DispatchPolicy::ArchTag;
|
||||
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
using SmemLayoutA = decltype(tile_to_shape(
|
||||
SmemLayoutAtomA{},
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
using SmemLayoutB = decltype(tile_to_shape(
|
||||
SmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
|
||||
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
|
||||
static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
|
||||
cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
|
||||
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
|
||||
|
||||
struct SharedStorage
|
||||
{
|
||||
cute::array_aligned<ElementA, cute::cosize_v<SmemLayoutA>> smem_a;
|
||||
cute::array_aligned<ElementB, cute::cosize_v<SmemLayoutB>> smem_b;
|
||||
};
|
||||
|
||||
struct Arguments {
|
||||
ElementA const* ptr_A;
|
||||
StrideA dA;
|
||||
ElementB const* ptr_B;
|
||||
StrideB dB;
|
||||
};
|
||||
|
||||
using Params = Arguments;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
CollectiveMma() = default;
|
||||
|
||||
template <class ProblemShape>
|
||||
static constexpr Params
|
||||
to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) {
|
||||
(void) workspace;
|
||||
return args;
|
||||
}
|
||||
|
||||
template<class ProblemShape>
|
||||
CUTLASS_HOST_DEVICE static bool
|
||||
can_implement(
|
||||
[[maybe_unused]] ProblemShape const& problem_shape,
|
||||
[[maybe_unused]] Arguments const& args) {
|
||||
return true;
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
template <
|
||||
class FrgTensorD,
|
||||
class TensorA,
|
||||
class TensorB,
|
||||
class FrgTensorC,
|
||||
class KTileIterator,
|
||||
class ResidueMNK
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
operator() (
|
||||
FrgTensorD &accum,
|
||||
TensorA gA_in,
|
||||
TensorB gB_in,
|
||||
FrgTensorC const &src_accum,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
ResidueMNK residue_mnk,
|
||||
int thread_idx,
|
||||
char *smem_buf)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
static_assert(is_rmem<FrgTensorD>::value, "D tensor must be rmem resident.");
|
||||
static_assert(is_gmem<TensorA>::value, "A tensor must be gmem resident.");
|
||||
static_assert(is_gmem<TensorB>::value, "B tensor must be gmem resident.");
|
||||
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2.");
|
||||
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2.");
|
||||
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(cute::is_same<TransformA, cute::identity>::value,
|
||||
"SM90 warpgroup MMA must specify transforms through MMA_Atom.");
|
||||
static_assert(cute::is_same<TransformB, cute::identity>::value,
|
||||
"SM90 warpgroup MMA must specify transforms through MMA_Atom.");
|
||||
static_assert(cute::is_same<SmemCopyAtomA, void>::value,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
static_assert(cute::is_same<SmemCopyAtomA, void>::value,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
// Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k)
|
||||
// This aligns the tensor with BLK_K for all but the 0th k_tile
|
||||
Tensor gA = domain_offset(make_coord(0, get<2>(residue_mnk), 0), gA_in);
|
||||
Tensor gB = domain_offset(make_coord(0, get<2>(residue_mnk), 0), gB_in);
|
||||
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
GmemTiledCopyA gmem_tiled_copy_a;
|
||||
GmemTiledCopyB gmem_tiled_copy_b;
|
||||
auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx);
|
||||
auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx);
|
||||
|
||||
Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
|
||||
Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
|
||||
Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k)
|
||||
Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE)
|
||||
|
||||
//
|
||||
// PREDICATES
|
||||
//
|
||||
|
||||
// Allocate predicate tensors for m and n
|
||||
Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});
|
||||
Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
|
||||
|
||||
// Construct identity layout for sA and sB
|
||||
Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k)
|
||||
Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k)
|
||||
|
||||
// Repeat the partitioning with identity layouts
|
||||
Tensor tAcA = gmem_thr_copy_a.partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k)
|
||||
Tensor tBcB = gmem_thr_copy_b.partition_S(cB); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k)
|
||||
|
||||
// Set predicates for m bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < size<0>(tApA); ++m) {
|
||||
tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk); // blk_m coord < residue_m
|
||||
}
|
||||
// Set predicates for n bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < size<0>(tBpB); ++n) {
|
||||
tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk); // blk_n coord < residue_n
|
||||
}
|
||||
|
||||
//
|
||||
// Prologue/PREFETCH
|
||||
//
|
||||
|
||||
// Clear the smem tiles to account for predicated off loads
|
||||
clear(tAsA);
|
||||
clear(tBsB);
|
||||
|
||||
// Start async loads for 0th k-tile, where we take care of the k residue
|
||||
{
|
||||
constexpr int k_pipe = 0;
|
||||
|
||||
Tensor tAgAk = tAgA(_,_,_,*k_tile_iter);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tAsA); ++k) {
|
||||
if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gA shifted)
|
||||
copy_if(gmem_tiled_copy_a, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,k_pipe));
|
||||
}
|
||||
}
|
||||
Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tBsB); ++k) {
|
||||
if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gB shifted)
|
||||
copy_if(gmem_tiled_copy_b, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,k_pipe));
|
||||
}
|
||||
}
|
||||
cp_async_fence();
|
||||
++k_tile_iter;
|
||||
--k_tile_count;
|
||||
}
|
||||
|
||||
// Start async loads for 1st k-tile onwards, no k-residue handling needed
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_pipe = 1; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) {
|
||||
if (k_tile_count <= 0) {
|
||||
clear(tApA);
|
||||
clear(tBpB);
|
||||
}
|
||||
copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe)); // CpAsync
|
||||
copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe)); // CpAsync
|
||||
cp_async_fence();
|
||||
++k_tile_iter;
|
||||
--k_tile_count;
|
||||
}
|
||||
|
||||
//
|
||||
// MMA Atom partitioning
|
||||
//
|
||||
|
||||
// Tile MMA atom and compute thread partitions across A, B and C
|
||||
TiledMma tiled_mma;
|
||||
auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
|
||||
|
||||
// Allocate registers for pipelining
|
||||
Tensor tCsA = thr_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCsB = thr_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_N,PIPE)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(src_accum)); // M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(src_accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
|
||||
|
||||
// Current pipe index in smem to read from
|
||||
int smem_pipe_read = 0;
|
||||
// Current pipe index in smem to write to
|
||||
int smem_pipe_write = DispatchPolicy::Stages-1;
|
||||
|
||||
//
|
||||
// Pipelined Main Loop
|
||||
//
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > -(DispatchPolicy::Stages-1); --k_tile_count)
|
||||
{
|
||||
//
|
||||
// Copy gmem to smem for *k_tile_iter
|
||||
//
|
||||
if (k_tile_count <= 0) {
|
||||
clear(tApA);
|
||||
clear(tBpB);
|
||||
}
|
||||
copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write)); // CpAsync
|
||||
copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write)); // CpAsync
|
||||
cp_async_fence();
|
||||
++k_tile_iter;
|
||||
|
||||
//
|
||||
// Compute on k_tile
|
||||
//
|
||||
warpgroup_fence_operand(accum);
|
||||
warpgroup_arrive();
|
||||
|
||||
cp_async_wait<DispatchPolicy::Stages-2>();
|
||||
cute::gemm(tiled_mma, accum, tCrA(_,_,_,smem_pipe_read), tCrB(_,_,_,smem_pipe_read), src_accum);
|
||||
warpgroup_commit_batch();
|
||||
|
||||
//
|
||||
// Advance the pipe
|
||||
//
|
||||
++smem_pipe_read;
|
||||
smem_pipe_read = (smem_pipe_read == DispatchPolicy::Stages) ? smem_pipe_read = 0 : smem_pipe_read;
|
||||
|
||||
++smem_pipe_write;
|
||||
smem_pipe_write = (smem_pipe_write == DispatchPolicy::Stages) ? smem_pipe_write = 0 : smem_pipe_write;
|
||||
|
||||
// Wait for the pipeline MMAs to drain
|
||||
warpgroup_wait<0>();
|
||||
warpgroup_fence_operand(accum);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -0,0 +1,483 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cute/arch/copy_sm90.hpp"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
|
||||
#include "cute/algorithm/functional.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
#include "cute/algorithm/gemm.hpp"
|
||||
#include "cute/tensor_predicate.hpp"
|
||||
#include "cute/numeric/arithmetic_tuple.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cutlass/trace.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// WarpSpecialized Mainloop
|
||||
template <
|
||||
int Stages,
|
||||
class ClusterShape_,
|
||||
class TileShape_,
|
||||
class KernelSchedule,
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
class ElementB_,
|
||||
class StrideB_,
|
||||
class TiledMma_,
|
||||
class GmemTiledCopyA_,
|
||||
class SmemLayoutAtomA_,
|
||||
class SmemCopyAtomA_,
|
||||
class TransformA_,
|
||||
class GmemTiledCopyB_,
|
||||
class SmemLayoutAtomB_,
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90CpAsyncGmmaWarpSpecialized<Stages,ClusterShape_,KernelSchedule>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
ElementB_,
|
||||
StrideB_,
|
||||
TiledMma_,
|
||||
GmemTiledCopyA_,
|
||||
SmemLayoutAtomA_,
|
||||
SmemCopyAtomA_,
|
||||
TransformA_,
|
||||
GmemTiledCopyB_,
|
||||
SmemLayoutAtomB_,
|
||||
SmemCopyAtomB_,
|
||||
TransformB_>
|
||||
{
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90CpAsyncGmmaWarpSpecialized<Stages,ClusterShape_,KernelSchedule>;
|
||||
using TileShape = TileShape_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ElementA = ElementA_;
|
||||
using StrideA = StrideA_;
|
||||
using ElementB = ElementB_;
|
||||
using StrideB = StrideB_;
|
||||
using TiledMma = TiledMma_;
|
||||
using ElementAccumulator = typename TiledMma::ValTypeC;
|
||||
using GmemTiledCopyA = GmemTiledCopyA_;
|
||||
using GmemTiledCopyB = GmemTiledCopyB_;
|
||||
using SmemLayoutAtomA = SmemLayoutAtomA_;
|
||||
using SmemLayoutAtomB = SmemLayoutAtomB_;
|
||||
using SmemCopyAtomA = SmemCopyAtomA_;
|
||||
using SmemCopyAtomB = SmemCopyAtomB_;
|
||||
using TransformA = TransformA_;
|
||||
using TransformB = TransformB_;
|
||||
using ArchTag = typename DispatchPolicy::ArchTag;
|
||||
|
||||
using MainloopPipeline = cutlass::PipelineAsync<DispatchPolicy::Stages>;
|
||||
using PipelineState = typename MainloopPipeline::PipelineState;
|
||||
using PipelineParams = typename MainloopPipeline::Params;
|
||||
|
||||
static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
using SmemLayoutA = decltype(tile_to_shape(
|
||||
SmemLayoutAtomA{},
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
using SmemLayoutB = decltype(tile_to_shape(
|
||||
SmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
|
||||
|
||||
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more.");
|
||||
static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
|
||||
cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
|
||||
"MMA atom must source both A and B operand from smem_desc for this mainloop.");
|
||||
|
||||
struct SharedStorage
|
||||
{
|
||||
struct TensorStorage : cute::aligned_struct<128> {
|
||||
cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>> smem_A;
|
||||
cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
|
||||
} tensors;
|
||||
|
||||
using PipelineStorage = typename MainloopPipeline::SharedStorage;
|
||||
PipelineStorage pipeline;
|
||||
};
|
||||
using TensorStorage = typename SharedStorage::TensorStorage;
|
||||
using PipelineStorage = typename SharedStorage::PipelineStorage;
|
||||
|
||||
// Host side kernel arguments
|
||||
struct Arguments {
|
||||
ElementA const* ptr_A = nullptr;
|
||||
StrideA dA{};
|
||||
ElementB const* ptr_B = nullptr;
|
||||
StrideB dB{};
|
||||
uint32_t mma_promotion_interval = 4;
|
||||
};
|
||||
|
||||
// Device side kernel params
|
||||
using Params = Arguments;
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
template <class ProblemShape>
|
||||
static constexpr Params
|
||||
to_underlying_arguments(
|
||||
[[maybe_unused]] ProblemShape const& problem_shape,
|
||||
Arguments const& args,
|
||||
[[maybe_unused]] void* workspace) {
|
||||
return args;
|
||||
}
|
||||
|
||||
template<class ProblemShape>
|
||||
CUTLASS_HOST_DEVICE static bool
|
||||
can_implement(
|
||||
ProblemShape const& problem_shape,
|
||||
[[maybe_unused]] Arguments const& args) {
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
bool implementable = true;
|
||||
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyA::NumValSrc>(cute::make_shape(M,K,L), StrideA{});
|
||||
implementable = implementable && cutlass::detail::check_alignment<GmemTiledCopyB::NumValSrc>(cute::make_shape(N,K,L), StrideB{});
|
||||
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
|
||||
}
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
|
||||
static constexpr int K_PIPE_MMAS = 1;
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class TensorA,
|
||||
class TensorB,
|
||||
class KTileIterator,
|
||||
class ResidueMNK
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
load(
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write,
|
||||
TensorA const& gA_in,
|
||||
TensorB const& gB_in,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
ResidueMNK residue_mnk,
|
||||
int thread_idx,
|
||||
TensorStorage& shared_tensors)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
static_assert(is_gmem<TensorA>::value, "A tensor must be gmem resident.");
|
||||
static_assert(is_gmem<TensorB>::value, "B tensor must be gmem resident.");
|
||||
|
||||
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
// Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k)
|
||||
// This aligns the tensor with BLK_K for all but the 0th k_tile
|
||||
Tensor gA = domain_offset(make_coord(0, get<2>(residue_mnk), 0), gA_in);
|
||||
Tensor gB = domain_offset(make_coord(0, get<2>(residue_mnk), 0), gB_in);
|
||||
|
||||
// Partition the copying of A and B tiles across the threads
|
||||
GmemTiledCopyA gmem_tiled_copy_a;
|
||||
GmemTiledCopyB gmem_tiled_copy_b;
|
||||
auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx);
|
||||
auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx);
|
||||
|
||||
Tensor tAgA = gmem_thr_copy_a.partition_S(gA); // (ACPY,ACPY_M,ACPY_K,k)
|
||||
Tensor tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K,PIPE)
|
||||
Tensor tBgB = gmem_thr_copy_b.partition_S(gB); // (BCPY,BCPY_N,BCPY_K,k)
|
||||
Tensor tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K,PIPE)
|
||||
|
||||
// Allocate predicate tensors for m and n
|
||||
Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});
|
||||
Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
|
||||
|
||||
// Construct identity layout for sA and sB
|
||||
Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); // (BLK_M,BLK_K) -> (blk_m,blk_k)
|
||||
Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); // (BLK_N,BLK_K) -> (blk_n,blk_k)
|
||||
|
||||
// Repeat the partitioning with identity layouts
|
||||
Tensor tAcA = gmem_thr_copy_a.partition_S(cA); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k)
|
||||
Tensor tBcB = gmem_thr_copy_b.partition_S(cB); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k)
|
||||
|
||||
// Set predicates for m bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < size<0>(tApA); ++m) {
|
||||
tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk); // blk_m coord < residue_m
|
||||
}
|
||||
// Set predicates for n bounds
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < size<0>(tBpB); ++n) {
|
||||
tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk); // blk_n coord < residue_n
|
||||
}
|
||||
|
||||
// 0-th stage with predication on k to account for residue
|
||||
{
|
||||
// LOCK smem_pipe_write for _writing_
|
||||
pipeline.producer_acquire(smem_pipe_write);
|
||||
int write_stage = smem_pipe_write.index();
|
||||
|
||||
// Copy gmem to smem for *k_tile_iter, predicating for k residue
|
||||
Tensor tAgAk = tAgA(_,_,_,*k_tile_iter);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tAsA); ++k) {
|
||||
if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gA shifted)
|
||||
copy_if(gmem_tiled_copy_a, tApA(_,k), tAgAk(_,_,k), tAsA(_,_,k,write_stage));
|
||||
}
|
||||
else {
|
||||
clear(tAsA(_,_,k,write_stage));
|
||||
}
|
||||
}
|
||||
Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k = 0; k < size<2>(tBsB); ++k) {
|
||||
if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) { // blk_k coord < residue_k (gB shifted)
|
||||
copy_if(gmem_tiled_copy_b, tBpB(_,k), tBgBk(_,_,k), tBsB(_,_,k,write_stage));
|
||||
}
|
||||
else {
|
||||
clear(tBsB(_,_,k,write_stage));
|
||||
}
|
||||
}
|
||||
++k_tile_iter;
|
||||
--k_tile_count;
|
||||
|
||||
// UNLOCK smem_pipe_write
|
||||
pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive);
|
||||
|
||||
// Advance smem_pipe_write
|
||||
++smem_pipe_write;
|
||||
}
|
||||
|
||||
// Mainloop
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > 0; --k_tile_count) {
|
||||
// LOCK smem_pipe_write for _writing_
|
||||
pipeline.producer_acquire(smem_pipe_write);
|
||||
int write_stage = smem_pipe_write.index();
|
||||
|
||||
// Copy gmem to smem for *k_tile_iter
|
||||
copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
++k_tile_iter;
|
||||
|
||||
// UNLOCK smem_pipe_write
|
||||
pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive);
|
||||
|
||||
// Advance smem_pipe_write
|
||||
++smem_pipe_write;
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
|
||||
CUTLASS_DEVICE void
|
||||
load_tail(
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write) {
|
||||
// Issue the epilogue waits
|
||||
/* This helps avoid early exit of blocks in Cluster
|
||||
* Waits for all stages to either be released (all
|
||||
* Consumer UNLOCKs), or if the stage was never used
|
||||
* then would just be acquired since the phase was
|
||||
* still inverted from make_producer_start_state
|
||||
*/
|
||||
pipeline.producer_tail(smem_pipe_write);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Consumer Perspective
|
||||
template <
|
||||
class FrgTensorC
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
mma(MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_read,
|
||||
FrgTensorC& accum,
|
||||
int k_tile_count,
|
||||
int thread_idx,
|
||||
TensorStorage& shared_tensors,
|
||||
Params const& mainloop_params)
|
||||
{
|
||||
using namespace cute;
|
||||
|
||||
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
|
||||
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(cute::is_void_v<SmemCopyAtomA>,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
static_assert(cute::is_void_v<SmemCopyAtomB>,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
//
|
||||
// Define C accumulators and A/B partitioning
|
||||
//
|
||||
|
||||
TiledMma tiled_mma;
|
||||
auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
|
||||
|
||||
Tensor tCsA = thread_mma.partition_A(sA); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCsB = thread_mma.partition_B(sB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
// Allocate "fragments/descriptors"
|
||||
Tensor tCrA = thread_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum)); // M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
|
||||
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
//
|
||||
static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX),
|
||||
"ERROR : Incorrect number of MMAs in flight");
|
||||
|
||||
// We release buffers to producer warps(dma load) with some mmas in flight
|
||||
PipelineState smem_pipe_release = smem_pipe_read;
|
||||
|
||||
// Prologue GMMAs
|
||||
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
|
||||
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue) {
|
||||
|
||||
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
|
||||
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
warpgroup_arrive();
|
||||
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
|
||||
// (V,M,K) x (V,N,K) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
}
|
||||
|
||||
warpgroup_commit_batch();
|
||||
|
||||
++smem_pipe_read;
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
// Mainloop GMMAs
|
||||
k_tile_count -= prologue_mma_count;
|
||||
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > 0; --k_tile_count) {
|
||||
|
||||
// WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
|
||||
auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
warpgroup_arrive();
|
||||
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
|
||||
// (V,M,K) x (V,N,K) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
}
|
||||
warpgroup_commit_batch();
|
||||
|
||||
/// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
|
||||
warpgroup_wait<K_PIPE_MMAS>();
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
// UNLOCK smem_pipe_release, done _computing_ on it
|
||||
pipeline.consumer_release(smem_pipe_release);
|
||||
|
||||
// Advance smem_pipe_read and smem_pipe_release
|
||||
++smem_pipe_read;
|
||||
++smem_pipe_release;
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
}
|
||||
|
||||
/// Perform a Consumer Epilogue to release all buffers
|
||||
CUTLASS_DEVICE void
|
||||
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
|
||||
// Prologue GMMAs
|
||||
int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
|
||||
k_tile_count -= prologue_mma_count;
|
||||
|
||||
smem_pipe_release.advance(k_tile_count);
|
||||
|
||||
// Wait on all GMMAs to complete
|
||||
warpgroup_wait<0>();
|
||||
|
||||
for (int count = 0; count < prologue_mma_count; ++count) {
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -59,7 +59,6 @@ template <
|
||||
int Stages,
|
||||
class ClusterShape,
|
||||
class KernelSchedule,
|
||||
int PipelineAsyncMmaStages,
|
||||
class TileShape_,
|
||||
class ElementA_,
|
||||
class StrideA_,
|
||||
@@ -75,7 +74,7 @@ template <
|
||||
class SmemCopyAtomB_,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90TmaGmmaRmemAWarpSpecialized<Stages, ClusterShape, KernelSchedule, PipelineAsyncMmaStages>,
|
||||
MainloopSm90TmaGmmaRmemAWarpSpecialized<Stages, ClusterShape, KernelSchedule>,
|
||||
TileShape_,
|
||||
ElementA_,
|
||||
StrideA_,
|
||||
@@ -94,7 +93,7 @@ struct CollectiveMma<
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecialized<Stages, ClusterShape, KernelSchedule, PipelineAsyncMmaStages>;
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecialized<Stages, ClusterShape, KernelSchedule>;
|
||||
using TileShape = TileShape_;
|
||||
using ElementA = ElementA_;
|
||||
using StrideA = StrideA_;
|
||||
@@ -190,9 +189,19 @@ struct CollectiveMma<
|
||||
static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time.");
|
||||
static_assert(TransposeB xor (cute::is_same_v<SmemLayoutB, GmmaSmemLayoutB>),
|
||||
"Should be same layout if not TransposeB.");
|
||||
static_assert(!TransposeB || size<1>(SmemLayoutB{}) * sizeof(InternalElementB) == 128,
|
||||
static_assert(!TransposeB || (((size<1>(SmemLayoutB{}) * sizeof_bits<InternalElementB>::value)) / 8) == 128,
|
||||
"SmemLayoutB K must be 128bytes to be transposed.");
|
||||
static_assert(!transform::collective::detail::use_universal_transposition<InternalSmemLayoutAtomB, InternalElementB>(),
|
||||
|
||||
static constexpr bool uses_universal_transposition() {
|
||||
if constexpr (TransposeB) {
|
||||
return transform::collective::detail::use_universal_transposition<InternalSmemLayoutAtomB, InternalElementB>();
|
||||
}
|
||||
else {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
static_assert(!uses_universal_transposition(),
|
||||
"Warp specialized ARF kernels have not supported universal B transposition yet.");
|
||||
|
||||
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
|
||||
@@ -216,10 +225,10 @@ struct CollectiveMma<
|
||||
|
||||
// Host side kernel arguments
|
||||
struct Arguments {
|
||||
ElementA const* ptr_A;
|
||||
StrideA dA;
|
||||
ElementB const* ptr_B;
|
||||
StrideB dB;
|
||||
ElementA const* ptr_A = nullptr;
|
||||
StrideA dA{};
|
||||
ElementB const* ptr_B = nullptr;
|
||||
StrideB dB{};
|
||||
uint32_t mma_promotion_interval = 4;
|
||||
};
|
||||
|
||||
@@ -321,11 +330,9 @@ struct CollectiveMma<
|
||||
}
|
||||
|
||||
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
|
||||
static constexpr int K_PIPE_MMAS = DispatchPolicy::PipelineAsyncMmaStages;
|
||||
static_assert(K_PIPE_MMAS == 0, "no MMA stage should be asynchronous for this mainloop for now.");
|
||||
static constexpr uint32_t TmaTransactionBytes =
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof(InternalElementA)))+
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof(InternalElementB)));
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<InternalElementA>::value)) / 8 +
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<InternalElementB>::value)) / 8 ;
|
||||
|
||||
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
|
||||
CUTLASS_DEVICE
|
||||
@@ -335,19 +342,45 @@ struct CollectiveMma<
|
||||
cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
|
||||
}
|
||||
|
||||
/// Set up the data needed by this collective for load and mma.
|
||||
/// Returns a tuple of tensors. The collective and the kernel layer have the contract
|
||||
/// that the tuple must contain at least two elements, with the first two elements being:
|
||||
/// gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l)
|
||||
/// gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l)
|
||||
/// The rest of the tensors can be specified as needed by this collective.
|
||||
template <class ProblemShape_MNKL,
|
||||
class TileShapeMNK>
|
||||
CUTLASS_DEVICE auto
|
||||
tile_input_tensors(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params, TileShapeMNK const& tileshape_mnk) {
|
||||
using X = Underscore;
|
||||
// Separate out problem shape for convenience
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, tileshape_mnk, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, tileshape_mnk, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
return cute::make_tuple(gA_mkl, gB_nkl);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class TensorA, class TMA_LOAD_A,
|
||||
class TensorB, class TMA_LOAD_B,
|
||||
class KTileIterator
|
||||
class TensorA, class TensorB,
|
||||
class KTileIterator, class BlockCoord
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
load(
|
||||
Params const& mainloop_params,
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write,
|
||||
TensorA const& gA, TMA_LOAD_A& tma_load_a,
|
||||
TensorB const& gB, TMA_LOAD_B& tma_load_b,
|
||||
cute::tuple<TensorA, TensorB> const& tiled_tensors,
|
||||
BlockCoord const& blk_coord,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
int thread_idx,
|
||||
uint32_t block_rank_in_cluster,
|
||||
@@ -372,8 +405,16 @@ struct CollectiveMma<
|
||||
constexpr uint32_t cluster_shape_x = get<0>(DispatchPolicy::ClusterShape());
|
||||
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
|
||||
|
||||
auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
|
||||
// Partition the inputs based on the current block coordinates.
|
||||
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Applies the mapping from block_tma_a
|
||||
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
|
||||
@@ -415,8 +456,8 @@ struct CollectiveMma<
|
||||
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
|
||||
|
||||
int write_stage = smem_pipe_write.index();
|
||||
copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
++k_tile_iter;
|
||||
|
||||
// Advance smem_pipe_write
|
||||
@@ -508,9 +549,12 @@ struct CollectiveMma<
|
||||
auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(thread_idx);
|
||||
|
||||
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA); // (CPY,CPY_M,CPY_K)
|
||||
Tensor tCsA_copy_view = smem_thr_copy_A.partition_S(sA); // (CPY,CPY_M,CPY_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA_copy_view) == size<1>(tCrA_copy_view)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA_copy_view) == size<2>(tCrA_copy_view)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
|
||||
@@ -521,8 +565,6 @@ struct CollectiveMma<
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
//
|
||||
static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS < K_PIPE_MAX),
|
||||
"ERROR : Incorrect number of MMAs in flight");
|
||||
|
||||
// We release buffers to producer warps(dma load) with some mmas in flight
|
||||
PipelineState smem_pipe_release = smem_pipe_read;
|
||||
@@ -548,43 +590,41 @@ struct CollectiveMma<
|
||||
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
|
||||
// copy smem->rmem for A operand
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,read_stage), tCrA_copy_view(_,_,0));
|
||||
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,0,read_stage), tCrA_copy_view(_,_,0));
|
||||
// transpose B operand in SMEM
|
||||
transpose(sB, gmma_sB, read_stage, 0);
|
||||
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
transpose.synchronize(k_block);
|
||||
transpose(sB, gmma_sB, read_stage, k_block + 1);
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
if(k_block == 0) {
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
}
|
||||
warpgroup_commit_batch();
|
||||
}
|
||||
|
||||
warpgroup_wait<2>();
|
||||
|
||||
--k_tile_count;
|
||||
if (k_tile_count > 0) {
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
}
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
--k_tile_count;
|
||||
if(k_tile_count == 0) {
|
||||
return;
|
||||
}
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
warpgroup_wait<2>();
|
||||
}
|
||||
|
||||
if (k_tile_count == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// Mainloop GMMAs
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
@@ -606,12 +646,12 @@ struct CollectiveMma<
|
||||
}
|
||||
if (k_block == size<2>(tCrA) - 1) {
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
// transpose B operand in SMEM
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
}
|
||||
else {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
// transpose B operand in SMEM
|
||||
transpose.synchronize(k_block); // make transpose of k_block available
|
||||
transpose(sB, gmma_sB, read_stage, k_block + 1);
|
||||
@@ -620,7 +660,6 @@ struct CollectiveMma<
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
if (k_block == 1) {
|
||||
@@ -647,8 +686,7 @@ struct CollectiveMma<
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA) - 1; ++k_block) {
|
||||
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
copy(smem_tiled_copy_A, tCsA_copy_view(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
transpose.synchronize(k_block); // make k_block transpose available
|
||||
transpose(sB, gmma_sB, read_stage, k_block + 1);
|
||||
warpgroup_arrive();
|
||||
@@ -667,7 +705,6 @@ struct CollectiveMma<
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA(_,_,size<2>(tCrA) - 1), tCrB(_,_,size<2>(tCrA) - 1,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,830 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_conversion.h"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cute/arch/copy_sm90.hpp"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/detail/dependent_false.hpp"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/detail/layout.hpp"
|
||||
|
||||
#include "cute/algorithm/functional.hpp"
|
||||
#include "cute/atom/mma_atom.hpp"
|
||||
#include "cute/atom/copy_traits_sm90_tma.hpp"
|
||||
#include "cute/algorithm/gemm.hpp"
|
||||
#include "cute/tensor_predicate.hpp"
|
||||
#include "cute/numeric/arithmetic_tuple.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cutlass/transform/collective/sm90_wgmma_transpose.hpp"
|
||||
#include "cutlass/trace.h"
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::collective {
|
||||
using namespace cute;
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// WarpSpecialized Mainloop that source A operand from registers
|
||||
template <
|
||||
int Stages,
|
||||
class ClusterShape,
|
||||
class KernelSchedule,
|
||||
class TileShape_,
|
||||
class ElementAOptionalTuple,
|
||||
class StrideAOptionalTuple,
|
||||
class ElementBOptionalTuple,
|
||||
class StrideBOptionalTuple,
|
||||
class TiledMma_,
|
||||
class GmemTiledCopyAOptionalTuple,
|
||||
class SmemLayoutAtomAOptionalTuple,
|
||||
class SmemCopyAtomAOptionalTuple,
|
||||
class TransformA_,
|
||||
class GmemTiledCopyBOptionalTuple,
|
||||
class SmemLayoutAtomBOptionalTuple,
|
||||
class SmemCopyAtomBOptionalTuple,
|
||||
class TransformB_>
|
||||
struct CollectiveMma<
|
||||
MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<Stages, ClusterShape, KernelSchedule>,
|
||||
TileShape_,
|
||||
ElementAOptionalTuple,
|
||||
StrideAOptionalTuple,
|
||||
ElementBOptionalTuple,
|
||||
StrideBOptionalTuple,
|
||||
TiledMma_,
|
||||
GmemTiledCopyAOptionalTuple,
|
||||
SmemLayoutAtomAOptionalTuple,
|
||||
SmemCopyAtomAOptionalTuple,
|
||||
TransformA_,
|
||||
GmemTiledCopyBOptionalTuple,
|
||||
SmemLayoutAtomBOptionalTuple,
|
||||
SmemCopyAtomBOptionalTuple,
|
||||
TransformB_>
|
||||
{
|
||||
private:
|
||||
template <class PointerType>
|
||||
static constexpr auto
|
||||
get_logical_ptr(PointerType const* ptr) {
|
||||
if constexpr (cute::sizeof_bits_v<PointerType> < 8) {
|
||||
return subbyte_iterator<PointerType const>(ptr);
|
||||
}
|
||||
else {
|
||||
return ptr;
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using DispatchPolicy = MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<Stages, ClusterShape, KernelSchedule>;
|
||||
using TileShape = TileShape_;
|
||||
|
||||
using ElementA = ElementAOptionalTuple;
|
||||
using ElementB = ElementBOptionalTuple;
|
||||
static constexpr bool IsATransformed = cute::sizeof_bits_v<ElementA> < cute::sizeof_bits_v<ElementB>;
|
||||
using ElementScale = void;
|
||||
|
||||
using StrideA = StrideAOptionalTuple;
|
||||
using StrideB = StrideBOptionalTuple;
|
||||
using StrideScale = void;
|
||||
static constexpr int AlignmentScale = cute::Int<0>{};
|
||||
|
||||
using TiledMma = TiledMma_;
|
||||
using ElementAccumulator = typename TiledMma::ValTypeC;
|
||||
|
||||
using GmemTiledCopyA = GmemTiledCopyAOptionalTuple;
|
||||
using GmemTiledCopyB = GmemTiledCopyBOptionalTuple;
|
||||
using GmemTiledCopyScale = void;
|
||||
|
||||
using SmemLayoutAtomA = SmemLayoutAtomAOptionalTuple;
|
||||
using SmemLayoutAtomB = SmemLayoutAtomBOptionalTuple;
|
||||
using SmemLayoutAtomScale = void;
|
||||
|
||||
using SmemCopyAtomA = SmemCopyAtomAOptionalTuple;
|
||||
using SmemCopyAtomB = SmemCopyAtomBOptionalTuple;
|
||||
using SmemCopyAtomScale = void;
|
||||
|
||||
// Swap and transpose A/B for A k-major layout and B mn-major layout since WGMMA is k-major only (e.g. tf32, Fp32, Int8, Fp8 WGMMA)
|
||||
static constexpr bool IsLayoutAkBmn =
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> &&
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor>;
|
||||
|
||||
static constexpr bool IsInputSizeTwoBytes = sizeof(ElementA) == 2 && sizeof(ElementB) == 2;
|
||||
|
||||
// We must ensure the type to be scaled goes to RF
|
||||
static constexpr bool SwapAB = !IsATransformed;
|
||||
using InternalSmemLayoutAtomA = cute::conditional_t<!SwapAB, SmemLayoutAtomA, SmemLayoutAtomB>;
|
||||
using InternalSmemLayoutAtomB = cute::conditional_t<!SwapAB, SmemLayoutAtomB, SmemLayoutAtomA>;
|
||||
using InternalSmemCopyAtomA = cute::conditional_t<!SwapAB, SmemCopyAtomA, SmemCopyAtomB>;
|
||||
using InternalSmemCopyAtomB = cute::conditional_t<!SwapAB, SmemCopyAtomB, SmemCopyAtomA>;
|
||||
// TMA converts f32 input to tf32 when copying from GMEM to SMEM
|
||||
// For all other types, cast to size equivalent uint type to avoid any rounding by TMA.
|
||||
static constexpr bool ConvertF32toTF32A = cute::is_same_v<float, ElementA>;
|
||||
static constexpr bool ConvertF32toTF32B = cute::is_same_v<float, ElementB>;
|
||||
using ConvertedElementA = cute::conditional_t<ConvertF32toTF32A, tfloat32_t, uint_bit_t<sizeof_bits_v<ElementA>>>;
|
||||
using ConvertedElementB = cute::conditional_t<ConvertF32toTF32B, tfloat32_t, uint_bit_t<sizeof_bits_v<ElementB>>>;
|
||||
using RealInternalElementA = cute::conditional_t<!SwapAB, ElementA, ElementB>;
|
||||
using RealInternalElementB = cute::conditional_t<!SwapAB, ElementB, ElementA>;
|
||||
using InternalElementA = cute::conditional_t<!SwapAB, ConvertedElementA, ConvertedElementB>;
|
||||
using InternalElementB = cute::conditional_t<!SwapAB, ConvertedElementB, ConvertedElementA>;
|
||||
using InternalStrideA = cute::conditional_t<!SwapAB, StrideA, StrideB>;
|
||||
using InternalStrideB = cute::conditional_t<!SwapAB, StrideB, StrideA>;
|
||||
|
||||
using TransformA = TransformA_;
|
||||
using TransformB = TransformB_;
|
||||
using InternalTransformA = cute::conditional_t<!SwapAB, TransformA, TransformB>;
|
||||
using InternalTransformB = cute::conditional_t<!SwapAB, TransformB, TransformA>;
|
||||
|
||||
static_assert(sizeof(InternalElementB) == 2 ||
|
||||
(cute::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::RowMajor> &&
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::ColumnMajor>),
|
||||
"B operand after swap must be 2 bytes OR K-major.");
|
||||
static constexpr int IsSubbyteA = cute::sizeof_bits_v<InternalElementA> < 8;
|
||||
using TmaElementA = cute::conditional_t<IsSubbyteA, uint8_t, InternalElementA>;
|
||||
|
||||
using ArchTag = typename DispatchPolicy::ArchTag;
|
||||
|
||||
using MainloopPipeline = cutlass::PipelineTmaAsync<
|
||||
DispatchPolicy::Stages,
|
||||
typename DispatchPolicy::ClusterShape>;
|
||||
using PipelineState = cutlass::PipelineState<DispatchPolicy::Stages>;
|
||||
|
||||
using PipelineParams = typename MainloopPipeline::Params;
|
||||
|
||||
static_assert(rank(InternalSmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<0>(TileShape{}) % size<0>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
static_assert(rank(InternalSmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
|
||||
static_assert((size<1>(TileShape{}) % size<0>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
static_assert((size<2>(TileShape{}) % size<1>(InternalSmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
|
||||
|
||||
static_assert(cute::is_same_v<GmemTiledCopyA, GmemTiledCopyScale> || cute::is_same_v<GmemTiledCopyScale, void>,
|
||||
"The TMA mcast for A must match the mcast for scales or the scale tiled copy must be void.");
|
||||
// Tile along modes in a way that maximizes the TMA box size.
|
||||
using SmemLayoutA = decltype(tile_to_shape(
|
||||
InternalSmemLayoutAtomA{},
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
using SmemLayoutB = decltype(tile_to_shape(
|
||||
InternalSmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
|
||||
// If A mn-layout and B mn-layout, transposing B matrix since WGMMA is k-major only (e.g. tf32, fp32, fp8, int8).
|
||||
static constexpr bool IsLayoutAmnBmn =
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagA_t<StrideA>, layout::ColumnMajor> &&
|
||||
cute::is_same_v<gemm::detail::StrideToLayoutTagB_t<StrideB>, layout::RowMajor>;
|
||||
static constexpr bool TransposeB = !IsInputSizeTwoBytes && IsLayoutAmnBmn;
|
||||
using TransposeOperandB = decltype(cutlass::transform::collective::detail::make_transpose_operand_b(
|
||||
0, 0, TiledMma{}, SmemLayoutB{}, InternalSmemLayoutAtomB{},
|
||||
InternalElementB{}, cute::bool_constant<TransposeB>{}));
|
||||
|
||||
static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 2 or more.");
|
||||
static_assert(not cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
|
||||
cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
|
||||
"MMA atom must source A from rmem and B operand from smem_desc for this mainloop.");
|
||||
static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>,
|
||||
"GmemTiledCopy - invalid SM90 TMA copy atom specified.");
|
||||
static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>,
|
||||
"GmemTiledCopy - invalid SM90 TMA copy atom specified.");
|
||||
|
||||
using GmmaSmemLayoutAtomB = decltype(transform::collective::detail::gmma_smem_transpose_or_passthrough<
|
||||
TransposeB, InternalSmemLayoutAtomB, InternalElementB>());
|
||||
|
||||
// SmemLayoutB for GMMA is different from SmemLayoutB for TMA if TransposeB
|
||||
using GmmaSmemLayoutB = decltype(tile_to_shape(
|
||||
GmmaSmemLayoutAtomB{},
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
|
||||
conditional_t< ::cutlass::gemm::detail::is_major<0,InternalStrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
|
||||
|
||||
static_assert(!SwapAB || !TransposeB, "Cannot SwapAB and TransposeB at the same time.");
|
||||
static_assert(TransposeB xor (cute::is_same_v<SmemLayoutB, GmmaSmemLayoutB>),
|
||||
"Should be same layout if not TransposeB.");
|
||||
static_assert(!TransposeB || size<1>(SmemLayoutB{}) * cute::sizeof_bits_v<InternalElementB> / 8 == 128,
|
||||
"SmemLayoutB K must be 128bytes to be transposed.");
|
||||
|
||||
static constexpr size_t SmemAlignmentA = cutlass::detail::alignment_for_swizzle(SmemLayoutA{});
|
||||
|
||||
static constexpr size_t SmemAlignmentB = cutlass::detail::alignment_for_swizzle(SmemLayoutB{});
|
||||
|
||||
static_assert(SmemAlignmentA >= 128 and SmemAlignmentB >= 128, "Require at least 128B alignment");
|
||||
|
||||
struct SharedStorage
|
||||
{
|
||||
struct TensorStorage : cute::aligned_struct<cute::max(SmemAlignmentA, SmemAlignmentB)> {
|
||||
cute::ArrayEngine<InternalElementA, cute::cosize_v<SmemLayoutA>> smem_A;
|
||||
cute::ArrayEngine<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>> smem_B;
|
||||
} tensors;
|
||||
|
||||
using PipelineStorage = typename MainloopPipeline::SharedStorage;
|
||||
PipelineStorage pipeline;
|
||||
};
|
||||
using TensorStorage = typename SharedStorage::TensorStorage;
|
||||
using PipelineStorage = typename SharedStorage::PipelineStorage;
|
||||
|
||||
// Host side kernel arguments
|
||||
struct Arguments {
|
||||
ElementA const* ptr_A = nullptr;
|
||||
StrideA dA{};
|
||||
ElementB const* ptr_B = nullptr;
|
||||
StrideB dB{};
|
||||
uint32_t mma_promotion_interval = 4;
|
||||
};
|
||||
|
||||
// Device side kernel params
|
||||
struct Params {
|
||||
private:
|
||||
using Outer = CollectiveMma<DispatchPolicy, TileShape_,
|
||||
ElementAOptionalTuple, StrideAOptionalTuple,
|
||||
ElementBOptionalTuple, StrideBOptionalTuple,
|
||||
TiledMma_,
|
||||
GmemTiledCopyAOptionalTuple, SmemLayoutAtomAOptionalTuple, SmemCopyAtomAOptionalTuple,
|
||||
TransformA,
|
||||
GmemTiledCopyBOptionalTuple, SmemLayoutAtomBOptionalTuple, SmemCopyAtomBOptionalTuple,
|
||||
TransformB>;
|
||||
|
||||
public:
|
||||
// Assumption: StrideA is congruent with Problem_MK
|
||||
using TMA_A = decltype(make_tma_copy<TmaElementA>(
|
||||
GmemTiledCopyA{},
|
||||
make_tensor(Outer::get_logical_ptr(static_cast<InternalElementA const*>(nullptr)), repeat_like(InternalStrideA{}, int32_t(0)), InternalStrideA{}),
|
||||
SmemLayoutA{}(_,_,cute::Int<0>{}),
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
|
||||
size<1>(ClusterShape{}))); // mcast along N mode for this M load, if any
|
||||
// Assumption: StrideB is congruent with Problem_NK
|
||||
using TMA_B = decltype(make_tma_copy(
|
||||
GmemTiledCopyB{},
|
||||
make_tensor(Outer::get_logical_ptr(static_cast<InternalElementB const*>(nullptr)), repeat_like(InternalStrideB{}, int32_t(0)), InternalStrideB{}),
|
||||
SmemLayoutB{}(_,_,cute::Int<0>{}),
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
|
||||
size<0>(ClusterShape{}))); // mcast along M mode for this N load, if any
|
||||
TMA_A tma_load_a;
|
||||
TMA_B tma_load_b;
|
||||
};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
template <class ProblemShape>
|
||||
static constexpr Params
|
||||
to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
|
||||
(void) workspace;
|
||||
|
||||
// Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
if constexpr (SwapAB) {
|
||||
M = get<1>(problem_shape_MNKL);
|
||||
N = get<0>(problem_shape_MNKL);
|
||||
}
|
||||
|
||||
InternalElementA const* ptr_A;
|
||||
InternalStrideA dA;
|
||||
InternalElementB const* ptr_B;
|
||||
InternalStrideB dB;
|
||||
|
||||
if constexpr (not SwapAB) {
|
||||
ptr_A = reinterpret_cast<InternalElementA const*>(args.ptr_A);
|
||||
ptr_B = reinterpret_cast<InternalElementB const*>(args.ptr_B);
|
||||
dA = args.dA;
|
||||
dB = args.dB;
|
||||
}
|
||||
else {
|
||||
ptr_A = reinterpret_cast<InternalElementA const*>(args.ptr_B);
|
||||
ptr_B = reinterpret_cast<InternalElementB const*>(args.ptr_A);
|
||||
dA = args.dB;
|
||||
dB = args.dA;
|
||||
}
|
||||
|
||||
Tensor tensor_a = make_tensor(get_logical_ptr(ptr_A), make_layout(make_shape(M,K,L), dA));
|
||||
Tensor tensor_b = make_tensor(get_logical_ptr(ptr_B), make_layout(make_shape(N,K,L), dB));
|
||||
typename Params::TMA_A tma_load_a = make_tma_copy<TmaElementA>(
|
||||
GmemTiledCopyA{},
|
||||
tensor_a,
|
||||
SmemLayoutA{}(_,_,cute::Int<0>{}),
|
||||
make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})),
|
||||
size<1>(ClusterShape{})); // mcast along N mode for this M load, if any
|
||||
typename Params::TMA_B tma_load_b = make_tma_copy(
|
||||
GmemTiledCopyB{},
|
||||
tensor_b,
|
||||
SmemLayoutB{}(_,_,cute::Int<0>{}),
|
||||
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})),
|
||||
size<0>(ClusterShape{})); // mcast along M mode for this N load, if any
|
||||
return {
|
||||
tma_load_a,
|
||||
tma_load_b
|
||||
};
|
||||
}
|
||||
|
||||
template<class ProblemShape>
|
||||
CUTLASS_HOST_DEVICE static bool
|
||||
can_implement(
|
||||
ProblemShape const& problem_shape,
|
||||
[[maybe_unused]] Arguments const& args) {
|
||||
constexpr int tma_alignment_bits = 128;
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
bool implementable = true;
|
||||
constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
|
||||
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
|
||||
constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
|
||||
implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
|
||||
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
|
||||
}
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
|
||||
static constexpr uint32_t TmaTransactionBytes =
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(cute::sizeof_bits_v<InternalElementA>) / 8)+
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(cute::sizeof_bits_v<InternalElementB>) / 8);
|
||||
|
||||
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
|
||||
CUTLASS_DEVICE
|
||||
static void prefetch_tma_descriptors(Params const& mainloop_params) {
|
||||
cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor());
|
||||
cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
|
||||
}
|
||||
|
||||
/// Set up the data needed by this collective for load and mma.
|
||||
/// Returns a tuple of tensors. The collective and the kernel layer have the contract
|
||||
/// that the tuple must contain at least two elements, with the first two elements being:
|
||||
/// gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l)
|
||||
/// gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l)
|
||||
/// The rest of the tensors can be specified as needed by this collective.
|
||||
template <class ProblemShape_MNKL,
|
||||
class TileShapeMNK>
|
||||
CUTLASS_DEVICE auto
|
||||
tile_input_tensors(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params, TileShapeMNK const& tileshape_mnk) {
|
||||
using X = Underscore;
|
||||
// Separate out problem shape for convenience
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, tileshape_mnk, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, tileshape_mnk, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
return cute::make_tuple(gA_mkl, gB_nkl);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class TensorA, class TensorB,
|
||||
class KTileIterator, class BlockCoord
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
load(
|
||||
Params const& mainloop_params,
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write,
|
||||
cute::tuple<TensorA, TensorB> const& tiled_tensors,
|
||||
BlockCoord const& blk_coord,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
int thread_idx,
|
||||
uint32_t block_rank_in_cluster,
|
||||
TensorStorage& shared_tensors) {
|
||||
|
||||
using namespace cute;
|
||||
int warp_idx = canonical_warp_idx_sync();
|
||||
int warp_idx_in_warp_group = warp_idx % 4;
|
||||
int lane_predicate = cute::elect_one_sync();
|
||||
|
||||
if (warp_idx_in_warp_group == 0 and lane_predicate) {
|
||||
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
//
|
||||
// Prepare the TMA loads for A and B
|
||||
//
|
||||
|
||||
constexpr uint32_t cluster_shape_x = get<0>(DispatchPolicy::ClusterShape());
|
||||
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
|
||||
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
|
||||
// Partition the inputs based on the current block coordinates.
|
||||
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Applies the mapping from block_tma_a
|
||||
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
|
||||
Tensor tAsA = block_tma_a.partition_D(sA); // (TMA,TMA_M,TMA_K,PIPE)
|
||||
|
||||
Tensor tBgB = block_tma_b.partition_S(gB); // (TMA,TMA_N,TMA_K,k)
|
||||
Tensor tBsB = block_tma_b.partition_D(sB); // (TMA,TMA_N,TMA_K,PIPE)
|
||||
|
||||
uint16_t mcast_mask_a = 0;
|
||||
uint16_t mcast_mask_b = 0;
|
||||
|
||||
// Issue TmaLoads
|
||||
// Maps the tile -> block, value
|
||||
if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) {
|
||||
auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) -> block_id
|
||||
for (int n = 0; n < size<1>(block_layout); ++n) {
|
||||
mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{}));
|
||||
}
|
||||
}
|
||||
|
||||
if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) {
|
||||
auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{}; // (m,n) -> block_id
|
||||
for (int m = 0; m < size<0>(block_layout); ++m) {
|
||||
mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{}));
|
||||
}
|
||||
}
|
||||
|
||||
// Mainloop
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > 0; --k_tile_count) {
|
||||
// LOCK smem_pipe_write for _writing_
|
||||
pipeline.producer_acquire(smem_pipe_write);
|
||||
|
||||
//
|
||||
// Copy gmem to smem for *k_tile_iter
|
||||
//
|
||||
|
||||
using BarrierType = typename MainloopPipeline::ProducerBarrierType;
|
||||
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
|
||||
|
||||
int write_stage = smem_pipe_write.index();
|
||||
copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
++k_tile_iter;
|
||||
|
||||
// Advance smem_pipe_write
|
||||
++smem_pipe_write;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
|
||||
CUTLASS_DEVICE void
|
||||
load_tail(MainloopPipeline pipeline, PipelineState smem_pipe_write) {
|
||||
int warp_idx = canonical_warp_idx_sync();
|
||||
int warp_idx_in_warp_group = warp_idx % 4;
|
||||
int lane_predicate = cute::elect_one_sync();
|
||||
|
||||
// Issue the epilogue waits
|
||||
if (warp_idx_in_warp_group == 0 and lane_predicate) {
|
||||
/* This helps avoid early exit of blocks in Cluster
|
||||
* Waits for all stages to either be released (all
|
||||
* Consumer UNLOCKs), or if the stage was never used
|
||||
* then would just be acquired since the phase was
|
||||
* still inverted from make_producer_start_state
|
||||
*/
|
||||
pipeline.producer_tail(smem_pipe_write);
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Consumer Perspective
|
||||
template <
|
||||
class FrgTensorC
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
mma(MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_read,
|
||||
FrgTensorC& accum,
|
||||
int k_tile_count,
|
||||
int thread_idx,
|
||||
TensorStorage& shared_tensors,
|
||||
Params const& mainloop_params) {
|
||||
using namespace cute;
|
||||
static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
|
||||
static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
|
||||
static_assert(rank(InternalSmemLayoutAtomA{}) == 2, "InternalSmemLayoutAtomA must be rank 2.");
|
||||
static_assert(rank(InternalSmemLayoutAtomB{}) == 2, "InternalSmemLayoutAtomB must be rank 2.");
|
||||
static_assert(!cute::is_void_v<InternalSmemCopyAtomA>,
|
||||
"SM90 GMMA mainloops must specify a non-void copy atom for RF sourced instructions.");
|
||||
static_assert(cute::is_void_v<InternalSmemCopyAtomB>,
|
||||
"SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
|
||||
|
||||
// Obtain warp index
|
||||
int warp_idx = canonical_warp_idx_sync();
|
||||
[[maybe_unused]] int warp_group_thread_idx = thread_idx % 128;
|
||||
|
||||
Tensor sA_ = make_tensor(make_smem_ptr(shared_tensors.smem_A.begin()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
|
||||
Tensor sA = as_position_independent_swizzle_tensor(sA_); // (BLK_M,BLK_K,PIPE)
|
||||
|
||||
Tensor sB_ = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
Tensor sB = as_position_independent_swizzle_tensor(sB_); // (BLK_M,BLK_K,PIPE)
|
||||
|
||||
// If TransposeB, GMMA will read from transposed B layout SMEM
|
||||
Tensor gmma_sB_position_dependent = make_tensor(make_smem_ptr(shared_tensors.smem_B.begin()),
|
||||
GmmaSmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
|
||||
Tensor gmma_sB = as_position_independent_swizzle_tensor(gmma_sB_position_dependent); // (BLK_N,BLK_K,PIPE)
|
||||
|
||||
//
|
||||
// Define C accumulators and A/B partitioning
|
||||
//
|
||||
|
||||
TiledMma tiled_mma;
|
||||
auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
|
||||
Tensor tCsA = thread_mma.partition_A(sA);
|
||||
|
||||
// Allocate fragments and descriptors
|
||||
Tensor tCrA_mma = thread_mma.partition_fragment_A(sA(_,_,Int<0>{})); // (MMA,MMA_M,MMA_K,PIPE)
|
||||
Tensor tCrA_load = make_fragment_like<InternalElementA>(tCrA_mma);
|
||||
|
||||
Tensor tCsB = thread_mma.partition_B(gmma_sB_position_dependent); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
Tensor tCrB = thread_mma.make_fragment_B(tCsB); // (MMA,MMA_N,MMA_K,PIPE)
|
||||
|
||||
//
|
||||
// Copy Atom A retiling
|
||||
//
|
||||
|
||||
auto smem_tiled_copy_A = make_tiled_copy_A(InternalSmemCopyAtomA{}, tiled_mma);
|
||||
auto smem_thr_copy_A = smem_tiled_copy_A.get_thread_slice(warp_group_thread_idx);
|
||||
|
||||
Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA_load); // (CPY,CPY_M,CPY_K)
|
||||
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view)); // CPY_M
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrA_copy_view)); // CPY_K
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCrA_mma) == size<1>(accum)); // MMA_M
|
||||
CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum)); // N
|
||||
CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB)); // K
|
||||
CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA)); // PIPE
|
||||
CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB)); // PIPE
|
||||
|
||||
//
|
||||
// PIPELINED MAIN LOOP
|
||||
//
|
||||
|
||||
// We release buffers to producer warps(dma load) with some mmas in flight
|
||||
PipelineState smem_pipe_release = smem_pipe_read;
|
||||
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
|
||||
|
||||
TransposeOperandB transpose = cutlass::transform::collective::detail::make_transpose_operand_b(
|
||||
warp_idx, warp_group_thread_idx, tiled_mma, SmemLayoutB{},
|
||||
InternalSmemLayoutAtomB{}, InternalElementB{},
|
||||
cute::bool_constant<TransposeB>{});
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
ConsumerToken barrier_token = {BarrierStatus::WaitAgain};
|
||||
// first k tile
|
||||
{
|
||||
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
++smem_pipe_read;
|
||||
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
|
||||
// copy smem->rmem for A operand
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,read_stage), tCrA_copy_view(_,_,0));
|
||||
transform_internal_A(tCrA_load(_, _, 0), tCrA_mma(_, _, 0));
|
||||
// transpose B operand in SMEM
|
||||
transpose(sB, gmma_sB, read_stage, 0);
|
||||
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA_load) - 1; ++k_block) {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
transform_internal_A(tCrA_load(_, _, k_block + 1), tCrA_mma(_, _, k_block + 1));
|
||||
transpose.synchronize(k_block);
|
||||
transpose(sB, gmma_sB, read_stage, k_block + 1);
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
}
|
||||
|
||||
warpgroup_wait<2>();
|
||||
|
||||
--k_tile_count;
|
||||
if (k_tile_count > 0) {
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
transform_internal_A(tCrA_load(_, _, 0), tCrA_mma(_, _, 0));
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
}
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
const int final_k = size<2>(tCrA_load) - 1;
|
||||
cute::gemm(tiled_mma, tCrA_mma(_,_, final_k), tCrB(_,_,final_k,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
}
|
||||
|
||||
if (k_tile_count == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// Mainloop GMMAs
|
||||
CUTLASS_PRAGMA_NO_UNROLL
|
||||
for ( ; k_tile_count > 1; --k_tile_count) {
|
||||
|
||||
//
|
||||
// Compute on k_tile
|
||||
//
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
++smem_pipe_read;
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA_load); ++k_block) {
|
||||
if (k_block == 0) {
|
||||
barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
|
||||
}
|
||||
if (k_block == size<2>(tCrA_load) - 1) {
|
||||
pipeline.consumer_wait(smem_pipe_read, barrier_token);
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,0,smem_pipe_read.index()), tCrA_copy_view(_,_,0));
|
||||
transform_internal_A(tCrA_load(_, _, 0), tCrA_mma(_, _, 0));
|
||||
// transpose B operand in SMEM
|
||||
transpose(sB, gmma_sB, smem_pipe_read.index(), 0);
|
||||
}
|
||||
else {
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
transform_internal_A(tCrA_load(_, _, k_block + 1), tCrA_mma(_, _, k_block + 1));
|
||||
// transpose B operand in SMEM
|
||||
transpose.synchronize(k_block); // make transpose of k_block available
|
||||
transpose(sB, gmma_sB, read_stage, k_block + 1);
|
||||
}
|
||||
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
if (k_block == 1) {
|
||||
// release prior barrier
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
{
|
||||
//
|
||||
// Compute on k_tile
|
||||
//
|
||||
|
||||
int read_stage = smem_pipe_read.index();
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
|
||||
// Unroll the K mode manually to set scale D to 1
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int k_block = 0; k_block < size<2>(tCrA_load) - 1; ++k_block) {
|
||||
|
||||
copy(smem_tiled_copy_A, tCsA(_,_,k_block + 1,read_stage), tCrA_copy_view(_,_,k_block + 1));
|
||||
transform_internal_A(tCrA_load(_, _, k_block + 1), tCrA_mma(_, _, k_block + 1));
|
||||
transpose.synchronize(k_block); // make k_block transpose available
|
||||
transpose(sB, gmma_sB, read_stage, k_block + 1);
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
cute::gemm(tiled_mma, tCrA_mma(_,_,k_block), tCrB(_,_,k_block,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
warpgroup_wait<2>();
|
||||
if (k_block == 1) {
|
||||
// release prior barrier
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
|
||||
warpgroup_arrive();
|
||||
// (V,M) x (V,N) => (V,M,N)
|
||||
const int final_k = size<2>(tCrA_load) - 1;
|
||||
cute::gemm(tiled_mma, tCrA_mma(_,_,final_k), tCrB(_,_,final_k,read_stage), accum);
|
||||
tiled_mma.accumulate_ = GMMA::ScaleOut::One;
|
||||
warpgroup_commit_batch();
|
||||
}
|
||||
|
||||
warpgroup_fence_operand(accum);
|
||||
}
|
||||
|
||||
/// Perform a Consumer Epilogue to release all buffers
|
||||
CUTLASS_DEVICE void
|
||||
mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
|
||||
// Prologue GMMAs
|
||||
int prologue_mma_count = 1;
|
||||
k_tile_count -= prologue_mma_count;
|
||||
|
||||
smem_pipe_release.advance(k_tile_count);
|
||||
|
||||
// Wait on all GMMAs to complete
|
||||
warpgroup_wait<0>();
|
||||
|
||||
for (int count = 0; count < prologue_mma_count; ++count) {
|
||||
pipeline.consumer_release(smem_pipe_release); // UNLOCK smem_pipe_release, done _computing_ on it
|
||||
++smem_pipe_release;
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
template <class Converter>
|
||||
static constexpr bool
|
||||
is_fast_converter_exact() {
|
||||
using DstType = typename Converter::result_type;
|
||||
using SrcType = typename Converter::source_type;
|
||||
|
||||
constexpr bool IsIntToFP32Exact = cute::is_same_v<float, DstType> &&
|
||||
(cute::numeric_limits<SrcType>::is_integer && cute::sizeof_bits_v<SrcType> <= 16);
|
||||
|
||||
constexpr bool IsIntToFP16orBF16Exact = (cute::is_same_v<cute::half_t, DstType> || cute::is_same_v<cute::bfloat16_t, DstType>) &&
|
||||
(cute::numeric_limits<SrcType>::is_integer && cute::sizeof_bits_v<SrcType> <= 8);
|
||||
|
||||
return IsIntToFP32Exact || IsIntToFP16orBF16Exact;
|
||||
}
|
||||
|
||||
template <class EngineIn, class EngineOut, class TensorLayout, int N = cosize_v<TensorLayout>>
|
||||
CUTLASS_DEVICE void
|
||||
transform_internal_A(Tensor<EngineIn,TensorLayout>&& in, Tensor<EngineOut,TensorLayout>&& out) {
|
||||
/// This is an element-wise conversion where we expect both tensors to have the same layout.
|
||||
/// As a result, we can cast as a cutlass array to use the fast numeric converters without
|
||||
/// worrying about indexing into the layout.
|
||||
|
||||
/// The inputs must be backed by registers & be statically sized so we can unroll the conversion loops.
|
||||
static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
|
||||
static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
|
||||
static_assert(cute::is_same_v<typename EngineIn::value_type, InternalElementA>, "Input engine must be same type as the A operand");
|
||||
static_assert(cute::is_same_v<typename EngineOut::value_type, typename TiledMma::ValTypeA>, "Output engine must be same type as the Mma input");
|
||||
static_assert(is_static_v<TensorLayout>, "Tensor layout for the conversion must be static");
|
||||
|
||||
using SrcArray = cutlass::Array<RealInternalElementA, N>;
|
||||
using DstArray = cutlass::Array<RealInternalElementB, N>;
|
||||
|
||||
constexpr cutlass::FloatRoundStyle RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
|
||||
using DefaultConverterAToB = cutlass::NumericArrayConverter<RealInternalElementB, RealInternalElementA, N, RoundStyle>;
|
||||
using FastConverterAToB = cutlass::FastNumericArrayConverter<RealInternalElementB, RealInternalElementA, N, RoundStyle>;
|
||||
|
||||
using ConverterAToB = cute::conditional_t<is_fast_converter_exact<FastConverterAToB>(), FastConverterAToB, DefaultConverterAToB>;
|
||||
|
||||
SrcArray* src_array_ptr = reinterpret_cast<SrcArray*>(raw_pointer_cast(in.data()));
|
||||
DstArray* dst_array_ptr = reinterpret_cast<DstArray*>(raw_pointer_cast(out.data()));
|
||||
*dst_array_ptr = std::move(ConverterAToB::convert(*src_array_ptr));
|
||||
}
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::collective
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
@@ -321,8 +321,8 @@ struct CollectiveMma<
|
||||
|
||||
// Set the bytes transferred in this TMA transaction (may involve multiple issues)
|
||||
constexpr uint32_t TmaTransactionBytes = static_cast<uint32_t>(
|
||||
(size<0>(sA) * size<1>(sA) * sizeof(InternalElementA)) +
|
||||
(size<0>(sB) * size<1>(sB) * sizeof(InternalElementB)));
|
||||
(size<0>(sA) * size<1>(sA) * sizeof_bits<InternalElementA>::value) / 8 +
|
||||
(size<0>(sB) * size<1>(sB) * sizeof_bits<InternalElementB>::value) / 8);
|
||||
|
||||
|
||||
// Obtain warp index
|
||||
|
||||
@@ -250,8 +250,8 @@ struct CollectiveMma<
|
||||
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
|
||||
static constexpr int K_PIPE_MMAS = 1;
|
||||
static constexpr uint32_t TmaTransactionBytes =
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof(ElementA)))+
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof(ElementB)));
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<ElementA>::value)) / 8+
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<ElementB>::value)) / 8;
|
||||
|
||||
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
|
||||
CUTLASS_DEVICE
|
||||
@@ -261,19 +261,45 @@ struct CollectiveMma<
|
||||
cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
|
||||
}
|
||||
|
||||
/// Set up the data needed by this collective for load and mma.
|
||||
/// Returns a tuple of tensors. The collective and the kernel layer have the contract
|
||||
/// that the tuple must contain at least two elements, with the first two elements being:
|
||||
/// gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l)
|
||||
/// gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l)
|
||||
/// The rest of the tensors can be specified as needed by this collective.
|
||||
template <class ProblemShape_MNKL,
|
||||
class TileShapeMNK>
|
||||
CUTLASS_DEVICE auto
|
||||
tile_input_tensors(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params, TileShapeMNK const& tileshape_mnk) {
|
||||
using X = Underscore;
|
||||
// Separate out problem shape for convenience
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, tileshape_mnk, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, tileshape_mnk, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
return cute::make_tuple(gA_mkl, gB_nkl);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class TensorA, class TMA_LOAD_A,
|
||||
class TensorB, class TMA_LOAD_B,
|
||||
class KTileIterator
|
||||
class TensorA, class TensorB,
|
||||
class KTileIterator, class BlockCoord
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
load(
|
||||
Params const& mainloop_params,
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write,
|
||||
TensorA const& gA, TMA_LOAD_A& tma_load_a,
|
||||
TensorB const& gB, TMA_LOAD_B& tma_load_b,
|
||||
cute::tuple<TensorA, TensorB> const& tiled_tensors,
|
||||
BlockCoord const& blk_coord,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
int thread_idx,
|
||||
uint32_t block_rank_in_cluster,
|
||||
@@ -296,8 +322,16 @@ struct CollectiveMma<
|
||||
constexpr uint32_t cluster_shape_x = get<0>(DispatchPolicy::ClusterShape());
|
||||
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
|
||||
|
||||
auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
|
||||
// Partition the inputs based on the current block coordinates.
|
||||
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Applies the mapping from block_tma_a
|
||||
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
|
||||
@@ -340,8 +374,8 @@ struct CollectiveMma<
|
||||
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
|
||||
|
||||
int write_stage = smem_pipe_write.index();
|
||||
copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
++k_tile_iter;
|
||||
|
||||
// Advance smem_pipe_write
|
||||
|
||||
@@ -248,8 +248,8 @@ struct CollectiveMma<
|
||||
static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
|
||||
static constexpr int K_PIPE_MMAS = 1;
|
||||
static constexpr uint32_t TmaTransactionBytes =
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof(ElementA)))+
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof(ElementB)));
|
||||
(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<ElementA>::value)) / 8+
|
||||
(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<ElementB>::value)) / 8;
|
||||
|
||||
/// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
|
||||
CUTLASS_DEVICE
|
||||
@@ -259,19 +259,45 @@ struct CollectiveMma<
|
||||
cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
|
||||
}
|
||||
|
||||
/// Set up the data needed by this collective for load and mma.
|
||||
/// Returns a tuple of tensors. The collective and the kernel layer have the contract
|
||||
/// that the tuple must contain at least two elements, with the first two elements being:
|
||||
/// gA_mkl - The tma tensor, A after a local tile so it has shape (BLK_M,BLK_K,m,k,l)
|
||||
/// gB_nkl - The tma tensor, B after a local tile so it has shape (BLK_N,BLK_K,n,k,l)
|
||||
/// The rest of the tensors can be specified as needed by this collective.
|
||||
template <class ProblemShape_MNKL,
|
||||
class TileShapeMNK>
|
||||
CUTLASS_DEVICE auto
|
||||
tile_input_tensors(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params, TileShapeMNK const& tileshape_mnk) {
|
||||
using X = Underscore;
|
||||
// Separate out problem shape for convenience
|
||||
auto [M,N,K,L] = problem_shape_MNKL;
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, tileshape_mnk, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, tileshape_mnk, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
return cute::make_tuple(gA_mkl, gB_nkl);
|
||||
}
|
||||
|
||||
/// Perform a collective-scoped matrix multiply-accumulate
|
||||
/// Producer Perspective
|
||||
template <
|
||||
class TensorA, class TMA_LOAD_A,
|
||||
class TensorB, class TMA_LOAD_B,
|
||||
class KTileIterator
|
||||
class TensorA, class TensorB,
|
||||
class KTileIterator, class BlockCoord
|
||||
>
|
||||
CUTLASS_DEVICE void
|
||||
load(
|
||||
Params const& mainloop_params,
|
||||
MainloopPipeline pipeline,
|
||||
PipelineState smem_pipe_write,
|
||||
TensorA const& gA, TMA_LOAD_A& tma_load_a,
|
||||
TensorB const& gB, TMA_LOAD_B& tma_load_b,
|
||||
cute::tuple<TensorA, TensorB> const& tiled_tensors,
|
||||
BlockCoord const& blk_coord,
|
||||
KTileIterator k_tile_iter, int k_tile_count,
|
||||
int thread_idx,
|
||||
uint32_t block_rank_in_cluster,
|
||||
@@ -294,8 +320,16 @@ struct CollectiveMma<
|
||||
constexpr uint32_t cluster_shape_x = get<0>(DispatchPolicy::ClusterShape());
|
||||
uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
|
||||
|
||||
auto block_tma_a = tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
|
||||
auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
|
||||
|
||||
// Partition the inputs based on the current block coordinates.
|
||||
auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Applies the mapping from block_tma_a
|
||||
Tensor tAgA = block_tma_a.partition_S(gA); // (TMA,TMA_M,TMA_K,k)
|
||||
@@ -338,8 +372,8 @@ struct CollectiveMma<
|
||||
BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
|
||||
|
||||
int write_stage = smem_pipe_write.index();
|
||||
copy(tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
|
||||
copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
|
||||
++k_tile_iter;
|
||||
|
||||
// Advance smem_pipe_write
|
||||
|
||||
@@ -41,6 +41,7 @@
|
||||
#include "cutlass/device_kernel.h"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/detail/layout.hpp"
|
||||
#include "cutlass/detail/mma.hpp"
|
||||
|
||||
#if !defined(__CUDACC_RTC__)
|
||||
#include "cutlass/cluster_launch.hpp"
|
||||
@@ -51,6 +52,7 @@
|
||||
#include "cutlass/gemm/device/gemm_universal_base.h"
|
||||
#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
|
||||
#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
|
||||
#include "cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h"
|
||||
|
||||
// 3.x
|
||||
#include "cutlass/gemm/kernel/gemm_universal.hpp"
|
||||
@@ -111,10 +113,7 @@ public:
|
||||
// Legacy: Assume MultiplyAdd only since we do not use this tag type in 3.0
|
||||
using MathOperator = cutlass::arch::OpMultiplyAdd;
|
||||
|
||||
// All tensorop operations have atom shape's M >= 8
|
||||
using OperatorClass = cute::conditional_t<
|
||||
cute::size<0>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{}) >= 8,
|
||||
cutlass::arch::OpClassTensorOp, cutlass::arch::OpClassSimt>;
|
||||
using OperatorClass = cutlass::detail::get_operator_class_t<typename CollectiveMainloop::TiledMma>;
|
||||
|
||||
using ArchTag = typename GemmKernel::ArchTag;
|
||||
|
||||
@@ -398,6 +397,7 @@ public:
|
||||
using GemmKernel = GemmKernel_;
|
||||
|
||||
static bool const kInternalTranspose =
|
||||
!cutlass::epilogue::threadblock::detail::is_2x_evt_v<typename GemmKernel::Epilogue> && // 2.x EVT does not require internal transpose
|
||||
cute::is_same<typename GemmKernel::LayoutC, cutlass::layout::RowMajor>::value;
|
||||
|
||||
using ThreadblockShape = typename GemmKernel::Mma::Shape;
|
||||
|
||||
@@ -43,23 +43,42 @@ using namespace cute;
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Policies for categorical dispatch of mainloop against kernel grid schedules
|
||||
// Kernel schedule policies (the base class tags, one for each kernel layer file)
|
||||
//
|
||||
struct KernelMultistage { };
|
||||
struct KernelCpAsyncWarpSpecialized { };
|
||||
struct KernelCpAsyncWarpSpecializedPingpong { };
|
||||
struct KernelCpAsyncWarpSpecializedCooperative { };
|
||||
struct KernelTma { };
|
||||
struct KernelTmaWarpSpecialized { };
|
||||
struct KernelTmaWarpSpecializedPingpong { };
|
||||
struct KernelTmaWarpSpecializedCooperative { };
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Builder dispatch policies (not a part of the main CUTLASS layers, simply used to opt into
|
||||
// specific collective builder dispatches)
|
||||
//
|
||||
|
||||
// FP8 related policies (including Fast Accumulation)
|
||||
struct KernelTmaWarpSpecializedFP8FastAccum : KernelTmaWarpSpecialized { };
|
||||
struct KernelTmaWarpSpecializedPingpongFP8FastAccum : KernelTmaWarpSpecializedPingpong { };
|
||||
struct KernelTmaWarpSpecializedCooperativeFP8FastAccum: KernelTmaWarpSpecializedCooperative { };
|
||||
|
||||
// Policies to opt into mixed type GEMMs
|
||||
struct KernelTmaWarpSpecializedMixedInput : KernelTmaWarpSpecialized { };
|
||||
struct KernelTmaWarpSpecializedPingpongMixedInput : KernelTmaWarpSpecializedPingpong { };
|
||||
struct KernelTmaWarpSpecializedCooperativeMixedInput: KernelTmaWarpSpecializedCooperative { };
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
// Policies for dispatch of epilogue
|
||||
struct EpilogueDefault { };
|
||||
struct EpilogueTransposed { };
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
//
|
||||
// Collective Mainloop Policies
|
||||
//
|
||||
@@ -98,28 +117,30 @@ struct MainloopSm80CpAsync {
|
||||
using ClusterShape = Shape<_1,_1,_1>;
|
||||
};
|
||||
|
||||
// n-buffer in smem (cp.async), pipelined with Hopper GMMA, WITHOUT predicated gmem loads
|
||||
// n-buffer in smem (cp.async), pipelined with Hopper GMMA, with predicated gmem loads, warp specialized dynamic schedule
|
||||
template<
|
||||
int Stages_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
class KernelSchedule = KernelCpAsyncWarpSpecialized
|
||||
>
|
||||
struct MainloopSm90CpAsyncGmmaUnpredicated {
|
||||
struct MainloopSm90CpAsyncGmmaWarpSpecialized {
|
||||
constexpr static int Stages = Stages_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm90;
|
||||
using Schedule = KernelMultistage;
|
||||
using Schedule = KernelSchedule;
|
||||
};
|
||||
|
||||
// n-buffer in smem (cp.async), pipelined with Hopper GMMA, with predicated gmem loads
|
||||
// n-buffer in smem (cp.async), pipelined with Hopper GMMA, with predicated gmem loads, warp specialized dynamic schedule
|
||||
template<
|
||||
int Stages_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
class KernelSchedule = KernelCpAsyncWarpSpecialized
|
||||
>
|
||||
struct MainloopSm90CpAsyncGmma {
|
||||
struct MainloopSm90CpAsyncGmmaRmemAWarpSpecialized {
|
||||
constexpr static int Stages = Stages_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm90;
|
||||
using Schedule = KernelMultistage;
|
||||
using Schedule = KernelSchedule;
|
||||
};
|
||||
|
||||
// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, static schedule between TMA and GMMA
|
||||
@@ -154,13 +175,11 @@ struct MainloopSm90TmaGmmaWarpSpecialized {
|
||||
template<
|
||||
int Stages_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
class KernelSchedule = KernelTmaWarpSpecialized,
|
||||
int PipelineAsyncMmaStages_ = 0
|
||||
class KernelSchedule = KernelTmaWarpSpecialized
|
||||
>
|
||||
struct MainloopSm90TmaGmmaRmemAWarpSpecialized {
|
||||
constexpr static int Stages = Stages_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
constexpr static int PipelineAsyncMmaStages = PipelineAsyncMmaStages_;
|
||||
using ArchTag = arch::Sm90;
|
||||
using Schedule = KernelSchedule;
|
||||
static_assert(
|
||||
@@ -170,6 +189,26 @@ struct MainloopSm90TmaGmmaRmemAWarpSpecialized {
|
||||
"KernelSchedule must be one of the warp specialized policies");
|
||||
};
|
||||
|
||||
template<
|
||||
int Stages_,
|
||||
class ClusterShape_ = Shape<_1,_1,_1>,
|
||||
class KernelSchedule = KernelTmaWarpSpecialized
|
||||
>
|
||||
struct MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput {
|
||||
constexpr static int Stages = Stages_;
|
||||
using ClusterShape = ClusterShape_;
|
||||
using ArchTag = arch::Sm90;
|
||||
using Schedule = KernelSchedule;
|
||||
static_assert(
|
||||
cute::is_same_v<Schedule, KernelTmaWarpSpecialized> ||
|
||||
cute::is_same_v<Schedule, KernelTmaWarpSpecializedMixedInput> ||
|
||||
cute::is_same_v<Schedule, KernelTmaWarpSpecializedPingpong> ||
|
||||
cute::is_same_v<Schedule, KernelTmaWarpSpecializedPingpongMixedInput> ||
|
||||
cute::is_same_v<Schedule, KernelTmaWarpSpecializedCooperative> ||
|
||||
cute::is_same_v<Schedule, KernelTmaWarpSpecializedCooperativeMixedInput>,
|
||||
"KernelSchedule must be one of the warp specialized policies");
|
||||
};
|
||||
|
||||
// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp specialized dynamic schedule
|
||||
// For FP8 kernels
|
||||
template<
|
||||
|
||||
@@ -375,14 +375,17 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false>
|
||||
|
||||
// Initialize the block mapping structure
|
||||
block_mapping = ThreadblockSwizzle(
|
||||
typename ThreadblockSwizzle::template KernelTraits<GemmStreamkWithFusedEpilogue>(),
|
||||
args.mode,
|
||||
args.problem_size,
|
||||
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
|
||||
args.batch_count,
|
||||
sm_occupancy,
|
||||
device_sms,
|
||||
avail_sms);
|
||||
avail_sms,
|
||||
sizeof(ElementA),
|
||||
sizeof(ElementB),
|
||||
sizeof(ElementC),
|
||||
Epilogue::kAccumulatorFragments);
|
||||
}
|
||||
|
||||
/// Returns the workspace size (in bytes) needed for these parameters
|
||||
|
||||
@@ -69,6 +69,9 @@ class GemmUniversal;
|
||||
|
||||
#include "cutlass/gemm/kernel/sm70_gemm.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_tma.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp"
|
||||
|
||||
@@ -357,14 +357,17 @@ public:
|
||||
|
||||
// Initialize the block mapping structure
|
||||
block_mapping = ThreadblockSwizzle(
|
||||
typename ThreadblockSwizzle::template KernelTraits<GemmUniversalStreamk>(),
|
||||
args.mode,
|
||||
args.problem_size,
|
||||
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
|
||||
args.batch_count,
|
||||
sm_occupancy,
|
||||
device_sms,
|
||||
avail_sms);
|
||||
avail_sms,
|
||||
sizeof(ElementA),
|
||||
sizeof(ElementB),
|
||||
sizeof(ElementC),
|
||||
Epilogue::kAccumulatorFragments);
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -233,14 +233,17 @@ public:
|
||||
|
||||
// Initialize the block mapping structure
|
||||
block_mapping = ThreadblockSwizzle(
|
||||
typename ThreadblockSwizzle::template KernelTraits<GemmWithEpilogueVisitorStreamk>(),
|
||||
args.mode,
|
||||
args.problem_size,
|
||||
{ThreadblockShape::kM, ThreadblockShape::kN, ThreadblockShape::kK},
|
||||
args.batch_count,
|
||||
sm_occupancy,
|
||||
device_sms,
|
||||
avail_sms);
|
||||
avail_sms,
|
||||
sizeof(ElementA),
|
||||
sizeof(ElementB),
|
||||
sizeof(ElementC),
|
||||
Epilogue::kAccumulatorFragments);
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -323,26 +323,26 @@ public:
|
||||
static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
// Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
TiledMma tiled_mma;
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
||||
|
||||
// Prepare and partition the input tensors. Expects a tuple of tensors where:
|
||||
// get<0>(tiled_tensors) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
|
||||
// get<1>(tiled_tensors) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
|
||||
auto tiled_tensors = collective_mainloop.tile_input_tensors(problem_shape_MNKL, params.mainloop, blk_shape);
|
||||
static_assert(tuple_size_v<decltype(tiled_tensors)> >= 2, "Output of tile_input_tensors must have at least two elements (A, B)");
|
||||
|
||||
// Extract out partitioned A and B.
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
// Compute m_coord, n_coord, and l_coord with their post-tiled shapes
|
||||
auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mkl));
|
||||
@@ -350,28 +350,21 @@ public:
|
||||
auto l_coord = idx2crd(int(blockIdx.z), shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with m_coord and n_coord
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Get pipeline iterators and increments from tensor shapes
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
|
||||
auto k_tile_count = size<2>(gA);
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<3>(gA_mkl));
|
||||
auto k_tile_count = size<3>(gA_mkl);
|
||||
|
||||
// Wait for all thread blocks in the Cluster
|
||||
cluster_wait_fn();
|
||||
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
||||
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
if (producer_warp_role == ProducerWarpRole::MainloopEpilogue) {
|
||||
collective_mainloop.load(
|
||||
params.mainloop,
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA, params.mainloop.tma_load_a,
|
||||
gB, params.mainloop.tma_load_b,
|
||||
tiled_tensors,
|
||||
blk_coord,
|
||||
k_tile_iter, k_tile_count,
|
||||
lane_idx,
|
||||
block_rank_in_cluster,
|
||||
@@ -408,7 +401,7 @@ public:
|
||||
mainloop_pipe_consumer_state,
|
||||
accumulators,
|
||||
k_tile_count,
|
||||
thread_idx,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.mainloop,
|
||||
params.mainloop
|
||||
);
|
||||
|
||||
@@ -31,6 +31,7 @@
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/workspace.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
@@ -187,16 +188,29 @@ public:
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
|
||||
|
||||
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
|
||||
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
|
||||
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
|
||||
|
||||
// Calculate workspace pointers
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
||||
size_t workspace_offset = 0;
|
||||
|
||||
void* scheduler_workspace = workspace_ptr;
|
||||
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
|
||||
void* epilogue_workspace = workspace_ptr + workspace_offset;
|
||||
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
|
||||
void* mainloop_workspace = nullptr;
|
||||
|
||||
return {
|
||||
args.mode,
|
||||
problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace),
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, mainloop_workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, epilogue_workspace),
|
||||
hw_info,
|
||||
scheduler,
|
||||
TileScheduler::to_underlying_arguments(problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace),
|
||||
workspace
|
||||
};
|
||||
}
|
||||
@@ -215,19 +229,42 @@ public:
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static int
|
||||
static size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
TileScheduler t;
|
||||
return t.template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
size_t workspace_size = 0;
|
||||
workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
||||
|
||||
workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
||||
|
||||
return workspace_size;
|
||||
}
|
||||
|
||||
static
|
||||
cutlass::Status
|
||||
static cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
TileScheduler t;
|
||||
return t.template initialize_workspace<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, workspace, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
Status status = Status::kSuccess;
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
||||
size_t workspace_offset = 0;
|
||||
|
||||
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream);
|
||||
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
// Computes the kernel launch grid shape based on runtime parameters
|
||||
@@ -368,30 +405,13 @@ public:
|
||||
}
|
||||
} ();
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
// Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
TiledMma tiled_mma;
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
// Get pipeline stage increments from tensor shapes
|
||||
auto k_tile_count = size<3>(gA_mkl);
|
||||
|
||||
TileScheduler scheduler{params.scheduler};
|
||||
auto work_tile_info = scheduler.get_current_work();
|
||||
|
||||
@@ -399,6 +419,19 @@ public:
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
||||
|
||||
// Prepare and partition the input tensors. Expects a tuple of tensors where:
|
||||
// get<0>(tiled_tensors) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
|
||||
// get<1>(tiled_tensors) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
|
||||
auto tiled_tensors = collective_mainloop.tile_input_tensors(problem_shape_MNKL, params.mainloop, blk_shape);
|
||||
static_assert(tuple_size_v<decltype(tiled_tensors)> >= 2, "Output of tile_input_tensors must have at least two elements (A, B)");
|
||||
|
||||
// Extract out partitioned A and B.
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
// Get pipeline stage increments from tensor shapes
|
||||
auto k_tile_count = size<3>(gA_mkl);
|
||||
|
||||
// Wait for all thread blocks in the Cluster
|
||||
cluster_wait_fn();
|
||||
|
||||
@@ -408,27 +441,24 @@ public:
|
||||
// Mainloop Producer Warp
|
||||
if (producer_warp_role == ProducerWarpRole::Mainloop) {
|
||||
bool do_load_order_arrive = true;
|
||||
while (work_tile_info.is_valid_tile) {
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with our work tile coordinates to construct mainloop tensor views
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
|
||||
auto work_k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
|
||||
auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info);
|
||||
auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<2>(gA)), shape<2>(gA));
|
||||
auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl));
|
||||
|
||||
collective_mainloop.load(
|
||||
params.mainloop,
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA, params.mainloop.tma_load_a,
|
||||
gB, params.mainloop.tma_load_b,
|
||||
tiled_tensors,
|
||||
blk_coord,
|
||||
k_tile_iter, work_k_tile_count,
|
||||
lane_idx,
|
||||
block_rank_in_cluster,
|
||||
@@ -454,8 +484,8 @@ public:
|
||||
// Epilogue Producer Warp
|
||||
else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
|
||||
load_order_barrier.wait();
|
||||
while (work_tile_info.is_valid_tile) {
|
||||
if (TileScheduler::compute_epilogue(work_tile_info)) {
|
||||
while (work_tile_info.is_valid()) {
|
||||
if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
@@ -489,7 +519,7 @@ public:
|
||||
|
||||
// Do we potentially issue tail arrives for TMA stores, if epilogue load is waiting for it
|
||||
bool do_store_tail = false;
|
||||
while (work_tile_info.is_valid_tile) {
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
@@ -529,7 +559,7 @@ public:
|
||||
TileScheduler::fixup(
|
||||
params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, consumer_warp_group_idx);
|
||||
|
||||
if (TileScheduler::compute_epilogue(work_tile_info)) {
|
||||
if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
|
||||
// Epilogue and write to gD
|
||||
auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] =
|
||||
collective_epilogue.store(
|
||||
|
||||
@@ -31,6 +31,7 @@
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/workspace.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
@@ -196,13 +197,28 @@ public:
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
|
||||
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
|
||||
|
||||
// Calculate workspace pointers
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
||||
size_t workspace_offset = 0;
|
||||
|
||||
void* scheduler_workspace = workspace_ptr;
|
||||
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
|
||||
void* epilogue_workspace = workspace_ptr + workspace_offset;
|
||||
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
|
||||
void* mainloop_workspace = nullptr;
|
||||
|
||||
return {
|
||||
args.mode,
|
||||
problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace),
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, mainloop_workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, epilogue_workspace),
|
||||
hw_info,
|
||||
TileScheduler::to_underlying_arguments(problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler)
|
||||
TileScheduler::to_underlying_arguments(problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace)
|
||||
};
|
||||
}
|
||||
|
||||
@@ -220,16 +236,42 @@ public:
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static
|
||||
int
|
||||
static size_t
|
||||
get_workspace_size(Arguments const& args) {
|
||||
return 0;
|
||||
size_t workspace_size = 0;
|
||||
workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
||||
|
||||
workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
||||
|
||||
return workspace_size;
|
||||
}
|
||||
|
||||
static
|
||||
cutlass::Status
|
||||
static cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
return Status::kSuccess;
|
||||
Status status = Status::kSuccess;
|
||||
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
||||
size_t workspace_offset = 0;
|
||||
|
||||
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream);
|
||||
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
||||
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
||||
if (status != Status::kSuccess) {
|
||||
return status;
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
// Computes the kernel launch grid shape based on runtime parameters
|
||||
@@ -371,25 +413,26 @@ public:
|
||||
} ();
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
// Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// TMA requires special handling of strides to deal with coord codomain mapping
|
||||
// Represent the full tensors -- get these from TMA
|
||||
Tensor mA_mkl = params.mainloop.tma_load_a.get_tma_tensor(make_shape(M,K,L)); // (m,k,l)
|
||||
Tensor mB_nkl = params.mainloop.tma_load_b.get_tma_tensor(make_shape(N,K,L)); // (n,k,l)
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
TiledMma tiled_mma;
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
||||
|
||||
// Prepare and partition the input tensors. Expects a tuple of tensors where:
|
||||
// get<0>(tiled_tensors) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
|
||||
// get<1>(tiled_tensors) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
|
||||
auto tiled_tensors = collective_mainloop.tile_input_tensors(problem_shape_MNKL, params.mainloop, blk_shape);
|
||||
static_assert(tuple_size_v<decltype(tiled_tensors)> >= 2, "Output of tile_input_tensors must have at least two elements (A, B)");
|
||||
|
||||
// Extract out partitioned A and B.
|
||||
Tensor gA_mkl = get<0>(tiled_tensors);
|
||||
Tensor gB_nkl = get<1>(tiled_tensors);
|
||||
|
||||
// Get pipeline stage increments from tensor shapes
|
||||
auto k_tile_count = size<3>(gA_mkl);
|
||||
@@ -408,10 +451,6 @@ public:
|
||||
}
|
||||
auto work_tile_info = scheduler.get_current_work();
|
||||
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
||||
|
||||
// Wait for all thread blocks in the Cluster
|
||||
cluster_wait_fn();
|
||||
|
||||
@@ -421,24 +460,21 @@ public:
|
||||
// Mainloop Producer Warp
|
||||
if (producer_warp_role == ProducerWarpRole::Mainloop) {
|
||||
bool do_load_order_arrive = true;
|
||||
while (work_tile_info.is_valid_tile) {
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with our work tile coordinates to construct mainloop tensor views
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<3>(gA_mkl));
|
||||
|
||||
collective_mainloop.load(
|
||||
params.mainloop,
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA, params.mainloop.tma_load_a,
|
||||
gB, params.mainloop.tma_load_b,
|
||||
tiled_tensors,
|
||||
blk_coord,
|
||||
k_tile_iter, k_tile_count,
|
||||
lane_idx,
|
||||
block_rank_in_cluster,
|
||||
@@ -465,7 +501,7 @@ public:
|
||||
// Epilogue Producer Warp
|
||||
else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
|
||||
load_order_barrier.wait();
|
||||
while (work_tile_info.is_valid_tile) {
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
@@ -497,7 +533,7 @@ public:
|
||||
else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
|
||||
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
|
||||
|
||||
while (work_tile_info.is_valid_tile) {
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
@@ -515,7 +551,7 @@ public:
|
||||
mainloop_pipe_consumer_state,
|
||||
accumulators,
|
||||
k_tile_count,
|
||||
thread_idx,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.mainloop,
|
||||
params.mainloop
|
||||
);
|
||||
|
||||
@@ -0,0 +1,417 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cutlass/arch/reg_reconfig.h"
|
||||
#include "cutlass/arch/mma_sm90.h"
|
||||
#include "cutlass/epilogue/collective/detail.hpp"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cute/tensor.hpp"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::kernel {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ProblemShape_,
|
||||
class CollectiveMainloop_,
|
||||
class CollectiveEpilogue_,
|
||||
class TileScheduler_
|
||||
>
|
||||
class GemmUniversal<
|
||||
ProblemShape_,
|
||||
CollectiveMainloop_,
|
||||
CollectiveEpilogue_,
|
||||
TileScheduler_,
|
||||
cute::enable_if_t<cute::is_base_of_v<KernelCpAsyncWarpSpecialized, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
|
||||
{
|
||||
public:
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using ProblemShape = ProblemShape_;
|
||||
static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
|
||||
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
|
||||
|
||||
// Mainloop derived types
|
||||
using CollectiveMainloop = CollectiveMainloop_;
|
||||
using TileShape = typename CollectiveMainloop::TileShape;
|
||||
using TiledMma = typename CollectiveMainloop::TiledMma;
|
||||
using ArchTag = typename CollectiveMainloop::ArchTag;
|
||||
using ElementA = typename CollectiveMainloop::ElementA;
|
||||
using StrideA = typename CollectiveMainloop::StrideA;
|
||||
using ElementB = typename CollectiveMainloop::ElementB;
|
||||
using StrideB = typename CollectiveMainloop::StrideB;
|
||||
using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
|
||||
using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
|
||||
using ClusterShape = typename DispatchPolicy::ClusterShape;
|
||||
using MainloopArguments = typename CollectiveMainloop::Arguments;
|
||||
using MainloopParams = typename CollectiveMainloop::Params;
|
||||
static_assert(ArchTag::kMinComputeCapability >= 90);
|
||||
|
||||
// Epilogue derived types
|
||||
using CollectiveEpilogue = CollectiveEpilogue_;
|
||||
using ElementC = typename CollectiveEpilogue::ElementC;
|
||||
using StrideC = typename CollectiveEpilogue::StrideC;
|
||||
using ElementD = typename CollectiveEpilogue::ElementD;
|
||||
using StrideD = typename CollectiveEpilogue::StrideD;
|
||||
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
|
||||
using EpilogueParams = typename CollectiveEpilogue::Params;
|
||||
|
||||
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
|
||||
"Non-persistent warp-specialized kernel does not support specializing the tile scheduler.");
|
||||
using TileSchedulerTag = TileScheduler_;
|
||||
using TileScheduler = typename detail::TileSchedulerSelector<
|
||||
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
|
||||
using TileSchedulerArguments = typename TileScheduler::Arguments;
|
||||
|
||||
// Kernel level shared memory storage
|
||||
struct SharedStorage {
|
||||
union TensorStorage {
|
||||
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
|
||||
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
|
||||
|
||||
MainloopTensorStorage mainloop;
|
||||
EpilogueTensorStorage epilogue;
|
||||
} tensors;
|
||||
|
||||
struct PipelineStorage : cute::aligned_struct<16> {
|
||||
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
|
||||
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
|
||||
|
||||
alignas(16) MainloopPipelineStorage mainloop;
|
||||
alignas(16) EpiLoadPipelineStorage epi_load;
|
||||
} pipelines;
|
||||
};
|
||||
|
||||
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
||||
|
||||
using GmemTiledCopyA = typename CollectiveMainloop::GmemTiledCopyA;
|
||||
using GmemTiledCopyB = typename CollectiveMainloop::GmemTiledCopyB;
|
||||
static_assert(cute::size(GmemTiledCopyA{}) == cute::size(GmemTiledCopyB{}), "Number of threads in A/B tiled copies must be the same.");
|
||||
|
||||
static constexpr uint32_t NumLoadWarpGroups = cute::size(GmemTiledCopyA{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumMmaWarpGroups = cute::size(TiledMma{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumWarpGroups = NumLoadWarpGroups + NumMmaWarpGroups;
|
||||
static_assert(NumWarpGroups == 2 || NumWarpGroups == 3, "Number of warp groups must be 2 or 3 for good performance.");
|
||||
|
||||
static constexpr uint32_t MaxThreadsPerBlock = NumWarpGroups * NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
|
||||
// Device side arguments
|
||||
struct Arguments {
|
||||
GemmUniversalMode mode{};
|
||||
ProblemShape problem_shape{};
|
||||
MainloopArguments mainloop{};
|
||||
EpilogueArguments epilogue{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
TileSchedulerArguments scheduler{};
|
||||
};
|
||||
|
||||
// Kernel entry point API
|
||||
struct Params {
|
||||
GemmUniversalMode mode;
|
||||
ProblemShape problem_shape;
|
||||
MainloopParams mainloop;
|
||||
EpilogueParams epilogue;
|
||||
};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
// Convert to underlying arguments. In this case, a simple copy for the aliased type.
|
||||
static
|
||||
Params
|
||||
to_underlying_arguments(Arguments const& args, void* workspace) {
|
||||
(void) workspace;
|
||||
auto problem_shape = args.problem_shape;
|
||||
if constexpr (detail::IF_SWAP_AB<CollectiveMainloop>::value) {
|
||||
// swap M/N
|
||||
get<0>(problem_shape) = get<1>(args.problem_shape);
|
||||
get<1>(problem_shape) = get<0>(args.problem_shape);
|
||||
}
|
||||
return {
|
||||
args.mode,
|
||||
problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace)
|
||||
};
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE static
|
||||
bool
|
||||
can_implement(Arguments const& args) {
|
||||
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
|
||||
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
|
||||
return implementable;
|
||||
}
|
||||
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
|
||||
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static
|
||||
int
|
||||
get_workspace_size(Arguments const& args) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static
|
||||
cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
// Computes the kernel launch grid shape based on runtime parameters
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
auto cluster_shape = Shape<_1,_1,_1>{};
|
||||
auto tile_shape = TileShape{};
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
return TileScheduler::get_tiled_cta_shape_mnl(
|
||||
problem_shape_MNKL, tile_shape, cluster_shape);
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_block_shape() {
|
||||
return dim3(MaxThreadsPerBlock, 1, 1);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
operator()(Params const& params, char* smem_buf) {
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
|
||||
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
|
||||
if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) {
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
enum class WarpGroupRole {
|
||||
Producer = 0,
|
||||
Consumer = 1,
|
||||
};
|
||||
|
||||
// Kernel level shared memory storage
|
||||
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
|
||||
int thread_idx = int(threadIdx.x);
|
||||
int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
|
||||
int warp_group_idx = canonical_warp_group_idx();
|
||||
CUTLASS_ASSERT(warp_group_idx < NumWarpGroups);
|
||||
WarpGroupRole warp_group_role = warp_group_idx < NumLoadWarpGroups ? WarpGroupRole::Producer : WarpGroupRole::Consumer;
|
||||
|
||||
// Mainloop Load pipeline
|
||||
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
|
||||
typename MainloopPipeline::Params mainloop_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
mainloop_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
mainloop_pipeline_params.consumer_arv_count = NumMmaWarpGroups * NumThreadsPerWarpGroup;
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
|
||||
|
||||
// Epilogue Load pipeline
|
||||
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
|
||||
typename EpiLoadPipeline::Params epi_load_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
epi_load_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
epi_load_pipeline_params.consumer_arv_count = NumMmaWarpGroups * NumThreadsPerWarpGroup;
|
||||
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
|
||||
|
||||
// Epilogue Store pipeline
|
||||
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
|
||||
typename EpiStorePipeline::Params epi_store_pipeline_params;
|
||||
epi_store_pipeline_params.always_wait = true;
|
||||
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
|
||||
|
||||
// Initialize starting pipeline states for the collectives
|
||||
// Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
|
||||
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
|
||||
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
|
||||
|
||||
// For the DMA Load (producer) we start with an opposite phase
|
||||
// i.e., we skip all waits since we know that the buffer is indeed empty
|
||||
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
|
||||
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
|
||||
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
|
||||
|
||||
// Preconditions
|
||||
static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA_mkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_A), make_shape(M,K,L), params.mainloop.dA); //(m,k,l)
|
||||
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_B), make_shape(N,K,L), params.mainloop.dB); //(n,k,l)
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
TiledMma tiled_mma;
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
// Compute m_coord, n_coord, and l_coord with their post-tiled shapes
|
||||
auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(int(blockIdx.z), shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with m_coord and n_coord
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Get pipeline iterators and increments from tensor shapes
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
|
||||
auto k_tile_count = size<2>(gA);
|
||||
auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape);
|
||||
auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape);
|
||||
|
||||
// Wait for all threads in the thread block
|
||||
__syncthreads();
|
||||
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue{params.epilogue, shared_storage.tensors.epilogue};
|
||||
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
// Compute tile residues for predication
|
||||
auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord); // M - BLK_M * m_coord
|
||||
auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord); // N - BLK_N * n_coord
|
||||
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
|
||||
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
|
||||
|
||||
collective_mainloop.load(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA,
|
||||
gB,
|
||||
k_tile_iter, k_tile_count,
|
||||
residue_mnk,
|
||||
thread_idx,
|
||||
shared_storage.tensors.mainloop
|
||||
);
|
||||
// Update starting mainloop pipeline state for the pipeline drain
|
||||
mainloop_pipe_producer_state.advance(k_tile_count);
|
||||
// Make sure mainloop consumer has been waited upon before issuing epilogue load
|
||||
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
|
||||
|
||||
if (collective_epilogue.is_producer_load_needed()) {
|
||||
epi_load_pipe_producer_state =
|
||||
collective_epilogue.load(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
tiled_mma,
|
||||
thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
|
||||
}
|
||||
}
|
||||
else if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
collective_mainloop.mma(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
accumulators,
|
||||
k_tile_count,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.mainloop,
|
||||
params.mainloop
|
||||
);
|
||||
|
||||
// Make sure the math instructions are done and free buffers before entering the epilogue
|
||||
collective_mainloop.mma_tail(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
k_tile_count
|
||||
);
|
||||
|
||||
// Epilogue and write to gD
|
||||
collective_epilogue.store(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline,
|
||||
epi_store_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
accumulators,
|
||||
tiled_mma,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::kernel
|
||||
@@ -0,0 +1,518 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cutlass/arch/reg_reconfig.h"
|
||||
#include "cutlass/arch/mma_sm90.h"
|
||||
#include "cutlass/epilogue/collective/detail.hpp"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cute/tensor.hpp"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::kernel {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ProblemShape_,
|
||||
class CollectiveMainloop_,
|
||||
class CollectiveEpilogue_,
|
||||
class TileScheduler_
|
||||
>
|
||||
class GemmUniversal<
|
||||
ProblemShape_,
|
||||
CollectiveMainloop_,
|
||||
CollectiveEpilogue_,
|
||||
TileScheduler_,
|
||||
cute::enable_if_t<cute::is_base_of_v<KernelCpAsyncWarpSpecializedCooperative, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
|
||||
{
|
||||
public:
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using ProblemShape = ProblemShape_;
|
||||
static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
|
||||
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
|
||||
|
||||
// Mainloop derived types
|
||||
using CollectiveMainloop = CollectiveMainloop_;
|
||||
using TileShape = typename CollectiveMainloop::TileShape;
|
||||
using TiledMma = typename CollectiveMainloop::TiledMma;
|
||||
using ArchTag = typename CollectiveMainloop::ArchTag;
|
||||
using ElementA = typename CollectiveMainloop::ElementA;
|
||||
using StrideA = typename CollectiveMainloop::StrideA;
|
||||
using ElementB = typename CollectiveMainloop::ElementB;
|
||||
using StrideB = typename CollectiveMainloop::StrideB;
|
||||
using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
|
||||
using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
|
||||
using ClusterShape = typename DispatchPolicy::ClusterShape;
|
||||
using MainloopArguments = typename CollectiveMainloop::Arguments;
|
||||
using MainloopParams = typename CollectiveMainloop::Params;
|
||||
static_assert(ArchTag::kMinComputeCapability >= 90);
|
||||
|
||||
// Epilogue derived types
|
||||
using CollectiveEpilogue = CollectiveEpilogue_;
|
||||
using ElementC = typename CollectiveEpilogue::ElementC;
|
||||
using StrideC = typename CollectiveEpilogue::StrideC;
|
||||
using ElementD = typename CollectiveEpilogue::ElementD;
|
||||
using StrideD = typename CollectiveEpilogue::StrideD;
|
||||
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
|
||||
using EpilogueParams = typename CollectiveEpilogue::Params;
|
||||
|
||||
using TileSchedulerTag = TileScheduler_;
|
||||
using TileScheduler = typename detail::TileSchedulerSelector<
|
||||
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
|
||||
using TileSchedulerArguments = typename TileScheduler::Arguments;
|
||||
using TileSchedulerParams = typename TileScheduler::Params;
|
||||
|
||||
using GmemTiledCopyA = typename CollectiveMainloop::GmemTiledCopyA;
|
||||
using GmemTiledCopyB = typename CollectiveMainloop::GmemTiledCopyB;
|
||||
static_assert(cute::size(GmemTiledCopyA{}) == cute::size(GmemTiledCopyB{}), "Number of threads in A/B tiled copies must be the same");
|
||||
|
||||
static constexpr uint32_t NumLoadWarpGroups = cute::size(GmemTiledCopyA{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumMmaWarpGroups = cute::size(TiledMma{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumWarpGroups = NumLoadWarpGroups + NumMmaWarpGroups;
|
||||
static_assert(NumWarpGroups == 2 || NumWarpGroups == 3, "Number of warp groups must be 2 or 3 for good performance.");
|
||||
|
||||
static constexpr uint32_t MaxThreadsPerBlock = NumWarpGroups * NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
|
||||
// Kernel level shared memory storage
|
||||
struct SharedStorage {
|
||||
struct TensorStorage : cute::aligned_struct<128> {
|
||||
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
|
||||
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
|
||||
|
||||
MainloopTensorStorage mainloop;
|
||||
EpilogueTensorStorage epilogue;
|
||||
} tensors;
|
||||
|
||||
struct PipelineStorage : cute::aligned_struct<16> {
|
||||
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
|
||||
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
|
||||
|
||||
alignas(16) MainloopPipelineStorage mainloop;
|
||||
alignas(16) EpiLoadPipelineStorage epi_load;
|
||||
} pipelines;
|
||||
};
|
||||
|
||||
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
||||
|
||||
// Device side arguments
|
||||
struct Arguments {
|
||||
GemmUniversalMode mode{};
|
||||
ProblemShape problem_shape{};
|
||||
MainloopArguments mainloop{};
|
||||
EpilogueArguments epilogue{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
TileSchedulerArguments scheduler{};
|
||||
};
|
||||
|
||||
// Kernel entry point API
|
||||
struct Params {
|
||||
GemmUniversalMode mode;
|
||||
ProblemShape problem_shape;
|
||||
MainloopParams mainloop;
|
||||
EpilogueParams epilogue;
|
||||
KernelHardwareInfo hw_info;
|
||||
TileSchedulerParams scheduler;
|
||||
};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
// Convert to underlying arguments. In this case, a simple copy for the aliased type.
|
||||
static
|
||||
Params
|
||||
to_underlying_arguments(Arguments const& args, void* workspace) {
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments():");
|
||||
|
||||
auto problem_shape = args.problem_shape;
|
||||
if constexpr (detail::IF_SWAP_AB<CollectiveMainloop>::value) {
|
||||
// swap M/N
|
||||
get<0>(problem_shape) = get<1>(args.problem_shape);
|
||||
get<1>(problem_shape) = get<0>(args.problem_shape);
|
||||
}
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
|
||||
// Get SM count if needed, otherwise use user supplied SM count
|
||||
int sm_count = args.hw_info.sm_count;
|
||||
if (sm_count <= 0) {
|
||||
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
|
||||
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
|
||||
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
|
||||
}
|
||||
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
|
||||
|
||||
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
|
||||
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
|
||||
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
|
||||
|
||||
return {
|
||||
args.mode,
|
||||
problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace),
|
||||
hw_info,
|
||||
scheduler
|
||||
};
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE static
|
||||
bool
|
||||
can_implement(Arguments const& args) {
|
||||
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
|
||||
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
|
||||
return implementable;
|
||||
}
|
||||
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
|
||||
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static
|
||||
int
|
||||
get_workspace_size(Arguments const& args) {
|
||||
TileScheduler t;
|
||||
return t.template get_workspace_size<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
}
|
||||
|
||||
static
|
||||
cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
TileScheduler t;
|
||||
return t.template initialize_workspace<ProblemShape, ElementAccumulator>(
|
||||
args.scheduler, workspace, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups);
|
||||
}
|
||||
|
||||
// Computes the kernel launch grid shape based on runtime parameters
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
// Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently
|
||||
TileSchedulerArguments args{};
|
||||
if constexpr (!std::is_const_v<decltype(args.max_swizzle_size)>) {
|
||||
args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
|
||||
}
|
||||
return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_block_shape() {
|
||||
return dim3(MaxThreadsPerBlock, 1, 1);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
operator()(Params const& params, char* smem_buf) {
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
|
||||
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
|
||||
if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) {
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
|
||||
/* In the Cooperative kernel, one or multiple Consumers collaborate on the same tile */
|
||||
enum class WarpGroupRole {
|
||||
Producer = 0,
|
||||
Consumer = 1,
|
||||
};
|
||||
|
||||
// Kernel level shared memory storage
|
||||
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
|
||||
int thread_idx = int(threadIdx.x);
|
||||
int mma_thread_idx = thread_idx % size(TiledMma{});
|
||||
int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
|
||||
int warp_group_idx = canonical_warp_group_idx();
|
||||
CUTLASS_ASSERT(warp_group_idx < NumWarpGroups);
|
||||
WarpGroupRole warp_group_role = warp_group_idx < NumLoadWarpGroups ? WarpGroupRole::Producer : WarpGroupRole::Consumer;
|
||||
|
||||
// Mainloop Load pipeline
|
||||
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
|
||||
typename MainloopPipeline::Params mainloop_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
mainloop_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
mainloop_pipeline_params.consumer_arv_count = NumMmaWarpGroups * NumThreadsPerWarpGroup;
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
|
||||
|
||||
// Epilogue Load pipeline
|
||||
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
|
||||
typename EpiLoadPipeline::Params epi_load_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
epi_load_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
epi_load_pipeline_params.consumer_arv_count = NumMmaWarpGroups * NumThreadsPerWarpGroup;
|
||||
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
|
||||
|
||||
// Epilogue Store pipeline
|
||||
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
|
||||
typename EpiStorePipeline::Params epi_store_pipeline_params;
|
||||
epi_store_pipeline_params.always_wait = true;
|
||||
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
|
||||
|
||||
// Initialize starting pipeline states for the collectives
|
||||
// Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
|
||||
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
|
||||
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
|
||||
|
||||
// For the DMA Load (producer) we start with an opposite phase
|
||||
// i.e., we skip all waits since we know that the buffer is indeed empty
|
||||
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
|
||||
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
|
||||
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA_mkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_A), make_shape(M,K,L), params.mainloop.dA); //(m,k,l)
|
||||
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_B), make_shape(N,K,L), params.mainloop.dB); //(n,k,l)
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
TiledMma tiled_mma;
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
TileScheduler scheduler{params.scheduler};
|
||||
auto work_tile_info = scheduler.get_current_work();
|
||||
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue{params.epilogue, shared_storage.tensors.epilogue};
|
||||
|
||||
// Wait for all threads in the thread block
|
||||
__syncthreads();
|
||||
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with our work tile coordinates to construct mainloop tensor views
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
|
||||
auto work_k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
|
||||
auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info);
|
||||
auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<2>(gA)), shape<2>(gA));
|
||||
|
||||
// Compute tile residues for predication
|
||||
auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord); // M - BLK_M * m_coord
|
||||
auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord); // N - BLK_N * n_coord
|
||||
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
|
||||
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
|
||||
|
||||
collective_mainloop.load(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA,
|
||||
gB,
|
||||
k_tile_iter, work_k_tile_count,
|
||||
residue_mnk,
|
||||
thread_idx,
|
||||
shared_storage.tensors.mainloop
|
||||
);
|
||||
// Update starting pipeline state for the next tile
|
||||
mainloop_pipe_producer_state.advance(work_k_tile_count);
|
||||
|
||||
if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler) &&
|
||||
collective_epilogue.is_producer_load_needed()) {
|
||||
epi_load_pipe_producer_state =
|
||||
collective_epilogue.load(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
tiled_mma,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
}
|
||||
|
||||
// Get next work tile
|
||||
work_tile_info = fetch_next_work(work_tile_info, scheduler);
|
||||
} // Scheduler work fetch loop
|
||||
|
||||
// Make sure all Consumer Warp Groups have been waited upon
|
||||
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
|
||||
|
||||
if (collective_epilogue.is_producer_load_needed()) {
|
||||
collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
|
||||
}
|
||||
} // Producer Warp Group End
|
||||
|
||||
else if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
|
||||
bool do_store_tail = false;
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
auto work_k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
|
||||
|
||||
// Allocate the the accumulators for the (M,N) blk_shape
|
||||
//
|
||||
// MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead.
|
||||
auto accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
collective_mainloop.mma(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
accumulators,
|
||||
work_k_tile_count,
|
||||
mma_thread_idx,
|
||||
shared_storage.tensors.mainloop,
|
||||
params.mainloop
|
||||
);
|
||||
|
||||
// Make sure the math instructions are done and free buffers before entering the epilogue
|
||||
collective_mainloop.mma_tail(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
work_k_tile_count
|
||||
);
|
||||
|
||||
// Update starting mainloop pipeline state for the next tile
|
||||
mainloop_pipe_consumer_state.advance(work_k_tile_count);
|
||||
|
||||
// Index of warp group within consumer warp groups
|
||||
int consumer_warp_group_idx = canonical_warp_group_idx() - NumLoadWarpGroups;
|
||||
|
||||
// Perform reduction across splits, if needed
|
||||
TileScheduler::fixup(
|
||||
params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, consumer_warp_group_idx);
|
||||
|
||||
if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
|
||||
// Epilogue and write to gD
|
||||
auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] =
|
||||
collective_epilogue.store(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline,
|
||||
epi_store_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
accumulators,
|
||||
tiled_mma,
|
||||
mma_thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next;
|
||||
epi_store_pipe_producer_state = epi_store_pipe_producer_state_next;
|
||||
do_store_tail = true;
|
||||
}
|
||||
|
||||
// Get next work tile
|
||||
work_tile_info = fetch_next_work(work_tile_info, scheduler);
|
||||
} // Scheduler work fetch loop
|
||||
|
||||
if (do_store_tail) {
|
||||
collective_epilogue.store_tail(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline,
|
||||
epi_store_pipe_producer_state
|
||||
);
|
||||
}
|
||||
} // Consumer Warp Groups End
|
||||
}
|
||||
|
||||
private:
|
||||
// Kernel helper function to get next work unit
|
||||
CUTLASS_DEVICE
|
||||
typename TileScheduler::WorkTileInfo
|
||||
fetch_next_work(
|
||||
typename TileScheduler::WorkTileInfo& work_tile_info,
|
||||
TileScheduler& scheduler) const {
|
||||
// Check whether we should continue on with the current work unit. If this is the case,
|
||||
// the work unit will have been updated in continue_current_work to reflect the new
|
||||
// tile to be computed.
|
||||
if (scheduler.continue_current_work(work_tile_info)) {
|
||||
return work_tile_info;
|
||||
}
|
||||
|
||||
// Get next work tile
|
||||
scheduler.advance_to_next_work();
|
||||
return scheduler.get_current_work();
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::kernel
|
||||
@@ -0,0 +1,516 @@
|
||||
/***************************************************************************************************
|
||||
* 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.
|
||||
*
|
||||
**************************************************************************************************/
|
||||
#pragma once
|
||||
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/kernel_hardware_info.hpp"
|
||||
#include "cutlass/fast_math.h"
|
||||
#include "cute/arch/cluster_sm90.hpp"
|
||||
#include "cutlass/arch/reg_reconfig.h"
|
||||
#include "cutlass/arch/mma_sm90.h"
|
||||
#include "cutlass/epilogue/collective/detail.hpp"
|
||||
#include "cutlass/gemm/gemm.h"
|
||||
#include "cutlass/gemm/dispatch_policy.hpp"
|
||||
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
|
||||
#include "cutlass/pipeline/pipeline.hpp"
|
||||
#include "cutlass/trace.h"
|
||||
|
||||
#include "cute/tensor.hpp"
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
namespace cutlass::gemm::kernel {
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
template <
|
||||
class ProblemShape_,
|
||||
class CollectiveMainloop_,
|
||||
class CollectiveEpilogue_,
|
||||
class TileScheduler_
|
||||
>
|
||||
class GemmUniversal<
|
||||
ProblemShape_,
|
||||
CollectiveMainloop_,
|
||||
CollectiveEpilogue_,
|
||||
TileScheduler_,
|
||||
cute::enable_if_t<cute::is_base_of_v<KernelCpAsyncWarpSpecializedPingpong, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
|
||||
{
|
||||
public:
|
||||
//
|
||||
// Type Aliases
|
||||
//
|
||||
using ProblemShape = ProblemShape_;
|
||||
static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
|
||||
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
|
||||
|
||||
// Mainloop derived types
|
||||
using CollectiveMainloop = CollectiveMainloop_;
|
||||
using TileShape = typename CollectiveMainloop::TileShape;
|
||||
using TiledMma = typename CollectiveMainloop::TiledMma;
|
||||
using ArchTag = typename CollectiveMainloop::ArchTag;
|
||||
using ElementA = typename CollectiveMainloop::ElementA;
|
||||
using StrideA = typename CollectiveMainloop::StrideA;
|
||||
using ElementB = typename CollectiveMainloop::ElementB;
|
||||
using StrideB = typename CollectiveMainloop::StrideB;
|
||||
using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
|
||||
using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
|
||||
using ClusterShape = typename DispatchPolicy::ClusterShape;
|
||||
using MainloopArguments = typename CollectiveMainloop::Arguments;
|
||||
using MainloopParams = typename CollectiveMainloop::Params;
|
||||
static_assert(ArchTag::kMinComputeCapability >= 90);
|
||||
|
||||
// Epilogue derived types
|
||||
using CollectiveEpilogue = CollectiveEpilogue_;
|
||||
using ElementC = typename CollectiveEpilogue::ElementC;
|
||||
using StrideC = typename CollectiveEpilogue::StrideC;
|
||||
using ElementD = typename CollectiveEpilogue::ElementD;
|
||||
using StrideD = typename CollectiveEpilogue::StrideD;
|
||||
using EpilogueArguments = typename CollectiveEpilogue::Arguments;
|
||||
using EpilogueParams = typename CollectiveEpilogue::Params;
|
||||
|
||||
static_assert(!cute::is_same_v<TileScheduler_, StreamKScheduler>, "Ping-pong kernel does not currently support stream-K scheduler.");
|
||||
using TileSchedulerTag = TileScheduler_;
|
||||
using TileScheduler = typename detail::TileSchedulerSelector<
|
||||
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
|
||||
using TileSchedulerArguments = typename TileScheduler::Arguments;
|
||||
using TileSchedulerParams = typename TileScheduler::Params;
|
||||
|
||||
using GmemTiledCopyA = typename CollectiveMainloop::GmemTiledCopyA;
|
||||
using GmemTiledCopyB = typename CollectiveMainloop::GmemTiledCopyB;
|
||||
static_assert(cute::size(GmemTiledCopyA{}) == cute::size(GmemTiledCopyB{}), "Number of threads in A/B tiled copies must be the same");
|
||||
|
||||
static constexpr uint32_t NumLoadWarpGroups = cute::size(GmemTiledCopyA{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumMmaWarpGroups = 2 * cute::size(TiledMma{}) / NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t NumWarpGroups = NumLoadWarpGroups + NumMmaWarpGroups;
|
||||
static_assert(NumWarpGroups == 2 || NumWarpGroups == 3, "Number of warp groups must be 2 or 3 for good performance.");
|
||||
static_assert(NumMmaWarpGroups == 2, "Pingpong kernel requires 2 MMA warp groups.");
|
||||
|
||||
static constexpr uint32_t MaxThreadsPerBlock = NumWarpGroups * NumThreadsPerWarpGroup;
|
||||
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
||||
|
||||
// Order Sequence barrier with two stages: one for Mainloop and one for Epilogue
|
||||
static constexpr uint32_t StagesPerMathWarpGroup = 2;
|
||||
using MathWarpGroupOrderBarrier = cutlass::OrderedSequenceBarrier<
|
||||
StagesPerMathWarpGroup, NumMmaWarpGroups>;
|
||||
|
||||
// Kernel level shared memory storage
|
||||
struct SharedStorage {
|
||||
struct TensorStorage : cute::aligned_struct<128> {
|
||||
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
|
||||
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
|
||||
|
||||
MainloopTensorStorage mainloop;
|
||||
EpilogueTensorStorage epilogue;
|
||||
} tensors;
|
||||
|
||||
struct PipelineStorage : cute::aligned_struct<16> {
|
||||
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
|
||||
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
|
||||
using MathWarpGroupOrderBarrierStorage = typename MathWarpGroupOrderBarrier::SharedStorage;
|
||||
|
||||
alignas(16) MainloopPipelineStorage mainloop;
|
||||
alignas(16) EpiLoadPipelineStorage epi_load;
|
||||
alignas(16) MathWarpGroupOrderBarrierStorage math_wg_order;
|
||||
} pipelines;
|
||||
};
|
||||
|
||||
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
||||
|
||||
// Device side arguments
|
||||
struct Arguments {
|
||||
GemmUniversalMode mode{};
|
||||
ProblemShape problem_shape{};
|
||||
MainloopArguments mainloop{};
|
||||
EpilogueArguments epilogue{};
|
||||
KernelHardwareInfo hw_info{};
|
||||
TileSchedulerArguments scheduler{};
|
||||
};
|
||||
|
||||
// Kernel entry point API
|
||||
struct Params {
|
||||
GemmUniversalMode mode;
|
||||
ProblemShape problem_shape;
|
||||
MainloopParams mainloop;
|
||||
EpilogueParams epilogue;
|
||||
KernelHardwareInfo hw_info;
|
||||
TileSchedulerParams scheduler;
|
||||
};
|
||||
|
||||
//
|
||||
// Methods
|
||||
//
|
||||
|
||||
// Convert to underlying arguments. In this case, a simple copy for the aliased type.
|
||||
static
|
||||
Params
|
||||
to_underlying_arguments(Arguments const& args, void* workspace) {
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments():");
|
||||
|
||||
(void) workspace;
|
||||
auto problem_shape = args.problem_shape;
|
||||
if constexpr (detail::IF_SWAP_AB<CollectiveMainloop>::value) {
|
||||
// swap M/N
|
||||
get<0>(problem_shape) = get<1>(args.problem_shape);
|
||||
get<1>(problem_shape) = get<0>(args.problem_shape);
|
||||
}
|
||||
auto problem_shape_MNKL = append<4>(problem_shape, 1);
|
||||
|
||||
// Get SM count if needed, otherwise use user supplied SM count
|
||||
int sm_count = args.hw_info.sm_count;
|
||||
if (sm_count <= 0) {
|
||||
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
|
||||
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
|
||||
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
|
||||
}
|
||||
|
||||
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
|
||||
|
||||
KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
|
||||
TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
|
||||
problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, workspace);
|
||||
|
||||
return {
|
||||
args.mode,
|
||||
problem_shape,
|
||||
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
|
||||
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace),
|
||||
hw_info,
|
||||
scheduler
|
||||
};
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE static
|
||||
bool
|
||||
can_implement(Arguments const& args) {
|
||||
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
|
||||
(args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
|
||||
if (!implementable) {
|
||||
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
|
||||
return implementable;
|
||||
}
|
||||
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
|
||||
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
|
||||
return implementable;
|
||||
}
|
||||
|
||||
static
|
||||
int
|
||||
get_workspace_size(Arguments const& args) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static
|
||||
cutlass::Status
|
||||
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
|
||||
return Status::kSuccess;
|
||||
}
|
||||
|
||||
// Computes the kernel launch grid shape based on runtime parameters
|
||||
static dim3
|
||||
get_grid_shape(Params const& params) {
|
||||
// Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently
|
||||
TileSchedulerArguments args{};
|
||||
if constexpr (!std::is_const_v<decltype(args.max_swizzle_size)>) {
|
||||
args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
|
||||
}
|
||||
return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
|
||||
}
|
||||
|
||||
static dim3
|
||||
get_block_shape() {
|
||||
return dim3(MaxThreadsPerBlock, 1, 1);
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
operator()(Params const& params, char* smem_buf) {
|
||||
using namespace cute;
|
||||
using X = Underscore;
|
||||
|
||||
// Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
|
||||
#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
|
||||
if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) {
|
||||
printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
// Preconditions
|
||||
static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
|
||||
|
||||
enum class WarpGroupRole {
|
||||
Producer = 0,
|
||||
Consumer = 1,
|
||||
};
|
||||
|
||||
// Kernel level shared memory storage
|
||||
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
||||
|
||||
int thread_idx = int(threadIdx.x);
|
||||
int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
|
||||
int warp_group_idx = canonical_warp_group_idx();
|
||||
CUTLASS_ASSERT(warp_group_idx < NumWarpGroups);
|
||||
WarpGroupRole warp_group_role = warp_group_idx < NumLoadWarpGroups ? WarpGroupRole::Producer : WarpGroupRole::Consumer;
|
||||
int warp_group_consumer_idx = warp_group_idx - NumLoadWarpGroups;
|
||||
|
||||
// Mainloop Load pipeline
|
||||
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
|
||||
typename MainloopPipeline::Params mainloop_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
mainloop_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
mainloop_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; // only 1 WG consumes at a time
|
||||
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
|
||||
|
||||
// Epilogue Load pipeline
|
||||
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
|
||||
typename EpiLoadPipeline::Params epi_load_pipeline_params;
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
|
||||
}
|
||||
if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
|
||||
}
|
||||
epi_load_pipeline_params.producer_arv_count = NumLoadWarpGroups * NumThreadsPerWarpGroup;
|
||||
epi_load_pipeline_params.consumer_arv_count = NumThreadsPerWarpGroup; // only 1 WG consumes at a time
|
||||
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
|
||||
|
||||
// Epilogue Store pipeline
|
||||
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
|
||||
typename EpiStorePipeline::Params epi_store_pipeline_params;
|
||||
epi_store_pipeline_params.always_wait = true;
|
||||
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
|
||||
|
||||
typename MathWarpGroupOrderBarrier::Params params_math_wg_order_barrier;
|
||||
// DMA Load WG will not participate in these Ordered Barrier syncs
|
||||
params_math_wg_order_barrier.group_id = warp_group_consumer_idx;
|
||||
params_math_wg_order_barrier.group_size = NumThreadsPerWarpGroup; // Number of threads / participants in a group
|
||||
MathWarpGroupOrderBarrier math_wg_order_barrier(shared_storage.pipelines.math_wg_order, params_math_wg_order_barrier);
|
||||
|
||||
// Initialize starting pipeline states for the collectives
|
||||
// Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
|
||||
typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
|
||||
typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
|
||||
|
||||
// For the DMA Load (producer) we start with an opposite phase
|
||||
// i.e., we skip all waits since we know that the buffer is indeed empty
|
||||
PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
|
||||
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
|
||||
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
|
||||
|
||||
// Separate out problem shape for convenience
|
||||
// Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
|
||||
auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
|
||||
auto M = get<0>(problem_shape_MNKL);
|
||||
auto N = get<1>(problem_shape_MNKL);
|
||||
auto K = get<2>(problem_shape_MNKL);
|
||||
auto L = get<3>(problem_shape_MNKL);
|
||||
|
||||
// Represent the full tensors
|
||||
Tensor mA_mkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_A), make_shape(M,K,L), params.mainloop.dA); //(m,k,l)
|
||||
Tensor mB_nkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_B), make_shape(N,K,L), params.mainloop.dB); //(n,k,l)
|
||||
|
||||
// Get the appropriate blocks for this thread block -- potential for thread block locality
|
||||
TiledMma tiled_mma;
|
||||
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
|
||||
|
||||
// Make tiled views, defer the slice
|
||||
Tensor gA_mkl = local_tile(mA_mkl, blk_shape, make_coord(_,_,_), Step<_1, X,_1>{}); // (BLK_M,BLK_K,m,k,l)
|
||||
Tensor gB_nkl = local_tile(mB_nkl, blk_shape, make_coord(_,_,_), Step< X,_1,_1>{}); // (BLK_N,BLK_K,n,k,l)
|
||||
|
||||
// Get pipeline stage increments from tensor shapes
|
||||
auto k_tile_count = size<3>(gA_mkl);
|
||||
auto c_tile_count = CollectiveEpilogue::get_load_pipe_increment(blk_shape);
|
||||
auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape);
|
||||
|
||||
TileScheduler scheduler{params.scheduler};
|
||||
|
||||
if (warp_group_consumer_idx == 1) {
|
||||
// Advance 2nd Math WG to the next work tile for the startup
|
||||
scheduler.advance_to_next_work();
|
||||
// Advance 2nd Math WG pipeline states to the end of 1st Math WG
|
||||
mainloop_pipe_consumer_state.advance(k_tile_count);
|
||||
epi_load_pipe_consumer_state.advance(c_tile_count);
|
||||
epi_store_pipe_producer_state.advance(d_tile_count);
|
||||
}
|
||||
auto work_tile_info = scheduler.get_current_work();
|
||||
|
||||
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
||||
CollectiveMainloop collective_mainloop;
|
||||
CollectiveEpilogue collective_epilogue{params.epilogue, shared_storage.tensors.epilogue};
|
||||
|
||||
// Wait for all threads in the thread block
|
||||
__syncthreads();
|
||||
|
||||
if (warp_group_role == WarpGroupRole::Producer) {
|
||||
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Slice with our work tile coordinates to construct mainloop tensor views
|
||||
Tensor gA = gA_mkl(_,_,m_coord,_,l_coord); // (BLK_M,BLK_K,k)
|
||||
Tensor gB = gB_nkl(_,_,n_coord,_,l_coord); // (BLK_N,BLK_K,k)
|
||||
|
||||
auto k_tile_iter = cute::make_coord_iterator(shape<2>(gA));
|
||||
|
||||
// Compute tile residues for predication
|
||||
auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord); // M - BLK_M * m_coord
|
||||
auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord); // N - BLK_N * n_coord
|
||||
auto k_residue = K - size<1>(gA) * size<2>(gA); // K - BLK_K * k_coord_max
|
||||
auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
|
||||
|
||||
collective_mainloop.load(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_producer_state,
|
||||
gA,
|
||||
gB,
|
||||
k_tile_iter, k_tile_count,
|
||||
residue_mnk,
|
||||
thread_idx,
|
||||
shared_storage.tensors.mainloop
|
||||
);
|
||||
// Update starting pipeline state for the next tile
|
||||
mainloop_pipe_producer_state.advance(k_tile_count);
|
||||
|
||||
if (collective_epilogue.is_producer_load_needed()) {
|
||||
collective_epilogue.load(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
tiled_mma,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
// Update starting pipeline state for the next tile
|
||||
epi_load_pipe_producer_state.advance(c_tile_count);
|
||||
}
|
||||
|
||||
// Get next work tile
|
||||
scheduler.advance_to_next_work();
|
||||
work_tile_info = scheduler.get_current_work();
|
||||
} // Scheduler work fetch loop
|
||||
|
||||
// Make sure all Consumer Warp Groups have been waited upon
|
||||
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
|
||||
if (collective_epilogue.is_producer_load_needed()) {
|
||||
collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
|
||||
}
|
||||
} // Producer Warp Group End
|
||||
|
||||
else if (warp_group_role == WarpGroupRole::Consumer) {
|
||||
|
||||
while (work_tile_info.is_valid()) {
|
||||
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
|
||||
auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
|
||||
auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
|
||||
auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
|
||||
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
|
||||
|
||||
// Allocate the the accumulators for the (M,N) blk_shape
|
||||
Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
|
||||
|
||||
// Order two Math WG's MMA one after the other, helps hide Epilogue
|
||||
math_wg_order_barrier.wait();
|
||||
|
||||
collective_mainloop.mma(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
accumulators,
|
||||
k_tile_count,
|
||||
thread_idx,
|
||||
shared_storage.tensors.mainloop,
|
||||
params.mainloop
|
||||
);
|
||||
|
||||
// Cue for next Math WG's MMA to start
|
||||
math_wg_order_barrier.arrive();
|
||||
|
||||
// Make sure the math instructions are done and free buffers before entering the epilogue
|
||||
collective_mainloop.mma_tail(
|
||||
mainloop_pipeline,
|
||||
mainloop_pipe_consumer_state,
|
||||
k_tile_count
|
||||
);
|
||||
// Update starting mainloop pipeline state for the next tile
|
||||
mainloop_pipe_consumer_state.advance(k_tile_count * NumMmaWarpGroups);
|
||||
|
||||
// Order two Math WG's Epilogue one after the other
|
||||
math_wg_order_barrier.wait();
|
||||
|
||||
// Epilogue and write to gD
|
||||
collective_epilogue.store(
|
||||
epi_load_pipeline,
|
||||
epi_load_pipe_consumer_state,
|
||||
epi_store_pipeline,
|
||||
epi_store_pipe_producer_state,
|
||||
problem_shape_MNKL,
|
||||
blk_shape,
|
||||
blk_coord,
|
||||
accumulators,
|
||||
tiled_mma,
|
||||
warp_group_thread_idx,
|
||||
shared_storage.tensors.epilogue
|
||||
);
|
||||
// Update starting load/store pipeline states for the next tile
|
||||
epi_load_pipe_consumer_state.advance(c_tile_count * NumMmaWarpGroups);
|
||||
epi_store_pipe_producer_state.advance(d_tile_count * NumMmaWarpGroups);
|
||||
|
||||
// Wait for all TMA stores to complete
|
||||
epi_store_pipeline.producer_tail(epi_store_pipe_producer_state);
|
||||
|
||||
// Cue for next Math WG's Epilogue to start
|
||||
math_wg_order_barrier.arrive();
|
||||
|
||||
// Get next work tile
|
||||
scheduler.advance_to_next_work(NumMmaWarpGroups);
|
||||
work_tile_info = scheduler.get_current_work();
|
||||
} // Scheduler work fetch loop
|
||||
} // Consumer Warp Groups End
|
||||
}
|
||||
};
|
||||
|
||||
///////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace cutlass::gemm::kernel
|
||||
@@ -50,6 +50,7 @@ class PersistentTileSchedulerSm90 {
|
||||
|
||||
private:
|
||||
uint64_t current_work_linear_idx_;
|
||||
uint64_t total_grid_size_;
|
||||
|
||||
public:
|
||||
struct WorkTileInfo {
|
||||
@@ -57,12 +58,29 @@ public:
|
||||
int32_t N_idx = 0;
|
||||
int32_t L_idx = 0;
|
||||
bool is_valid_tile = false;
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool
|
||||
is_valid() const {
|
||||
return is_valid_tile;
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
static WorkTileInfo
|
||||
invalid_work_tile() {
|
||||
return {-1, -1, -1, false};
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool
|
||||
is_final_split(uint32_t k_tiles_per_output_tile) const {
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
using Params = PersistentTileSchedulerSm90Params;
|
||||
using RasterOrder = typename Params::RasterOrder;
|
||||
using RasterOrderOptions = typename Params::RasterOrderOptions;
|
||||
|
||||
struct Arguments {
|
||||
int max_swizzle_size = 1;
|
||||
RasterOrderOptions raster_order = RasterOrderOptions::Heuristic;
|
||||
@@ -116,6 +134,8 @@ public:
|
||||
else {
|
||||
current_work_linear_idx_ = uint64_t(blockIdx.x) * uint64_t(gridDim.y) + uint64_t(blockIdx.y);
|
||||
}
|
||||
|
||||
total_grid_size_ = uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z);
|
||||
#else
|
||||
CUTLASS_ASSERT(false && "This line should never be reached");
|
||||
#endif
|
||||
@@ -130,6 +150,10 @@ public:
|
||||
CUTLASS_DEVICE
|
||||
WorkTileInfo
|
||||
get_current_work_for_linear_idx(uint64_t linear_idx) const {
|
||||
if (linear_idx >= scheduler_params.blocks_per_problem_) {
|
||||
return WorkTileInfo::invalid_work_tile();
|
||||
}
|
||||
|
||||
// Map worker's linear index into the CTA tiled problem shape to the corresponding MNL indices
|
||||
uint64_t work_idx_l, remainder;
|
||||
scheduler_params.divmod_batch_(work_idx_l, remainder, linear_idx);
|
||||
@@ -143,19 +167,13 @@ public:
|
||||
scheduler_params.log_swizzle_size_,
|
||||
scheduler_params.raster_order_);
|
||||
|
||||
return {work_idx_m, work_idx_n, static_cast<int32_t>(work_idx_l), linear_idx < scheduler_params.blocks_per_problem_};
|
||||
return {work_idx_m, work_idx_n, static_cast<int32_t>(work_idx_l), true};
|
||||
}
|
||||
|
||||
CUTLASS_DEVICE
|
||||
void
|
||||
advance_to_next_work(uint32_t advance_count = 1) {
|
||||
// MSVC requires protecting use of CUDA-specific nonstandard syntax,
|
||||
// like blockIdx and gridDim, with __CUDA_ARCH__.
|
||||
#if defined(__CUDA_ARCH__)
|
||||
current_work_linear_idx_ += uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z) * uint64_t(advance_count);
|
||||
#else
|
||||
CUTLASS_ASSERT(false && "This line should never be reached");
|
||||
#endif
|
||||
current_work_linear_idx_ += total_grid_size_ * uint64_t(advance_count);
|
||||
}
|
||||
|
||||
// get work_idx_m, work_idx_n from blk_per_grid_dim while applying swizzle
|
||||
@@ -163,8 +181,8 @@ public:
|
||||
cute::tuple<int32_t, int32_t>
|
||||
get_work_idx_m_and_n(
|
||||
uint64_t blk_per_grid_dim,
|
||||
FastDivmodU64 const& divmod_cluster_shape_major,
|
||||
FastDivmodU64 const& divmod_cluster_shape_minor,
|
||||
FastDivmodU64Pow2 const& divmod_cluster_shape_major,
|
||||
FastDivmodU64Pow2 const& divmod_cluster_shape_minor,
|
||||
FastDivmodU64 const& divmod_cluster_blk_major,
|
||||
int32_t log_swizzle_size,
|
||||
RasterOrder raster_order) {
|
||||
@@ -205,6 +223,46 @@ public:
|
||||
|
||||
}
|
||||
|
||||
// Computes the linear index within a batch given M and N tile offsets within the batch.
|
||||
// This essentially inverts the mapping performed in get_work_idx_m_and_n
|
||||
static CUTLASS_DEVICE
|
||||
uint64_t
|
||||
get_linear_idx_from_m_and_n(
|
||||
int32_t tile_m,
|
||||
int32_t tile_n,
|
||||
FastDivmodU64Pow2 const& divmod_cluster_shape_major,
|
||||
FastDivmodU64Pow2 const& divmod_cluster_shape_minor,
|
||||
FastDivmodU64 const& divmod_cluster_blk_major,
|
||||
int32_t log_swizzle_size,
|
||||
RasterOrder raster_order) {
|
||||
|
||||
auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
|
||||
|
||||
uint64_t minor_work_idx, major_work_idx, cluster_minor_offset;
|
||||
if (raster_order == RasterOrder::AlongN) {
|
||||
minor_work_idx = static_cast<uint64_t>(tile_m);
|
||||
major_work_idx = static_cast<uint64_t>(tile_n);
|
||||
cluster_minor_offset = cta_m_in_cluster;
|
||||
}
|
||||
else {
|
||||
major_work_idx = static_cast<uint64_t>(tile_m);
|
||||
minor_work_idx = static_cast<uint64_t>(tile_n);
|
||||
cluster_minor_offset = cta_n_in_cluster;
|
||||
}
|
||||
|
||||
uint64_t cluster_idx_minor, cluster_idx_major, cluster_major_offset;
|
||||
cluster_idx_minor = divmod_cluster_shape_minor.divide(minor_work_idx - cluster_minor_offset);
|
||||
divmod_cluster_shape_major(cluster_idx_major, cluster_major_offset, major_work_idx);
|
||||
|
||||
uint64_t cluster_idx_minor_div_swizzle = cluster_idx_minor >> log_swizzle_size;
|
||||
uint64_t offset = cluster_idx_minor & ((1 << log_swizzle_size) - 1);
|
||||
|
||||
uint64_t extra = cluster_idx_minor_div_swizzle * divmod_cluster_blk_major.divisor + cluster_idx_major;
|
||||
|
||||
uint64_t cluster_id = (extra << log_swizzle_size) | offset;
|
||||
return (cluster_id * divmod_cluster_shape_major.divisor + cluster_major_offset) * divmod_cluster_shape_minor.divisor + cluster_minor_offset;
|
||||
}
|
||||
|
||||
// Given the inputs, computes the total number of output blocks this problem will compute over
|
||||
// Note that this is only the logical size of our grid, not the physical grid we will actually launch.
|
||||
template<class ProblemShapeMNKL, class BlockShape, class ClusterShape>
|
||||
@@ -250,7 +308,7 @@ public:
|
||||
// output tile. For the basic tile scheduler, this is always true.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static bool
|
||||
compute_epilogue(WorkTileInfo const&) {
|
||||
compute_epilogue(WorkTileInfo const&, Params const&) {
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
@@ -64,33 +64,44 @@ public:
|
||||
|
||||
using RasterOrder = UnderlyingScheduler::RasterOrder;
|
||||
using RasterOrderOptions = UnderlyingScheduler::RasterOrderOptions;
|
||||
|
||||
// Use a dummy barrier manager to simply get the type used to store the barrier
|
||||
using BarrierType = typename NamedBarrierManager<1>::T;
|
||||
|
||||
using Params = PersistentTileSchedulerSm90StreamKParams;
|
||||
using ReductionMode = Params::ReductionMode;
|
||||
|
||||
struct WorkTileInfo {
|
||||
int32_t M_idx = 0;
|
||||
int32_t N_idx = 0;
|
||||
int32_t K_idx = 0;
|
||||
int32_t L_idx = 0;
|
||||
bool is_valid_tile = false;
|
||||
|
||||
// Number of splits to be used in computing the {L_idx, M_idx, N_idx} output tile.
|
||||
// Splits = 1 indicates that this is a data-parallel block.
|
||||
uint32_t splits = 1;
|
||||
|
||||
// Number of k iterations to compute for the current tile
|
||||
// Number of k tiles to compute for this unit of work. For stream-K, this
|
||||
// can indicate the number of K tiles across multiple output tiles.
|
||||
uint32_t k_tile_count = 0;
|
||||
|
||||
// Number of k iterations remaining for the work unit as a whole
|
||||
// Number of k tiles remaining for the work unit as a whole
|
||||
uint32_t k_tile_remaining = 0;
|
||||
|
||||
// Whether this unit of work is the final split for the given tile
|
||||
bool is_final_split = true;
|
||||
};
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool
|
||||
is_valid() const {
|
||||
// Use negative indices to denote invalid work
|
||||
return M_idx >= 0;
|
||||
}
|
||||
|
||||
using Params = PersistentTileSchedulerSm90StreamKParams;
|
||||
using ReductionMode = Params::ReductionMode;
|
||||
CUTLASS_HOST_DEVICE
|
||||
static WorkTileInfo
|
||||
invalid_work_tile() {
|
||||
return {-1, -1, -1, -1, 0};
|
||||
}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
bool
|
||||
is_final_split(uint32_t k_tiles_per_output_tile) const {
|
||||
return (K_idx + k_tile_count) == k_tiles_per_output_tile;
|
||||
}
|
||||
};
|
||||
|
||||
struct Arguments {
|
||||
|
||||
@@ -117,6 +128,12 @@ public:
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments(int splits_) : splits(splits_) {}
|
||||
|
||||
CUTLASS_HOST_DEVICE
|
||||
Arguments(int splits_, int max_swizzle_size_, RasterOrderOptions raster_order_) :
|
||||
splits(splits_),
|
||||
max_swizzle_size(max_swizzle_size_),
|
||||
raster_order(raster_order_) {}
|
||||
|
||||
// The splitting factor to be used in a split-K decomposition of the problem.
|
||||
// If this is set to a value greater than 1, stream-K decomposition logic
|
||||
// is bypassed in favor of a split-K decomposition.
|
||||
@@ -187,26 +204,20 @@ public:
|
||||
CUTLASS_DEVICE
|
||||
static WorkTileInfo
|
||||
get_current_work_for_linear_idx(uint64_t linear_idx, Params const& params) {
|
||||
if (linear_idx >= params.units_per_problem_) {
|
||||
// The maximum number of work units is units_per_problem_ * splits_.
|
||||
// The multiplication by splits_ is used for handling split-K, in which
|
||||
// units_per_problem_ is equal to the total number of output tiles. To account
|
||||
// for the fact that we have splits_ peers per output tile, we multiply this
|
||||
// value by splits_. For stream-K, this multiplication ends up being a no-op
|
||||
// because splits_ is set to 1 for stream-K.
|
||||
if (linear_idx >= params.units_per_problem_ * params.splits_) {
|
||||
// Invalid work. Return an empty result.
|
||||
return {0, 0, 0, 0, false, 0};
|
||||
return WorkTileInfo::invalid_work_tile();
|
||||
}
|
||||
|
||||
// Determine whether this work unit is a data-parallel or stream-K work unit
|
||||
bool is_stream_k_unit = linear_idx < params.sk_units_;
|
||||
|
||||
bool is_split_k = params.splits_ > 1;
|
||||
|
||||
if (is_split_k || !is_stream_k_unit) {
|
||||
// Bypass the stream-K scheduling logic for basic data-parallel or split-K work
|
||||
return set_non_stream_k_work(linear_idx, params, is_split_k);
|
||||
}
|
||||
else {
|
||||
// This is a stream-K work unit
|
||||
WorkTileInfo work_tile_info;
|
||||
set_stream_k_work(params, linear_idx, work_tile_info, /*new_unit = */ true);
|
||||
return work_tile_info;
|
||||
}
|
||||
WorkTileInfo work_tile_info;
|
||||
assign_work(params, linear_idx, work_tile_info);
|
||||
return work_tile_info;
|
||||
}
|
||||
|
||||
// Returns whether the current work_tile_info passed in should continue to be used. This
|
||||
@@ -233,7 +244,7 @@ public:
|
||||
return false;
|
||||
}
|
||||
|
||||
set_stream_k_work(params, linear_idx, work_tile_info, /* new_unit = */ false);
|
||||
assign_work(params, linear_idx, work_tile_info);
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -280,7 +291,7 @@ public:
|
||||
static bool
|
||||
requires_fixup(Params const& params, WorkTileInfo const& work_tile_info) {
|
||||
// Fixup is not needed for data-parallel tiles
|
||||
return work_tile_info.k_tile_count != params.k_tiles_per_output_tile_;
|
||||
return work_tile_info.k_tile_count != params.divmod_tiles_per_output_tile_.divisor;
|
||||
}
|
||||
|
||||
// Performs the reduction across splits for a given output tile.
|
||||
@@ -293,7 +304,9 @@ public:
|
||||
FrgTensorC& accumulators,
|
||||
uint32_t num_barriers,
|
||||
uint32_t barrier_idx) {
|
||||
using BarrierManager = NamedBarrierManager<NumThreadsPerWarpGroup, 2>;
|
||||
static constexpr uint32_t Offset = 2;
|
||||
static constexpr uint32_t MaxNumNamedBarriers = 2;
|
||||
using BarrierManager = NamedBarrierManager<NumThreadsPerWarpGroup, Offset, MaxNumNamedBarriers>;
|
||||
return fixup_helper<FrgTensorC, BarrierManager>(
|
||||
params, work_tile_info, accumulators, num_barriers, barrier_idx);
|
||||
}
|
||||
@@ -331,24 +344,23 @@ public:
|
||||
using AccumulatorArrayT = Array<typename FrgTensorC::value_type, size(FrgTensorC{})>;
|
||||
using BlockStripedReduceT = BlockStripedReduce<BarrierManager::ThreadCount, AccumulatorArrayT>;
|
||||
|
||||
AccumulatorArrayT* reduction_workspace_array = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace);
|
||||
AccumulatorArrayT* accumulator_array = reinterpret_cast<AccumulatorArrayT*>(&accumulators);
|
||||
|
||||
int barrier_group_thread_idx = threadIdx.x % BarrierManager::ThreadCount;
|
||||
|
||||
// The number of tiles for which reduction is required is either:
|
||||
// (a) the total number of output tiles (in the case of split-K)
|
||||
// (b) the number of stream-K tiles
|
||||
// To calcualte the the total number of output tiles in the split-K case, we
|
||||
// To calculate the the total number of output tiles in the split-K case, we
|
||||
// note that, in the split-K case, the units_per_problem_ member of Params will be
|
||||
// the total number of output tiles multiplied by the number of splits.
|
||||
auto reduction_tiles = params.splits_ > 1 ? (params.units_per_problem_ / params.splits_) : params.sk_tiles_;
|
||||
// the total number of output tiles.
|
||||
auto reduction_tiles = params.splits_ > 1 ? params.units_per_problem_ : params.sk_tiles_;
|
||||
auto reduction_workspace_size = Params::get_reduction_workspace_size(
|
||||
reduction_tiles, to_gemm_coord(TileShape{}), sizeof_bits<ElementAccumulator>::value);
|
||||
BarrierType* lock_workspace = reinterpret_cast<BarrierType*>(
|
||||
reinterpret_cast<uint8_t*>(params.reduction_workspace_) + reduction_workspace_size);
|
||||
|
||||
if (!work_tile_info.is_final_split) {
|
||||
AccumulatorArrayT* reduction_workspace_array = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace);
|
||||
AccumulatorArrayT* accumulator_array = reinterpret_cast<AccumulatorArrayT*>(&accumulators);
|
||||
int barrier_group_thread_idx = threadIdx.x % BarrierManager::ThreadCount;
|
||||
|
||||
if (!work_tile_info.is_final_split(params.divmod_tiles_per_output_tile_.divisor)) {
|
||||
if (work_tile_info.K_idx == 0) {
|
||||
// First peer initializes the workspace partials
|
||||
BlockStripedReduceT::store(reduction_workspace_array, *accumulator_array, barrier_group_thread_idx);
|
||||
@@ -359,7 +371,12 @@ public:
|
||||
BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
|
||||
}
|
||||
else {
|
||||
// Wait unitl the first split has stored its accumulators
|
||||
// Wait until the first split has stored its accumulators. Note that the first split will have
|
||||
// accumulated a value into the lock potentially greater than one (since the locked value is
|
||||
// incremented by work_tile_info.k_tile_count below for both the deterministic and non-deterministic)
|
||||
// cases. For non-deterministic reductions, all that non-first or last splits care about is whether
|
||||
// the first split has been written, so we only wait while the locked value is less than 1. This
|
||||
// avoids having to add logic to determine the work_tile_info.k_tile_count for the first split.
|
||||
BarrierManager::wait_lt(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, 1);
|
||||
}
|
||||
|
||||
@@ -371,7 +388,11 @@ public:
|
||||
BarrierManager::arrive_inc(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.k_tile_count);
|
||||
}
|
||||
else {
|
||||
// Wait until the preceding split added its accumulators
|
||||
// Wait until the preceding split added its accumulators.
|
||||
// For both the deterministic and non-deterministic case, each preceding split will have incremented
|
||||
// the locked value by work_tile_info.k_tile_count. Thus, the final split konws that it can begin
|
||||
// loading the partially-reduced value when the locked value reaches its starting K tile index (i.e.,
|
||||
// work_tile_info.K_idx).
|
||||
BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
|
||||
|
||||
// The block computing the final split for the tile adds previously-reduced partials
|
||||
@@ -384,54 +405,25 @@ public:
|
||||
// output tile. For the case of stream-K, this should only occur if the work is marked as the final split.
|
||||
CUTLASS_HOST_DEVICE
|
||||
static bool
|
||||
compute_epilogue(WorkTileInfo const& work_tile_info) {
|
||||
return work_tile_info.is_final_split;
|
||||
compute_epilogue(WorkTileInfo const& work_tile_info, Params const& params) {
|
||||
return work_tile_info.is_final_split(params.divmod_tiles_per_output_tile_.divisor);
|
||||
}
|
||||
|
||||
// Returns the linearized index of the output tile corresponding to the tile with offset [L, M, K]
|
||||
CUTLASS_DEVICE
|
||||
static int
|
||||
output_tile_index(Params const& params, WorkTileInfo const& work_tile_info) {
|
||||
if (params.splits_ > 1) {
|
||||
auto tiles_mn = params.divmod_batch_.divisor / params.splits_;
|
||||
if (params.raster_order_ == RasterOrder::AlongN) {
|
||||
return
|
||||
(tiles_mn * work_tile_info.L_idx) +
|
||||
(params.divmod_cluster_shape_major_.divisor *
|
||||
params.divmod_cluster_blk_major_.divisor * work_tile_info.M_idx) +
|
||||
work_tile_info.N_idx;
|
||||
}
|
||||
else {
|
||||
return
|
||||
(tiles_mn * work_tile_info.L_idx) +
|
||||
(params.divmod_cluster_shape_major_.divisor *
|
||||
params.divmod_cluster_blk_major_.divisor * work_tile_info.N_idx) +
|
||||
work_tile_info.M_idx;
|
||||
}
|
||||
}
|
||||
else {
|
||||
auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
|
||||
uint64_t linear_idx_in_batch = UnderlyingScheduler::get_linear_idx_from_m_and_n(
|
||||
work_tile_info.M_idx, work_tile_info.N_idx,
|
||||
params.divmod_cluster_shape_major_,
|
||||
params.divmod_cluster_shape_minor_,
|
||||
params.divmod_cluster_blk_major_,
|
||||
params.log_swizzle_size_,
|
||||
params.raster_order_
|
||||
);
|
||||
|
||||
uint64_t cta_per_grid_dim;
|
||||
uint64_t cluster_dim_idx;
|
||||
if (params.raster_order_ == RasterOrder::AlongN) {
|
||||
uint64_t block_idx_m = (work_tile_info.M_idx - cta_m_in_cluster) / params.divmod_cluster_shape_minor_.divisor;
|
||||
uint64_t block_idx_n = work_tile_info.N_idx;
|
||||
cta_per_grid_dim = (params.divmod_cluster_shape_major_.divisor *
|
||||
params.divmod_cluster_blk_major_.divisor * block_idx_m) + block_idx_n;
|
||||
cluster_dim_idx = cta_m_in_cluster;
|
||||
}
|
||||
else {
|
||||
uint64_t block_idx_m = work_tile_info.M_idx;
|
||||
uint64_t block_idx_n = (work_tile_info.N_idx - cta_n_in_cluster) / params.divmod_cluster_shape_minor_.divisor;
|
||||
cta_per_grid_dim = (params.divmod_cluster_shape_major_.divisor *
|
||||
params.divmod_cluster_blk_major_.divisor * block_idx_n) + block_idx_m;
|
||||
cluster_dim_idx = cta_n_in_cluster;
|
||||
}
|
||||
|
||||
uint64_t tile_in_batch = params.divmod_cluster_shape_minor_.divisor * cta_per_grid_dim;
|
||||
return params.divmod_batch_.divisor * work_tile_info.L_idx + tile_in_batch + cluster_dim_idx;
|
||||
}
|
||||
uint64_t tiles_mn = params.divmod_batch_.divisor;
|
||||
return tiles_mn * work_tile_info.L_idx + linear_idx_in_batch;
|
||||
}
|
||||
|
||||
template <class ProblemShape, class ElementAccumulator>
|
||||
@@ -518,106 +510,125 @@ public:
|
||||
// iterations) is used to find the next tile in the current work unit.
|
||||
CUTLASS_DEVICE
|
||||
static void
|
||||
set_stream_k_work(
|
||||
assign_work(
|
||||
Params const& params,
|
||||
uint64_t linear_idx,
|
||||
WorkTileInfo& work_tile_info,
|
||||
bool new_unit) {
|
||||
// In the CUTLASS 2.x implementation of stream K, stream-K work is assigned to each stream-K
|
||||
// threadblock individually. For the most part, the set of K iterations corresponding to stream-K
|
||||
// work was divided amongst stream-K threadblocks, and a threadblock determined which tile
|
||||
// it would compute a (potentially-partial) output tile for based on the space of k iterations
|
||||
// assigned to it. This often results in stream-K threadblocks processing tiles with different
|
||||
// offsets in the K dimension from one another. This can reduce locality, but is lmitied to the
|
||||
// (generally few) waves of threadblocks assigned to compute stream-K work.
|
||||
//
|
||||
// With the introduction of threadblock clusters, there is additional benefit to maintaining
|
||||
// locality in the K dimension: shared portions of operands can be multicasted to threadblocks
|
||||
// within a cluster. Thus, we would like to ensure that the assignment of stream-K work to
|
||||
// threadblocks respects the ability to perform multicasting.
|
||||
//
|
||||
// To do so, we divide up the linearized stream-K units into clusters and share the same K
|
||||
// offsets for work within clusters.
|
||||
auto cluster_size = params.divmod_cluster_shape_major_.divisor * params.divmod_cluster_shape_minor_.divisor;
|
||||
auto cluster_linear_work_idx = linear_idx / cluster_size;
|
||||
WorkTileInfo& work_tile_info) {
|
||||
|
||||
// Determine the starting k iteration computed by this stream-K work unit
|
||||
uint32_t unit_iter_start = params.k_tiles_per_sk_unit_ * cluster_linear_work_idx;
|
||||
|
||||
// Adjust the starting position and number of k iterations for "big units," which
|
||||
// compute one extra iteration. These are the first big_units_ units in the
|
||||
// linearized ID space.
|
||||
bool is_big_unit = cluster_linear_work_idx < params.big_units_;
|
||||
if (is_big_unit) {
|
||||
// Since the "big units" are the first units in the linearized ID space, each
|
||||
// of the units preceding this big unit computed one extra iteration. Thus,
|
||||
// we must offset our start iteration by the number of units that precede
|
||||
// the current unit in the linearized ID space.
|
||||
unit_iter_start += cluster_linear_work_idx;
|
||||
} else {
|
||||
// Increment by one for each of the big clusters (since all big units precede this unit)
|
||||
unit_iter_start += params.big_units_;
|
||||
uint64_t true_tile_id = linear_idx;
|
||||
if (linear_idx >= params.sk_units_ && params.splits_ == 1) {
|
||||
// Data-parallel work
|
||||
true_tile_id = linear_idx - params.sk_units_ + params.sk_tiles_;
|
||||
work_tile_info.K_idx = 0;
|
||||
work_tile_info.k_tile_count = params.divmod_tiles_per_output_tile_.divisor;
|
||||
work_tile_info.k_tile_remaining = params.divmod_tiles_per_output_tile_.divisor;
|
||||
}
|
||||
else {
|
||||
// In the CUTLASS 2.x implementation of stream K, stream-K work is assigned to each stream-K
|
||||
// threadblock individually. For the most part, the set of K iterations corresponding to stream-K
|
||||
// work was divided amongst stream-K threadblocks, and a threadblock determined which tile
|
||||
// it would compute a (potentially-partial) output tile for based on the space of k iterations
|
||||
// assigned to it. This often results in stream-K threadblocks processing tiles with different
|
||||
// offsets in the K dimension from one another. This can reduce locality, but is lmitied to the
|
||||
// (generally few) waves of threadblocks assigned to compute stream-K work.
|
||||
//
|
||||
// With the introduction of threadblock clusters, there is additional benefit to maintaining
|
||||
// locality in the K dimension: shared portions of operands can be multicasted to threadblocks
|
||||
// within a cluster. Thus, we would like to ensure that the assignment of stream-K work to
|
||||
// threadblocks respects the ability to perform multicasting.
|
||||
//
|
||||
// To do so, we divide up the linearized stream-K units into clusters and share the same K
|
||||
// offsets for work within clusters.
|
||||
|
||||
uint32_t unit_iters;
|
||||
if (new_unit) {
|
||||
unit_iters = params.k_tiles_per_sk_unit_;
|
||||
// Equivalent to linear_idx / cluster_size
|
||||
auto cluster_linear_work_idx = params.divmod_cluster_shape_minor_.divide(
|
||||
params.divmod_cluster_shape_major_.divide(linear_idx)
|
||||
);
|
||||
|
||||
// Only adjust iteration count for big unit if we are initializing this
|
||||
// work unit. For existing work units, the extra iteration for big units
|
||||
// has already been accounted for in k_tiles_reamaining
|
||||
uint64_t split;
|
||||
params.divmod_clusters_mnl_(split, cluster_linear_work_idx, cluster_linear_work_idx);
|
||||
auto big_unit_cmp = params.splits_ > 1 ? split : cluster_linear_work_idx;
|
||||
auto linear_idx_mult = params.splits_ > 1 ? params.divmod_tiles_per_output_tile_.divisor : params.k_tiles_per_sk_unit_;
|
||||
|
||||
// Determine the starting k iteration computed by this stream-K work unit
|
||||
uint32_t unit_iter_start = (linear_idx_mult * cluster_linear_work_idx) + (params.k_tiles_per_sk_unit_ * split);
|
||||
|
||||
// Adjust the starting position and number of k iterations for "big units," which
|
||||
// compute one extra iteration. These are the first big_units_ units in the
|
||||
// linearized ID space.
|
||||
bool is_big_unit = big_unit_cmp < params.big_units_;
|
||||
if (is_big_unit) {
|
||||
++unit_iters;
|
||||
// Since the "big units" are the first units in the linearized ID space, each
|
||||
// of the units preceding this big unit computed one extra iteration. Thus,
|
||||
// we must offset our start iteration by the number of units that precede
|
||||
// the current unit in the linearized ID space.
|
||||
unit_iter_start += big_unit_cmp;
|
||||
}
|
||||
else {
|
||||
// Increment by one for each of the big clusters (since all big units precede this unit)
|
||||
unit_iter_start += params.big_units_;
|
||||
}
|
||||
|
||||
if (work_tile_info.k_tile_count == 0) {
|
||||
// This is a new unit
|
||||
work_tile_info.k_tile_remaining = params.k_tiles_per_sk_unit_;
|
||||
|
||||
// Only adjust iteration count for big unit if we are initializing this
|
||||
// work unit. For existing work units, the extra iteration for big units
|
||||
// has already been accounted for in k_tiles_reamaining
|
||||
if (is_big_unit) {
|
||||
++work_tile_info.k_tile_remaining;
|
||||
}
|
||||
}
|
||||
|
||||
// Find the output tile corresponding to the final k iteration covered by this
|
||||
// work unit. Stream-K work units will work backwards in terms of the tiles they
|
||||
// are responsible computing. This is beneficial because the final (partial)
|
||||
// tile computed by a stream-K block is typically the beginning of the output
|
||||
// tile, while the beginning (partial) tile is typically the ending of another
|
||||
// output tile. Since ending portions of an output tile must reduce across
|
||||
// other work units computing portions of that output tile, it is preferable
|
||||
// for them to be computed later, so as to reduce the likelihood of blocking
|
||||
// on other work.
|
||||
uint32_t unit_iter_end = unit_iter_start + work_tile_info.k_tile_remaining - 1;
|
||||
|
||||
true_tile_id = params.divmod_tiles_per_output_tile_.divide(unit_iter_end);
|
||||
uint32_t true_tile_iter_start = true_tile_id * params.divmod_tiles_per_output_tile_.divisor;
|
||||
uint32_t true_tile_iter_end = true_tile_iter_start + params.divmod_tiles_per_output_tile_.divisor;
|
||||
|
||||
// Bring the linearized tile ID back into the space of tiles, rather than clusters
|
||||
true_tile_id *= params.divmod_cluster_shape_major_.divisor * params.divmod_cluster_shape_minor_.divisor;
|
||||
|
||||
auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
|
||||
|
||||
// The final linearized tile ID is in units of the cluster dimension over which we rasterize.
|
||||
if (params.raster_order_ == RasterOrder::AlongN) {
|
||||
true_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
|
||||
}
|
||||
else {
|
||||
true_tile_id += cta_m_in_cluster * params.divmod_cluster_shape_minor_.divisor;
|
||||
}
|
||||
|
||||
// The unit's starting k iteration in the current tile is either the starting
|
||||
// iteration for the tile as a whole, or the starting k iteration for the unit
|
||||
// as a whole (if the latter is greater than the former).
|
||||
uint32_t tile_iter_start = max(true_tile_iter_start, unit_iter_start);
|
||||
|
||||
// Similarly, the unit's ending k iteration (exclusive) is either the end of
|
||||
// the current tile it is assigned, or the ending iteration of the unit as a whole
|
||||
// (if the latter is less than the former).
|
||||
uint32_t tile_iter_end = min(true_tile_iter_end, unit_iter_end + 1);
|
||||
|
||||
// Set the k offset to be the starting k tile for this output tile
|
||||
work_tile_info.K_idx = static_cast<int32_t>(tile_iter_start - true_tile_iter_start);
|
||||
|
||||
work_tile_info.k_tile_count = tile_iter_end - tile_iter_start;
|
||||
}
|
||||
else {
|
||||
unit_iters = work_tile_info.k_tile_remaining;
|
||||
}
|
||||
|
||||
// Find the output tile corresponding to the final k iteration covered by this
|
||||
// work unit. Stream-K work units will work backwards in terms of the tiles they
|
||||
// are responsible computing. This is beneficial because the final (partial)
|
||||
// tile computed by a stream-K block is typically the beginning of the output
|
||||
// tile, while the beginning (partial) tile is typically the ending of another
|
||||
// output tile. Since ending portions of an output tile must reduce across
|
||||
// other work units computing portions of that output tile, it is preferable
|
||||
// for them to be computed later, so as to reduce the likelihood of blocking
|
||||
// on other work.
|
||||
uint32_t unit_iter_end = unit_iter_start + unit_iters - 1;
|
||||
uint32_t true_tile_id = unit_iter_end / params.k_tiles_per_output_tile_;
|
||||
uint32_t true_tile_iter_start = true_tile_id * params.k_tiles_per_output_tile_;
|
||||
uint32_t true_tile_iter_end = true_tile_iter_start + params.k_tiles_per_output_tile_;
|
||||
|
||||
// Bring the linearized tile ID back into the space of tiles, rather than clusters
|
||||
true_tile_id *= cluster_size;
|
||||
|
||||
auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
|
||||
|
||||
// The final linearized tile ID is in units of the cluster dimension over which we rasterize.
|
||||
if (params.raster_order_ == RasterOrder::AlongN) {
|
||||
true_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
|
||||
}
|
||||
else {
|
||||
true_tile_id += cta_m_in_cluster * params.divmod_cluster_shape_minor_.divisor;
|
||||
}
|
||||
|
||||
// The unit's starting k iteration in the current tile is either the starting
|
||||
// iteration for the tile as a whole, or the starting k iteration for the unit
|
||||
// as a whole (if the latter is greater than the former).
|
||||
uint32_t tile_iter_start = max(true_tile_iter_start, unit_iter_start);
|
||||
|
||||
// Similarly, the unit's ending k iteration (exclusive) is either the end of
|
||||
// the current tile it is assigned, or the ending iteration of the unit as a whole
|
||||
// (if the latter is less than the former).
|
||||
uint32_t tile_iter_end = min(true_tile_iter_end, unit_iter_end + 1);
|
||||
|
||||
uint32_t tile_iters = tile_iter_end - tile_iter_start;
|
||||
|
||||
uint64_t work_idx_l, remainder;
|
||||
params.divmod_batch_(work_idx_l, remainder, true_tile_id);
|
||||
|
||||
uint64_t cta_per_grid_dim, dontcare;
|
||||
params.divmod_cluster_shape_minor_(cta_per_grid_dim, dontcare, remainder);
|
||||
uint64_t cta_per_grid_dim = params.divmod_cluster_shape_minor_.divide(remainder);
|
||||
|
||||
auto [work_idx_m, work_idx_n] = UnderlyingScheduler::get_work_idx_m_and_n(
|
||||
cta_per_grid_dim,
|
||||
@@ -627,113 +638,11 @@ public:
|
||||
params.log_swizzle_size_,
|
||||
params.raster_order_);
|
||||
|
||||
//
|
||||
// Update the work_tile_info
|
||||
//
|
||||
|
||||
// Set the M, N, and L block offsets
|
||||
work_tile_info.M_idx = work_idx_m;
|
||||
work_tile_info.N_idx = work_idx_n;
|
||||
work_tile_info.L_idx = static_cast<int32_t>(work_idx_l);
|
||||
|
||||
// Set the k offset to be the starting k tile for this output tile
|
||||
work_tile_info.K_idx = static_cast<int32_t>(tile_iter_start - true_tile_iter_start);
|
||||
|
||||
// Set the split count to be the number of k tiles in the output tile
|
||||
work_tile_info.splits = params.k_tiles_per_output_tile_;
|
||||
|
||||
// Any checks for invalid work units should be done prior to this call
|
||||
work_tile_info.is_valid_tile = true;
|
||||
|
||||
work_tile_info.k_tile_count = tile_iters;
|
||||
work_tile_info.k_tile_remaining = unit_iters;
|
||||
|
||||
// Compute the epilogue if this unit of work contains the ending k iteration for
|
||||
// the output tile in question
|
||||
work_tile_info.is_final_split = (tile_iter_end == true_tile_iter_end);
|
||||
}
|
||||
|
||||
// Returns a WorkTileInfo to be computed for either the data-parallel or split-K
|
||||
// work unit identified by the provided linear ID.
|
||||
CUTLASS_DEVICE
|
||||
static WorkTileInfo
|
||||
set_non_stream_k_work(uint64_t linear_idx, Params const& params, bool is_split_k) {
|
||||
|
||||
// The linearized ID space is in terms of work units, rather than tiles. However,
|
||||
// to compute the correct block offset for a data-parallel tile, we must convert
|
||||
// the current ID to the data-parallel tile it corresponds to. Each data-parallel
|
||||
// unit maps to a single data-parallel tile, but each stream-K unit can map to more
|
||||
// than one tile. Thus, we must offset the work-unit ID among the data-parallel units
|
||||
// by the total number of output tiles that will be computed by stream-K units.
|
||||
//
|
||||
// The logic below also works for the split-K case, in which sk_units_ and sk_tiles_
|
||||
// are each 0.
|
||||
uint64_t linear_work_idx = linear_idx - params.sk_units_ + params.sk_tiles_;
|
||||
|
||||
// Map worker's linear index into the CTA-tiled problem shape to the corresponding MNL indices
|
||||
uint64_t work_idx_l, remainder;
|
||||
params.divmod_batch_(work_idx_l, remainder, linear_work_idx);
|
||||
|
||||
uint64_t work_idx_k = 0;
|
||||
if (is_split_k) {
|
||||
params.divmod_k_(work_idx_k, remainder, remainder);
|
||||
}
|
||||
|
||||
uint64_t cta_per_grid_dim, dontcare;
|
||||
params.divmod_cluster_shape_minor_(cta_per_grid_dim, dontcare, remainder);
|
||||
|
||||
auto [work_idx_m, work_idx_n] = UnderlyingScheduler::get_work_idx_m_and_n(
|
||||
cta_per_grid_dim,
|
||||
params.divmod_cluster_shape_major_,
|
||||
params.divmod_cluster_shape_minor_,
|
||||
params.divmod_cluster_blk_major_,
|
||||
params.log_swizzle_size_,
|
||||
params.raster_order_);
|
||||
|
||||
bool is_final_split = (work_idx_k == params.splits_ - 1);
|
||||
|
||||
uint32_t k_tiles = params.k_tiles_per_output_tile_;
|
||||
if (is_split_k) {
|
||||
// Determine the number of iterations and starting iteration of this split.
|
||||
// Doing so requires accounting for residual iterations, which are handled
|
||||
// by the first big_units_ splits (with big_units_ = tiles % sm_count).
|
||||
|
||||
// Offsets for "normal" units. No additional k iterations are performed,
|
||||
// and big_units_ "big" units preceded us, each of which performed one
|
||||
// additional iteration. Thus, we must increase our split starting offset
|
||||
// by big_units_.
|
||||
int additional_k_tiles = 0;
|
||||
int split_start_offset = params.big_units_;
|
||||
|
||||
if (work_idx_k < params.big_units_) {
|
||||
// Offsets for "big" units. One additional k iteration is performed,
|
||||
// and each split preceding us was a big unit, so we must increase
|
||||
// our split starting offset by our split ID (work_idx_k).
|
||||
additional_k_tiles = 1;
|
||||
split_start_offset = work_idx_k;
|
||||
}
|
||||
|
||||
// Set up k iteration count and split starting iteration assuming the
|
||||
// iteration space is evenly split.
|
||||
k_tiles /= params.splits_;
|
||||
work_idx_k *= k_tiles;
|
||||
|
||||
// Apply any fixup needed to handle residuals
|
||||
work_idx_k += split_start_offset;
|
||||
k_tiles += additional_k_tiles;
|
||||
}
|
||||
|
||||
return {
|
||||
work_idx_m,
|
||||
work_idx_n,
|
||||
static_cast<int32_t>(work_idx_k),
|
||||
static_cast<int32_t>(work_idx_l),
|
||||
true,
|
||||
params.k_tiles_per_output_tile_,
|
||||
k_tiles,
|
||||
k_tiles, // remaining iterations
|
||||
is_final_split
|
||||
};
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -78,8 +78,8 @@ struct PersistentTileSchedulerSm90Params {
|
||||
AlongN
|
||||
};
|
||||
|
||||
FastDivmodU64 divmod_cluster_shape_major_{};
|
||||
FastDivmodU64 divmod_cluster_shape_minor_{};
|
||||
FastDivmodU64Pow2 divmod_cluster_shape_major_{};
|
||||
FastDivmodU64Pow2 divmod_cluster_shape_minor_{};
|
||||
FastDivmodU64 divmod_batch_{};
|
||||
FastDivmodU64 divmod_cluster_blk_major_{};
|
||||
|
||||
@@ -143,13 +143,13 @@ struct PersistentTileSchedulerSm90Params {
|
||||
divmod_batch_ = FastDivmodU64(problem_blocks_m * problem_blocks_n);
|
||||
|
||||
if (raster_order == RasterOrder::AlongN) {
|
||||
divmod_cluster_shape_major_ = FastDivmodU64(cluster_shape.n());
|
||||
divmod_cluster_shape_minor_ = FastDivmodU64(cluster_shape.m());
|
||||
divmod_cluster_shape_major_ = FastDivmodU64Pow2(cluster_shape.n());
|
||||
divmod_cluster_shape_minor_ = FastDivmodU64Pow2(cluster_shape.m());
|
||||
divmod_cluster_blk_major_ = FastDivmodU64(problem_blocks_n / cluster_shape.n());
|
||||
}
|
||||
else {
|
||||
divmod_cluster_shape_major_ = FastDivmodU64(cluster_shape.m());
|
||||
divmod_cluster_shape_minor_ = FastDivmodU64(cluster_shape.n());
|
||||
divmod_cluster_shape_major_ = FastDivmodU64Pow2(cluster_shape.m());
|
||||
divmod_cluster_shape_minor_ = FastDivmodU64Pow2(cluster_shape.n());
|
||||
divmod_cluster_blk_major_ = FastDivmodU64(problem_blocks_m / cluster_shape.m());
|
||||
}
|
||||
}
|
||||
@@ -374,15 +374,22 @@ struct PersistentTileSchedulerSm90StreamKParams {
|
||||
using RasterOrder = UnderlyingParams::RasterOrder;
|
||||
using RasterOrderOptions = UnderlyingParams::RasterOrderOptions;
|
||||
|
||||
FastDivmodU64 divmod_cluster_shape_major_{};
|
||||
FastDivmodU64 divmod_cluster_shape_minor_{};
|
||||
// Cluster dimensions are typically always a power of 2, so use
|
||||
// the power-of-two variants of FastDivmod for these.
|
||||
FastDivmodU64Pow2 divmod_cluster_shape_major_{};
|
||||
FastDivmodU64Pow2 divmod_cluster_shape_minor_{};
|
||||
|
||||
FastDivmodU64 divmod_batch_{};
|
||||
FastDivmodU64 divmod_k_{};
|
||||
FastDivmodU64 divmod_cluster_blk_major_{};
|
||||
|
||||
int32_t log_swizzle_size_ = 0;
|
||||
// Total number of cluster-sized output tiles (i.e., not including any
|
||||
// splitting factors). This is primarily used for split-K decompositions,
|
||||
// and may be overridden in other decompositions.
|
||||
FastDivmodU64 divmod_clusters_mnl_{};
|
||||
|
||||
uint64_t units_per_problem_ = 0;
|
||||
FastDivmod divmod_tiles_per_output_tile_{};
|
||||
int32_t log_swizzle_size_ = 0;
|
||||
RasterOrder raster_order_ = RasterOrder::AlongN;
|
||||
|
||||
// The splitting factor to be used in a split-K decomposition of the problem.
|
||||
@@ -390,9 +397,6 @@ struct PersistentTileSchedulerSm90StreamKParams {
|
||||
// is bypassed in favor of a split-K decomposition.
|
||||
uint32_t splits_ = 1;
|
||||
|
||||
// Number of tiled k iterations required to compute a single output tile.
|
||||
uint32_t k_tiles_per_output_tile_ = 0;
|
||||
|
||||
// Number of stream-K or split-K work units that compute an extra k iteration.
|
||||
// This is done to handle residuals in dividing up the k iteration space.
|
||||
// For stream-K, since the actual assignment of work to stream-K units will be done
|
||||
@@ -475,10 +479,10 @@ struct PersistentTileSchedulerSm90StreamKParams {
|
||||
raster_order_option
|
||||
);
|
||||
|
||||
auto problem_blocks_m = problem_blocks.x;
|
||||
auto problem_blocks_n = problem_blocks.y;
|
||||
auto problem_blocks_l = problem_blocks.z;
|
||||
|
||||
auto problem_blocks_m = round_up(problem_blocks.x, (1 << underlying_params.log_swizzle_size_) * cluster_shape.m());
|
||||
auto problem_blocks_n = round_up(problem_blocks.y, (1 << underlying_params.log_swizzle_size_) * cluster_shape.n());
|
||||
uint64_t output_tiles = problem_blocks_m * problem_blocks_n * problem_blocks_l;
|
||||
|
||||
// Reduction workspace is at the beginning of the workspace. Lock workspace follows.
|
||||
@@ -620,13 +624,17 @@ struct PersistentTileSchedulerSm90StreamKParams {
|
||||
divmod_cluster_shape_major_ = underlying_params.divmod_cluster_shape_major_;
|
||||
divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_;
|
||||
divmod_batch_ = underlying_params.divmod_batch_;
|
||||
divmod_k_ = FastDivmodU64(problem_blocks_m * problem_blocks_n); // Static k-splitting divmod. Unused for stream-K.
|
||||
divmod_tiles_per_output_tile_ = FastDivmod(k_tiles_per_output_tile);
|
||||
divmod_cluster_blk_major_ = underlying_params.divmod_cluster_blk_major_;
|
||||
|
||||
// Override divmod_clusters_mnl_ to be the number of cluster-sized stream-K units.
|
||||
// This setting ensures that the use of this divmod for stream-K decompositions
|
||||
// is essentially a no-op.
|
||||
divmod_clusters_mnl_ = FastDivmodU64(sk_units / cluster_size);
|
||||
splits_ = 1;
|
||||
log_swizzle_size_ = underlying_params.log_swizzle_size_;
|
||||
units_per_problem_ = static_cast<uint32_t>(dp_units + sk_units);
|
||||
raster_order_ = underlying_params.raster_order_;
|
||||
splits_ = 1; // Static k-splitting factor. Unused for stream-K.
|
||||
k_tiles_per_output_tile_ = k_tiles_per_output_tile;
|
||||
big_units_ = static_cast<uint32_t>(sk_big_units_per_cluster);
|
||||
reduction_workspace_ = reduction_workspace;
|
||||
sk_tiles_ = sk_tiles;
|
||||
@@ -755,6 +763,10 @@ struct PersistentTileSchedulerSm90StreamKParams {
|
||||
uint32_t barrier_bits,
|
||||
uint32_t accumulator_bits) {
|
||||
|
||||
auto log_swizzle_size = UnderlyingParams::get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle);
|
||||
problem_blocks.x = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());
|
||||
problem_blocks.y = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n());
|
||||
|
||||
// Workspace is needed only for output tiles that will be split. Thus, we first determine the number
|
||||
// of output tiles that will be split, and then calculate the workspace needed to cover these.
|
||||
uint64_t output_tiles = problem_blocks.x * problem_blocks.y * problem_blocks.z;
|
||||
@@ -966,24 +978,25 @@ struct PersistentTileSchedulerSm90StreamKParams {
|
||||
void* reduction_workspace,
|
||||
ReductionMode reduction_mode) {
|
||||
|
||||
divmod_cluster_shape_major_ = underlying_params.divmod_cluster_shape_major_,
|
||||
divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_,
|
||||
divmod_batch_ = FastDivmodU64(blocks_m * blocks_n * splits),
|
||||
divmod_k_ = FastDivmodU64(blocks_m * blocks_n),
|
||||
divmod_cluster_blk_major_ = underlying_params.divmod_cluster_blk_major_,
|
||||
log_swizzle_size_ = underlying_params.log_swizzle_size_,
|
||||
units_per_problem_ = blocks_m * blocks_n * blocks_l * splits,
|
||||
raster_order_ = underlying_params.raster_order_,
|
||||
splits_ = splits,
|
||||
k_tiles_per_output_tile_ = k_tiles_per_output_tile,
|
||||
big_units_ = k_tiles_per_output_tile % splits,
|
||||
divmod_cluster_shape_major_ = underlying_params.divmod_cluster_shape_major_;
|
||||
divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_;
|
||||
divmod_batch_ = FastDivmodU64(blocks_m * blocks_n);
|
||||
divmod_tiles_per_output_tile_ = FastDivmod(k_tiles_per_output_tile);
|
||||
auto cluster_size = underlying_params.divmod_cluster_shape_major_.divisor * underlying_params.divmod_cluster_shape_minor_.divisor;
|
||||
divmod_clusters_mnl_ = FastDivmodU64((blocks_m * blocks_n * blocks_l) / cluster_size);
|
||||
splits_ = splits;
|
||||
divmod_cluster_blk_major_ = underlying_params.divmod_cluster_blk_major_;
|
||||
log_swizzle_size_ = underlying_params.log_swizzle_size_;
|
||||
units_per_problem_ = blocks_m * blocks_n * blocks_l;
|
||||
raster_order_ = underlying_params.raster_order_;
|
||||
big_units_ = k_tiles_per_output_tile % splits;
|
||||
reduction_workspace_ = reduction_workspace;
|
||||
reduction_mode_ = reduction_mode;
|
||||
k_tiles_per_sk_unit_ = k_tiles_per_output_tile / splits;
|
||||
|
||||
// No stream-K work is performed for "basic" data-parallel and split-K decompositions
|
||||
sk_tiles_ = 0;
|
||||
sk_units_ = 0;
|
||||
k_tiles_per_sk_unit_ = 0;
|
||||
}
|
||||
|
||||
private:
|
||||
|
||||
@@ -643,9 +643,9 @@ public:
|
||||
// we can start right away on mma instructions
|
||||
if (warp_mma_k + 1 == Base::kWarpGemmIterations)
|
||||
warp_mma.transform(warp_transformed_frag_A[(warp_mma_k + 1) % 2],
|
||||
warp_transformed_frag_B[(warp_mma_k + 1) % 2],
|
||||
warp_transformed_frag_B[(warp_mma_k + 1) % Detail::kBBufferSize],
|
||||
warp_loaded_frag_A[(warp_mma_k + 1) % 2],
|
||||
warp_loaded_frag_B[(warp_mma_k + 1) % 2]);
|
||||
warp_loaded_frag_B[(warp_mma_k + 1) % Detail::kBBufferSize]);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -408,19 +408,21 @@ struct ThreadblockSwizzleStreamK {
|
||||
}
|
||||
|
||||
/// Constructor: *Gemm* problem size (m, n, k)
|
||||
template <typename GemmKernel>
|
||||
ThreadblockSwizzleStreamK(
|
||||
KernelTraits<GemmKernel> const kernel_traits_,
|
||||
GemmUniversalMode const mode_,
|
||||
GemmCoord const problem_size_,
|
||||
GemmCoord const tile_size_,
|
||||
int const batch_split_, /// Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor (1 defaults to StreamK, >1 emulates Split-K)
|
||||
int const sm_occupancy_,
|
||||
int const device_sms_,
|
||||
int const avail_sms_) /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
|
||||
int const avail_sms_, /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
|
||||
size_t const element_A_bytes_,
|
||||
size_t const element_B_bytes_,
|
||||
size_t const element_C_bytes_,
|
||||
int const epilogue_acc_fragments_)
|
||||
:
|
||||
problem_size(problem_size_),
|
||||
batch_count((mode_ == GemmUniversalMode::kBatched) ? batch_split_ : 1),
|
||||
batch_count((mode_ == GemmUniversalMode::kBatched || mode_ == GemmUniversalMode::kArray) ? batch_split_ : 1),
|
||||
reduction_blocks(0),
|
||||
dp_blocks(0),
|
||||
dp_first_wave_tiles(1), // Default: one tile per DP-block in the first wave of DP blocks
|
||||
@@ -446,17 +448,17 @@ struct ThreadblockSwizzleStreamK {
|
||||
batch_count);
|
||||
|
||||
size_t problem_bytes =
|
||||
(sizeof(typename GemmKernel::ElementC) * problem_size.m() * problem_size.n()) +
|
||||
(sizeof(typename GemmKernel::ElementA) * problem_size.m() * problem_size.k()) +
|
||||
(sizeof(typename GemmKernel::ElementB) * problem_size.k() * problem_size.n());
|
||||
(element_C_bytes_ * problem_size.m() * problem_size.n()) +
|
||||
(element_A_bytes_ * problem_size.m() * problem_size.k()) +
|
||||
(element_B_bytes_ * problem_size.k() * problem_size.n());
|
||||
|
||||
size_t problem_flops = size_t(problem_size.m()) * size_t(problem_size.n()) * size_t(problem_size.k()) * 2;
|
||||
|
||||
float flops_per_byte = float(problem_flops) / float(problem_bytes);
|
||||
[[maybe_unused]] float flops_per_byte = float(problem_flops) / float(problem_bytes);
|
||||
|
||||
int output_tiles = tiled_shape.m() * tiled_shape.n();
|
||||
int waves = (output_tiles + avail_sms - 1) / avail_sms;
|
||||
float dp_efficiency = float(output_tiles) / float(waves * avail_sms);
|
||||
[[maybe_unused]] float dp_efficiency = float(output_tiles) / float(waves * avail_sms);
|
||||
|
||||
//
|
||||
// Determine dispatch composition of DP-tiles and SK-blocks
|
||||
@@ -528,8 +530,7 @@ struct ThreadblockSwizzleStreamK {
|
||||
(sk_blocks > 2 * sk_tiles))
|
||||
{
|
||||
// Launch a reduction block for every accumulator fragment in each SK-tile
|
||||
static const int kAccumulatorFragments = GemmKernel::Epilogue::kAccumulatorFragments;
|
||||
reduction_blocks = sk_tiles * kAccumulatorFragments;
|
||||
reduction_blocks = sk_tiles * epilogue_acc_fragments_;
|
||||
|
||||
}
|
||||
|
||||
|
||||
@@ -294,6 +294,8 @@ struct DefaultMmaTensorOp<
|
||||
Policy, PartitionsK, AccumulatorsInRowMajor>;
|
||||
};
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
} // namespace warp
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
@@ -132,10 +132,10 @@ public:
|
||||
static int const kNumElementsInWarpFragment = NumElementsInWarpFragment;
|
||||
static int const kNumElementsInMmaFragment = NumElementsInMmaFragment;
|
||||
static Operand const kOperand = Operand::kA;
|
||||
|
||||
|
||||
using WarpFragment = Array<ElementLoad, kNumElementsInWarpFragment>;
|
||||
using MmaFragment = Array<ElementLoad, kNumElementsInMmaFragment>;
|
||||
|
||||
|
||||
static uint32_t const kSelectBytesEvenThread = 0x5410;
|
||||
static uint32_t const kSelectBytesOddThread = 0x7632;
|
||||
|
||||
@@ -168,7 +168,7 @@ public:
|
||||
|
||||
uint32_t const* src_ptr = reinterpret_cast<uint32_t const *>(&mma_frag_src_ptr[n]);
|
||||
uint32_t *dst_ptr = reinterpret_cast<uint32_t *>(&mma_frag_dst_ptr[n]);
|
||||
|
||||
|
||||
// Shuffle data within the warp, pull from other threads within the warp
|
||||
uint32_t tmp0 = __shfl_up_sync(0xFFFFFFFF, src_ptr[0], delta_up_);
|
||||
uint32_t tmp1 = __shfl_down_sync(0xFFFFFFFF, src_ptr[0], delta_down_);
|
||||
@@ -218,7 +218,7 @@ public:
|
||||
|
||||
using WarpFragment = Array<ElementLoad, kNumElementsInWarpFragment>;
|
||||
using MmaFragment = Array<ElementLoad, kNumElementsInMmaFragment>;
|
||||
|
||||
|
||||
static uint32_t const kSelectBytesEvenThread = 0x5410;
|
||||
static uint32_t const kSelectBytesOddThread = 0x7632;
|
||||
|
||||
@@ -260,7 +260,7 @@ public:
|
||||
// Reorder the data within the 32-bit word (4x8b) required for mma.sync
|
||||
dst_ptr[0] = __byte_perm(tmp0, tmp1, byte_selector_);
|
||||
}
|
||||
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
@@ -279,7 +279,7 @@ template <
|
||||
///
|
||||
typename Enable = void>
|
||||
struct FragmentConverter {
|
||||
|
||||
|
||||
using ElementDst = ElementDst_;
|
||||
using ElementSrc = ElementSrc_;
|
||||
|
||||
@@ -522,17 +522,6 @@ public:
|
||||
void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B,
|
||||
FragmentA const &A, FragmentB const &B) const {
|
||||
|
||||
// Shuffle data within warp to obtain the mma.sync operand layout
|
||||
detail::FragmentShuffler<MmaElementA, ElementA, MmaIterations::kRow,
|
||||
FragmentA::kElements, MmaOperandA::kElements, Operand::kA> shuffler_A;
|
||||
FragmentA tmp_A;
|
||||
tmp_A = shuffler_A(A);
|
||||
|
||||
// Convert the A operand to the Mma Instruction operand type
|
||||
detail::FragmentConverter<MmaElementA, ElementA, FragmentA::kElements> convert_A;
|
||||
dst_A = convert_A(tmp_A);
|
||||
|
||||
|
||||
// Shuffle data within warp to obtain the mma.sync operand layout
|
||||
detail::FragmentShuffler<MmaElementB, ElementB, MmaIterations::kColumn,
|
||||
FragmentB::kElements, MmaOperandB::kElements, Operand::kB> shuffler_B;
|
||||
@@ -542,6 +531,27 @@ public:
|
||||
// Convert the B operand to the Mma Instruction operand type
|
||||
detail::FragmentConverter<MmaElementB, ElementB, FragmentB::kElements> convert_B;
|
||||
dst_B = convert_B(tmp_B);
|
||||
|
||||
FragmentA tmp_A;
|
||||
|
||||
Array<ElementA, FragmentA::kElements / 2> *
|
||||
ptr_tmp_A = reinterpret_cast<Array<ElementA,
|
||||
FragmentA::kElements / 2> *>(&tmp_A);
|
||||
Array<MmaElementA, FragmentA::kElements / 2> *
|
||||
ptr_dst_A = reinterpret_cast<Array<MmaElementA,
|
||||
FragmentA::kElements / 2> *>(&dst_A);
|
||||
|
||||
// Shuffle data within warp to obtain the mma.sync operand layout
|
||||
detail::FragmentShuffler<MmaElementA, ElementA, MmaIterations::kRow,
|
||||
FragmentA::kElements, MmaOperandA::kElements, Operand::kA> shuffler_A;
|
||||
|
||||
// Convert the A operand to the Mma Instruction operand type
|
||||
detail::FragmentConverter<MmaElementA, ElementA, FragmentA::kElements / 2> convert_A;
|
||||
|
||||
tmp_A = shuffler_A(A);
|
||||
ptr_dst_A[0] = convert_A(ptr_tmp_A[0]);
|
||||
|
||||
ptr_dst_A[1] = convert_A(ptr_tmp_A[1]);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -551,4 +561,4 @@ public:
|
||||
} // namespace gemm
|
||||
} // namespace cutlass
|
||||
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
/////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
@@ -158,6 +158,7 @@ public:
|
||||
/// Max ID2
|
||||
static int const kMaxID2 = Policy::Operator::kMaxID2;
|
||||
|
||||
static int const kVerticalVisit = false;
|
||||
/// Data type of meta E that is moved at the same time
|
||||
using ElementE =
|
||||
typename cutlass::platform::conditional<kMaxID2 == 1, uint32_t,
|
||||
@@ -251,8 +252,6 @@ public:
|
||||
using MmaOperandC = typename Policy::Operator::FragmentC;
|
||||
using MmaOperandE = typename Policy::Operator::FragmentE;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
|
||||
D = C;
|
||||
|
||||
MmaOperandA const *ptr_A = reinterpret_cast<MmaOperandA const *>(&A);
|
||||
@@ -260,6 +259,36 @@ public:
|
||||
MmaOperandC *ptr_D = reinterpret_cast<MmaOperandC *>(&D);
|
||||
MmaOperandE const *ptr_E = reinterpret_cast<MmaOperandE const *>(&E);
|
||||
|
||||
if (kVerticalVisit) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < MmaIterations::kColumn; ++n) {
|
||||
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < MmaIterations::kRow; ++m) {
|
||||
|
||||
int m_serpentine = ((n % 2) ? (MmaIterations::kRow - 1 - m) : m);
|
||||
int id2 = m_serpentine % kMaxID2;
|
||||
|
||||
if (AccumulatorsInRowMajor) { // matrix B is reordered
|
||||
mma(
|
||||
ptr_D[n + m_serpentine * MmaIterations::kColumn],
|
||||
ptr_A[m_serpentine],
|
||||
ptr_B[n],
|
||||
ptr_D[n + m_serpentine * MmaIterations::kColumn],
|
||||
ptr_E[(m_serpentine / kMaxID2)],
|
||||
id2);
|
||||
} else {
|
||||
mma(
|
||||
ptr_D[m_serpentine + n * MmaIterations::kRow],
|
||||
ptr_A[m_serpentine],
|
||||
ptr_B[n],
|
||||
ptr_D[m_serpentine + n * MmaIterations::kRow],
|
||||
ptr_E[(m_serpentine / kMaxID2)],
|
||||
id2);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < MmaIterations::kRow; ++m) {
|
||||
|
||||
@@ -288,9 +317,7 @@ public:
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
assert(0);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/// Transform the mma operands to the required types
|
||||
@@ -298,7 +325,6 @@ public:
|
||||
void transform(TransformedFragmentA &dst_A, TransformedFragmentB &dst_B,
|
||||
FragmentA const &A, FragmentB const &B) const {
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
//
|
||||
// Define conversions from source type to instruction type
|
||||
//
|
||||
@@ -308,25 +334,42 @@ public:
|
||||
FloatRoundStyle const kRoundB =
|
||||
PreferredRoundingMode<typename ArchMmaOperator::ElementB,
|
||||
ElementB>::kRound;
|
||||
detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
|
||||
FragmentA::kElements / 2, kRoundA>
|
||||
convert_A;
|
||||
NumericArrayConverter<typename ArchMmaOperator::ElementB, ElementB,
|
||||
FragmentB::kElements, kRoundB>
|
||||
convert_B;
|
||||
Array<ElementA, FragmentA::kElements / 2> const *ptr_A =
|
||||
reinterpret_cast<Array<ElementA, FragmentA::kElements / 2> const *>(&A);
|
||||
Array<typename ArchMmaOperator::ElementA, FragmentA::kElements / 2> *
|
||||
ptr_dst_A = reinterpret_cast<Array<typename ArchMmaOperator::ElementA,
|
||||
FragmentA::kElements / 2> *>(&dst_A);
|
||||
|
||||
dst_B = convert_B(B);
|
||||
|
||||
ptr_dst_A[0] = convert_A(ptr_A[0]);
|
||||
ptr_dst_A[1] = convert_A(ptr_A[1]);
|
||||
#else
|
||||
assert(0);
|
||||
#endif
|
||||
if (kVerticalVisit) {
|
||||
detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
|
||||
FragmentA::kElements, kRoundA>
|
||||
convert_A;
|
||||
NumericArrayConverter<typename ArchMmaOperator::ElementB, ElementB,
|
||||
FragmentB::kElements / 2, kRoundB>
|
||||
convert_B;
|
||||
Array<ElementB, FragmentB::kElements / 2> const *ptr_B =
|
||||
reinterpret_cast<Array<ElementB, FragmentB::kElements / 2> const *>(&B);
|
||||
Array<typename ArchMmaOperator::ElementB, FragmentB::kElements / 2> *
|
||||
ptr_dst_B = reinterpret_cast<Array<typename ArchMmaOperator::ElementB,
|
||||
FragmentB::kElements / 2> *>(&dst_B);
|
||||
|
||||
dst_A = convert_A(A);
|
||||
|
||||
ptr_dst_B[0] = convert_B(ptr_B[0]);
|
||||
ptr_dst_B[1] = convert_B(ptr_B[1]);
|
||||
} else {
|
||||
detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
|
||||
FragmentA::kElements / 2, kRoundA>
|
||||
convert_A;
|
||||
NumericArrayConverter<typename ArchMmaOperator::ElementB, ElementB,
|
||||
FragmentB::kElements, kRoundB>
|
||||
convert_B;
|
||||
Array<ElementA, FragmentA::kElements / 2> const *ptr_A =
|
||||
reinterpret_cast<Array<ElementA, FragmentA::kElements / 2> const *>(&A);
|
||||
Array<typename ArchMmaOperator::ElementA, FragmentA::kElements / 2> *
|
||||
ptr_dst_A = reinterpret_cast<Array<typename ArchMmaOperator::ElementA,
|
||||
FragmentA::kElements / 2> *>(&dst_A);
|
||||
|
||||
dst_B = convert_B(B);
|
||||
|
||||
ptr_dst_A[0] = convert_A(ptr_A[0]);
|
||||
ptr_dst_A[1] = convert_A(ptr_A[1]);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
@@ -217,6 +217,12 @@ public:
|
||||
/// Number of partitions along K dimension
|
||||
static int const kPartitionsK = PartitionsK_;
|
||||
|
||||
#if defined(__CUDA_ARCH__) && ((__CUDA_ARCH__ < 800) || (__CUDA_ARCH__ == 890))
|
||||
static int const kVerticalVisit = true;
|
||||
#else
|
||||
static int const kVerticalVisit = false;
|
||||
#endif
|
||||
|
||||
public:
|
||||
|
||||
/// Iterates over the A operand in memory
|
||||
@@ -293,16 +299,8 @@ public:
|
||||
MmaOperandB const *ptr_B = reinterpret_cast<MmaOperandB const *>(&B);
|
||||
MmaOperandC *ptr_D = reinterpret_cast<MmaOperandC *>(&D);
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
|
||||
// Serpentine visitation order maximizing reuse of Rb
|
||||
// The visitation order is like
|
||||
// _
|
||||
// | | | |
|
||||
// | | | |
|
||||
// |_| |_|
|
||||
//
|
||||
// Down Up Down Up
|
||||
|
||||
|
||||
if (kVerticalVisit) {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int n = 0; n < MmaIterations::kColumn; ++n) {
|
||||
|
||||
@@ -326,16 +324,7 @@ public:
|
||||
}
|
||||
}
|
||||
}
|
||||
#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
// Serpentine visitation order maximizing reuse of Ra
|
||||
// The visitation order is like
|
||||
// _________
|
||||
// _________|
|
||||
// |_________
|
||||
// __________|
|
||||
//
|
||||
// Right Left Right Left
|
||||
|
||||
} else {
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int m = 0; m < MmaIterations::kRow; ++m) {
|
||||
|
||||
@@ -358,9 +347,7 @@ public:
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
assert(0);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
/// Transform the mma operands to the required types
|
||||
@@ -377,7 +364,7 @@ public:
|
||||
FloatRoundStyle const kRoundB =
|
||||
PreferredRoundingMode<typename ArchMmaOperator::ElementB,
|
||||
ElementB>::kRound;
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800)
|
||||
if (kVerticalVisit) {
|
||||
detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
|
||||
FragmentA::kElements, kRoundA>
|
||||
convert_A;
|
||||
@@ -394,8 +381,7 @@ public:
|
||||
|
||||
ptr_dst_B[0] = convert_B(ptr_B[0]);
|
||||
ptr_dst_B[1] = convert_B(ptr_B[1]);
|
||||
|
||||
#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
|
||||
} else {
|
||||
detail::ConvertAndPack<typename ArchMmaOperator::ElementA, ElementA,
|
||||
FragmentA::kElements / 2, kRoundA>
|
||||
convert_A;
|
||||
@@ -412,9 +398,7 @@ public:
|
||||
|
||||
ptr_dst_A[0] = convert_A(ptr_A[0]);
|
||||
ptr_dst_A[1] = convert_A(ptr_A[1]);
|
||||
#else
|
||||
assert(0);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
Reference in New Issue
Block a user