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

CUTLASS 3.1 (#915)

Browse Source 
Co-authored-by: Aniket Shivam <ashivam@nvidia.com>
This commit is contained in:
ANIKET SHIVAM
2023-04-14 23:19:34 -04:00
committed by GitHub
co-authored by Aniket Shivam
parent 9b8166e3f0
commit d572cc1aab
482 changed files with 37175 additions and 16410 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/pipeline/pipeline_async.cu
+6 -6
View File
@@ -51,7 +51,7 @@
#include "cutlass/util/GPU_Clock.hpp"
#include "testbed.h"
#include "cutlass/pipeline.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/barrier.h"
#include "cute/arch/cluster_sm90.hpp"
@@ -98,21 +98,21 @@ void pipeline_async_basic_device(uint32_t const num_iterations)
cute::cluster_wait();
__syncthreads();
if (lane_predicate) {
// Producer Warps
if (warp_idx==0 || warp_idx==1) {
PipelineState smem_pipe_write = cutlass::make_producer_start_state<MainloopPipeline>();
int prologue_iterations = min(NumStages, num_iterations);
for ( int i = 0; i < prologue_iterations; ++i) {
// Can also specify stage to commit directly
pipeline.producer_commit(i);
pipeline.producer_commit(smem_pipe_write);
++smem_pipe_write;
}
int mainloop_iterations = num_iterations - prologue_iterations;
// Only the mainloop needs a PipelineState because this is where we start "waiting" (acquiring)
PipelineState smem_pipe_write;
for ( ; mainloop_iterations > 0; --mainloop_iterations) {
pipeline.producer_acquire(smem_pipe_write);
pipeline.producer_commit(smem_pipe_write);
@@ -123,7 +123,7 @@ void pipeline_async_basic_device(uint32_t const num_iterations)
PipelineState smem_pipe_read;
for (int iter=0 ; iter < num_iterations; ++iter) {
pipeline.consumer_wait(smem_pipe_read);
pipeline.consumer_release(smem_pipe_read.index());
pipeline.consumer_release(smem_pipe_read);
++smem_pipe_read;
}
}
test/unit/pipeline/pipeline_tma_async.cu
+12 -18
View File
@@ -41,7 +41,7 @@
#include <thrust/device_vector.h>
#include <cute/tensor.hpp>
#include <cute/arch/cluster_sm90.hpp>
#include <cute/arch/cluster_sm90.hpp>
#include <cutlass/util/reference/host/gemm.h>
#include <cutlass/cluster_launch.hpp>
@@ -52,7 +52,7 @@
#include "cutlass/util/GPU_Clock.hpp"
#include "testbed.h"
#include "cutlass/pipeline.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/barrier.h"
#include "cute/arch/cluster_sm90.hpp"
@@ -68,12 +68,11 @@ struct SharedStorage
// Goal of this kernel is to complete deadlock-free
template <class ClusterShape, uint32_t NumStages>
__global__ static
__global__ static
void pipeline_device(uint32_t const NumIterations)
{
extern __shared__ char shared_memory[];
using DispatchPolicy = cutlass::gemm::MainloopSm90TmaGmma<NumStages, ClusterShape>;
using MainloopPipeline = cutlass::PipelineTmaAsync<NumStages, ClusterShape>;
using PipelineState = cutlass::PipelineState<NumStages>;
@@ -86,8 +85,8 @@ void pipeline_device(uint32_t const NumIterations)
dim3 block_id_in_cluster = cute::block_id_in_cluster();
auto cluster_shape = ClusterShape{};
// #Producers = #RowsInCluster + #ColsInCluster - 1
// #Producers = #RowsInCluster + #ColsInCluster - 1
uint32_t const NumProducers = cute::size<0>(cluster_shape) + cute::size<1>(cluster_shape) - 1;
uint32_t const TmaTransactionBytes = sizeof(uint32_t) * NumProducers;
uint32_t const per_cta_bytes = sizeof(uint32_t);
@@ -104,7 +103,7 @@ void pipeline_device(uint32_t const NumIterations)
__syncthreads();
// Ensure All CTAs in Cluster have completed init before issuing commits
cute::cluster_arrive_relaxed();
cute::cluster_arrive_relaxed();
cute::cluster_wait();
// Total number of gemm_k_iterations
@@ -126,7 +125,7 @@ void pipeline_device(uint32_t const NumIterations)
for(int i = 0; i < k_pipe_tma_prologue; ++i) {
pipeline.producer_acquire(smem_pipe_write);
// cp.async.bulk.tensor would typically happen here
pipeline.producer_commit(smem_pipe_write.index(), per_cta_bytes);
pipeline.producer_commit(smem_pipe_write, per_cta_bytes);
++smem_pipe_write;
}
tma_k_iterations -= k_pipe_tma_prologue;
@@ -156,7 +155,7 @@ void pipeline_device(uint32_t const NumIterations)
if (lane_predicate && (warp_idx == 0) && (tma_k_iterations > 0)) {
pipeline.producer_acquire(smem_pipe_write);
// cp.async.bulk.tensor would typically happen here
pipeline.producer_commit(smem_pipe_write.index(), per_cta_bytes);
pipeline.producer_commit(smem_pipe_write, per_cta_bytes);
++smem_pipe_write;
--tma_k_iterations;
}
@@ -167,7 +166,7 @@ void pipeline_device(uint32_t const NumIterations)
}
// To make sure remote SMEM doesn't get destoryed
cute::cluster_arrive();
cute::cluster_arrive();
cute::cluster_wait();
}
/////////////////////////////////////////////////////
@@ -224,11 +223,6 @@ struct PipelineTest {
}
for (int iter = 0; iter < iterations; ++iter) {
// Define the tiled MMA layout (static, 4warps)
using DispatchPolicy = cutlass::gemm::MainloopSm90TmaGmma<Stages, decltype(cluster_shape)>;
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, decltype(cluster_shape)>;
int smem_size = int(sizeof(SharedStorage<Stages, decltype(cluster_shape)>));
result = cudaFuncSetAttribute(
@@ -237,15 +231,15 @@ struct PipelineTest {
smem_size);
// Launch a single Cluster, with 128 thread per CTA
dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), 1);
dim3 dimGrid(size<0>(cluster_shape), size<1>(cluster_shape), 1);
dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), 1);
dim3 dimGrid(size<0>(cluster_shape), size<1>(cluster_shape), 1);
dim3 dimBlock(kBlockSize,1,1);
const void* kernel = (const void*)pipeline_device<decltype(cluster_shape), Stages>;
int iters = kNumIters;
void* kernel_params[] = {reinterpret_cast<void*>(&iters)};
cutlass::ClusterLauncher::launch(dimGrid, dimCluster, dimBlock, smem_size, stream, kernel, kernel_params);
} // profiling loop ends
result = cudaEventRecord(events[1]);
test/unit/pipeline/pipeline_tma_async_warp_specialized.cu
+3 -3
View File
@@ -50,7 +50,7 @@
#include "cutlass/util/GPU_Clock.hpp"
#include "testbed.h"
#include "cutlass/pipeline.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/barrier.h"
#include "cute/arch/cluster_sm90.hpp"
#include "cutlass/arch/barrier.h"
@@ -138,7 +138,7 @@ void pipeline_device(KernelParams const kernel_params)
for(int i = 0; i < tma_k_prologue; ++i) {
pipeline.producer_acquire(smem_pipe_write);
// Simulating cp.async.bulk.tensor behavior
pipeline.producer_commit(smem_pipe_write.index(), per_cta_bytes);
pipeline.producer_commit(smem_pipe_write, per_cta_bytes);
++smem_pipe_write;
}
int tma_k_iter = kernel_params.num_iterations - tma_k_prologue;
@@ -150,7 +150,7 @@ void pipeline_device(KernelParams const kernel_params)
pipeline.producer_acquire(smem_pipe_write);
// Simulating cp.async.bulk.tensor behavior
pipeline.producer_commit(smem_pipe_write.index(), per_cta_bytes);
pipeline.producer_commit(smem_pipe_write, per_cta_bytes);
// Advance write stage
++smem_pipe_write;
test/unit/pipeline/pipeline_tma_async_warp_specialized_persistent.cu
+3 -9
View File
@@ -50,7 +50,7 @@
#include "cutlass/util/GPU_Clock.hpp"
#include "testbed.h"
#include "cutlass/pipeline.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/barrier.h"
#include "cute/arch/cluster_sm90.hpp"
#include "cutlass/arch/barrier.h"
@@ -90,7 +90,7 @@ struct CollectiveSimulation {
for(int i = 0; i < tma_k_prologue; ++i) {
pipeline.producer_acquire(tile_start_state_pipe);
// Simulating cp.async.bulk.tensor behavior
pipeline.producer_commit(tile_start_state_pipe.index(), per_cta_bytes);
pipeline.producer_commit(tile_start_state_pipe, per_cta_bytes);
++tile_start_state_pipe;
}
int tma_k_iter = num_iterations - tma_k_prologue;
@@ -103,7 +103,7 @@ struct CollectiveSimulation {
pipeline.producer_acquire(wr_pipe);
// Simulating cp.async.bulk.tensor behavior
pipeline.producer_commit(wr_pipe.index(), per_cta_bytes);
pipeline.producer_commit(wr_pipe, per_cta_bytes);
// Advance write stage
++wr_pipe;
@@ -198,9 +198,6 @@ __global__ static
void pipeline_device(KernelParams params)
{
extern __shared__ char shared_memory[];
using DispatchPolicy = cutlass::gemm::MainloopSm90TmaGmmaWarpSpecialized<Stages,
ClusterShape,
cutlass::gemm::KernelTmaWarpSpecializedPersistent>;
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, ClusterShape>;
using PipelineState = typename cutlass::PipelineState<Stages>;
@@ -345,9 +342,6 @@ struct PipelineTest {
}
for (int iter = 0; iter < iterations; ++iter) {
using MainloopPipeline = typename cutlass::PipelineTmaAsync<Stages, decltype(cluster_shape)>;
constexpr int StagesPerMathWarpGroup = 2;
constexpr int MathWarpGroupCountPersistent = 2;
int smem_size = int(sizeof(SharedStorage<Stages, decltype(cluster_shape),
test/unit/pipeline/sequence_barrier.cu
+3 -3
View File
@@ -49,7 +49,7 @@
#include "cutlass/util/GPU_Clock.hpp"
#include "testbed.h"
#include "cutlass/pipeline.hpp"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/barrier.h"
#include "cute/arch/cluster_sm90.hpp"
@@ -96,7 +96,7 @@ void ordered_sequence_device(uint32_t const num_iterations)
#ifndef NDEBUG
int thread_idx_in_group = threadIdx.x % ThreadsPerGroup;
if (thread_idx_in_group == 0) {
printf("STAGE 0 : Group_IDX : %d, id = %d, iter = %d, tidx = %d\n", group_idx, params.id, i, threadIdx.x);
printf("STAGE 0 : Group_IDX : %d, id = %d, iter = %d, tidx = %d\n", group_idx, params.group_id, i, threadIdx.x);
}
#endif
// Simulates long running stage
@@ -109,7 +109,7 @@ void ordered_sequence_device(uint32_t const num_iterations)
// STAGE 2 CODE...
#ifndef NDEBUG
if (thread_idx_in_group == 0) {
printf("STAGE 1 : Group_IDX : %d, id = %d, iter = %d, tidx = %d\n", group_idx, params.id, i, threadIdx.x);
printf("STAGE 1 : Group_IDX : %d, id = %d, iter = %d, tidx = %d\n", group_idx, params.group_id, i, threadIdx.x);
}
#endif
// Simulates long running stage