From 4dfbc9de68bb51b7b280a49d20204969aff78703 Mon Sep 17 00:00:00 2001 From: kyle-256 Date: Tue, 20 Jan 2026 09:09:02 +0000 Subject: [PATCH] update gg configs for mi355 --- example/ck_tile/03_gemm/gemm_utils.hpp | 51 +-- example/ck_tile/03_gemm/universal_gemm.cpp | 24 +- .../ck_tile/17_grouped_gemm/grouped_gemm.cpp | 35 +- .../ck_tile/17_grouped_gemm/grouped_gemm.hpp | 347 ++++-------------- tile_engine/ops/gemm/gemm_instance_builder.py | 46 ++- .../configs/default_config.json | 19 +- .../ops/gemm/gemm_universal/gemm_benchmark.py | 24 +- .../gemm_universal_instance_builder.py | 7 +- 8 files changed, 192 insertions(+), 361 deletions(-) diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp index c3614f661b..1ea51ba386 100644 --- a/example/ck_tile/03_gemm/gemm_utils.hpp +++ b/example/ck_tile/03_gemm/gemm_utils.hpp @@ -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 -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(); - - 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 -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(); - - static constexpr bool DoubleSmemBuffer = false; - static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3; - - static constexpr int kBlockPerCu = 2; -}; - template 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; diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp index 1b8d328aae..5f6c149539 100644 --- a/example/ck_tile/03_gemm/universal_gemm.cpp +++ b/example/ck_tile/03_gemm/universal_gemm.cpp @@ -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, - 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, + 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(arg_parser); + { run_gemm_example(arg_parser); + run_gemm_example(arg_parser); run_gemm_example(arg_parser); return 0; } diff --git a/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp b/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp index f618e0a96a..781e0c0932 100644 --- a/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp +++ b/example/ck_tile/17_grouped_gemm/grouped_gemm.cpp @@ -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, + ck_tile::bf16_t>(a_layout, b_layout, argc, argv); + } + else if(config == "compute_v3_256x256_w16_k1") + { + return run_gemm_example_prec_type, + 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>(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>(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>(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 diff --git a/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp b/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp index 5a69c523a9..6ad6648fca 100644 --- a/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp +++ b/example/ck_tile/17_grouped_gemm/grouped_gemm.hpp @@ -213,6 +213,28 @@ struct GemmConfigComputeV3_2 : public GemmConfigBase static constexpr int kBlockPerCu = 1; }; +// 256x256, warp tile 16x16x16 +template +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(); + + 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 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 -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(); - - 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 -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(); + 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 +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 +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 -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(); - - 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 -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(); - - 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 -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(); - - 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 -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(); - - 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 -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(); - - 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 -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(); - - 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 using GemmConfigComputeV3_128x128 = GemmConfigComputeV3_128x128_16_k2; -// ============================================================================ -// 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 -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(); - - 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 -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(); - - 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 -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(); - - 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 -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(); - - static constexpr bool DoubleSmemBuffer = false; - static constexpr ck_tile::GemmPipeline Pipeline = ck_tile::GemmPipeline::COMPUTE_V3; - - static constexpr int kBlockPerCu = 2; -}; - - template struct GemmConfigComputeV4 : public GemmConfigBase { @@ -759,7 +555,10 @@ std::pair 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); diff --git a/tile_engine/ops/gemm/gemm_instance_builder.py b/tile_engine/ops/gemm/gemm_instance_builder.py index 9c60c565de..9b7792a7f2 100644 --- a/tile_engine/ops/gemm/gemm_instance_builder.py +++ b/tile_engine/ops/gemm/gemm_instance_builder.py @@ -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()}" diff --git a/tile_engine/ops/gemm/gemm_universal/configs/default_config.json b/tile_engine/ops/gemm/gemm_universal/configs/default_config.json index 2447428158..03ab6ef67f 100644 --- a/tile_engine/ops/gemm/gemm_universal/configs/default_config.json +++ b/tile_engine/ops/gemm/gemm_universal/configs/default_config.json @@ -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 ] } diff --git a/tile_engine/ops/gemm/gemm_universal/gemm_benchmark.py b/tile_engine/ops/gemm/gemm_universal/gemm_benchmark.py index b7424c6d1d..cc8f8ce91f 100644 --- a/tile_engine/ops/gemm/gemm_universal/gemm_benchmark.py +++ b/tile_engine/ops/gemm/gemm_universal/gemm_benchmark.py @@ -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, diff --git a/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py b/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py index 08f0e7e284..857be77d99 100644 --- a/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py +++ b/tile_engine/ops/gemm/gemm_universal/gemm_universal_instance_builder.py @@ -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