Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add TT, TN, NT, NN tests for HopperMultipleMatmulScheduler #3310

Merged
merged 5 commits into from
Nov 3, 2024
Merged

Add TT, TN, NT, NN tests for HopperMultipleMatmulScheduler #3310

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
3ae7baa
Add guard to skip ampere matmul tests on hopper
rdspring1 Nov 1, 2024
1ce800a
create TN and NT tests
rdspring1 Oct 30, 2024
7c8f375
Create TT and NT tests
rdspring1 Nov 1, 2024
43b2433
hopper-only
rdspring1 Nov 1, 2024
9956259
disable
rdspring1 Nov 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 35 additions & 16 deletions tests/cpp/test_matmul.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,8 @@ using namespace at::indexing;

// Matmul test for Ampere MMA: across supported layouts
TEST_P(MatmulTestWithLayout, AmpereMatmul) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -141,6 +143,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmul) {
}

TEST_P(MatmulTestWithLayout, AmperePrologueFusionBroadcast) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -195,6 +199,8 @@ TEST_P(MatmulTestWithLayout, AmperePrologueFusionBroadcast) {
}

TEST_P(MatmulTestWithLayout, AmpereProloguePointwise) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -256,6 +262,8 @@ TEST_P(MatmulTestWithLayout, AmpereProloguePointwise) {
}

TEST_P(MatmulTestWithLayout, AmpereMatmulBFloat16) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -314,6 +322,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulBFloat16) {

// Matmul test for Ampere MMA: with pipelined gmem load
TEST_P(MatmulTestWithLayout, AmpereMatmulPipelineGmem) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;
REQUIRE_DEVICE_SMEM_SIZE(70 << 10, 0);
Expand Down Expand Up @@ -376,6 +386,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulPipelineGmem) {

// Matmul test for Ampere MMA: checking CTA Swizzles
TEST_P(MatmulTestWithLayout, AmpereSwizzle) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int dim = 8192;
int M = dim, N = dim, K = dim;
Expand Down Expand Up @@ -522,6 +534,8 @@ TEST_P(MatmulTestWithLayout, AmpereSwizzle) {
}

TEST_P(MatmulTestWithLayout, AmpereMatmulRegCircularBuffer) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;
REQUIRE_DEVICE_SMEM_SIZE(70 << 10, 0);
Expand Down Expand Up @@ -1254,6 +1268,8 @@ TEST_F(MatmulTest, MatmulSoftmaxMatmulAmpere) {

// Matmul test for Turing MMA: across supported layouts
TEST_P(MatmulTestWithLayout, TuringMatmul) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -1973,6 +1989,8 @@ TEST_F(MatmulTest, AmpereMatmulTNSwizzled) {

// Matmul test on Ampere using ldmatrix.x4 to load operands
TEST_P(MatmulTestWithLayout, AmpereMatmulLargeLoad) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

REQUIRE_DEVICE_SMEM_SIZE(98384, 0);
// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;
Expand Down Expand Up @@ -2030,6 +2048,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulLargeLoad) {

// Matmul test for Turing MMA: across supported layouts
TEST_P(MatmulTestWithLayout, TuringMatmulLargeLoad) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -2084,6 +2104,8 @@ TEST_P(MatmulTestWithLayout, TuringMatmulLargeLoad) {

// Tile layout check for symmetric 4-warp recipes
TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck4warp) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

REQUIRE_DEVICE_SMEM_SIZE(98384, 0);
// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;
Expand Down Expand Up @@ -2162,6 +2184,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck4warp) {
}

TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck8warp) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

REQUIRE_DEVICE_SMEM_SIZE(98384, 0);
// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;
Expand Down Expand Up @@ -2235,6 +2259,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck8warp) {
}

TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck6warp) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

REQUIRE_DEVICE_SMEM_SIZE(98384, 0);
// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;
Expand Down Expand Up @@ -2303,6 +2329,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulTileCheck6warp) {

// Matmul test on Ampere using ldmatrix.x4 to load operands
TEST_P(MatmulTestWithLayout, AmpereMatmulLargeLoadLargeK) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 2048;
Fusion fusion;
Expand Down Expand Up @@ -2360,6 +2388,8 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulLargeLoadLargeK) {

// Matmul test for Ampere MMA: across supported layouts
TEST_P(MatmulTestWithLayout, AmpereSplitKLikeStridedBatchedMatmul) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(7, 5, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int B = 2, M = 504, N = 136, K = 248;

Expand Down Expand Up @@ -2853,10 +2883,7 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulSmemEpilogueRelu) {

// Test the matmul scheduler's single-kernel split-K support
TEST_P(MatmulTestWithLayout, FusionAmpereMatmulSplitK_CUDA) {
// requires Ampere or higher GPU
if (!deviceMajorMinorCheck(8)) {
GTEST_SKIP() << "skipping tests on pre-AMPERE GPUs";
}
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 8096;
Expand Down Expand Up @@ -2931,10 +2958,7 @@ TEST_P(MatmulTestWithLayout, FusionAmpereMatmulSplitK_CUDA) {

// Test splitk with bias epilogue
TEST_P(MatmulTestWithLayout, FusionAmpereMatmulSplitKBias_CUDA) {
// requires Ampere or higher GPU
if (!deviceMajorMinorCheck(8)) {
GTEST_SKIP() << "skipping tests on pre-AMPERE GPUs";
}
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 8096;
Expand Down Expand Up @@ -3000,10 +3024,7 @@ TEST_P(MatmulTestWithLayout, FusionAmpereMatmulSplitKBias_CUDA) {

// Same as above but has a batch dimension and splitk
TEST_P(MatmulTestWithLayout, AmpereMatmulBatchSplitK) {
// requires Ampere or higher GPU
if (!deviceMajorMinorCheck(8)) {
GTEST_SKIP() << "skipping tests on pre-AMPERE GPUs";
}
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int B = 2, M = 504, N = 136, K = 2048;
Expand Down Expand Up @@ -3065,10 +3086,7 @@ TEST_P(MatmulTestWithLayout, AmpereMatmulBatchSplitK) {

// Test batch splitk with bias epilogue
TEST_P(MatmulTestWithLayout, AmpereMatmulBatchSplitKBias) {
// requires Ampere or higher GPU
if (!deviceMajorMinorCheck(8)) {
GTEST_SKIP() << "skipping tests on pre-AMPERE GPUs";
}
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int B = 2, M = 504, N = 136, K = 2048;
Expand Down Expand Up @@ -3196,6 +3214,7 @@ TEST_F(MatmulTest, ReproIssue1808) {

// Test matmul with sizes that are not divisible by 8 and with misaligned inputs
TEST_P(MatmulTestWithLayout, MisalignedVectorization) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);
for (bool add_2d_bias : {false, true}) {
for (bool downcast_output : {false, true}) {
for (const auto& [M, N, K, alignA, alignB, alignBias] : std::vector<
Expand Down
2 changes: 2 additions & 0 deletions tests/cpp/test_matmul_sass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,8 @@ sass::Container getBinaryOpMulEpilogueSASSFor(
} // namespace

TEST_P(MatmulSASSTestWithLayout, AmpereSanity) {
NVFUSER_TEST_CUDA_ARCH_RANGE_GUARD(8, 0, 9, 0);

// Keep multiples of 8 to keep vectorizable.
int M = 504, N = 136, K = 248;

Expand Down
Loading
Loading