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:
co-authored by
Pradeep Ramani
parent
b7508e3379
commit
8236f30675
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user