CUTLASS 3.5.1 (#1623)

* CUTLASS 3.5.1

* updates, optimizations, fixes
This commit is contained in:
Vijay Thakkar
2024-07-29 08:46:24 -04:00
committed by GitHub
parent 56b46e2d13
commit be60a0b272
312 changed files with 19793 additions and 6775 deletions
@@ -179,6 +179,111 @@ TEST(SM80_sparse_gemm_threadblock_congruous,
////////////////////////////////////////////////////////////////////////////////
TEST(SM80_sparse_gemm_threadblock_congruous,
tensor_op_128x32x64_32x32x64_16x8x32_4stage) {
using ElementA = cutlass::half_t;
using LayoutA = cutlass::layout::ColumnMajor;
using ElementB = cutlass::half_t;
using LayoutB = cutlass::layout::RowMajor;
using ElementC = float;
using LayoutC = cutlass::layout::ColumnMajor;
cutlass::gemm::GemmCoord problem_size(128, 32, 512);
using ThreadblockShape = cutlass::gemm::GemmShape<128, 32, 64>;
using WarpShape = cutlass::gemm::GemmShape<32, 32, 64>;
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
float alpha = 1.f;
float beta = 0.0f;
int const Stages = 4;
// Define the MmaCore components
using MmaCore = typename cutlass::gemm::threadblock::DefaultSparseMmaCore<
ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA,
ElementB, LayoutB, ElementC, LayoutC, cutlass::arch::OpClassTensorOp,
Stages>;
dim3 grid(1, 1);
dim3 block(32, 4, 1);
test::gemm::threadblock::SparseTestbed<MmaCore>(
problem_size.m(), problem_size.n(), problem_size.k(), alpha, beta)
.run(grid, block);
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM80_sparse_gemm_threadblock_congruous,
tensor_op_32x256x128_32x64x128_16x8x32_4stage) {
using ElementA = cutlass::half_t;
using LayoutA = cutlass::layout::ColumnMajor;
using ElementB = cutlass::half_t;
using LayoutB = cutlass::layout::RowMajor;
using ElementC = float;
using LayoutC = cutlass::layout::ColumnMajor;
cutlass::gemm::GemmCoord problem_size(32, 256, 512);
using ThreadblockShape = cutlass::gemm::GemmShape<32, 256, 128>;
using WarpShape = cutlass::gemm::GemmShape<32, 64, 128>;
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
float alpha = 1.f;
float beta = 0.0f;
int const Stages = 4;
// Define the MmaCore components
using MmaCore = typename cutlass::gemm::threadblock::DefaultSparseMmaCore<
ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA,
ElementB, LayoutB, ElementC, LayoutC, cutlass::arch::OpClassTensorOp,
Stages>;
dim3 grid(1, 1);
dim3 block(32, 4, 1);
test::gemm::threadblock::SparseTestbed<MmaCore>(
problem_size.m(), problem_size.n(), problem_size.k(), alpha, beta)
.run(grid, block);
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM80_sparse_gemm_threadblock_congruous,
tensor_op_128x16x64_32x16x64_16x8x32_4stage) {
using ElementA = cutlass::half_t;
using LayoutA = cutlass::layout::ColumnMajor;
using ElementB = cutlass::half_t;
using LayoutB = cutlass::layout::RowMajor;
using ElementC = float;
using LayoutC = cutlass::layout::ColumnMajor;
cutlass::gemm::GemmCoord problem_size(128, 16, 512);
using ThreadblockShape = cutlass::gemm::GemmShape<128, 16, 64>;
using WarpShape = cutlass::gemm::GemmShape<32, 16, 64>;
using InstructionShape = cutlass::gemm::GemmShape<16, 8, 32>;
float alpha = 1.f;
float beta = 0.0f;
int const Stages = 4;
// Define the MmaCore components
using MmaCore = typename cutlass::gemm::threadblock::DefaultSparseMmaCore<
ThreadblockShape, WarpShape, InstructionShape, ElementA, LayoutA,
ElementB, LayoutB, ElementC, LayoutC, cutlass::arch::OpClassTensorOp,
Stages>;
dim3 grid(1, 1);
dim3 block(32, 4, 1);
test::gemm::threadblock::SparseTestbed<MmaCore>(
problem_size.m(), problem_size.n(), problem_size.k(), alpha, beta)
.run(grid, block);
}
////////////////////////////////////////////////////////////////////////////////
TEST(SM80_sparse_gemm_threadblock_congruous,
tensor_op_128x128x64_64x64x64_16x8x32_4stage) {
using ElementA = cutlass::half_t;