879 lines
38 KiB
C++
879 lines
38 KiB
C++
/***************************************************************************************************
|
|
* Copyright (c) 2023 - 2026 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/tensor.hpp"
|
|
#include "cute/arch/tmem_allocator_sm100.hpp"
|
|
#include "cute/arch/cluster_sm90.hpp"
|
|
|
|
#include "cutlass/arch/arch.h"
|
|
#include "cutlass/arch/grid_dependency_control.h"
|
|
#include "cutlass/conv/detail.hpp"
|
|
#include "cutlass/conv/convolution.h"
|
|
#include "cutlass/conv/dispatch_policy.hpp"
|
|
#include "cutlass/gemm/kernel/tile_scheduler.hpp"
|
|
#include "cutlass/pipeline/sm100_pipeline.hpp"
|
|
#include "cutlass/detail/sm100_tmem_helper.hpp"
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
|
|
namespace cutlass::conv::kernel {
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
|
|
template <
|
|
class ProblemShape_,
|
|
class CollectiveMainloop_,
|
|
class CollectiveEpilogue_,
|
|
class TileSchedulerTag_
|
|
>
|
|
class ConvUniversal<
|
|
ProblemShape_,
|
|
CollectiveMainloop_,
|
|
CollectiveEpilogue_,
|
|
TileSchedulerTag_,
|
|
cute::enable_if_t<cute::is_base_of_v<KernelImplicitTmaWarpSpecializedSm100,
|
|
typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
|
|
{
|
|
public:
|
|
//
|
|
// Type Aliases
|
|
//
|
|
|
|
// Mainloop derived types
|
|
using ProblemShape = ProblemShape_;
|
|
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;
|
|
using CtaShape_MNK = typename CollectiveMainloop::CtaShape_MNK;
|
|
using AtomThrShapeMNK = typename CollectiveMainloop::AtomThrShapeMNK;
|
|
static constexpr int NumSpatialDimensions = CollectiveMainloop::NumSpatialDimensions;
|
|
static constexpr bool is_grouped_wgrad = CollectiveMainloop::is_grouped_wgrad;
|
|
static constexpr bool IsComplex = false;
|
|
static_assert(ArchTag::kMinComputeCapability >= 100);
|
|
|
|
// 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 constexpr bool IsGdcEnabled = cutlass::arch::IsGdcGloballyEnabled;
|
|
// TileID scheduler
|
|
// CLC pipeline depth determines how many waves (stages-1) the scheduler can race ahead
|
|
static constexpr uint32_t SchedulerPipelineStageCount = DispatchPolicy::Schedule::SchedulerPipelineStageCount;
|
|
static constexpr uint32_t AccumulatorPipelineStageCount = DispatchPolicy::Schedule::AccumulatorPipelineStageCount;
|
|
|
|
using TileSchedulerTag = TileSchedulerTag_;
|
|
using TileScheduler = typename cutlass::gemm::kernel::detail::TileSchedulerSelector<
|
|
TileSchedulerTag, ArchTag, CtaShape_MNK, ClusterShape, SchedulerPipelineStageCount>::Scheduler;
|
|
using TileSchedulerArguments = typename TileScheduler::Arguments;
|
|
using TileSchedulerParams = typename TileScheduler::Params;
|
|
|
|
static constexpr bool IsDynamicCluster = not cute::is_static_v<ClusterShape>;
|
|
|
|
// Warp specialization thread count per threadblock
|
|
static constexpr uint32_t NumSchedThreads = NumThreadsPerWarp; // 1 warp
|
|
static constexpr uint32_t NumMMAThreads = NumThreadsPerWarp; // 1 warp
|
|
static constexpr uint32_t NumMainloopLoadThreads = NumThreadsPerWarp; // 1 warp
|
|
static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp; // 1 warp
|
|
static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount;
|
|
static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp;
|
|
|
|
static constexpr uint32_t MaxThreadsPerBlock = NumSchedThreads +
|
|
NumMainloopLoadThreads + NumMMAThreads +
|
|
NumEpilogueLoadThreads + NumEpilogueThreads;
|
|
static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
|
|
static constexpr uint32_t NumFixupBarriers = 1;
|
|
|
|
// Pipelines and pipeline states
|
|
static constexpr uint32_t CLCResponseSize = sizeof(typename TileScheduler::CLCResponse);
|
|
|
|
// Pipeline and pipeline state types
|
|
using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
|
|
using MainloopPipelineState = typename CollectiveMainloop::MainloopPipelineState;
|
|
|
|
using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
|
|
using EpiLoadPipelineState = typename CollectiveEpilogue::LoadPipelineState;
|
|
|
|
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
|
|
using EpiStorePipelineState = typename CollectiveEpilogue::StorePipelineState;
|
|
|
|
using LoadOrderBarrier = cutlass::OrderedSequenceBarrier<1,2>;
|
|
|
|
using AccumulatorPipeline = cutlass::PipelineUmmaAsync<AccumulatorPipelineStageCount, AtomThrShapeMNK>;
|
|
using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState;
|
|
|
|
using CLCPipeline = cutlass::PipelineCLCFetchAsync<SchedulerPipelineStageCount, ClusterShape>;
|
|
using CLCPipelineState = cutlass::PipelineDetail::PipelineCLCFetchAsyncPipelineState<SchedulerPipelineStageCount>;
|
|
using CLCPipelineSharedStorage = cutlass::PipelineDetail::PipelineCLCFetchAsyncSharedStorage<SchedulerPipelineStageCount>;
|
|
|
|
using TmemAllocator = cute::conditional_t<cute::size(cute::shape<0>(typename TiledMma::ThrLayoutVMNK{})) == 1,
|
|
cute::TMEM::Allocator1Sm, cute::TMEM::Allocator2Sm>;
|
|
|
|
// Kernel level shared memory storage
|
|
struct SharedStorage {
|
|
struct PipelineStorage : cute::aligned_struct<16, _1> {
|
|
using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
|
|
using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
|
|
using LoadOrderBarrierStorage = typename LoadOrderBarrier::SharedStorage;
|
|
using CLCPipelineStorage = CLCPipelineSharedStorage;
|
|
using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage;
|
|
|
|
alignas(16) MainloopPipelineStorage mainloop;
|
|
alignas(16) EpiLoadPipelineStorage epi_load;
|
|
alignas(16) LoadOrderBarrierStorage load_order;
|
|
alignas(16) CLCPipelineStorage clc;
|
|
alignas(16) AccumulatorPipelineStorage accumulator;
|
|
alignas(16) arch::ClusterBarrier tmem_dealloc;
|
|
} pipelines;
|
|
|
|
alignas(16) typename TileScheduler::CLCResponse clc_response[SchedulerPipelineStageCount];
|
|
uint32_t tmem_base_ptr;
|
|
|
|
struct TensorStorage : cute::aligned_struct<128, _1> {
|
|
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
|
|
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
|
|
|
|
EpilogueTensorStorage epilogue;
|
|
MainloopTensorStorage mainloop;
|
|
} tensors;
|
|
};
|
|
|
|
static constexpr int SharedStorageSize = sizeof(SharedStorage);
|
|
static_assert(SharedStorageSize <= cutlass::arch::sm100_smem_capacity_bytes, "SMEM usage exceeded capacity.");
|
|
|
|
// Host facing host arguments
|
|
struct Arguments {
|
|
ProblemShape problem_shape{};
|
|
MainloopArguments mainloop{};
|
|
EpilogueArguments epilogue{};
|
|
KernelHardwareInfo hw_info{};
|
|
TileSchedulerArguments scheduler{};
|
|
};
|
|
|
|
// Kernel device entry point API
|
|
struct Params {
|
|
using ProblemShapeMNKL = decltype(CollectiveMainloop::get_problem_shape_MNKL(ProblemShape{}));
|
|
ProblemShapeMNKL problem_shape;
|
|
MainloopParams mainloop;
|
|
EpilogueParams epilogue;
|
|
TileSchedulerParams scheduler;
|
|
KernelHardwareInfo hw_info{};
|
|
};
|
|
|
|
enum class WarpCategory : int32_t {
|
|
MMA = 0,
|
|
Sched = 1,
|
|
MainloopLoad = 2,
|
|
EpilogueLoad = 3,
|
|
Epilogue = 4
|
|
};
|
|
|
|
struct IsParticipant {
|
|
uint32_t mma = false;
|
|
uint32_t sched = false;
|
|
uint32_t main_load = false;
|
|
uint32_t epi_load = false;
|
|
uint32_t epilogue = false;
|
|
};
|
|
|
|
//
|
|
// Methods
|
|
//
|
|
// Map user facing arguments to device facing params
|
|
CUTLASS_HOST
|
|
static Params
|
|
to_underlying_arguments(Arguments const& args, void* workspace) {
|
|
static constexpr uint32_t NumEpilogueSubTiles = 1;
|
|
|
|
auto problem_shape_mnkl = CollectiveMainloop::get_problem_shape_MNKL(args.problem_shape);
|
|
|
|
auto mainloop_params = CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace, args.hw_info);
|
|
|
|
// Calculate workspace pointers
|
|
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
|
size_t workspace_offset = 0;
|
|
|
|
// Epilogue
|
|
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);
|
|
|
|
// Tile scheduler
|
|
void* scheduler_workspace = workspace_ptr + workspace_offset;
|
|
workspace_offset += TileScheduler::template get_workspace_size<decltype(problem_shape_mnkl), ElementAccumulator>(
|
|
args.scheduler, problem_shape_mnkl, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
|
|
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
|
|
|
return {
|
|
problem_shape_mnkl,
|
|
mainloop_params,
|
|
CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, epilogue_workspace),
|
|
TileScheduler::to_underlying_arguments(
|
|
args.problem_shape, TileShape{}, AtomThrShapeMNK{}, ClusterShape{},
|
|
args.hw_info, args.scheduler, scheduler_workspace),
|
|
args.hw_info
|
|
};
|
|
}
|
|
|
|
CUTLASS_HOST
|
|
static bool
|
|
can_implement(Arguments const& args) {
|
|
bool implementable = true;
|
|
implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
|
|
implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
|
|
implementable &= TileScheduler::can_implement(args.scheduler);
|
|
|
|
if constexpr (IsDynamicCluster) {
|
|
static constexpr int MaxClusterSize = 16;
|
|
implementable &= size(args.hw_info.cluster_shape) <= MaxClusterSize;
|
|
implementable &= size(args.hw_info.cluster_shape_fallback) <= MaxClusterSize;
|
|
// Early return if cluster shape validation failed to avoid division by zero below
|
|
if (not cutlass::detail::preferred_cluster_can_implement<AtomThrShapeMNK>(args.hw_info.cluster_shape, args.hw_info.cluster_shape_fallback)) {
|
|
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: Invalid dynamic cluster shape\n");
|
|
return false;
|
|
}
|
|
}
|
|
|
|
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape);
|
|
auto cluster_shape_fallback = cutlass::detail::select_cluster_shape(ClusterShape{}, args.hw_info.cluster_shape_fallback);
|
|
|
|
// implicit gemm B tile can be small for conv, ensure multicast smem offsets are 128B aligned
|
|
int multicast_b_bits = (size<1>(TileShape{}) * size<2>(TileShape{}) / size<0>(cluster_shape)) * sizeof_bits_v<ElementB>;
|
|
int multicast_b_fallback_bits = (size<1>(TileShape{}) * size<2>(TileShape{}) / size<0>(cluster_shape_fallback)) * sizeof_bits_v<ElementB>;
|
|
implementable &= multicast_b_bits % (128*8) == 0 && multicast_b_fallback_bits % (128*8) == 0;
|
|
if (not implementable) {
|
|
CUTLASS_TRACE_HOST(" CAN IMPLEMENT: multicast size too large for B tile\n");
|
|
return false;
|
|
}
|
|
|
|
if constexpr (is_grouped_wgrad) {
|
|
implementable &= size<0>(cluster_shape) == 1 && size<0>(cluster_shape_fallback) == 1;
|
|
|
|
if (!implementable) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return implementable;
|
|
}
|
|
|
|
CUTLASS_HOST
|
|
static size_t
|
|
get_workspace_size(Arguments const& args) {
|
|
static constexpr uint32_t NumEpilogueSubTiles = 1;
|
|
size_t workspace_size = 0;
|
|
auto linear_problem_shape_MNKL = cutlass::conv::detail::get_linearized_problem_shape_MNKL(args.problem_shape);
|
|
|
|
// Epilogue
|
|
workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
|
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
|
|
|
// Tile scheduler
|
|
workspace_size += TileScheduler::template get_workspace_size<decltype(linear_problem_shape_MNKL), ElementAccumulator>(
|
|
args.scheduler, linear_problem_shape_MNKL, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs);
|
|
workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
|
|
|
|
return workspace_size;
|
|
}
|
|
|
|
CUTLASS_HOST
|
|
static cutlass::Status
|
|
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
|
|
CudaHostAdapter* cuda_adapter = nullptr) {
|
|
static constexpr uint32_t NumEpilogueSubTiles = 1;
|
|
auto linear_problem_shape_MNKL = cutlass::conv::detail::get_linearized_problem_shape_MNKL(args.problem_shape);
|
|
Status status = Status::kSuccess;
|
|
uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
|
|
size_t workspace_offset = 0;
|
|
|
|
// Epilogue
|
|
status = CollectiveEpilogue::initialize_workspace(
|
|
args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
|
|
|
|
workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
|
|
workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
|
|
if (status != Status::kSuccess) {
|
|
return status;
|
|
}
|
|
|
|
// Tile scheduler
|
|
status = TileScheduler::template initialize_workspace
|
|
<decltype(linear_problem_shape_MNKL), ElementAccumulator>(
|
|
args.scheduler, workspace_ptr + workspace_offset, stream, linear_problem_shape_MNKL,
|
|
args.hw_info, NumFixupBarriers, NumEpilogueSubTiles, CollectiveEpilogue::NumAccumulatorMtxs, cuda_adapter);
|
|
|
|
workspace_offset += TileScheduler::template get_workspace_size
|
|
<decltype(linear_problem_shape_MNKL), ElementAccumulator>(
|
|
args.scheduler, linear_problem_shape_MNKL, args.hw_info, NumFixupBarriers, NumEpilogueSubTiles,
|
|
CollectiveEpilogue::NumAccumulatorMtxs);
|
|
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
|
|
CUTLASS_HOST
|
|
static dim3
|
|
get_grid_shape(Params const& params) {
|
|
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{}, params.hw_info.cluster_shape);
|
|
|
|
return TileScheduler::get_grid_shape(
|
|
params.scheduler,
|
|
params.problem_shape,
|
|
TileShape{},
|
|
AtomThrShapeMNK{},
|
|
cluster_shape
|
|
,params.hw_info
|
|
);
|
|
}
|
|
|
|
CUTLASS_HOST
|
|
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;
|
|
|
|
// Separate out problem shape for convenience
|
|
auto problem_shape_MNKL = append<4>(params.problem_shape, _1{});
|
|
auto [M, N, K, L] = problem_shape_MNKL;
|
|
|
|
// Account for more than one epilogue warp
|
|
int warp_idx = canonical_warp_idx_sync();
|
|
WarpCategory warp_category = warp_idx < static_cast<int>(WarpCategory::Epilogue) ? WarpCategory(warp_idx)
|
|
: WarpCategory::Epilogue;
|
|
|
|
uint32_t lane_predicate = cute::elect_one_sync();
|
|
auto cluster_shape = cutlass::detail::select_cluster_shape(ClusterShape{});
|
|
int cluster_size = size(cluster_shape);
|
|
uint32_t cta_rank_in_cluster = cute::block_rank_in_cluster();
|
|
bool is_first_cta_in_cluster = cta_rank_in_cluster == 0;
|
|
int cta_coord_v = cta_rank_in_cluster % size<0>(typename TiledMma::AtomThrID{});
|
|
bool is_mma_leader_cta = cta_coord_v == 0;
|
|
constexpr bool has_mma_peer_cta = size(AtomThrShapeMNK{}) == 2;
|
|
[[maybe_unused]] uint32_t mma_peer_cta_rank = has_mma_peer_cta ? cta_rank_in_cluster ^ 1 : cta_rank_in_cluster;
|
|
|
|
// Kernel level shared memory storage
|
|
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
|
|
|
|
// In a warp specialized kernel, collectives expose data movement and compute operations separately
|
|
CollectiveMainloop collective_mainloop(params.mainloop, cluster_shape, cta_rank_in_cluster);
|
|
CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
|
|
|
|
// Issue Tma Descriptor Prefetch from a single thread
|
|
if ((warp_category == WarpCategory::Sched) && lane_predicate) {
|
|
collective_mainloop.prefetch_tma_descriptors();
|
|
}
|
|
if ((warp_category == WarpCategory::EpilogueLoad) && lane_predicate) {
|
|
collective_epilogue.prefetch_tma_descriptors(params.epilogue);
|
|
}
|
|
|
|
// Do we load source tensor C or other aux inputs
|
|
bool is_epi_load_needed = collective_epilogue.is_producer_load_needed();
|
|
IsParticipant is_participant = {
|
|
(warp_category == WarpCategory::MMA), // mma
|
|
(warp_category == WarpCategory::Sched) && is_first_cta_in_cluster, // sched
|
|
(warp_category == WarpCategory::MainloopLoad), // main_load
|
|
(warp_category == WarpCategory::EpilogueLoad) && is_epi_load_needed, // epi_load
|
|
(warp_category == WarpCategory::Epilogue) // epilogue
|
|
};
|
|
|
|
// Mainloop Load pipeline
|
|
typename MainloopPipeline::Params mainloop_pipeline_params;
|
|
if (WarpCategory::MainloopLoad == warp_category) {
|
|
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
|
|
}
|
|
if (WarpCategory::MMA == warp_category) {
|
|
mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
|
|
}
|
|
mainloop_pipeline_params.is_leader = lane_predicate && is_mma_leader_cta && is_participant.main_load;
|
|
mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes;
|
|
mainloop_pipeline_params.initializing_warp = 0;
|
|
MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop,
|
|
mainloop_pipeline_params,
|
|
cluster_shape,
|
|
cute::true_type{}, // Perform barrier init
|
|
cute::false_type{}); // Delay mask calculation
|
|
|
|
// Epilogue Load pipeline
|
|
typename EpiLoadPipeline::Params epi_load_pipeline_params;
|
|
if (WarpCategory::EpilogueLoad == warp_category) {
|
|
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
|
|
}
|
|
if (WarpCategory::Epilogue == warp_category) {
|
|
epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
|
|
}
|
|
epi_load_pipeline_params.dst_blockid = cta_rank_in_cluster;
|
|
epi_load_pipeline_params.producer_arv_count = NumEpilogueLoadThreads;
|
|
epi_load_pipeline_params.consumer_arv_count = NumEpilogueThreads;
|
|
epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes;
|
|
epi_load_pipeline_params.initializing_warp = 1;
|
|
EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
|
|
|
|
// Epilogue Store pipeline
|
|
typename EpiStorePipeline::Params epi_store_pipeline_params;
|
|
epi_store_pipeline_params.always_wait = true;
|
|
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
|
|
|
|
// Load order barrier
|
|
typename LoadOrderBarrier::Params load_order_barrier_params;
|
|
load_order_barrier_params.group_id = (warp_category == WarpCategory::MainloopLoad) ? 0 : 1;
|
|
load_order_barrier_params.group_size = NumMainloopLoadThreads;
|
|
load_order_barrier_params.initializing_warp = 3;
|
|
LoadOrderBarrier load_order_barrier(shared_storage.pipelines.load_order, load_order_barrier_params);
|
|
|
|
// CLC pipeline
|
|
typename CLCPipeline::Params clc_pipeline_params;
|
|
if (WarpCategory::Sched == warp_category) {
|
|
clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer;
|
|
}
|
|
else {
|
|
clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer;
|
|
}
|
|
clc_pipeline_params.producer_blockid = 0;
|
|
clc_pipeline_params.producer_arv_count = 1;
|
|
clc_pipeline_params.consumer_arv_count = NumSchedThreads + cluster_size *
|
|
(NumMainloopLoadThreads + NumEpilogueThreads + NumMMAThreads);
|
|
if (is_epi_load_needed) {
|
|
clc_pipeline_params.consumer_arv_count += cluster_size * NumEpilogueLoadThreads;
|
|
}
|
|
clc_pipeline_params.transaction_bytes = CLCResponseSize;
|
|
clc_pipeline_params.initializing_warp = 4;
|
|
CLCPipeline clc_pipeline(shared_storage.pipelines.clc, clc_pipeline_params, cluster_shape);
|
|
|
|
// Mainloop-Epilogue pipeline
|
|
typename AccumulatorPipeline::Params accumulator_pipeline_params;
|
|
if (WarpCategory::MMA == warp_category) {
|
|
accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer;
|
|
}
|
|
if (WarpCategory::Epilogue == warp_category) {
|
|
accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer;
|
|
}
|
|
// Only one producer thread arrives on this barrier.
|
|
accumulator_pipeline_params.producer_arv_count = 1;
|
|
accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * NumEpilogueThreads;
|
|
accumulator_pipeline_params.initializing_warp = 5;
|
|
AccumulatorPipeline accumulator_pipeline(shared_storage.pipelines.accumulator,
|
|
accumulator_pipeline_params,
|
|
cluster_shape,
|
|
cute::true_type{}, // Perform barrier init
|
|
cute::false_type{}); // Delay mask calculation
|
|
|
|
// Tmem allocator
|
|
TmemAllocator tmem_allocator{};
|
|
|
|
// Sync allocation status between MMA and epilogue warps within CTA
|
|
arch::NamedBarrier tmem_allocation_result_barrier(NumMMAThreads + NumEpilogueThreads, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);
|
|
// Sync deallocation status between MMA warps of peer CTAs
|
|
arch::ClusterBarrier& tmem_deallocation_result_barrier = shared_storage.pipelines.tmem_dealloc;
|
|
[[maybe_unused]] uint32_t dealloc_barrier_phase = 0;
|
|
if (WarpCategory::MMA == warp_category && has_mma_peer_cta && lane_predicate) {
|
|
tmem_deallocation_result_barrier.init(NumMMAThreads);
|
|
}
|
|
|
|
// We need this to guarantee that the Pipeline init is visible
|
|
// To all producers and consumer threadblocks in the cluster
|
|
pipeline_init_arrive_relaxed(cluster_size);
|
|
|
|
auto load_inputs = collective_mainloop.load_init(
|
|
problem_shape_MNKL, params.mainloop, shared_storage.tensors.mainloop);
|
|
|
|
uint32_t tmem_stage_ptrs[AccumulatorPipelineStageCount];
|
|
MainloopPipelineState mainloop_pipe_consumer_state;
|
|
MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
|
|
|
|
EpiLoadPipelineState epi_load_pipe_consumer_state;
|
|
EpiLoadPipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
|
|
|
|
// epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
|
|
EpiStorePipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
|
|
|
|
CLCPipelineState clc_pipe_consumer_state;
|
|
CLCPipelineState clc_pipe_producer_state = cutlass::make_producer_start_state<CLCPipeline>();
|
|
|
|
AccumulatorPipelineState accumulator_pipe_consumer_state;
|
|
AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state<AccumulatorPipeline>();
|
|
|
|
dim3 block_id_in_cluster = cute::block_id_in_cluster();
|
|
|
|
// Calculate mask after cluster barrier arrival
|
|
mainloop_pipeline.init_masks(cluster_shape, block_id_in_cluster);
|
|
accumulator_pipeline.init_masks(cluster_shape, block_id_in_cluster);
|
|
|
|
// TileID scheduler
|
|
TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, problem_shape_MNKL, TileShape{}, block_id_in_cluster);
|
|
typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
|
|
auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
|
auto acc_shape = collective_mainloop.partition_accumulator_shape();
|
|
auto accumulators = TiledMma::make_fragment_C(acc_shape);
|
|
|
|
int TmemColumnsPerAccumulatorTile = cutlass::detail::find_tmem_tensor_col_offset(accumulators);
|
|
pipeline_init_wait(cluster_size);
|
|
|
|
if (is_participant.main_load) {
|
|
// 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 = is_epi_load_needed;
|
|
Tensor gA_mk = get<0>(load_inputs);
|
|
|
|
do {
|
|
// Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
|
|
auto k_tile_iter = scheduler.get_k_tile_iterator(work_tile_info, problem_shape_MNKL, TileShape{}, shape<3>(gA_mk));
|
|
auto k_tile_count = scheduler.get_work_k_tile_count(work_tile_info, problem_shape_MNKL, TileShape{});
|
|
auto k_tile_prologue = min(MainloopPipeline::Stages, k_tile_count);
|
|
|
|
auto [mainloop_producer_state_next, k_tile_iter_next] = collective_mainloop.load(
|
|
params.mainloop,
|
|
mainloop_pipeline,
|
|
mainloop_pipe_producer_state,
|
|
load_inputs,
|
|
cta_coord_mnkl,
|
|
k_tile_iter, k_tile_prologue
|
|
);
|
|
mainloop_pipe_producer_state = mainloop_producer_state_next;
|
|
|
|
if (do_load_order_arrive) {
|
|
load_order_barrier.arrive();
|
|
do_load_order_arrive = false;
|
|
}
|
|
|
|
auto [mainloop_producer_state_next_, unused_] = collective_mainloop.load(
|
|
params.mainloop,
|
|
mainloop_pipeline,
|
|
mainloop_pipe_producer_state,
|
|
load_inputs,
|
|
cta_coord_mnkl,
|
|
k_tile_iter_next, k_tile_count - k_tile_prologue
|
|
);
|
|
mainloop_pipe_producer_state = mainloop_producer_state_next_;
|
|
|
|
// Sync warp to prevent non-participating threads entering next wave early
|
|
__syncwarp();
|
|
|
|
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
|
work_tile_info,
|
|
clc_pipeline,
|
|
clc_pipe_consumer_state
|
|
);
|
|
work_tile_info = next_work_tile_info;
|
|
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
|
if (increment_pipe) {
|
|
++clc_pipe_consumer_state;
|
|
}
|
|
} while (work_tile_info.is_valid());
|
|
collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
|
|
|
|
}
|
|
|
|
else if (is_participant.sched) {
|
|
// Whether a new CLC query must be performed.
|
|
// See comment below where this variable is updated for a description of
|
|
// why this variable is needed.
|
|
bool requires_clc_query = true;
|
|
|
|
do {
|
|
if (requires_clc_query) {
|
|
// Query next clcID and update producer state
|
|
clc_pipe_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipe_producer_state);
|
|
}
|
|
|
|
// Fetch next work tile
|
|
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
|
work_tile_info,
|
|
clc_pipeline,
|
|
clc_pipe_consumer_state
|
|
);
|
|
|
|
// Only perform a new CLC query if we consumed a new CLC query result in
|
|
// `fetch_next_work`. An example of a case in which CLC `fetch_next_work` does
|
|
// not consume a new CLC query response is when processing stream-K units.
|
|
// The current stream-K scheduler uses single WorkTileInfo to track multiple
|
|
// (potentially-partial) tiles to be computed via stream-K. In this case,
|
|
// `fetch_next_work` simply performs in-place updates on the existing WorkTileInfo,
|
|
// rather than consuming a CLC query response.
|
|
requires_clc_query = increment_pipe;
|
|
if (increment_pipe) {
|
|
++clc_pipe_consumer_state;
|
|
}
|
|
|
|
work_tile_info = next_work_tile_info;
|
|
} while (work_tile_info.is_valid());
|
|
clc_pipeline.producer_tail(clc_pipe_producer_state);
|
|
}
|
|
|
|
else if (is_participant.mma) {
|
|
// Tmem allocation sequence
|
|
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr);
|
|
__syncwarp();
|
|
tmem_allocation_result_barrier.arrive();
|
|
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
|
|
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int acc_stage = 0; acc_stage < AccumulatorPipelineStageCount; acc_stage++) {
|
|
tmem_stage_ptrs[acc_stage] = tmem_base_ptr + (TmemColumnsPerAccumulatorTile * acc_stage) & cutlass::detail::TmemColMask;
|
|
}
|
|
auto mma_inputs = collective_mainloop.mma_init(shared_storage.tensors.mainloop);
|
|
|
|
do {
|
|
auto k_tile_count = scheduler.get_work_k_tile_count(work_tile_info, problem_shape_MNKL, TileShape{});
|
|
|
|
// Fetch next work tile
|
|
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
|
work_tile_info,
|
|
clc_pipeline,
|
|
clc_pipe_consumer_state
|
|
);
|
|
|
|
if (increment_pipe) {
|
|
++clc_pipe_consumer_state;
|
|
}
|
|
|
|
// Wait for tmem accumulator buffer to become empty with a flipped phase
|
|
if (is_mma_leader_cta) {
|
|
accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state);
|
|
}
|
|
|
|
// Accumulator stage slice
|
|
int acc_stage = accumulator_pipe_producer_state.index();
|
|
accumulators.data() = tmem_stage_ptrs[acc_stage];
|
|
|
|
if (is_mma_leader_cta) {
|
|
mainloop_pipe_consumer_state = collective_mainloop.mma(
|
|
mainloop_pipeline,
|
|
mainloop_pipe_consumer_state,
|
|
accumulators,
|
|
mma_inputs,
|
|
k_tile_count
|
|
);
|
|
accumulator_pipeline.producer_commit(accumulator_pipe_producer_state);
|
|
}
|
|
++accumulator_pipe_producer_state;
|
|
work_tile_info = next_work_tile_info;
|
|
} while (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();
|
|
|
|
// Release the right to allocate before deallocations so that the next CTA can rasterize
|
|
tmem_allocator.release_allocation_lock();
|
|
|
|
// Leader MMA waits for leader + peer epilogues to release accumulator stage
|
|
if (is_mma_leader_cta) {
|
|
accumulator_pipeline.producer_tail(accumulator_pipe_producer_state);
|
|
}
|
|
// Signal to peer MMA that entire tmem allocation can be deallocated
|
|
if constexpr (has_mma_peer_cta) {
|
|
// Leader does wait + arrive, follower does arrive + wait
|
|
tmem_deallocation_result_barrier.arrive(mma_peer_cta_rank, not is_mma_leader_cta);
|
|
tmem_deallocation_result_barrier.wait(dealloc_barrier_phase);
|
|
tmem_deallocation_result_barrier.arrive(mma_peer_cta_rank, is_mma_leader_cta);
|
|
}
|
|
|
|
// Free entire tmem allocation
|
|
tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns);
|
|
}
|
|
|
|
else if (is_participant.epi_load) {
|
|
// 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_wait = true;
|
|
bool do_tail_load = false;
|
|
|
|
do {
|
|
bool compute_epilogue = TileScheduler::compute_epilogue(work_tile_info, params.scheduler);
|
|
|
|
// Get current work tile and fetch next work tile
|
|
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
|
work_tile_info,
|
|
clc_pipeline,
|
|
clc_pipe_consumer_state
|
|
);
|
|
work_tile_info = next_work_tile_info;
|
|
|
|
if (increment_pipe) {
|
|
++clc_pipe_consumer_state;
|
|
}
|
|
|
|
if (compute_epilogue) {
|
|
if (do_load_order_wait) {
|
|
load_order_barrier.wait();
|
|
do_load_order_wait = false;
|
|
}
|
|
|
|
epi_load_pipe_producer_state = collective_epilogue.load(
|
|
epi_load_pipeline,
|
|
epi_load_pipe_producer_state,
|
|
problem_shape_MNKL,
|
|
CtaShape_MNK{},
|
|
cta_coord_mnkl,
|
|
TileShape{},
|
|
TiledMma{},
|
|
shared_storage.tensors.epilogue
|
|
);
|
|
|
|
do_tail_load = true;
|
|
}
|
|
|
|
// Calculate the cta coordinates of the next work tile
|
|
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
|
} while (work_tile_info.is_valid());
|
|
|
|
// Only perform a tail load if one of the work units processed performed
|
|
// an epilogue load. An example of a case in which a tail load should not be
|
|
// performed is in split-K if a cluster is only assigned non-final splits (for which
|
|
// the cluster does not compute the epilogue).
|
|
if (do_tail_load) {
|
|
collective_epilogue.load_tail(
|
|
epi_load_pipeline, epi_load_pipe_producer_state,
|
|
epi_store_pipeline, epi_store_pipe_producer_state);
|
|
}
|
|
}
|
|
|
|
else if (is_participant.epilogue) {
|
|
// Wait for tmem allocate here
|
|
tmem_allocation_result_barrier.arrive_and_wait();
|
|
uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr;
|
|
CUTLASS_PRAGMA_UNROLL
|
|
for (int acc_stage = 0; acc_stage < AccumulatorPipelineStageCount; acc_stage++) {
|
|
tmem_stage_ptrs[acc_stage] = tmem_base_ptr + (TmemColumnsPerAccumulatorTile * acc_stage) & cutlass::detail::TmemColMask;
|
|
}
|
|
|
|
bool do_tail_store = false;
|
|
do {
|
|
// Fetch next work tile
|
|
auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(
|
|
work_tile_info,
|
|
clc_pipeline,
|
|
clc_pipe_consumer_state
|
|
);
|
|
|
|
if (increment_pipe) {
|
|
++clc_pipe_consumer_state;
|
|
}
|
|
|
|
// Accumulator stage slice after making sure allocation has been performed
|
|
int acc_stage = accumulator_pipe_consumer_state.index();
|
|
accumulators.data() = tmem_stage_ptrs[acc_stage];
|
|
|
|
accumulator_pipe_consumer_state = scheduler.template fixup<IsComplex>(
|
|
TiledMma{},
|
|
work_tile_info,
|
|
accumulators,
|
|
accumulator_pipeline,
|
|
accumulator_pipe_consumer_state,
|
|
typename CollectiveEpilogue::CopyOpT2R{}
|
|
);
|
|
|
|
//
|
|
// Epilogue and write to gD
|
|
//
|
|
if (scheduler.compute_epilogue(work_tile_info)) {
|
|
auto [load_state_next, store_state_next, acc_state_next] = collective_epilogue.store(
|
|
epi_load_pipeline,
|
|
epi_load_pipe_consumer_state,
|
|
epi_store_pipeline,
|
|
epi_store_pipe_producer_state,
|
|
accumulator_pipeline,
|
|
accumulator_pipe_consumer_state,
|
|
problem_shape_MNKL,
|
|
CtaShape_MNK{},
|
|
cta_coord_mnkl,
|
|
TileShape{},
|
|
TiledMma{},
|
|
accumulators,
|
|
shared_storage.tensors.epilogue
|
|
);
|
|
epi_load_pipe_consumer_state = load_state_next;
|
|
epi_store_pipe_producer_state = store_state_next;
|
|
accumulator_pipe_consumer_state = acc_state_next;
|
|
do_tail_store = true;
|
|
}
|
|
work_tile_info = next_work_tile_info;
|
|
cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
|
|
} while (work_tile_info.is_valid());
|
|
|
|
// Only perform a tail store if one of the work units processed performed
|
|
// an epilogue. An example of a case in which a tail load should not be
|
|
// performed is in split-K if a cluster is only assigned non-final splits (for which
|
|
// the cluster does not compute the epilogue).
|
|
if (do_tail_store) {
|
|
collective_epilogue.store_tail(
|
|
epi_load_pipeline, epi_load_pipe_consumer_state,
|
|
epi_store_pipeline, epi_store_pipe_producer_state,
|
|
CtaShape_MNK{});
|
|
}
|
|
}
|
|
|
|
else {
|
|
}
|
|
}
|
|
};
|
|
|
|
///////////////////////////////////////////////////////////////////////////////
|
|
|
|
} // namespace cutlass::gemm::kernel
|