mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-13 18:51:13 +00:00
update gg configs for mi355
This commit is contained in:
kyle-256
@@ -103,7 +103,7 @@ struct GemmConfigComputeV3_1 : public GemmConfigBase
|
||||
static constexpr ck_tile::index_t N_Tile = 256;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); // 64 for bf16
|
||||
|
||||
static constexpr ck_tile::index_t M_Warp = 2; // 4 warps in M
|
||||
static constexpr ck_tile::index_t M_Warp = 2; // 2 warps in M
|
||||
static constexpr ck_tile::index_t N_Warp = 2; // 2 warps in N
|
||||
static constexpr ck_tile::index_t K_Warp = 1;
|
||||
|
||||
@@ -118,53 +118,6 @@ struct GemmConfigComputeV3_1 : public GemmConfigBase
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// Optimized config: 256x128x64 with COMPUTE_V3
|
||||
// Smaller N tile for better L2 cache utilization
|
||||
template <typename PrecType>
|
||||
struct GemmConfigOptimized : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 256;
|
||||
static constexpr ck_tile::index_t N_Tile = 128;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType); // 64 for 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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2; // Better occupancy with smaller tile
|
||||
};
|
||||
|
||||
// Alternative: 128x128x128 config for higher K throughput
|
||||
template <typename PrecType>
|
||||
struct GemmConfigOptimized2 : 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 = 256 / sizeof(PrecType); // 128 for 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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_2 : public GemmConfigBase
|
||||
{
|
||||
@@ -256,7 +209,7 @@ struct GemmConfigComputeAsync : 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); // 64 for bf16 (same as V3_1)
|
||||
static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType); // 64 for bf16
|
||||
|
||||
static constexpr ck_tile::index_t M_Warp = 2;
|
||||
static constexpr ck_tile::index_t N_Warp = 2;
|
||||
|
||||
@@ -23,22 +23,11 @@ int run_gemm_example(ck_tile::ArgParser& arg_parser)
|
||||
|
||||
if(data_type == "bf16")
|
||||
{
|
||||
// Only support RC layout (A=Row, B=Column) to reduce compile time
|
||||
if(a_layout != "R" || b_layout != "C")
|
||||
{
|
||||
throw std::runtime_error("Only RC layout (A=Row, B=Column) is supported!");
|
||||
}
|
||||
using Row = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
|
||||
// Directly call with fixed layout to avoid std::visit instantiating all combinations
|
||||
return run_gemm_example_with_layouts<GemmConfig<ck_tile::bf16_t>,
|
||||
Invoker,
|
||||
ck_tile::bf16_t,
|
||||
ck_tile::bf16_t,
|
||||
ck_tile::bf16_t,
|
||||
Row,
|
||||
Col,
|
||||
Row>(arg_parser, Row{}, Col{}, Row{});
|
||||
return run_gemm_example_prec_type<GemmConfig<ck_tile::bf16_t>,
|
||||
Invoker,
|
||||
ck_tile::bf16_t,
|
||||
ck_tile::bf16_t,
|
||||
ck_tile::bf16_t>(a_layout, b_layout, arg_parser);
|
||||
}
|
||||
else
|
||||
{
|
||||
@@ -55,7 +44,8 @@ int main(int argc, char* argv[])
|
||||
return -1;
|
||||
|
||||
try
|
||||
{ run_gemm_example<GemmConfigComputeV4>(arg_parser);
|
||||
{ run_gemm_example<GemmConfigComputeV3_1>(arg_parser);
|
||||
run_gemm_example<GemmConfigComputeV4>(arg_parser);
|
||||
run_gemm_example<GemmConfigComputeAsync>(arg_parser);
|
||||
return 0;
|
||||
}
|
||||
|
||||
@@ -316,9 +316,42 @@ int run_grouped_gemm_example_with_n_check(int argc, char* argv[])
|
||||
throw std::runtime_error("Unsupported N alignment for compute_v3 config.");
|
||||
}
|
||||
}
|
||||
else if(config == "compute_v3_128x128")
|
||||
{
|
||||
return run_gemm_example_prec_type<GemmConfigComputeV3_128x128<ck_tile::bf16_t>,
|
||||
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
|
||||
}
|
||||
else if(config == "compute_v3_256x256_w16_k1")
|
||||
{
|
||||
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_128x128x64_2x2x1_32x32x16")
|
||||
{
|
||||
return run_gemm_example_prec_type<
|
||||
GemmConfigComputeV3_128x128x64_2x2x1_32x32x16<ck_tile::bf16_t>,
|
||||
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
|
||||
}
|
||||
else if(config == "compute_v3_128x128x128_2x2x1_32x32x16")
|
||||
{
|
||||
return run_gemm_example_prec_type<
|
||||
GemmConfigComputeV3_128x128x128_2x2x1_32x32x16<ck_tile::bf16_t>,
|
||||
ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
|
||||
}
|
||||
else if(config == "compute_v3_128x256x64_1x4x1_32x32x16")
|
||||
{
|
||||
return run_gemm_example_prec_type<
|
||||
GemmConfigComputeV3_128x256x64_1x4x1_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_32x128, compute_v3_128x128");
|
||||
throw std::runtime_error(
|
||||
"Unknown config: " + config +
|
||||
". Use: compute_v3, compute_v3_128x128, compute_v3_256x256_w16_k1, "
|
||||
"compute_v3_128x128x64_2x2x1_32x32x16, "
|
||||
"compute_v3_128x128x128_2x2x1_32x32x16, "
|
||||
"compute_v3_128x256x64_1x4x1_32x32x16");
|
||||
}
|
||||
}
|
||||
else
|
||||
|
||||
@@ -213,6 +213,28 @@ struct GemmConfigComputeV3_2 : public GemmConfigBase
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 256x256, warp tile 16x16x16
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_256x256_w16_k1 : 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 = 128 / sizeof(PrecType);
|
||||
|
||||
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 = 16;
|
||||
static constexpr ck_tile::index_t N_Warp_Tile = 16;
|
||||
static constexpr ck_tile::index_t K_Warp_Tile =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 256x128 tile config for N % 128 == 0 but N % 256 != 0
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_256x128 : public GemmConfigBase
|
||||
@@ -236,32 +258,9 @@ struct GemmConfigComputeV3_256x128 : public GemmConfigBase
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 32x128 tile config (small M, from FBGEMM)
|
||||
// 128x128x64, warp 2x2x1, warp tile 32x32x16
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_32x128 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 32;
|
||||
static constexpr ck_tile::index_t N_Tile = 128;
|
||||
static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType);
|
||||
|
||||
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 = 32;
|
||||
static constexpr ck_tile::index_t N_Warp_Tile = 32;
|
||||
static constexpr ck_tile::index_t K_Warp_Tile =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 128x128, M_Warp_Tile=16, kBlockPerCu=1
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x128_16_k1 : public GemmConfigBase
|
||||
struct GemmConfigComputeV3_128x128x64_2x2x1_32x32x16 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 128;
|
||||
static constexpr ck_tile::index_t N_Tile = 128;
|
||||
@@ -271,15 +270,55 @@ struct GemmConfigComputeV3_128x128_16_k1 : public GemmConfigBase
|
||||
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 = 16;
|
||||
static constexpr ck_tile::index_t N_Warp_Tile = 16;
|
||||
static constexpr ck_tile::index_t K_Warp_Tile =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
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 = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
// 128x128x128, warp 2x2x1, warp tile 32x32x16
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x128x128_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 = 256 / sizeof(PrecType);
|
||||
|
||||
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 = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 128x256x64, warp 1x4x1, warp tile 32x32x16
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x256x64_1x4x1_32x32x16 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 128;
|
||||
static constexpr ck_tile::index_t N_Tile = 256;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
|
||||
|
||||
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 = 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 = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 128x128, M_Warp_Tile=16, kBlockPerCu=2
|
||||
@@ -305,253 +344,10 @@ struct GemmConfigComputeV3_128x128_16_k2 : public GemmConfigBase
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
// 128x128, M_Warp_Tile=32, kBlockPerCu=1
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x128_32_k1 : 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);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 128x128, M_Warp_Tile=32, kBlockPerCu=2
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x128_32_k2 : 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);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
// 64x64, M_Warp_Tile=16, kBlockPerCu=1
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_64x64_16_k1 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 64;
|
||||
static constexpr ck_tile::index_t N_Tile = 64;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
|
||||
|
||||
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 = 16;
|
||||
static constexpr ck_tile::index_t N_Warp_Tile = 16;
|
||||
static constexpr ck_tile::index_t K_Warp_Tile =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 64x64, M_Warp_Tile=16, kBlockPerCu=2
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_64x64_16_k2 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 64;
|
||||
static constexpr ck_tile::index_t N_Tile = 64;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
|
||||
|
||||
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 = 16;
|
||||
static constexpr ck_tile::index_t N_Warp_Tile = 16;
|
||||
static constexpr ck_tile::index_t K_Warp_Tile =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
// 64x64, M_Warp_Tile=32, kBlockPerCu=1
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_64x64_32_k1 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 64;
|
||||
static constexpr ck_tile::index_t N_Tile = 64;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 64x64, M_Warp_Tile=32, kBlockPerCu=2
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_64x64_32_k2 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 64;
|
||||
static constexpr ck_tile::index_t N_Tile = 64;
|
||||
static constexpr ck_tile::index_t K_Tile = 128 / sizeof(PrecType);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
// Backward compatible alias
|
||||
template <typename PrecType>
|
||||
using GemmConfigComputeV3_128x128 = GemmConfigComputeV3_128x128_16_k2<PrecType>;
|
||||
|
||||
// ============================================================================
|
||||
// Configs for grad_B (small K scenario: K' = M_original, typically 512-2048)
|
||||
// CR layout, large M' and N', small K'
|
||||
// ============================================================================
|
||||
|
||||
// ============================================================================
|
||||
// grad_B configs - using kbatch to handle small K (K' = M_original)
|
||||
// grad_B: A^T(K,M) x grad_out(M,N) -> grad_B(K,N)
|
||||
// M_gemm = K_original (1408-16384), N_gemm = N_original (2048-32768), K_gemm = M_original (512-2048, small!)
|
||||
// Strategy: use kbatch=2/4 to split K_gemm, and use configs with good N coverage
|
||||
// ============================================================================
|
||||
|
||||
// 256x256 with kBlockPerCu=2 for better occupancy on small K
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_256x256_k2 : 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);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
// 128x256 for large N (grad_B typically has large N)
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x256 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 128;
|
||||
static constexpr ck_tile::index_t N_Tile = 256;
|
||||
static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
};
|
||||
|
||||
// 256x128 - good for moderate N
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_256x128_k2 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 256;
|
||||
static constexpr ck_tile::index_t N_Tile = 128;
|
||||
static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
// 128x256 with kBlockPerCu=2
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV3_128x256_k2 : public GemmConfigBase
|
||||
{
|
||||
static constexpr ck_tile::index_t M_Tile = 128;
|
||||
static constexpr ck_tile::index_t N_Tile = 256;
|
||||
static constexpr ck_tile::index_t K_Tile = 64 / sizeof(PrecType);
|
||||
|
||||
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 =
|
||||
ck_tile::get_k_warp_tile<PrecType, M_Warp_Tile>();
|
||||
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3;
|
||||
|
||||
static constexpr int kBlockPerCu = 2;
|
||||
};
|
||||
|
||||
|
||||
template <typename PrecType>
|
||||
struct GemmConfigComputeV4 : public GemmConfigBase
|
||||
{
|
||||
@@ -759,7 +555,10 @@ std::pair<bool, ck_tile::ArgParser> create_args(int argc, char* argv[])
|
||||
.insert("kbatch", "1", "kbatch for SplitK")
|
||||
.insert("json", "0", "0: No Json, 1: Dump Results in Json format")
|
||||
.insert("jsonfile", "grouped_gemm.json", "json file name to dump results")
|
||||
.insert("config", "", "Tile config: compute_v3 (default), memory_interwave, memory_intrawave");
|
||||
.insert("config",
|
||||
"",
|
||||
"Tile config: compute_v3 (auto), compute_v3_* (manual), memory_interwave, "
|
||||
"memory_intrawave");
|
||||
|
||||
bool result = arg_parser.parse(argc, argv);
|
||||
return std::make_pair(result, arg_parser);
|
||||
|
||||
@@ -88,6 +88,26 @@ class GemmKernelBuilder:
|
||||
|
||||
kernel_name += f"_{tile_str}"
|
||||
|
||||
# Validate tile config against the current pipeline
|
||||
if not is_tile_config_valid(
|
||||
tile_config["tile_m"],
|
||||
tile_config["tile_n"],
|
||||
tile_config["tile_k"],
|
||||
tile_config["warp_m"],
|
||||
tile_config["warp_n"],
|
||||
tile_config["warp_k"],
|
||||
tile_config["warp_tile_m"],
|
||||
tile_config["warp_tile_n"],
|
||||
tile_config["warp_tile_k"],
|
||||
self.datatype,
|
||||
self.datatype,
|
||||
self.datatype,
|
||||
pipeline,
|
||||
self.layout,
|
||||
self.gpu_target,
|
||||
):
|
||||
continue
|
||||
|
||||
kernel_list.append(
|
||||
{
|
||||
"name": kernel_name,
|
||||
@@ -163,9 +183,11 @@ class GemmKernelBuilder:
|
||||
# Generate all combinations
|
||||
default_pipeline = ""
|
||||
if self.kernel_name_prefix == "gemm_universal":
|
||||
default_pipeline = "compv4"
|
||||
pipelines = self.config.get("trait_config", {}).get("pipeline", {}).get("values", [])
|
||||
default_pipeline = "compv3" if "compv3" in pipelines else "compv4"
|
||||
elif self.kernel_name_prefix == "gemm_multi_d":
|
||||
default_pipeline = "compv4"
|
||||
pipelines = self.config.get("trait_config", {}).get("pipeline", {}).get("values", [])
|
||||
default_pipeline = "compv3" if "compv3" in pipelines else "compv4"
|
||||
elif self.kernel_name_prefix == "gemm_preshuffle":
|
||||
default_pipeline = "preshufflev2"
|
||||
|
||||
@@ -321,6 +343,26 @@ class GemmKernelBuilder:
|
||||
persistent,
|
||||
) = trait_combo
|
||||
|
||||
# Skip invalid tile/pipeline combinations (e.g., LDS size limits per pipeline)
|
||||
if not is_tile_config_valid(
|
||||
tile_config["tile_m"],
|
||||
tile_config["tile_n"],
|
||||
tile_config["tile_k"],
|
||||
tile_config["warp_m"],
|
||||
tile_config["warp_n"],
|
||||
tile_config["warp_k"],
|
||||
tile_config["warp_tile_m"],
|
||||
tile_config["warp_tile_n"],
|
||||
tile_config["warp_tile_k"],
|
||||
self.datatype,
|
||||
self.datatype,
|
||||
self.datatype,
|
||||
pipeline,
|
||||
self.layout,
|
||||
self.gpu_target,
|
||||
):
|
||||
return None
|
||||
|
||||
# Create kernel name with proper boolean capitalization
|
||||
kernel_name = f"{self.kernel_name_prefix}_{self.datatype}_{self.layout}_{pipeline}_{epilogue}_{scheduler}_{str(pad_m).capitalize()}_{str(pad_n).capitalize()}_{str(pad_k).capitalize()}_{str(persistent).capitalize()}"
|
||||
|
||||
|
||||
@@ -2,18 +2,18 @@
|
||||
"tile_config": {
|
||||
"tile_m": {
|
||||
"max": 256,
|
||||
"min": 64,
|
||||
"min": 128,
|
||||
"step": 64
|
||||
},
|
||||
"tile_n": {
|
||||
"max": 256,
|
||||
"min": 64,
|
||||
"min": 128,
|
||||
"step": 64
|
||||
},
|
||||
"tile_k": {
|
||||
"max": 256,
|
||||
"min": 64,
|
||||
"step": 64
|
||||
"min": 32,
|
||||
"step": 32
|
||||
},
|
||||
"warp_m": {
|
||||
"values": [
|
||||
@@ -36,7 +36,6 @@
|
||||
},
|
||||
"warp_tile_m": {
|
||||
"values": [
|
||||
4,
|
||||
16,
|
||||
32
|
||||
]
|
||||
@@ -62,20 +61,17 @@
|
||||
"pipeline": {
|
||||
"values": [
|
||||
"compv3",
|
||||
"compv4",
|
||||
"mem"
|
||||
"compv4"
|
||||
]
|
||||
},
|
||||
"scheduler": {
|
||||
"values": [
|
||||
"intrawave",
|
||||
"interwave"
|
||||
"intrawave"
|
||||
]
|
||||
},
|
||||
"epilogue": {
|
||||
"values": [
|
||||
"cshuffle",
|
||||
"default"
|
||||
"cshuffle"
|
||||
]
|
||||
},
|
||||
"pad_m": {
|
||||
@@ -95,7 +91,6 @@
|
||||
},
|
||||
"persistent": {
|
||||
"values": [
|
||||
false,
|
||||
true
|
||||
]
|
||||
}
|
||||
|
||||
@@ -249,7 +249,13 @@ class GemmBenchmark:
|
||||
with open(json_file, "r") as f:
|
||||
content = f.read().strip()
|
||||
|
||||
# Parse the JSON directly since executables produce clean JSON
|
||||
# Extract JSON payload if stdout includes extra logs
|
||||
json_start = content.find("{")
|
||||
json_end = content.rfind("}")
|
||||
if json_start != -1 and json_end != -1 and json_end > json_start:
|
||||
content = content[json_start : json_end + 1]
|
||||
|
||||
# Parse the JSON payload
|
||||
data = json.loads(content)
|
||||
|
||||
# Return the complete JSON data as-is, just add some convenience fields
|
||||
@@ -597,7 +603,14 @@ def main():
|
||||
parser.add_argument(
|
||||
"--split-k", nargs="+", type=int, default=[1], help="Split-K values to test"
|
||||
)
|
||||
parser.add_argument("--verify", action="store_true", help="Enable verification")
|
||||
parser.add_argument("--verify", action="store_true", help="Enable GPU verification (verify=2)")
|
||||
parser.add_argument(
|
||||
"--verify-level",
|
||||
type=int,
|
||||
choices=[0, 1, 2],
|
||||
default=None,
|
||||
help="Verification level: 0 none, 1 CPU, 2 GPU (overrides --verify)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--csv", default="gemm_benchmark_results.csv", help="CSV output filename"
|
||||
)
|
||||
@@ -650,10 +663,15 @@ def main():
|
||||
print("Starting GEMM kernel benchmark sweep...")
|
||||
start_time = time.time()
|
||||
|
||||
if args.verify_level is not None:
|
||||
verify = args.verify_level
|
||||
else:
|
||||
verify = 2 if args.verify else 0
|
||||
|
||||
best_kernels = benchmark.benchmark_sweep(
|
||||
problem_sizes=problem_sizes,
|
||||
split_k_values=args.split_k,
|
||||
verify=args.verify,
|
||||
verify=verify,
|
||||
warmup=args.warmup,
|
||||
repeat=args.repeat,
|
||||
flush_cache=args.flush_cache,
|
||||
|
||||
@@ -141,9 +141,10 @@ def _generate_single_kernel_individual(work_item):
|
||||
)
|
||||
|
||||
try:
|
||||
kernel_name, instance_code = builder._generate_kernel_instance(
|
||||
tile_config, trait_combo
|
||||
)
|
||||
result = builder._generate_kernel_instance(tile_config, trait_combo)
|
||||
if result is None:
|
||||
return None
|
||||
kernel_name, instance_code = result
|
||||
|
||||
# Create simplified filename without the "gemm_universal_" prefix
|
||||
# Remove "gemm_universal_" from the beginning of kernel_name for the filename
|
||||
|
||||
Reference in New Issue
Block a user