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

CUTLASS 3.6.0 (#1850)

Browse Source 
* v3.6

* update changelog

* update readme

* fix typo

* fixing typos

* hopper gemm with weight prefetch

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>
This commit is contained in:
Yujia Zhai
2024-10-09 15:33:27 -04:00
committed by GitHub
co-authored by yuzhai Haicheng Wu
parent 0837a2a00a
commit cc3c29a81a
354 changed files with 105937 additions and 8197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/cutlass/gemm/kernel/sm70_gemm.hpp
+2
View File
@@ -81,6 +81,8 @@ public:
TileScheduler_, ArchTag, TileShape,
cute::Shape<cute::Int<1>, cute::Int<1>, cute::Int<1>>>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
static constexpr bool IsGdcEnabled = false;
static constexpr bool is_valid_tile_scheduler =
cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>;
static_assert(is_valid_tile_scheduler, "SM70 kernel does not support specializing the tile scheduler.");
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
+30 -25
View File
@@ -48,6 +48,7 @@
#include "cutlass/trace.h"
#include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
#include "cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
@@ -78,6 +79,8 @@ public:
static_assert(cute::is_base_of_v<KernelPtrArrayTmaWarpSpecializedCooperative, typename CollectiveMainloop_::DispatchPolicy::Schedule>);
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
@@ -138,7 +141,7 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128> {
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -146,7 +149,7 @@ public:
EpilogueTensorStorage epilogue;
} tensors;
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
@@ -155,7 +158,7 @@ public:
alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order;
} pipelines;
struct TensorMapStorage : cute::aligned_struct<128> {
struct TensorMapStorage : cute::aligned_struct<128, _1> {
using MainloopTensorMapStorage = typename CollectiveMainloop::TensorMapStorage;
using EpilogueTensorMapStorage = typename CollectiveEpilogue::TensorMapStorage;
@@ -228,18 +231,14 @@ public:
workspace_offset += CollectiveMainloop::get_workspace_size(problem_shapes, args.mainloop, sm_count);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
// Precompute the sub tiles numbers in epilogue, pass into tile scheduler. Therefore it will be used
// in separate reduction scheme for streamk case, NumEpilogueSubTiles default value is 1, which means
// subtile will not be used, therefore separate reduction will not be enabled.
constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
TileSchedulerParams scheduler;
if constexpr (IsGroupedGemmKernel) {
scheduler = TileScheduler::to_underlying_arguments(
problem_shapes, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles);
problem_shapes, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace);
}
else {
scheduler = TileScheduler::to_underlying_arguments(
problem_shapes.get_host_problem_shape(), TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles);
problem_shapes.get_host_problem_shape(), TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace);
}
return {
@@ -305,9 +304,10 @@ public:
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
static constexpr uint32_t NumAccumulatorMtxs = 1;
status = TileScheduler::template initialize_workspace<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, cuda_adapter);
args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
@@ -341,10 +341,10 @@ public:
args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
dim3 grid_shape;
if constexpr (IsGroupedGemmKernel) {
grid_shape = TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
grid_shape = TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
else {
grid_shape = TileScheduler::get_grid_shape(params.problem_shape.get_host_problem_shape(), TileShape{}, ClusterShape{}, params.hw_info, args);
grid_shape = TileScheduler::get_grid_shape(params.scheduler, params.problem_shape.get_host_problem_shape(), TileShape{}, ClusterShape{}, params.hw_info, args);
}
return grid_shape;
}
@@ -372,7 +372,7 @@ public:
static_assert(NumMmaWarpGroups == 2, "Cooperative kernels currently only support NumMmaWarpGroups == 2");
if constexpr (cutlass::epilogue::collective::detail::sm90_is_ptr_array_tma_dispatch_policy_v<typename CollectiveEpilogue::DispatchPolicy>) {
static_assert(NumMmaWarpGroups == CollectiveEpilogue::NumEpilogueWarpGroups,
static_assert(NumMmaWarpGroups == CollectiveEpilogue::NumEpilogueWarpGroups,
"Tiled MmA does not match expected warp groups performing the epilogue");
}
@@ -546,7 +546,8 @@ public:
bool did_batch_change = true;
while (work_tile_info.is_valid()) {
if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
continue;
}
@@ -587,7 +588,8 @@ public:
}
// Get next work tile
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
auto next_batch = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx
did_batch_change = next_batch != curr_batch;
if (work_tile_info.is_valid() && did_batch_change) {
@@ -646,11 +648,12 @@ public:
}
load_order_barrier.wait();
while (work_tile_info.is_valid()) {
int32_t curr_batch = work_tile_info.L_idx;
// Get next work tile
auto next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
if constexpr (IsGroupedGemmKernel) {
@@ -682,7 +685,6 @@ public:
work_tile_info.reduction_subtile_idx(),
wait
);
}
work_tile_info = next_work_tile_info;
@@ -734,8 +736,8 @@ public:
constexpr bool IsEpiLoad = false;
if (work_tile_info.is_valid()) {
if (warp_idx_in_warp_group == 0) {
if (warp_idx_in_warp_group == 0) {
collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
shared_storage.tensormaps.epilogue,
params.epilogue,
@@ -747,9 +749,9 @@ public:
// Converge before issuing tensormap fence release since fence is aligned
__syncwarp();
collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
lane_predicate,
collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
lane_predicate,
consumer_warp_group_idx);
}
}
@@ -776,7 +778,8 @@ public:
static_assert(cute::is_any_of_v<TileScheduler,
detail::PersistentTileSchedulerSm90Group<ProblemShape>,
detail::PersistentTileSchedulerSm90>);
if(TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
@@ -825,13 +828,15 @@ public:
epi_store_tensormap,
work_tile_info.reduction_subtile_idx()
);
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 = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
did_batch_change = curr_batch != work_tile_info.L_idx;
if (work_tile_info.is_valid() && did_batch_change) {
@@ -850,8 +855,8 @@ public:
// Converge before issuing tensormap fence release since fence is aligned
__syncwarp();
collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
epi_store_tensormap,
lane_predicate,
consumer_warp_group_idx);
}
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp
+15 -16
View File
@@ -140,16 +140,14 @@ public:
// 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>;
using MathWarpGroupOrderBarrierSharedStorage = cutlass::PipelineDetail::OrderedSequenceBarrierSharedStorage<
MathWarpGroupOrderBarrier::SequenceDepth,
MathWarpGroupOrderBarrier::SequenceLength>;
// Kernel level shared memory storage
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128> {
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -157,7 +155,7 @@ public:
EpilogueTensorStorage epilogue;
} tensors;
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
using MathWarpGroupOrderBarrierStorage = MathWarpGroupOrderBarrierSharedStorage;
@@ -168,7 +166,7 @@ public:
alignas(16) MathWarpGroupOrderBarrierStorage math_wg_order;
} pipelines;
struct TensorMapStorage : cute::aligned_struct<128> {
struct TensorMapStorage : cute::aligned_struct<128, _1> {
using MainloopTensorMapStorage = typename CollectiveMainloop::TensorMapStorage;
using EpilogueTensorMapStorage = typename CollectiveEpilogue::TensorMapStorage;
@@ -318,9 +316,10 @@ public:
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
static constexpr uint32_t NumAccumulatorMtxs = 1;
status = TileScheduler::template initialize_workspace<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, cuda_adapter);
args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
@@ -354,10 +353,10 @@ public:
args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
dim3 grid_shape;
if constexpr (IsGroupedGemmKernel) {
grid_shape = TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
grid_shape = TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
else {
grid_shape = TileScheduler::get_grid_shape(params.problem_shape.get_host_problem_shape(), TileShape{}, ClusterShape{}, params.hw_info, args);
grid_shape = TileScheduler::get_grid_shape(params.scheduler, params.problem_shape.get_host_problem_shape(), TileShape{}, ClusterShape{}, params.hw_info, args);
}
return grid_shape;
}
@@ -521,7 +520,7 @@ public:
// Advance 2nd Math WG to the next work tile for the startup
const auto k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
auto next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
if (!work_tile_info.is_valid()) {
return;
@@ -580,7 +579,7 @@ public:
bool did_batch_change = true;
while (work_tile_info.is_valid()) {
if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
auto next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
continue;
}
@@ -622,7 +621,7 @@ public:
}
// Get next work tile
auto next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
auto next_batch = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx
did_batch_change = next_batch != curr_batch;
@@ -687,7 +686,7 @@ public:
int32_t curr_batch = work_tile_info.L_idx;
// Get next work tile
auto next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
if constexpr (IsGroupedGemmKernel) {
@@ -772,7 +771,6 @@ public:
if (work_tile_info.is_valid()) {
if (warp_idx_in_warp_group == 0) {
collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
shared_storage.tensormaps.epilogue,
params.epilogue,
@@ -876,7 +874,7 @@ public:
}
// Get next work tile
auto next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
// Skip a tile for pingpong
@@ -888,9 +886,10 @@ public:
mainloop_pipe_consumer_state.advance(work_k_tile_count);
// Go to next tile
auto next_next_work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_next_work_tile_info, next_increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_next_work_tile_info;
increment_pipe = next_increment_pipe;
}
did_batch_change = curr_batch != work_tile_info.L_idx;
@@ -946,4 +945,4 @@ public:
///////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::gemm::kernel
} // namespace cutlass::gemm::kernel
include/cutlass/gemm/kernel/sm90_gemm_tma.hpp
+2
View File
@@ -70,6 +70,8 @@ public:
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp
+119 -43
View File
@@ -33,7 +33,6 @@
#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"
@@ -43,7 +42,14 @@
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/trace.h"
#include "cutlass/conv/detail.hpp"
#include "cute/tensor.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
@@ -61,15 +67,22 @@ class GemmUniversal<
CollectiveMainloop_,
CollectiveEpilogue_,
TileScheduler_,
cute::enable_if_t<cute::is_base_of_v<KernelTmaWarpSpecialized, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
cute::enable_if_t<cute::is_base_of_v<cutlass::gemm::KernelTmaWarpSpecialized, typename CollectiveMainloop_::DispatchPolicy::Schedule>>
>
{
public:
//
// Type Aliases
//
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
// Handles the static_assert placed inside the operator()
// This is also used to decide whether the load_init inside collective mainloop returns rank 4 tensors or rank 5 tensors
static constexpr bool IsConvProblemShape = not (cute::is_tuple_v<ProblemShape>|| IsCutlass3ArrayKernel<ProblemShape>::value);
static_assert( IsConvProblemShape || (cute::rank(ProblemShape{}) == 3 || cute::rank(ProblemShape{}) == 4), "ProblemShape{} should be <M,N,K> or <M,N,K,L> for Gemm");
static constexpr bool IsGdcEnabled = cutlass::arch::IsGdcGloballyEnabled;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
@@ -99,7 +112,8 @@ public:
"TMA warp-specialized kernel does not support specializing the tile scheduler.");
using TileSchedulerTag = TileScheduler_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
TileSchedulerTag, ArchTag, TileShape, ClusterShape>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
// Kernel level shared memory storage
@@ -113,7 +127,7 @@ public:
EpilogueTensorStorage epilogue;
} tensors;
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
@@ -123,7 +137,6 @@ public:
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
static constexpr uint32_t NumLoadWarpGroups = 1;
static constexpr uint32_t NumMmaWarpGroups = 1;
static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
@@ -131,18 +144,54 @@ public:
// Device side arguments
struct Arguments {
GemmUniversalMode mode{};
cutlass::gemm::GemmUniversalMode mode{}; //maintained here for backward compatibility
ProblemShape problem_shape{};
MainloopArguments mainloop{};
EpilogueArguments epilogue{};
KernelHardwareInfo hw_info{};
TileSchedulerArguments scheduler{};
// Default constructor
Arguments() = default;
// Constructor with specified mode
// It is used for Gemm
Arguments(
cutlass::gemm::GemmUniversalMode mode_,
ProblemShape problem_shape_,
MainloopArguments mainloop_,
EpilogueArguments epilogue_,
KernelHardwareInfo hw_info_ = KernelHardwareInfo(),
TileSchedulerArguments scheduler_ = TileSchedulerArguments())
: mode(mode_)
, problem_shape(problem_shape_)
, mainloop(mainloop_)
, epilogue(epilogue_)
, hw_info(hw_info_)
, scheduler(scheduler_) {}
// Constructor with default value for 'mode'
// This allows us to set GemmUniversal mode as kGemm for Conv right away
// while keeping the testbeds unchanged
Arguments(
ProblemShape problem_shape_,
MainloopArguments mainloop_,
EpilogueArguments epilogue_,
KernelHardwareInfo hw_info_ = KernelHardwareInfo(),
TileSchedulerArguments scheduler_ = TileSchedulerArguments())
: mode(cutlass::gemm::GemmUniversalMode::kGemm) // Default mode
, problem_shape(problem_shape_)
, mainloop(mainloop_)
, epilogue(epilogue_)
, hw_info(hw_info_)
, scheduler(scheduler_) {}
};
// Kernel entry point API
struct Params {
GemmUniversalMode mode{};
ProblemShape problem_shape{};
using ProblemShapeMNKL = decltype(cutlass::conv::detail::get_problem_shape_MNKL_helper<CollectiveMainloop>(ProblemShape{}, cute::conditional_t<IsConvProblemShape, cute::true_type, cute::false_type>{}));
ProblemShapeMNKL problem_shape{};
MainloopParams mainloop{};
EpilogueParams epilogue{};
};
@@ -152,47 +201,49 @@ public:
//
// Convert to underlying arguments. In this case, a simple copy for the aliased type.
static
Params
static Params
to_underlying_arguments(Arguments const& args, void* workspace) {
(void) workspace;
auto problem_shape = args.problem_shape;
auto problem_shape_mnkl = cutlass::conv::detail::get_problem_shape_MNKL_helper<CollectiveMainloop>(args.problem_shape, cute::conditional_t<IsConvProblemShape, cute::true_type, cute::false_type>{});
auto transformed_problem_shape = cutlass::conv::detail::get_transformed_problem_shape_MNKL(args.problem_shape);
auto swapped_problem_shape = problem_shape_mnkl;
if constexpr (detail::Has_SwapAB_v<CollectiveMainloop>) {
// swap M/N
get<0>(problem_shape) = get<1>(args.problem_shape);
get<1>(problem_shape) = get<0>(args.problem_shape);
get<0>(swapped_problem_shape) = get<1>(problem_shape_mnkl);
get<1>(swapped_problem_shape) = get<0>(problem_shape_mnkl);
}
return {
args.mode,
problem_shape,
swapped_problem_shape,
CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace)
CollectiveEpilogue::to_underlying_arguments(transformed_problem_shape, args.epilogue, workspace)
};
}
static bool
can_implement(Arguments const& args) {
bool implementable = (args.mode == GemmUniversalMode::kGemm) or
(args.mode == GemmUniversalMode::kBatched && cute::rank(ProblemShape{}) == 4);
bool implementable = true;
auto transformed_problem_shape = cutlass::conv::detail::get_transformed_problem_shape_MNKL(args.problem_shape);
if (!implementable) {
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
return 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);
implementable &= CollectiveEpilogue::can_implement(transformed_problem_shape, args.epilogue);
implementable &= TileScheduler::can_implement(args.scheduler);
return implementable;
}
static
size_t
static size_t
get_workspace_size(Arguments const& args) {
return 0;
}
static
cutlass::Status
static cutlass::Status
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
CudaHostAdapter* cuda_adapter = nullptr) {
return Status::kSuccess;
@@ -318,38 +369,55 @@ public:
return [] () {}; // do nothing
}
} ();
// Preconditions
static_assert(cute::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(cute::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(cute::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(cute::rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
// 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>{});
// Preconditions only valid for Gemm
static_assert(IsConvProblemShape || cute::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(IsConvProblemShape || cute::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(IsConvProblemShape || cute::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(IsConvProblemShape || cute::rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
// Get the appropriate blocks for this thread block -- potential for thread block locality
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
auto blk_shape = TileShape{}; // (BLK_M,BLK_N,BLK_K)
TiledMma tiled_mma;
// Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
// Using constexpr if (C++17 and later)
auto problem_shape_MNKL = append<4>(params.problem_shape, cute::Int<1>{});
// 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>(load_inputs) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
// get<1>(load_inputs) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
// Prepare and partition the input tensors.
// Expects a tuple of tensors for conv where:
// get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k)
// get<1>(load_inputs) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k)
auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop);
static_assert(cute::tuple_size_v<decltype(load_inputs)> >= 2, "Output of load_init must have at least two elements (A, B)");
// Extract out partitioned A and B.
Tensor gA_mkl = get<0>(load_inputs);
Tensor gB_nkl = get<1>(load_inputs);
// 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 n_coord = idx2crd(int(blockIdx.y), shape<2>(gB_nkl), compact_col_major(shape<2>(gB_nkl)));
// handles the difference between the rank of Tensor returned by load_input in case they do not have a batch mode
auto l_coord = [&] (auto const& gB_nkl_) {
// gB_nkl needs to be passed into the lambda because C++17
// does not permit lambda capture of structured bindings.
if constexpr (not IsConvProblemShape) {
// This needs to be inside an `if constexpr`,
// because shape<4>(gB_nkl) is not well-formed otherwise.
return idx2crd(int(blockIdx.z), shape<4>(gB_nkl_));
}
else {
return Int<0>{};
}
} (gB_nkl);
auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
// Get pipeline iterators and increments from tensor shapes
@@ -361,6 +429,9 @@ public:
if (warp_group_role == WarpGroupRole::Producer) {
if (producer_warp_role == ProducerWarpRole::MainloopEpilogue) {
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
collective_mainloop.load(
params.mainloop,
mainloop_pipeline,
@@ -414,6 +485,11 @@ public:
k_tile_count
);
// Hint on an early release of global memory resources.
// The timing of calling this function only influences performance,
// not functional correctness.
cutlass::arch::launch_dependent_grids();
// Epilogue and write to gD
auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] =
collective_epilogue.store(
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp
+75 -29
View File
@@ -45,6 +45,8 @@
#include "cute/tensor.hpp"
#include "cutlass/trace.h"
#include "cutlass/gemm/kernel/gemm_universal_decl.h"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
@@ -55,13 +57,13 @@ template <
class ProblemShape_,
class CollectiveMainloop_,
class CollectiveEpilogue_,
class TileScheduler_
class TileSchedulerTag_
>
class GemmUniversal<
ProblemShape_,
CollectiveMainloop_,
CollectiveEpilogue_,
TileScheduler_,
TileSchedulerTag_,
cute::enable_if_t<cute::is_base_of_v<KernelTmaWarpSpecializedCooperative, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
{
public:
@@ -71,6 +73,7 @@ public:
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::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;
@@ -85,7 +88,6 @@ public:
using ClusterShape = typename DispatchPolicy::ClusterShape;
using MainloopArguments = typename CollectiveMainloop::Arguments;
using MainloopParams = typename CollectiveMainloop::Params;
// Epilogue derived types
using CollectiveEpilogue = CollectiveEpilogue_;
using ElementC = typename CollectiveEpilogue::ElementC;
@@ -97,15 +99,28 @@ public:
static_assert(ArchTag::kMinComputeCapability >= 90);
using TileSchedulerTag = TileScheduler_;
using TileSchedulerTag = TileSchedulerTag_;
using TileScheduler = typename detail::TileSchedulerSelector<
TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
TileSchedulerTag,
ArchTag,
TileShape,
ClusterShape
>::Scheduler;
using TileSchedulerArguments = typename TileScheduler::Arguments;
using TileSchedulerParams = typename TileScheduler::Params;
// Warp specialization thread count per threadblock
static constexpr uint32_t NumMMAThreads = size(TiledMma{}); // 8 warps
static constexpr uint32_t NumMainloopLoadThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp; // 1 warp for C
static constexpr bool IsGdcEnabled = cutlass::arch::IsGdcGloballyEnabled;
static constexpr uint32_t NumLoadWarpGroups = 1;
static constexpr uint32_t NumMmaWarpGroups = CUTE_STATIC_V(size(TiledMma{})) / NumThreadsPerWarpGroup;
static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
static constexpr uint32_t NumMmaWarpGroups = NumMMAThreads / NumThreadsPerWarpGroup;
static constexpr uint32_t MaxThreadsPerBlock = NumMMAThreads + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
/// Register requirement for Load and Math WGs
@@ -117,7 +132,7 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
@@ -126,7 +141,7 @@ public:
alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order;
} pipelines;
struct TensorStorage : cute::aligned_struct<128> {
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -256,9 +271,10 @@ public:
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
static constexpr uint32_t NumAccumulatorMtxs = 1;
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, cuda_adapter);
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
@@ -285,7 +301,7 @@ public:
args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
}
args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
static dim3
@@ -298,7 +314,6 @@ public:
operator()(Params const& params, char* smem_buf) {
using namespace cute;
using X = Underscore;
#if defined(__CUDA_ARCH_FEAT_SM90_ALL)
# define ENABLE_SM90_KERNEL_LEVEL 1
#endif
@@ -308,7 +323,7 @@ public:
#else
// Preconditions
static_assert(size(TiledMma{}) == 256, "Cooperative kernel must have TiledMMA operating using 256 threads.");
static_assert(NumMMAThreads == 256, "Cooperative kernel must have TiledMMA operating using 256 threads.");
static_assert(size<0>(TileShape{}) >= 128,
"Cooperative kernel requires Tile Size to be greater than or equal to 128 along the M-dimension.");
@@ -330,6 +345,8 @@ public:
Warp3 = 3
};
// Kernel level shared memory storage
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
@@ -338,7 +355,7 @@ public:
int warp_idx = canonical_warp_idx_sync();
int warp_idx_in_warp_group = warp_idx % NumWarpsPerWarpGroup;
int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
int mma_thread_idx = thread_idx % size(TiledMma{});
int mma_thread_idx = thread_idx % NumMMAThreads;
auto warp_group_role = WarpGroupRole(canonical_warp_group_idx());
auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group);
int lane_predicate = cute::elect_one_sync();
@@ -350,6 +367,8 @@ public:
CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue);
}
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
bool is_epi_load_needed = collective_epilogue.is_producer_load_needed();
// Mainloop Load pipeline
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
typename MainloopPipeline::Params mainloop_pipeline_params;
@@ -360,7 +379,7 @@ public:
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
}
mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
mainloop_pipeline_params.num_consumers = size(TiledMma{});
mainloop_pipeline_params.num_consumers = NumMMAThreads;
mainloop_pipeline_params.transaction_bytes = params.mainloop.tma_transaction_bytes;
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
@@ -369,13 +388,13 @@ public:
typename EpiLoadPipeline::Params epi_load_pipeline_params;
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::Epilogue) {
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
}
}
if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
}
epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster();
epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp;
epi_load_pipeline_params.consumer_arv_count = size(TiledMma{});
epi_load_pipeline_params.producer_arv_count = NumEpilogueLoadThreads;
epi_load_pipeline_params.consumer_arv_count = NumMMAThreads;
if constexpr (CollectiveEpilogue::RequiresTransactionBytes) {
epi_load_pipeline_params.transaction_bytes = params.epilogue.tma_transaction_bytes;
}
@@ -403,6 +422,7 @@ public:
PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
auto cluster_wait_fn = [] () {
// We need this to guarantee that the Pipeline init is visible
// To all producers and consumer thread blocks in the Cluster
@@ -425,7 +445,7 @@ public:
TileScheduler scheduler{params.scheduler};
auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{});
// In a warp specialized kernel, collectives expose data movement and compute operations separately
CollectiveMainloop collective_mainloop;
@@ -445,14 +465,16 @@ public:
if (warp_group_role == WarpGroupRole::Producer) {
cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
// Mainloop Producer Warp
if (producer_warp_role == ProducerWarpRole::Mainloop) {
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
bool do_load_order_arrive = true;
while (work_tile_info.is_valid()) {
if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
continue;
}
@@ -465,6 +487,7 @@ public:
// 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<3>(gA_mkl)), shape<3>(gA_mkl));
collective_mainloop.load(
@@ -486,9 +509,11 @@ public:
load_order_barrier.arrive();
do_load_order_arrive = false;
}
// Get next work tile
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info
);
work_tile_info = next_work_tile_info;
} // Scheduler work fetch loop
// Make sure all Consumer Warp Groups have been waited upon
@@ -497,11 +522,18 @@ public:
} // Mainloop Producer Warp End
// Epilogue Producer Warp
else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
else if (producer_warp_role == ProducerWarpRole::Epilogue && is_epi_load_needed) {
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
if (!TileScheduler::requires_separate_reduction(params.scheduler) && work_tile_info.is_valid()) {
load_order_barrier.wait();
}
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
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
@@ -509,7 +541,7 @@ public:
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);
epi_load_pipe_producer_state =
collective_epilogue.load(
epi_load_pipeline,
@@ -525,7 +557,9 @@ public:
}
// Get next work tile
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info
);
work_tile_info = next_work_tile_info;
} // Scheduler work fetch loop
// Make sure all Consumer Warp Groups have been waited upon
@@ -552,7 +586,7 @@ public:
//
// 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)
if(TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
if (TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
collective_mainloop.mma(
mainloop_pipeline,
mainloop_pipe_consumer_state,
@@ -573,6 +607,16 @@ public:
// Update starting mainloop pipeline state for the next tile
mainloop_pipe_consumer_state.advance(work_k_tile_count);
}
#ifdef CUTLASS_ENABLE_GDC_FOR_SM90
if (scheduler.is_last_tile(work_tile_info)) {
// Hint on an early release of global memory resources.
// The timing of calling this function only influences performance,
// not functional correctness.
cutlass::arch::launch_dependent_grids();
}
#endif
// Index of warp group within consumer warp groups
int consumer_warp_group_idx = canonical_warp_group_idx() - NumLoadWarpGroups;
@@ -603,7 +647,9 @@ public:
}
// Get next work tile
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info
);
work_tile_info = next_work_tile_info;
} // Scheduler work fetch loop
if (do_store_tail) {
include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
+40 -4
View File
@@ -47,6 +47,8 @@
#include "cutlass/trace.h"
#include "cute/tensor.hpp"
#include "cutlass/arch/grid_dependency_control.h"
///////////////////////////////////////////////////////////////////////////////
namespace cutlass::gemm::kernel {
@@ -73,6 +75,8 @@ public:
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = cutlass::arch::IsGdcGloballyEnabled;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
@@ -128,7 +132,7 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
using MathWarpGroupOrderBarrierStorage = MathWarpGroupOrderBarrierSharedStorage;
@@ -139,7 +143,7 @@ public:
alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order;
} pipelines;
struct TensorStorage : cute::aligned_struct<128> {
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -259,9 +263,11 @@ public:
Status status = Status::kSuccess;
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
size_t workspace_offset = 0;
static constexpr uint32_t NumEpilogueSubTiles = 1;
static constexpr uint32_t NumAccumulatorMtxs = 1;
status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, 1, cuda_adapter);
args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
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);
@@ -288,7 +294,7 @@ public:
args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
}
args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
static dim3
@@ -463,6 +469,9 @@ public:
// Mainloop Producer Warp
if (producer_warp_role == ProducerWarpRole::Mainloop) {
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
bool do_load_order_arrive = true;
while (work_tile_info.is_valid()) {
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
@@ -506,6 +515,10 @@ public:
// Epilogue Producer Warp
else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
// Ensure that the prefetched kernel does not touch
// unflushed global memory prior to this instruction
cutlass::arch::wait_on_dependent_grids();
load_order_barrier.wait();
while (work_tile_info.is_valid()) {
// Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
@@ -539,6 +552,19 @@ public:
else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
#ifdef CUTLASS_ENABLE_GDC_FOR_SM90
// It is possible to have work tiles start off invalid,
// so we have to check that first.
if (not work_tile_info.is_valid()) {
// Hint on an early release of global memory resources.
// The timing of calling this function only influences performance,
// not functional correctness.
cutlass::arch::launch_dependent_grids();
return;
}
#endif
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));
@@ -574,6 +600,16 @@ public:
// Update starting mainloop pipeline state for the next tile
mainloop_pipe_consumer_state.advance(k_tile_count * NumMmaWarpGroups);
#ifdef CUTLASS_ENABLE_GDC_FOR_SM90
if (scheduler.is_last_tile(work_tile_info, NumMmaWarpGroups)) {
// Hint on an early release of global memory resources.
// The timing of calling this function only influences performance,
// not functional correctness.
cutlass::arch::launch_dependent_grids();
}
#endif
// Order two Math WG's Epilogue one after the other
math_wg_order_barrier.wait();
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp
+3 -1
View File
@@ -68,6 +68,8 @@ public:
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
@@ -110,7 +112,7 @@ public:
EpilogueTensorStorage epilogue;
} tensors;
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp
+11 -6
View File
@@ -69,6 +69,7 @@ public:
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
@@ -114,7 +115,7 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128> {
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -122,7 +123,7 @@ public:
EpilogueTensorStorage epilogue;
} tensors;
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
@@ -223,8 +224,10 @@ public:
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
CudaHostAdapter* cuda_adapter = nullptr) {
TileScheduler t;
static constexpr uint32_t NumEpilogueSubTiles = 1;
static constexpr uint32_t NumAccumulatorMtxs = 1;
return t.template initialize_workspace<ProblemShape, ElementAccumulator>(
args.scheduler, workspace, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, 1, cuda_adapter);
args.scheduler, workspace, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
}
// Computes the kernel launch grid shape based on runtime parameters
@@ -235,7 +238,7 @@ public:
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);
return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
static dim3
@@ -401,7 +404,8 @@ public:
}
// Get next work tile
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
} // Scheduler work fetch loop
// Make sure all Consumer Warp Groups have been waited upon
@@ -477,7 +481,8 @@ public:
}
// Get next work tile
work_tile_info = scheduler.fetch_next_work(work_tile_info);
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info);
work_tile_info = next_work_tile_info;
} // Scheduler work fetch loop
if (do_store_tail) {
include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp
+4 -3
View File
@@ -71,6 +71,7 @@ public:
using ProblemShape = ProblemShape_;
static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
"ProblemShape{} should be <M,N,K> or <M,N,K,L>");
static constexpr bool IsGdcEnabled = false;
// Mainloop derived types
using CollectiveMainloop = CollectiveMainloop_;
using TileShape = typename CollectiveMainloop::TileShape;
@@ -123,7 +124,7 @@ public:
// Kernel level shared memory storage
struct SharedStorage {
struct TensorStorage : cute::aligned_struct<128> {
struct TensorStorage : cute::aligned_struct<128, _1> {
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -131,7 +132,7 @@ public:
EpilogueTensorStorage epilogue;
} tensors;
struct PipelineStorage : cute::aligned_struct<16> {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
using MathWarpGroupOrderBarrierStorage = typename MathWarpGroupOrderBarrier::SharedStorage;
@@ -243,7 +244,7 @@ public:
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);
return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
}
static dim3
include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp
+4 -14
View File
@@ -48,6 +48,8 @@ public:
using RasterOrderOptions = typename Params::RasterOrderOptions;
using Arguments = BaseScheduler::Arguments;
static constexpr bool IsDynamicPersistent = false;
// get work_idx_m, work_idx_n from blk_per_grid_dim while applying swizzle
static CUTLASS_DEVICE
cute::tuple<int32_t, int32_t>
@@ -121,29 +123,17 @@ public:
// The basic tile scheduler does not require any additional workspace
template <class ProblemShape, class ElementAccumulator>
static size_t
get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1) {
get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1, uint32_t = 1) {
return 0;
}
template <class ProblemShape, class ElementAccumulator>
static cutlass::Status
initialize_workspace(Arguments const&, void*, cudaStream_t, ProblemShape, KernelHardwareInfo const&,
uint32_t, const uint32_t = 1, CudaHostAdapter* cuda_adapter = nullptr) {
uint32_t, const uint32_t = 1, uint32_t = 1, CudaHostAdapter* cuda_adapter = nullptr) {
return Status::kSuccess;
}
// Kernel helper function to get next work tile
CUTLASS_DEVICE
auto
fetch_next_work(WorkTileInfo work_tile_info) {
if (continue_current_work(work_tile_info)) {
return work_tile_info;
}
advance_to_next_work();
return get_current_work();
}
};
}
include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp
+11 -6
View File
@@ -96,6 +96,8 @@ public:
using Params = PersistentTileSchedulerSm90GroupParams<ProblemShape>;
using RasterOrder = typename Params::RasterOrder;
using RasterOrderOptions = typename Params::RasterOrderOptions;
static constexpr bool IsDynamicPersistent = false;
struct Arguments {
int max_swizzle_size = 1;
// Not applying Heuristics for Grouped problems, since largest dimension can change per group
@@ -118,7 +120,9 @@ public:
KernelHardwareInfo const& hw_info,
Arguments const& arguments,
[[maybe_unused]] void* workspace=nullptr,
[[maybe_unused]] const uint32_t epilogue_subtile = 1) {
[[maybe_unused]] const uint32_t epilogue_subtile = 1,
[[maybe_unused]] uint32_t ktile_start_alignment_count = 1u
) {
// We only need the tile and cluster shape during scheduler setup, so let FTAD do the magic
static_assert(cute::is_static<TileShape>::value);
@@ -151,6 +155,7 @@ public:
CUTLASS_HOST_DEVICE static
dim3
get_grid_shape(
[[maybe_unused]] Params const& params,
GroupProblemShape problem_shapes,
TileShape tile_shape,
ClusterShape cluster_shape,
@@ -400,14 +405,14 @@ public:
// The basic tile scheduler does not require any additional workspace
template <class ProblemShape, class ElementAccumulator>
static size_t
get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1) {
get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1, uint32_t = 1) {
return 0;
}
template <class ProblemShape, class ElementAccumulator>
static cutlass::Status
initialize_workspace(Arguments const&, void*, cudaStream_t, ProblemShape, KernelHardwareInfo const&,
uint32_t, const uint32_t = 1, CudaHostAdapter* cuda_adapter = nullptr) {
uint32_t, const uint32_t = 1, uint32_t = 1, CudaHostAdapter* cuda_adapter = nullptr) {
return Status::kSuccess;
}
@@ -485,13 +490,13 @@ public:
auto
fetch_next_work(WorkTileInfo work_tile_info) {
if (continue_current_work(work_tile_info)) {
return work_tile_info;
return cute::make_tuple(work_tile_info, true);
}
advance_to_next_work();
return get_current_work();
return cute::make_tuple(get_current_work(), true);
}
// Returns the initial work tile info that will be computed over
template <class ClusterShape>
CUTLASS_DEVICE
include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp
+53 -46
View File
@@ -64,6 +64,8 @@ public:
using RasterOrder = UnderlyingScheduler::RasterOrder;
using RasterOrderOptions = UnderlyingScheduler::RasterOrderOptions;
static constexpr bool IsDynamicPersistent = false;
// Use a dummy barrier manager to simply get the type used to store the barrier
using BarrierType = typename NamedBarrierManager<1>::T;
@@ -194,13 +196,14 @@ public:
template <class ProblemShape>
static Params
to_underlying_arguments(
ProblemShape problem_shape,
TileShape tile_shape,
ClusterShape cluster_shape,
KernelHardwareInfo const& hw_info,
Arguments const& args,
void* workspace,
const uint32_t epilogue_subtile = 1) {
ProblemShape problem_shape,
TileShape tile_shape,
ClusterShape cluster_shape,
KernelHardwareInfo const& hw_info,
Arguments const& args,
void* workspace,
const uint32_t epilogue_subtile = 1,
[[maybe_unused]] uint32_t ktile_start_alignment_count = 1u) {
static_assert(cute::is_static<TileShape>::value);
static_assert(cute::is_static<ClusterShape>::value);
@@ -319,6 +322,7 @@ public:
CUTLASS_HOST_DEVICE static
dim3
get_grid_shape(
[[maybe_unused]] Params const& params,
ProblemShape problem_shape,
TileShape tile_shape,
ClusterShape cluster_shape,
@@ -393,28 +397,30 @@ public:
if (!requires_fixup(params, work_tile_info)) {
return;
}
auto tile_idx = output_tile_index(params, work_tile_info);
uint64_t tile_idx = output_tile_index(params, work_tile_info);
// Index of the lock on which to wait
auto lock_idx = (tile_idx * num_barriers) + barrier_idx;
uint64_t lock_idx = (tile_idx * num_barriers) + barrier_idx;
auto reduction_tile_idx = tile_idx;
auto [first_peer_id, my_peer_id, last_peer_id] = tile_peer_range(params, tile_idx, static_cast<uint32_t>(work_tile_info.K_idx));
auto reduction_peer_offset = 0;
uint64_t reduction_tile_idx = tile_idx;
uint64_t num_peers = 0;
uint64_t reduction_peer_offset = 0;
if (params.requires_separate_reduction()) {
// If separate reduction is to be performed, each stream-K unit writes its partials
// to a separate portion of the workspace. There are as many of these portions as there
// are peers for a given output tile, so we multiply the tile index by the maximum peer count.
auto [first_peer_id, my_peer_id, last_peer_id] = tile_peer_range(params, tile_idx, static_cast<uint32_t>(work_tile_info.K_idx));
num_peers = last_peer_id - first_peer_id + 1;
reduction_tile_idx *= Params::max_peers_per_tile(params.sk_units_, params.sk_tiles_);
reduction_peer_offset = my_peer_id * cute::size<0>(TileShape{}) * cute::size<1>(TileShape{});
}
// Reductions use BlockStripedReduce with a width of BarrierManager::ThreadCount under the hood.
// Thus, the start of the reduction space is the same across all threads in a warp group.
int reduction_offset =
(cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * reduction_tile_idx * num_accumulator_mtxs) +
uint64_t reduction_offset =
(static_cast<uint64_t>(cute::size<0>(TileShape{})) * static_cast<uint64_t>(cute::size<1>(TileShape{})) * reduction_tile_idx * num_accumulator_mtxs) +
reduction_peer_offset +
(size(accumulators) * barrier_idx * BarrierManager::ThreadCount);
(static_cast<uint64_t>(size(accumulators)) * barrier_idx * BarrierManager::ThreadCount);
ElementAccumulator* group_reduction_workspace = reinterpret_cast<ElementAccumulator*>(params.reduction_workspace_) + reduction_offset;
@@ -424,7 +430,7 @@ public:
AccumulatorArrayT* reduction_workspace_array = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace);
AccumulatorArrayT* accumulator_array = reinterpret_cast<AccumulatorArrayT*>(accumulators.data());
int barrier_group_thread_idx = threadIdx.x % BarrierManager::ThreadCount;
uint32_t 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)
@@ -443,24 +449,23 @@ public:
reduction_tiles = params.sk_tiles_;
}
auto reduction_workspace_size = Params::get_reduction_workspace_size(
uint64_t reduction_workspace_size = Params::get_reduction_workspace_size(
reduction_tiles, to_gemm_coord(TileShape{}), sizeof_bits<ElementAccumulator>::value, num_accumulator_mtxs);
BarrierType* lock_workspace = reinterpret_cast<BarrierType*>(
reinterpret_cast<uint8_t*>(params.reduction_workspace_) + reduction_workspace_size);
if (work_tile_info.is_reduction_unit()) {
plus<AccumulatorArrayT> add_fragments;
auto peer_offset = size(accumulators) * num_barriers * BarrierManager::ThreadCount;
uint64_t peer_offset = size(accumulators) * num_barriers * BarrierManager::ThreadCount;
// Wait until the peers collaborating on this output tile have all written
// their accumulators to workspace.
uint32_t num_peers = last_peer_id - first_peer_id + 1;
BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, num_peers);
// Load the first peer's data
BlockStripedReduceT::load(*accumulator_array, reduction_workspace_array, barrier_group_thread_idx);
for (int i = 1; i < num_peers; ++i) {
for (uint64_t i = 1; i < num_peers; ++i) {
// Load peer fragment
AccumulatorArrayT addend_fragment;
auto peer_reduction_workspace = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace + (i * peer_offset));
@@ -487,7 +492,7 @@ public:
// If separate reduction is being performed, each participating stream-K unit increments the barrier
// by only 1. Otherwise, increment by the K tile count that this unit has processed.
int32_t increment = params.requires_separate_reduction() ? 1 : work_tile_info.k_tile_count;
uint32_t increment = params.requires_separate_reduction() ? 1 : work_tile_info.k_tile_count;
// Signal our arrival
BarrierManager::arrive_inc(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, increment);
@@ -524,7 +529,7 @@ public:
// Returns the linearized index of the output tile corresponding to the tile with offset [L, M, K]
CUTLASS_DEVICE
static int
static uint64_t
output_tile_index(Params const& params, WorkTileInfo const& work_tile_info) {
uint64_t linear_idx_in_batch = UnderlyingScheduler::get_linear_idx_from_m_and_n(
work_tile_info.M_idx, work_tile_info.N_idx,
@@ -546,7 +551,8 @@ public:
ProblemShape problem_shape,
KernelHardwareInfo const& hw_info,
uint32_t mma_warp_groups,
const uint32_t epilogue_subtile = 1) {
const uint32_t epilogue_subtile = 1,
[[maybe_unused]] uint32_t num_accumulator_mtxs = 1) {
auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
@@ -583,6 +589,7 @@ public:
KernelHardwareInfo const& hw_info,
uint32_t mma_warp_groups,
const uint32_t epilogue_subtile = 1,
[[maybe_unused]] uint32_t num_accumulator_mtxs = 1,
CudaHostAdapter* cuda_adapter = nullptr) {
auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
@@ -616,7 +623,7 @@ public:
template <class ProblemShape>
CUTLASS_HOST_DEVICE
static int
static uint32_t
get_work_k_tile_count(WorkTileInfo const& work_tile_info, ProblemShape, TileShape) {
return work_tile_info.k_tile_count;
}
@@ -632,11 +639,11 @@ public:
auto
fetch_next_work(WorkTileInfo work_tile_info) {
if (continue_current_work(work_tile_info)) {
return work_tile_info;
return cute::make_tuple(work_tile_info, true);
}
advance_to_next_work();
return get_current_work();
return cute::make_tuple(get_current_work(), true);
}
// Returns the initial work tile info that will be computed over
@@ -665,7 +672,7 @@ private:
// Separate-reduction work
auto cluster_size = params.get_cluster_size();
// Divide up the linearized separate reduction units into clusters
auto cluster_linear_reduction_unit_idx = params.div_cluster_size((linear_idx - params.units_per_problem_));
uint64_t cluster_linear_reduction_unit_idx = params.div_cluster_size((linear_idx - params.units_per_problem_));
uint64_t cluster_tile_idx, epi_subtile_idx;
params.divmod_epilogue_subtile_(cluster_tile_idx, epi_subtile_idx, cluster_linear_reduction_unit_idx);
// Bring the linearized tile ID back into the space of tiles, rather than clusters
@@ -697,35 +704,35 @@ private:
// 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_linear_work_idx = params.div_cluster_size(linear_idx);
uint64_t cluster_linear_work_idx = params.div_cluster_size(linear_idx);
uint64_t group_idx;
params.divmod_sk_groups_(cluster_linear_work_idx, group_idx, cluster_linear_work_idx);
// Determine whether we are in a "big group" that will process an additional
// stream-K cluster tile.
auto sk_cluster_tiles = params.div_cluster_size(params.sk_tiles_);
auto sk_cluster_tiles_in_group = params.divmod_sk_groups_.divide(sk_cluster_tiles);
uint64_t sk_cluster_tiles = params.div_cluster_size(params.sk_tiles_);
uint64_t sk_cluster_tiles_in_group = params.divmod_sk_groups_.divide(sk_cluster_tiles);
if (group_idx < params.big_groups_) {
++sk_cluster_tiles_in_group;
}
// Determine whether we are in a "big unit" within the group, that will process
// an additional K chunk in the group.
auto sk_tiles_in_group = sk_cluster_tiles_in_group * params.get_cluster_size();
auto k_tiles_in_group = sk_tiles_in_group * params.divmod_tiles_per_output_tile_.divisor;
auto k_tiles_per_unit_in_group = params.divmod_sk_units_per_group_.divide(k_tiles_in_group);
auto big_units_in_group = params.div_cluster_size(
uint64_t sk_tiles_in_group = sk_cluster_tiles_in_group * params.get_cluster_size();
uint64_t k_tiles_in_group = sk_tiles_in_group * params.divmod_tiles_per_output_tile_.divisor;
uint64_t k_tiles_per_unit_in_group = params.divmod_sk_units_per_group_.divide(k_tiles_in_group);
uint64_t big_units_in_group = params.div_cluster_size(
k_tiles_in_group - (k_tiles_per_unit_in_group * params.divmod_sk_units_per_group_.divisor));
uint64_t split;
params.divmod_clusters_mnl_(split, cluster_linear_work_idx, cluster_linear_work_idx);
bool is_split_k = params.divmod_splits_.divisor > 1;
auto big_unit_cmp_lhs = is_split_k ? split : cluster_linear_work_idx;
auto big_unit_cmp_rhs = is_split_k ? params.big_units_ : big_units_in_group;
auto linear_idx_mult = is_split_k ? params.divmod_tiles_per_output_tile_.divisor : k_tiles_per_unit_in_group;
auto k_tiles_per_split = is_split_k ? params.divmod_k_tiles_per_sk_unit_.divisor : k_tiles_per_unit_in_group;
uint64_t big_unit_cmp_lhs = is_split_k ? split : cluster_linear_work_idx;
uint64_t big_unit_cmp_rhs = is_split_k ? params.big_units_ : big_units_in_group;
uint64_t linear_idx_mult = is_split_k ? params.divmod_tiles_per_output_tile_.divisor : k_tiles_per_unit_in_group;
uint64_t k_tiles_per_split = is_split_k ? params.divmod_k_tiles_per_sk_unit_.divisor : k_tiles_per_unit_in_group;
// Determine the starting k iteration computed by this stream-K work unit
uint32_t unit_iter_start = (linear_idx_mult * cluster_linear_work_idx) +
@@ -889,14 +896,14 @@ private:
CUTLASS_HOST_DEVICE
static auto
tile_peer_range(Params const& params, uint32_t tile_idx, uint32_t cur_k_tile) {
auto tile_idx_in_cluster_path = params.div_cluster_size(tile_idx);
auto start_k_tile = params.divmod_tiles_per_output_tile_.divisor * tile_idx_in_cluster_path;
auto end_k_tile = start_k_tile + params.divmod_tiles_per_output_tile_.divisor - 1;
auto big_unit_k_tiles = params.big_units_ * (params.divmod_k_tiles_per_sk_unit_.divisor + 1);
uint32_t tile_idx_in_cluster_path = params.div_cluster_size(tile_idx);
uint32_t start_k_tile = params.divmod_tiles_per_output_tile_.divisor * tile_idx_in_cluster_path;
uint32_t end_k_tile = start_k_tile + params.divmod_tiles_per_output_tile_.divisor - 1;
uint32_t big_unit_k_tiles = params.big_units_ * (params.divmod_k_tiles_per_sk_unit_.divisor + 1);
auto adjust_unit = [&](uint32_t k_tile, uint32_t unit_idx, uint32_t k_tiles_per_unit) {
auto unit_k_start = unit_idx * k_tiles_per_unit;
auto unit_k_end = unit_k_start + k_tiles_per_unit;
uint32_t unit_k_start = unit_idx * k_tiles_per_unit;
uint32_t unit_k_end = unit_k_start + k_tiles_per_unit;
if (k_tile - start_k_tile < Params::min_iters_per_sk_unit_ &&
unit_k_end - start_k_tile < Params::min_iters_per_sk_unit_) {
// k_tile is within the first min_iters_per_sk_unit_ K tiles of this output tile,
@@ -920,13 +927,13 @@ private:
auto find_unit = [&](uint32_t k_tile) {
if (k_tile < big_unit_k_tiles) {
// The tile is within the "big unit range"
auto unit_idx = params.divmod_k_tiles_per_sk_big_unit_.divide(k_tile);
uint32_t unit_idx = params.divmod_k_tiles_per_sk_big_unit_.divide(k_tile);
return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, params.divmod_k_tiles_per_sk_big_unit_.divisor));
}
else {
// The tile is after the "big unit range." Account for this by finding the "normal unit"
// that it belongs to, and then offsetting by the number of big units
auto unit_idx = params.divmod_k_tiles_per_sk_unit_.divide(k_tile - big_unit_k_tiles) + params.big_units_;
uint32_t unit_idx = params.divmod_k_tiles_per_sk_unit_.divide(k_tile - big_unit_k_tiles) + params.big_units_;
return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, params.divmod_k_tiles_per_sk_unit_.divisor));
}
};
include/cutlass/gemm/kernel/static_tile_scheduler.hpp
+82 -28
View File
@@ -46,9 +46,6 @@ namespace cutlass::gemm::kernel::detail {
// This is a CRTP base class for the actual tile schedulers.
template<class Subclass>
class StaticPersistentTileScheduler {
//
// Data members
//
private:
uint64_t current_work_linear_idx_;
@@ -89,6 +86,8 @@ public:
using Params = PersistentTileSchedulerSm90Params;
using RasterOrder = typename Params::RasterOrder;
using RasterOrderOptions = typename Params::RasterOrderOptions;
static constexpr bool IsDynamicPersistent = false;
public:
struct Arguments {
int max_swizzle_size = 1;
@@ -98,13 +97,14 @@ public:
template <class ProblemShapeMNKL, class TileShape, class ClusterShape>
static Params
to_underlying_arguments(
ProblemShapeMNKL problem_shape_mnkl,
TileShape tile_shape,
ClusterShape cluster_shape,
[[maybe_unused]] KernelHardwareInfo const& hw_info,
Arguments const& arguments,
[[maybe_unused]] void* workspace=nullptr,
[[maybe_unused]] const uint32_t epilogue_subtile = 1) {
ProblemShapeMNKL problem_shape_mnkl,
TileShape tile_shape,
ClusterShape cluster_shape,
[[maybe_unused]] KernelHardwareInfo const& hw_info,
Arguments const& arguments,
[[maybe_unused]] void* workspace=nullptr,
[[maybe_unused]] const uint32_t epilogue_subtile = 1,
[[maybe_unused]] uint32_t ktile_start_alignment_count = 1u) {
// We only need the tile and cluster shape during scheduler setup, so let FTAD do the magic
static_assert(cute::is_static<TileShape>::value);
@@ -249,6 +249,39 @@ public:
);
}
// Reloaded interface that receives WorkTileInfo to deduce next work.
// Kernel helper function to get next work tile
CUTLASS_DEVICE
auto
fetch_next_work(WorkTileInfo work_tile_info) {
if (continue_current_work(work_tile_info)) {
return cute::make_tuple(work_tile_info, true);
}
advance_to_next_work();
return cute::make_tuple(get_current_work(), true);
}
// Given the inputs, computes the total number of output blocks over which this problem will compute.
// Note that this is only the logical size of our grid, not the physical grid we will actually launch.
template<class ProblemShapeMNKL, class TileShape, class AtomThrShape, class ClusterShape>
CUTLASS_HOST_DEVICE static
dim3
get_tiled_cta_shape_mnl(ProblemShapeMNKL problem_shape_mnkl,
TileShape tile_shape_mnk,
AtomThrShape atom_thr_shape_mnk,
ClusterShape cluster_shape_mnk) {
auto [tiles_m, tiles_n, tiles_l] = product_each(ceil_div(select<0,1,3>(problem_shape_mnkl), take<0,2>(tile_shape_mnk)));
auto cta_m = round_nearest(tiles_m * size<0>(atom_thr_shape_mnk), size<0>(cluster_shape_mnk));
auto cta_n = round_nearest(tiles_n * size<1>(atom_thr_shape_mnk), size<1>(cluster_shape_mnk));
return Params::get_tiled_cta_shape_mnl(
to_gemm_coord(problem_shape_mnkl),
to_gemm_coord(cluster_shape_mnk),
cta_m, cta_n
);
}
CUTLASS_DEVICE
static auto
work_tile_to_cta_coord(WorkTileInfo work_tile_info) {
@@ -262,17 +295,31 @@ public:
);
}
CUTLASS_DEVICE
static auto
work_tile_to_cta_coord(WorkTileInfo work_tile_info, dim3 block_id_in_cluster) {
// Get every cta coord in three dimensions of the cluster
auto [cta_m_in_cluster, cta_n_in_cluster, cta_l_in_cluster] = block_id_in_cluster;
return make_coord(
work_tile_info.M_idx + static_cast<int32_t>(cta_m_in_cluster),
work_tile_info.N_idx + static_cast<int32_t>(cta_n_in_cluster),
_,
work_tile_info.L_idx + static_cast<int32_t>(cta_l_in_cluster)
);
}
// Given the inputs, computes the physical grid we should launch.
template<class ProblemShapeMNKL, class BlockShape, class ClusterShape>
CUTLASS_HOST_DEVICE static
dim3
get_grid_shape(
ProblemShapeMNKL problem_shape_mnk,
BlockShape cta_shape,
ClusterShape cluster_shape,
KernelHardwareInfo hw_info,
Arguments arguments,
bool truncate_by_problem_size=true) {
[[maybe_unused]] Params const& params,
ProblemShapeMNKL problem_shape_mnk,
BlockShape cta_shape,
ClusterShape cluster_shape,
KernelHardwareInfo hw_info,
Arguments arguments = Arguments{},
bool truncate_by_problem_size=true) {
auto problem_shape_mnkl = cute::append<4>(problem_shape_mnk, cute::Int<1>{});
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, cta_shape, cluster_shape);
@@ -288,19 +335,17 @@ public:
}
// Given the inputs, computes the physical grid we should launch.
template<class ProblemShapeMNKL, class BlockShape, class ClusterShape>
CUTLASS_HOST_DEVICE static
dim3
template<class ProblemShapeMNKL, class TileShape, class AtomThrShape, class ClusterShape>
static dim3
get_grid_shape(
Params const& params,
ProblemShapeMNKL problem_shape_mnk,
BlockShape cta_shape,
ClusterShape cluster_shape,
KernelHardwareInfo hw_info) {
auto problem_shape_mnkl = cute::append<4>(problem_shape_mnk, cute::Int<1>{});
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, cta_shape, cluster_shape);
Params const& params,
ProblemShapeMNKL problem_shape_mnkl,
TileShape tile_shape_mnk,
AtomThrShape atom_thr_shape_mnk,
ClusterShape cluster_shape_mnk,
KernelHardwareInfo hw_info) {
dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape_mnk, atom_thr_shape_mnk, cluster_shape_mnk);
Arguments args{};
if constexpr (!std::is_const_v<decltype(args.max_swizzle_size)>) {
args.max_swizzle_size = 1 << params.log_swizzle_size_;
@@ -309,7 +354,7 @@ public:
return Params::get_grid_shape(
problem_blocks,
to_gemm_coord(cluster_shape),
to_gemm_coord(cluster_shape_mnk),
hw_info,
args.max_swizzle_size,
args.raster_order,
@@ -368,6 +413,14 @@ public:
return false;
}
template <class ProblemShapeMNKL, class TileShape, class Shape>
CUTLASS_DEVICE
auto
get_k_tile_iterator(WorkTileInfo const& work_tile_info, ProblemShapeMNKL problem_shape_MNKL, TileShape tile_shape, Shape) {
auto k_tiles = cute::ceil_div(cute::get<2>(problem_shape_MNKL), cute::get<2>(tile_shape));
return cute::make_coord_iterator(k_tiles);
}
template <class ProblemShape, class TileShape>
CUTLASS_HOST_DEVICE
static int
@@ -430,6 +483,7 @@ public:
requires_separate_reduction(Params const& params) {
return false;
}
public:
// Sink scheduler params as a member
Params scheduler_params;
include/cutlass/gemm/kernel/tile_scheduler.hpp
+20 -20
View File
@@ -86,10 +86,10 @@ template <
class ClusterShape
>
struct TileSchedulerSelector<
PersistentScheduler,
ArchTag,
TileShape,
ClusterShape
PersistentScheduler,
ArchTag,
TileShape,
ClusterShape
> {
using Scheduler = PersistentTileSchedulerSm90;
};
@@ -101,16 +101,16 @@ template <
class ClusterShape
>
struct TileSchedulerSelector<
void,
ArchTag,
TileShape,
ClusterShape
> {
using Scheduler = typename TileSchedulerSelector<
PersistentScheduler,
void,
ArchTag,
TileShape,
ClusterShape
> {
using Scheduler = typename TileSchedulerSelector<
PersistentScheduler,
ArchTag,
TileShape,
ClusterShape
>::Scheduler;
};
@@ -119,10 +119,10 @@ template <
class ClusterShape
>
struct TileSchedulerSelector<
StreamKScheduler,
arch::Sm90,
TileShape,
ClusterShape
StreamKScheduler,
arch::Sm90,
TileShape,
ClusterShape
> {
using Scheduler = PersistentTileSchedulerSm90StreamK<TileShape, ClusterShape>;
};
@@ -133,11 +133,11 @@ template <
, class GroupProblemShape
>
struct TileSchedulerSelector<
GroupScheduler,
arch::Sm90,
TileShape,
ClusterShape
, GroupProblemShape
GroupScheduler,
arch::Sm90,
TileShape,
ClusterShape
, GroupProblemShape
> {
using Scheduler = PersistentTileSchedulerSm90Group<GroupProblemShape>;
};
include/cutlass/gemm/kernel/tile_scheduler_params.h
+3 -6
View File
@@ -1046,12 +1046,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
// Returns whether the kernel is configured in a manner for which separate reduction should be used
CUTLASS_HOST_DEVICE
static bool
should_perform_separate_reduction(uint32_t epilogue_subtile, uint64_t sk_units, uint64_t sk_tiles, uint64_t dp_tiles, uint64_t ctas_per_wave) {
// We perform separate reduction if we have fewer than one wave of output tiles
// and each output tile is covered by at least to stream-K units. When sk_units is
// multiple of sk_tiles, will choose basic split-k path instead of separate reduction for now.
return (epilogue_subtile != 1) && (dp_tiles == 0) && (sk_units > 2u * sk_tiles) &&
(sk_units + sk_tiles * epilogue_subtile <= ctas_per_wave);
should_perform_separate_reduction(uint32_t, uint64_t, uint64_t, uint64_t, uint64_t) {
// Separate reduction is temporarily disabled, pending fixes
return false;
}
// Get the amount of scratch workspace needed for the kernel. This variant of the method should only be used when