Files
composable_kernel/test/ck_tile/gemm/test_gemm_pipeline_compiler.cpp

901 lines
25 KiB
C++

// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
// SPDX-License-Identifier: MIT
#include "test_gemm_pipeline_kernel_types.hpp"
#include "test_gemm_pipeline_util.hpp"
#include "gtest/gtest.h"
// ============================================================================
// Comprehensive GEMM Compiler Validation Test Suite
// This file consolidates all GEMM pipeline tests for compiler validation
// Covers essential combinations of data types, layouts, and pipeline types
// ============================================================================
// ----------------------------------------------------------------------------
// Test Class Definitions for Different Pipeline Types
// ----------------------------------------------------------------------------
template <typename T>
class TestGemmMem : public TestCkTileGemmPipeline<T, TestGemmMem<T>>
{
};
#if defined(CK_TILE_USE_WMMA)
template <typename T>
class TestGemmMemWmma : public TestCkTileGemmPipeline<T, TestGemmMemWmma<T>>
{
};
#endif
template <typename T>
class TestGemmCompV3 : public TestCkTileGemmPipeline<T, TestGemmCompV3<T>>
{
};
#if defined(CK_TILE_USE_WMMA)
template <typename T>
class TestGemmCompV3Wmma : public TestCkTileGemmPipeline<T, TestGemmCompV3Wmma<T>>
{
};
#endif
template <typename T>
class TestGemmCompV4 : public TestCkTileGemmPipeline<T, TestGemmCompV4<T>>
{
};
#if defined(CK_TILE_USE_WMMA)
template <typename T>
class TestGemmCompV4Wmma : public TestCkTileGemmPipeline<T, TestGemmCompV4Wmma<T>>
{
};
#endif
template <typename T>
class TestGemmCompV6 : public TestCkTileGemmPipeline<T, TestGemmCompV6<T>>
{
};
template <typename T>
class TestGemmPersistent : public TestCkTileGemmPipeline<T, TestGemmPersistent<T>>
{
};
#if defined(CK_TILE_USE_WMMA)
template <typename T>
class TestGemmPersistentWmma : public TestCkTileGemmPipeline<T, TestGemmPersistentWmma<T>>
{
};
#endif
// ----------------------------------------------------------------------------
// Type Definitions for Each Pipeline Configuration
// ----------------------------------------------------------------------------
// Memory Pipeline Types
using MemTestTypes = ::testing::Types<
// Parameters: ALayout, BLayout, CLayout, ADataType, BDataType, AccDataType, CDataType,
// M_BlockSize, N_BlockSize, K_BlockSize, M_TileSize, N_TileSize, K_TileSize, Scheduler,
// PipelineType
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Interwave, Mem>,
std::tuple<Row, Row, Row, BF16, BF16, F32, BF16, I64, I64, I32, I16, I16, I16, Interwave, Mem>>;
#if defined(CK_TILE_USE_WMMA)
// Memory Pipeline WMMA Types
using MemWmmaTestTypes = ::testing::Types<
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Interwave, Mem>,
std::tuple<Row, Row, Row, BF16, BF16, F32, BF16, I64, I64, I32, I16, I16, I16, Interwave, Mem>>;
#endif
// CompV3 Pipeline Types
using CompV3TestTypes = ::testing::Types<
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Intrawave, CompV3>,
std::tuple<Row,
Row,
Row,
BF16,
BF16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV3>>;
#if defined(CK_TILE_USE_WMMA)
// CompV3 Pipeline WMMA Types
using CompV3WmmaTestTypes = ::testing::Types<
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Intrawave, CompV3>,
std::tuple<Row,
Row,
Row,
BF16,
BF16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV3>>;
#endif
// CompV4 Pipeline Types
using CompV4TestTypes = ::testing::Types<
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Intrawave, CompV4>,
std::tuple<Row,
Row,
Row,
BF16,
BF16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV4>>;
#if defined(CK_TILE_USE_WMMA)
// CompV4 Pipeline WMMA Types
using CompV4WmmaTestTypes = ::testing::Types<
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Intrawave, CompV4>,
std::tuple<Row,
Row,
Row,
BF16,
BF16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV4>>;
#endif
// CompV6 Pipeline Types
using CompV6TestTypes = ::testing::Types<
std::tuple<Row, Row, Row, F16, F16, F32, F16, I64, I64, I32, I16, I16, I16, Intrawave, CompV6>,
std::tuple<Row,
Row,
Row,
BF16,
BF16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV6>>;
// Persistent CompV3 Pipeline Types
using PersistentTestTypes = ::testing::Types<std::tuple<Row,
Col,
Row,
F16,
F16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV3,
Persistent>,
std::tuple<Row,
Col,
Row,
F16,
F16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV3,
NonPersistent>>;
#if defined(CK_TILE_USE_WMMA)
// Persistent CompV3 Pipeline WMMA Types
using PersistentWmmaTestTypes = ::testing::Types<std::tuple<Row,
Col,
Row,
F16,
F16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV3,
Persistent>,
std::tuple<Row,
Col,
Row,
F16,
F16,
F32,
F16,
I64,
I64,
I32,
I16,
I16,
I16,
Intrawave,
CompV3,
NonPersistent>>;
#endif
// ----------------------------------------------------------------------------
// Test Suite Registrations
// ----------------------------------------------------------------------------
TYPED_TEST_SUITE(TestGemmMem, MemTestTypes);
#if defined(CK_TILE_USE_WMMA)
TYPED_TEST_SUITE(TestGemmMemWmma, MemWmmaTestTypes);
#endif
TYPED_TEST_SUITE(TestGemmCompV3, CompV3TestTypes);
#if defined(CK_TILE_USE_WMMA)
TYPED_TEST_SUITE(TestGemmCompV3Wmma, CompV3WmmaTestTypes);
#endif
TYPED_TEST_SUITE(TestGemmCompV4, CompV4TestTypes);
#if defined(CK_TILE_USE_WMMA)
TYPED_TEST_SUITE(TestGemmCompV4Wmma, CompV4WmmaTestTypes);
#endif
TYPED_TEST_SUITE(TestGemmCompV6, CompV6TestTypes);
TYPED_TEST_SUITE(TestGemmPersistent, PersistentTestTypes);
#if defined(CK_TILE_USE_WMMA)
TYPED_TEST_SUITE(TestGemmPersistentWmma, PersistentWmmaTestTypes);
#endif
// ============================================================================
// Memory Pipeline Tests (Mem)
// ============================================================================
#define TEST_SUITE_NAME TestGemmMem
TYPED_TEST(TEST_SUITE_NAME, SmallM_SingleRow)
{
std::vector<int> Ms{1};
constexpr int N = 1024;
constexpr int K = TestFixture::K_Tile * 2;
for(int M : Ms)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, ExactlyTwoTiles_M)
{
this->Run(TestFixture::M_Tile * 2, TestFixture::N_Tile, TestFixture::K_Tile * 2);
}
TYPED_TEST(TEST_SUITE_NAME, ExactlyTwoTiles_N)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile * 2, TestFixture::K_Tile * 2);
}
TYPED_TEST(TEST_SUITE_NAME, ExactlyTwoTiles_K)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile * 2);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_512x1024x512)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, Square_1024x1024x1024)
{
constexpr int M = 1024;
constexpr int N = 1024;
constexpr int K = 1024;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_2048x2048x2048)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, VeryLargeMatrix_4096x4096x4096)
{
constexpr int M = 4096;
constexpr int N = 4096;
constexpr int K = 4096;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, TallSkinny_4096x128x1024)
{
constexpr int M = 4096;
constexpr int N = 128;
constexpr int K = 1024;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, ShortWide_128x4096x1024)
{
constexpr int M = 128;
constexpr int N = 4096;
constexpr int K = 1024;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, DeepNarrow_2048x2048x8192)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 8192;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, StressTest_ExtremelyTallMatrix)
{
constexpr int M = 16384;
constexpr int N = 64;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, StressTest_ExtremelyWideMatrix)
{
constexpr int M = 64;
constexpr int N = 16384;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, StressTest_VeryDeepK)
{
constexpr int M = 1024;
constexpr int N = 1024;
constexpr int K = 16384;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#if defined(CK_TILE_USE_WMMA)
// ============================================================================
// Memory Pipeline Tests with WMMA
// ============================================================================
#define TEST_SUITE_NAME TestGemmMemWmma
TYPED_TEST(TEST_SUITE_NAME, SingleTile_WMMA)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_WMMA)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_WMMA)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#endif // CK_TILE_USE_WMMA
// ============================================================================
// Compute V3 Pipeline Tests
// ============================================================================
#define TEST_SUITE_NAME TestGemmCompV3
TYPED_TEST(TEST_SUITE_NAME, SmallM_CompV3)
{
std::vector<int> Ms{1, 2};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile_CompV3)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, MidLargeM_CompV3)
{
std::vector<int> Ms{127, 255};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
constexpr int VecLoadSize = (std::is_same_v<typename TestFixture::ADataType, ck_tile::fp8_t> ||
std::is_same_v<typename TestFixture::ADataType, ck_tile::bf8_t> ||
std::is_same_v<typename TestFixture::ADataType, ck_tile::int8_t>)
? 16
: 8;
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
if(M % VecLoadSize == 0)
{
this->Run(M, N, K);
}
else
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, Regular_CompV3)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_CompV3)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, BatchedSmall_CompV3)
{
constexpr int M = 256;
constexpr int N = 256;
constexpr int K = 256;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#if defined(CK_TILE_USE_WMMA)
// ============================================================================
// Compute V3 Pipeline Tests with WMMA
// ============================================================================
#define TEST_SUITE_NAME TestGemmCompV3Wmma
TYPED_TEST(TEST_SUITE_NAME, SmallM_CompV3Wmma)
{
std::vector<int> Ms{1, 2};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile_CompV3Wmma)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_CompV3Wmma)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_CompV3Wmma)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#endif // CK_TILE_USE_WMMA
// ============================================================================
// Compute V4 Pipeline Tests
// ============================================================================
#define TEST_SUITE_NAME TestGemmCompV4
TYPED_TEST(TEST_SUITE_NAME, SmallM_CompV4)
{
std::vector<int> Ms{1, 2};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile_CompV4)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_CompV4)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_CompV4)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#if defined(CK_TILE_USE_WMMA)
// ============================================================================
// Compute V4 Pipeline Tests with WMMA
// ============================================================================
#define TEST_SUITE_NAME TestGemmCompV4Wmma
TYPED_TEST(TEST_SUITE_NAME, SingleTile_CompV4Wmma)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_CompV4Wmma)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_CompV4Wmma)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#endif // CK_TILE_USE_WMMA
// ============================================================================
// Compute V6 Pipeline Tests
// ============================================================================
#define TEST_SUITE_NAME TestGemmCompV6
TYPED_TEST(TEST_SUITE_NAME, SmallM_CompV6)
{
std::vector<int> Ms{1, 2};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile_CompV6)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, MidLargeM_CompV6)
{
std::vector<int> Ms{127, 255};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
constexpr int VecLoadSize = (std::is_same_v<typename TestFixture::ADataType, ck_tile::fp8_t> ||
std::is_same_v<typename TestFixture::ADataType, ck_tile::bf8_t> ||
std::is_same_v<typename TestFixture::ADataType, ck_tile::int8_t>)
? 16
: 8;
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
if(M % VecLoadSize == 0)
{
this->Run(M, N, K);
}
else
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, Regular_CompV6)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_CompV6)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
// ============================================================================
// Persistent Kernel Tests
// ============================================================================
#define TEST_SUITE_NAME TestGemmPersistent
TYPED_TEST(TEST_SUITE_NAME, SmallM_Persistent)
{
std::vector<int> Ms{1, 2};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile_Persistent)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_Persistent)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_Persistent)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#if defined(CK_TILE_USE_WMMA)
// ============================================================================
// Persistent Kernel Tests with WMMA
// ============================================================================
#define TEST_SUITE_NAME TestGemmPersistentWmma
TYPED_TEST(TEST_SUITE_NAME, SmallM_PersistentWmma)
{
std::vector<int> Ms{1, 2};
constexpr int N = 1024;
std::vector<int> Ks;
for(auto K_count : {2, 4})
{
Ks.push_back(K_count * TestFixture::K_Tile);
}
for(int M : Ms)
{
for(int K : Ks)
{
if constexpr(std::is_same_v<typename TestFixture::ALayout,
ck_tile::tensor_layout::gemm::ColumnMajor>)
{
EXPECT_THROW((this->Run(M, N, K)), std::runtime_error);
}
else
{
this->Run(M, N, K);
}
}
}
}
TYPED_TEST(TEST_SUITE_NAME, SingleTile_PersistentWmma)
{
this->Run(TestFixture::M_Tile, TestFixture::N_Tile, TestFixture::K_Tile);
}
TYPED_TEST(TEST_SUITE_NAME, Regular_PersistentWmma)
{
constexpr int M = 512;
constexpr int N = 1024;
constexpr int K = 512;
this->Run(M, N, K);
}
TYPED_TEST(TEST_SUITE_NAME, LargeMatrix_PersistentWmma)
{
constexpr int M = 2048;
constexpr int N = 2048;
constexpr int K = 2048;
this->Run(M, N, K);
}
#undef TEST_SUITE_NAME
#endif // CK_TILE_USE_WMMA