CUTLASS 3.4.0 (#1286)

* CUTLASS 3.4.0

* Update CHANGELOG.md

---------

Co-authored-by: Pradeep Ramani <prramani@nvidia.com>
This commit is contained in:
Pradeep Ramani
2023-12-29 15:21:31 -05:00
committed by GitHub
co-authored by Pradeep Ramani
parent b7508e3379
commit 8236f30675
211 changed files with 11409 additions and 2763 deletions
+6 -6
View File
@@ -60,10 +60,10 @@ using namespace cute;
//////////////////// KERNEL /////////////////////////
template <uint32_t Stages, typename ClusterShape>
template <uint32_t Stages>
struct SharedStorage
{
typename cutlass::PipelineTmaAsync<Stages, ClusterShape>::SharedStorage storage;
typename cutlass::PipelineTmaAsync<Stages>::SharedStorage storage;
};
// Goal of this kernel is to complete deadlock-free
@@ -73,10 +73,10 @@ void pipeline_device(uint32_t const NumIterations)
{
extern __shared__ char shared_memory[];
using MainloopPipeline = cutlass::PipelineTmaAsync<NumStages, ClusterShape>;
using MainloopPipeline = cutlass::PipelineTmaAsync<NumStages>;
using PipelineState = cutlass::PipelineState<NumStages>;
using SharedStorage = SharedStorage<NumStages, ClusterShape>;
using SharedStorage = SharedStorage<NumStages>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
[[maybe_unused]] auto cta_layout = Layout<ClusterShape>{}; // (m,n) -> cta_id
@@ -98,7 +98,7 @@ void pipeline_device(uint32_t const NumIterations)
params.is_leader = warp_group_thread_idx == 0;
params.num_consumers = 128;
MainloopPipeline pipeline(shared_storage.storage, params);
MainloopPipeline pipeline(shared_storage.storage, params, cluster_shape);
__syncthreads();
@@ -223,7 +223,7 @@ struct PipelineTest {
}
for (int iter = 0; iter < iterations; ++iter) {
int smem_size = int(sizeof(SharedStorage<Stages, decltype(cluster_shape)>));
int smem_size = int(sizeof(SharedStorage<Stages>));
result = cudaFuncSetAttribute(
pipeline_device<decltype(cluster_shape), Stages>,
@@ -62,10 +62,10 @@ using namespace cutlass;
//////////////////// KERNEL /////////////////////////
template <uint32_t Stages, typename ClusterShape>
template <uint32_t Stages>
struct SharedStorage
{
typename cutlass::PipelineTmaAsync<Stages, ClusterShape>::SharedStorage storage ;
typename cutlass::PipelineTmaAsync<Stages>::SharedStorage storage ;
};
struct KernelParams
@@ -81,10 +81,10 @@ __global__ static
void pipeline_device(KernelParams const kernel_params)
{
extern __shared__ char shared_memory[];
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, ClusterShape>;
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages>;
using PipelineState = typename cutlass::PipelineState<Stages>;
using SharedStorage = SharedStorage<Stages, ClusterShape>;
using SharedStorage = SharedStorage<Stages>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
[[maybe_unused]] auto cta_layout = Layout<ClusterShape>{}; // (m,n) -> cta_id
@@ -112,7 +112,7 @@ void pipeline_device(KernelParams const kernel_params)
params.is_leader = warp_group_thread_idx == 0;
params.num_consumers = 128;
MainloopPipeline pipeline(shared_storage.storage, params);
MainloopPipeline pipeline(shared_storage.storage, params, cluster_shape);
__syncthreads();
@@ -292,9 +292,9 @@ struct PipelineTest {
for (int iter = 0; iter < iterations; ++iter) {
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, decltype(cluster_shape)>;
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages>;
int smem_size = int(sizeof(SharedStorage<Stages, decltype(cluster_shape)>));
int smem_size = int(sizeof(SharedStorage<Stages>));
result = cudaFuncSetAttribute(
pipeline_device<decltype(cluster_shape), Stages>,
@@ -62,16 +62,16 @@ using namespace cutlass;
//////////////////// KERNEL /////////////////////////
template <uint32_t Stages, typename ClusterShape, typename PingPongBarrier>
template <uint32_t Stages, typename PingPongBarrier>
struct SharedStorage
{
typename cutlass::PipelineTmaAsync<Stages, ClusterShape>::SharedStorage pipeline_storage;
typename cutlass::PipelineTmaAsync<Stages>::SharedStorage pipeline_storage;
typename PingPongBarrier::SharedStorage pingpong_storage;
};
template <typename ClusterShape, uint32_t Stages>
struct CollectiveSimulation {
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, ClusterShape>;
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages>;
using PipelineState = typename cutlass::PipelineState<Stages>;
CUTLASS_DEVICE
@@ -198,7 +198,7 @@ __global__ static
void pipeline_device(KernelParams params)
{
extern __shared__ char shared_memory[];
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, ClusterShape>;
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages>;
using PipelineState = typename cutlass::PipelineState<Stages>;
/* One for Mainloop and one for Epilogue */
@@ -206,7 +206,7 @@ void pipeline_device(KernelParams params)
constexpr int MathWarpGroupCountPersistent = 2;
using PingPongBarrier = typename cutlass::OrderedSequenceBarrier<StagesPerMathWarpGroup, MathWarpGroupCountPersistent>;
using SharedStorage = SharedStorage<Stages, ClusterShape, PingPongBarrier>;
using SharedStorage = SharedStorage<Stages, PingPongBarrier>;
SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(shared_memory);
[[maybe_unused]] auto cta_layout = Layout<ClusterShape>{}; // (m,n) -> cta_id
@@ -232,7 +232,7 @@ void pipeline_device(KernelParams params)
pipeline_params.is_leader = warp_group_thread_idx == 0;
pipeline_params.num_consumers = NumThreadsPerWarpGroup;
MainloopPipeline pipeline(shared_storage.pipeline_storage, pipeline_params);
MainloopPipeline pipeline(shared_storage.pipeline_storage, pipeline_params, cluster_shape);
PipelineState tile_start_state_pipe;
int tiles_per_cluster = params.tiles_per_cluster;
@@ -343,7 +343,7 @@ struct PipelineTest {
for (int iter = 0; iter < iterations; ++iter) {
constexpr int StagesPerMathWarpGroup = 2;
constexpr int MathWarpGroupCountPersistent = 2;
int smem_size = int(sizeof(SharedStorage<Stages, decltype(cluster_shape),
int smem_size = int(sizeof(SharedStorage<Stages,
typename cutlass::OrderedSequenceBarrier<StagesPerMathWarpGroup, MathWarpGroupCountPersistent>>));
result = cudaFuncSetAttribute(