update test codes

This commit is contained in:
kyle-256
2026-01-23 08:59:50 +00:00
parent 4dfbc9de68
commit 797aff8d49
3 changed files with 64 additions and 2 deletions

View File

@@ -46,7 +46,7 @@ int main(int argc, char* argv[])
try
{ run_gemm_example<GemmConfigComputeV3_1>(arg_parser);
run_gemm_example<GemmConfigComputeV4>(arg_parser);
run_gemm_example<GemmConfigComputeAsync>(arg_parser);
// run_gemm_example<GemmConfigComputeAsync>(arg_parser);
return 0;
}
catch(const std::runtime_error& e)

View File

@@ -326,6 +326,12 @@ int run_grouped_gemm_example_with_n_check(int argc, char* argv[])
return run_gemm_example_prec_type<GemmConfigComputeV3_256x256_w16_k1<ck_tile::bf16_t>,
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
}
else if(config == "compute_v3_256x256x32_1x4x1_16x16x32")
{
return run_gemm_example_prec_type<
GemmConfigComputeV3_256x256x32_1x4x1_16x16x32<ck_tile::bf16_t>,
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
}
else if(config == "compute_v3_128x128x64_2x2x1_32x32x16")
{
return run_gemm_example_prec_type<
@@ -344,14 +350,27 @@ int run_grouped_gemm_example_with_n_check(int argc, char* argv[])
GemmConfigComputeV3_128x256x64_1x4x1_32x32x16<ck_tile::bf16_t>,
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
}
else if(config == "compute_v4")
{
return run_gemm_example_prec_type<GemmConfigComputeV4<ck_tile::bf16_t>, ck_tile::bf16_t>(
a_layout, b_layout, argc, argv);
}
else if(config == "compute_v4_128x128x64_2x2x1_32x32x16")
{
return run_gemm_example_prec_type<
GemmConfigComputeV4_128x128x64_2x2x1_32x32x16<ck_tile::bf16_t>,
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
}
else
{
throw std::runtime_error(
"Unknown config: " + config +
". Use: compute_v3, compute_v3_128x128, compute_v3_256x256_w16_k1, "
"compute_v3_256x256x32_1x4x1_16x16x32, "
"compute_v3_128x128x64_2x2x1_32x32x16, "
"compute_v3_128x128x128_2x2x1_32x32x16, "
"compute_v3_128x256x64_1x4x1_32x32x16");
"compute_v3_128x256x64_1x4x1_32x32x16, "
"compute_v4, compute_v4_128x128x64_2x2x1_32x32x16");
}
}
else

View File

@@ -258,6 +258,27 @@ struct GemmConfigComputeV3_256x128 : public GemmConfigBase
static constexpr int kBlockPerCu = 1;
};
// 256x256x32, warp 1x4x1, warp tile 16x16x32
template <typename PrecType>
struct GemmConfigComputeV3_256x256x32_1x4x1_16x16x32 : public GemmConfigBase
{
static constexpr ck_tile::index_t M_Tile = 256;
static constexpr ck_tile::index_t N_Tile = 256;
static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType); // 32 for fp16/bf16
static constexpr ck_tile::index_t M_Warp = 1;
static constexpr ck_tile::index_t N_Warp = 4;
static constexpr ck_tile::index_t K_Warp = 1;
static constexpr ck_tile::index_t M_Warp_Tile = 16;
static constexpr ck_tile::index_t N_Warp_Tile = 16;
static constexpr ck_tile::index_t K_Warp_Tile = 32;
static constexpr bool DoubleSmemBuffer = false;
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
static constexpr int kBlockPerCu = 1;
};
// 128x128x64, warp 2x2x1, warp tile 32x32x16
template <typename PrecType>
struct GemmConfigComputeV3_128x128x64_2x2x1_32x32x16 : public GemmConfigBase
@@ -372,6 +393,28 @@ struct GemmConfigComputeV4 : public GemmConfigBase
static constexpr int kBlockPerCu = 2;
};
// 128x128x64, warp 2x2x1, warp tile 32x32x16 (V4)
template <typename PrecType>
struct GemmConfigComputeV4_128x128x64_2x2x1_32x32x16 : public GemmConfigBase
{
static constexpr ck_tile::index_t M_Tile = 128;
static constexpr ck_tile::index_t N_Tile = 128;
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); // 64 for fp16/bf16
static constexpr ck_tile::index_t M_Warp = 2;
static constexpr ck_tile::index_t N_Warp = 2;
static constexpr ck_tile::index_t K_Warp = 1;
static constexpr ck_tile::index_t M_Warp_Tile = 32;
static constexpr ck_tile::index_t N_Warp_Tile = 32;
static constexpr ck_tile::index_t K_Warp_Tile = 16;
static constexpr bool DoubleSmemBuffer = true;
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V4;
static constexpr int kBlockPerCu = 1;
};
template <typename PrecType>
struct GemmConfigComputeV4_V2 : public GemmConfigBase
{