From 8e3d84aba3be5e851de5d6c6c3e9c08cadbce1da Mon Sep 17 00:00:00 2001 From: Yi DING Date: Wed, 28 Jan 2026 15:46:49 +0800 Subject: [PATCH 01/11] [CK_TILE] ABQuant New Preshuffle (#3638) * Refactor * Gemm quant improvement * Change preshuffle * Fix * Fix grouped gemm ut * Fix --------- Co-authored-by: Thomas Ning --- .../38_block_scale_gemm/CMakeLists.txt | 1 + .../gemm_abquant_quantgrouped.cpp | 8 +-- .../gemm_aquant_quantgrouped.cpp | 8 +-- ...mm_aquant_quantgrouped_preshufflequant.cpp | 8 +-- .../gemm_bquant_quantgrouped_bf16mxfp4.cpp | 8 +-- .../gemm_bquant_quantgrouped_bf8.cpp | 8 +-- .../gemm_bquant_quantgrouped_bf8i4.cpp | 8 +-- .../gemm_bquant_quantgrouped_fp8.cpp | 8 +-- .../gemm_bquant_quantgrouped_fp8i4.cpp | 8 +-- ...mm_bquant_quantgrouped_preshuffleb_bf8.cpp | 8 +-- ..._bquant_quantgrouped_preshuffleb_bf8i4.cpp | 8 +-- ...mm_bquant_quantgrouped_preshuffleb_fp8.cpp | 8 +-- ..._bquant_quantgrouped_preshuffleb_fp8i4.cpp | 8 +-- ...rouped_preshuffleb_preshufflequant_bf8.cpp | 8 +-- ...uped_preshuffleb_preshufflequant_bf8i4.cpp | 8 +-- ...rouped_preshuffleb_preshufflequant_fp8.cpp | 8 +-- ...uped_preshuffleb_preshufflequant_fp8i4.cpp | 8 +-- ...quant_quantgrouped_preshufflequant_bf8.cpp | 8 +-- ...ant_quantgrouped_preshufflequant_bf8i4.cpp | 8 +-- ...quant_quantgrouped_preshufflequant_fp8.cpp | 8 +-- ...ant_quantgrouped_preshufflequant_fp8i4.cpp | 8 +-- .../38_block_scale_gemm/gemm_quant.cpp | 70 +------------------ .../38_block_scale_gemm/gemm_quant_rowcol.cpp | 8 +-- .../38_block_scale_gemm/gemm_quant_tensor.cpp | 8 +-- .../38_block_scale_gemm/gemm_utils.hpp | 8 +++ .../run_gemm_quant_example.inc | 16 ++--- include/ck_tile/core/arch/arch.hpp | 2 +- include/ck_tile/host/tensor_shuffle_utils.hpp | 42 ++++++----- ...emm_universal_pipeline_ag_bg_cr_policy.hpp | 8 +-- ..._pipeline_agmem_bgmem_creg_base_policy.hpp | 42 +++++++---- .../gemm_quant/kernel/gemm_quant_kernel.hpp | 10 +-- ..._abquant_pipeline_ag_bg_cr_base_policy.hpp | 20 ++++-- 32 files changed, 182 insertions(+), 213 deletions(-) diff --git a/example/ck_tile/38_block_scale_gemm/CMakeLists.txt b/example/ck_tile/38_block_scale_gemm/CMakeLists.txt index ec536f7287..13cbcc8b55 100644 --- a/example/ck_tile/38_block_scale_gemm/CMakeLists.txt +++ b/example/ck_tile/38_block_scale_gemm/CMakeLists.txt @@ -6,6 +6,7 @@ if(CK_USE_OCP_FP8) list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -DCK_TILE_USE_OCP_FP8) endif() +list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -Wno-global-constructors) # use global constructors to add kernel instances list(APPEND EXAMPLE_GEMM_COMPILE_OPTIONS -mllvm -enable-noalias-to-md-conversion=0) if(GPU_TARGETS MATCHES "gfx94|gfx95|gfx12") diff --git a/example/ck_tile/38_block_scale_gemm/gemm_abquant_quantgrouped.cpp b/example/ck_tile/38_block_scale_gemm/gemm_abquant_quantgrouped.cpp index b1cd1a52a7..153ab4845a 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_abquant_quantgrouped.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_abquant_quantgrouped.cpp @@ -12,9 +12,8 @@ using GemmConfigPreshuffleB = GemmConfigPreshuffleB_ABQuant_Prefill; // template // using GemmConfigPreshuffleB = GemmConfigPreshuffleB_ABQuant_Decode; -void abquant_quantgrouped_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); lut[hash_multiple_strings({"fp8", "abquant", "non-preshuffleb", @@ -135,4 +134,5 @@ void abquant_quantgrouped_instance_factory( BQuantGroupSize, ck_tile::QuantType::ABQuantGrouped>(arg_parser); }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped.cpp b/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped.cpp index e037be5a18..a6cc5e011d 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped.cpp @@ -10,9 +10,8 @@ using GemmConfig = GemmConfigQuantDecodeInterwave; // template // using GemmConfig = GemmConfigQuantPrefill; -void aquant_quantgrouped_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using QuantGroupSize = ck_tile::QuantGroupShape>; lut[hash_multiple_strings( {"fp8", "aquant", "non-preshufflequant", "1x1x128"})] = [](const ck_tile::ArgParser& @@ -56,4 +55,5 @@ void aquant_quantgrouped_instance_factory( QuantGroupSize, ck_tile::QuantType::AQuantGrouped>(arg_parser); }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped_preshufflequant.cpp b/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped_preshufflequant.cpp index 45e8c28a4e..32e2b3d603 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped_preshufflequant.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_aquant_quantgrouped_preshufflequant.cpp @@ -6,9 +6,8 @@ template using GemmConfig = GemmConfigPreshuffleQuantDecode; -void aquant_quantgrouped_preshufflequant_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using QuantGroupSize = ck_tile::QuantGroupShape>; lut[hash_multiple_strings( {"fp8", "aquant", "preshufflequant", "1x1x128"})] = [](const ck_tile::ArgParser& @@ -52,4 +51,5 @@ void aquant_quantgrouped_preshufflequant_instance_factory( QuantGroupSize, ck_tile::QuantType::AQuantGrouped>(arg_parser); }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf16mxfp4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf16mxfp4.cpp index 31d263ea1d..b8eb670135 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf16mxfp4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf16mxfp4.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_bf16fp4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8.cpp index 82e30e56d2..a95c0346cf 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_bf8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); #ifndef CK_GFX950_SUPPORT @@ -55,4 +54,5 @@ void bquant_quantgrouped_bf8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8i4.cpp index 515e6eb027..d2b95d3263 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_bf8i4.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_bf8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8.cpp index eaf10f057c..a8c13c1b3d 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_fp8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); #ifndef CK_GFX950_SUPPORT @@ -55,4 +54,5 @@ void bquant_quantgrouped_fp8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8i4.cpp index c91867534f..6576b22c03 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_fp8i4.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_fp8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8.cpp index 7166a5647e..e0c112e3b7 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_BQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_bf8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); lut[hash_multiple_strings({"bf8", "bquant", "preshuffleb", "non-preshufflequant", "1x1x128"})] = @@ -50,4 +49,5 @@ void bquant_quantgrouped_preshuffleb_bf8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8i4.cpp index 85599864db..3ffcfdac69 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_bf8i4.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_BQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_bf8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8.cpp index 87cb4c9d10..de7e290eeb 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_BQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_fp8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); lut[hash_multiple_strings({"fp8", "bquant", "preshuffleb", "non-preshufflequant", "1x1x128"})] = @@ -50,4 +49,5 @@ void bquant_quantgrouped_preshuffleb_fp8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8i4.cpp index 0cb16441a9..d36c20e700 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_fp8i4.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_BQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_fp8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8.cpp index 640757a956..12e23ba722 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_PreshuffleBQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_preshufflequant_bf8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); lut[hash_multiple_strings({"bf8", "bquant", "preshuffleb", "preshufflequant", "1x1x128"})] = @@ -47,4 +46,5 @@ void bquant_quantgrouped_preshuffleb_preshufflequant_bf8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4.cpp index 575a43afd8..cb8beee11d 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_PreshuffleBQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8.cpp index 9e40fbaa87..edfa144341 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_PreshuffleBQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_preshufflequant_fp8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); lut[hash_multiple_strings({"fp8", "bquant", "preshuffleb", "preshufflequant", "1x1x128"})] = @@ -47,4 +46,5 @@ void bquant_quantgrouped_preshuffleb_preshufflequant_fp8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4.cpp index 2552a1d134..c83dc0a396 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4.cpp @@ -17,9 +17,8 @@ using GemmConfig = GemmConfigPreshuffleB_PreshuffleBQuant_Prefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8.cpp index edb28236af..2aa54fa72b 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigPreshuffleBQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshufflequant_bf8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); lut[hash_multiple_strings({"bf8", "bquant", "non-preshuffleb", "preshufflequant", "1x1x128"})] = @@ -52,4 +51,5 @@ void bquant_quantgrouped_preshufflequant_bf8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8i4.cpp index 59da63447e..2ace775216 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_bf8i4.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigPreshuffleBQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshufflequant_bf8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8.cpp index 29c88001e8..aba9a146ca 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigPreshuffleBQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshufflequant_fp8_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig{}); lut[hash_multiple_strings({"fp8", "bquant", "non-preshuffleb", "preshufflequant", "1x1x128"})] = @@ -52,4 +51,5 @@ void bquant_quantgrouped_preshufflequant_fp8_instance_factory( using QuantGroupSize = ck_tile::QuantGroupShape>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8i4.cpp b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8i4.cpp index f487132557..e0e21cef13 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8i4.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_bquant_quantgrouped_preshufflequant_fp8i4.cpp @@ -12,9 +12,8 @@ using GemmConfig = GemmConfigPreshuffleBQuantPrefill; QuantGroupSize, \ ck_tile::QuantType::BQuantGrouped>(arg_parser); -void bquant_quantgrouped_preshufflequant_fp8i4_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); using TypeConfig = decltype(GemmQuantTypeConfig>; return RUN_GEMM_EXAMPLE_PREC_TYPE; }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant.cpp index 8de58b0a30..1fbe4d7b47 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_quant.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_quant.cpp @@ -95,51 +95,6 @@ auto gen_lut_key(const ck_tile::ArgParser& arg_parser) return hash_multiple_strings(params); } -void abquant_quantgrouped_instance_factory( - std::unordered_map>& lut); -void aquant_quantgrouped_instance_factory( - std::unordered_map>& lut); -void aquant_quantgrouped_preshufflequant_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_fp8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_bf8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_fp8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_bf8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_bf16fp4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_fp8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_bf8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_fp8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_bf8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshufflequant_fp8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshufflequant_bf8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshufflequant_fp8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshufflequant_bf8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_preshufflequant_fp8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_preshufflequant_bf8_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4_instance_factory( - std::unordered_map>& lut); -void bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4_instance_factory( - std::unordered_map>& lut); -void quant_rowcol_instance_factory( - std::unordered_map>& lut); -void quant_tensor_instance_factory( - std::unordered_map>& lut); - int main(int argc, char* argv[]) { auto [result, arg_parser] = create_args(argc, argv); @@ -153,29 +108,8 @@ int main(int argc, char* argv[]) std::cout << "Device ID: " << device_id << std::endl; ck_tile::hip_check_error(hipSetDevice(device_id)); - std::unordered_map> lut; - abquant_quantgrouped_instance_factory(lut); - aquant_quantgrouped_instance_factory(lut); - aquant_quantgrouped_preshufflequant_instance_factory(lut); - bquant_quantgrouped_fp8_instance_factory(lut); - bquant_quantgrouped_bf8_instance_factory(lut); - bquant_quantgrouped_fp8i4_instance_factory(lut); - bquant_quantgrouped_bf8i4_instance_factory(lut); - bquant_quantgrouped_bf16fp4_instance_factory(lut); - bquant_quantgrouped_preshuffleb_fp8_instance_factory(lut); - bquant_quantgrouped_preshuffleb_bf8_instance_factory(lut); - bquant_quantgrouped_preshuffleb_fp8i4_instance_factory(lut); - bquant_quantgrouped_preshuffleb_bf8i4_instance_factory(lut); - bquant_quantgrouped_preshufflequant_fp8_instance_factory(lut); - bquant_quantgrouped_preshufflequant_bf8_instance_factory(lut); - bquant_quantgrouped_preshufflequant_fp8i4_instance_factory(lut); - bquant_quantgrouped_preshufflequant_bf8i4_instance_factory(lut); - bquant_quantgrouped_preshuffleb_preshufflequant_fp8_instance_factory(lut); - bquant_quantgrouped_preshuffleb_preshufflequant_bf8_instance_factory(lut); - bquant_quantgrouped_preshuffleb_preshufflequant_fp8i4_instance_factory(lut); - bquant_quantgrouped_preshuffleb_preshufflequant_bf8i4_instance_factory(lut); - quant_rowcol_instance_factory(lut); - quant_tensor_instance_factory(lut); + auto& lut = get_kernel_lut(); + std::cout << "Available kernels: " << lut.size() << std::endl; auto key = gen_lut_key(arg_parser); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant_rowcol.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant_rowcol.cpp index 19c02b7ae2..d450a36f84 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_quant_rowcol.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_quant_rowcol.cpp @@ -6,9 +6,8 @@ template using GemmConfig = GemmConfigQuantDecode; -void quant_rowcol_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); // NOTE: QuantGroupSize is a place holder. rowcol pipeline does not use QuantGroupSize using QuantGroupSize = ck_tile::QuantGroupShape>; lut[hash_multiple_strings({"fp8", "rowcol"})] = [](const ck_tile::ArgParser& arg_parser) { @@ -27,4 +26,5 @@ void quant_rowcol_instance_factory( QuantGroupSize, ck_tile::QuantType::RowColQuant>(arg_parser); }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_quant_tensor.cpp b/example/ck_tile/38_block_scale_gemm/gemm_quant_tensor.cpp index 0deb3d890b..71b193d8f4 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_quant_tensor.cpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_quant_tensor.cpp @@ -6,9 +6,8 @@ template using GemmConfig = GemmConfigQuantDecode; -void quant_tensor_instance_factory( - std::unordered_map>& lut) -{ +static auto _ = []() { + auto& lut = get_kernel_lut(); // NOTE: QuantGroupSize is a place holder. tensor pipeline does not use QuantGroupSize using QuantGroupSize = ck_tile::QuantGroupShape>; lut[hash_multiple_strings({"fp8", "tensor"})] = [](const ck_tile::ArgParser& arg_parser) { @@ -27,4 +26,5 @@ void quant_tensor_instance_factory( QuantGroupSize, ck_tile::QuantType::TensorQuant>(arg_parser); }; -} + return 0; +}(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp index 37117eaa0f..0feba8e471 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp @@ -11,6 +11,14 @@ #include "ck_tile/ops/gemm.hpp" #include "ck_tile/ops/gemm_quant.hpp" +inline auto& get_kernel_lut() +{ + // In an inline function, function-local static objects in all function definitions are shared + // across all translation units. + static std::unordered_map> lut; + return lut; +} + inline size_t hash_multiple_strings(const std::vector& inputs) { std::hash hasher; diff --git a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc index ed1709a9ae..270653d408 100644 --- a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc +++ b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc @@ -80,10 +80,9 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str ck_tile::BaseGemmPipelineAgBgCrMem, ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2>>>; - const ck_tile::index_t K_split = - (args.K + GemmConfig::K_Tile - 1) / GemmConfig::K_Tile * GemmConfig::K_Tile; - const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); - const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); + const ck_tile::index_t K_split = ck_tile::integer_least_multiple(args.K, GemmConfig::K_Tile); + const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split); + const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop); const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop); const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) { @@ -553,8 +552,7 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, ck_tile::host_tensor_descriptor(1, 1, stride_BQ, is_row_major(bq_layout))); } - std::random_device rd; - std::mt19937 gen(rd()); + std::mt19937 gen(42); std::uniform_int_distribution fill_seed(0, 500); if(init_method == 0) @@ -630,7 +628,7 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, else if(init_method == 1) { std::cout << "Monotonic initialization is not supported." << std::endl; - return 0; + return -1; } else if(init_method == 2) { @@ -1078,10 +1076,10 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, else if(arg_parser.get_int("v") == 2) { std::cout << "GPU verification is not implemented yet. Re-run with -v=1" << std::endl; - return false; + return -1; } - return pass; + return pass ? 0 : -1; } // Usage of Two-Matrix Quantization (AB-Quant) template * t, int block_bq_k) } template -auto shuffle_b(const ck_tile::HostTensor& t, const GemmConfig& gemmConfig) +auto shuffle_b(const ck_tile::HostTensor& t, GemmConfig) { assert(t.get_lengths().size() == 2); int n_ = t.get_lengths()[1]; @@ -79,36 +79,40 @@ auto shuffle_b(const ck_tile::HostTensor& t, const GemmConfig& gemmConfig) { constexpr int divisor = 2; constexpr int kABK1PerLane = 8; - int kABK0PerLane = gemmConfig.K_Warp_Tile / divisor / kABK1PerLane; - ck_tile::HostTensor t_view({n_ / gemmConfig.N_Warp_Tile, - gemmConfig.N_Warp_Tile, - k_ / gemmConfig.K_Warp_Tile, + int kABK0PerLane = GemmConfig::K_Warp_Tile / divisor / kABK1PerLane; + ck_tile::HostTensor t_view({n_ / GemmConfig::N_Warp_Tile, + GemmConfig::N_Warp_Tile, + k_ / GemmConfig::K_Warp_Tile, kABK0PerLane, divisor, kABK1PerLane}); std::copy(t.begin(), t.end(), t_view.begin()); return ck_tile::reference_permute(t_view, {0, 2, 4, 1, 3, 5}); } - else + else if(ck_tile::is_gfx11_supported()) { int divisor = 1; - if(ck_tile::is_gfx11_supported()) - { - divisor = 1; - } - else - { - assert(is_wave32() == false); - divisor = get_warp_size() / gemmConfig.N_Warp_Tile; - } - ck_tile::HostTensor t_view({n_ / gemmConfig.N_Warp_Tile, - gemmConfig.N_Warp_Tile, - k_ / gemmConfig.K_Warp_Tile, + ck_tile::HostTensor t_view({n_ / GemmConfig::N_Warp_Tile, + GemmConfig::N_Warp_Tile, + k_ / GemmConfig::K_Warp_Tile, divisor, - gemmConfig.K_Warp_Tile / divisor}); + GemmConfig::K_Warp_Tile / divisor}); std::copy(t.begin(), t.end(), t_view.begin()); return ck_tile::reference_permute(t_view, {0, 2, 3, 1, 4}); } + else + { + constexpr int KLane = ck_tile::get_warp_size() / GemmConfig::N_Warp_Tile; + constexpr int ItemsPerAccess = + std::min(16 / static_cast(sizeof(T)), GemmConfig::K_Warp_Tile / KLane); + + ck_tile::HostTensor t_view({n_ / GemmConfig::N_Warp_Tile, + GemmConfig::N_Warp_Tile, + k_ / ItemsPerAccess, + ItemsPerAccess}); + std::copy(t.begin(), t.end(), t_view.begin()); + return ck_tile::reference_permute(t_view, {0, 2, 1, 3}); + } } template diff --git a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp index e123cee9e1..8074994fdd 100644 --- a/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp @@ -160,7 +160,7 @@ struct UniversalGemmBasePolicy constexpr auto K0PerThreadRead = AK0 / KThreadRead; // check if we exceed all LDS banks - constexpr auto LdsBanksWidth = get_n_lds_banks() * get_n_words_per_128b(); + constexpr auto LdsBanksWidth = get_n_lds_banks() * get_n_dwords_per_128b(); constexpr auto kfold = (AK1 * M0 * sizeof(ADataType) > LdsBanksWidth) ? 1 : LdsBanksWidth / (AK1 * M0 * sizeof(ADataType)); @@ -250,7 +250,7 @@ struct UniversalGemmBasePolicy constexpr uint64_t MinLdsLayer = 1ULL; constexpr auto MLdsLayer = max(MinLdsLayer, - get_n_lds_banks() * get_n_words_per_128b() / KPerBlock / DataTypeSize); + get_n_lds_banks() * get_n_dwords_per_128b() / KPerBlock / DataTypeSize); constexpr index_t NBanks = get_n_lds_banks(); static_assert(NBanks == 32 || NBanks == 64, "Unexpected LDS bank count"); @@ -357,7 +357,7 @@ struct UniversalGemmBasePolicy constexpr auto K0PerThreadRead = BK0 / KThreadRead; // check if we exceed all LDS banks - constexpr auto LdsBanksWidth = get_n_lds_banks() * get_n_words_per_128b(); + constexpr auto LdsBanksWidth = get_n_lds_banks() * get_n_dwords_per_128b(); constexpr auto kfold = (BK1 * N0 * sizeof(BDataType) > LdsBanksWidth) ? 1 : LdsBanksWidth / (BK1 * N0 * sizeof(BDataType)); @@ -450,7 +450,7 @@ struct UniversalGemmBasePolicy constexpr uint64_t MinLdsLayer = 1ULL; constexpr auto NLdsLayer = max(MinLdsLayer, - get_n_lds_banks() * get_n_words_per_128b() / KPerBlock / DataTypeSize); + get_n_lds_banks() * get_n_dwords_per_128b() / KPerBlock / DataTypeSize); constexpr index_t NBanks = get_n_lds_banks(); static_assert(NBanks == 32 || NBanks == 64, "Unexpected LDS bank count"); diff --git a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp index 1ff95b157c..1784436f87 100644 --- a/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp +++ b/include/ck_tile/ops/gemm/pipeline/wp_pipeline_agmem_bgmem_creg_base_policy.hpp @@ -151,6 +151,7 @@ struct UniversalWeightPreshufflePipelineAgBgCrPolicy CK_TILE_DEVICE static constexpr auto MakeBFlatDramTileDistribution() { using TileShape = typename Problem::BlockGemmShape; + using BDataType = typename Problem::BDataType; constexpr index_t kNPerBlock = TileShape::kN; constexpr index_t kKPerBlock = TileShape::kK; @@ -162,16 +163,18 @@ struct UniversalWeightPreshufflePipelineAgBgCrPolicy constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNum = BlockSize / WaveSize; - constexpr index_t KBPerLoad = GetKBPerLoad(); #if defined(__gfx11__) constexpr index_t KRepeatInWave = 2; #else constexpr index_t KRepeatInWave = 1; #endif + constexpr index_t KBPerLoad = min( + GetKBPerLoad(), KRepeatInWave * 16 / static_cast(sizeof(BDataType))); constexpr index_t KThdPerWave = WaveSize / KRepeatInWave; // threads cnt in K dim constexpr index_t KWavePerBlk = 1; constexpr index_t KRepeat = KIterPerWarp; - static_assert(TileShape::flatKPerWarp == KThdPerWave * KBPerLoad, "wrong"); + constexpr index_t KAccess = GetKBPerLoad() / KBPerLoad; + static_assert(TileShape::flatKPerWarp == KAccess * KThdPerWave * KBPerLoad, "wrong"); constexpr index_t NBPerLoad = 1; constexpr index_t NThdPerWave = 1; @@ -181,16 +184,16 @@ struct UniversalWeightPreshufflePipelineAgBgCrPolicy constexpr index_t WaveRepeat = WaveNum / TileShape::flatNPerWarp; return make_static_tile_distribution( tile_distribution_encoding< - sequence, // ? - tuple, // second direction - sequence>, // first direction + sequence, // ? + tuple, // second direction + sequence>, // wave in blk, // thd in wave // // tuple, sequence<0, 1, 2>>, // which direction - tuple, sequence<1, 2, 2>>, // which index + tuple, sequence<1, 2, 3>>, // which index // - sequence<1, 2, 1, 2>, - sequence<0, 0, 3, 3>>{}); + sequence<1, 2, 1, 2, 2>, + sequence<0, 0, 3, 1, 4>>{}); } template @@ -256,13 +259,22 @@ struct UniversalWeightPreshufflePipelineAgBgCrPolicy std::conditional_t, typename Problem::ADataType, typename Problem::BDataType>; - using WarpGemm = WarpGemmDispatcher; + constexpr index_t WaveSize = get_warp_size(); + constexpr index_t KLane = WarpTile::at(I2) * WarpTile::at(I0) / WaveSize; + using BDataType = typename Problem::BDataType; + constexpr index_t KLaneBytes = + KLane / numeric_traits::PackedSize * sizeof(BDataType); + constexpr auto NumAccess = static_cast(max(1, KLaneBytes / 16)); + using WarpGemm = WarpGemmDispatcher; using BlockWeightPreshufflePolicy = BlockWeightPreshuffleASmemBSmemCRegV1CustomPolicy{})); - index_t kFlatN = kargs.N * kargs.K / kFlatK; + constexpr auto warp_k = GemmPipeline::BlockGemmShape::WarpTile::at(I2); + index_t kFlatKSplit = GemmPipeline::flatKPerWarp * (k_size / warp_k); + index_t kFlatK = GemmPipeline::flatKPerWarp * (kargs.K / warp_k); + index_t kFlatN = kargs.N * kargs.K / kFlatK; return make_naive_tensor_view( b_ptr, - make_tuple(kFlatN, kFlatK), + make_tuple(kFlatN, kFlatKSplit), make_tuple(kFlatK, 1), number{}, number<1>{}); diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_base_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_base_policy.hpp index 80e41cad45..ae2a601f8a 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_base_policy.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_wp_abquant_pipeline_ag_bg_cr_base_policy.hpp @@ -52,11 +52,13 @@ struct GemmWPABQuantPipelineAgBgCrPolicy : public UniversalWeightPreshufflePipel CK_TILE_DEVICE static constexpr auto MakeBFlatDramTileDistribution() { using TileShape = typename Problem::BlockGemmShape; + using BDataType = typename Problem::BDataType; constexpr index_t BlockSize = Problem::kBlockSize; constexpr index_t WaveSize = get_warp_size(); constexpr index_t WaveNum = BlockSize / WaveSize; - constexpr index_t KBPerLoad = GetKBPerLoad(); + constexpr index_t KBPerLoad = + min(GetKBPerLoad(), 16 / static_cast(sizeof(BDataType))); #if defined(__gfx11__) constexpr index_t KRepeatInWave = 2; #else @@ -64,8 +66,8 @@ struct GemmWPABQuantPipelineAgBgCrPolicy : public UniversalWeightPreshufflePipel #endif constexpr index_t KThdPerWave = WaveSize / KRepeatInWave; // threads cnt in K dim constexpr index_t KWavePerBlk = 1; - constexpr index_t KRepeat = 1; - static_assert(TileShape::flatKPerWarp == KThdPerWave * KBPerLoad, "wrong"); + constexpr index_t KRepeat = GetKBPerLoad() / KBPerLoad; + static_assert(TileShape::flatKPerWarp == KRepeat * KThdPerWave * KBPerLoad, "wrong"); constexpr index_t NBPerLoad = 1; constexpr index_t NThdPerWave = 1; @@ -98,13 +100,23 @@ struct GemmWPABQuantPipelineAgBgCrPolicy : public UniversalWeightPreshufflePipel typename Problem::ADataType, typename Problem::BDataType>; + constexpr index_t WaveSize = get_warp_size(); + constexpr index_t KLane = WarpTile::at(I2) * WarpTile::at(I0) / WaveSize; + using BDataType = typename Problem::BDataType; + constexpr index_t KLaneBytes = + KLane / numeric_traits::PackedSize * sizeof(BDataType); + constexpr auto NumAccess = static_cast(max(1, KLaneBytes / 16)); + using WarpGemm = WarpGemmDispatcher; + Problem::TransposeC, + false, + false, + NumAccess>; // TODO : Use a custom block policy for AsBrCr using BlockGemmPolicy = From bc6083bdd466d1e060253e7a49626c923293c483 Mon Sep 17 00:00:00 2001 From: Johannes Graner Date: Wed, 28 Jan 2026 15:38:10 +0100 Subject: [PATCH 02/11] Update pytorch version in convolution dataset test generation (#3667) * Update torch version in dataset test gen --- .gitignore | 1 + test_data/generate_test_dataset.sh | 6 ++---- test_data/requirements.txt | 3 +++ 3 files changed, 6 insertions(+), 4 deletions(-) create mode 100644 test_data/requirements.txt diff --git a/.gitignore b/.gitignore index a2fb1473ab..17f93500bd 100644 --- a/.gitignore +++ b/.gitignore @@ -103,6 +103,7 @@ dispatcher/**/dispatcher_kernels.json test_data/* !test_data/*.py !test_data/*.sh +!test_data/requirements.txt # Exceptions to build* patterns above # The experimental/builder directory should be tracked despite matching build* diff --git a/test_data/generate_test_dataset.sh b/test_data/generate_test_dataset.sh index 27f45a3bc7..5cbc5514e6 100755 --- a/test_data/generate_test_dataset.sh +++ b/test_data/generate_test_dataset.sh @@ -50,10 +50,8 @@ if ! python3 -c "import torch" 2>/dev/null; then # Install PyTorch in virtual environment with ROCm support echo "Installing PyTorch and torchvision with ROCm support in virtual environment..." - # Since we're in a ROCm 6.4.1 environment, we need compatible PyTorch - # PyTorch doesn't have 6.4 wheels yet, so we use 6.2 which should be compatible - echo "Installing PyTorch with ROCm 6.2 support (compatible with ROCm 6.4)..." - pip install torch==2.5.1 torchvision==0.20.1 --index-url https://download.pytorch.org/whl/rocm6.2 || { + echo "Installing PyTorch with ROCm 7.1 support..." + pip install -r requirements.txt || { echo "ERROR: Failed to install PyTorch with ROCm support." echo "Creating empty CSV files as fallback..." echo "# 2D Convolution Test Cases" > conv_test_set_2d_dataset.csv diff --git a/test_data/requirements.txt b/test_data/requirements.txt new file mode 100644 index 0000000000..ecf05539f5 --- /dev/null +++ b/test_data/requirements.txt @@ -0,0 +1,3 @@ +-i https://download.pytorch.org/whl/rocm7.1 +torch==2.10.* +torchvision==0.25.* \ No newline at end of file From d6cccf6093807df5513338483b79b157378d57cc Mon Sep 17 00:00:00 2001 From: kabrahamAMD Date: Wed, 28 Jan 2026 17:33:45 +0100 Subject: [PATCH 03/11] [CK_BUILDER] Add reflection for wmma and bwd weight instances to ck builder reflection (#3592) * added reflection for conv_fwd_multiple_d_wmma_cshuffle.hpp * added reflection for device_grouped_conv_bwd_weight_xdl_cshuffle * added reflection for device_grouped_conv_bwd_weight_xdl_cshuffle v3 * added reflection of max_transpose parameters * fix printing of std optional parameters * fix use of undefined ck::index * added conv traits for device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle * added xdl two stage instance to reflection * added additional variables * added reflection for grouped_conv_bwd_weight_multiple_d_wmma_cshuffle, _v3, grouped_conv_two_stage_wmma_cshuffle_v3, * added reflection for device_grouped_conv_bwd_weigh_wmma_cshuffle_v3 * added reflection for bwd_weight_wmma_cshuffle * added comments back in * add printed output for optional parameters * update README * fix typo * added num_gemm_k_prefetch_stage and small fixes * modified test string due to reflection of new parameter --------- Co-authored-by: Kevin Abraham --- .../include/ck_tile/builder/reflect/README.md | 19 +- .../ck_tile/builder/reflect/conv_describe.hpp | 25 +- .../builder/reflect/conv_description.hpp | 201 ++- .../ck_tile/builder/reflect/conv_traits.hpp | 8 +- ...bwd_weight_multiple_d_wmma_cshuffle_v3.hpp | 46 + ...onv_bwd_weight_multiple_d_xdl_cshuffle.hpp | 53 + ..._bwd_weight_two_stage_wmma_cshuffle_v3.hpp | 50 + ...conv_bwd_weight_two_stage_xdl_cshuffle.hpp | 57 + ..._grouped_conv_bwd_weight_wmma_cshuffle.hpp | 48 + ...ouped_conv_bwd_weight_wmma_cshuffle_v3.hpp | 50 + ...e_grouped_conv_bwd_weight_xdl_cshuffle.hpp | 56 + ...rouped_conv_bwd_weight_xdl_cshuffle_v3.hpp | 53 + ...ped_conv_fwd_multiple_abd_xdl_cshuffle.hpp | 4 +- ..._conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp | 4 +- ...uped_conv_fwd_multiple_d_wmma_cshuffle.hpp | 46 + ...d_multiple_d_xdl_large_tensor_cshuffle.hpp | 4 +- .../builder/reflect/conv_traits_helpers.hpp | 365 +++--- .../reflect/instance_to_conv_traits.hpp | 12 + ...bwd_weight_multiple_d_wmma_cshuffle_v3.hpp | 71 +- ...onv_bwd_weight_multiple_d_xdl_cshuffle.hpp | 66 +- ..._bwd_weight_two_stage_wmma_cshuffle_v3.hpp | 71 +- ...conv_bwd_weight_two_stage_xdl_cshuffle.hpp | 68 +- ..._grouped_conv_bwd_weight_wmma_cshuffle.hpp | 71 +- ...ouped_conv_bwd_weight_wmma_cshuffle_v3.hpp | 71 +- ...e_grouped_conv_bwd_weight_xdl_cshuffle.hpp | 88 +- ...rouped_conv_bwd_weight_xdl_cshuffle_v3.hpp | 70 +- ...uped_conv_fwd_multiple_d_wmma_cshuffle.hpp | 7 + .../builder/test/conv/ck/test_conv_traits.cpp | 1102 +++++++++++++++++ .../builder/test/test_conv_description.cpp | 258 +++- 29 files changed, 2555 insertions(+), 489 deletions(-) create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp create mode 100644 experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp diff --git a/experimental/builder/include/ck_tile/builder/reflect/README.md b/experimental/builder/include/ck_tile/builder/reflect/README.md index 8bb9c89c80..43192b7c48 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/README.md +++ b/experimental/builder/include/ck_tile/builder/reflect/README.md @@ -9,6 +9,7 @@ See the [main builder documentation](../README.md) for an overview. The reflection system works by extracting properties from a convolution kernel *type* and formatting them into a string. This is useful for debugging, performance tuning, and generating documentation. 1. **Trait Extraction**: The `ConvTraits` template (in `conv_traits.hpp`) is specialized for each kernel instance. It extracts low-level details like tile sizes, data layouts, and pipeline versions from the kernel's type definition. +This template is common for xld and wmma, fwd and backwards weight kernels. std::optional is used for parameters that are only used by some kernels 2. **Description Generation**: The `describe()` function (in `conv_description.hpp`) uses `ConvTraits` to populate a `ConvDescription` (`Description`) object. @@ -48,6 +49,15 @@ The reflection system (`ckr::describe`) currently supports the following convolu - **Standard XDL Forward Convolution** (`DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle`) - **Large Tensor XDL Forward Convolution** (`DeviceGroupedConvFwdMultipleD_Xdl_CShuffle_Large_Tensor`) - **V3 XDL Forward Convolution** (`DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3`) +- **V3 WMMA Forward Convolution** (`DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3`) +- **XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Xdl_CShuffle`) +- **V3 XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Xdl_CShuffleV3`) +- **XDL Multiple D Backward Weight Convolution** (`DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle`) +- **Two Stage XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle`) +- **V3 Two Stage XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3`) +- **Wmma Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Wmma_CShuffle`) +- **V3 Wmma Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Wmma_CShuffleV3`) +- **V3 Wmma Multiple D Backward Weight Convolution** (`DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3`) These variants all share similar template parameter structures and are compatible with the current `ConvTraits` implementation. @@ -59,15 +69,6 @@ The following instance types are **not yet supported** by the reflection system: - Uses different internal structure with parameters like `K0PerBlock`, `K1`, `M1PerThread`, etc. - Missing standard members like `kKPerBlock`, `kMPerXDL`, `kAK1` -- **WMMA Variants** (`DeviceGroupedConvFwdMultipleD_Wmma_CShuffle`) - - Uses WMMA-specific parameters like `MPerWmma`, `NPerWmma`, `MRepeat`, `NRepeat` - - Different tile transfer structure incompatible with current `ConvTraits` - -- **Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Xdl_CShuffle`) - - Uses different layout naming: `InLayout`, `WeiLayout`, `OutLayout` instead of `ALayout`, `BLayout`, `ELayout` - - Different specialization type: `ConvBackwardWeightSpecialization` vs `ConvForwardSpecialization` - - Missing several members expected by forward convolution traits - ### Future Work To support these additional instance types, the reflection system would need: diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_describe.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_describe.hpp index 359b12c4a3..27e7dfb362 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_describe.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_describe.hpp @@ -29,30 +29,7 @@ conv::ConvDescription describe() const auto traits = conv::instance_to_conv_traits(); return conv::ConvDescription( - conv::ConvSignatureInfo{ - .spatial_dim = traits.spatial_dim, - .direction = traits.direction, - .input_layout = traits.layout[0], - .weight_layout = traits.layout[1], - .output_layout = traits.layout[2], - .data_type = traits.data_type, - .input_element_op = traits.input_element_op, - .weight_element_op = traits.weight_element_op, - .output_element_op = traits.output_element_op, - }, - conv::GemmAlgorithmInfo{ - .thread_block_size = traits.thread_block_size, - .tile_dims = traits.tile_dims, - .warp_gemm = traits.warp_gemm, - .a_tile_transfer = traits.a_tile_transfer, - .b_tile_transfer = traits.b_tile_transfer, - .c_tile_transfer = traits.c_tile_transfer, - .pipeline_version = traits.pipeline_version, - .pipeline_scheduler = traits.pipeline_scheduler, - .conv_specialization = traits.conv_specialization, - .padding = traits.gemm_padding, - }, - []() { return reflect::instance_string(); }); + traits, []() { return reflect::instance_string(); }); } } // namespace ck_tile::reflect diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_description.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_description.hpp index a7b6c60a73..5c09e4b735 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_description.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_description.hpp @@ -29,44 +29,12 @@ #include #include #include +#include namespace ck_tile::reflect { namespace conv { -/// @brief Signature information for a convolution operation -/// Contains high-level properties that define the convolution's interface, -/// including dimensionality, data layout, data types, and elementwise operations. -struct ConvSignatureInfo -{ - int spatial_dim; - builder::ConvDirection direction; - builder::TensorLayout input_layout; - builder::TensorLayout weight_layout; - builder::TensorLayout output_layout; - builder::DataType data_type; - builder::ElementwiseOperation input_element_op; - builder::ElementwiseOperation weight_element_op; - builder::ElementwiseOperation output_element_op; -}; - -/// @brief Algorithm configuration for a convolution kernel -/// Contains low-level implementation details including thread block configuration, -/// tile dimensions, memory access patterns, and pipeline settings. -struct GemmAlgorithmInfo -{ - int thread_block_size; - DataTileInfo tile_dims; - WarpGemmParams warp_gemm; - InputTileTransferInfo a_tile_transfer; - InputTileTransferInfo b_tile_transfer; - OutputTileTransferInfo c_tile_transfer; - builder::PipelineVersion pipeline_version; - builder::PipelineScheduler pipeline_scheduler; - builder::ConvSpecialization conv_specialization; - builder::GemmPadding padding; -}; - /// @brief Provides human-readable descriptions of convolution kernel instances /// Generates formatted text descriptions at various levels of detail for /// understanding and documenting convolution kernel configurations. @@ -74,16 +42,12 @@ class ConvDescription : public Description { public: /// @brief Constructor for ConvDescription - /// @param sig The signature information containing high-level convolution properties - /// @param algo The algorithm configuration containing low-level implementation details + /// @param traits The ConvTraits object containing all relevant signature and algorithm + /// information /// @param instance_string_getter A callable that returns a string representation of the /// instance - ConvDescription(ConvSignatureInfo sig, - GemmAlgorithmInfo algo, - std::function instance_string_getter) - : signature_(std::move(sig)), - algorithm_(std::move(algo)), - instance_string_getter_(std::move(instance_string_getter)) + ConvDescription(ConvTraits traits, std::function instance_string_getter) + : traits_(std::move(traits)), instance_string_getter_(std::move(instance_string_getter)) { } @@ -92,7 +56,7 @@ class ConvDescription : public Description std::string brief() const override { std::ostringstream oss; - oss << signature_.spatial_dim << "D " << signature_.direction << " convolution"; + oss << traits_.spatial_dim << "D " << traits_.direction << " convolution"; return oss.str(); } @@ -101,39 +65,42 @@ class ConvDescription : public Description std::string detailed() const override { TreeFormatter f; - f.writeLine(0, signature_.spatial_dim, "D ", signature_.direction, " Convolution Kernel"); + f.writeLine(0, traits_.spatial_dim, "D ", traits_.direction, " Convolution Kernel"); f.writeLine(1, "Signature"); - f.writeLine(2, "Tensor Type: ", signature_.data_type); - f.writeLine(2, "Input Layout: ", signature_.input_layout); - f.writeLine(2, "Weight Layout: ", signature_.weight_layout); - f.writeLine(2, "Output Layout: ", signature_.output_layout); - f.writeLine(2, "Input elementwise operation: ", signature_.input_element_op); - f.writeLine(2, "Weights elementwise operation: ", signature_.weight_element_op); - f.writeLast(2, "Output elementwise operation: ", signature_.output_element_op); + f.writeLine(2, "Tensor Type: ", traits_.data_type); + f.writeLine(2, "Input Layout: ", traits_.layout[0]); + f.writeLine(2, "Weight Layout: ", traits_.layout[1]); + f.writeLine(2, "Output Layout: ", traits_.layout[2]); + f.writeLine(2, "Input elementwise operation: ", traits_.input_element_op); + f.writeLine(2, "Weights elementwise operation: ", traits_.weight_element_op); + f.writeLast(2, "Output elementwise operation: ", traits_.output_element_op); f.writeLast(1, "Algorithm"); // Compute Block section - f.writeLine(2, "Thread block size: ", algorithm_.thread_block_size); + f.writeLine(2, "Thread block size: ", traits_.thread_block_size); f.writeLine(2, "Data tile size: ", - algorithm_.tile_dims.m, + traits_.tile_dims.m, "×", - algorithm_.tile_dims.n, + traits_.tile_dims.n, "×", - algorithm_.tile_dims.k); - f.writeLine(2, "Gemm padding: ", algorithm_.padding); - f.writeLine(2, "Convolution specialization: ", algorithm_.conv_specialization); + traits_.tile_dims.k); + if(traits_.gemm_padding) + f.writeLine( + 2, "Gemm padding: ", traits_.gemm_padding.value_or(builder::GemmPadding::DEFAULT)); + else + f.writeLine(2, "Struct does not contain optional gemm_padding argument"); + f.writeLine(2, "Convolution specialization: ", traits_.conv_specialization); // Pipeline section - f.writeLine(2, "Pipeline version: ", algorithm_.pipeline_version); - f.writeLine(2, "Pipeline scheduler: ", algorithm_.pipeline_scheduler); + f.writeLine(2, "Pipeline version: ", traits_.pipeline_version); + f.writeLine(2, "Pipeline scheduler: ", traits_.pipeline_scheduler); f.writeLine(2, "Warp Gemm parameters: "); - f.writeLine( - 3, "subtile size: ", algorithm_.warp_gemm.gemm_m, "×", algorithm_.warp_gemm.gemm_n); + f.writeLine(3, "subtile size: ", traits_.warp_gemm.gemm_m, "×", traits_.warp_gemm.gemm_n); f.writeLast(3, "Number of warp gemm iterations: ", - algorithm_.warp_gemm.m_iter, + traits_.warp_gemm.m_iter, "×", - algorithm_.warp_gemm.n_iter); + traits_.warp_gemm.n_iter); // Memory Access section f.writeLast(2, "Memory access:"); @@ -141,99 +108,126 @@ class ConvDescription : public Description f.writeLine(3, "A Tile transfer: "); f.writeLine(4, "Tile dimensions: ", - algorithm_.a_tile_transfer.tile_dimensions.k0, + traits_.a_tile_transfer.tile_dimensions.k0, "×", - algorithm_.a_tile_transfer.tile_dimensions.m_or_n, + traits_.a_tile_transfer.tile_dimensions.m_or_n, "×", - algorithm_.a_tile_transfer.tile_dimensions.k1, + traits_.a_tile_transfer.tile_dimensions.k1, "×"); - f.writeLine(4, - "The innermost K subdimension size: ", - algorithm_.a_tile_transfer.transfer_params.k1); + f.writeLine( + 4, "The innermost K subdimension size: ", traits_.a_tile_transfer.transfer_params.k1); f.writeLine(4, "Spatial thread distribution over the data tile: ", - algorithm_.a_tile_transfer.transfer_params.thread_cluster_order[0], + traits_.a_tile_transfer.transfer_params.thread_cluster_order[0], "×", - algorithm_.a_tile_transfer.transfer_params.thread_cluster_order[1], + traits_.a_tile_transfer.transfer_params.thread_cluster_order[1], "×", - algorithm_.a_tile_transfer.transfer_params.thread_cluster_order[2]); + traits_.a_tile_transfer.transfer_params.thread_cluster_order[2]); f.writeLine(4, "The order of accessing data tile axes: ", - algorithm_.a_tile_transfer.transfer_params.src_access_order[0], + traits_.a_tile_transfer.transfer_params.src_access_order[0], "×", - algorithm_.a_tile_transfer.transfer_params.src_access_order[1], + traits_.a_tile_transfer.transfer_params.src_access_order[1], "×", - algorithm_.a_tile_transfer.transfer_params.src_access_order[2]); + traits_.a_tile_transfer.transfer_params.src_access_order[2]); f.writeLine(4, "Vectorized memory access axis index (with contiguous memory): ", - algorithm_.a_tile_transfer.transfer_params.src_vector_dim); + traits_.a_tile_transfer.transfer_params.src_vector_dim); f.writeLine(4, "Vector access (GMEM read) instruction size: ", - algorithm_.a_tile_transfer.transfer_params.src_scalar_per_vector); + traits_.a_tile_transfer.transfer_params.src_scalar_per_vector); f.writeLine(4, "Vector access (LDS write) instruction size: ", - algorithm_.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1); + traits_.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1); f.writeLast(4, "LDS data layout padding (to prevent bank conflicts): ", - algorithm_.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1); + traits_.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1); f.writeLine(3, "B Tile transfer: "); f.writeLine(4, "Tile dimensions: ", - algorithm_.b_tile_transfer.tile_dimensions.k0, + traits_.b_tile_transfer.tile_dimensions.k0, "×", - algorithm_.b_tile_transfer.tile_dimensions.m_or_n, + traits_.b_tile_transfer.tile_dimensions.m_or_n, "×", - algorithm_.b_tile_transfer.tile_dimensions.k1, + traits_.b_tile_transfer.tile_dimensions.k1, "×"); - f.writeLine(4, - "The innermost K subdimension size: ", - algorithm_.b_tile_transfer.transfer_params.k1); + f.writeLine( + 4, "The innermost K subdimension size: ", traits_.b_tile_transfer.transfer_params.k1); f.writeLine(4, "Spatial thread distribution over the data tile: ", - algorithm_.b_tile_transfer.transfer_params.thread_cluster_order[0], + traits_.b_tile_transfer.transfer_params.thread_cluster_order[0], "×", - algorithm_.b_tile_transfer.transfer_params.thread_cluster_order[1], + traits_.b_tile_transfer.transfer_params.thread_cluster_order[1], "×", - algorithm_.b_tile_transfer.transfer_params.thread_cluster_order[2]); + traits_.b_tile_transfer.transfer_params.thread_cluster_order[2]); f.writeLine(4, "The order of accessing data tile axes: ", - algorithm_.b_tile_transfer.transfer_params.src_access_order[0], + traits_.b_tile_transfer.transfer_params.src_access_order[0], "×", - algorithm_.b_tile_transfer.transfer_params.src_access_order[1], + traits_.b_tile_transfer.transfer_params.src_access_order[1], "×", - algorithm_.b_tile_transfer.transfer_params.src_access_order[2]); + traits_.b_tile_transfer.transfer_params.src_access_order[2]); f.writeLine(4, "Vectorized memory access axis index (with contiguous memory): ", - algorithm_.b_tile_transfer.transfer_params.src_vector_dim); + traits_.b_tile_transfer.transfer_params.src_vector_dim); f.writeLine(4, "Vector access (GMEM read) instruction size: ", - algorithm_.b_tile_transfer.transfer_params.src_scalar_per_vector); + traits_.b_tile_transfer.transfer_params.src_scalar_per_vector); f.writeLine(4, "Vector access (LDS write) instruction size: ", - algorithm_.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1); + traits_.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1); f.writeLast(4, "LDS data layout padding (to prevent bank conflicts): ", - algorithm_.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1); + traits_.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1); f.writeLast(3, "C Tile transfer: "); f.writeLine(4, "Data shuffle (number of gemm instructions per iteration): ", - algorithm_.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, + traits_.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, "×", - algorithm_.c_tile_transfer.shuffle_params.n_gemms_per_shuffle); + traits_.c_tile_transfer.shuffle_params.n_gemms_per_shuffle); f.writeLine(4, "Spatial thread distribution used to store data: ", - algorithm_.c_tile_transfer.thread_cluster_dims[0], + traits_.c_tile_transfer.thread_cluster_dims[0], "×", - algorithm_.c_tile_transfer.thread_cluster_dims[1], + traits_.c_tile_transfer.thread_cluster_dims[1], "×", - algorithm_.c_tile_transfer.thread_cluster_dims[2], + traits_.c_tile_transfer.thread_cluster_dims[2], "×", - algorithm_.c_tile_transfer.thread_cluster_dims[3]); - f.writeLast(4, + traits_.c_tile_transfer.thread_cluster_dims[3]); + f.writeLine(4, "Vector access (GMEM write) instruction size: ", - algorithm_.c_tile_transfer.scalar_per_vector); + traits_.c_tile_transfer.scalar_per_vector); + if(traits_.num_gemm_k_prefetch_stage) + f.writeLine( + 2, "Num gemm k prefetch stage: ", traits_.num_gemm_k_prefetch_stage.value_or(0)); + else + f.writeLine(2, + "Struct does not contain optional " + "num_gemm_k_prefetch_stage parameter"); + + if(traits_.max_transpose_transfer_src_scalar_per_vector) + f.writeLine(2, + "Max Transpose transfer scr scalar per vector: ", + traits_.max_transpose_transfer_src_scalar_per_vector.value_or(0)); + else + f.writeLine(2, + "Struct does not contain optional " + "max_transpose_transfer_src_scalar_per_vector parameter"); + if(traits_.max_transpose_dst_scalar_per_vector) + f.writeLine(2, + "Max Transpose dst scalar per vector: ", + traits_.max_transpose_dst_scalar_per_vector.value_or(0)); + else + f.writeLine( + 2, + "Struct does not contain optional max_transpose_dst_scalar_per_vector parameter"); + if(traits_.num_groups_to_merge) + f.writeLast(2, "Num groups to merge: ", traits_.num_groups_to_merge.value_or(0)); + else + f.writeLast(2, "Struct does not contain optional num_groups_to_merge parameter"); + return f.getString(); } @@ -242,8 +236,7 @@ class ConvDescription : public Description std::string instance_string() const override { return instance_string_getter_(); } private: - ConvSignatureInfo signature_; - GemmAlgorithmInfo algorithm_; + ConvTraits traits_; std::function instance_string_getter_; }; diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits.hpp index 451a74be34..16a9c47f7e 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_traits.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits.hpp @@ -88,7 +88,7 @@ struct ConvTraits builder::ElementwiseOperation weight_element_op; builder::ElementwiseOperation output_element_op; - builder::GemmPadding gemm_padding; + std::optional gemm_padding = std::nullopt; builder::ConvSpecialization conv_specialization; // --- Algorithm Information --- @@ -102,8 +102,14 @@ struct ConvTraits OutputTileTransferInfo c_tile_transfer; + std::optional num_gemm_k_prefetch_stage = std::nullopt; + builder::PipelineVersion pipeline_version; builder::PipelineScheduler pipeline_scheduler; + + std::optional max_transpose_transfer_src_scalar_per_vector = std::nullopt; + std::optional max_transpose_dst_scalar_per_vector = std::nullopt; + std::optional num_groups_to_merge = std::nullopt; }; } // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp new file mode 100644 index 0000000000..f052a9701b --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp @@ -0,0 +1,46 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_multiple_d_Wmma_CShuffle_V3_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kKPerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kKPerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kKPerBlock), + .warp_gemm = conv_traits_wmma_warp_gemm_params(), + .c_tile_transfer = conv_traits_wmma_c_tile_transfer(), + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp new file mode 100644 index 0000000000..2f7c68458f --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp @@ -0,0 +1,53 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_Xdl_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_multiple_d_Xdl_CShuffle_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kK0PerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .warp_gemm = conv_traits_xdl_warp_gemm_params(), + .c_tile_transfer = + {.shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle, + .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle}, + .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0], + InstTraits::kCThreadClusterLengths[1], + InstTraits::kCThreadClusterLengths[2], + InstTraits::kCThreadClusterLengths[3]}, + .scalar_per_vector = InstTraits::kCBlockTransferScalarPerVector_NWaveNPerXdl}, + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp new file mode 100644 index 0000000000..4f39b00b5c --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp @@ -0,0 +1,50 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_wmma_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_two_stage_Wmma_CShuffle_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kKPerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kABK1, InstTraits::kKPerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kABK1, InstTraits::kKPerBlock), + .warp_gemm = conv_traits_wmma_warp_gemm_params(), + .c_tile_transfer = conv_traits_wmma_c_tile_transfer(), + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + .max_transpose_transfer_src_scalar_per_vector = + InstTraits::kTransposeTransferSrcScalarPerVector, + .max_transpose_dst_scalar_per_vector = InstTraits::kTransposeTransferDstScalarPerVector, + .num_groups_to_merge = InstTraits::kNumGroupsToMerge, + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp new file mode 100644 index 0000000000..5666233091 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp @@ -0,0 +1,57 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdTwoStage_Xdl_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_two_stage_Xdl_CShuffle_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kKPerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kKPerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kKPerBlock), + .warp_gemm = conv_traits_xdl_warp_gemm_params(), + .c_tile_transfer = + {.shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle, + .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle}, + .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0], + InstTraits::kCThreadClusterLengths[1], + InstTraits::kCThreadClusterLengths[2], + InstTraits::kCThreadClusterLengths[3]}, + .scalar_per_vector = InstTraits::kCBlockTransferScalarPerVector_NWaveNPerXdl}, + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + .max_transpose_transfer_src_scalar_per_vector = + InstTraits::kTransposeTransferSrcScalarPerVector, + .max_transpose_dst_scalar_per_vector = InstTraits::kTransposeTransferDstScalarPerVector, + .num_groups_to_merge = InstTraits::kNumGroupsToMerge, + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp new file mode 100644 index 0000000000..470a10d031 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp @@ -0,0 +1,48 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_Wmma_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_Wmma_CShuffle_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kK0PerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .warp_gemm = conv_traits_wmma_warp_gemm_params(), + .c_tile_transfer = conv_traits_wmma_c_tile_transfer(), + .num_gemm_k_prefetch_stage = InstTraits::kNumGemmKPrefetchStage, + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp new file mode 100644 index 0000000000..13625aa182 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp @@ -0,0 +1,50 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_Wmma_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kKPerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kKPerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kKPerBlock), + .warp_gemm = conv_traits_wmma_warp_gemm_params(), + .c_tile_transfer = conv_traits_wmma_c_tile_transfer(), + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + .max_transpose_transfer_src_scalar_per_vector = + InstTraits::kMaxTransposeTransferSrcScalarPerVector, + .max_transpose_dst_scalar_per_vector = InstTraits::kMaxTransposeTransferDstScalarPerVector, + + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp new file mode 100644 index 0000000000..39fde33217 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp @@ -0,0 +1,56 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_Xdl_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_Xdl_CShuffle_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kK0PerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .warp_gemm = conv_traits_xdl_warp_gemm_params(), + .c_tile_transfer = + {.shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle, + .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle}, + .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0], + InstTraits::kCThreadClusterLengths[1], + InstTraits::kCThreadClusterLengths[2], + InstTraits::kCThreadClusterLengths[3]}, + .scalar_per_vector = InstTraits::kCBlockTransferScalarPerVector_NWaveNPerXdl}, + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + .max_transpose_transfer_src_scalar_per_vector = + InstTraits::kMaxTransposeTransferSrcScalarPerVector, + .max_transpose_dst_scalar_per_vector = InstTraits::kMaxTransposeTransferDstScalarPerVector, + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp new file mode 100644 index 0000000000..de98645514 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -0,0 +1,53 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_Xdl_V3_CShuffle_Tag +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = bwd_wei_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(InstTraits::kK0PerBlock), + .a_tile_transfer = + conv_traits_a_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .b_tile_transfer = + conv_traits_b_transfer_params(InstTraits::kK1, InstTraits::kK0PerBlock), + .warp_gemm = conv_traits_xdl_warp_gemm_params(), + .c_tile_transfer = + {.shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle, + .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle}, + .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0], + InstTraits::kCThreadClusterLengths[1], + InstTraits::kCThreadClusterLengths[2], + InstTraits::kCThreadClusterLengths[3]}, + .scalar_per_vector = InstTraits::kCBlockTransferScalarPerVector_NWaveNPerXdl}, + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp index cdd238f36a..2f5d84a4a8 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp @@ -24,8 +24,8 @@ constexpr ConvTraits instance_to_conv_traits() return ConvTraits{ .spatial_dim = InstTraits::kSpatialDim, .direction = conv_direction(), - .layout = conv_layout(), - .data_type = conv_data_type(), + .layout = fwd_conv_layout(), + .data_type = conv_data_type(), .input_element_op = elementwise_op(), .weight_element_op = elementwise_op(), .output_element_op = elementwise_op(), diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp index 28c43c342f..2108c79054 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp @@ -24,8 +24,8 @@ constexpr ConvTraits instance_to_conv_traits() return ConvTraits{ .spatial_dim = InstTraits::kSpatialDim, .direction = conv_direction(), - .layout = conv_layout(), - .data_type = conv_data_type(), + .layout = fwd_conv_layout(), + .data_type = conv_data_type(), .input_element_op = elementwise_op(), .weight_element_op = elementwise_op(), .output_element_op = elementwise_op(), diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp new file mode 100644 index 0000000000..9413107df7 --- /dev/null +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp @@ -0,0 +1,46 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#pragma once + +#include + +#include "ck_tile/builder/reflect/conv_traits.hpp" +#include "ck_tile/builder/reflect/conv_traits_helpers.hpp" +#include "ck_tile/builder/reflect/instance_traits.hpp" +#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp" + +namespace ck_tile::reflect::conv { + +/// @brief Tag dispatch implementation for DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle +template + requires HasInstanceTraits && + std::same_as::device_kernel_tag, + DeviceGroupedConvFwdMultipleD_Wmma_CShuffle_Tag> +constexpr ConvTraits instance_to_conv_traits() +{ + using InstTraits = InstanceTraits; + + return ConvTraits{ + .spatial_dim = InstTraits::kSpatialDim, + .direction = conv_direction(), + .layout = fwd_conv_layout(), + .data_type = conv_data_type(), + .input_element_op = elementwise_op(), + .weight_element_op = elementwise_op(), + .output_element_op = elementwise_op(), + .gemm_padding = gemm_spec(), + .conv_specialization = conv_spec(), + .thread_block_size = InstTraits::kBlockSize, + .tile_dims = conv_traits_data_tile(), + .a_tile_transfer = conv_traits_a_transfer_params(InstTraits::kK1), + .b_tile_transfer = conv_traits_b_transfer_params(InstTraits::kK1), + .warp_gemm = conv_traits_wmma_warp_gemm_params(), + .c_tile_transfer = conv_traits_wmma_c_tile_transfer(), + .num_gemm_k_prefetch_stage = InstTraits::kNumGemmKPrefetchStage, + .pipeline_version = get_pipeline_version(), + .pipeline_scheduler = get_pipeline_scheduler(), + }; +} + +} // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_xdl_large_tensor_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_xdl_large_tensor_cshuffle.hpp index c4bed850eb..0cce3bf513 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_xdl_large_tensor_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_xdl_large_tensor_cshuffle.hpp @@ -24,8 +24,8 @@ constexpr ConvTraits instance_to_conv_traits() return ConvTraits{ .spatial_dim = InstTraits::kSpatialDim, .direction = conv_direction(), - .layout = conv_layout(), - .data_type = conv_data_type(), + .layout = fwd_conv_layout(), + .data_type = conv_data_type(), .input_element_op = elementwise_op(), .weight_element_op = elementwise_op(), .output_element_op = elementwise_op(), diff --git a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_helpers.hpp b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_helpers.hpp index 46c196e95a..4baf2423ee 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/conv_traits_helpers.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/conv_traits_helpers.hpp @@ -80,6 +80,22 @@ namespace ck_tile::reflect::conv { // SECTION 1: ENUM CONVERSIONS // ============================================================================ +// Forward convolution layout concept - checks for A/B/E layout types +template +concept HasFwdConvLayouts = requires { + typename T::ALayout; + typename T::BLayout; + typename T::ELayout; +}; + +// Backwards weight layout concept - checks for In, wei and out layouts +template +concept HasBwdWeiLayouts = requires { + typename T::InLayout; + typename T::WeiLayout; + typename T::OutLayout; +}; + /// @brief Converts a CK BlockGemmPipelineVersion enum to a builder PipelineVersion enum. /// @tparam ck_ver The CK BlockGemmPipelineVersion enum value to convert. /// @return The corresponding builder::PipelineVersion enum value. @@ -322,12 +338,25 @@ constexpr builder::ConvSpecialization conv_spec() // Tensor Layouts // ---------------------------------------------------------------------------- +// Helper variable template to check if CK layout enums match +template +inline constexpr bool layouts_are = + std::is_same_v && std::is_same_v && std::is_same_v; + /// @brief Helper function to report unsupported layout combinations with a clear error message. /// @details This consteval function uses throw (not static_assert) to ensure the error is not /// silently ignored during SFINAE. The thrown string becomes part of the compiler error message. +/// @details This consteval function is designed to fail at compile time with a descriptive +/// error message when an unsupported layout combination is encountered. template [[noreturn]] consteval void report_unsupported_layout_error() { + // This will produce a compile-time error with the exception message throw "Unsupported convolution layout combination detected!\n" "The combination of ALayout, BLayout, and ELayout template parameters\n" "is not recognized for the given spatial dimension.\n" @@ -335,111 +364,99 @@ template "Check the conv_layout() function for the list of supported layout combinations."; } -/// @brief Derives the grouped convolution layout from a device kernel Instance type. -/// @tparam Instance The device kernel instance type. -/// @return An std::array containing the layouts for: -/// - [0] Input tensor layout -/// - [1] Weight tensor layout -/// - [2] Output tensor layout -/// @details This function examines the Instance's ALayout, BLayout, and ELayout types -/// along with the spatial dimension to determine the appropriate layout configuration. -/// -/// Supported layout combinations vary by spatial dimension (1D, 2D, 3D convolutions). -/// Common patterns include GNHWC (grouped, batch, spatial, channels) and variants. -/// -/// @note Compilation will fail with a clear error message if the layout combination -/// is not supported for the given spatial dimension. -/// -/// TODO: If we don't check for supported layouts, this function can be simplified. -template -constexpr std::array conv_layout() +template +constexpr auto conv_layout() { - using InstTraits = InstanceTraits; - using A = typename InstTraits::ALayout; - using B = typename InstTraits::BLayout; - using E = typename InstTraits::ELayout; - namespace ctl = ck::tensor_layout::convolution; + + // Helper lambda to construct layout array + auto layouts = [](auto... Ls) { return std::array{Ls...}; }; + + namespace ctl = ck::tensor_layout::convolution; using enum builder::TensorLayout; - // Helper to check if layouts match expected types - constexpr auto layouts_match = []() { - return std::is_same_v && std::is_same_v && std::is_same_v; - }; + switch(kSpatialDim) + { + case 1: + if constexpr(layouts_are) + return layouts(GNWC, GKXC, GNWK); + if constexpr(layouts_are) + return layouts(GNWC, GKXC, GNWK); + if constexpr(layouts_are) + return layouts(NWGC, GKXC, NWGK); + if constexpr(layouts_are) + return layouts(NGCW, GKXC, NGKW); + if constexpr(layouts_are) + return layouts(NGCW, GKCX, NGKW); + break; + case 2: + if constexpr(layouts_are) + return layouts(GNHWC, GKYXC, GNHWK); + if constexpr(layouts_are) + return layouts(GNHWC, GKYXC, GNHWK); + if constexpr(layouts_are) + return layouts(NHWGC, GKYXC, NHWGK); + if constexpr(layouts_are) + return layouts(NHWGC, GKYXC, NHWGK); + if constexpr(layouts_are) + return layouts(NGCHW, GKYXC, NGKHW); + if constexpr(layouts_are) + return layouts(NGCHW, GKCYX, NGKHW); + break; + case 3: + if constexpr(layouts_are) + return layouts(GNDHWC, GKZYXC, GNDHWK); + if constexpr(layouts_are) + return layouts(GNDHWC, GKZYXC, GNDHWK); + if constexpr(layouts_are) + return layouts(NDHWGC, GKZYXC, NDHWGK); + if constexpr(layouts_are) + return layouts(NGCDHW, GKZYXC, NGKDHW); + if constexpr(layouts_are) + return layouts(NGCDHW, GKCZYX, NGKDHW); + break; + } - // Helper to construct layout array - constexpr auto make_layouts = [](auto in, auto weight, auto out) { - return std::array{in, weight, out}; - }; + // If we reach here, the layout combination is not supported + // Call consteval function to trigger a compile-time error with a clear message + report_unsupported_layout_error(); - constexpr int spatial_dim = InstTraits::kSpatialDim; + // This return is unreachable but needed to satisfy the compiler + return layouts(GNHWC, GKYXC, GNHWK); +} - if constexpr(spatial_dim == 1) - { - if constexpr(layouts_match.template operator()()) - return make_layouts(GNWC, GKXC, GNWK); - else if constexpr(layouts_match - .template operator()()) - return make_layouts(GNWC, GKXC, GNWK); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NWGC, GKXC, NWGK); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NGCW, GKXC, NGKW); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NGCW, GKCX, NGKW); - else - { - report_unsupported_layout_error(); - return make_layouts(GNWC, GKXC, GNWK); // Unreachable - } - } - else if constexpr(spatial_dim == 2) - { - if constexpr(layouts_match.template operator()()) - return make_layouts(GNHWC, GKYXC, GNHWK); - else if constexpr(layouts_match - .template operator()()) - return make_layouts(GNHWC, GKYXC, GNHWK); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NHWGC, GKYXC, NHWGK); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NHWGC, GKYXC, NHWGK); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NGCHW, GKYXC, NGKHW); - else if constexpr(layouts_match.template operator()()) - return make_layouts(NGCHW, GKCYX, NGKHW); - else - { - report_unsupported_layout_error(); - return make_layouts(GNHWC, GKYXC, GNHWK); // Unreachable - } - } - else if constexpr(spatial_dim == 3) - { - if constexpr(layouts_match.template operator()()) - return make_layouts(GNDHWC, GKZYXC, GNDHWK); - else if constexpr(layouts_match - .template operator()()) - return make_layouts(GNDHWC, GKZYXC, GNDHWK); - else if constexpr(layouts_match - .template operator()()) - return make_layouts(NDHWGC, GKZYXC, NDHWGK); - else if constexpr(layouts_match - .template operator()()) - return make_layouts(NGCDHW, GKZYXC, NGKDHW); - else if constexpr(layouts_match - .template operator()()) - return make_layouts(NGCDHW, GKCZYX, NGKDHW); - else - { - report_unsupported_layout_error(); - return make_layouts(GNDHWC, GKZYXC, GNDHWK); // Unreachable - } - } - else - { - report_unsupported_layout_error(); - return make_layouts(GNHWC, GKYXC, GNHWK); // Unreachable - } +/// @brief Derives the grouped convolution layout from a device kernel `Instance` type. +/// @tparam Instance The device kernel instance type. +/// @return An std::array corresponding to the tensor layouts: +/// index 0 -> Input layout +/// index 1 -> Weight layout +/// index 2 -> Output layout + +template +constexpr auto fwd_conv_layout() + requires HasFwdConvLayouts> +{ + + using A = typename InstanceTraits::ALayout; + using B = typename InstanceTraits::BLayout; + using E = typename InstanceTraits::ELayout; + return conv_layout::kSpatialDim>(); +} + +/// @brief Derives the grouped convolution layout from a device kernel `Instance` type. +/// @tparam Instance The device kernel instance type. +/// @return An std::array corresponding to the tensor layouts: +/// index 0 -> Input layout +/// index 1 -> Weight layout +/// index 2 -> Output layout +template +constexpr auto bwd_wei_conv_layout() + requires HasBwdWeiLayouts> +{ + + using A = typename InstanceTraits::InLayout; + using B = typename InstanceTraits::WeiLayout; + using E = typename InstanceTraits::OutLayout; + return conv_layout::kSpatialDim>(); } // ---------------------------------------------------------------------------- @@ -447,13 +464,11 @@ constexpr std::array conv_layout() // ---------------------------------------------------------------------------- /// @brief Helper function to report unsupported data type with a clear error message. -/// @details This consteval function uses throw (not static_assert) to ensure the error is not -/// silently ignored during SFINAE. The thrown string becomes part of the compiler error message. -template +template [[noreturn]] consteval void report_unsupported_data_type_error() { throw "Unsupported data type detected!\n" - "The ADataType is not recognized.\n" + "The DataTypeFromInstance is not recognized.\n" "Supported types are: ck::half_t (FP16), ck::Tuple (FP16_FP16), " "ck::bhalf_t (BF16), ck::Tuple (BF16_BF16), float (FP32), " "ck::Tuple (FP32_FP32), double (FP64), ck::f8_t (FP8), ck::bf8_fnuz_t " @@ -462,62 +477,44 @@ template "Please verify that your kernel instance uses a supported data type."; } -/// @brief Derives the data type from a device kernel Instance type. -/// @tparam Instance The device kernel instance type. -/// @return A builder::DataType enum value representing the input data type. -/// @details This function examines the Instance's ADataType to determine the data type -/// used for the input tensor. The function supports various floating-point and integer -/// types, including tuple types for mixed-precision operations. -/// -/// Supported data types include: -/// - FP16 (ck::half_t) -/// - FP16_FP16 (ck::Tuple) -/// - BF16 (ck::bhalf_t) -/// - BF16_BF16 (ck::Tuple) -/// - FP32 (float) -/// - FP32_FP32 (ck::Tuple) -/// - FP64 (double) -/// - FP8 (ck::f8_t) -/// - BF8 (ck::bf8_fnuz_t, ck::bf8_ocp_t) -/// - I8 (int8_t) -/// - I8_I8 (ck::Tuple) -/// - U8 (uint8_t) -template +/// @brief Derives the data type from a device kernel `Instance` type. +/// Returns a `builder::DataType` enum value (e.g., FP16, BF16, FP32, BF8). +// Note: maybe move to types.hpp? +template constexpr builder::DataType conv_data_type() + { - using InstTraits = InstanceTraits; - using ADataType = typename InstTraits::ADataType; using enum builder::DataType; - if constexpr(std::is_same_v) + if constexpr(std::is_same_v) return FP16; - else if constexpr(std::is_same_v>) + else if constexpr(std::is_same_v>) return FP16_FP16; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return BF16; - else if constexpr(std::is_same_v>) + else if constexpr(std::is_same_v>) return BF16_BF16; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return FP32; - else if constexpr(std::is_same_v>) + else if constexpr(std::is_same_v>) return FP32_FP32; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return FP64; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return FP8; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return BF8; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return BF8; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return I8; - else if constexpr(std::is_same_v>) + else if constexpr(std::is_same_v>) return I8_I8; - else if constexpr(std::is_same_v) + else if constexpr(std::is_same_v) return U8; else { - report_unsupported_data_type_error(); + report_unsupported_data_type_error(); return FP32; // Unreachable } } @@ -736,4 +733,92 @@ constexpr builder::PipelineScheduler get_pipeline_scheduler() } } +// ============================================================================ +// SECTION 4: Helper functions for common structures often used in reflection +// ============================================================================ + +template +constexpr DataTileInfo conv_traits_data_tile(int k_or_k0 = InstTraits::kKPerBlock) +{ + return DataTileInfo{.m = InstTraits::kMPerBlock, .n = InstTraits::kNPerBlock, .k = k_or_k0}; +} + +template +constexpr InputTileTransferInfo +conv_traits_a_transfer_params(int _k1, int kPerBlock = InstTraits::kKPerBlock) +{ + return InputTileTransferInfo{ + .tile_dimensions = {.k0 = kPerBlock / _k1, .m_or_n = InstTraits::kMPerBlock, .k1 = _k1}, + .transfer_params = {.k1 = _k1, + .thread_cluster_dims = InstTraits::kAThreadClusterLengths, + .thread_cluster_order = InstTraits::kAThreadClusterArrangeOrder, + .src_access_order = InstTraits::kABlockTransferSrcAccessOrder, + .src_vector_dim = InstTraits::kABlockTransferSrcVectorDim, + .src_scalar_per_vector = InstTraits::kABlockTransferSrcScalarPerVector, + .dst_scalar_per_vector_k1 = + InstTraits::kABlockTransferDstScalarPerVectorK1, + .lds_padding = static_cast(InstTraits::kABlockLdsExtraM)}}; +} + +template +constexpr InputTileTransferInfo +conv_traits_b_transfer_params(int _k1, int kPerBlock = InstTraits::kKPerBlock) +{ + return InputTileTransferInfo{ + .tile_dimensions = {.k0 = kPerBlock / _k1, .m_or_n = InstTraits::kNPerBlock, .k1 = _k1}, + .transfer_params = {.k1 = _k1, + .thread_cluster_dims = InstTraits::kBThreadClusterLengths, + .thread_cluster_order = InstTraits::kBThreadClusterArrangeOrder, + .src_access_order = InstTraits::kBBlockTransferSrcAccessOrder, + .src_vector_dim = InstTraits::kBBlockTransferSrcVectorDim, + .src_scalar_per_vector = InstTraits::kBBlockTransferSrcScalarPerVector, + .dst_scalar_per_vector_k1 = + InstTraits::kBBlockTransferDstScalarPerVectorK1, + .lds_padding = static_cast(InstTraits::kBBlockLdsExtraN)}}; +} + +template +constexpr WarpGemmParams conv_traits_wmma_warp_gemm_params() +{ + return WarpGemmParams{.gemm_m = InstTraits::kMPerWmma, + .gemm_n = InstTraits::kNPerWmma, + .m_iter = InstTraits::kMRepeat, + .n_iter = InstTraits::kNRepeat}; +} + +template +constexpr WarpGemmParams conv_traits_xdl_warp_gemm_params() +{ + return WarpGemmParams{.gemm_m = InstTraits::kMPerXDL, + .gemm_n = InstTraits::kNPerXDL, + .m_iter = InstTraits::kMXdlPerWave, + .n_iter = InstTraits::kNXdlPerWave}; +} + +template +constexpr OutputTileTransferInfo conv_traits_wmma_c_tile_transfer() +{ + return OutputTileTransferInfo{ + .shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMRepeatPerShuffle, + .n_gemms_per_shuffle = InstTraits::kCShuffleNRepeatPerShuffle}, + .thread_cluster_dims = {InstTraits::kCDEThreadClusterLengths[0], + InstTraits::kCDEThreadClusterLengths[1], + InstTraits::kCDEThreadClusterLengths[2], + InstTraits::kCDEThreadClusterLengths[3]}, + .scalar_per_vector = InstTraits::kCDEBlockTransferScalarPerVector}; +} + +template +constexpr OutputTileTransferInfo conv_traits_xdl_c_tile_transfer() +{ + return OutputTileTransferInfo{ + .shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle, + .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle}, + .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0], + InstTraits::kCThreadClusterLengths[1], + InstTraits::kCThreadClusterLengths[2], + InstTraits::kCThreadClusterLengths[3]}, + .scalar_per_vector = InstTraits::kCBlockTransferScalarPerVector}; +} + } // namespace ck_tile::reflect::conv diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_to_conv_traits.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_to_conv_traits.hpp index 00010e2d48..e10baaf712 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_to_conv_traits.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_to_conv_traits.hpp @@ -3,6 +3,18 @@ #pragma once +// Fwd instances #include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp" #include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp" #include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_xdl_large_tensor_cshuffle.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp" + +// Bwd weight instances +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp" +#include "ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp" diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp index cde1896993..c3a5f9df29 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp @@ -62,6 +62,10 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvFwdMultipleD_Wmma_CShuffle_v3 device kernel +struct DeviceGroupedConvBwdWeight_multiple_d_Wmma_CShuffle_V3_Tag +{ +}; template ::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_AK1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_AK1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_BK0_N_BK1 = BBlockTransferThreadClusterLengths_BK0_N_BK1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_BK1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_BK1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; using CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; - + static constexpr auto kCDEThreadClusterLengths = detail::SequenceToArray< + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; + static constexpr int kCDEBlockTransferScalarPerVector = + CShuffleBlockTransferScalarPerVector_NPerBlock; static constexpr ck::BlockGemmPipelineScheduler kBlkGemmPipeSched = BlkGemmPipeSched; static constexpr ck::BlockGemmPipelineVersion kBlkGemmPipelineVer = BlkGemmPipelineVer; @@ -231,7 +256,7 @@ struct InstanceTraits< oss << "DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3"; // Template parameters in exact order - oss << "<" << kNDimSpatial; // 1. NDimSpatial + oss << "<" << kSpatialDim; // 1. NDimSpatial oss << "," << detail::layout_name(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -251,30 +276,30 @@ struct InstanceTraits< // OutElementwiseOperation oss << "," << detail::conv_bwd_weight_spec_name( - kConvBackwardWeightSpecialization); // 14. ConvBackwardWeightSpecialization - oss << "," << kBlockSize; // 15. BlockSize - oss << "," << kMPerBlock; // 16. MPerBlock - oss << "," << kNPerBlock; // 17. NPerBlock - oss << "," << kKPerBlock; // 18. KPerBlock - oss << "," << kABK1; // 19. ABK1 - oss << "," << kMPerWmma; // 20. MPerWmma - oss << "," << kNPerWmma; // 21. NPerWmma - oss << "," << kMRepeat; // 22. MRepeat - oss << "," << kNRepeat; // 23. NRepeat + kConvBwdWeightSpecialization); // 14. ConvBackwardWeightSpecialization + oss << "," << kBlockSize; // 15. BlockSize + oss << "," << kMPerBlock; // 16. MPerBlock + oss << "," << kNPerBlock; // 17. NPerBlock + oss << "," << kKPerBlock; // 18. KPerBlock + oss << "," << kK1; // 19. ABK1 + oss << "," << kMPerWmma; // 20. MPerWmma + oss << "," << kNPerWmma; // 21. NPerWmma + oss << "," << kMRepeat; // 22. MRepeat + oss << "," << kNRepeat; // 23. NRepeat oss << "," << detail::sequence_name(); // 24. oss << "," << detail::sequence_name(); // 25. oss << "," << detail::sequence_name(); // 26. oss << "," << kABlockTransferSrcVectorDim; // 27. oss << "," << kABlockTransferSrcScalarPerVector; // 28. - oss << "," << kABlockTransferDstScalarPerVector_AK1; // 29. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 30. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 29. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << detail::sequence_name(); // 32. oss << "," << detail::sequence_name(); // 33. oss << "," << kBBlockTransferSrcVectorDim; // 34. oss << "," << kBBlockTransferSrcScalarPerVector; // 35. - oss << "," << kBBlockTransferDstScalarPerVector_BK1; // 36. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 37. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 36. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 37. oss << "," << kCShuffleMRepeatPerShuffle; // 38. oss << "," << kCShuffleNRepeatPerShuffle; // 39. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp index 6508ac7d6e..173da8268a 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp @@ -59,6 +59,10 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle device kernel +struct DeviceGroupedConvBwdWeight_multiple_d_Xdl_CShuffle_Tag +{ +}; template ::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_K1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_K0_N_K1 = BBlockTransferThreadClusterLengths_K0_N_K1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_K1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; using CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; @@ -211,6 +232,9 @@ struct InstanceTraits< using ComputeTypeA = ComputeTypeA_; using ComputeTypeB = ComputeTypeB_; + static constexpr auto kCThreadClusterLengths = detail::SequenceToArray< + CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; + // Static member function to generate instance string static std::string instance_string() { @@ -220,7 +244,7 @@ struct InstanceTraits< oss << "DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle"; // Template parameters in exact order - oss << "<" << kNDimSpatial; // 1. NDimSpatial + oss << "<" << kSpatialDim; // 1. NDimSpatial oss << "," << detail::layout_name(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -240,30 +264,30 @@ struct InstanceTraits< // OutElementwiseOperation oss << "," << detail::conv_bwd_weight_spec_name( - kConvBackwardWeightSpecialization); // 14. ConvBackwardWeightSpecialization - oss << "," << kBlockSize; // 15. BlockSize - oss << "," << kMPerBlock; // 16. MPerBlock - oss << "," << kNPerBlock; // 17. NPerBlock - oss << "," << kK0PerBlock; // 18. K0PerBlock - oss << "," << kK1; // 19. K1 - oss << "," << kMPerXDL; // 20. MPerXDL - oss << "," << kNPerXDL; // 21. NPerXDL - oss << "," << kMXdlPerWave; // 22. MXdlPerWave - oss << "," << kNXdlPerWave; // 23. NXdlPerWave + kConvBwdWeightSpecialization); // 14. ConvBackwardWeightSpecialization + oss << "," << kBlockSize; // 15. BlockSize + oss << "," << kMPerBlock; // 16. MPerBlock + oss << "," << kNPerBlock; // 17. NPerBlock + oss << "," << kK0PerBlock; // 18. K0PerBlock + oss << "," << kK1; // 19. K1 + oss << "," << kMPerXDL; // 20. MPerXDL + oss << "," << kNPerXDL; // 21. NPerXDL + oss << "," << kMXdlPerWave; // 22. MXdlPerWave + oss << "," << kNXdlPerWave; // 23. NXdlPerWave oss << "," << detail::sequence_name(); // 24. oss << "," << detail::sequence_name(); // 25. oss << "," << detail::sequence_name(); // 26. oss << "," << kABlockTransferSrcVectorDim; // 27. oss << "," << kABlockTransferSrcScalarPerVector; // 28. - oss << "," << kABlockTransferDstScalarPerVector_K1; // 29. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 30. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 29. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << detail::sequence_name(); // 32. oss << "," << detail::sequence_name(); // 33. oss << "," << kBBlockTransferSrcVectorDim; // 34. oss << "," << kBBlockTransferSrcScalarPerVector; // 35. - oss << "," << kBBlockTransferDstScalarPerVector_K1; // 36. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 37. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 36. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 37. oss << "," << kCShuffleMXdlPerWavePerShuffle; // 38. oss << "," << kCShuffleNXdlPerWavePerShuffle; // 39. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp index f1e40de7d2..4b90a6ab64 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp @@ -63,6 +63,11 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvBwdWeight_two_stage_Wmma_CShuffle_Tag device kernel +struct DeviceGroupedConvBwdWeight_two_stage_Wmma_CShuffle_Tag +{ +}; + template ::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_AK1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_AK1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_BK0_N_BK1 = BBlockTransferThreadClusterLengths_BK0_N_BK1_; @@ -215,13 +231,26 @@ struct InstanceTraits< static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_BK1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_BK1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; using CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; + static constexpr auto kCDEThreadClusterLengths = detail::SequenceToArray< + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; + static constexpr int kCDEBlockTransferScalarPerVector = + CShuffleBlockTransferScalarPerVector_NPerBlock; + static constexpr ck::BlockGemmPipelineScheduler kBlkGemmPipeSched = BlkGemmPipeSched; static constexpr ck::BlockGemmPipelineVersion kBlkGemmPipelineVer = BlkGemmPipelineVer; @@ -237,7 +266,7 @@ struct InstanceTraits< oss << "DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3"; // Template parameters in exact order - oss << "<" << kNDimSpatial; // 1. NDimSpatial + oss << "<" << kSpatialDim; // 1. NDimSpatial oss << "," << detail::layout_name(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -255,30 +284,30 @@ struct InstanceTraits< // OutElementwiseOperation oss << "," << detail::conv_bwd_weight_spec_name( - kConvBackwardWeightSpecialization); // 12. ConvBackwardWeightSpecialization - oss << "," << kBlockSize; // 13. BlockSize - oss << "," << kMPerBlock; // 14. MPerBlock - oss << "," << kNPerBlock; // 15. NPerBlock - oss << "," << kKPerBlock; // 16. KPerBlock - oss << "," << kABK1; // 17. ABK1 - oss << "," << kMPerWmma; // 18. MPerWmma - oss << "," << kNPerWmma; // 19. NPerWmma - oss << "," << kMRepeat; // 20. MRepeat - oss << "," << kNRepeat; // 21. NRepeat + kConvBwdWeightSpecialization); // 12. ConvBackwardWeightSpecialization + oss << "," << kBlockSize; // 13. BlockSize + oss << "," << kMPerBlock; // 14. MPerBlock + oss << "," << kNPerBlock; // 15. NPerBlock + oss << "," << kKPerBlock; // 16. KPerBlock + oss << "," << kABK1; // 17. ABK1 + oss << "," << kMPerWmma; // 18. MPerWmma + oss << "," << kNPerWmma; // 19. NPerWmma + oss << "," << kMRepeat; // 20. MRepeat + oss << "," << kNRepeat; // 21. NRepeat oss << "," << detail::sequence_name(); // 22. oss << "," << detail::sequence_name(); // 23. oss << "," << detail::sequence_name(); // 24. oss << "," << kABlockTransferSrcVectorDim; // 25. oss << "," << kABlockTransferSrcScalarPerVector; // 26. - oss << "," << kABlockTransferDstScalarPerVector_AK1; // 27. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 28. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 27. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 28. oss << "," << detail::sequence_name(); // 29. oss << "," << detail::sequence_name(); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << kBBlockTransferSrcVectorDim; // 32. oss << "," << kBBlockTransferSrcScalarPerVector; // 33. - oss << "," << kBBlockTransferDstScalarPerVector_BK1; // 34. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 35. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 34. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 35. oss << "," << kCShuffleMRepeatPerShuffle; // 36. oss << "," << kCShuffleNRepeatPerShuffle; // 37. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp index 460b49de93..999aff6f1e 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp @@ -63,6 +63,11 @@ struct DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle device kernel +struct DeviceGroupedConvBwdWeight_two_stage_Xdl_CShuffle_Tag +{ +}; + template ::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_K1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_K0_N_K1 = BBlockTransferThreadClusterLengths_K0_N_K1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_K1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; using CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; + static constexpr auto kCThreadClusterLengths = detail::SequenceToArray< + CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; static constexpr ck::BlockGemmPipelineScheduler kBlkGemmPipeSched = BlkGemmPipeSched; static constexpr ck::BlockGemmPipelineVersion kBlkGemmPipelineVer = BlkGemmPipelineVer; @@ -234,7 +260,7 @@ struct InstanceTraits(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -252,30 +278,30 @@ struct InstanceTraits(); // 22. oss << "," << detail::sequence_name(); // 23. oss << "," << detail::sequence_name(); // 24. oss << "," << kABlockTransferSrcVectorDim; // 25. oss << "," << kABlockTransferSrcScalarPerVector; // 26. - oss << "," << kABlockTransferDstScalarPerVector_K1; // 27. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 28. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 27. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 28. oss << "," << detail::sequence_name(); // 29. oss << "," << detail::sequence_name(); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << kBBlockTransferSrcVectorDim; // 32. oss << "," << kBBlockTransferSrcScalarPerVector; // 33. - oss << "," << kBBlockTransferDstScalarPerVector_K1; // 34. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 35. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 34. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 35. oss << "," << kCShuffleMXdlPerWavePerShuffle; // 36. oss << "," << kCShuffleNXdlPerWavePerShuffle; // 37. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp index f87e295159..eba422b85f 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle.hpp @@ -59,6 +59,11 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffle; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvFwdMultipleD_Wmma_CShuffle_v3 device kernel +struct DeviceGroupedConvBwdWeight_Wmma_CShuffle_Tag +{ +}; + template > // Use false to match with the default value { static constexpr auto kTensorOpName = "DeviceGroupedConvBwdWeight_Wmma_CShuffle"; + using device_kernel_tag = DeviceGroupedConvBwdWeight_Wmma_CShuffle_Tag; - static constexpr ck::index_t kNDimSpatial = NDimSpatial; + static constexpr ck::index_t kSpatialDim = NDimSpatial; using InLayout = InLayout_; using WeiLayout = WeiLayout_; @@ -164,15 +170,15 @@ struct InstanceTraits::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_K1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_K0_N_K1 = BBlockTransferThreadClusterLengths_K0_N_K1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_K1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; using CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; - + static constexpr auto kCDEThreadClusterLengths = detail::SequenceToArray< + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; + static constexpr int kCDEBlockTransferScalarPerVector = + CShuffleBlockTransferScalarPerVector_NPerBlock; static constexpr ck::LoopScheduler kLoopSched = LoopSched; static constexpr ck::PipelineVersion kPipelineVer = PipelineVer; @@ -216,7 +239,7 @@ struct InstanceTraits(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -234,30 +257,30 @@ struct InstanceTraits(); // 22. oss << "," << detail::sequence_name(); // 23. oss << "," << detail::sequence_name(); // 24. oss << "," << kABlockTransferSrcVectorDim; // 25. oss << "," << kABlockTransferSrcScalarPerVector; // 26. - oss << "," << kABlockTransferDstScalarPerVector_K1; // 27. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 28. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 27. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 28. oss << "," << detail::sequence_name(); // 29. oss << "," << detail::sequence_name(); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << kBBlockTransferSrcVectorDim; // 32. oss << "," << kBBlockTransferSrcScalarPerVector; // 33. - oss << "," << kBBlockTransferDstScalarPerVector_K1; // 34. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 35. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 34. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 35. oss << "," << kCShuffleMRepeatPerShuffle; // 36. oss << "," << kCShuffleNRepeatPerShuffle; // 37. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp index 29459d67b0..cfc8b4e05a 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp @@ -62,6 +62,11 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvFwdMultipleD_Wmma_CShuffle_v3 device kernel +struct DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3_Tag +{ +}; + template > { static constexpr auto kTensorOpName = "DeviceGroupedConvBwdWeight_Wmma_CShuffleV3"; + using device_kernel_tag = DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3_Tag; - static constexpr ck::index_t kNDimSpatial = NDimSpatial; + static constexpr ck::index_t kSpatialDim = NDimSpatial; using InLayout = InLayout_; using WeiLayout = WeiLayout_; @@ -172,13 +178,13 @@ struct InstanceTraits::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_AK1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_AK1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_BK0_N_BK1 = BBlockTransferThreadClusterLengths_BK0_N_BK1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_BK1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_BK1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; using CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; - + static constexpr auto kCDEThreadClusterLengths = detail::SequenceToArray< + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; + static constexpr int kCDEBlockTransferScalarPerVector = + CShuffleBlockTransferScalarPerVector_NPerBlock; static constexpr ck::BlockGemmPipelineScheduler kBlkGemmPipeSched = BlkGemmPipeSched; static constexpr ck::BlockGemmPipelineVersion kBlkGemmPipelineVer = BlkGemmPipelineVer; @@ -232,7 +257,7 @@ struct InstanceTraits(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -250,30 +275,30 @@ struct InstanceTraits(); // 22. oss << "," << detail::sequence_name(); // 23. oss << "," << detail::sequence_name(); // 24. oss << "," << kABlockTransferSrcVectorDim; // 25. oss << "," << kABlockTransferSrcScalarPerVector; // 26. - oss << "," << kABlockTransferDstScalarPerVector_AK1; // 27. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 28. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 27. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 28. oss << "," << detail::sequence_name(); // 29. oss << "," << detail::sequence_name(); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << kBBlockTransferSrcVectorDim; // 32. oss << "," << kBBlockTransferSrcScalarPerVector; // 33. - oss << "," << kBBlockTransferDstScalarPerVector_BK1; // 34. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 35. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 34. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 35. oss << "," << kCShuffleMRepeatPerShuffle; // 36. oss << "," << kCShuffleNRepeatPerShuffle; // 37. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp index 2c893b9c1d..1edf03740f 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle.hpp @@ -61,6 +61,11 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvBwdWeight_Xdl_CShuffle device kernel +struct DeviceGroupedConvBwdWeight_Xdl_CShuffle_Tag +{ +}; + template ::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_K1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_K0_N_K1 = BBlockTransferThreadClusterLengths_K0_N_K1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_K1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; static constexpr ck::index_t kCShuffleMXdlPerWavePerShuffle = CShuffleMXdlPerWavePerShuffle; static constexpr ck::index_t kCShuffleNXdlPerWavePerShuffle = CShuffleNXdlPerWavePerShuffle; using CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; + static constexpr auto kCThreadClusterLengths = detail::SequenceToArray< + CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; + static constexpr ck::index_t kCBlockTransferScalarPerVector_NWaveNPerXdl = CBlockTransferScalarPerVector_NWaveNPerXdl; @@ -224,7 +250,7 @@ struct InstanceTraits(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -242,30 +268,30 @@ struct InstanceTraits(); // 22. oss << "," << detail::sequence_name(); // 23. oss << "," << detail::sequence_name(); // 24. oss << "," << kABlockTransferSrcVectorDim; // 25. oss << "," << kABlockTransferSrcScalarPerVector; // 26. - oss << "," << kABlockTransferDstScalarPerVector_K1; // 27. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 28. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 27. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 28. oss << "," << detail::sequence_name(); // 29. oss << "," << detail::sequence_name(); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << kBBlockTransferSrcVectorDim; // 32. oss << "," << kBBlockTransferSrcScalarPerVector; // 33. - oss << "," << kBBlockTransferDstScalarPerVector_K1; // 34. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 35. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 34. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 35. oss << "," << kCShuffleMXdlPerWavePerShuffle; // 36. oss << "," << kCShuffleNXdlPerWavePerShuffle; // 37. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp index 147028f9cf..ce23dac1d7 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -61,6 +61,11 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3; namespace ck_tile { namespace reflect { +/// @brief Tag type for DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3 device kernel +struct DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3_Tag +{ +}; + template > { + + /// @brief Tag type identifying this device kernel variant + using device_kernel_tag = DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3_Tag; static constexpr auto kTensorOpName = "DeviceGroupedConvBwdWeight_Xdl_CShuffleV3"; - static constexpr ck::index_t kNDimSpatial = NDimSpatial; + static constexpr ck::index_t kSpatialDim = NDimSpatial; using InLayout = InLayout_; using WeiLayout = WeiLayout_; @@ -167,7 +175,7 @@ struct InstanceTraits::value; + static constexpr auto kAThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kABlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kABlockTransferSrcVectorDim = ABlockTransferSrcVectorDim; static constexpr ck::index_t kABlockTransferSrcScalarPerVector = ABlockTransferSrcScalarPerVector; - static constexpr ck::index_t kABlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kABlockTransferDstScalarPerVectorK1 = ABlockTransferDstScalarPerVector_K1; - static constexpr bool kABlockLdsAddExtraM = ABlockLdsAddExtraM; + static constexpr bool kABlockLdsExtraM = ABlockLdsAddExtraM; using BBlockTransferThreadClusterLengths_K0_N_K1 = BBlockTransferThreadClusterLengths_K0_N_K1_; using BBlockTransferThreadClusterArrangeOrder = BBlockTransferThreadClusterArrangeOrder_; using BBlockTransferSrcAccessOrder = BBlockTransferSrcAccessOrder_; + + // B block transfer thread cluster dimensions (converted to std::array) + static constexpr auto kBThreadClusterLengths = + detail::SequenceToArray::value; + static constexpr auto kBThreadClusterArrangeOrder = + detail::SequenceToArray::value; + static constexpr auto kBBlockTransferSrcAccessOrder = + detail::SequenceToArray::value; + static constexpr ck::index_t kBBlockTransferSrcVectorDim = BBlockTransferSrcVectorDim; static constexpr ck::index_t kBBlockTransferSrcScalarPerVector = BBlockTransferSrcScalarPerVector; - static constexpr ck::index_t kBBlockTransferDstScalarPerVector_K1 = + static constexpr ck::index_t kBBlockTransferDstScalarPerVectorK1 = BBlockTransferDstScalarPerVector_K1; - static constexpr bool kBBlockLdsAddExtraN = BBlockLdsAddExtraN; + static constexpr bool kBBlockLdsExtraN = BBlockLdsAddExtraN; static constexpr ck::index_t kCShuffleMXdlPerWavePerShuffle = CShuffleMXdlPerWavePerShuffle; static constexpr ck::index_t kCShuffleNXdlPerWavePerShuffle = CShuffleNXdlPerWavePerShuffle; using CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock = CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_; + static constexpr auto kCThreadClusterLengths = detail::SequenceToArray< + CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock>::value; static constexpr ck::index_t kCBlockTransferScalarPerVector_NWaveNPerXdl = CBlockTransferScalarPerVector_NWaveNPerXdl; @@ -222,7 +250,7 @@ struct InstanceTraits(); // 2. InLayout oss << "," << detail::layout_name(); // 3. WeiLayout oss << "," << detail::layout_name(); // 4. OutLayout @@ -240,30 +268,30 @@ struct InstanceTraits(); // 22. oss << "," << detail::sequence_name(); // 23. oss << "," << detail::sequence_name(); // 24. oss << "," << kABlockTransferSrcVectorDim; // 25. oss << "," << kABlockTransferSrcScalarPerVector; // 26. - oss << "," << kABlockTransferDstScalarPerVector_K1; // 27. - oss << "," << (kABlockLdsAddExtraM ? "true" : "false"); // 28. + oss << "," << kABlockTransferDstScalarPerVectorK1; // 27. + oss << "," << (kABlockLdsExtraM ? "true" : "false"); // 28. oss << "," << detail::sequence_name(); // 29. oss << "," << detail::sequence_name(); // 30. oss << "," << detail::sequence_name(); // 31. oss << "," << kBBlockTransferSrcVectorDim; // 32. oss << "," << kBBlockTransferSrcScalarPerVector; // 33. - oss << "," << kBBlockTransferDstScalarPerVector_K1; // 34. - oss << "," << (kBBlockLdsAddExtraN ? "true" : "false"); // 35. + oss << "," << kBBlockTransferDstScalarPerVectorK1; // 34. + oss << "," << (kBBlockLdsExtraN ? "true" : "false"); // 35. oss << "," << kCShuffleMXdlPerWavePerShuffle; // 36. oss << "," << kCShuffleNXdlPerWavePerShuffle; // 37. oss << "," diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp index 782fd158c5..645d75258e 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp @@ -79,6 +79,11 @@ struct DeviceGroupedConvFwdMultipleD_Wmma_CShuffle; } // namespace ck::tensor_operation::device +/// @brief Tag type for DeviceGroupedConvFwdMultipleD_Wmma_CShuffle device kernel +struct DeviceGroupedConvFwdMultipleD_Wmma_CShuffle_Tag +{ +}; + namespace ck_tile::reflect { // Specialization for DeviceGroupedConvFwdMultipleD_Wmma_CShuffle @@ -176,6 +181,8 @@ struct InstanceTraits> { + /// @brief Tag type identifying this device kernel variant + using device_kernel_tag = DeviceGroupedConvFwdMultipleD_Wmma_CShuffle_Tag; // Spatial dimension static constexpr int kSpatialDim = NDimSpatial; diff --git a/experimental/builder/test/conv/ck/test_conv_traits.cpp b/experimental/builder/test/conv/ck/test_conv_traits.cpp index 42235df2fe..3221113565 100644 --- a/experimental/builder/test/conv/ck/test_conv_traits.cpp +++ b/experimental/builder/test/conv/ck/test_conv_traits.cpp @@ -9,8 +9,17 @@ #include #include #include +#include #include +#include +#include +#include +#include +#include +#include +#include + namespace { using ck_tile::builder::ConvDirection; @@ -26,6 +35,1099 @@ class ConvTraitsTest : public ::testing::Test { }; +// Test ConvTraits with DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3 +TEST_F(ConvTraitsTest, ConvBwdWeightCshuffleWmmaTraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Wmma_CShuffle< + 3, // NDimSpatial + ck::tensor_layout::convolution::GNDHWC, // InLayout + ck::tensor_layout::convolution::GKZYXC, // WeiLayout + ck::tensor_layout::convolution::GNDHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerWmma + 32, // NPerWmma + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + 1, // NummGemmKPrefetchStage + ck::LoopScheduler::Default, // BlkGemmPipeSched + ck::PipelineVersion::v1, // BlkGemmPipelineVer + false>; // BComputeDataType + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 3); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNDHWC, TensorLayout::GKZYXC, TensorLayout::GNDHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration +} + +// Test ConvTraits with DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3 +TEST_F(ConvTraitsTest, ConvBwdWeightCshuffleWmmaV3TraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerWmma + 32, // NPerWmma + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched + ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer + ck::half_t, // AComputeDataType + ck::half_t, // BComputeDataType + 1, // MaxTransposeTransferSrcScalarPerVector + 1>; // MaxTransposeTransferDstScalarPerVector> + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration +} + +// Test ConvTraits with DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_V3 +TEST_F(ConvTraitsTest, ConvBwdWeightMultipleDCshuffleWmmaV3TraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = + ck::tensor_operation::device::DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::Tuple<>, // DsLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::Tuple<>, // DsDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerWmma + 32, // NPerWmma + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched + ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer + ck::half_t, // AComputeDataType + ck::half_t>; // BComputeDataType + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration +} + +// Test ConvTraits with DeviceGroupedConvBwdWeight_two_stage_Xdl_CShuffleV3 +TEST_F(ConvTraitsTest, ConvBwdWeightTwoStageWmmaCshuffleTraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = + ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // AK1 + 32, // MPerWMMA + 32, // NPerXDL + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched + ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer + 4, // NumGroupsToMerge + ck::half_t, // AComputeDataType + ck::half_t, // BComputeDataType + 1, // MaxTransposeTransferSrcScalarPerVector + 1>; // MaxTransposeTransferDstScalarPerVector> + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + EXPECT_EQ(traits.pipeline_version, PipelineVersion::V1); +} + +// Test ConvTraits with DeviceGroupedConvBwdWeight_two_stage_Xdl_CShuffleV3 +TEST_F(ConvTraitsTest, ConvBwdWeightTwoStageXdlCshuffleTraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = + ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerXDL + 32, // NPerXDL + 4, // MXdlPerWave + 4, // NXdlPerWave + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched + ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer + 4, // NumGroupsToMerge + ck::half_t, // AComputeDataType + ck::half_t, // BComputeDataType + 1, // MaxTransposeTransferSrcScalarPerVector + 1>; // MaxTransposeTransferDstScalarPerVector> + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + EXPECT_EQ(traits.pipeline_version, PipelineVersion::V1); +} + +// Test ConvTraits with DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle +TEST_F(ConvTraitsTest, ConvBwdWeightMultipleDCshuffleXDLTraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = + ck::tensor_operation::device::DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::Tuple<>, // DsLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::Tuple<>, // DsDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerXDL + 32, // NPerXDL + 4, // MXdlPerWave + 4, // NXdlPerWave + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::half_t, // AComputeDataType + ck::half_t>; // BComputeDataType + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration +} + +// Test ConvTraits with DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 +TEST_F(ConvTraitsTest, ConvBwdWeightXdlCshuffleV3TraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerXDL + 32, // NPerXDL + 4, // MXdlPerWave + 4, // NXdlPerWave + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched + ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer + ck::half_t, // AComputeDataType + ck::half_t>; // BComputeDataType + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + EXPECT_EQ(traits.pipeline_version, PipelineVersion::V1); +} + +// Test ConvTraits with DeviceGroupedConvBwdWeight_Xdl_CShuffle +TEST_F(ConvTraitsTest, ConvBwdWeightXdlCshuffleTraitsExtraction) +{ + // Define a concrete instance type with specific template parameters + using DeviceInstance = ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Xdl_CShuffle< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerXDL + 32, // NPerXDL + 4, // MXdlPerWave + 4, // NXdlPerWave + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::half_t, // AComputeDataType + ck::half_t, // BComputeDataType + 1, // MaxTransposeTransferSrcScalarPerVector + 1>; // MaxTransposeTransferDstScalarPerVector> + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::BACKWARD_WEIGHT); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration + EXPECT_EQ(traits.pipeline_scheduler, PipelineScheduler::DEFAULT); + EXPECT_EQ(traits.pipeline_version, PipelineVersion::V1); +} + +// test conv traits device_grouped_conv_fwd_multiple_d_wmma_cshuffle.hpp +TEST_F(ConvTraitsTest, ConvFwdTraitsMultipleDCshuffleWmmaExtraction) +{ + using DeviceInstance = + ck::tensor_operation::device::DeviceGroupedConvFwdMultipleD_Wmma_CShuffle< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // ALayout + ck::tensor_layout::convolution::GKYXC, // BLayout + ck::Tuple<>, // DsLayout + ck::tensor_layout::convolution::GNHWK, // ELayout + ck::half_t, // ADataType + ck::half_t, // BDataType + float, // AccDataType + ck::half_t, // CShuffleDataType + ck::Tuple<>, // DsDataType + ck::half_t, // EDataType + ck::tensor_operation::element_wise::PassThrough, // AElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // BElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // CDEElementwiseOperation + ck::tensor_operation::device::ConvolutionForwardSpecialization:: + Default, // ConvForwardSpecialization + ck::tensor_operation::device::GemmSpecialization::Default, // GemmSpec + 1, // NummGemmKPrefetchStage + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // KPerBlock + 8, // K1 + 32, // MPerWmma + 32, // NPerWmma + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_AK1 + 1, // ABlockLdsExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_BK1 + 1, // BBlockLdsExtraN + 1, // CShuffleMRepeatPerShuffle + 1, // CShuffleNRepeatPerShuffle + ck::Sequence< + 1, + 32, + 1, + 8>, // CDEShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock + 8, // CDEShuffleBlockTransferScalarPerVector_NPerBlock + ck::LoopScheduler::Default, // BlkGemmPipeSched + ck::PipelineVersion::v1>; // BlkGemmPipelineVer + + // Use ConvTraitsTmpl to extract compile-time information + const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); + + // Verify signature information + EXPECT_EQ(traits.spatial_dim, 2); + EXPECT_EQ(traits.direction, ConvDirection::FORWARD); + EXPECT_THAT(traits.layout, + ElementsAre(TensorLayout::GNHWC, TensorLayout::GKYXC, TensorLayout::GNHWK)); + EXPECT_EQ(traits.data_type, DataType::FP16); + EXPECT_EQ(traits.input_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.weight_element_op, ElementwiseOperation::PASS_THROUGH); + EXPECT_EQ(traits.output_element_op, ElementwiseOperation::PASS_THROUGH); + + // Verify specializations + EXPECT_EQ(traits.gemm_padding, ck_tile::builder::GemmPadding::DEFAULT); + EXPECT_EQ(traits.conv_specialization, ck_tile::builder::ConvSpecialization::DEFAULT); + EXPECT_EQ(traits.num_gemm_k_prefetch_stage, 1); + + // Verify algorithm information + EXPECT_EQ(traits.thread_block_size, 256); + + // Verify tile dimensions + EXPECT_EQ(traits.tile_dims.m, 128); + EXPECT_EQ(traits.tile_dims.n, 128); + EXPECT_EQ(traits.tile_dims.k, 16); + + // Verify A tile transfer info + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.a_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.a_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.a_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.a_tile_transfer.transfer_params.lds_padding); + + // Verify B tile transfer info + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k0, 2); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.m_or_n, 128); + EXPECT_EQ(traits.b_tile_transfer.tile_dimensions.k1, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.k1, 8); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_dims, ElementsAre(4, 64, 1)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.thread_cluster_order, ElementsAre(1, 0, 2)); + EXPECT_THAT(traits.b_tile_transfer.transfer_params.src_access_order, ElementsAre(1, 0, 2)); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_vector_dim, 2); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.src_scalar_per_vector, 8); + EXPECT_EQ(traits.b_tile_transfer.transfer_params.dst_scalar_per_vector_k1, 8); + EXPECT_TRUE(traits.b_tile_transfer.transfer_params.lds_padding); + + // Verify warp GEMM params + EXPECT_EQ(traits.warp_gemm.gemm_m, 32); + EXPECT_EQ(traits.warp_gemm.gemm_n, 32); + EXPECT_EQ(traits.warp_gemm.m_iter, 4); + EXPECT_EQ(traits.warp_gemm.n_iter, 4); + + // Verify output tile transfer info + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.m_gemms_per_shuffle, 1); + EXPECT_EQ(traits.c_tile_transfer.shuffle_params.n_gemms_per_shuffle, 1); + EXPECT_THAT(traits.c_tile_transfer.thread_cluster_dims, ElementsAre(1, 32, 1, 8)); + EXPECT_EQ(traits.c_tile_transfer.scalar_per_vector, 8); + + // Verify pipeline configuration + EXPECT_EQ(traits.pipeline_scheduler, + ck_tile::reflect::conv::convert_pipeline_scheduler()); + EXPECT_EQ(traits.pipeline_version, PipelineVersion::V1); +} + // Test ConvTraits with DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3 TEST_F(ConvTraitsTest, ConvFwdTraitsExtraction) { diff --git a/experimental/builder/test/test_conv_description.cpp b/experimental/builder/test/test_conv_description.cpp index bcea406fa7..4cbde73bcb 100644 --- a/experimental/builder/test/test_conv_description.cpp +++ b/experimental/builder/test/test_conv_description.cpp @@ -259,9 +259,118 @@ TEST(ConvDescriptionTest, DefaultInstanceHasDetailedDescription) static constexpr const ConvSignature SIGNATURE; static constexpr const DefaultAlgorithm ALGORITHM; using Instance = ckb::ConvBuilder::Instance; + EXPECT_THAT( + ckr::describe().detailed(), + ckt::StringEqWithDiff( // + "2D Forward Convolution Kernel\n" + "├─ Signature\n" + "│ ├─ Tensor Type: FP16\n" + "│ ├─ Input Layout: GNHWC\n" + "│ ├─ Weight Layout: GKYXC\n" + "│ ├─ Output Layout: GNHWK\n" + "│ ├─ Input elementwise operation: PASS_THROUGH\n" + "│ ├─ Weights elementwise operation: PASS_THROUGH\n" + "│ └─ Output elementwise operation: PASS_THROUGH\n" + "└─ Algorithm\n" + " ├─ Thread block size: 256\n" + " ├─ Data tile size: 256×256×32\n" + " ├─ Gemm padding: DEFAULT\n" + " ├─ Convolution specialization: DEFAULT\n" + " ├─ Pipeline version: V4\n" + " ├─ Pipeline scheduler: INTRAWAVE\n" + " ├─ Warp Gemm parameters: \n" + " │ ├─ subtile size: 16×16\n" + " │ └─ Number of warp gemm iterations: 8×8\n" + " └─ Memory access:\n" + " ├─ A Tile transfer: \n" + " │ ├─ Tile dimensions: 4×256×8×\n" + " │ ├─ The innermost K subdimension size: 8\n" + " │ ├─ Spatial thread distribution over the data tile: 0×1×2\n" + " │ ├─ The order of accessing data tile axes: 0×1×2\n" + " │ ├─ Vectorized memory access axis index (with contiguous memory): 2\n" + " │ ├─ Vector access (GMEM read) instruction size: 2\n" + " │ ├─ Vector access (LDS write) instruction size: 2\n" + " │ └─ LDS data layout padding (to prevent bank conflicts): 2\n" + " ├─ B Tile transfer: \n" + " │ ├─ Tile dimensions: 4×256×8×\n" + " │ ├─ The innermost K subdimension size: 8\n" + " │ ├─ Spatial thread distribution over the data tile: 0×1×2\n" + " │ ├─ The order of accessing data tile axes: 0×1×2\n" + " │ ├─ Vectorized memory access axis index (with contiguous memory): 2\n" + " │ ├─ Vector access (GMEM read) instruction size: 2\n" + " │ ├─ Vector access (LDS write) instruction size: 2\n" + " │ └─ LDS data layout padding (to prevent bank conflicts): 2\n" + " └─ C Tile transfer: \n" + " ├─ Data shuffle (number of gemm instructions per iteration): 1×1\n" + " ├─ Spatial thread distribution used to store data: 1×32×1×8\n" + " ├─ Vector access (GMEM write) instruction size: 2\n" + " ├─ Struct does not contain optional num_gemm_k_prefetch_stage parameter\n" + " ├─ Struct does not contain optional max_transpose_transfer_src_scalar_per_vector " + "parameter\n" + " ├─ Struct does not contain optional max_transpose_dst_scalar_per_vector parameter\n" + " └─ Struct does not contain optional num_groups_to_merge parameter")); +} + +// Test printing of optional parameters num_groups_to_merge, +// nax_transose_transfer_src_scalar_per_vector and max_transpose_dst_scalar_per_vector +TEST(ConvDescriptionTest, BwdWeightTwoStageWmmaV3DescriptionTest) +{ + using Instance = + ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3< + 2, // NDimSpatial + ck::tensor_layout::convolution::GNHWC, // InLayout + ck::tensor_layout::convolution::GKYXC, // WeiLayout + ck::tensor_layout::convolution::GNHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // AK1 + 32, // MPerWMMA + 32, // NPerXDL + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_BK0_N_BK1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched + ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer + 4, // NumGroupsToMerge + ck::half_t, // AComputeDataType + ck::half_t, // BComputeDataType + 1, // MaxTransposeTransferSrcScalarPerVector + 1>; // MaxTransposeTransferDstScalarPerVector> + EXPECT_THAT(ckr::describe().detailed(), ckt::StringEqWithDiff( // - "2D Forward Convolution Kernel\n" + "2D Backward Weight Convolution Kernel\n" "├─ Signature\n" "│ ├─ Tensor Type: FP16\n" "│ ├─ Input Layout: GNHWC\n" @@ -272,37 +381,146 @@ TEST(ConvDescriptionTest, DefaultInstanceHasDetailedDescription) "│ └─ Output elementwise operation: PASS_THROUGH\n" "└─ Algorithm\n" " ├─ Thread block size: 256\n" - " ├─ Data tile size: 256×256×32\n" - " ├─ Gemm padding: DEFAULT\n" + " ├─ Data tile size: 128×128×16\n" + " ├─ Struct does not contain optional gemm_padding argument\n" " ├─ Convolution specialization: DEFAULT\n" - " ├─ Pipeline version: V4\n" - " ├─ Pipeline scheduler: INTRAWAVE\n" + " ├─ Pipeline version: V1\n" + " ├─ Pipeline scheduler: DEFAULT\n" " ├─ Warp Gemm parameters: \n" - " │ ├─ subtile size: 16×16\n" - " │ └─ Number of warp gemm iterations: 8×8\n" + " │ ├─ subtile size: 32×32\n" + " │ └─ Number of warp gemm iterations: 4×4\n" " └─ Memory access:\n" " ├─ A Tile transfer: \n" - " │ ├─ Tile dimensions: 4×256×8×\n" + " │ ├─ Tile dimensions: 2×128×8×\n" " │ ├─ The innermost K subdimension size: 8\n" - " │ ├─ Spatial thread distribution over the data tile: 0×1×2\n" - " │ ├─ The order of accessing data tile axes: 0×1×2\n" + " │ ├─ Spatial thread distribution over the data tile: 1×0×2\n" + " │ ├─ The order of accessing data tile axes: 1×0×2\n" " │ ├─ Vectorized memory access axis index (with contiguous memory): 2\n" - " │ ├─ Vector access (GMEM read) instruction size: 2\n" - " │ ├─ Vector access (LDS write) instruction size: 2\n" - " │ └─ LDS data layout padding (to prevent bank conflicts): 2\n" + " │ ├─ Vector access (GMEM read) instruction size: 8\n" + " │ ├─ Vector access (LDS write) instruction size: 8\n" + " │ └─ LDS data layout padding (to prevent bank conflicts): 8\n" " ├─ B Tile transfer: \n" - " │ ├─ Tile dimensions: 4×256×8×\n" + " │ ├─ Tile dimensions: 2×128×8×\n" " │ ├─ The innermost K subdimension size: 8\n" - " │ ├─ Spatial thread distribution over the data tile: 0×1×2\n" - " │ ├─ The order of accessing data tile axes: 0×1×2\n" + " │ ├─ Spatial thread distribution over the data tile: 1×0×2\n" + " │ ├─ The order of accessing data tile axes: 1×0×2\n" " │ ├─ Vectorized memory access axis index (with contiguous memory): 2\n" - " │ ├─ Vector access (GMEM read) instruction size: 2\n" - " │ ├─ Vector access (LDS write) instruction size: 2\n" - " │ └─ LDS data layout padding (to prevent bank conflicts): 2\n" + " │ ├─ Vector access (GMEM read) instruction size: 8\n" + " │ ├─ Vector access (LDS write) instruction size: 8\n" + " │ └─ LDS data layout padding (to prevent bank conflicts): 8\n" " └─ C Tile transfer: \n" " ├─ Data shuffle (number of gemm instructions per iteration): 1×1\n" " ├─ Spatial thread distribution used to store data: 1×32×1×8\n" - " └─ Vector access (GMEM write) instruction size: 2")); + " ├─ Vector access (GMEM write) instruction size: 8\n" + " ├─ Struct does not contain optional num_gemm_k_prefetch_stage parameter\n" + " ├─ Max Transpose transfer scr scalar per vector: 1\n" + " ├─ Max Transpose dst scalar per vector: 1\n" + " └─ Num groups to merge: 4")); +} + +// Test printing of optional parameters num_groups_to_merge, +// nax_transose_transfer_src_scalar_per_vector and max_transpose_dst_scalar_per_vector +TEST(ConvDescriptionTest, BwdWeightWmmaCshuffleV3DescriptionTest) +{ + using Instance = ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Wmma_CShuffle< + 3, // NDimSpatial + ck::tensor_layout::convolution::GNDHWC, // InLayout + ck::tensor_layout::convolution::GKZYXC, // WeiLayout + ck::tensor_layout::convolution::GNDHWK, // OutLayout + ck::half_t, // InDataType + ck::half_t, // WeiDataType + ck::half_t, // OutDataType + float, // AccDataType + ck::tensor_operation::element_wise::PassThrough, // InElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // WeiElementwiseOperation + ck::tensor_operation::element_wise::PassThrough, // OutElementwiseOperation + ck::tensor_operation::device::ConvolutionBackwardWeightSpecialization:: + Default, // ConvBackwardWeightSpecialization + 256, // BlockSize + 128, // MPerBlock + 128, // NPerBlock + 16, // K0PerBlock + 8, // K1 + 32, // MPerWmma + 32, // NPerWmma + 4, // MRepeat + 4, // NRepeat + ck::Sequence<4, 64, 1>, // ABlockTransferThreadClusterLengths_K0_M_K1 + ck::Sequence<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // ABlockTransferSrcAccessOrder + 2, // ABlockTransferSrcVectorDim + 8, // ABlockTransferSrcScalarPerVector + 8, // ABlockTransferDstScalarPerVector_K1 + 1, // ABlockLdsAddExtraM + ck::Sequence<4, 64, 1>, // BBlockTransferThreadClusterLengths_K0_N_K1 + ck::Sequence<1, 0, 2>, // BBlockTransferThreadClusterArrangeOrder_ + ck::Sequence<1, 0, 2>, // BBlockTransferSrcAccessOrder_ + 2, // BBlockTransferSrcVectorDim + 8, // BBlockTransferSrcScalarPerVector + 8, // BBlockTransferDstScalarPerVector_K1 + 1, // BBlockLdsAddExtraN + 1, // CShuffleMXdlPerWavePerShuffle + 1, // CShuffleNXdlPerWavePerShuffle + ck::Sequence<1, + 32, + 1, + 8>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock_ + 8, // CDEBlockTransferScalarPerVector_NPerBlock_ + 1, // NummGemmKPrefetchStage + ck::LoopScheduler::Default, // BlkGemmPipeSched + ck::PipelineVersion::v1, // BlkGemmPipelineVer + false>; // BComputeDataType + + EXPECT_THAT( + ckr::describe().detailed(), + ckt::StringEqWithDiff( // + "3D Backward Weight Convolution Kernel\n" + "├─ Signature\n" + "│ ├─ Tensor Type: FP16\n" + "│ ├─ Input Layout: GNDHWC\n" + "│ ├─ Weight Layout: GKZYXC\n" + "│ ├─ Output Layout: GNDHWK\n" + "│ ├─ Input elementwise operation: PASS_THROUGH\n" + "│ ├─ Weights elementwise operation: PASS_THROUGH\n" + "│ └─ Output elementwise operation: PASS_THROUGH\n" + "└─ Algorithm\n" + " ├─ Thread block size: 256\n" + " ├─ Data tile size: 128×128×16\n" + " ├─ Struct does not contain optional gemm_padding argument\n" + " ├─ Convolution specialization: DEFAULT\n" + " ├─ Pipeline version: V1\n" + " ├─ Pipeline scheduler: DEFAULT\n" + " ├─ Warp Gemm parameters: \n" + " │ ├─ subtile size: 32×32\n" + " │ └─ Number of warp gemm iterations: 4×4\n" + " └─ Memory access:\n" + " ├─ A Tile transfer: \n" + " │ ├─ Tile dimensions: 2×128×8×\n" + " │ ├─ The innermost K subdimension size: 8\n" + " │ ├─ Spatial thread distribution over the data tile: 1×0×2\n" + " │ ├─ The order of accessing data tile axes: 1×0×2\n" + " │ ├─ Vectorized memory access axis index (with contiguous memory): 2\n" + " │ ├─ Vector access (GMEM read) instruction size: 8\n" + " │ ├─ Vector access (LDS write) instruction size: 8\n" + " │ └─ LDS data layout padding (to prevent bank conflicts): 8\n" + " ├─ B Tile transfer: \n" + " │ ├─ Tile dimensions: 2×128×8×\n" + " │ ├─ The innermost K subdimension size: 8\n" + " │ ├─ Spatial thread distribution over the data tile: 1×0×2\n" + " │ ├─ The order of accessing data tile axes: 1×0×2\n" + " │ ├─ Vectorized memory access axis index (with contiguous memory): 2\n" + " │ ├─ Vector access (GMEM read) instruction size: 8\n" + " │ ├─ Vector access (LDS write) instruction size: 8\n" + " │ └─ LDS data layout padding (to prevent bank conflicts): 8\n" + " └─ C Tile transfer: \n" + " ├─ Data shuffle (number of gemm instructions per iteration): 1×1\n" + " ├─ Spatial thread distribution used to store data: 1×32×1×8\n" + " ├─ Vector access (GMEM write) instruction size: 8\n" + " ├─ Num gemm k prefetch stage: 1\n" + " ├─ Struct does not contain optional max_transpose_transfer_src_scalar_per_vector " + "parameter\n" + " ├─ Struct does not contain optional max_transpose_dst_scalar_per_vector parameter\n" + " └─ Struct does not contain optional num_groups_to_merge parameter")); } TEST(ConvDescriptionTest, DefaultInstanceHasInstanceString) From 42048bdb7d8d931966af76c6dacfedce1c9da90a Mon Sep 17 00:00:00 2001 From: Robin Voetter Date: Wed, 28 Jan 2026 17:41:02 +0100 Subject: [PATCH 04/11] [CK_BUILDER] Integrate CKB validation with CK verification (#3649) * ck-builder: tensor copy function This function copies one tensor to another, so that the memory layout can be changed between them. * ck-builder: fix ck::bhalf literals These types don't work properly. * ck-builder: abstract compare_elements in gpu_verification.hpp and make builder use it This reduces the amount of duplicated code a bit. * ck-builder: add flat tensor iterator This "iterator" type pretends to be a pointer, useful for passing tensors to functions expecting pointer-like types. * ck-builder: integrate validation with ck gpu verification By templating the gpu_verify function over iterators, we can use the new FlatTensorIterator to adapt the function to multi- dimensional tensors without changing either implementation too much. * ck-builder: add check_by_accumulations This changes the gpu_verification.hpp code to also accept "iterator" types for the relevant gpu_verify and gpu_reduce_max functions. * ck: fix test_gpu_verification GenerateRandomData for bhalf is_integer_it yields true, but it is not actually an integer. * ck: make gpu_verification kernels be proper persistent kernels Previously these were using a hardcoded value for the grid size. This commit changes that so that the grid size is automatically derived from the kernel's occupancy and the number of multiprocessors on the GPU. * ck: clean up gpu_verification.hpp using block_reduce This implements a small generic block reduce function, and rewrites the rest of gpu_verification.hpp using that function to clean it up a bit. * ck-builder: doc typos * ck-builder: update testing readme with validation interface. * ck-builder: rebase fixes + review comments * ck-builder: fix device integer generation with float types Passing bfloat here causes a nans due to type_convert performing a bitcast. * ck: another bhalf_t bug CK expects that int-generation with ck::bhalf_t yields bhalf integers, not unsigned integers. This makes the logic of FillUniformRandInteger compatible with GeneratorTensor_2, however idiotic that may be. --- .../include/ck_tile/builder/testing/README.md | 30 +- .../builder/testing/tensor_foreach.hpp | 112 +++++ .../ck_tile/builder/testing/validation.hpp | 167 +++---- experimental/builder/test/testing_utils.hpp | 3 +- .../builder/test/unit_conv_fwd_testing.cpp | 12 +- .../builder/test/unit_tensor_foreach.cpp | 96 +++++ experimental/builder/test/unit_validation.cpp | 65 ++- .../utility/device_tensor_generator.hpp | 8 +- .../ck/library/utility/gpu_verification.hpp | 406 +++++++++++------- test/device_memory/test_device_prng.cpp | 26 +- .../test_gpu_verification.cpp | 2 +- 11 files changed, 636 insertions(+), 291 deletions(-) diff --git a/experimental/builder/include/ck_tile/builder/testing/README.md b/experimental/builder/include/ck_tile/builder/testing/README.md index c6662c2b04..cbac281aca 100644 --- a/experimental/builder/include/ck_tile/builder/testing/README.md +++ b/experimental/builder/include/ck_tile/builder/testing/README.md @@ -12,7 +12,8 @@ The core components are: - **`Args`**: A struct template that holds runtime parameters for a specific test case. - **`Input`** and **`Output`**: Helper classes that groups operation inputs and outputs. -- **`Validator`**: A utility that performs on-GPU validation and integrates with GoogleTest/GoogleMock. +- **`run()`**: Invokes an algorithm on the GPU. +- **`validate()`**: A utility that performs on-GPU validation and integrates with GoogleTest/GoogleMock. Together, these components enable a structured approach to kernel testing that mirrors the Given-When-Then pattern commonly used in behavior-driven development. @@ -200,26 +201,27 @@ auto reference_outputs = ck_tile::builder::test::allocate_outputs(args); ck_tile::builder::test::run(conv, args, inputs.get(), reference_outputs.get()); ``` -#### `Validator` +#### Validating Results -The `Validator` class encapsulates the validation logic. It performs on-GPU correctness checks by comparing two instances of the `Outputs` structure. +In order to actually verify that the results of the executed device operation are correct, they are compared against the reference output obtained in the previous step. This is done by calling `validate()` with the runtime arguments of the operation, as well as both the actual and reference output. This then yields a *`ValidationReport`*, a type which holds information about which tensors of the output are considered to be equivalent and which are considered to be different. Actually comparing the tensor elements is performed on the GPU to keep the tests fast. ```cpp -ck_tile::builder::test::Validator validator(outputs.get(), reference_outputs.get()); +const auto report = ck_tile::builder::test::validate(args, outputs.get(), reference_outputs.get()); ``` -The `Validator` provides methods that return GoogleMock matchers, enabling clean integration with GoogleTest: +`ValidationReport::get_errors()` returns a vector of tensors from both outputs which are considered to be incorrect, each error case exposes some information about what went wrong. ```cpp -EXPECT_THAT(validator.result(), validator.matches_reference_output()); +for (const auto& e : report.get_errors()) { + std::cout << "error: " << e.tensor_name << " was incorrect!" << std::endl; +} ``` -The `matches_reference_output()` matcher checks that the output is numerically correct within acceptable tolerances. The `Validator` can also provide more detailed diagnostics, such as: +GoogleTest/GoogleMock integration is provided using the `MatchesReference` matcher. This invokes `validate()` internally, and then turns the result into a GoogleMock-comparible value. Note that this function is closely tied to GoogleMock and the test setup that CK-Builder uses internally, and so is not exposed through the CK-Builder public interface. -- Maximum absolute error -- Maximum relative error -- Number of mismatched elements -- Specific locations of errors +```cpp +EXPECT_THAT(outputs.get(), MatchesReference(args, reference_outputs.get())); +``` ## Complete Example @@ -232,6 +234,7 @@ Here's a complete test that demonstrates the Given-When-Then pattern: #include "ck_tile/builder/conv_builder.hpp" #include "ck_tile/testing/tensor_memory_manager.hpp" #include "ck_tile/testing/validator.hpp" +#include "testing_utils.hpp" // Define the convolution signature struct ConvSignature { @@ -318,8 +321,7 @@ TEST(ConvolutionTest, Forward2D_FP16) { ck_tile::builder::test::run(conv, args, inputs.get(), reference_outputs.get()); // Check the results - ck_tile::builder::test::Validator validator(outputs.get(), reference_outputs.get()); - EXPECT_THAT(validator.result(), validator.is_ok()); + EXPECT_THAT(outputs.get(), ck_tile::test::MatchesReference(args, reference_outputs.get())); } ``` @@ -333,7 +335,7 @@ TEST(ConvolutionTest, Forward2D_FP16) { 4. **Flexibility**: The `Args` struct can be easily extended to support different test scenarios, `Inputs` and `Outputs` can be modified to support additional tensors where necessary, and alternatives to `init_inputs()` can be provided to support additional testing strategies. -5. **Integration**: The `Validator` integrates seamlessly with GoogleTest/GoogleMock, providing familiar assertion syntax. +5. **Integration**: `validate()` integrates seamlessly with GoogleTest/GoogleMock through `MatchesReference`, providing familiar assertion syntax. 6. **Maintainability**: Changes to the testing infrastructure are localized to the utility classes, not scattered across individual tests. diff --git a/experimental/builder/include/ck_tile/builder/testing/tensor_foreach.hpp b/experimental/builder/include/ck_tile/builder/testing/tensor_foreach.hpp index 28ab954de9..8503157bb9 100644 --- a/experimental/builder/include/ck_tile/builder/testing/tensor_foreach.hpp +++ b/experimental/builder/include/ck_tile/builder/testing/tensor_foreach.hpp @@ -6,6 +6,7 @@ #include "ck_tile/builder/testing/tensor_descriptor.hpp" #include "ck_tile/builder/factory/helpers/ck/conv_tensor_type.hpp" #include +#include #include #include @@ -348,4 +349,115 @@ void clear_tensor_buffer(const TensorDescriptor& desc, fill_tensor_buffer(desc, buffer, [value]([[maybe_unused]] size_t i) { return value; }); } +/// @brief Utility for copying a tensor from one layout to another +/// +/// This function copies tensor data from `src_buffer` to `dst_buffer`, +/// changing the layout from `src_desc` to `dst_desc`. Note that the src and +/// dst tensor lengths must be compatible, otherwise this function may write +/// out of bounds. +/// +/// @tparam DT The element datatype of both tensors. +/// @tparam RANK The rank (number of spatial dimensions) of the tensors. +/// +/// @param src_desc The descriptor of the source tensor to copy from. +/// @param src_buffer The memory of the source tensor. +/// @param dst_desc The descriptor of the destination tensor to copy to. +/// @param dst_buffer The memory of the destination tensor. +template +void copy_tensor(const TensorDescriptor& src_desc, + const void* src_buffer, + const TensorDescriptor& dst_desc, + void* dst_buffer) +{ + assert(src_desc.get_lengths() == dst_desc.get_lengths()); + const auto src_strides = src_desc.get_strides(); + const auto dst_strides = dst_desc.get_strides(); + tensor_foreach(dst_desc.get_lengths(), + [src_buffer, dst_buffer, src_strides, dst_strides](const auto& index) { + using T = detail::cpp_type_t
; + const auto* src = static_cast(src_buffer); + auto* dst = static_cast(dst_buffer); + const auto src_off = calculate_offset(index, src_strides); + const auto dst_off = calculate_offset(index, dst_strides); + dst[dst_off] = src[src_off]; + }); +} + +/// @brief Simple iterator implementation over tensors. +/// +/// This type implements a simple "iterator" type for tensor types, +/// basically exposing operator[] for flat indices. This type is useful +/// to be able to provide a "pointer-like" object to API that does not +/// expect higher dimensional tensor types, and expects linear pointers +/// instead. Ideally, one just needs to replace the `T* ptr` with +/// `Iterator it` to update those API to be compatible with this type. +/// +/// @note This is not intended to be a full implementation of the C++ +/// iterator concept. For example, it does not really hold any state, +/// because that is not really useful anyway. +/// +/// @tparam DT The datatype of the tensor to iterate over. Note that this +/// is only here for reference purposes, the actual data type of the backing +/// memory is provided via the backing iterator type. +/// @tparam RANK The rank (number of spatial dimensions) of the tensors. +/// @tparam Iterator The backing iterator type. This can be a (non-void) +/// pointer type. +template +struct FlatTensorIterator +{ + /// @brief Construct a FlatTensorIterator. + /// + /// Construct a FlatTensorIterator from a tensor descriptor and a backing + /// iterator. The backing iterator can just be a non-void pointer type, + /// note that the result of FlatTensorIterator::operator[] is the same as + /// that of Iterator::operator[]. + /// + /// @param desc The descriptor of the tensor to iterate. + /// @param inner The inner iterator, for example a (non-void) pointer. + FlatTensorIterator(const TensorDescriptor& desc, Iterator inner) + : iter_(desc.get_lengths()), strides_(desc.get_strides()), inner_(inner) + { + } + + /// @brief Return the value at a particular flat index. + /// + /// This function returns the value of the tensor at flat coordinate + /// `flat_index`. This index is then unflattened into a multi-dimensional + /// index according to the way described in `NdIter`, and a tensor offset + /// is computed from that according to `calculate_offset`. The value at + /// that offset in the inner iterator is then the return value of this + /// function. + /// + /// @note NdIter iterates such that the inner dimension (right-most value + /// in the tensor shape) changes fastest. + /// + /// @note This function performs no bounds checking. + /// + /// @param flat_index The flat index into this tensor. + /// + /// @pre flat_index < numel() + /// + /// @see NdIter + __host__ __device__ auto& operator[](size_t flat_index) const + { + const auto index = iter_(flat_index); + const auto offset = calculate_offset(index, strides_); + return inner_[offset]; + } + + /// @brief Return the total number of elements to iterate over. + /// + /// @see NdIter::numel() + __host__ __device__ size_t numel() const { return iter_.numel(); } + + private: + NdIter iter_; + Extent strides_; + Iterator inner_; +}; + +template +FlatTensorIterator(const TensorDescriptor&, + Iterator) -> FlatTensorIterator; + } // namespace ck_tile::builder::test diff --git a/experimental/builder/include/ck_tile/builder/testing/validation.hpp b/experimental/builder/include/ck_tile/builder/testing/validation.hpp index 8410a71b15..d510b96264 100644 --- a/experimental/builder/include/ck_tile/builder/testing/validation.hpp +++ b/experimental/builder/include/ck_tile/builder/testing/validation.hpp @@ -8,6 +8,7 @@ #include "ck_tile/builder/testing/tensor_foreach.hpp" #include "ck_tile/builder/factory/helpers/ck/conv_tensor_type.hpp" #include "ck/utility/type_convert.hpp" +#include "ck/library/utility/gpu_verification.hpp" #include #include #include @@ -48,8 +49,8 @@ struct ValidationReport /// The total number of elements in each tensor. uint64_t total_elements; - /// The number of elements which were bitwise 0. - uint64_t zero_elements; + /// Set to true if both tensors have all their elements be 0. + bool both_all_zero; // Max error. double max_error; @@ -59,7 +60,7 @@ struct ValidationReport /// If both tensors are all zero, it indicates either an incorrect testing setup /// or an issue with the testing framework. For that reason we also consider that /// a failure. - bool is_all_zero() const { return zero_elements == total_elements; } + bool is_all_zero() const { return both_all_zero; } /// @brief Return whether the check associated to this case was successful. /// @@ -86,7 +87,7 @@ struct ValidationReport /// @brief Compare two tensors and record the results in the report. /// - /// This is the main function used to compare two tensors. The results of this + /// This is one of the main function used to compare two tensors. The results of this /// comparison, including any supplemental information, is recorded into the report. /// /// @returns `false` if the comparison failed. If so, the details can be found via @@ -111,8 +112,45 @@ struct ValidationReport const TensorDescriptor& descriptor, const void* actual, const void* expected, - double rtol = 1e-3, - double atol = 1e-3); + float rtol = 1e-3f, + float atol = 1e-3f); + + /// @brief Compare two tensors and record the results in the report, with automatic + /// computation of tolerances. + /// + /// This variant computes the tolerances automatically based on the compute + /// (accumulation) type, and the number of accumulations required per result value. + /// This is one of the main function used to compare two tensors. The results of this + /// comparison, including any supplemental information, is recorded into the report. + /// @returns `false` if the comparison failed. If so, the details can be found via + /// `get_errors()`. + /// + /// @tparam OutDataType The data type of the tensors to check. This is the type of the + /// values in tensor memory. + /// @tparam ComputeType The data type that tensor operations are computed with internally. + /// @tparam AccType The data type that tensor values are accumulated with internally. + /// @tparam RANK The rank (number of spatial dimensions) of the tensor to check. + /// + /// @param tensor_name The name of the tensors to check. This should be a value by which + /// whoever is debugging the associated test later can easily find out which of the + /// outputs of a device operation was incorrect. + /// @param descriptor The descriptor (memory layout) of the tensor. + /// @param actual The device buffer with the values of the tensor to-be-tested, ie, the + /// results of the device operation. + /// @param expected The device buffer with the values of the reference tensor. These are + /// treated as a "golden standard", and should usually be generated by a reference + /// implementation. + /// @param number_of_accumulations The maximum number of accumulations required to compute + /// a value of the result tensor. + template + bool check_by_accumulations(std::string_view tensor_name, + const TensorDescriptor& descriptor, + const void* actual, + const void* expected, + const size_t number_of_accumulations); private: std::vector reports_; @@ -121,89 +159,58 @@ struct ValidationReport template bool ValidationReport::check(std::string_view tensor_name, const TensorDescriptor& descriptor, - const void* actual_data, - const void* expected_data, - double rtol, - double atol) + const void* actual, + const void* expected, + float rtol, + float atol) { - const auto strides = descriptor.get_strides(); + using CKType = detail::cpp_type_t
; - // During development and CI, only the kernels that were changed would fail, and so we can - // assume that the average case does not have errors. Therefore, split out testing into a - // quick test which just counts the incorrect elements, and a more in-depth test that also - // returns the indices of the incorrect items. + const auto a_it = FlatTensorIterator(descriptor, static_cast(actual)); + const auto e_it = FlatTensorIterator(descriptor, static_cast(expected)); + const auto numel = a_it.numel(); - // Initial pass: count errors - - // Allocate and reset counter - auto d_counters = alloc_buffer(sizeof(uint64_t) * 3); - check_hip(hipMemset(d_counters.get(), 0, sizeof(uint64_t) * 3)); - - auto d_error_count = &reinterpret_cast(d_counters.get())[0]; - auto d_zero_count = &reinterpret_cast(d_counters.get())[1]; - auto d_max_error = &reinterpret_cast(d_counters.get())[2]; - - tensor_foreach(descriptor.get_lengths(), [=](auto index) { - using CKType = typename factory::internal::DataTypeToCK
::type; - - const auto* actual = static_cast(actual_data); - const auto* expected = static_cast(expected_data); - - static_assert(!std::is_same_v, - "TODO implement compare_kernel() for double"); - - const auto offset = calculate_offset(index, strides); - - const auto a = actual[offset]; - const auto b = expected[offset]; - - const auto o = static_cast(type_convert(a)); - const auto r = static_cast(type_convert(b)); - const auto err = std::abs(o - r); - - atomicMax(d_max_error, err); - if(err > atol + rtol * std::abs(r) || !std::isfinite(o) || !std::isfinite(r)) - { - // We expect the number of errors to be very low, so just use an atomic - // for now. - atomicAdd(d_error_count, 1); - } - - // Now compare the numbers as bitwise too. - // Update the counter if they're both zero. - using Bytes = std::array; - bool all_zero = true; - for(auto x : std::bit_cast(a)) - { - if(x != std::byte{0}) - all_zero = false; - } - for(auto x : std::bit_cast(b)) - { - if(x != std::byte{0}) - all_zero = false; - } - if(all_zero) - { - atomicAdd(d_zero_count, 1); - } - }); - - uint64_t error_count = 0; - check_hip(hipMemcpy(&error_count, d_error_count, sizeof(uint64_t), hipMemcpyDeviceToHost)); - uint64_t zero_count = 0; - check_hip(hipMemcpy(&zero_count, d_zero_count, sizeof(uint64_t), hipMemcpyDeviceToHost)); - double max_error = 0; - check_hip(hipMemcpy(&max_error, d_max_error, sizeof(double), hipMemcpyDeviceToHost)); + const auto result = ck::profiler::gpu_verify(a_it, e_it, rtol, atol, numel); // TODO: Gather detailed coordinates. reports_.push_back(Case{ .tensor_name = std::string(tensor_name), - .wrong_elements = error_count, + .wrong_elements = result.error_count, .total_elements = descriptor.get_element_size(), - .zero_elements = zero_count, - .max_error = max_error, + .both_all_zero = result.all_zero, + .max_error = result.max_error, + }); + + return reports_.back().is_ok(); +} + +template +bool ValidationReport::check_by_accumulations(std::string_view tensor_name, + const TensorDescriptor& descriptor, + const void* actual, + const void* expected, + const size_t number_of_accumulations) +{ + using CKComputeType = detail::cpp_type_t; + using CKAccType = detail::cpp_type_t; + using CKOutDataType = detail::cpp_type_t; + + const auto a_it = FlatTensorIterator(descriptor, static_cast(actual)); + const auto e_it = FlatTensorIterator(descriptor, static_cast(expected)); + const auto numel = a_it.numel(); + + const auto result = ck::profiler::gpu_verify( + a_it, e_it, static_cast(number_of_accumulations), numel); + + // TODO: Gather detailed coordinates. + + reports_.push_back(Case{ + .tensor_name = std::string(tensor_name), + .wrong_elements = result.error_count, + .total_elements = descriptor.get_element_size(), + .both_all_zero = result.all_zero, + .max_error = result.max_error, }); return reports_.back().is_ok(); diff --git a/experimental/builder/test/testing_utils.hpp b/experimental/builder/test/testing_utils.hpp index 55de133a2a..415e3dfe43 100644 --- a/experimental/builder/test/testing_utils.hpp +++ b/experimental/builder/test/testing_utils.hpp @@ -209,7 +209,8 @@ struct ReferenceOutputMatcher // Round to 2 digits const float percentage = e.wrong_elements * 10000 / e.total_elements / 100.f; *listener << e.wrong_elements << "/" << e.total_elements - << " incorrect elements (~" << percentage << "%)"; + << " incorrect elements (~" << percentage << "%)," << " max error " + << e.max_error; } } } diff --git a/experimental/builder/test/unit_conv_fwd_testing.cpp b/experimental/builder/test/unit_conv_fwd_testing.cpp index 9fc07568b4..cbb31dc495 100644 --- a/experimental/builder/test/unit_conv_fwd_testing.cpp +++ b/experimental/builder/test/unit_conv_fwd_testing.cpp @@ -98,8 +98,10 @@ TEST(ConvFwdTesting, Validate) [&]([[maybe_unused]] std::string_view name, const auto& desc, void* ckt::Outputs::*ptr) { - ckt::clear_tensor_buffer(desc, a.get().*ptr, ck::bhalf_t{123}); - ckt::clear_tensor_buffer(desc, b.get().*ptr, ck::bhalf_t{123}); + ckt::clear_tensor_buffer( + desc, a.get().*ptr, ck::type_convert(123)); + ckt::clear_tensor_buffer( + desc, b.get().*ptr, ck::type_convert(123)); }); const auto report = ckt::validate(ARGS, a.get(), b.get()); @@ -115,8 +117,10 @@ TEST(ConvFwdTesting, Validate) const auto& desc, void* ckt::Outputs::*ptr) { ++field_count; - ckt::clear_tensor_buffer(desc, a.get().*ptr, ck::bhalf_t{2}); - ckt::clear_tensor_buffer(desc, b.get().*ptr, ck::bhalf_t{1}); + ckt::clear_tensor_buffer( + desc, a.get().*ptr, ck::type_convert(2)); + ckt::clear_tensor_buffer( + desc, b.get().*ptr, ck::type_convert(1)); }); const auto report = ckt::validate(ARGS, a.get(), b.get()); diff --git a/experimental/builder/test/unit_tensor_foreach.cpp b/experimental/builder/test/unit_tensor_foreach.cpp index f689d3c82f..70777e87cf 100644 --- a/experimental/builder/test/unit_tensor_foreach.cpp +++ b/experimental/builder/test/unit_tensor_foreach.cpp @@ -225,3 +225,99 @@ TEST(TensorForeach, ClearTensorZeros) EXPECT_THAT(actual, Eq(0)); } + +TEST(TensorForeach, CopyTensor) +{ + constexpr auto dt = ckb::DataType::I32; + const ckt::Extent shape = {10, 3, 45, 23, 6}; + using Counter = uint32_t; + + const auto src_desc = ckt::make_descriptor
(shape, ckt::PackedRightLayout{}); + const auto dst_desc = ckt::make_descriptor
(shape, ckt::PackedLeftLayout{}); + + auto src_buffer = ckt::alloc_tensor_buffer(src_desc); + auto dst_buffer = ckt::alloc_tensor_buffer(dst_desc); + + const auto gen = [](const auto& index, const auto& lengths) { + // Simple incrementing counter + return static_cast(ckt::calculate_offset(index, lengths)); + }; + + ckt::fill_tensor( + src_desc, src_buffer.get(), [lengths = src_desc.get_lengths(), gen](const auto& index) { + return gen(index, lengths); + }); + ckt::clear_tensor_buffer(dst_desc, dst_buffer.get()); + + // Perform the actual test + + ckt::copy_tensor(src_desc, src_buffer.get(), dst_desc, dst_buffer.get()); + + // Check that the dst tensor has the same data + + auto d_invalid = ckt::alloc_buffer(sizeof(Counter)); + ckt::check_hip(hipMemset(d_invalid.get(), 0, sizeof(Counter))); + + ckt::tensor_foreach(shape, + [lengths = dst_desc.get_lengths(), + gen, + dst = dst_buffer.get(), + invalid = reinterpret_cast(d_invalid.get()), + strides = dst_desc.get_strides()](const auto& index) { + const auto offset = ckt::calculate_offset(index, strides); + const auto expected = gen(index, lengths); + const auto actual = reinterpret_cast(dst)[offset]; + + if(expected != actual) + atomicAdd(invalid, 1); + }); + + Counter invalid = 0; + ckt::check_hip(hipMemcpy(&invalid, d_invalid.get(), sizeof(Counter), hipMemcpyDeviceToHost)); + + EXPECT_THAT(invalid, Eq(0)); +} + +TEST(TensorForeach, FlatTensorIterator) +{ + using Counter = uint32_t; + + constexpr auto dt = ckb::DataType::I32; + const ckt::Extent shape = {10, 9, 8, 7, 6, 5, 4, 3, 2, 1}; + const ckt::Extent packed_strides = ckt::PackedRightLayout{}(shape); + + const auto desc = ckt::make_descriptor
(shape, ckt::PackedLeftLayout{}); + + auto buffer = ckt::alloc_tensor_buffer(desc); + + // Fill the tensor with random values according to the *flat* index. The + // FlatTensorIterator iterates over flat values even if the strides are not + // packed, so indexing these elements according to the flat index in the + // iterator should yield again this value. + ckt::fill_tensor(desc, buffer.get(), [packed_strides](const auto& index) { + const auto flat_index = ckt::calculate_offset(index, packed_strides); + return static_cast(flat_index * 10001 % 1001); + }); + + auto iterator = ckt::FlatTensorIterator(desc, reinterpret_cast(buffer.get())); + + auto d_invalid = ckt::alloc_buffer(sizeof(Counter)); + ckt::check_hip(hipMemset(d_invalid.get(), 0, sizeof(Counter))); + + ckt::tensor_foreach(shape, + [iterator, + packed_strides, + strides = desc.get_strides(), + data = reinterpret_cast(buffer.get()), + invalid = reinterpret_cast(d_invalid.get())](const auto& index) { + const auto flat_index = ckt::calculate_offset(index, packed_strides); + const auto offset = ckt::calculate_offset(index, strides); + if(iterator[flat_index] != data[offset]) + atomicAdd(invalid, 1); + }); + + Counter invalid = 0; + ckt::check_hip(hipMemcpy(&invalid, d_invalid.get(), sizeof(Counter), hipMemcpyDeviceToHost)); + + EXPECT_THAT(invalid, Eq(0)); +} diff --git a/experimental/builder/test/unit_validation.cpp b/experimental/builder/test/unit_validation.cpp index 0dad8593fb..178488c33e 100644 --- a/experimental/builder/test/unit_validation.cpp +++ b/experimental/builder/test/unit_validation.cpp @@ -74,7 +74,8 @@ TYPED_TEST(ValidationReportTests, SingleCorrect) ckt::fill_tensor(desc, b.get(), generator); ckt::ValidationReport report; - report.check("correct", desc, b.get(), a.get()); + report.check("correct - explicit tolerance", desc, b.get(), a.get()); + report.check_by_accumulations("correct - implicit tolerance", desc, b.get(), a.get(), 0); EXPECT_THAT(report.get_errors().size(), Eq(0)); } @@ -97,17 +98,22 @@ TYPED_TEST(ValidationReportTests, SingleIncorrect) }); ckt::ValidationReport report; - report.check("incorrect", desc, b.get(), a.get()); + report.check("incorrect - explicit tolerance", desc, b.get(), a.get()); + report.check_by_accumulations("incorrect - implicit tolerance", desc, b.get(), a.get(), 0); const auto errors = report.get_errors(); const auto flat_size = desc.get_element_size(); const auto expected_errors = flat_size >= 999999 ? 3 : flat_size >= 12345 ? 2 : 1; - ASSERT_THAT(errors.size(), Eq(1)); - EXPECT_THAT(errors[0].tensor_name, StrEq("incorrect")); - EXPECT_THAT(errors[0].wrong_elements, Eq(expected_errors)); - EXPECT_THAT(errors[0].total_elements, Eq(desc.get_element_size())); + ASSERT_THAT(errors.size(), Eq(2)); + EXPECT_THAT(errors[0].tensor_name, StrEq("incorrect - explicit tolerance")); + EXPECT_THAT(errors[1].tensor_name, StrEq("incorrect - implicit tolerance")); + for(int i = 0; i < 2; ++i) + { + EXPECT_THAT(errors[i].wrong_elements, Eq(expected_errors)); + EXPECT_THAT(errors[i].total_elements, Eq(desc.get_element_size())); + } } TYPED_TEST(ValidationReportTests, ZeroIsIncorrect) @@ -121,14 +127,20 @@ TYPED_TEST(ValidationReportTests, ZeroIsIncorrect) ckt::clear_tensor_buffer(desc, b.get()); ckt::ValidationReport report; - report.check("zero_is_incorrect", desc, b.get(), a.get()); + report.check("zero_is_incorrect - explicit tolerance", desc, b.get(), a.get()); + report.check_by_accumulations( + "zero_is_incorrect - implicit tolerance", desc, b.get(), a.get(), 0); const auto errors = report.get_errors(); - ASSERT_THAT(errors.size(), Eq(1)); - EXPECT_THAT(errors[0].tensor_name, StrEq("zero_is_incorrect")); - EXPECT_THAT(errors[0].wrong_elements, Eq(0)); - EXPECT_THAT(errors[0].total_elements, Eq(desc.get_element_size())); - EXPECT_THAT(errors[0].zero_elements, Eq(desc.get_element_size())); + ASSERT_THAT(errors.size(), Eq(2)); + EXPECT_THAT(errors[0].tensor_name, StrEq("zero_is_incorrect - explicit tolerance")); + EXPECT_THAT(errors[1].tensor_name, StrEq("zero_is_incorrect - implicit tolerance")); + for(int i = 0; i < 2; ++i) + { + EXPECT_THAT(errors[i].wrong_elements, Eq(0)); + EXPECT_THAT(errors[i].total_elements, Eq(desc.get_element_size())); + EXPECT_THAT(errors[i].both_all_zero, Eq(true)); + } } TEST(ValidationReportTests, MultipleSomeIncorrect) @@ -143,11 +155,12 @@ TEST(ValidationReportTests, MultipleSomeIncorrect) auto b = ckt::alloc_tensor_buffer(desc); ckt::fill_tensor_buffer( - desc, a.get(), [](size_t i) { return ck::type_convert(i % 100); }); + desc, a.get(), [](size_t i) { return ck::type_convert(float(i % 100)); }); ckt::fill_tensor_buffer( - desc, b.get(), [](size_t i) { return ck::type_convert(i % 101); }); + desc, b.get(), [](size_t i) { return ck::type_convert(float(i % 101)); }); - report.check("incorrect 1", desc, b.get(), a.get()); + report.check("incorrect 1 - explicit tolerance", desc, b.get(), a.get()); + report.check("incorrect 1 - implicit tolerance", desc, b.get(), a.get(), 0); } { @@ -169,7 +182,8 @@ TEST(ValidationReportTests, MultipleSomeIncorrect) } }); - report.check("correct", desc, b.get(), a.get()); + report.check("correct - explicit tolerance", desc, b.get(), a.get()); + report.check("correct - implicit tolerance", desc, b.get(), a.get(), 0); } { @@ -182,16 +196,21 @@ TEST(ValidationReportTests, MultipleSomeIncorrect) ckt::fill_tensor_buffer(desc, a.get(), []([[maybe_unused]] size_t i) { return 1; }); ckt::fill_tensor_buffer(desc, b.get(), []([[maybe_unused]] size_t i) { return 555; }); - report.check("incorrect 2", desc, b.get(), a.get()); + report.check("incorrect 2 - explicit tolerance", desc, b.get(), a.get()); + report.check("incorrect 2 - implicit tolerance", desc, b.get(), a.get(), 0); } const auto errors = report.get_errors(); - ASSERT_THAT(errors.size(), Eq(2)); - EXPECT_THAT(errors[0].tensor_name, StrEq("incorrect 1")); + ASSERT_THAT(errors.size(), Eq(4)); + EXPECT_THAT(errors[0].tensor_name, StrEq("incorrect 1 - explicit tolerance")); EXPECT_THAT(errors[0].wrong_elements, Eq(46840334)); - EXPECT_THAT(errors[1].tensor_name, StrEq("incorrect 2")); - EXPECT_THAT(errors[1].wrong_elements, Eq(482800)); + EXPECT_THAT(errors[1].tensor_name, StrEq("incorrect 1 - implicit tolerance")); + EXPECT_THAT(errors[1].wrong_elements, Eq(46840334)); + EXPECT_THAT(errors[2].tensor_name, StrEq("incorrect 2 - explicit tolerance")); + EXPECT_THAT(errors[2].wrong_elements, Eq(482800)); + EXPECT_THAT(errors[3].tensor_name, StrEq("incorrect 2 - implicit tolerance")); + EXPECT_THAT(errors[3].wrong_elements, Eq(482800)); } // MatchesReference operates on the types defined in testing.hpp, so just @@ -234,7 +253,7 @@ ValidationReport validate(const Args& args, { ValidationReport report; report.check("a", args.make_a_descriptor(), actual.a, expected.a); - report.check("b", args.make_b_descriptor(), actual.b, expected.b); + report.check_by_accumulations("b", args.make_b_descriptor(), actual.b, expected.b, 0); return report; } @@ -299,5 +318,5 @@ TEST(MatchesReference, Incorrect) EXPECT_THAT(listener.str(), StringEqWithDiff( // "1 tensors failed to validate\n" - " - a: 625/625 incorrect elements (~100%)")); + " - a: 625/625 incorrect elements (~100%), max error 1")); } diff --git a/include/ck/library/utility/device_tensor_generator.hpp b/include/ck/library/utility/device_tensor_generator.hpp index 60bc3110d4..4d1bcc9af9 100644 --- a/include/ck/library/utility/device_tensor_generator.hpp +++ b/include/ck/library/utility/device_tensor_generator.hpp @@ -67,8 +67,12 @@ __global__ void fill_tensor_uniform_rand_int_values(T* p, } else { - p[i] = ck::type_convert( - static_cast((ran_gen_round_u32(s)) % (max_value - min_value)) + min_value); + const auto value = + static_cast((ran_gen_round_u32(s)) % (max_value - min_value)) + min_value; + if constexpr(std::is_integral_v && !std::is_same_v) + p[i] = ck::type_convert(value); + else + p[i] = ck::type_convert(value); } } } diff --git a/include/ck/library/utility/gpu_verification.hpp b/include/ck/library/utility/gpu_verification.hpp index e4a444ecb9..6be7e1886b 100644 --- a/include/ck/library/utility/gpu_verification.hpp +++ b/include/ck/library/utility/gpu_verification.hpp @@ -5,10 +5,15 @@ #include #include +#include +#include +#include +#include #include "ck/utility/data_type.hpp" #include "ck/utility/type_convert.hpp" #include "ck/utility/type.hpp" +#include "ck/utility/env.hpp" #include "ck/host_utility/device_prop.hpp" #include "ck/host_utility/hip_check_error.hpp" #include "ck/library/utility/check_err.hpp" @@ -106,6 +111,102 @@ inline float compute_relative_tolerance(const int number_of_accumulations = 1) } } +/// @brief Turn an iterator type into an iterator that can be dereferenced. +/// +/// In gpu_verify and gpu_reduce_max, it is valid to pass a void pointer and +/// have the function automatically derive the "concrete" pointer type to +/// be used in the kernel. This function does that: depending on whether +/// the `Iterator` is a void pointer or not, it returns either the iterator +/// (assuming that it is already concrete), or returns the pointer casted +/// to the concrete type. +/// +/// @tparam T The value type of the pointer, when dereferenced. +/// @tparam Iterator The abstract iterator, can be void* or an actual pointer. +/// +/// @param it The iterator to make concrete. +template +__device__ Iterator make_concrete_iterator(Iterator it) +{ + return it; +} + +template +__device__ const T* make_concrete_iterator(const void* it) +{ + return reinterpret_cast(it); +} + +template +__device__ const T* make_concrete_iterator(void* it) +{ + return reinterpret_cast(it); +} + +/// @brief Utility to launch persistent kernels. +/// +/// This function launches a GPU kernel with a grid size derived from the kernel's +/// occupancy and the total number of multiprocessors on the GPU. +/// +/// @tparam Kernel The type of the kernel function. +/// @tparam Args The types of the kernel arguments. +/// +/// @param kernel An instance of the kernel function. This should be a __global__ function. +/// @param block_size The kernel's (1D) block size. +/// @param stream The stream to launch the kernel on. +/// @param args The kernel launch arguments. +template +void launch_persistent_kernel(Kernel kernel, + int block_size, + hipStream_t stream, + const Args&... args) +{ + int occupancy; + hip_check_error( + hipOccupancyMaxActiveBlocksPerMultiprocessor(&occupancy, kernel, block_size, 0)); + + int device; + hip_check_error(hipGetDevice(&device)); + + int multiprocessors; + hip_check_error( + hipDeviceGetAttribute(&multiprocessors, hipDeviceAttributeMultiprocessorCount, device)); + + kernel<<>>(args...); + + hip_check_error(hipGetLastError()); +} + +/// @brief Simple block reduce kernel. +/// +/// This function reduces all `value`s across a block according to `reduce`. This function +/// is a relatively simple implementation as its primary purpose is to be correct and +/// readable: No special cases are done for warp reductions, and the function allocates +/// its own shared memory. The result is broadcasted to all threads. +/// +/// @tparam BlockSize The number of threads in a block. +/// @tparam T The value type to reduce over. +/// @tparam F The reduction functor type. +/// +/// @param value This thread's value to reduce over. +/// @param reduce The reduction functor, used to combine two values. Should be associative. +template +__device__ T block_reduce(const T& value, F reduce) +{ + __shared__ T workspace[BlockSize]; + + workspace[threadIdx.x] = value; + __syncthreads(); + + for(unsigned int s = BlockSize / 2; s >= 1; s >>= 1) + { + if(threadIdx.x < s) + workspace[threadIdx.x] = reduce(workspace[threadIdx.x], workspace[threadIdx.x + s]); + __syncthreads(); + } + + return workspace[0]; +} + // Device-side result structure for kernel output // Packed into a single struct to minimize device memory allocations struct GpuVerifyDeviceResult @@ -113,121 +214,142 @@ struct GpuVerifyDeviceResult unsigned long long error_count; // Number of errors found float max_error; // Maximum error value int all_zero; // 1 = device result is all zeros, 0 = has non-zero values + + /// @brief Return the neutral element of a GpuVerifyDeviceResult + /// + /// This function returns the "neutral element", the element which does nothing + /// when reduced with another with `reduce_results`. Good to be used as an + /// initial value. + __host__ __device__ static GpuVerifyDeviceResult identity() + { + GpuVerifyDeviceResult result; + result.error_count = 0; // No errors yet + result.max_error = 0.0f; // No error observed + result.all_zero = 1; // Start assuming all zeros (will be cleared if nonzero found) + return result; + } }; +/// @brief Combine two device verify results. +/// +/// This function returns the "combined" version of two GpuVerifyDeviceResult values, which +/// adds the total amount of errors, sets the correct max error, and records whether +/// any of the values had any zeros. +__device__ GpuVerifyDeviceResult reduce_results(const GpuVerifyDeviceResult& a, + const GpuVerifyDeviceResult& b) +{ + GpuVerifyDeviceResult result; + result.error_count = a.error_count + b.error_count; + result.max_error = std::max(a.max_error, b.max_error); + result.all_zero = a.all_zero & b.all_zero; + return result; +} + +/// @brief Compare individual tensor elements. +/// +/// This function is what actually does the comparison between two tensor +/// elements. The function returns a tuple of three elements. +/// - The absolute maximum difference. +/// - If the second value is set to false, it indicates either that the elements are not +/// equal according to the thresholds `rtol` and `atol`, or that either value is not +/// finite (NaN/Infinity). If set to true, the values are considered equal. +/// - If the third value is set to true, it indicates that both elements are bitwise +/// equal to zero. +template +__device__ std::tuple +compare_elements(const T& actual, const T& expected, const float rtol, const float atol) +{ + static_assert(!std::is_same_v, "TODO: implement compare_elements() for double"); + + const auto o = type_convert(actual); + const auto r = type_convert(expected); + const auto e = std::abs(o - r); + + const auto inequal = e > atol + rtol * std::abs(r) || !std::isfinite(o) || !std::isfinite(r); + + using Bytes = std::array; + const auto o_bytes = *reinterpret_cast(&actual); + const auto r_bytes = *reinterpret_cast(&expected); + bool all_zero = true; + for(const auto x : o_bytes) + { + if(x != std::byte{0}) + all_zero = false; + } + + for(const auto x : r_bytes) + { + if(x != std::byte{0}) + all_zero = false; + } + + return std::make_tuple(e, inequal, all_zero); +} + // GPU verification kernel - compares device result against reference using relative and absolute // tolerance. Tracks all errors (no early exit) to provide detailed error reporting. // // Uses LDS (shared memory) for block-level reduction to minimize atomic contention. // This reduces atomic operations from O(errors) to O(blocks), providing massive speedup // when there are many errors. -// -// Assumption: Block size is 256 -template -__global__ void gpu_verify_kernel(const T* __restrict__ device_result, - const T* __restrict__ reference_result, - float rtol, - float atol, - long long size, - GpuVerifyDeviceResult* result) +template +__global__ __launch_bounds__(BlockSize) // + void gpu_verify_kernel(IteratorA device_result_it, + IteratorB reference_result_it, + float rtol, + float atol, + long long size, + GpuVerifyDeviceResult* result) { - constexpr int block_size = 256; + auto device_result = make_concrete_iterator(device_result_it); + auto reference_result = make_concrete_iterator(reference_result_it); - // Shared memory for block-level reduction - __shared__ unsigned long long shared_error_count[block_size]; - __shared__ float shared_max_error[block_size]; - __shared__ int shared_has_error[block_size]; - __shared__ int shared_has_nonzero[block_size]; - - // Thread-local accumulators (in registers) - unsigned long long local_error_count = 0; - float local_max_error = 0.0f; - int local_has_error = 0; - int local_has_nonzero = 0; + auto local_result = GpuVerifyDeviceResult::identity(); // Grid-stride loop to handle any tensor size - long long idx = blockIdx.x * blockDim.x + threadIdx.x; - long long stride = blockDim.x * gridDim.x; + long long idx = blockIdx.x * BlockSize + threadIdx.x; + long long stride = BlockSize * gridDim.x; for(long long i = idx; i < size; i += stride) { - // Convert to float for comparison - float dev_val = type_convert(device_result[i]); - float ref_val = type_convert(reference_result[i]); + const auto [abs_diff, inequal, bitwise_zero] = + compare_elements(device_result[i], reference_result[i], rtol, atol); - // Check if device value is non-zero - if(dev_val != 0.0f) - { - local_has_nonzero = 1; - } - - // Compute absolute difference - float abs_diff = fabsf(dev_val - ref_val); - - // Check tolerance (matches CPU check_err logic: err > atol + rtol * abs(ref)) - if(abs_diff > atol + rtol * fabsf(ref_val)) - { - local_has_error = 1; - local_error_count++; - local_max_error = fmaxf(local_max_error, abs_diff); - } + local_result = reduce_results(local_result, + GpuVerifyDeviceResult{ + static_cast(inequal), // error_count + abs_diff, // max_error + bitwise_zero // all_zero + }); } - // Store thread-local results to shared memory - shared_error_count[threadIdx.x] = local_error_count; - shared_max_error[threadIdx.x] = local_max_error; - shared_has_error[threadIdx.x] = local_has_error; - shared_has_nonzero[threadIdx.x] = local_has_nonzero; - __syncthreads(); - - // Block-level reduction: 256 -> 128 -> 64 -> 32 - for(unsigned int s = block_size / 2; s >= 32; s >>= 1) - { - if(threadIdx.x < s) - { - shared_error_count[threadIdx.x] += shared_error_count[threadIdx.x + s]; - shared_max_error[threadIdx.x] = - fmaxf(shared_max_error[threadIdx.x], shared_max_error[threadIdx.x + s]); - shared_has_error[threadIdx.x] |= shared_has_error[threadIdx.x + s]; - shared_has_nonzero[threadIdx.x] |= shared_has_nonzero[threadIdx.x + s]; - } - __syncthreads(); - } + const auto block_result = block_reduce(local_result, reduce_results); // Final reduction of remaining 32 elements in thread 0 if(threadIdx.x == 0) { - for(int i = 1; i < 32; ++i) + // Single atomic update per block (reduces contention from O(errors) to O(blocks)) + if(block_result.error_count > 0) { - shared_error_count[0] += shared_error_count[i]; - shared_max_error[0] = fmaxf(shared_max_error[0], shared_max_error[i]); - shared_has_error[0] |= shared_has_error[i]; - shared_has_nonzero[0] |= shared_has_nonzero[i]; + atomicAdd(&result->error_count, block_result.error_count); + atomicMax(&result->max_error, block_result.max_error); } - // Single atomic update per block (reduces contention from O(errors) to O(blocks)) - if(shared_has_error[0]) + if(!block_result.all_zero) { - atomicAdd(&result->error_count, shared_error_count[0]); - atomicMax(&result->max_error, shared_max_error[0]); - } - // Update all_zero flag: if no nonzero values found, mark as all zero - if(!shared_has_nonzero[0]) - { - atomicMin(&result->all_zero, 1); - } - else - { - atomicMin(&result->all_zero, 0); + // A nonzero was found, so set the global value to false. + // Note: this is a benign race condition; technically a race condition but + // all blocks write the same value, so its fine. + result->all_zero = 0; } } } // Host-side wrapper for GPU verification with explicit tolerances // Returns GpuVerifyResult with detailed error information -template -GpuVerifyResult gpu_verify(const void* device_result, - const void* reference_result, +template +GpuVerifyResult gpu_verify(IteratorA device_result, + IteratorB reference_result, float rtol, float atol, std::size_t size, @@ -238,31 +360,25 @@ GpuVerifyResult gpu_verify(const void* device_result, hip_check_error(hipMalloc(&result_dev, sizeof(GpuVerifyDeviceResult))); // Initialize result struct - GpuVerifyDeviceResult result_host; - result_host.error_count = 0; // No errors yet - result_host.max_error = 0.0f; // No error observed - result_host.all_zero = 1; // Start assuming all zeros (will be cleared if nonzero found) + auto result_host = GpuVerifyDeviceResult::identity(); hip_check_error( hipMemcpy(result_dev, &result_host, sizeof(GpuVerifyDeviceResult), hipMemcpyHostToDevice)); - // Launch kernel with grid-stride loop - // Use 65535 as max grid size (hardware limit for grid dimension in x) - // Grid-stride loop handles any tensor size regardless of grid dimensions + // Launch persistent kernel. + // automatically derive the optimal grid size from the kernel's occupancy and the + // number of multiprocessors. constexpr int block_size = 256; - int grid_size = std::min(65535, (size + block_size - 1) / block_size); + const auto kernel = gpu_verify_kernel; - gpu_verify_kernel - <<>>(static_cast(device_result), - static_cast(reference_result), - rtol, - atol, - static_cast(size), - result_dev); - - hip_check_error(hipGetLastError()); - - // Synchronize the stream to ensure kernel completion before reading results - hip_check_error(hipStreamSynchronize(stream)); + launch_persistent_kernel(kernel, + block_size, + stream, + device_result, + reference_result, + rtol, + atol, + static_cast(size), + result_dev); // Get result hip_check_error( @@ -276,23 +392,25 @@ GpuVerifyResult gpu_verify(const void* device_result, result.error_count = result_host.error_count; result.max_error = result_host.max_error; result.total = size; - result.all_zero = (result_host.all_zero == 1); + result.all_zero = result_host.all_zero == 1; return result; } // Forward declaration of gpu_reduce_max -template -float gpu_reduce_max(const void* device_buffer, std::size_t size, hipStream_t stream = nullptr); +template +float gpu_reduce_max(Iterator device_buffer, std::size_t size, hipStream_t stream = nullptr); // Host-side wrapper for GPU verification with automatic tolerance computation // Computes max value on GPU, then computes tolerances and verifies // Returns GpuVerifyResult with detailed error information template -GpuVerifyResult gpu_verify(const void* device_result, - const void* reference_result, + typename AccDataType = ComputeDataType, + typename IteratorA, + typename IteratorB> +GpuVerifyResult gpu_verify(IteratorA device_result, + IteratorB reference_result, int number_of_accumulations, std::size_t size, hipStream_t stream = nullptr) @@ -323,23 +441,26 @@ GpuVerifyResult gpu_verify(const void* device_result, max_abs_value, number_of_accumulations)); } + if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING))) + { + std::cout << "verify: accumulations=" << number_of_accumulations << " rtol = " << rtol + << " atol=" << atol << std::endl; + } + // Call the explicit tolerance version return gpu_verify(device_result, reference_result, rtol, atol, size, stream); } // GPU reduction kernel for computing max(abs(data)) // This is an internal kernel called only by gpu_reduce_max() wrapper. -// -// Assumption: Block size is 256 -template -__global__ void -gpu_reduce_max_kernel(const T* __restrict__ data, long long size, float* __restrict__ max_val) +template +__global__ __launch_bounds__((BlockSize)) // + void gpu_reduce_max_kernel(Iterator it, long long size, float* __restrict__ max_val) { - constexpr int block_size = 256; - __shared__ float shared_max[block_size]; + auto data = make_concrete_iterator(it); - long long idx = blockIdx.x * blockDim.x + threadIdx.x; - long long stride = blockDim.x * gridDim.x; + long long idx = blockIdx.x * BlockSize + threadIdx.x; + long long stride = BlockSize * gridDim.x; float local_max = 0.0f; @@ -349,37 +470,18 @@ gpu_reduce_max_kernel(const T* __restrict__ data, long long size, float* __restr local_max = fmaxf(local_max, val); } - shared_max[threadIdx.x] = local_max; - __syncthreads(); + const auto block_max = block_reduce( + local_max, [](const auto& a, const auto& b) { return std::max(a, b); }); - // Block-level reduction: 256 -> 128 -> 64 -> 32 - for(unsigned int s = block_size / 2; s >= 32; s >>= 1) - { - if(threadIdx.x < s) - { - shared_max[threadIdx.x] = fmaxf(shared_max[threadIdx.x], shared_max[threadIdx.x + s]); - } - __syncthreads(); - } - - // Final reduction of remaining 32 elements in thread 0 if(threadIdx.x == 0) - { - for(int i = 1; i < 32; ++i) - { - shared_max[0] = fmaxf(shared_max[0], shared_max[i]); - } - - // Single atomic update per block - atomicMax(max_val, shared_max[0]); - } + atomicMax(max_val, block_max); } // Host-side wrapper for GPU max reduction // Computes max(abs(data)) and returns as float // Only transfers 4 bytes (the final max value) instead of entire tensor -template -float gpu_reduce_max(const void* device_buffer, std::size_t size, hipStream_t stream) +template +float gpu_reduce_max(Iterator device_buffer, std::size_t size, hipStream_t stream) { if(size == 0) { @@ -394,22 +496,14 @@ float gpu_reduce_max(const void* device_buffer, std::size_t size, hipStream_t st float init_val = 0.0f; hip_check_error(hipMemcpy(max_dev, &init_val, sizeof(float), hipMemcpyHostToDevice)); - // Launch reduction kernel - // Use 1024 blocks max for reduction to balance occupancy vs. grid-stride iterations - // For very large tensors (>256M elements), grid-stride loop handles the remainder + // Launch persistent kernel. + // automatically derive the optimal grid size from the kernel's occupancy and the + // number of multiprocessors. constexpr int block_size = 256; - int grid_size = std::min(1024, (size + block_size - 1) / block_size); + const auto kernel = gpu_reduce_max_kernel; - gpu_reduce_max_kernel<<>>( - static_cast(device_buffer), static_cast(size), max_dev); - - hip_check_error(hipGetLastError()); - - // Synchronize if using default stream - if(stream == nullptr) - { - hip_check_error(hipDeviceSynchronize()); - } + launch_persistent_kernel( + kernel, block_size, stream, device_buffer, static_cast(size), max_dev); // Copy result to host (only 4 bytes!) float max_host; diff --git a/test/device_memory/test_device_prng.cpp b/test/device_memory/test_device_prng.cpp index 39fa77237d..a871f9a59c 100644 --- a/test/device_memory/test_device_prng.cpp +++ b/test/device_memory/test_device_prng.cpp @@ -11,32 +11,37 @@ #include "ck/utility/common_header.hpp" #include "ck/ck.hpp" -template -void convertTypeFromDevice(std::vector& fromDevice, - std::vector& res, +template +void convertTypeFromDevice(std::vector& fromDevice, + std::vector& res, uint64_t num_elements) { - for(uint64_t i = 0; i < num_elements / ck::packed_size_v; i++) + for(uint64_t i = 0; i < num_elements / ck::packed_size_v; i++) { // since the CPU dosen't have non-standard data types, we need to convert to float - if constexpr(ck::is_same_v, ck::f4x2_pk_t>) + if constexpr(ck::is_same_v, ck::f4x2_pk_t>) { ck::float2_t tmp = ck::type_convert(fromDevice[i]); res[i * 2] = tmp.x; res[i * 2 + 1] = tmp.y; } - else if constexpr(ck::is_same_v, ck::pk_i4_t>) + else if constexpr(ck::is_same_v, ck::pk_i4_t>) { uint8_t packed = fromDevice[i].data; int hi = (packed >> 4) & 0x0f; int lo = packed & 0x0f; - res[i * 2] = static_cast(hi - 8); - res[i * 2 + 1] = static_cast(lo - 8); + res[i * 2] = static_cast(hi - 8); + res[i * 2 + 1] = static_cast(lo - 8); + } + else if constexpr(ck::is_same_v) + { + res[i] = ck::type_convert( + ck::type_convert(fromDevice[i])); } else { - res[i] = ck::type_convert(fromDevice[i]); + res[i] = ck::type_convert(fromDevice[i]); } } } @@ -198,12 +203,13 @@ void TDevRanNormGenFp(double sigma, } TEST(TDevIntegerRanUniGen, U8) { TDevRanUniGenInt(0, 2, 15000); } -TEST(TDevIntegerRanUniGen, U16) { TDevRanUniGenInt(0, 100, 100000); } +// Note: U16 conflicts with ck::bhalf_t TEST(TDevIntegerRanUniGen, U32) { TDevRanUniGenInt(0, 10000, 10000000); } TEST(TDevIntegerRanUniGen, I4) { TDevRanUniGenInt(-2, 2, 10000000); } TEST(TDevIntegerRanUniGen, F32) { TDevRanUniGenInt(-2, 2, 10000000); } TEST(TDevIntegerRanUniGen, F16) { TDevRanUniGenInt(-2, 2, 1000000); } +TEST(TDevIntegerRanUniGen, BF16) { TDevRanUniGenInt(-2, 2, 1000000); } TEST(TDevFpRanUniGen, F32_1) { TDevRanUniGenFp(0, 1, 100000); } TEST(TDevFpRanUniGen, F32_2) { TDevRanUniGenFp(0, 37, 73000); } diff --git a/test/gpu_verification/test_gpu_verification.cpp b/test/gpu_verification/test_gpu_verification.cpp index 3be950c2b6..fcd0efe76f 100644 --- a/test/gpu_verification/test_gpu_verification.cpp +++ b/test/gpu_verification/test_gpu_verification.cpp @@ -83,7 +83,7 @@ class GPUVerificationTest : public ::testing::Test // Use test fixture's RNG (rng_) for reproducibility // RNG is seeded in SetUp() with fixed seed or CK_TEST_SEED environment variable - if constexpr(std::is_integral::value) + if constexpr(std::is_integral_v && !std::is_same_v) { std::uniform_int_distribution dis(static_cast(min_val), static_cast(max_val)); From 654bec3362e825c27f0374e9e4f4e5b970e0f86f Mon Sep 17 00:00:00 2001 From: ltqin Date: Thu, 29 Jan 2026 04:37:15 +0800 Subject: [PATCH 05/11] Fix block scale init value (#3666) * Make blockscale descale range adaptive to data type max value * format --- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index b6287245a0..1227724d40 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -750,9 +750,21 @@ fwd_result fmha_fwd_run(mode_enum mode, } else if(qscale.type == quant_scale_enum::blockscale) { - ck_tile::FillUniformDistribution{0.012f, 0.015f, next_seed()}(q_descale_host); - ck_tile::FillUniformDistribution{0.012f, 0.015f, next_seed()}(k_descale_host); - ck_tile::FillUniformDistribution{0.012f, 0.015f, next_seed()}(v_descale_host); + float q_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float k_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + float v_dtype_max = ck_tile::type_convert(ck_tile::numeric::max()); + + float qkv_max = 3.f; + float max_descale_q = qkv_max / q_dtype_max; + float max_descale_k = qkv_max / k_dtype_max; + float max_descale_v = qkv_max / v_dtype_max; + + ck_tile::FillUniformDistribution{max_descale_q * 0.8f, max_descale_q, next_seed()}( + q_descale_host); + ck_tile::FillUniformDistribution{max_descale_k * 0.8f, max_descale_k, next_seed()}( + k_descale_host); + ck_tile::FillUniformDistribution{max_descale_v * 0.8f, max_descale_v, next_seed()}( + v_descale_host); } iota_shuffle(block_table_host.begin(), block_table_host.end(), 0, random_engine); From 83b58bb0c3ff12f426d45383900a6fd91b4116a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bart=C5=82omiej=20Kocot?= Date: Wed, 28 Jan 2026 22:31:54 +0100 Subject: [PATCH 06/11] Grouped Conv Bwd Weight Direct Load (#3648) * Grouped Conv Bwd Weight Direct Load * Update gridwise_gemm_xdl_cshuffle_conv_v3.hpp * Implement group merging for bwd_weight and add instances * Link direct load instances * builder fixes * fix * fixes * fix --------- Co-authored-by: Graner, Johannes --- .../builder/factory/conv_algorithms.hpp | 2 +- ...rouped_conv_bwd_weight_xdl_cshuffle_v3.hpp | 17 +- ...st_ckb_conv_bwd_weight_xdl_cshuffle_v3.cpp | 3 +- .../test/impl/conv_algorithm_types.hpp | 3 +- ...ance_string_bwd_weight_grp_conv_xdl_v3.cpp | 2 + .../blockwise_gemm_pipeline_xdlops_base.hpp | 7 +- ...lockwise_gemm_pipeline_xdlops_selector.hpp | 11 +- .../blockwise_gemm_pipeline_xdlops_v1.hpp | 18 +- ...roup_tensor_slice_transfer_direct_load.hpp | 4 - ...rouped_conv_bwd_weight_xdl_cshuffle_v3.hpp | 126 ++++-- ..._conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp | 1 + .../gridwise_gemm_xdl_cshuffle_conv_v3.hpp | 417 ++++++++++++------ ...rouped_conv_bwd_weight_v3_xdl_instance.hpp | 49 ++ .../grouped_convolution_backward_weight.hpp | 6 + ...rouped_convolution_backward_weight_xdl.inc | 24 + .../grouped_conv2d_bwd_weight/CMakeLists.txt | 2 + ...xdl_nhwgc_gkyxc_nhwgk_bf16_direct_load.cpp | 40 ++ ..._xdl_nhwgc_gkyxc_nhwgk_f16_direct_load.cpp | 40 ++ 18 files changed, 578 insertions(+), 194 deletions(-) create mode 100644 library/src/tensor_operation_instance/gpu/grouped_conv2d_bwd_weight/xdl/nhwgc_gkyxc_nhwgk/device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_bf16_direct_load.cpp create mode 100644 library/src/tensor_operation_instance/gpu/grouped_conv2d_bwd_weight/xdl/nhwgc_gkyxc_nhwgk/device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_direct_load.cpp diff --git a/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp b/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp index 79b818555e..c508126adb 100644 --- a/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp +++ b/experimental/builder/include/ck_tile/builder/factory/conv_algorithms.hpp @@ -35,7 +35,7 @@ template concept BwdXdlV3AlgorithmBase = ConvAlgorithmDescriptor && SpecifiesThreadBlock && SpecifiesTileTransferParameters3D && SpecifiesGridwiseBwdXdlGemm && SpecifiesBwdWeightConvSpecialization && - SpecifiesBlockGemm; + SpecifiesBlockGemm && SpecifiesNumGroupsToMerge; template concept BwdWmmaAlgorithmBase = diff --git a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp index ce23dac1d7..249071c195 100644 --- a/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp +++ b/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -53,7 +53,9 @@ template + typename ComputeTypeB, + bool DirectLoad, + index_t NumGroupsToMerge> struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3; } // namespace ck::tensor_operation::device @@ -109,7 +111,9 @@ template + typename ComputeTypeB_, + bool DirectLoad, + index_t NumGroupsToMerge> struct InstanceTraits> + ComputeTypeB_, + DirectLoad, + NumGroupsToMerge>> { /// @brief Tag type identifying this device kernel variant @@ -241,6 +247,9 @@ struct InstanceTraits(); // 42. oss << "," << detail::type_name(); // 43. + oss << "," << kDirectLoad; // 44. + oss << "," << kNumGroupsToMerge; // 45. oss << ">"; return oss.str(); diff --git a/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_xdl_cshuffle_v3.cpp b/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_xdl_cshuffle_v3.cpp index a3f4a988ef..80c6646211 100644 --- a/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_xdl_cshuffle_v3.cpp +++ b/experimental/builder/test/conv/ck/test_ckb_conv_bwd_weight_xdl_cshuffle_v3.cpp @@ -32,7 +32,8 @@ constexpr auto ALGORITHM = .with_gemm_config(cku::BwdGemmParams_Xdl_1x1_per_wave) .with_transfer(cku::BwdTransfer_4x8x1_4x16x1_v3) .with_bwd_specialization(ckb::ConvSpecialization::FILTER_1X1_STRIDE1_PAD0) - .with_block_gemm(cku::BlockGemmDesc_v2_intrawave); + .with_block_gemm(cku::BlockGemmDesc_v2_intrawave) + .with_num_conv_groups_to_merge(1); using Builder = ckb::ConvBuilder; using Instance = Builder::Instance; diff --git a/experimental/builder/test/impl/conv_algorithm_types.hpp b/experimental/builder/test/impl/conv_algorithm_types.hpp index b775505a26..f5b9bdc3b5 100644 --- a/experimental/builder/test/impl/conv_algorithm_types.hpp +++ b/experimental/builder/test/impl/conv_algorithm_types.hpp @@ -632,7 +632,8 @@ using ConvAlgorithm_DeviceGroupedConvBwdWeight_Xdl_CShuffle_V3 = BwdXdlGemm_, Transfer_<>, ConvSpecializationBwdWeight_, - BlockGemm_>; + BlockGemm_, + GemmBatchOptions_>; using ConvAlgorithm_DeviceGroupedConvBwdWeight_Dl = ConvAlgorithmTemplate"; // Test describe() through base class pointer for XDL V3 variant diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp index 512a019ec8..ffd728f259 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_base.hpp @@ -30,7 +30,8 @@ template + bool TransposeC = false, + bool LdsScalarLoadToVgpr = false> struct BlockwiseGemmXdlops_pipeline_base { static constexpr auto I0 = Number<0>{}; @@ -385,7 +386,7 @@ struct BlockwiseGemmXdlops_pipeline_base Sequence<1, 1, 1, KPack>, Sequence<0, 1, 2, 3>, 3, - A_K1, + LdsScalarLoadToVgpr ? 1 : A_K1, A_K1>; using BThreadCopy = ThreadwiseTensorSliceTransfer_v4, Sequence<0, 1, 2, 3>, 3, - B_K1, + LdsScalarLoadToVgpr ? 1 : B_K1, B_K1>; AThreadCopy a_thread_copy_; diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_selector.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_selector.hpp index 7ddb4e9b74..461ca513f9 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_selector.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_selector.hpp @@ -32,9 +32,15 @@ template + bool DirectLoad = false, + bool LdsScalarLoadToVgpr = false> constexpr auto BlockGemmPipeline_Selector() { + // Supported for Direct Load and V1 + if constexpr(LdsScalarLoadToVgpr) + { + static_assert(DirectLoad && BlkGemmPipelineVer == BlockGemmPipelineVersion::v1); + } if constexpr(DirectLoad) { if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1) @@ -58,7 +64,8 @@ constexpr auto BlockGemmPipeline_Selector() NPerXDL, MRepeat, NRepeat, - KPack>{}; + KPack, + LdsScalarLoadToVgpr>{}; } else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v4) { diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1.hpp index 5604a31091..ae4504d6ba 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1.hpp @@ -758,7 +758,8 @@ template + index_t KPacks, + bool LdsScalarLoadToVgpr = false> struct BlockwiseGemmXdlopsDirectLoad_pipeline_v1 { }; @@ -781,9 +782,9 @@ template + bool LdsScalarLoadToVgpr> struct BlockwiseGemmXdlopsDirectLoad_pipeline_v1 + KPack, + LdsScalarLoadToVgpr> : BlockwiseGemmXdlops_pipeline_base + KPack, + false /*TransposeC*/, + LdsScalarLoadToVgpr> { using Base = BlockwiseGemmXdlops_pipeline_base; + KPack, + false /*TransposeC*/, + LdsScalarLoadToVgpr>; using Base::I0; using Base::KRepeat; using Base::xdlops_gemm; diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp index ad74ee847e..a31c9101a1 100644 --- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp +++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp @@ -140,10 +140,6 @@ struct ThreadGroupTensorSliceTransfer_DirectLoad "Direct load transfer does not support datatypes conversion. Source and " "destination data types must be the same."); - static_assert( - DstVectorDim == nDim - 1, - "Direct load transfer requires the destination vector dimension to be the last one."); - static_assert(ScalarPerVector == 1 || SrcVectorDim == DstVectorDim, "When loading more than one element per thread at once, the contiguous " "dimension must be the same between source and destination."); diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp index 175b4625ba..26cf586017 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -82,23 +82,48 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy) __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte(get_device_arch())]; - DispatchSplitKHack(karg.p_a_grid + a_batch_offset + split_k_offset_a, - karg.p_b_grid + b_batch_offset + split_k_offset_b, - karg.p_c_grid + e_batch_offset, - p_shared, - karg, - a_grid_desc_ak0_m_ak1, - b_grid_desc_bk0_n_bk1, - c_grid_desc_mblock_mperblock_nblock_nperblock, - k_idx * num_k_per_block, - gridDim.y, - split_k_offset_hack); + if constexpr(GridwiseGemm::DirectLoadEnabled) + { +#if defined(__gfx950__) + DispatchSplitKHack(karg.p_a_grid + a_batch_offset + split_k_offset_a, + karg.p_b_grid + b_batch_offset + split_k_offset_b, + karg.p_c_grid + e_batch_offset, + p_shared, + karg, + a_grid_desc_ak0_m_ak1, + b_grid_desc_bk0_n_bk1, + c_grid_desc_mblock_mperblock_nblock_nperblock, + k_idx * num_k_per_block, + gridDim.y, + split_k_offset_hack); +#endif + } + else + { + DispatchSplitKHack(karg.p_a_grid + a_batch_offset + split_k_offset_a, + karg.p_b_grid + b_batch_offset + split_k_offset_b, + karg.p_c_grid + e_batch_offset, + p_shared, + karg, + a_grid_desc_ak0_m_ak1, + b_grid_desc_bk0_n_bk1, + c_grid_desc_mblock_mperblock_nblock_nperblock, + k_idx * num_k_per_block, + gridDim.y, + split_k_offset_hack); + } } #else ignore = karg; @@ -236,7 +261,9 @@ template + typename ComputeTypeB = ComputeTypeA, + bool DirectLoad = false, + index_t NumGroupsToMerge = 1> struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 : public DeviceGroupedConvBwdWeight{}; template ::type = false> @@ -371,6 +398,16 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 using BGridDesc_K0_N_K1 = remove_cvref_t; using CGridDesc_M_N = remove_cvref_t; + // Disable vector load = 4. It is not supported for Direct Load. Align to 2 in such case. + static constexpr index_t ABlockTransferSrcScalarPerVectorAligned = + ABlockTransferSrcScalarPerVector * sizeof(ADataType) == 8 + ? 4 / sizeof(ADataType) + : ABlockTransferSrcScalarPerVector; + static constexpr index_t BBlockTransferSrcScalarPerVectorAligned = + BBlockTransferSrcScalarPerVector * sizeof(BDataType) == 8 + ? 4 / sizeof(BDataType) + : BBlockTransferSrcScalarPerVector; + template using GridwiseGemmBase = GridwiseGemm_xdl_cshuffle_conv_v3< tensor_layout::gemm::RowMajor, @@ -399,7 +436,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, - ABlockTransferSrcScalarPerVector, + DirectLoad ? ABlockTransferSrcScalarPerVectorAligned : ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_K1, false, ABlockLdsAddExtraM, @@ -407,7 +444,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, - BBlockTransferSrcScalarPerVector, + DirectLoad ? BBlockTransferSrcScalarPerVectorAligned : BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_K1, false, BBlockLdsAddExtraN, @@ -418,7 +455,8 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 BlkGemmPipeSched, BlkGemmPipelineVer, ComputeTypeA, - ComputeTypeB>; + ComputeTypeB, + DirectLoad>; using GridwiseGemm64 = GridwiseGemmBase; using GridwiseGemm32 = GridwiseGemmBase; @@ -653,15 +691,16 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 if(split_k_offset_hack_) split_k_stride_b_ /= k_batch_; - // A/B/C Batch Stride - compute_ptr_offset_of_batch_.BatchStrideA_ = a_g_n_k_wos_strides[0]; - compute_ptr_offset_of_batch_.BatchStrideB_ = b_g_n_c_wis_strides[0]; + // A/B/C Batch Stride (multiply by NumGroupsToMerge for group merging) + compute_ptr_offset_of_batch_.BatchStrideA_ = a_g_n_k_wos_strides[0] * NumGroupsToMerge; + compute_ptr_offset_of_batch_.BatchStrideB_ = b_g_n_c_wis_strides[0] * NumGroupsToMerge; compute_ptr_offset_of_batch_.BatchStrideC_ = Conv_K_ * Conv_C_ * std::accumulate(begin(filter_spatial_lengths_), end(filter_spatial_lengths_), index_t{1}, - std::multiplies<>{}); + std::multiplies<>{}) * + NumGroupsToMerge; const index_t GemmM = a_grid_desc_k0_m_k1_.GetLength(I1); const index_t GemmN = b_grid_desc_k0_n_k1_.GetLength(I1); @@ -743,7 +782,7 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 index_t gdx, gdy, gdz; std::tie(gdx, gdy, gdz) = GridwiseGemm::CalculateGridSize( - gemm_arg.M, gemm_arg.N, gemm_arg.KBatch, arg.Conv_G_); + gemm_arg.M, gemm_arg.N, gemm_arg.KBatch, arg.Conv_G_ / NumGroupsToMerge); float ave_time = 0; @@ -1367,6 +1406,30 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 } #endif + // check device + if constexpr(DirectLoad) + { + if(get_device_name() != "gfx950") + { + return false; + } + } + + // Check that NumGroupsToMerge divides Conv_G evenly + if constexpr(NumGroupsToMerge > 1) + { + if(arg.Conv_G_ % NumGroupsToMerge != 0) + { + if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING))) + { + std::cout << "Unsupported! Conv_G_ % NumGroupsToMerge != 0: Conv_G_=" + << arg.Conv_G_ << ", NumGroupsToMerge=" << NumGroupsToMerge + << std::endl; + } + return false; + } + } + const index_t GemmM = arg.a_grid_desc_k0_m_k1_.GetLength(I1); const index_t GemmN = arg.b_grid_desc_k0_n_k1_.GetLength(I1); const index_t GemmK = @@ -1617,8 +1680,13 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 auto str = std::stringstream(); // clang-format off - str << "DeviceGroupedConvBwdWeight_Xdl_CShuffleV3" - << "<" + str << "DeviceGroupedConvBwdWeight_Xdl_CShuffleV3"; + + if constexpr(DirectLoad) { + str << "_DirectLoad"; + } + + str << "<" << BlockSize << ", " << MPerBlock << ", " << NPerBlock << ", " diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp index d8cb5f4a8c..15a5e08803 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp @@ -567,6 +567,7 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3 using DsGridDesc_M_N = remove_cvref_t; + // Disable vector load = 4. It is not supported for Direct Load. Align to 2 in such case. static constexpr index_t ABlockTransferSrcScalarPerVectorAligned = ABlockTransferSrcScalarPerVector * sizeof(ADataType) == 8 ? 4 / sizeof(ADataType) diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_conv_v3.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_conv_v3.hpp index cfbfaf3262..5289e209fb 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_conv_v3.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_conv_v3.hpp @@ -12,6 +12,7 @@ #include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp" #include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" +#include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp" #include "ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_common.hpp" namespace ck { @@ -61,7 +62,8 @@ template + typename ComputeTypeB = ComputeTypeA, + bool DirectLoad = false> struct GridwiseGemm_xdl_cshuffle_conv_v3 : public GridwiseGemm_xdl_cshuffle_base< ALayout, @@ -109,6 +111,10 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 ComputeTypeB, false> // ForceNaiveLayout { + static_assert((is_same_v && + is_same_v) || + !DirectLoad); + using Base = GridwiseGemm_xdl_cshuffle_base< ALayout, BLayout, @@ -164,6 +170,8 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 using Base::I2; using ThisThreadBlock = typename Base::ThisThreadBlock; + static constexpr bool DirectLoadEnabled = DirectLoad; + static constexpr auto lcm_AK1_BK1 = math::lcm(AK1Number, BK1Number); static constexpr bool is_single_rate_mfma = (((is_same::value || is_same::value) && @@ -353,7 +361,13 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 template __device__ static constexpr auto GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(DeviceArch) { - if constexpr(is_same_v) + if constexpr(DirectLoad) + { + return make_naive_tensor_descriptor( + make_tuple(AK0Number, Number{}, AK1Number), + make_tuple(Number{}, I1, Number{})); + } + else if constexpr(is_same_v) { // Force use padded layout on gfx950 to reduce bank conflicts constexpr index_t ABlockLdsExtraM = 1; @@ -370,7 +384,13 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 template __device__ static constexpr auto GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(DeviceArch) { - if constexpr(is_same_v) + if constexpr(DirectLoad) + { + return make_naive_tensor_descriptor( + make_tuple(BK0Number, Number{}, BK1Number), + make_tuple(Number{}, I1, Number{})); + } + else if constexpr(is_same_v) { constexpr index_t BBlockLdsExtraN = 1; return make_naive_tensor_descriptor( @@ -385,31 +405,36 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 IS_VALID_COMPILATION_PARAMETER_IMPL(CDataType) - using BlockwiseGemmPipe = remove_cvref_t< - decltype(BlockGemmPipeline_Selector< - BlkGemmPipelineVer, - BlkGemmPipeSched, - BlockSize, - ADataType, - BDataType, - ComputeTypeA, - AccDataType, - decltype(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(get_device_arch())), - decltype(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch())), - decltype(MakeAMmaTileDescriptor_M0_M1_M2_K( + // Disable vector load from lds to vgpr for direct load (backward weight store with continous M + // or N dimension) + static constexpr bool LdsScalarLoadToVgpr = DirectLoad; + using BlockwiseGemmPipe = remove_cvref_t< + decltype(BlockGemmPipeline_Selector< + BlkGemmPipelineVer, + BlkGemmPipeSched, + BlockSize, + ADataType, + BDataType, + ComputeTypeA, + AccDataType, + decltype(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(get_device_arch())), + decltype(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch())), + decltype(MakeAMmaTileDescriptor_M0_M1_M2_K( GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1(get_device_arch()))), - decltype(MakeBMmaTileDescriptor_N0_N1_N2_K( + decltype(MakeBMmaTileDescriptor_N0_N1_N2_K( GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch()))), - ABlockTransferSrcScalarPerVector, - BBlockTransferSrcScalarPerVector, - MPerBlock, - NPerBlock, - KPerBlock, - MPerXdl, - NPerXdl, - MXdlPerWave, - NXdlPerWave, - KPack>())>; + ABlockTransferSrcScalarPerVector, + BBlockTransferSrcScalarPerVector, + MPerBlock, + NPerBlock, + KPerBlock, + MPerXdl, + NPerXdl, + MXdlPerWave, + NXdlPerWave, + KPack, + DirectLoad, + LdsScalarLoadToVgpr>())>; template __device__ static constexpr index_t GetSharedMemoryNumberOfByte(DeviceArch) @@ -539,67 +564,119 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch()); - // A matrix blockwise copy - auto a_blockwise_copy = - ThreadGroupTensorSliceTransfer_v4r1, - ABlockTransferThreadClusterLengths_AK0_M_AK1, - ABlockTransferThreadClusterArrangeOrder, - ADataType, - ADataType, - decltype(a_grid_desc_ak0_m_ak1), - decltype(a_block_desc_ak0_m_ak1), - ABlockTransferSrcAccessOrder, - Sequence<0, 1, 2>, - ABlockTransferSrcVectorDim, - 2, - ABlockTransferSrcScalarPerVector, - ABlockTransferDstScalarPerVector_AK1, - 1, - 1, - AThreadTransferSrcResetCoordinateAfterRun, - true, - BlockwiseGemmPipe::GlobalBufferNum>( - a_grid_desc_ak0_m_ak1, - make_multi_index(SplitKOffsetHack ? 0 : k_id, m_block_data_idx_on_grid, 0), - a_element_op, - a_block_desc_ak0_m_ak1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}); + auto get_a_blockwise_copy = [&]() { + if constexpr(DirectLoad) + { + return ThreadGroupTensorSliceTransfer_DirectLoad< + ThisThreadBlock, + Sequence, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ADataType, + ADataType, + decltype(a_grid_desc_ak0_m_ak1), + decltype(a_block_desc_ak0_m_ak1), + ABlockTransferSrcAccessOrder, + ABlockTransferSrcVectorDim, + 1, + ABlockTransferSrcScalarPerVector>( + a_grid_desc_ak0_m_ak1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, m_block_data_idx_on_grid, 0), + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0)); + } + else + { + return ThreadGroupTensorSliceTransfer_v4r1< + ThisThreadBlock, + AElementwiseOperation, + ck::tensor_operation::element_wise::PassThrough, + InMemoryDataOperationEnum::Set, + Sequence, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ADataType, + ADataType, + decltype(a_grid_desc_ak0_m_ak1), + decltype(a_block_desc_ak0_m_ak1), + ABlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + ABlockTransferSrcVectorDim, + 2, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + 1, + 1, + AThreadTransferSrcResetCoordinateAfterRun, + true, + BlockwiseGemmPipe::GlobalBufferNum>( + a_grid_desc_ak0_m_ak1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, m_block_data_idx_on_grid, 0), + a_element_op, + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + } + }; // B matrix blockwise copy - auto b_blockwise_copy = - ThreadGroupTensorSliceTransfer_v4r1, - BBlockTransferThreadClusterLengths_BK0_N_BK1, - BBlockTransferThreadClusterArrangeOrder, - BDataType, - BDataType, - decltype(b_grid_desc_bk0_n_bk1), - decltype(b_block_desc_bk0_n_bk1), - BBlockTransferSrcAccessOrder, - Sequence<0, 1, 2>, - BBlockTransferSrcVectorDim, - 2, - BBlockTransferSrcScalarPerVector, - BBlockTransferDstScalarPerVector_BK1, - 1, - 1, - BThreadTransferSrcResetCoordinateAfterRun, - true, - BlockwiseGemmPipe::GlobalBufferNum>( - b_grid_desc_bk0_n_bk1, - make_multi_index(SplitKOffsetHack ? 0 : k_id, n_block_data_idx_on_grid, 0), - b_element_op, - b_block_desc_bk0_n_bk1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}); + auto get_b_blockwise_copy = [&]() { + if constexpr(DirectLoad) + { + return ThreadGroupTensorSliceTransfer_DirectLoad< + ThisThreadBlock, + Sequence, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BDataType, + BDataType, + decltype(b_grid_desc_bk0_n_bk1), + decltype(b_block_desc_bk0_n_bk1), + BBlockTransferSrcAccessOrder, + BBlockTransferSrcVectorDim, + 1, + BBlockTransferSrcScalarPerVector>( + b_grid_desc_bk0_n_bk1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, n_block_data_idx_on_grid, 0), + b_block_desc_bk0_n_bk1, + make_multi_index(0, 0, 0)); + } + else + { + return ThreadGroupTensorSliceTransfer_v4r1< + ThisThreadBlock, + BElementwiseOperation, + ck::tensor_operation::element_wise::PassThrough, + InMemoryDataOperationEnum::Set, + Sequence, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BDataType, + BDataType, + decltype(b_grid_desc_bk0_n_bk1), + decltype(b_block_desc_bk0_n_bk1), + BBlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + BBlockTransferSrcVectorDim, + 2, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + 1, + 1, + BThreadTransferSrcResetCoordinateAfterRun, + true, + BlockwiseGemmPipe::GlobalBufferNum>( + b_grid_desc_bk0_n_bk1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, n_block_data_idx_on_grid, 0), + b_element_op, + b_block_desc_bk0_n_bk1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + } + }; + + auto a_blockwise_copy = get_a_blockwise_copy(); + auto b_blockwise_copy = get_b_blockwise_copy(); // LDS allocation for A and B: be careful of alignment constexpr auto a_block_space_size_aligned = math::integer_least_multiple( @@ -722,67 +799,119 @@ struct GridwiseGemm_xdl_cshuffle_conv_v3 constexpr auto b_block_desc_bk0_n_bk1 = GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1(get_device_arch()); - // A matrix blockwise copy - auto a_blockwise_copy = - ThreadGroupTensorSliceTransfer_v4r1, - ABlockTransferThreadClusterLengths_AK0_M_AK1, - ABlockTransferThreadClusterArrangeOrder, - ADataType, - ADataType, - decltype(a_grid_desc_ak0_m_ak1), - decltype(a_block_desc_ak0_m_ak1), - ABlockTransferSrcAccessOrder, - Sequence<0, 1, 2>, - ABlockTransferSrcVectorDim, - 2, - ABlockTransferSrcScalarPerVector, - ABlockTransferDstScalarPerVector_AK1, - 1, - 1, - AThreadTransferSrcResetCoordinateAfterRun, - true, - BlockwiseGemmPipe::GlobalBufferNum>( - a_grid_desc_ak0_m_ak1, - make_multi_index(SplitKOffsetHack ? 0 : k_id, m_block_data_idx_on_grid, 0), - a_element_op, - a_block_desc_ak0_m_ak1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}); + auto get_a_blockwise_copy = [&]() { + if constexpr(DirectLoad) + { + return ThreadGroupTensorSliceTransfer_DirectLoad< + ThisThreadBlock, + Sequence, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ADataType, + ADataType, + decltype(a_grid_desc_ak0_m_ak1), + decltype(a_block_desc_ak0_m_ak1), + ABlockTransferSrcAccessOrder, + ABlockTransferSrcVectorDim, + 1, + ABlockTransferSrcScalarPerVector>( + a_grid_desc_ak0_m_ak1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, m_block_data_idx_on_grid, 0), + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0)); + } + else + { + return ThreadGroupTensorSliceTransfer_v4r1< + ThisThreadBlock, + AElementwiseOperation, + ck::tensor_operation::element_wise::PassThrough, + InMemoryDataOperationEnum::Set, + Sequence, + ABlockTransferThreadClusterLengths_AK0_M_AK1, + ABlockTransferThreadClusterArrangeOrder, + ADataType, + ADataType, + decltype(a_grid_desc_ak0_m_ak1), + decltype(a_block_desc_ak0_m_ak1), + ABlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + ABlockTransferSrcVectorDim, + 2, + ABlockTransferSrcScalarPerVector, + ABlockTransferDstScalarPerVector_AK1, + 1, + 1, + AThreadTransferSrcResetCoordinateAfterRun, + true, + BlockwiseGemmPipe::GlobalBufferNum>( + a_grid_desc_ak0_m_ak1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, m_block_data_idx_on_grid, 0), + a_element_op, + a_block_desc_ak0_m_ak1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + } + }; // B matrix blockwise copy - auto b_blockwise_copy = - ThreadGroupTensorSliceTransfer_v4r1, - BBlockTransferThreadClusterLengths_BK0_N_BK1, - BBlockTransferThreadClusterArrangeOrder, - BDataType, - BDataType, - decltype(b_grid_desc_bk0_n_bk1), - decltype(b_block_desc_bk0_n_bk1), - BBlockTransferSrcAccessOrder, - Sequence<0, 1, 2>, - BBlockTransferSrcVectorDim, - 2, - BBlockTransferSrcScalarPerVector, - BBlockTransferDstScalarPerVector_BK1, - 1, - 1, - BThreadTransferSrcResetCoordinateAfterRun, - true, - BlockwiseGemmPipe::GlobalBufferNum>( - b_grid_desc_bk0_n_bk1, - make_multi_index(SplitKOffsetHack ? 0 : k_id, n_block_data_idx_on_grid, 0), - b_element_op, - b_block_desc_bk0_n_bk1, - make_multi_index(0, 0, 0), - ck::tensor_operation::element_wise::PassThrough{}); + auto get_b_blockwise_copy = [&]() { + if constexpr(DirectLoad) + { + return ThreadGroupTensorSliceTransfer_DirectLoad< + ThisThreadBlock, + Sequence, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BDataType, + BDataType, + decltype(b_grid_desc_bk0_n_bk1), + decltype(b_block_desc_bk0_n_bk1), + BBlockTransferSrcAccessOrder, + BBlockTransferSrcVectorDim, + 1, + BBlockTransferSrcScalarPerVector>( + b_grid_desc_bk0_n_bk1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, n_block_data_idx_on_grid, 0), + b_block_desc_bk0_n_bk1, + make_multi_index(0, 0, 0)); + } + else + { + return ThreadGroupTensorSliceTransfer_v4r1< + ThisThreadBlock, + BElementwiseOperation, + ck::tensor_operation::element_wise::PassThrough, + InMemoryDataOperationEnum::Set, + Sequence, + BBlockTransferThreadClusterLengths_BK0_N_BK1, + BBlockTransferThreadClusterArrangeOrder, + BDataType, + BDataType, + decltype(b_grid_desc_bk0_n_bk1), + decltype(b_block_desc_bk0_n_bk1), + BBlockTransferSrcAccessOrder, + Sequence<0, 1, 2>, + BBlockTransferSrcVectorDim, + 2, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_BK1, + 1, + 1, + BThreadTransferSrcResetCoordinateAfterRun, + true, + BlockwiseGemmPipe::GlobalBufferNum>( + b_grid_desc_bk0_n_bk1, + make_multi_index(SplitKOffsetHack ? 0 : k_id, n_block_data_idx_on_grid, 0), + b_element_op, + b_block_desc_bk0_n_bk1, + make_multi_index(0, 0, 0), + ck::tensor_operation::element_wise::PassThrough{}); + } + }; + + auto a_blockwise_copy = get_a_blockwise_copy(); + auto b_blockwise_copy = get_b_blockwise_copy(); // LDS allocation for A and B: be careful of alignment constexpr auto a_block_space_size_aligned = math::integer_least_multiple( diff --git a/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_xdl_instance.hpp b/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_xdl_instance.hpp index 143d857333..3a3dc156ec 100644 --- a/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_xdl_instance.hpp +++ b/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_xdl_instance.hpp @@ -101,6 +101,55 @@ using device_grouped_conv_bwd_weight_v3_xdl_c_shuffle_f16_instances = std::tuple // clang-format on >; +template +using device_grouped_conv_bwd_weight_v3_xdl_c_shuffle_f16_direct_load_instances = std::tuple< + // clang-format off + //#########################################| Num| InLayout| WeiLayout| OutLayout| InData| WeiData| OutData| AccData| In| Wei| Out| ConvBackward| Block| MPer| NPer| K0Per| K1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransfer| CBlockTransfer| BlockGemm| BlockGemm| Compute| Compute| Direct| + //#########################################| Dim| | | | Type| Type| Type| Type| Elementwise| Elementwise| Elementwise| Weight| Size| Block| Block| Block| | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| ClusterLengths| ScalarPerVector| Pipeline| Pipeline| Data| Data| Load| + //#########################################| Spatial| | | | | | | | Operation| Operation| Operation| Specialization| | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| MBlock_MPerBlock| NWaveNPerXdl| Scheduler| Version| Type| Type| | + //#########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | NBlock_NPerBlock| | | | | | | + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 128, 16, 32, 64, 8, 16, 16, 1, 1, S<8, 2, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 4, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, Scheduler, PipelineVersion, F16, F16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 128, 16, 32, 64, 8, 16, 16, 1, 1, S<8, 2, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 4, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, Scheduler, PipelineVersion, F16, F16, true, 2>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 128, 16, 64, 64, 8, 16, 16, 1, 2, S<8, 2, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<2, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, Scheduler, PipelineVersion, F16, F16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 128, 16, 64, 64, 8, 16, 16, 1, 2, S<8, 2, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<2, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, Scheduler, PipelineVersion, F16, F16, true, 2>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 256, 32, 64, 8, 32, 32, 2, 1, S<4, 32, 2>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<8, 4, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 64, 1, 4>, 8, Scheduler, PipelineVersion, F16, F16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 128, 64, 8, 32, 32, 1, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 16, 4>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion, F16, F16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 128, 64, 8, 32, 32, 1, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 16, 4>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion, F16, F16, true, 2>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 4>, 8, Scheduler, PipelineVersion, F16, F16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 4>, 2, Scheduler, PipelineVersion, F16, F16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 32, 2>, S<0, 2, 1>,S<0, 2, 1>, 1, 2, 1, 0, S<4, 32, 2>, S<0, 2, 1>,S<0, 2, 1>, 1, 2, 1, 0, 1, 1, S<1, 32, 1, 4>, 4, Scheduler, PipelineVersion, F16, F16, true> + // clang-format on + >; + +template +using device_grouped_conv_bwd_weight_v3_xdl_c_shuffle_bf16_direct_load_instances = std::tuple< + // clang-format off + //#########################################| Num| InLayout| WeiLayout| OutLayout| InData| WeiData| OutData| AccData| In| Wei| Out| ConvBackward| Block| MPer| NPer| K0Per| K1| MPer| NPer| MXdl| NXdl| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransfer| CBlockTransfer| BlockGemm| BlockGemm| Compute| Compute| Direct| + //#########################################| Dim| | | | Type| Type| Type| Type| Elementwise| Elementwise| Elementwise| Weight| Size| Block| Block| Block| | XDL| XDL| Per| Per| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| ClusterLengths| ScalarPerVector| Pipeline| Pipeline| Data| Data| Load| + //#########################################| Spatial| | | | | | | | Operation| Operation| Operation| Specialization| | | | | | | | Wave| Wave| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| MBlock_MPerBlock| NWaveNPerXdl| Scheduler| Version| Type| Type| | + //#########################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | NBlock_NPerBlock| | | | | | | + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 128, 16, 32, 64, 8, 16, 16, 1, 1, S<8, 2, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 4, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, Scheduler, PipelineVersion, BF16, BF16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 128, 16, 64, 64, 8, 16, 16, 1, 2, S<8, 2, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<2, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 8>, 4, Scheduler, PipelineVersion, BF16, BF16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 256, 32, 64, 8, 32, 32, 2, 1, S<4, 32, 2>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<8, 4, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 64, 1, 4>, 8, Scheduler, PipelineVersion, BF16, BF16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 128, 64, 8, 32, 32, 1, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 16, 4>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion, BF16, BF16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 4>, 8, Scheduler, PipelineVersion, BF16, BF16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, S<4, 8, 8>, S<0, 2, 1>,S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 4>, 2, Scheduler, PipelineVersion, BF16, BF16, true>, + DeviceGroupedConvBwdWeight_Xdl_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 64, 64, 8, 16, 16, 2, 2, S<4, 32, 2>, S<0, 2, 1>,S<0, 2, 1>, 1, 2, 1, 0, S<4, 32, 2>, S<0, 2, 1>,S<0, 2, 1>, 1, 2, 1, 0, 1, 1, S<1, 32, 1, 4>, 4, Scheduler, PipelineVersion, BF16, BF16, true> + // clang-format on + >; + template >>& instances); +void add_device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_bf16_direct_load_instances( + std::vector>>& instances); + void add_device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_bf16_default_pipev5_instances( std::vector>>& instances); +void add_device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_direct_load_instances( + std::vector>>& instances); + void add_device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_default_pipev5_instances( std::vector>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_conv_bwd_weight_v3_xdl_c_shuffle_bf16_direct_load_instances< + 2, + NHWGC, + GKYXC, + NHWGK, + ConvBwdWeightDefault, + BlockGemmPipelineScheduler::Intrawave, + BlockGemmPipelineVersion::v1>{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck diff --git a/library/src/tensor_operation_instance/gpu/grouped_conv2d_bwd_weight/xdl/nhwgc_gkyxc_nhwgk/device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_direct_load.cpp b/library/src/tensor_operation_instance/gpu/grouped_conv2d_bwd_weight/xdl/nhwgc_gkyxc_nhwgk/device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_direct_load.cpp new file mode 100644 index 0000000000..7035a4cf71 --- /dev/null +++ b/library/src/tensor_operation_instance/gpu/grouped_conv2d_bwd_weight/xdl/nhwgc_gkyxc_nhwgk/device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_direct_load.cpp @@ -0,0 +1,40 @@ +// Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +// SPDX-License-Identifier: MIT + +#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" +#include "ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_xdl_instance.hpp" + +namespace ck { +namespace tensor_operation { +namespace device { +namespace instance { + +// Compilation parameters for in[n, hi, wi, g, c] * wei[g, k, y, x, c] = out[n, ho, wo, g, k] +void add_device_grouped_conv2d_bwd_weight_xdl_nhwgc_gkyxc_nhwgk_f16_direct_load_instances( + std::vector>>& instances) +{ + add_device_operation_instances( + instances, + device_grouped_conv_bwd_weight_v3_xdl_c_shuffle_f16_direct_load_instances< + 2, + NHWGC, + GKYXC, + NHWGK, + ConvBwdWeightDefault, + BlockGemmPipelineScheduler::Intrawave, + BlockGemmPipelineVersion::v1>{}); +} + +} // namespace instance +} // namespace device +} // namespace tensor_operation +} // namespace ck From e3556fed0453e66cdebc5dad6b903f5e902cd9b4 Mon Sep 17 00:00:00 2001 From: Jeff Huang Date: Thu, 29 Jan 2026 07:18:41 +0800 Subject: [PATCH 07/11] Optimize batch prefill kernel performance for VECTORIZED_LAYOUT KV cache (#3657) - Add multi-dimensional page index support (YsGatherDims) in tile_scatter_gather - Add is_gather_dim() and get_gather_index() for multi-dim page lookup - Override MakeVDramTileDistribution() for VECTORIZED_LAYOUT to match GEMM's BWarpDstrEncoding (K decomposition: {K2, K0, K1}) - Add GetGemmKDecomposition() to retrieve kABKLane and kKPerThread - Add static_assert for RowMajor VLayout requirement in batch prefill Co-authored-by: Po Yen Chen --- .../core/tensor/tile_scatter_gather.hpp | 166 ++++++++++-- ..._batch_prefill_pipeline_qr_ks_vs_async.hpp | 227 ++++++++++++----- ...pipeline_qr_ks_vs_async_default_policy.hpp | 241 +++++++++++++++++- .../pipeline/block_fmha_pipeline_problem.hpp | 3 + 4 files changed, 553 insertions(+), 84 deletions(-) diff --git a/include/ck_tile/core/tensor/tile_scatter_gather.hpp b/include/ck_tile/core/tensor/tile_scatter_gather.hpp index 2ffaff2973..aa29345892 100644 --- a/include/ck_tile/core/tensor/tile_scatter_gather.hpp +++ b/include/ck_tile/core/tensor/tile_scatter_gather.hpp @@ -26,17 +26,26 @@ namespace ck_tile { * * @tparam BottomTensorView_ Class describing & holding device tensor memory. * @tparam WindowLengths_ Spatial sizes of windowed view on tensor. - * @tparam StaticTileDistribution_ Thread distribution (mapping) into Tile dimensions - * @tparam NumCoord TBD + * @tparam StaticTileDistribution_ Thread distribution (mapping) into Tile dimensions. + * @tparam StaticPageIndexArray_ Array type holding page indices for scatter/gather. + * @tparam StaticValidArray_ Array type holding validity flags (nullptr_t if unused). + * @tparam HsGatherDim H-space dimension index used for gather lookup (default: 0). + * @tparam NumCoord Number of pre-computed coordinates for pipelining (default: 1). + * @tparam YsGatherDims Sequence of Y-space dimension indices used for page lookup. + * For single dimension: sequence<0> (default). + * For multiple dimensions: sequence where + * the combined index is computed as: + * idx[dim0] + idx[dim1] * len[dim0] + idx[dim2] * len[dim0] * + * len[dim1] + ... */ template + index_t HsGatherDim = 0, + index_t NumCoord = 1, + typename YsGatherDims = sequence<0>> struct tile_scatter_gather { using BottomTensorView = remove_reference_t; @@ -77,6 +86,75 @@ struct tile_scatter_gather using BottomTensorCoord = decltype(make_tensor_coordinate(BottomTensorDesc{}, BottomTensorIndex{})); + /** + * @brief Check if a given Y-space dimension index is a gather dimension. + * + * Gather dimensions are those specified in YsGatherDims template parameter. + * When computing forward_step_scatter, gather dimensions are set to 0 + * because page offset lookup handles address calculation for these dimensions. + * + * @param i Y-space dimension index to check + * @return true if dimension i is in YsGatherDims, false otherwise + */ + CK_TILE_DEVICE static constexpr bool is_gather_dim(index_t i) + { + return sequence_any_of(YsGatherDims{}, [i](auto k) { return i == k; }); + } + + /** + * @brief Compute the linearized gather index from Y-space indices for page lookup. + * + * This function converts multi-dimensional Y-space indices (specified by YsGatherDims) + * into a single linearized index used to look up the page offset in page_idx_ array. + * + * For single gather dimension (YsGatherDims::size() == 1): + * Simply returns idx_ys_start[YsGatherDims::at(0)] + * + * For multiple gather dimensions (e.g., YsGatherDims = sequence<0, 2>): + * Computes: idx[dim0] + idx[dim1] * len[dim0] + idx[dim2] * len[dim0] * len[dim1] + ... + * This is row-major linearization where earlier dimensions are inner (faster-varying). + * + * @tparam YsIndex Type of the Y-space index tuple/array + * @param idx_ys_start Current Y-space indices from space-filling curve iteration + * @return Linearized index for page_idx_ array lookup + */ + template + CK_TILE_DEVICE static constexpr auto get_gather_index(const YsIndex& idx_ys_start) + { + // TODO: Consider making ys_lengths_ part of public API or adding accessor + static_assert(sizeof(TileDstr::DstrEncode::detail::ys_lengths_) > 0, + "Relies on internal detail::ys_lengths_"); + + constexpr index_t num_gather_dims = YsGatherDims::size(); + + if constexpr(num_gather_dims == 1) + { + return idx_ys_start[number{}]; + } + else + { + // Recursive lambda to compute index as a compile-time number + // Uses row-major linearization: idx[0] + idx[1] * len[0] + idx[2] * len[0] * len[1] + + // ... + auto recurse = [&](auto self, auto i_constant) { + constexpr index_t i = decltype(i_constant)::value; + constexpr index_t dim = YsGatherDims::at(i); + auto current_val = idx_ys_start[number{}]; + + if constexpr(i + 1 < num_gather_dims) + { + constexpr index_t len = TileDstr::DstrEncode::detail::ys_lengths_[dim]; + return current_val + self(self, number{}) * number{}; + } + else + { + return current_val; + } + }; + return recurse(recurse, number<0>{}); + } + } + struct load_store_traits { private: @@ -375,7 +453,7 @@ struct tile_scatter_gather // data index [y0, y1, ...] constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto idx_gather = idx_ys_start[number{}]; + constexpr auto idx_gather = get_gather_index(idx_ys_start); const auto page_offset = page_idx_[idx_gather]; // read from bottom tensor @@ -427,7 +505,7 @@ struct tile_scatter_gather constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); constexpr auto forward_step_scatter = generate_tuple( - [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + [&](auto i) { return is_gather_dim(i) ? 0 : idx_diff_ys[i]; }, number{}); constexpr auto idx_diff_ps_ys = container_concat( @@ -485,7 +563,7 @@ struct tile_scatter_gather // data index [y0, y1, ...] constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto idx_gather = idx_ys_start[number{}]; + constexpr auto idx_gather = get_gather_index(idx_ys_start); const auto page_offset = page_idx_[idx_gather]; // merge page_offset into bottom_coord @@ -513,7 +591,7 @@ struct tile_scatter_gather constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); constexpr auto forward_step_scatter = generate_tuple( - [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + [&](auto i) { return is_gather_dim(i) ? 0 : idx_diff_ys[i]; }, number{}); constexpr auto idx_diff_ps_ys = container_concat( @@ -598,7 +676,7 @@ struct tile_scatter_gather }(); constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto idx_gather = idx_ys_start[number{}]; + constexpr auto idx_gather = get_gather_index(idx_ys_start); const auto page_offset = page_idx_[idx_gather]; // read from bottom tensor @@ -624,7 +702,7 @@ struct tile_scatter_gather constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); constexpr auto forward_step_scatter = generate_tuple( - [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + [&](auto i) { return is_gather_dim(i) ? 0 : idx_diff_ys[i]; }, number{}); constexpr auto idx_diff_ps_ys = container_concat( @@ -718,7 +796,7 @@ struct tile_scatter_gather }(); constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto idx_gather = idx_ys_start[number{}]; + constexpr auto idx_gather = get_gather_index(idx_ys_start); const auto page_offset = page_idx_[idx_gather]; auto mixed_bottom_thread_coord = bottom_tensor_thread_coord; @@ -748,7 +826,7 @@ struct tile_scatter_gather constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); constexpr auto forward_step_scatter = generate_tuple( - [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + [&](auto i) { return is_gather_dim(i) ? 0 : idx_diff_ys[i]; }, number{}); constexpr auto idx_diff_ps_ys = container_concat( @@ -791,7 +869,7 @@ struct tile_scatter_gather // data index [y0, y1, ...] constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto idx_gather = idx_ys_start[number<0>{}]; + constexpr auto idx_gather = get_gather_index(idx_ys_start); const auto page_offset = page_idx_[idx_gather]; // read from distributed tensor @@ -837,7 +915,7 @@ struct tile_scatter_gather constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); constexpr auto forward_step_scatter = generate_tuple( - [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + [&](auto i) { return is_gather_dim(i) ? 0 : idx_diff_ys[i]; }, number{}); constexpr auto idx_diff_ps_ys = container_concat( @@ -874,11 +952,11 @@ struct tile_scatter_gather // data index [y0, y1, ...] constexpr auto idx_ys_start = SFC_Ys::get_index(iAccess); - constexpr auto idx_gather = idx_ys_start[number<0>{}]; + constexpr auto idx_gather = get_gather_index(idx_ys_start); const auto page_offset = page_idx_[idx_gather]; // printf("idx_ys_start[0], idx_ys_start[1](%d, %d) \n", - // idx_ys_start[number<0>{}]+0, idx_ys_start[number<1>{}]+0); + // get_gather_index(idx_ys_start)+0, idx_ys_start[number<1>{}]+0); // read from distributed tensor // vector_type_t vec; @@ -928,7 +1006,7 @@ struct tile_scatter_gather constexpr auto idx_diff_ys = SFC_Ys::get_forward_step(iAccess); constexpr auto forward_step_scatter = generate_tuple( - [&](auto i) { return i == YsGatherDim ? 0 : idx_diff_ys[i]; }, + [&](auto i) { return is_gather_dim(i) ? 0 : idx_diff_ys[i]; }, number{}); constexpr auto idx_diff_ps_ys = container_concat( @@ -1076,6 +1154,53 @@ struct tile_scatter_gather }; // TODO: use strategy +/** + * @brief Factory function to create tile_scatter_gather with multi-dimensional gather support. + * + * This overload accepts a sequence to specify multiple Y-space dimensions + * for page lookup. Use this when the tile distribution decomposes the paged dimension + * into multiple Y-space dimensions (e.g., VECTORIZED_LAYOUT V tensor with K decomposition + * {K2, K0, K1} where both Y0 and Y2 contribute to page index). + * + * @tparam HsGatherDim H-space dimension for gather + * @tparam NumCoord Number of pre-computed coordinates + * @tparam YsGatherDims Parameter pack specifying which Y-dimensions are used for page lookup + * + * @param tensor_view The underlying tensor view for device memory access + * @param window_lengths Static window sizes for each dimension + * @param origin Window origin coordinates on the bottom tensor + * @param tile_distribution Thread-to-tile mapping distribution + * @param page_idx Array of page offsets (in bytes) for scatter/gather + */ +template +CK_TILE_DEVICE constexpr auto +make_tile_scatter_gather(const TensorView_& tensor_view, + const WindowLengths_& window_lengths, + const multi_index& origin, + const StaticTileDistribution_& tile_distribution, + const StaticPageIndexArray_& page_idx, + number, + number, + sequence) +{ + return tile_scatter_gather, + remove_cvref_t, + remove_cvref_t, + remove_cvref_t, + std::nullptr_t, + HsGatherDim, + NumCoord, + sequence>{ + tensor_view, window_lengths, origin, tile_distribution, page_idx, nullptr}; +} + +// Legacy overload (compatible with original API) template & origin, const StaticTileDistribution_& tile_distribution, - const StaticPageIndexArray_& page_idx, // perbytes + const StaticPageIndexArray_& page_idx, number = {}, number = {}) { @@ -1097,7 +1222,8 @@ make_tile_scatter_gather(const TensorView_& tensor_view, remove_cvref_t, std::nullptr_t, HsGatherDim, - NumCoord>{ + NumCoord, + sequence<0>>{ tensor_view, window_lengths, origin, tile_distribution, page_idx, nullptr}; } diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp index c75f5d58c4..48e8f75ae7 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp @@ -533,32 +533,170 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync auto randval_dram_window = dropout.template MakeRandvalDramWindow( randval_dram_block_window_tmp, seqlen_k_start); - auto v_dist = Policy::template MakeVDramTileDistribution(); - auto v_coord = v_dist.calculate_index(); - const auto VPageIndexDim = I1; - using VDstrEncode = typename decltype(v_dist)::DstrEncode; - constexpr index_t V_KRepeat = VDstrEncode::hs_lengthss_[I1][I3]; - statically_indexed_array v_offsets; - kv_offset_array_transform, - decltype(v_coord), - VPageIndexDim, - kPageBlockSize, - 0, - V_KRepeat, - 1, - kKVMemoryLayout, - false, - kN0, - kVectorSize>( - page_idx, stride_v, page_stride_v, v_coord, v_offsets, current_seq_k); + auto v_dist = Policy::template MakeVDramTileDistribution(); + auto v_coord = v_dist.calculate_index(); + using VDstrEncode = typename decltype(v_dist)::DstrEncode; + // V tensor K-dimension decomposition for page index computation + // ============================================================ + // The K dimension (seqlen_k) in V distribution is decomposed into multiple sub-dimensions. + // This decomposition determines how threads iterate over the K dimension and how page + // indices are computed for paged KV cache. + // + // The decomposition pattern differs by memory layout: + // + // VECTORIZED_LAYOUT (ColumnMajor, custom distribution): + // 3D decomposition: K = K2 × K0 × K1 + // - K2 (V_KIterOuter): Outer iteration count + // - K0 (V_KLanes): Lanes for K dimension (matches GEMM kABKLane) + // - K1 (V_KIterInner): Vector load size (matches GEMM kKPerThread) + // - hs_lengthss_[I1] = {K2, K0, K1}, size = 3 (or {K0, K1} size = 2 if no outer iter) + // + // LINEAR_LAYOUT ColumnMajor (base class distribution): + // 2D decomposition: K = K0 × K1 + // - K0: Lanes for K dimension (may not match GEMM kABKLane) + // - K1: Vector load size + // - hs_lengthss_[I1] = {K0, K1}, size = 2 + // + // LINEAR_LAYOUT RowMajor (base class distribution): + // 4D decomposition: K = K0 × K1 × K2 × K3 (uses shuffle_tile for GEMM alignment) + // 3D decomposition: K = K0 × K1 × K2 (fallback case) + // - Page lookup uses Y-space's last dimension only (inner iteration) + // + // V_PageIdxRepeat = total number of page lookups per thread = V_KIterOuter × V_KIterInner + constexpr index_t V_KIterInner = VDstrEncode::hs_lengthss_[I1].back(); + + // Compute V_KIterOuter and V_KLanes based on memory layout and K decomposition + constexpr index_t V_KIterOuter = [] { + if constexpr(kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + // VECTORIZED_LAYOUT: 3D decomposition {K2, K0, K1} when outer iteration is needed + if constexpr(VDstrEncode::hs_lengthss_[I1].size() == 3) + return static_cast(VDstrEncode::hs_lengthss_[I1][I0]); + else + return index_t{1}; + } + else + { + // LINEAR_LAYOUT: No outer iteration for page lookup + // RowMajor uses shuffle_tile, ColumnMajor has simple 2D decomposition + // Both cases use single-dimension Y-space page lookup + return index_t{1}; + } + }(); + + constexpr index_t V_KLanes = [] { + if constexpr(kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + // VECTORIZED_LAYOUT: K0 is the lanes dimension + if constexpr(V_KIterOuter > 1) + return static_cast(VDstrEncode::hs_lengthss_[I1][I1]); + else + return static_cast(VDstrEncode::hs_lengthss_[I1][I0]); + } + else + { + // LINEAR_LAYOUT: First dimension is K0 (lanes) + return static_cast(VDstrEncode::hs_lengthss_[I1][I0]); + } + }(); + + // This affects page offset computation - need to track offsets for each (k2, k1) + // combination + constexpr index_t V_PageIdxRepeat = V_KIterInner * V_KIterOuter; + + // VPageIndexYDims: Y-space dimension indices that participate in page index computation + // ================================================================================ + // In tile_scatter_gather, the gather index is computed from Y-space coordinates. + // This sequence specifies which Y dimensions should be linearized to form the page lookup + // index. + // + // VECTORIZED_LAYOUT with outer iteration: sequence + // - Both K1 and K2 are in Y-space (thread iteration dimensions) + // - gather_index = y_k1 + y_k2 * len(Y_K1) (linearized 2D -> 1D) + // + // VECTORIZED_LAYOUT without outer iteration / LINEAR_LAYOUT: sequence + // - Only the innermost K dimension is used for page lookup (single dimension) + // + constexpr auto VPageIndexYDims = []() { + // K1Minor is always the last element index in hs_lengthss_[I1] + constexpr index_t K1Minor = VDstrEncode::hs_lengthss_[I1].size() - 1; + constexpr index_t Y_K1 = VDstrEncode::detail::rhs_major_minor_to_ys_[2][K1Minor]; + + if constexpr(kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT && + V_KIterOuter > 1) + { + // VECTORIZED_LAYOUT with outer iteration: need 2D page lookup + constexpr index_t Y_K2 = VDstrEncode::detail::rhs_major_minor_to_ys_[2][I0]; + return sequence{}; + } + else + { + // LINEAR_LAYOUT or VECTORIZED_LAYOUT without outer iteration: 1D page lookup + return sequence{}; + } + }(); + + static_assert(decltype(VPageIndexYDims)::at(0) < VDstrEncode::NDimY, + "V page-index Y dim must be valid"); + + statically_indexed_array v_offsets; + auto update_v_offsets = [&](auto k_loop_start) { + constexpr index_t kLoopStart = decltype(k_loop_start)::value; + // For 3D K decomposition (K2, K0, K1), compute offsets for each K2 slice + // The global K offset for (k2, k1) is: kLoopStart + k2 * (K0 * K1) + k1 + // We iterate K2 outer, K1 inner, and merge into 1D v_offsets array + if constexpr(V_KIterOuter > 1) + { + static_for<0, V_KIterOuter, 1>{}([&](auto k2) { + statically_indexed_array v_offsets_k2; + kv_offset_array_transform, + decltype(v_coord), + I1, + kPageBlockSize, + kLoopStart + k2.value * V_KLanes * V_KIterInner, + V_KIterInner, + 1, + kKVMemoryLayout, + false, + kN0, + kVectorSize>( + page_idx, stride_v, page_stride_v, v_coord, v_offsets_k2, current_seq_k); + static_for<0, V_KIterInner, 1>{}([&](auto k1) { + constexpr auto idx = number{}; + v_offsets[idx] = v_offsets_k2[k1]; + }); + }); + } + else + { + kv_offset_array_transform, + decltype(v_coord), + I1, + kPageBlockSize, + kLoopStart, + V_KIterInner, + 1, + kKVMemoryLayout, + false, + kN0, + kVectorSize>( + page_idx, stride_v, page_stride_v, v_coord, v_offsets, current_seq_k); + } + }; + update_v_offsets(number<0>{}); auto v_dram_window = make_tile_scatter_gather(v_dram_block_window_tmp.get_bottom_tensor_view(), v_dram_block_window_tmp.get_window_lengths(), {0, seqlen_k_start}, // TODO: hdim split? v_dist, v_offsets, - VPageIndexDim); + number<1>{}, // HsGatherDim + number<1>{}, // NumCoord + VPageIndexYDims); // prefetch K tile async_load_tile_raw( @@ -625,18 +763,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync __builtin_amdgcn_sched_barrier(1); auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant{}); - kv_offset_array_transform, - decltype(v_coord), - VPageIndexDim, - kPageBlockSize, - kK1, - V_KRepeat, - 1, - kKVMemoryLayout, - false, - kN0, - kVectorSize>( - page_idx, stride_v, page_stride_v, v_coord, v_offsets, current_seq_k); + update_v_offsets(number{}); v_dram_window.update_page_idx(v_offsets); const auto p = [&]() { @@ -766,7 +893,9 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync __builtin_amdgcn_sched_barrier(0x7F); // store & prefetch next v, after the max reduction - if constexpr(std::is_same_v) + if constexpr(std::is_same_v && + kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT) { auto v_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledVRegBlockDescriptor()); @@ -787,8 +916,8 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync get_slice_tile(v_lds_window, sequence<(LdsSeq.at(number{})) * kN1, 0>{}, sequence<(LdsSeq.at(number{}) + 1) * kN1, kK1>{}); - store_tile(v_lds_window_tmp, - tile_elementwise_in(v_element_func, v_buf)); // store the prefetch + const auto v_store_tile = tile_elementwise_in(v_element_func, v_buf); + store_tile(v_lds_window_tmp, v_store_tile); // store the prefetch } if constexpr(k1_loops > 1) @@ -799,18 +928,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync kK1}); // will have scratch if move this right after load_tile(v_dram)... v_buf = load_tile( v_dram_window, number<-1>{}, bool_constant{}); // load next v_buf - kv_offset_array_transform, - decltype(v_coord), - VPageIndexDim, - kPageBlockSize, - 2 * kK1, - V_KRepeat, - 1, - kKVMemoryLayout, - false, - kN0, - kVectorSize>( - page_idx, stride_v, page_stride_v, v_coord, v_offsets, current_seq_k); + update_v_offsets(number<2 * kK1>{}); v_dram_window.update_page_idx(v_offsets); } __builtin_amdgcn_sched_barrier(0); @@ -938,18 +1056,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync { v_buf = load_tile( v_dram_window, number<-1>{}, bool_constant{}); // load next v_buf - kv_offset_array_transform, - decltype(v_coord), - VPageIndexDim, - kPageBlockSize, - (2 + i_k1.value) * kK1, - V_KRepeat, - 1, - kKVMemoryLayout, - false, - kN0, - kVectorSize>( - page_idx, stride_v, page_stride_v, v_coord, v_offsets, current_seq_k); + update_v_offsets(number<(2 + i_k1.value) * kK1>{}); v_dram_window.update_page_idx(v_offsets); } block_sync_lds(); @@ -961,7 +1068,9 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync sequence<(LdsSeq.at(number{})) * kN1, 0>{}, sequence<(LdsSeq.at(number{}) + 1) * kN1, kK1>{})); - if constexpr(std::is_same_v) + if constexpr(std::is_same_v && + kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT) { auto v_shuffle_tmp = make_static_distributed_tensor( Policy::template MakeShuffledVRegBlockDescriptor()); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async_default_policy.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async_default_policy.hpp index 33e6ad006a..45b7356dfa 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async_default_policy.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async_default_policy.hpp @@ -4,15 +4,246 @@ #pragma once #include "ck_tile/core.hpp" +#include "ck_tile/ops/fmha/block/block_attention_kvcache_layout_enum.hpp" #include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qx_ks_vs_custom_policy.hpp" namespace ck_tile { // This pipeline is qkv all located in LDS -using BlockFmhaBatchPrefillPipelineQRKSVSAsyncDefaultPolicy = - BlockFmhaPipelineQXKSVSCustomPolicy; +struct BlockFmhaBatchPrefillPipelineQRKSVSAsyncDefaultPolicy + : BlockFmhaPipelineQXKSVSCustomPolicy +{ + using Base = BlockFmhaPipelineQXKSVSCustomPolicy; + + template + CK_TILE_HOST_DEVICE static constexpr auto GetAlignmentV() + { + if constexpr(Problem::kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + using VDataType = remove_cvref_t; + constexpr index_t kDwordx4Bytes = 16; + return kDwordx4Bytes / sizeof(VDataType); + } + else + { + return Base::template GetAlignmentV(); + } + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackV() + { + if constexpr(Problem::kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + // For VECTORIZED_LAYOUT, kKPack should match GEMM's kKPerThread + // to ensure correct LDS access pattern + constexpr auto gemm_k_decomp = GetGemmKDecomposition(); + constexpr index_t kKPerThread = gemm_k_decomp.template at<1>(); + return kKPerThread; + } + else + { + return Base::template GetSmemKPackV(); + } + } + + template + CK_TILE_HOST_DEVICE static constexpr auto GetSingleSmemElementSpaceSize() + { + if constexpr(Problem::kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + // For VECTORIZED_LAYOUT, we need to use our GetSmemKPackV for V size calculation + constexpr index_t SingleKSize = [&]() { + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN0; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + constexpr index_t NumWarps = Problem::BlockFmhaShape::NumWarps; + constexpr index_t WarpSize = ck_tile::get_warp_size(); + + constexpr index_t KPack = Base::template GetSmemKPackK(); + constexpr index_t KVector = Base::template GetAlignmentK(); + constexpr index_t kPad = KPack; + + static_assert(WarpSize * KVector >= kKPerBlock && + WarpSize * KVector % kKPerBlock == 0); + constexpr index_t LanesPerK = kKPerBlock / KVector; + constexpr index_t LaneGroups = WarpSize / LanesPerK; + constexpr index_t NumIssues = kNPerBlock / (LaneGroups * NumWarps); + + return NumIssues * NumWarps * (WarpSize * KVector + kPad); + }(); + + constexpr index_t SingleVSize = [&]() { + using VDataType = remove_cvref_t; + constexpr index_t Banks = get_n_lds_banks(); + constexpr index_t PixelsPerRow = Banks * 4 / sizeof(VDataType); + constexpr index_t kKPack = GetSmemKPackV(); // Use our override! + static_assert(PixelsPerRow % kKPack == 0); + constexpr index_t NPerRow = PixelsPerRow / kKPack; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN1; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + static_assert(kNPerBlock % NPerRow == 0); + static_assert(kKPerBlock % kKPack == 0); + + return (kKPerBlock / kKPack) * (kNPerBlock / NPerRow) * (PixelsPerRow + kKPack); + }(); + + return max(SingleKSize, SingleVSize); + } + else + { + return Base::template GetSingleSmemElementSpaceSize(); + } + } + + template + CK_TILE_HOST_DEVICE static constexpr auto MakeVLdsBlockDescriptor() + { + if constexpr(Problem::kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + using VDataType = remove_cvref_t; + constexpr index_t Banks = get_n_lds_banks(); + constexpr index_t PixelsPerRow = Banks * 4 / sizeof(VDataType); + constexpr index_t kKPack = GetSmemKPackV(); + static_assert(PixelsPerRow % kKPack == 0); + constexpr index_t NPerRow = PixelsPerRow / kKPack; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN1; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + static_assert(kNPerBlock % NPerRow == 0); + static_assert(kKPerBlock % kKPack == 0); + + constexpr auto v_lds_block_desc_0 = make_naive_tensor_descriptor( + make_tuple(number{}, + number{}, + number{}, + number{}, + number{}), + make_tuple(number()>{}, + number<(kNPerBlock / NPerRow) * (PixelsPerRow + kKPack)>{}, + number{}, + number{}, + number<1>{}), + number{}, + number<1>{}); + + constexpr auto v_lds_block_desc = transform_tensor_descriptor( + v_lds_block_desc_0, + make_tuple(make_merge_transform(make_tuple(number{}, + number{}, + number{})), + make_merge_transform( + make_tuple(number{}, number{}))), + make_tuple(sequence<0, 2, 3>{}, sequence<1, 4>{}), + make_tuple(sequence<0>{}, sequence<1>{})); + + return v_lds_block_desc; + } + else + { + return Base::template MakeVLdsBlockDescriptor(); + } + } + + // Helper to get GEMM's K decomposition parameters (kABKLane, kKPerThread) + template + CK_TILE_HOST_DEVICE static constexpr auto GetGemmKDecomposition() + { + // Get the KV block GEMM and extract warp gemm's K decomposition + constexpr auto gemm = Base::template GetKVBlockGemm(); + using BlockGemm = remove_cvref_t; + constexpr auto config = + BlockGemm::Policy::template GetWarpGemmMWarpNWarp(); + using WG = remove_cvref_t())>; + + // Return kABKLane and kKPerThread from warp gemm + return make_tuple(number{}, + number{}); + } + + template + CK_TILE_DEVICE static constexpr auto MakeVDramTileDistribution() + { + if constexpr(Problem::kKVMemoryLayout == + BlockAttentionKVCacheMemoryLayoutEnum::VECTORIZED_LAYOUT) + { + // For VECTORIZED_LAYOUT, use column-major distribution (K direction vector load) + // The K decomposition must match GEMM's BWarpDstrEncoding to ensure correct LDS access + constexpr index_t kBlockSize = Problem::kBlockSize; + constexpr index_t kNPerBlock = Problem::BlockFmhaShape::kN1; + constexpr index_t kKPerBlock = Problem::BlockFmhaShape::kK1; + + // Get GEMM's K decomposition (kABKLane, kKPerThread) + constexpr auto gemm_k_decomp = GetGemmKDecomposition(); + constexpr index_t kABKLane = gemm_k_decomp.template at<0>(); + constexpr index_t kKPerThread = gemm_k_decomp.template at<1>(); + + // K1 = kKPerThread (inner K dimension, matches GEMM's expectation) + // K0 = kKPerBlock / K1 (outer K dimension) + // But we need K0 to match kABKLane for the per-warp iteration + constexpr index_t K1 = kKPerThread; + constexpr index_t K0 = kABKLane; + + // Verify K decomposition matches GEMM's BWarpDstrEncoding requirements + static_assert(K0 == kABKLane, "K0 must match GEMM's kABKLane for correct LDS access"); + static_assert(K1 == kKPerThread, + "K1 must match GEMM's kKPerThread for correct LDS access"); + + // K0 * K1 may be less than kKPerBlock, so we need outer iteration + constexpr index_t KPerIter = K0 * K1; + constexpr index_t KOuterIter = kKPerBlock / KPerIter; + + constexpr index_t N2 = get_warp_size() / K0; + constexpr index_t N1 = kBlockSize / get_warp_size(); + static_assert(N2 != 0, "N2 is zero, which will lead to a division by zero error."); + static_assert(N1 != 0, "N1 is zero, which will lead to a division by zero error."); + constexpr index_t N0 = kNPerBlock / (N2 * N1); + static_assert(N0 != 0, "N0 is zero"); + + if constexpr(KOuterIter == 1) + { + // Simple case: K decomposition matches exactly + constexpr auto dstr = make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<2, 1>, + sequence<1, 0>>{}); + static_assert(container_reduce(dstr.get_lengths(), std::multiplies{}, 1) == + kNPerBlock * kKPerBlock); + return dstr; + } + else + { + // Need outer K iteration + constexpr index_t K2 = KOuterIter; + constexpr auto dstr = make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 1>>, + sequence<2, 1, 2>, + sequence<2, 0, 0>>{}); + static_assert(container_reduce(dstr.get_lengths(), std::multiplies{}, 1) == + kNPerBlock * kKPerBlock); + return dstr; + } + } + else + { + // For non-VECTORIZED_LAYOUT, use base class implementation + return Base::template MakeVDramTileDistribution(); + } + } +}; } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp index a489eabb73..eabf74faf8 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_problem.hpp @@ -121,6 +121,9 @@ struct BlockFmhaBatchPrefillPipelineProblem static_assert(!kIsVectorizedLayout || kPageBlockSize % kVectorSize == 0, "kPageBlockSize must be divisible by kVectorSize for vectorized layout"); static_assert(kIsGroupMode_, "Batch prefill requires group mode"); + + static_assert(BlockFmhaShape_::IsVLayoutRowMajor, + "Batch prefill kernel requires RowMajor VLayout"); }; template Date: Wed, 28 Jan 2026 19:45:09 -0800 Subject: [PATCH 08/11] [CK_Tile] Adding support for preshuffleQuant in AB quant Block Scale Gemm (#3629) * initial commit * preshuffleQuant support for ABQuant * fix mxfp4 to use correct QuantGroupSize * addressing review comments and seperated Preshufflequant for A and B * updated grouped gemm example for updated traits definition * fix for CI failure * updated grouped_gemm_abquant test for updated traits definition * updated grouped_gemm_abquant test for updated traits definition --- .../17_grouped_gemm/abquant_grouped_gemm.cpp | 6 +- .../quant_invoke_grouped_gemm_kernel.hpp | 6 +- .../gemm_abquant_quantgrouped.cpp | 30 ++++ .../38_block_scale_gemm/gemm_utils.hpp | 12 +- .../run_gemm_quant_example.inc | 45 +++-- .../block/block_gemm_quant_common.hpp | 6 +- ...versal_gemm_ar_aquant_flatbr_bquant_cr.hpp | 17 +- ...ock_universal_gemm_ar_flatbr_bquant_cr.hpp | 19 +- ..._universal_gemm_as_aquant_bs_bquant_cr.hpp | 32 +++- .../block_universal_gemm_as_aquant_bs_cr.hpp | 20 +-- .../block_universal_gemm_as_bs_bquant_cr.hpp | 45 +++-- .../gemm_quant/kernel/gemm_quant_kernel.hpp | 168 +++++++++--------- .../gemm_abquant_pipeline_ag_bg_cr_v3.hpp | 25 ++- .../gemm_aquant_pipeline_ag_bg_cr_base.hpp | 18 +- .../gemm_aquant_pipeline_ag_bg_cr_mem.hpp | 26 +-- .../gemm_aquant_pipeline_ag_bg_cr_policy.hpp | 36 ++-- .../gemm_aquant_pipeline_ag_bg_cr_v3.hpp | 26 +-- .../gemm_bquant_pipeline_ag_bg_cr_base.hpp | 30 ++-- .../gemm_bquant_pipeline_ag_bg_cr_policy.hpp | 20 +-- .../gemm_bquant_pipeline_ag_bg_cr_v3.hpp | 32 ++-- .../pipeline/gemm_group_quant_utils.hpp | 15 +- .../gemm_mxfp4_pipeline_ag_bg_cr_base.hpp | 30 ++-- .../gemm_mxfp4_pipeline_ag_bg_cr_policy.hpp | 8 +- .../gemm_mxfp4_pipeline_ag_bg_cr_v3.hpp | 26 +-- .../pipeline/gemm_quant_pipeline_problem.hpp | 8 +- .../gemm_wp_bquant_pipeline_ag_bg_cr_v2.hpp | 26 +-- .../pipeline/tile_gemm_quant_traits.hpp | 8 +- test/ck_tile/gemm_block_scale/CMakeLists.txt | 39 ++-- ...est_gemm_quant_abquant_preshuffleQuant.cpp | 43 +++++ .../gemm_block_scale/test_gemm_quant_base.hpp | 6 +- .../test_gemm_quant_fixtures.hpp | 25 +-- .../test_grouped_gemm_abquant_util.hpp | 2 + .../test_grouped_gemm_util_quant.hpp | 2 + 33 files changed, 490 insertions(+), 367 deletions(-) create mode 100644 test/ck_tile/gemm_block_scale/test_gemm_quant_abquant_preshuffleQuant.cpp diff --git a/example/ck_tile/17_grouped_gemm/abquant_grouped_gemm.cpp b/example/ck_tile/17_grouped_gemm/abquant_grouped_gemm.cpp index 84da1e26da..703751b760 100644 --- a/example/ck_tile/17_grouped_gemm/abquant_grouped_gemm.cpp +++ b/example/ck_tile/17_grouped_gemm/abquant_grouped_gemm.cpp @@ -59,7 +59,8 @@ float grouped_gemm_abquant(const std::vector& gemm_descs, using GemmUniversalTraits = ck_tile::TileGemmQuantTraits& gemm_descs, using GemmUniversalTraits = ck_tile::TileGemmQuantTraits(arg_parser); }; + lut[hash_multiple_strings({"fp8", + "abquant", + "non-preshuffleb", + "preshufflequant", + "1x1x128"})] = [](const ck_tile::ArgParser& arg_parser) { + using AQuantGroupSize = ck_tile::QuantGroupShape>; + using BQuantGroupSize = ck_tile::QuantGroupShape>; + using TypeConfig = + decltype(GemmQuantTypeConfig{}); + return run_gemm_example_prec_type, + TypeConfig, + AQuantGroupSize, + BQuantGroupSize, + ck_tile::QuantType::ABQuantGrouped>(arg_parser); + }; + lut[hash_multiple_strings({"fp8", + "abquant", + "non-preshuffleb", + "preshufflequant", + "1x128x128"})] = [](const ck_tile::ArgParser& arg_parser) { + using AQuantGroupSize = ck_tile::QuantGroupShape>; + using BQuantGroupSize = ck_tile::QuantGroupShape>; + using TypeConfig = + decltype(GemmQuantTypeConfig{}); + return run_gemm_example_prec_type, + TypeConfig, + AQuantGroupSize, + BQuantGroupSize, + ck_tile::QuantType::ABQuantGrouped>(arg_parser); + }; return 0; }(); diff --git a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp index 0feba8e471..085d634441 100644 --- a/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp +++ b/example/ck_tile/38_block_scale_gemm/gemm_utils.hpp @@ -80,7 +80,8 @@ struct GemmConfigBase static constexpr ck_tile::index_t TileParitionerGroupNum = 8; static constexpr ck_tile::index_t TileParitionerM01 = 4; - static constexpr bool PreshuffleQuant = false; + static constexpr bool APreshuffleQuant = false; + static constexpr bool BPreshuffleQuant = false; static constexpr bool PreshuffleB = false; static constexpr bool DoubleSmemBuffer = false; static constexpr bool TiledMMAPermuteN = false; @@ -157,7 +158,8 @@ struct GemmConfigPreshuffleQuantDecode : public GemmConfigBase static constexpr ck_tile::index_t K_Warp_Tile = ck_tile::get_k_warp_tile(); - static constexpr bool PreshuffleQuant = true; + static constexpr bool APreshuffleQuant = true; + static constexpr bool BPreshuffleQuant = true; }; template @@ -187,7 +189,7 @@ template struct GemmConfigPreshuffleB_PreshuffleBQuant_Decode : public GemmConfigPreshuffleB_BQuant_Decode { - static constexpr bool PreshuffleQuant = true; + static constexpr bool BPreshuffleQuant = true; }; template @@ -218,7 +220,7 @@ template struct GemmConfigPreshuffleB_PreshuffleBQuant_Prefill : public GemmConfigPreshuffleB_BQuant_Prefill { - static constexpr bool PreshuffleQuant = true; + static constexpr bool BPreshuffleQuant = true; }; template @@ -272,7 +274,7 @@ struct GemmConfigABQuantPrefill : public GemmConfigQuantPrefill template struct GemmConfigPreshuffleBQuantPrefill : public GemmConfigQuantPrefill { - static constexpr bool PreshuffleQuant = true; + static constexpr bool BPreshuffleQuant = true; }; template diff --git a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc index 270653d408..665c7828ad 100644 --- a/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc +++ b/example/ck_tile/38_block_scale_gemm/run_gemm_quant_example.inc @@ -33,7 +33,8 @@ template ); - constexpr bool transpose_c = QuantMode == ck_tile::QuantType::ABQuantGrouped; + constexpr bool transpose_c = + GemmConfig::TransposeC; // QuantMode == ck_tile::QuantType::ABQuantGrouped; using ComputeDataType = std::conditional_t; @@ -73,7 +75,7 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str GemmConfig::PreshuffleB == true, ck_tile::BaseWeightPreshufflePipelineAGmemBGmemCRegV2, std::conditional_t< - QuantMode == ck_tile::QuantType::AQuantGrouped && GemmConfig::PreshuffleQuant == true, + QuantMode == ck_tile::QuantType::AQuantGrouped && GemmConfig::APreshuffleQuant == true, ck_tile::BaseGemmPipelineAgBgCrCompV3, std::conditional_t< QuantMode == ck_tile::QuantType::AQuantGrouped, @@ -146,7 +148,7 @@ float gemm_calc_quant(const ck_tile::QuantGemmHostArgs& args, const ck_tile::str has_hot_loop_v, tail_number_v>>>>; using AQuantPipeline = - std::conditional_t, ck_tile::AQuantGemmPipelineAgBgCrMem>; @@ -390,8 +392,10 @@ float invoke_gemm(ck_tile::DeviceMem& a_m_k_dev_buf, std::cout << " Acc_Type = " << ck_tile::DataTypeTraits::name << " C_Type = " << ck_tile::DataTypeTraits::name << " QuantMode = " << quant_type_to_string(QuantMode) - << " PreshuffleQuant = " << (GemmConfig::PreshuffleQuant ? "true" : "false") << " : " - << " PreshuffleB = " << (GemmConfig::PreshuffleB ? "true" : "false") << " : " + << " APreshuffleQuant = " << (GemmConfig::APreshuffleQuant ? "true" : "false") + << " : " + << " BPreshuffleQuant = " << (GemmConfig::BPreshuffleQuant ? "true" : "false") + << " : " << " PreshuffleB = " << (GemmConfig::PreshuffleB ? "true" : "false") << " : " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, " << std::endl; @@ -536,21 +540,13 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, // Create BQ tensor with appropriate shape std::unique_ptr> bq_tensor_ptr = nullptr; if constexpr(QuantMode == ck_tile::QuantType::BQuantGrouped || - QuantMode == ck_tile::QuantType::RowColQuant) + QuantMode == ck_tile::QuantType::ABQuantGrouped || + QuantMode == ck_tile::QuantType::RowColQuant || + QuantMode == ck_tile::QuantType::TensorQuant) { bq_tensor_ptr = std::make_unique>( ck_tile::host_tensor_descriptor(BQK, BQN, stride_BQ, is_row_major(bq_layout))); } - else if constexpr(QuantMode == ck_tile::QuantType::ABQuantGrouped) - { - bq_tensor_ptr = std::make_unique>( - ck_tile::host_tensor_descriptor(BQK, BQN, stride_BQ, is_row_major(bq_layout))); - } - else if constexpr(QuantMode == ck_tile::QuantType::TensorQuant) - { - bq_tensor_ptr = std::make_unique>( - ck_tile::host_tensor_descriptor(1, 1, stride_BQ, is_row_major(bq_layout))); - } std::mt19937 gen(42); std::uniform_int_distribution fill_seed(0, 500); @@ -870,7 +866,7 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, QuantMode == ck_tile::QuantType::RowColQuant || QuantMode == ck_tile::QuantType::TensorQuant) { - if constexpr(GemmConfig::PreshuffleQuant) + if constexpr(GemmConfig::APreshuffleQuant) { ck_tile::HostTensor aq_shuffle_host = ck_tile::shuffle_aq(aq_tensor_ptr.get(), GemmConfig::K_Tile / AQuantGroupSize::kK); @@ -929,7 +925,7 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, ck_tile::HostTensor bq_permuted_host = ck_tile::bq_permuteN(*bq_tensor_ptr, BQuantGroupSize::kN); - if constexpr(GemmConfig::PreshuffleQuant) + if constexpr(GemmConfig::BPreshuffleQuant) { ck_tile::HostTensor bq_shuffle_host = ck_tile::shuffle_bq( &bq_permuted_host, GemmConfig::K_Tile / BQuantGroupSize::kK); @@ -940,7 +936,7 @@ int run_gemm_example_with_layouts(const ck_tile::ArgParser& arg_parser, bq_dev_buf_ptr->ToDevice(bq_permuted_host.data()); } } - else if constexpr(GemmConfig::PreshuffleQuant) + else if constexpr(GemmConfig::BPreshuffleQuant) { ck_tile::HostTensor bq_shuffle_host = ck_tile::shuffle_bq(bq_tensor_ptr.get(), GemmConfig::K_Tile / BQuantGroupSize::kK); @@ -1121,7 +1117,7 @@ int run_gemm_example_prec_type(const ck_tile::ArgParser& arg_parser) if constexpr((QuantMode == ck_tile::QuantType::AQuantGrouped || QuantMode == ck_tile::QuantType::ABQuantGrouped) && - !GemmConfig::PreshuffleQuant && !GemmConfig::PreshuffleB) + !GemmConfig::APreshuffleQuant && !GemmConfig::PreshuffleB) { if(a_layout == "R" && b_layout == "R") { @@ -1142,7 +1138,8 @@ int run_gemm_example_prec_type(const ck_tile::ArgParser& arg_parser) arg_parser, Col{}, Row{}, Row{}, Col{}, Row{}); } } - if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped && !GemmConfig::PreshuffleQuant) + if constexpr(QuantMode == ck_tile::QuantType::AQuantGrouped && + !GemmConfig::APreshuffleQuant) { if(a_layout == "C" && b_layout == "C") { diff --git a/include/ck_tile/ops/gemm_quant/block/block_gemm_quant_common.hpp b/include/ck_tile/ops/gemm_quant/block/block_gemm_quant_common.hpp index fb4a701eac..fcf1261754 100644 --- a/include/ck_tile/ops/gemm_quant/block/block_gemm_quant_common.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_gemm_quant_common.hpp @@ -96,9 +96,9 @@ struct AQPickerCommon : public BlockGemmQuantBase if constexpr(Traits::TransposeC) // transposed C { index_t reg_offset = - Traits::PreshuffleQuant ? mIter : mIter * Traits::AQPerBlock + kQScale; + Traits::APreshuffleQuant ? mIter : mIter * Traits::AQPerBlock + kQScale; auto scale_reg = aq_block_tensor.get_thread_buffer()[reg_offset]; - if constexpr(Traits::PreshuffleQuant) + if constexpr(Traits::APreshuffleQuant) { auto pull_from_lane = (__lane_id() & (Traits::WarpGemm::kN - 1)) * Traits::AQPerBlock + kQScale; @@ -121,7 +121,7 @@ struct AQPickerCommon : public BlockGemmQuantBase } else { - if constexpr(Traits::PreshuffleQuant) + if constexpr(Traits::APreshuffleQuant) { // A view is created on top of the preshuffled AQ, where each row of // the view is composed of a row from a warp tile within an AQ block diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_aquant_flatbr_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_aquant_flatbr_bquant_cr.hpp index b4a1bf886e..abb9de6cf3 100644 --- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_aquant_flatbr_bquant_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_aquant_flatbr_bquant_cr.hpp @@ -69,7 +69,8 @@ struct BlockGemmWeightPreshuffleABQuantARegBRegCReg static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + static constexpr bool APreshuffleQuant = Problem::Traits::APreshuffleQuant; + static constexpr bool BPreshuffleQuant = Problem::Traits::BPreshuffleQuant; static constexpr index_t QScalesPerBlockRow = integer_divide_ceil(KPerBlock, BQuantGroupSize::kK); @@ -127,9 +128,9 @@ struct BlockGemmWeightPreshuffleABQuantARegBRegCReg using CDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; // TileFlatmmShape - using QuantGroupSize = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; - static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant preshuffle kernel!"); + static_assert(BQuantGroupSize::kM == 1, "only N/K blocks for BQuant preshuffle kernel!"); static constexpr auto I0 = number<0>(); static constexpr auto I1 = number<1>(); @@ -162,12 +163,12 @@ struct BlockGemmWeightPreshuffleABQuantARegBRegCReg static constexpr auto MIter_2nd_last = (MIterPerWarp >= 2) ? MIterPerWarp - 2 : MIterPerWarp - 1; - static constexpr index_t KPerBlockBQ = KPerBlock / QuantGroupSize::kK; + static constexpr index_t KPerBlockBQ = KPerBlock / BQuantGroupSize::kK; static constexpr index_t QScalesPerBlockRow = - integer_divide_ceil(KPerBlock, QuantGroupSize::kK); // 128 / 128 = 1 + integer_divide_ceil(KPerBlock, BQuantGroupSize::kK); // 128 / 128 = 1 static constexpr index_t QScalesPerWarpGemmRow = - integer_divide_ceil(WG::kK, QuantGroupSize::kK); + integer_divide_ceil(WG::kK, BQuantGroupSize::kK); static constexpr index_t KIterPerQScale = KIterPerWarp / QScalesPerBlockRow; // 8 / 1 = 8 static constexpr index_t DsReadPreload = 2; // default 2, preload 2 ds read @@ -289,9 +290,9 @@ struct BlockGemmWeightPreshuffleABQuantARegBRegCReg CBlockTensor::PackedSize>{}; index_t reg_offset = [&]() { - if constexpr(QuantGroupSize::kN >= (NWarp * WG::kN)) + if constexpr(BQuantGroupSize::kN >= (NWarp * WG::kN)) { - return (nIter * NWarp * WG::kN) / QuantGroupSize::kN * KPerBlockBQ + + return (nIter * NWarp * WG::kN) / BQuantGroupSize::kN * KPerBlockBQ + kQScale; } else diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp index 2b67b65856..d2cfaca7b7 100644 --- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_ar_flatbr_bquant_cr.hpp @@ -25,9 +25,9 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg using CDataType = remove_cvref_t; using ComputeDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; // TileFlatmmShape - using QuantGroupSize = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; - static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant preshuffle kernel!"); + static_assert(BQuantGroupSize::kM == 1, "only N/K blocks for BQuant preshuffle kernel!"); static constexpr auto I0 = number<0>(); static constexpr auto I1 = number<1>(); @@ -53,7 +53,7 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg static constexpr index_t kBlockSize = Problem::kBlockSize; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + static constexpr bool BPreshuffleQuant = Problem::Traits::BPreshuffleQuant; static constexpr index_t MIterPerWarp = MPerBlock / (MWarp * WG::kM); static constexpr index_t NIterPerWarp = @@ -63,12 +63,12 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg static constexpr auto MIter_2nd_last = (MIterPerWarp >= 2) ? MIterPerWarp - 2 : MIterPerWarp - 1; - static constexpr index_t KPerBlockBQ = KPerBlock / QuantGroupSize::kK; + static constexpr index_t KPerBlockBQ = KPerBlock / BQuantGroupSize::kK; static constexpr index_t QScalesPerBlockRow = - integer_divide_ceil(KPerBlock, QuantGroupSize::kK); + integer_divide_ceil(KPerBlock, BQuantGroupSize::kK); static constexpr index_t QScalesPerWarpGemmRow = - integer_divide_ceil(WG::kK, QuantGroupSize::kK); + integer_divide_ceil(WG::kK, BQuantGroupSize::kK); static constexpr index_t KIterPerQScale = KIterPerWarp / QScalesPerBlockRow; static constexpr index_t DsReadPreload = 2; // default 2, preload 2 ds read @@ -173,7 +173,7 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg c_warp_y_index_zeros)) / CBlockTensor::PackedSize>{}; - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { constexpr index_t reg_offset = nIter; auto pull_from_lane = (__lane_id() & (WG::kN - 1)) * KPerBlockBQ + kQScale; @@ -205,9 +205,10 @@ struct BlockGemmWeightPreshuffleBQuantARegBRegCReg else { index_t reg_offset = [&]() { - if constexpr(QuantGroupSize::kN >= (NWarp * WG::kN)) + if constexpr(BQuantGroupSize::kN >= (NWarp * WG::kN)) { - return (nIter * NWarp * WG::kN) / QuantGroupSize::kN * KPerBlockBQ + + return (nIter * NWarp * WG::kN) / BQuantGroupSize::kN * + KPerBlockBQ + kQScale; } else diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_bquant_cr.hpp index 3fb80c21ff..2d28b813bf 100644 --- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_bquant_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_bquant_cr.hpp @@ -33,6 +33,7 @@ struct ABQuantBlockUniversalGemmAsBsCr : public BlockGemmQuantBase using AQDataType = remove_cvref_t; using BDataType = remove_cvref_t; using BQDataType = remove_cvref_t; + using BLayout = remove_cvref_t; using BQLayout = remove_cvref_t; using ComputeDataType = remove_cvref_t; using CDataType = remove_cvref_t; @@ -75,7 +76,8 @@ struct ABQuantBlockUniversalGemmAsBsCr : public BlockGemmQuantBase static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + static constexpr bool APreshuffleQuant = Problem::Traits::APreshuffleQuant; + static constexpr bool BPreshuffleQuant = Problem::Traits::BPreshuffleQuant; static constexpr index_t QScalesPerBlockRow = integer_divide_ceil(KPerBlock, BQuantGroupSize::kK); @@ -134,8 +136,12 @@ struct ABQuantBlockUniversalGemmAsBsCr : public BlockGemmQuantBase using CDataType = remove_cvref_t; // BDataType gets converted from PkInt4 during loading - using OverrideBDataType = - std::conditional_t, ADataType, BDataType>; + using OverrideBDataType = std::conditional_t< + std::is_same_v && + std::is_same_v, + ADataType, + BDataType>; + using Base = BlockGemmQuantBase; using WarpGemm = remove_cvref_t; @@ -156,7 +162,8 @@ struct ABQuantBlockUniversalGemmAsBsCr : public BlockGemmQuantBase using BWarpTensor = typename WarpGemm::BWarpTensor; using CWarpTensor = typename WarpGemm::CWarpTensor; - static constexpr bool PreshuffleQuant = Traits::PreshuffleQuant; + static constexpr bool APreshuffleQuant = Traits::APreshuffleQuant; + static constexpr bool BPreshuffleQuant = Traits::BPreshuffleQuant; static_assert(std::is_same_v); @@ -354,11 +361,24 @@ struct ABQuantBlockUniversalGemmAsBsCr : public BlockGemmQuantBase AQPickerCommon aq_picker( aq_block_tensor); - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { - constexpr index_t reg_offset = nIter; + constexpr index_t reg_offset = [&]() { + if constexpr(GemmTraits::BQuantGroupSize::kN > + (NWarp * WarpGemm::kN) && + Traits::NPerBlock == GemmTraits::BQuantGroupSize::kN) + { + return kQScale; + } + else + { + return nIter; + } + }(); + auto pull_from_lane = (__lane_id() & (WarpGemm::kN - 1)) * Traits::KQPerBlock + kQScale; + auto& scale_reg = bq_block_tensor.get_thread_buffer()[reg_offset]; // cross lane ops uint32_t scale_reg_dword; diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp index 9d19e902e5..22563da498 100644 --- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_aquant_bs_cr.hpp @@ -34,7 +34,7 @@ struct AQuantBlockUniversalGemmAsBsCr using ComputeDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using AQuantGroupSize = remove_cvref_t; static constexpr index_t kBlockSize = Problem::kBlockSize; static constexpr auto Scheduler = Problem::Scheduler; @@ -43,7 +43,7 @@ struct AQuantBlockUniversalGemmAsBsCr static constexpr index_t MPerBlock = BlockGemmShape::kM; static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; - static constexpr index_t AQPerBlock = KPerBlock / QuantGroupSize::kK; + static constexpr index_t AQPerBlock = KPerBlock / AQuantGroupSize::kK; static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); using WarpGemm = remove_cvref_t())>; @@ -69,20 +69,20 @@ struct AQuantBlockUniversalGemmAsBsCr static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; static constexpr index_t QScalesPerBlockRow = - integer_divide_ceil(KPerBlock, QuantGroupSize::kK); + integer_divide_ceil(KPerBlock, AQuantGroupSize::kK); static constexpr index_t QScalesPerWarpGemmRow = - integer_divide_ceil(WarpGemm::kK, QuantGroupSize::kK); + integer_divide_ceil(WarpGemm::kK, AQuantGroupSize::kK); static constexpr index_t KIterPerQScale = KIterPerWarp / QScalesPerBlockRow; - static_assert(QuantGroupSize::kK % WarpGemm::kK == 0, - "Error! WarpGemm::kK should be a multiple of QuantGroupSize"); + static_assert(AQuantGroupSize::kK % WarpGemm::kK == 0, + "Error! WarpGemm::kK should be a multiple of AQuantGroupSize"); static_assert(QScalesPerWarpGemmRow == 1, - "Error! QuantGroupSize shouldn't be smaller than WarpGemm::kK"); + "Error! AQuantGroupSize shouldn't be smaller than WarpGemm::kK"); static_assert(KIterPerWarp % QScalesPerBlockRow == 0, "Error! KItersPerWarp should be a multiple of QscalesPerBlockRow"); - static_assert(KPerBlock / QuantGroupSize::kK > 0, + static_assert(KPerBlock / AQuantGroupSize::kK > 0, "Error! Each row of blockgemm should have a separate scale"); static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, @@ -110,8 +110,8 @@ struct AQuantBlockUniversalGemmAsBsCr static constexpr index_t KPack = WarpGemm::kKPerThread; static constexpr index_t KPerThread = KIterPerWarp * WarpGemm::kKPerThread; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; - static constexpr bool TransposeC = Problem::TransposeC; + static constexpr bool APreshuffleQuant = Problem::Traits::APreshuffleQuant; + static constexpr bool TransposeC = Problem::TransposeC; }; public: diff --git a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp index 03b9dfe34d..9d711c4862 100644 --- a/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp +++ b/include/ck_tile/ops/gemm_quant/block/block_universal_gemm_as_bs_bquant_cr.hpp @@ -36,7 +36,7 @@ struct BQuantBlockUniversalGemmAsBsCr using ComputeDataType = remove_cvref_t; using CDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; static constexpr index_t kBlockSize = Problem::kBlockSize; static constexpr auto Scheduler = Problem::Scheduler; @@ -46,8 +46,8 @@ struct BQuantBlockUniversalGemmAsBsCr static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; - static constexpr index_t NQPerBlock = NPerBlock / QuantGroupSize::kN; - static constexpr index_t KQPerBlock = KPerBlock / QuantGroupSize::kK; + static constexpr index_t NQPerBlock = NPerBlock / BQuantGroupSize::kN; + static constexpr index_t KQPerBlock = KPerBlock / BQuantGroupSize::kK; static constexpr auto config = Policy::template GetWarpGemmMWarpNWarp(); using WarpGemm = remove_cvref_t())>; @@ -72,23 +72,23 @@ struct BQuantBlockUniversalGemmAsBsCr static constexpr index_t NIterPerWarp = NPerBlock / (NWarp * WarpGemm::kN); static constexpr index_t KIterPerWarp = KPerBlock / WarpGemm::kK; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + static constexpr bool BPreshuffleQuant = Problem::Traits::BPreshuffleQuant; static constexpr index_t QScalesPerBlockRow = - integer_divide_ceil(KPerBlock, QuantGroupSize::kK); + integer_divide_ceil(KPerBlock, BQuantGroupSize::kK); static constexpr index_t QScalesPerWarpGemmRow = - integer_divide_ceil(WarpGemm::kK, QuantGroupSize::kK); + integer_divide_ceil(WarpGemm::kK, BQuantGroupSize::kK); static constexpr index_t KIterPerQScale = KIterPerWarp / QScalesPerBlockRow; - static_assert(QuantGroupSize::kK % WarpGemm::kK == 0, - "Error! WarpGemm::kK should be a multiple of QuantGroupSize"); + static_assert(BQuantGroupSize::kK % WarpGemm::kK == 0, + "Error! WarpGemm::kK should be a multiple of BQuantGroupSize"); static_assert(QScalesPerWarpGemmRow == 1, - "Error! QuantGroupSize shouldn't be smaller than WarpGemm::kK"); + "Error! BQuantGroupSize shouldn't be smaller than WarpGemm::kK"); static_assert(KIterPerWarp % QScalesPerBlockRow == 0, "Error! KItersPerWarp should be a multiple of QscalesPerBlockRow"); - static_assert(KPerBlock / QuantGroupSize::kK > 0, + static_assert(KPerBlock / BQuantGroupSize::kK > 0, "Error! Each row of blockgemm should have a separate scale"); static_assert(MIterPerWarp * MWarp * WarpGemm::kM == MPerBlock, @@ -153,7 +153,7 @@ struct BQuantBlockUniversalGemmAsBsCr using BWarpTensor = typename WarpGemm::BWarpTensor; using CWarpTensor = typename WarpGemm::CWarpTensor; - static constexpr bool PreshuffleQuant = Traits::PreshuffleQuant; + static constexpr bool BPreshuffleQuant = Traits::BPreshuffleQuant; static_assert(std::is_same_v); @@ -317,25 +317,21 @@ struct BQuantBlockUniversalGemmAsBsCr c_warp_y_index_zeros)) / CBlockTensor::PackedSize>{}; - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { - // constexpr index_t reg_offset = nIter; constexpr index_t reg_offset = [&]() { - if constexpr(GemmTraits::QuantGroupSize::kN > - (NWarp * WarpGemm::kN)) + if constexpr(GemmTraits::BQuantGroupSize::kN > + (NWarp * WarpGemm::kN) && + Traits::NPerBlock == GemmTraits::BQuantGroupSize::kN) { - if constexpr(Traits::NPerBlock == - GemmTraits::QuantGroupSize::kN) - return kQScale; - else - return nIter; // for prefill needs kQscale, for decode needs - // nIter + return kQScale; // prefill: one quant group per block } else { - return nIter; + return nIter; // decode or multiple groups per warp } }(); + auto pull_from_lane = (__lane_id() & (WarpGemm::kN - 1)) * Traits::KQPerBlock + kQScale; @@ -370,10 +366,11 @@ struct BQuantBlockUniversalGemmAsBsCr { // Multiply bquant with accumulated C constexpr index_t reg_offset = [&]() { - if constexpr(GemmTraits::QuantGroupSize::kN >= + if constexpr(GemmTraits::BQuantGroupSize::kN >= (NWarp * WarpGemm::kN)) return (nIter * NWarp * WarpGemm::kN) / - GemmTraits::QuantGroupSize::kN * Traits::KQPerBlock + + GemmTraits::BQuantGroupSize::kN * + Traits::KQPerBlock + kQScale; else { diff --git a/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp index ce8eb94a1b..21bd691b49 100644 --- a/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp +++ b/include/ck_tile/ops/gemm_quant/kernel/gemm_quant_kernel.hpp @@ -67,15 +67,27 @@ struct get_bq_data_type_or> }; template -struct is_quantpreshuffle_enabled +struct is_Aquantpreshuffle_enabled { static constexpr bool value = false; }; template -struct is_quantpreshuffle_enabled> +struct is_Aquantpreshuffle_enabled> { - static constexpr bool value = T::PreshuffleQuant; + static constexpr bool value = T::APreshuffleQuant; +}; + +template +struct is_Bquantpreshuffle_enabled +{ + static constexpr bool value = false; +}; + +template +struct is_Bquantpreshuffle_enabled> +{ + static constexpr bool value = T::BPreshuffleQuant; }; template @@ -206,8 +218,10 @@ struct QuantGemmKernel typename detail::get_bq_layout_or::type>; static constexpr index_t kBlockSize = GemmPipeline::BlockSize; - static constexpr bool PreshuffleQuant = - detail::is_quantpreshuffle_enabled::value; + static constexpr bool APreshuffleQuant = + detail::is_Aquantpreshuffle_enabled::value; + static constexpr bool BPreshuffleQuant = + detail::is_Bquantpreshuffle_enabled::value; static constexpr bool PreshuffleB = detail::is_preshuffleB_enabled::value; using ADataType = remove_cvref_t; @@ -476,7 +490,7 @@ struct QuantGemmKernel { // Step 1: Create tensor view for AQ const auto& aq_tensor_view = [&]() { - if constexpr(kQuantType == QuantType::AQuantGrouped && PreshuffleQuant) + if constexpr(kQuantType == QuantType::AQuantGrouped && APreshuffleQuant) { static_assert(std::is_same_v); const auto aq_x = kargs.M * GemmPipeline::KPerBlockAQ; @@ -533,7 +547,7 @@ struct QuantGemmKernel } else if constexpr((kQuantType == QuantType::AQuantGrouped || kQuantType == QuantType::ABQuantGrouped) && - !PreshuffleQuant) + !APreshuffleQuant) { if constexpr(std::is_same_v) { @@ -571,13 +585,13 @@ struct QuantGemmKernel // Step 2: Create tile window (no padding for AQ) const auto& aq_block_window = [&]() { - if constexpr(kQuantType == QuantType::AQuantGrouped && PreshuffleQuant) + if constexpr(kQuantType == QuantType::AQuantGrouped && APreshuffleQuant) { static_assert(std::is_same_v); - using QuantGroupSize = remove_cvref_t; + using AQuantGroupSize = remove_cvref_t; constexpr auto block_m = TilePartitioner::MPerBlock; constexpr auto warp_m = GemmPipeline::BlockGemmShape::WarpTile::at(I0); - constexpr auto aqk_per_block = TilePartitioner::KPerBlock / QuantGroupSize::kK; + constexpr auto aqk_per_block = TilePartitioner::KPerBlock / AQuantGroupSize::kK; constexpr auto tile_window_width = ck_tile::integer_least_multiple(warp_m * aqk_per_block, get_warp_size()); constexpr auto tile_window_height = block_m / warp_m; @@ -587,11 +601,19 @@ struct QuantGemmKernel make_tuple(number{}, number{}), {block_m_idx * tile_window_height, 0}); } - else if constexpr(kQuantType == QuantType::AQuantGrouped && !PreshuffleQuant) + else if constexpr((kQuantType == QuantType::AQuantGrouped || + kQuantType == QuantType::ABQuantGrouped) && + !APreshuffleQuant) { - using QuantGroupSize = remove_cvref_t; - constexpr auto aqk_per_block = TilePartitioner::KPerBlock / QuantGroupSize::kK; + + using AQuantGroupSize = remove_cvref_t; + constexpr auto aqk_per_block = TilePartitioner::KPerBlock / AQuantGroupSize::kK; constexpr auto block_m = TilePartitioner::MPerBlock; + if constexpr(kQuantType == QuantType::ABQuantGrouped) + { + static_assert(std::is_same_v, + "ABQuantGrouped requires RowMajor AQ layout"); + } if constexpr(std::is_same_v) { return make_tile_window(aq_tensor_view, @@ -605,17 +627,6 @@ struct QuantGemmKernel {0, i_m}); } } - else if constexpr(kQuantType == QuantType::ABQuantGrouped && !PreshuffleQuant) - { - static_assert(std::is_same_v); - using QuantGroupSize = remove_cvref_t; - constexpr auto block_m = TilePartitioner::MPerBlock; - constexpr auto block_k = TilePartitioner::KPerBlock; - return make_tile_window( - aq_tensor_view, - make_tuple(number{}, number{}), - {i_m, 0}); - } else if constexpr(kQuantType == QuantType::RowColQuant) { return make_tile_window(aq_tensor_view, @@ -808,14 +819,15 @@ struct QuantGemmKernel number<1>{}, number<1>{}); } - else if constexpr(kQuantType == QuantType::BQuantGrouped) + else if constexpr(kQuantType == QuantType::BQuantGrouped || + kQuantType == QuantType::ABQuantGrouped) { - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { static_assert(std::is_same_v, "PreshuffleQuant with BQuantGrouped currently only supports " "ColumnMajor BQ layout"); - using QuantGroupSize = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; return MakePreshuffledQuantTensorView< GemmPipeline::KPerBlockBQ, @@ -824,48 +836,42 @@ struct QuantGemmKernel TilePartitioner::BlockGemmShape::WarpTile::at(I1), GemmPipeline::GetVectorSizeBQ()>( bq_ptr, - ck_tile::integer_divide_ceil(kargs.N, QuantGroupSize::kN), - QuantGroupSize::kN, + ck_tile::integer_divide_ceil(kargs.N, BQuantGroupSize::kN), + BQuantGroupSize::kN, kargs.QK_B); } else { - using QuantGroupSize = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; + + if constexpr(kQuantType == QuantType::ABQuantGrouped) + { + static_assert(std::is_same_v, + "ABQuantGrouped requires ColumnMajor BQ layout"); + } if constexpr(std::is_same_v) { return make_naive_tensor_view( bq_ptr, - make_tuple(integer_divide_ceil(kargs.K, QuantGroupSize::kK), - integer_divide_ceil(kargs.N, QuantGroupSize::kN)), - make_tuple(integer_divide_ceil(kargs.N, QuantGroupSize::kN), 1), + make_tuple(integer_divide_ceil(kargs.K, BQuantGroupSize::kK), + integer_divide_ceil(kargs.N, BQuantGroupSize::kN)), + make_tuple(integer_divide_ceil(kargs.N, BQuantGroupSize::kN), 1), number{}, number<1>{}); } else { - static_assert(std::is_same_v); return make_naive_tensor_view( bq_ptr, - make_tuple(integer_divide_ceil(kargs.N, QuantGroupSize::kN), - integer_divide_ceil(kargs.K, QuantGroupSize::kK)), - make_tuple(integer_divide_ceil(kargs.K, QuantGroupSize::kK), 1), + make_tuple(integer_divide_ceil(kargs.N, BQuantGroupSize::kN), + integer_divide_ceil(kargs.K, BQuantGroupSize::kK)), + make_tuple(integer_divide_ceil(kargs.K, BQuantGroupSize::kK), 1), number{}, number<1>{}); } } } - else if constexpr(kQuantType == QuantType::ABQuantGrouped) - { - static_assert(std::is_same_v); - using QuantGroupSize = remove_cvref_t; - return make_naive_tensor_view( - bq_ptr, - make_tuple(integer_divide_ceil(kargs.N, QuantGroupSize::kN), kargs.QK_B), - make_tuple(kargs.stride_BQ, 1), - number{}, - number<1>{}); - } else { return nullptr; @@ -881,28 +887,29 @@ struct QuantGemmKernel number{}), {i_m, i_n}); } - else if constexpr(kQuantType == QuantType::BQuantGrouped) + else if constexpr(kQuantType == QuantType::BQuantGrouped || + kQuantType == QuantType::ABQuantGrouped) { - using QuantGroupSize = remove_cvref_t; - if constexpr(PreshuffleQuant) + using BQuantGroupSize = remove_cvref_t; + if constexpr(BPreshuffleQuant) { static_assert(std::is_same_v); // Number of N-dimension quantization groups per block - constexpr auto block_n = (QuantGroupSize::kN <= TilePartitioner::NPerBlock) - ? TilePartitioner::NPerBlock / QuantGroupSize::kN - : QuantGroupSize::kN / TilePartitioner::NPerBlock; + constexpr auto block_n = (BQuantGroupSize::kN <= TilePartitioner::NPerBlock) + ? TilePartitioner::NPerBlock / BQuantGroupSize::kN + : BQuantGroupSize::kN / TilePartitioner::NPerBlock; // Number of N-dimension elements per warp constexpr auto warp_n = TilePartitioner::BlockGemmShape::WarpTile::at(I1); // Determine how many warps share the same scale in N-dimension - constexpr auto warp_per_group = (QuantGroupSize::kN < warp_n) - ? (warp_n / QuantGroupSize::kN) - : (QuantGroupSize::kN / warp_n); + constexpr auto warp_per_group = (BQuantGroupSize::kN < warp_n) + ? (warp_n / BQuantGroupSize::kN) + : (BQuantGroupSize::kN / warp_n); // Number of K-dimension quantization groups per block - constexpr auto bqk_per_block = TilePartitioner::KPerBlock / QuantGroupSize::kK; + constexpr auto bqk_per_block = TilePartitioner::KPerBlock / BQuantGroupSize::kK; // The pre-shuffled layout flattens warp_n × // bqk_per_block scales per row, Padded up to warp_size @@ -911,25 +918,25 @@ struct QuantGemmKernel ck_tile::integer_least_multiple(warp_n * bqk_per_block, get_warp_size()); // Adapts based on fine vs coarse quantization granularity: - // - Fine-grained (QuantGroupSize::kN < warp_n): + // - Fine-grained (BQuantGroupSize::kN < warp_n): // Multiple quant groups per warp → fewer rows needed per block. // height = block_n / warp_per_group // - // - Coarse-grained (QuantGroupSize::kN >= warp_n): + // - Coarse-grained (BQuantGroupSize::kN >= warp_n): // Each row represents one quant group. // height = block_n constexpr auto tile_window_height = - (QuantGroupSize::kN < warp_n) ? block_n / warp_per_group : block_n; + (BQuantGroupSize::kN < warp_n) ? block_n / warp_per_group : block_n; auto block_n_idx = i_n / TilePartitioner::NPerBlock; // For decode shapes GN: 128, Blocks needs to repeat 0,0,1,1,2,2 ... - if(QuantGroupSize::kN > TilePartitioner::NPerBlock) + if(BQuantGroupSize::kN > TilePartitioner::NPerBlock) { block_n_idx = block_n_idx >> 1; } - if(QuantGroupSize::kN > TilePartitioner::NPerBlock) + if(BQuantGroupSize::kN > TilePartitioner::NPerBlock) { return make_tile_window( bq_tensor_view, @@ -946,17 +953,22 @@ struct QuantGemmKernel } else { + if constexpr(kQuantType == QuantType::ABQuantGrouped) + { + static_assert(std::is_same_v, + "ABQuantGrouped requires RowMajor AQ layout"); + } constexpr auto tensor_dim = - (QuantGroupSize::kN <= TilePartitioner::NPerBlock) - ? TilePartitioner::NPerBlock / QuantGroupSize::kN + (BQuantGroupSize::kN <= TilePartitioner::NPerBlock) + ? TilePartitioner::NPerBlock / BQuantGroupSize::kN : 1; if constexpr(std::is_same_v) { return make_tile_window( bq_tensor_view, - make_tuple(number{}, + make_tuple(number{}, number{}), - {0, i_n / QuantGroupSize::kN}); + {0, i_n / BQuantGroupSize::kN}); } else { @@ -964,21 +976,11 @@ struct QuantGemmKernel return make_tile_window( bq_tensor_view, make_tuple(number{}, - number{}), - {i_n / QuantGroupSize::kN, 0}); + number{}), + {i_n / BQuantGroupSize::kN, 0}); } } } - else if constexpr(kQuantType == QuantType::ABQuantGrouped) - { - static_assert(std::is_same_v); - using QuantGroupSize = remove_cvref_t; - return make_tile_window( - bq_tensor_view, - make_tuple(number{}, - number{}), - {i_n / QuantGroupSize::kN, 0}); - } else { return nullptr; @@ -1223,7 +1225,7 @@ struct QuantGemmKernel if constexpr(kQuantType == QuantType::AQuantGrouped) { index_t m = 0; - if constexpr(PreshuffleQuant) + if constexpr(APreshuffleQuant) { m = kargs.M; } @@ -1233,7 +1235,7 @@ struct QuantGemmKernel else if constexpr(kQuantType == QuantType::BQuantGrouped) { index_t n = 0; - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { n = kargs.N; } @@ -1244,9 +1246,9 @@ struct QuantGemmKernel { index_t m = 0; index_t n = 0; - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { - m = kargs.M; + // m = kargs.M; n = kargs.N; } return GemmPipeline{}.template operator()(a_block_window, diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_abquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_abquant_pipeline_ag_bg_cr_v3.hpp index cd70c2ca86..5902dd0c4f 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_abquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_abquant_pipeline_ag_bg_cr_v3.hpp @@ -72,7 +72,10 @@ struct ABQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3(); } @@ -95,7 +98,8 @@ struct ABQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3{}), - 0) + (BPreshuffleQuant) + ? make_array(((NPerBlockBQ <= BlockGemmShape::BlockWarps::at(number<1>{})) + ? ck_tile::integer_divide_ceil(n, BQuantGroupSize::kN) + : ck_tile::integer_least_multiple(n, NPerBlock) / + BlockGemmShape::WarpTile::at(number<1>{})), + 0) : is_bq_row_major ? make_array(KPerBlockBQ, 0) : make_array(0, KPerBlockBQ); @@ -484,7 +491,7 @@ struct ABQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3( Policy::template MakeShuffledARegTileDistribution()); @@ -495,7 +502,7 @@ struct ABQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3( diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp index 1acf0444cf..9ddb7eecac 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_base.hpp @@ -12,21 +12,21 @@ namespace ck_tile { template struct GemmAQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase { - using Base = GemmPipelineAgBgCrImplBase; - using ADataType = typename Base::ADataType; - using ALayout = typename Base::ALayout; - using BDataType = typename Base::BDataType; - using BLayout = typename Base::BLayout; - using BlockGemmShape = typename Base::BlockGemmShape; - using QuantGroupSize = remove_cvref_t; + using Base = GemmPipelineAgBgCrImplBase; + using ADataType = typename Base::ADataType; + using ALayout = typename Base::ALayout; + using BDataType = typename Base::BDataType; + using BLayout = typename Base::BLayout; + using BlockGemmShape = typename Base::BlockGemmShape; + using AQuantGroupSize = remove_cvref_t; static constexpr index_t MPerBlock = BlockGemmShape::kM; static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; - static constexpr index_t KPerBlockAQ = KPerBlock / QuantGroupSize::kK; + static constexpr index_t KPerBlockAQ = KPerBlock / AQuantGroupSize::kK; - static_assert(KPerBlock % QuantGroupSize::kK == 0, + static_assert(KPerBlock % AQuantGroupSize::kK == 0, "KPerBlock must be a multiple of QuantGroupSize"); // Create DRAM tile window for AQ diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp index b87c12c14a..ccd4758e34 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_mem.hpp @@ -23,19 +23,19 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem using Base = BaseGemmPipelineAgBgCrMem; using PipelineImplBase = GemmAQuantPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using AQDataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using ADataType = remove_cvref_t; + using AQDataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using AQuantGroupSize = remove_cvref_t; // When ADataType is pk_int4_t, use BDataType instead for transpose operations // since packed 4-bit integers cannot be directly transposed (requires at least 8-bit precision) using OverrideADataType = std::conditional_t, BDataType, ADataType>; - static_assert(QuantGroupSize::kM == 1, "no block for M supported yet!"); - static_assert(QuantGroupSize::kN == 1, "only M/K blocks for AQuant kernel!"); + static_assert(AQuantGroupSize::kM == 1, "no block for M supported yet!"); + static_assert(AQuantGroupSize::kN == 1, "only M/K blocks for AQuant kernel!"); using I0 = number<0>; using I1 = number<1>; @@ -60,7 +60,7 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem static constexpr index_t MPerBlock = BlockGemmShape::kM; static constexpr index_t NPerBlock = BlockGemmShape::kN; static constexpr index_t KPerBlock = BlockGemmShape::kK; - static constexpr index_t KPerBlockAQ = BlockGemmShape::kK / QuantGroupSize::kK; + static constexpr index_t KPerBlockAQ = BlockGemmShape::kK / AQuantGroupSize::kK; static constexpr index_t GetVectorSizeA() { return Policy::template GetVectorSizeA(); } static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } @@ -78,7 +78,7 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem static constexpr bool kPadK = Problem::kPadK; static constexpr bool DoubleSmemBuffer = Problem::DoubleSmemBuffer; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + static constexpr bool APreshuffleQuant = Problem::Traits::APreshuffleQuant; static constexpr bool HasHotLoop = Problem::HasHotLoop; static constexpr auto TailNum = Problem::TailNum; @@ -99,7 +99,7 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem BlockSize, concat('x', WaveNumM, WaveNumN), concat('x', BlockGemm::WarpGemm::kM, BlockGemm::WarpGemm::kN, BlockGemm::WarpGemm::kK), - concat('x', kPadM, kPadN, kPadK), QuantGroupSize::GetName(), + concat('x', kPadM, kPadN, kPadK), AQuantGroupSize::GetName(), Scheduler == GemmPipelineScheduler::Interwave ? "interwave" : "intrawave"); // else Intrawave // clang-format on } @@ -156,7 +156,7 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem << "\n" << "A/B LDS read inst: " << A_LDS_Read_Inst_Num << ", " << B_LDS_Read_Inst_Num << "\n" << "C MFMA inst: " << C_MFMA_Inst_Num << "\n" - << "QuantGroupSize: " << QuantGroupSize::GetName() << "\n" + << "AQuantGroupSize: " << AQuantGroupSize::GetName() << "\n" << "KPack: " << BlockGemm::Traits::KPack << "\n" << "PrefetchStages: " << PrefetchStages << "\n"; return str.str(); @@ -216,7 +216,7 @@ struct AQuantGemmPipelineAgBgCrMem : public BaseGemmPipelineAgBgCrMem std::is_same_v; constexpr bool is_b_row_major = std::is_same_v; - static_assert(!PreshuffleQuant, "Memory pipeline does not support PreshuffleQuant!"); + static_assert(!APreshuffleQuant, "Memory pipeline does not support APreshuffleQuant!"); static_assert(is_a_col_major ? (KPerBlock == ADramBlockWindowTmp{}.get_window_lengths()[I0{}] && diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp index 4485144f5f..27828cce63 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_policy.hpp @@ -32,22 +32,22 @@ struct GemmAQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC using AQLayout = remove_cvref_t; using BlockGemmShape = typename Problem::BlockGemmShape; - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPerBlockAQ = KPerBlock / Problem::AQuantGroupSize::kK; - constexpr index_t VecLoadSize = GetVectorSizeAQ(); - constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; - using WarpTile = typename Problem::BlockGemmShape::WarpTile; - using WarpGemm = WarpGemmDispatcher; + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t MPerBlock = Problem::BlockGemmShape::kM; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t KPerBlockAQ = KPerBlock / Problem::AQuantGroupSize::kK; + constexpr index_t VecLoadSize = GetVectorSizeAQ(); + constexpr bool APreshuffleQuant = Problem::Traits::APreshuffleQuant; + using WarpTile = typename Problem::BlockGemmShape::WarpTile; + using WarpGemm = WarpGemmDispatcher; - if constexpr(PreshuffleQuant) + if constexpr(APreshuffleQuant) { using TileEncodingPattern = tile_distribution_encoding_pattern_aq< BlockGemmShape, @@ -57,7 +57,7 @@ struct GemmAQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC ck_tile::integer_least_multiple(WarpGemm::kM * KPerBlockAQ, get_warp_size()), KPerBlockAQ, VecLoadSize, - PreshuffleQuant>; + APreshuffleQuant>; return TileEncodingPattern::make_2d_static_tile_distribution(); } @@ -89,7 +89,7 @@ struct GemmAQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC KPerBlockAQ, KPerBlockAQ, VecLoadSize, - PreshuffleQuant>; + APreshuffleQuant>; return TileEncodingPattern::make_2d_static_tile_distribution(); } @@ -103,7 +103,7 @@ struct GemmAQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC MPerBlock, // XPerTile KPerBlockAQ, VecLoadSize, - PreshuffleQuant>; + APreshuffleQuant>; return TileEncodingPattern::make_2d_static_tile_distribution_transposed(); } } diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp index 71e4a74400..76d8985fb1 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_aquant_pipeline_ag_bg_cr_v3.hpp @@ -20,19 +20,19 @@ struct AQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3; using PipelineImplBase = GemmAQuantPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using AQDataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using ADataType = remove_cvref_t; + using AQDataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using AQuantGroupSize = remove_cvref_t; // When ADataType is pk_int4_t, use BDataType instead for transpose operations // since packed 4-bit integers cannot be directly transposed (requires at least 8-bit precision) using OverrideADataType = std::conditional_t, BDataType, ADataType>; - static_assert(QuantGroupSize::kM == 1, "no block for M supported yet!"); - static_assert(QuantGroupSize::kN == 1, "only M/K blocks for AQuant kernel!"); + static_assert(AQuantGroupSize::kM == 1, "no block for M supported yet!"); + static_assert(AQuantGroupSize::kN == 1, "only M/K blocks for AQuant kernel!"); using I0 = number<0>; using I1 = number<1>; @@ -57,7 +57,7 @@ struct AQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3(); } static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } @@ -75,7 +75,7 @@ struct AQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3 struct GemmBQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase { - using Base = GemmPipelineAgBgCrImplBase; - using ADataType = typename Base::ADataType; - using ALayout = typename Base::ALayout; - using BDataType = typename Base::BDataType; - using BLayout = typename Base::BLayout; - using BlockGemmShape = typename Base::BlockGemmShape; - using QuantGroupSize = remove_cvref_t; + using Base = GemmPipelineAgBgCrImplBase; + using ADataType = typename Base::ADataType; + using ALayout = typename Base::ALayout; + using BDataType = typename Base::BDataType; + using BLayout = typename Base::BLayout; + using BlockGemmShape = typename Base::BlockGemmShape; + using BQuantGroupSize = remove_cvref_t; using BQLayout = remove_cvref_t; @@ -27,16 +27,16 @@ struct GemmBQuantPipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase= 1, "NPerBlock must be >= QuantGroupSize"); - static_assert(KPerBlockBQ >= 1, "KPerBlock must be >= QuantGroupSize"); + // static_assert(NPerBlockBQ >= 1, "NPerBlock must be >= BQuantGroupSize"); + static_assert(KPerBlockBQ >= 1, "KPerBlock must be >= BQuantGroupSize"); - // static_assert(NPerBlock % QuantGroupSize::kN == 0, - // "NPerBlock must be a multiple of QuantGroupSize::kN"); - static_assert(KPerBlock % QuantGroupSize::kK == 0, - "KPerBlock must be a multiple of QuantGroupSize::kK"); + // static_assert(NPerBlock % BQuantGroupSize::kN == 0, + // "NPerBlock must be a multiple of BQuantGroupSize::kN"); + static_assert(KPerBlock % BQuantGroupSize::kK == 0, + "KPerBlock must be a multiple of BQuantGroupSize::kK"); // Create DRAM tile window for BQ template diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp index 5c4dfd37c7..08139324fc 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_policy.hpp @@ -43,14 +43,14 @@ struct GemmBQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC using BQLayout = remove_cvref_t; using BlockGemmShape = typename Problem::BlockGemmShape; - constexpr index_t BlockSize = Problem::kBlockSize; - constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; - constexpr index_t NPerBlockBQ = (Problem::QuantGroupSize::kN <= NPerBlock) - ? NPerBlock / Problem::QuantGroupSize::kN - : 1; - constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPerBlockBQ = KPerBlock / Problem::BQuantGroupSize::kK; - constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + constexpr index_t BlockSize = Problem::kBlockSize; + constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; + constexpr index_t NPerBlockBQ = (Problem::BQuantGroupSize::kN <= NPerBlock) + ? NPerBlock / Problem::BQuantGroupSize::kN + : 1; + constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; + constexpr index_t KPerBlockBQ = KPerBlock / Problem::BQuantGroupSize::kK; + constexpr bool BPreshuffleQuant = Problem::Traits::BPreshuffleQuant; using WarpTile = typename Problem::BlockGemmShape::WarpTile; using WarpGemm = WarpGemmDispatcher; - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { using TileEncodingPattern = tile_distribution_encoding_pattern_bq< BlockGemmShape, @@ -72,7 +72,7 @@ struct GemmBQuantPipelineAgBgCrDefaultPolicy : public UniversalGemmPipelineAgBgC Problem::BQuantGroupSize::kN, Problem::BQuantGroupSize::kK, BQLayout, - PreshuffleQuant>; + BPreshuffleQuant>; return TileEncodingPattern::make_2d_static_tile_distribution(); } else diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp index be91002cdb..df94eb7273 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_bquant_pipeline_ag_bg_cr_v3.hpp @@ -26,12 +26,12 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3; using PipelineImplBase = GemmBQuantPipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using BQDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using BQDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; using ALayout = remove_cvref_t; using BQLayout = remove_cvref_t; @@ -45,7 +45,7 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3; - static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant kernel!"); + static_assert(BQuantGroupSize::kM == 1, "only N/K blocks for BQuant kernel!"); using I0 = number<0>; using I1 = number<1>; using I2 = number<2>; @@ -66,11 +66,11 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3(); } static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } @@ -88,7 +88,7 @@ struct BQuantGemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3{})) - ? ck_tile::integer_divide_ceil(n, QuantGroupSize::kN) + ? ck_tile::integer_divide_ceil(n, BQuantGroupSize::kN) : ck_tile::integer_least_multiple(n, NPerBlock) / BlockGemmShape::WarpTile::at(number<1>{})), 0) diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp index bde0be89c0..48c27945b3 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_group_quant_utils.hpp @@ -52,7 +52,7 @@ template + bool APreshuffleQuant> struct tile_distribution_encoding_pattern_aq : public tile_distribution_encoding_pattern { static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!"); @@ -72,7 +72,7 @@ struct tile_distribution_encoding_pattern_aq : public tile_distribution_encoding CK_TILE_HOST_DEVICE static constexpr auto make_2d_static_tile_distribution() { - if constexpr(PreshuffleQuant) + if constexpr(APreshuffleQuant) { // # of elements per thread static_assert(XPerTile >= warp_size && XPerTile % warp_size == 0); @@ -193,8 +193,8 @@ template + typename BQLayout = tensor_layout::gemm::ColumnMajor, + bool BPreshuffleQuant = false> struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding_pattern { static constexpr index_t warp_size = get_warp_size(); @@ -212,10 +212,11 @@ struct tile_distribution_encoding_pattern_bq : public tile_distribution_encoding CK_TILE_HOST_DEVICE static constexpr auto make_2d_static_tile_distribution() { // Preshuffle only supported for ColumnMajor currently - static_assert(!(PreshuffleQuant && std::is_same_v), - "PreshuffleQuant only supported for ColumnMajor BQLayout"); + static_assert( + !(BPreshuffleQuant && std::is_same_v), + "PreshuffleQuant only supported for ColumnMajor BQLayout"); - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { // ============================================================================= // PRE-SHUFFLED BQ SCALE TILE DISTRIBUTION diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_base.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_base.hpp index 95122630ee..facec252a3 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_base.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_base.hpp @@ -12,13 +12,13 @@ namespace ck_tile { template struct GemmMxFp4PipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase { - using Base = GemmPipelineAgBgCrImplBase; - using ADataType = typename Base::ADataType; - using ALayout = typename Base::ALayout; - using BDataType = typename Base::BDataType; - using BLayout = typename Base::BLayout; - using BlockGemmShape = typename Base::BlockGemmShape; - using QuantGroupSize = remove_cvref_t; + using Base = GemmPipelineAgBgCrImplBase; + using ADataType = typename Base::ADataType; + using ALayout = typename Base::ALayout; + using BDataType = typename Base::BDataType; + using BLayout = typename Base::BLayout; + using BlockGemmShape = typename Base::BlockGemmShape; + using BQuantGroupSize = remove_cvref_t; using BQLayout = remove_cvref_t; @@ -26,16 +26,16 @@ struct GemmMxFp4PipelineAgBgCrImplBase : public GemmPipelineAgBgCrImplBase= 1, "NPerBlock must be >= QuantGroupSize"); - static_assert(KPerBlockBQ >= 1, "KPerBlock must be >= QuantGroupSize"); + static_assert(NPerBlockBQ >= 1, "NPerBlock must be >= BQuantGroupSize"); + static_assert(KPerBlockBQ >= 1, "KPerBlock must be >= BQuantGroupSize"); - static_assert(NPerBlock % QuantGroupSize::kN == 0, - "NPerBlock must be a multiple of QuantGroupSize::kN"); - static_assert(KPerBlock % QuantGroupSize::kK == 0, - "KPerBlock must be a multiple of QuantGroupSize::kK"); + static_assert(NPerBlock % BQuantGroupSize::kN == 0, + "NPerBlock must be a multiple of BQuantGroupSize::kN"); + static_assert(KPerBlock % BQuantGroupSize::kK == 0, + "KPerBlock must be a multiple of BQuantGroupSize::kK"); // Create DRAM tile window for BQ template diff --git a/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_policy.hpp b/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_policy.hpp index 7a2d1db2c8..6cf9e22f41 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_policy.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/gemm_mxfp4_pipeline_ag_bg_cr_policy.hpp @@ -22,9 +22,9 @@ struct GemmMxFp4PipelineAgBgCrPolicy : public UniversalGemmPipelineAgBgCrPolicy using BQLayout = remove_cvref_t; using BQDataType = remove_cvref_t; constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; - constexpr index_t NPerBlockBQ = NPerBlock / Problem::QuantGroupSize::kN; + constexpr index_t NPerBlockBQ = NPerBlock / Problem::BQuantGroupSize::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KPerBlockBQ = KPerBlock / Problem::QuantGroupSize::kK; + constexpr index_t KPerBlockBQ = KPerBlock / Problem::BQuantGroupSize::kK; static_assert(std::is_same_v); return GetABQGlobalVectorLoadSize(); @@ -76,7 +76,7 @@ struct GemmMxFp4PipelineAgBgCrPolicy : public UniversalGemmPipelineAgBgCrPolicy constexpr index_t NPerBlock = Problem::BlockGemmShape::kN; constexpr index_t KPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t KScale = KPerBlock / Problem::QuantGroupSize::kK; // k_scale num //2 + constexpr index_t KScale = KPerBlock / Problem::BQuantGroupSize::kK; // k_scale num //2 constexpr index_t VecLoadSize = Problem::FixedVectorSize ? Problem::VectorSizeB : GetVectorSizeB(); constexpr index_t NumWaveGroups = Problem::NumWaveGroups; @@ -109,7 +109,7 @@ struct GemmMxFp4PipelineAgBgCrPolicy : public UniversalGemmPipelineAgBgCrPolicy using BlockWarps = typename Problem::BlockGemmShape::BlockWarps; using WarpTile = typename Problem::BlockGemmShape::WarpTile; - static_assert(Problem::QuantGroupSize::kK % WarpTile::at(I2) == 0, + static_assert(Problem::BQuantGroupSize::kK % WarpTile::at(I2) == 0, "KPerWarpGemm must be a multiple of QuantGroupSize!"); using WarpGemm = WarpGemmDispatcher; using PipelineImplBase = GemmMxFp4PipelineAgBgCrImplBase; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using BDqDataType = remove_cvref_t; - using BQDataType = remove_cvref_t; - using CDataType = remove_cvref_t; - using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using BDqDataType = remove_cvref_t; + using BQDataType = remove_cvref_t; + using CDataType = remove_cvref_t; + using BlockGemmShape = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; - static_assert(QuantGroupSize::kM == 1, "only N/K blocks for BQuant kernel!"); + static_assert(BQuantGroupSize::kM == 1, "only N/K blocks for BQuant kernel!"); using I0 = number<0>; using I1 = number<1>; @@ -58,8 +58,8 @@ struct MxFp4GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3(); } static constexpr index_t GetVectorSizeB() { return Policy::template GetVectorSizeB(); } @@ -93,7 +93,7 @@ struct MxFp4GemmPipelineAgBgCrCompV3 : public BaseGemmPipelineAgBgCrCompV3; using ComputeDataType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; - using QuantGroupSize = remove_cvref_t; + using BQuantGroupSize = remove_cvref_t; using ALayout = remove_cvref_t; using BLayout = remove_cvref_t; @@ -69,14 +69,14 @@ struct WPQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRegV using Base::m_preload; - static constexpr bool PreshuffleQuant = Problem::Traits::PreshuffleQuant; + static constexpr bool BPreshuffleQuant = Problem::Traits::BPreshuffleQuant; static constexpr index_t VectorLoadSize = Problem::VectorLoadSize; static constexpr index_t NPerBlockBQ = - integer_divide_ceil(BlockGemmShape::kN, QuantGroupSize::kN); + integer_divide_ceil(BlockGemmShape::kN, BQuantGroupSize::kN); static constexpr index_t KPerBlockBQ = - integer_divide_ceil(BlockGemmShape::kK, QuantGroupSize::kK); + integer_divide_ceil(BlockGemmShape::kK, BQuantGroupSize::kK); static constexpr index_t QScalesPerBlockRow = - integer_divide_ceil(kKPerBlock, QuantGroupSize::kK); + integer_divide_ceil(kKPerBlock, BQuantGroupSize::kK); static constexpr index_t GetVectorSizeBQ() { @@ -94,7 +94,7 @@ struct WPQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRegV BlockSize, concat('x', WaveNumM, WaveNumN), concat('x', Base::GetVectorSizeA(), Base::GetVectorSizeB(), GetVectorSizeBQ()), - concat('x', kPadM, kPadN, kPadK), QuantGroupSize::GetName()); + concat('x', kPadM, kPadN, kPadK), BQuantGroupSize::GetName()); // clang-format on } @@ -115,7 +115,7 @@ struct WPQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRegV // then by vector width to get an approximate number of vector loads. constexpr index_t BQload_inst = ck_tile::integer_divide_ceil( ck_tile::integer_divide_ceil(kKPerBlock * kNPerBlock * sizeof(BQDataType), - QuantGroupSize::kK * QuantGroupSize::kK), + BQuantGroupSize::kK * BQuantGroupSize::kK), VectorLoadSize); // ToDo: Hardcoded, need to change in future. How many instruction emit per iteration @@ -360,11 +360,11 @@ struct WPQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRegV BQBlockTile bq_block_tile, bq_block_tile_2; bq_block_tile = load_tile(bq_copy_dram_window); // move BQ to tile 1 - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { move_tile_window(bq_copy_dram_window, {((NPerBlockBQ <= BlockGemmShape::BlockWarps::at(number<1>{})) - ? ck_tile::integer_divide_ceil(n, QuantGroupSize::kN) + ? ck_tile::integer_divide_ceil(n, BQuantGroupSize::kN) : ck_tile::integer_least_multiple(n, kNPerBlock) / BlockGemmShape::WarpTile::at(number<1>{})), 0}); @@ -437,11 +437,11 @@ struct WPQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRegV move_tile_window(b_flat_dram_window, {0, BlockGemmShape::flatKPerBlock}); bq_block_tile_2 = load_tile(bq_copy_dram_window); - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { move_tile_window(bq_copy_dram_window, {((NPerBlockBQ <= BlockGemmShape::BlockWarps::at(number<1>{})) - ? ck_tile::integer_divide_ceil(n, QuantGroupSize::kN) + ? ck_tile::integer_divide_ceil(n, BQuantGroupSize::kN) : ck_tile::integer_least_multiple(n, kNPerBlock) / BlockGemmShape::WarpTile::at(number<1>{})), 0}); @@ -474,11 +474,11 @@ struct WPQuantBPipelineAgBgCrV2 : public WeightPreshufflePipelineAGmemBGmemCRegV move_tile_window(b_flat_dram_window, {0, BlockGemmShape::flatKPerBlock}); bq_block_tile = load_tile(bq_copy_dram_window); - if constexpr(PreshuffleQuant) + if constexpr(BPreshuffleQuant) { move_tile_window(bq_copy_dram_window, {((NPerBlockBQ <= BlockGemmShape::BlockWarps::at(number<1>{})) - ? ck_tile::integer_divide_ceil(n, QuantGroupSize::kN) + ? ck_tile::integer_divide_ceil(n, BQuantGroupSize::kN) : ck_tile::integer_least_multiple(n, kNPerBlock) / BlockGemmShape::WarpTile::at(number<1>{})), 0}); diff --git a/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp b/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp index b956caa14f..5db09a0c46 100644 --- a/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp +++ b/include/ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp @@ -33,7 +33,8 @@ inline std::string quant_type_to_string(QuantType quant_type) template +#include + +#include "test_gemm_quant_fixtures.hpp" + +// Type aliases for readability +using RowMajor = ck_tile::tensor_layout::gemm::RowMajor; +using ColumnMajor = ck_tile::tensor_layout::gemm::ColumnMajor; +using FP8 = ck_tile::fp8_t; +using BF8 = ck_tile::bf8_t; +using Half = ck_tile::half_t; +using PkInt4 = ck_tile::pk_int4_t; +using ABQuantGrouped = + std::integral_constant; +using GroupSize = ck_tile::QuantGroupShape>; + +// 2d block sizes for BQuant +using GroupSize2D128N = ck_tile::QuantGroupShape>; + +// Type combinations for ABQuant tests +// Tuple format: +// clang-format off +using ABQuantPreshuffleQuantTypes = ::testing::Types< + std::tuple, + std::tuple +>; +// clang-format on + +// Test suite for ABQuant +TYPED_TEST_SUITE(TestCkTileGemmABQuant, ABQuantPreshuffleQuantTypes); + +// AQuant tests +TYPED_TEST(TestCkTileGemmABQuant, ABQuantGroupedTest) +{ + this->run_test_with_validation(1024, 1024, 1024); +} diff --git a/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp index 8c9955da74..7be4131db4 100644 --- a/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp +++ b/test/ck_tile/gemm_block_scale/test_gemm_quant_base.hpp @@ -75,7 +75,8 @@ class TestCkTileGemmQuantBase : public ::testing::Test static constexpr ck_tile::index_t M_Warp_Tile = GemmConfig::M_Warp_Tile; static constexpr ck_tile::index_t N_Warp_Tile = GemmConfig::N_Warp_Tile; static constexpr ck_tile::index_t K_Warp_Tile = GemmConfig::K_Warp_Tile; - static constexpr bool PreshuffleQuant = GemmConfig::PreshuffleQuant; + static constexpr bool APreshuffleQuant = GemmConfig::APreshuffleQuant; + static constexpr bool BPreshuffleQuant = GemmConfig::BPreshuffleQuant; static constexpr bool PreshuffleB = GemmConfig::PreshuffleB; static constexpr bool TiledMMAPermuteN = GemmConfig::TiledMMAPermuteN; static constexpr bool DoubleSmemBuffer = GemmConfig::DoubleSmemBuffer; @@ -111,7 +112,8 @@ class TestCkTileGemmQuantBase : public ::testing::Test using CodegenGemmTraits = ck_tile::TileGemmQuantTraits @@ -244,7 +245,7 @@ class TestCkTileGemmAQuant : public TestCkTileGemmQuantBase aq_shuffle_host = ck_tile::shuffle_aq(&aq_m_aqk, Base::GemmConfig::K_Tile / QuantGroupSize::kK); @@ -481,7 +482,7 @@ class TestCkTileGemmAQuantMem a_m_k_dev_buf.ToDevice(a_m_k.data()); } // aq_m_aqk_dev_buf.ToDevice(aq_m_aqk.data()); - if constexpr(Base::GemmConfig::PreshuffleQuant) + if constexpr(Base::GemmConfig::APreshuffleQuant) { ck_tile::HostTensor aq_shuffle_host = ck_tile::shuffle_aq(&aq_m_aqk, Base::GemmConfig::K_Tile / QuantGroupSize::kK); @@ -727,7 +728,7 @@ class TestCkTileGemmBQuant : public TestCkTileGemmQuantBase(bq_bqk_bqn, QuantGroupSize::kN); bq_bqk_bqn_dev_buf.ToDevice(bq_shuffle_host.data()); } - else if constexpr(GemmConfig::PreshuffleQuant) + else if constexpr(GemmConfig::BPreshuffleQuant) { ck_tile::HostTensor bq_shuffle_host = ck_tile::shuffle_bq(&bq_bqk_bqn, GemmConfig::K_Tile / QuantGroupSize::kK); @@ -1024,7 +1025,7 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase aq_shuffle_host = ck_tile::shuffle_aq(&aq_m_aqk, Base::GemmConfig::K_Tile / AQuantGroupSize::kK); @@ -1041,7 +1042,7 @@ class TestCkTileGemmABQuant : public TestCkTileGemmQuantBase(bq_bqk_bqn, BQuantGroupSize::kN); bq_bqk_bqn_dev_buf.ToDevice(bq_shuffle_host.data()); } - else if constexpr(GemmConfig::PreshuffleQuant) + else if constexpr(GemmConfig::BPreshuffleQuant) { ck_tile::HostTensor bq_shuffle_host = ck_tile::shuffle_bq(&bq_bqk_bqn, GemmConfig::K_Tile / BQuantGroupSize::kK); diff --git a/test/ck_tile/grouped_gemm_abquant/test_grouped_gemm_abquant_util.hpp b/test/ck_tile/grouped_gemm_abquant/test_grouped_gemm_abquant_util.hpp index c7ed6f5472..b2da7f52a9 100644 --- a/test/ck_tile/grouped_gemm_abquant/test_grouped_gemm_abquant_util.hpp +++ b/test/ck_tile/grouped_gemm_abquant/test_grouped_gemm_abquant_util.hpp @@ -117,6 +117,7 @@ class TestCkTileGroupedGemmABQuant : public ::testing::Test Config::kPadN, Config::kPadK, false, + false, Config::PreshuffleB, ALayout, BLayout, @@ -241,6 +242,7 @@ class TestCkTileGroupedGemmABQuant : public ::testing::Test Config::kPadN, Config::kPadK, false, + false, Config::PreshuffleB, ALayout, BLayout, diff --git a/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_util_quant.hpp b/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_util_quant.hpp index 3d52bca9e0..b472df0cb0 100644 --- a/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_util_quant.hpp +++ b/test/ck_tile/grouped_gemm_quant/test_grouped_gemm_util_quant.hpp @@ -112,6 +112,7 @@ class TestCkTileGroupedGemmQuant : public ::testing::Test GroupedGemKernelParam::kPadN, GroupedGemKernelParam::kPadK, false, + false, PreshuffleB, ALayout, BLayout, @@ -289,6 +290,7 @@ class TestCkTileGroupedGemmQuant : public ::testing::Test GroupedGemKernelParam::kPadN, GroupedGemKernelParam::kPadK, false, + false, PreshuffleB, ALayout, BLayout, From e33f15709f8c1e05f5056edc7295276e121dc253 Mon Sep 17 00:00:00 2001 From: Robin Voetter Date: Thu, 29 Jan 2026 16:15:56 +0100 Subject: [PATCH 09/11] ck-builder: fix test related to changed xdl bwd cshuf v3 interface (#3677) Force merging because I verified this fix manually: git checkout develop git pull ninja smoke-builder (failed to build, as expected) git checkout rvoetter/ckb-fix ninja smoke-builder (passed!) --- experimental/builder/test/conv/ck/test_conv_traits.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/experimental/builder/test/conv/ck/test_conv_traits.cpp b/experimental/builder/test/conv/ck/test_conv_traits.cpp index 3221113565..7de7fae92d 100644 --- a/experimental/builder/test/conv/ck/test_conv_traits.cpp +++ b/experimental/builder/test/conv/ck/test_conv_traits.cpp @@ -812,7 +812,9 @@ TEST_F(ConvTraitsTest, ConvBwdWeightXdlCshuffleV3TraitsExtraction) ck::BlockGemmPipelineScheduler::Intrawave, // BlkGemmPipeSched ck::BlockGemmPipelineVersion::v1, // BlkGemmPipelineVer ck::half_t, // AComputeDataType - ck::half_t>; // BComputeDataType + ck::half_t, // BComputeDataType + false, // DirectLoad + 1>; // NumGroupsToMerge // Use ConvTraitsTmpl to extract compile-time information const auto traits = ck_tile::reflect::conv::instance_to_conv_traits(); From fabac7e2c38f134e70c4caab718579d4d44c2870 Mon Sep 17 00:00:00 2001 From: Johannes Graner Date: Thu, 29 Jan 2026 18:40:28 +0100 Subject: [PATCH 10/11] [Conv] Enable bwd weight splitk autodeduction with cap (#3656) * Enable bwd weight splitk autodeduction with cap * Fix error threshold calculations * Add missing logic to wmma multiple d kernel * Fix threshold calculation * Update test with new applicability --- .../device/device_grouped_conv_bwd_weight.hpp | 2 - ...ice_grouped_conv_bwd_weight_multiple_d.hpp | 2 - ...evice_grouped_conv_bwd_weight_explicit.hpp | 15 ++---- ...bwd_weight_multiple_d_wmma_cshuffle_v3.hpp | 53 +++++++++++++++---- ...onv_bwd_weight_multiple_d_xdl_cshuffle.hpp | 11 ++-- ...ouped_conv_bwd_weight_wmma_cshuffle_v3.hpp | 12 ++--- ...e_grouped_conv_bwd_weight_xdl_cshuffle.hpp | 11 ++-- ...rouped_conv_bwd_weight_xdl_cshuffle_v3.hpp | 12 ++--- .../profile_grouped_conv_bwd_weight_impl.hpp | 47 ++++++++++------ ...rouped_convnd_bwd_weight_interface_xdl.cpp | 2 +- 10 files changed, 91 insertions(+), 76 deletions(-) diff --git a/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp b/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp index 58da96e2f0..eadfa29c9f 100644 --- a/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp +++ b/include/ck/tensor_operation/gpu/device/device_grouped_conv_bwd_weight.hpp @@ -11,8 +11,6 @@ namespace ck { namespace tensor_operation { namespace device { -#define DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS 1 - template ()) diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp index bc072a7019..f662ff834f 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp @@ -22,6 +22,7 @@ #include #include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_utils.hpp" #include "ck/tensor_operation/gpu/device/impl/split_k_arg.hpp" +#include "ck/tensor_operation/gpu/device/impl/split_k_utils.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" #include "ck/host_utility/device_prop.hpp" @@ -524,6 +525,44 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 decltype(GridwiseGemm::MakeDEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock( CGridDesc_M_N{}, 1, 1)); + struct ActiveWorkgroupsPerCU + { + ActiveWorkgroupsPerCU() + { + if(!ck::is_gfx11_supported() && !ck::is_gfx12_supported()) + { + return; + } + constexpr int dynamic_smem_size = 0; + constexpr index_t minimum_occupancy = + BlkGemmPipeSched == BlockGemmPipelineScheduler::Intrawave ? 1 : 2; + int max_occupancy = 0; + + if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v4) + { + // TODO: implement + } + else + { + hip_check_error(hipOccupancyMaxActiveBlocksPerMultiprocessor( + &max_occupancy, + kernel_grouped_conv_bwd_weight_wmma_cshuffle_v3_multiple_d< + GridwiseGemm, + remove_reference_t, + remove_reference_t, + remove_reference_t, + ComputePtrOffsetOfStridedBatch, + true, + InMemoryDataOperationEnum::AtomicAdd, + minimum_occupancy>, + BlockSize, + dynamic_smem_size)); + } + max_occupancy_ = std::max(1, max_occupancy); + } + int max_occupancy_; + }; + struct Argument : public BaseArgument, public ArgumentSplitK { Argument( @@ -574,6 +613,8 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 input_left_pads_{input_left_pads}, input_right_pads_{input_right_pads} { + static ActiveWorkgroupsPerCU active_workgroups_per_cu; + constexpr index_t spatial_offset = 3; std::copy(begin(b_g_n_c_wis_lengths) + spatial_offset, end(b_g_n_c_wis_lengths), @@ -585,7 +626,6 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); -#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN, gemmK; @@ -602,6 +642,9 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 const auto k_batch_max = math::integer_divide_ceil((gemmK - 1), KPerBlock); k_batch_ = std::min(k_batch_, k_batch_max); + // Cap k_batch_ to 128 to avoid accuracy issues + k_batch_ = std::min(k_batch_, 128); + if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING))) { std::cout << "[SPLIT-K AUTODEDUCE] k_batch max value: " << k_batch_max @@ -611,7 +654,6 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 } } else -#endif { k_batch_ = split_k; } @@ -988,13 +1030,6 @@ struct DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3 static bool IsSupportedArgument(const Argument& arg) { -#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS - if(arg.k_batch_ < 0) - { - return false; - } -#endif - const index_t GemmM = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I1); const index_t GemmN = arg.b_grid_desc_kbatch_k0_n_k1_.GetLength(I1); const index_t GemmK = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I0) * diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp index 51dc56e306..1e23fef191 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp @@ -677,7 +677,6 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); -#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN; @@ -688,9 +687,11 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle calculate_mn_grid_size(gemmM, gemmN) * Conv_G_; k_batch_ = get_best_occupancy_k_batch_value(active_workgroups_per_cu.max_occupancy_, grid_size); + + // Cap k_batch_ to 128 to avoid accuracy issues + k_batch_ = std::min(k_batch_, 128); } else -#endif { k_batch_ = split_k; } @@ -947,12 +948,6 @@ struct DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle static bool IsSupportedArgument(const Argument& arg) { -#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS - if(arg.k_batch_ < 0) - { - return false; - } -#endif if(!ck::is_xdl_wmma_supported()) { return false; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp index 3f8093afe1..b2ae092c27 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp @@ -511,7 +511,7 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3 std::copy(begin(a_g_n_k_wos_lengths) + spatial_offset, end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); -#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS + if(split_k < 0) { ck::index_t gemmM, gemmN, gemmK; @@ -528,6 +528,9 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3 const auto k_batch_max = math::integer_divide_ceil((gemmK - 1), KPerBlock); k_batch_ = std::min(k_batch_, k_batch_max); + // Cap k_batch_ to 128 to avoid accuracy issues + k_batch_ = std::min(k_batch_, 128); + if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING))) { std::cout << "[SPLIT-K AUTODEDUCE] k_batch max value: " << k_batch_max @@ -537,7 +540,6 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3 } } else -#endif { k_batch_ = split_k; } @@ -1040,12 +1042,6 @@ struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3 static bool IsSupportedArgument(const Argument& arg) { -#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS - if(arg.k_batch_ < 0) - { - return false; - } -#endif const index_t GemmM = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I1); const index_t GemmN = arg.b_grid_desc_kbatch_k0_n_k1_.GetLength(I1); const index_t GemmK = arg.a_grid_desc_kbatch_k0_m_k1_.GetLength(I0) * diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp index 0ea94806d0..1f6f2fb789 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle.hpp @@ -651,7 +651,6 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle conv_ngchw_to_nhwgc_transformer.TransposeWeiStrides(e_g_k_c_xs_lengths, e_g_k_c_xs_strides); -#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN; @@ -662,9 +661,11 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle calculate_mn_grid_size(gemmM, gemmN) * Conv_G_; k_batch_ = get_best_occupancy_k_batch_value(active_workgroups_per_cu.max_occupancy_, grid_size); + + // Cap k_batch_ to 128 to avoid accuracy issues + k_batch_ = std::min(k_batch_, 128); } else -#endif { k_batch_ = split_k; } @@ -1083,12 +1084,6 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffle static bool IsSupportedArgument(const Argument& arg) { -#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS - if(arg.k_batch_ < 0) - { - return false; - } -#endif if(!ck::is_xdl_wmma_supported()) { return false; diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp index 26cf586017..ac83cee251 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_weight_xdl_cshuffle_v3.hpp @@ -594,7 +594,6 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 end(a_g_n_k_wos_lengths), begin(output_spatial_lengths_)); -#if !DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS if(split_k < 0) { ck::index_t gemmM, gemmN, gemmK; @@ -611,6 +610,9 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 const auto k_batch_max = static_cast((gemmK - 1) / K0PerBlock); k_batch_ = std::max(std::min(k_batch_, k_batch_max), 1); + // Cap k_batch_ to 128 to avoid accuracy issues + k_batch_ = std::min(k_batch_, 128); + if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING))) { std::cout << "[SPLIT-K AUTODEDUCE] k_batch max value: " << k_batch_max @@ -620,7 +622,6 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 } } else -#endif { k_batch_ = split_k; } @@ -1399,13 +1400,6 @@ struct DeviceGroupedConvBwdWeight_Xdl_CShuffleV3 static bool IsSupportedArgument(const Argument& arg) { -#if DISABLE_SPLIT_K_AUTODEDUCE_FOR_ONE_STAGE_KERNELS - if(arg.k_batch_ < 0) - { - return false; - } -#endif - // check device if constexpr(DirectLoad) { diff --git a/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp b/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp index 3a9f14e595..afc88150ed 100644 --- a/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp +++ b/profiler/include/profiler/profile_grouped_conv_bwd_weight_impl.hpp @@ -364,26 +364,39 @@ bool profile_grouped_conv_bwd_weight_impl(int do_verification, using AccDataType = std::conditional_t, int32_t, float>; - // Calculate number of accumulations accounting for split_k - const int num_accums = - static_cast(output.GetElementSize() / conv_param.K_ / split_k_value); - - // Additional tolerance for split_k accumulation if needed - int total_accums = num_accums; - if(split_k_value > 1) - { - total_accums = std::max(num_accums, static_cast(split_k_value)); - } - - // Perform GPU verification (max value computed internally on GPU) + const index_t num_accums = output.GetElementSize() / conv_param.K_; + const index_t num_accums_split_k = split_k_value; + // Get maximum accumulated value from reference const std::size_t tensor_size = weight_device_result.mDesc.GetElementSpaceSize(); + max_accumulated_value = + gpu_reduce_max(gpu_ref_wei_buf.GetDeviceBuffer(), tensor_size); + // Calculate thresholds + auto rtol = + ck::utils::get_relative_threshold( + num_accums / num_accums_split_k); + auto atol = + ck::utils::get_absolute_threshold( + max_accumulated_value / num_accums_split_k, + num_accums / num_accums_split_k); + // Calculate error due to split_k accumulation + auto rtol_split_k = + ck::utils::get_relative_threshold( + num_accums_split_k); + auto atol_split_k = + ck::utils::get_absolute_threshold( + max_accumulated_value, num_accums_split_k); + // Use higher threshold + rtol = std::max(rtol, rtol_split_k); + atol = std::max(atol, atol_split_k); + + // Perform GPU verification auto gpu_result = - ck::profiler::gpu_verify( - wei_device_buf.GetDeviceBuffer(), - gpu_ref_wei_buf.GetDeviceBuffer(), - total_accums, - tensor_size); + ck::profiler::gpu_verify(wei_device_buf.GetDeviceBuffer(), + gpu_ref_wei_buf.GetDeviceBuffer(), + rtol, + atol, + tensor_size); if(!gpu_result) { diff --git a/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp b/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp index bce6da4b68..5aa0b13c07 100644 --- a/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp +++ b/test/grouped_convnd_bwd_weight/test_grouped_convnd_bwd_weight_interface_xdl.cpp @@ -184,5 +184,5 @@ TYPED_TEST(TestGroupedConvndBwdWeightDefault, SingleStageAutoDeduce) this->conv_param = {2, 2, 128, 128, 256, {1, 1}, {3, 3}, {1, 1}, {1, 1}, {0, 0}, {0, 0}}; this->split_k_ = -1; bool is_supported = this->template Run<2>(); - EXPECT_FALSE(is_supported); + EXPECT_TRUE(is_supported); } From f16d9100e42a978261f76319c66a7995e5f6d555 Mon Sep 17 00:00:00 2001 From: Enrico Degregori <73224202+EnricoDeg@users.noreply.github.com> Date: Thu, 29 Jan 2026 19:29:40 +0100 Subject: [PATCH 11/11] Multi AB support for wave transfer (#3578) * Add multi AB support to wave transfer * Improviments to multi ABD examples * Add instances and use intrawave v1 instead of interwave * Apply changes to other transfers * Wave transfer: add support for multiple internal vgpr buffers * Fix compilation error gfx11 --- ...m_multi_ABD_wmma_bias_fastgelu_bf16_i8.cpp | 29 +- .../gemm_multi_ABD_wmma_fastgelu_bf16_i8.cpp | 30 ++- .../gemm_multi_ABD_wmma_fp16.cpp | 27 +- ...BD_wmma_multiply_bias_fastgelu_bf16_i8.cpp | 29 +- ...ead_group_tensor_slice_transfer_global.hpp | 247 +++++++++++------- .../gridwise_ab_transfer_thread_tiles.hpp | 13 + ...se_ab_transfer_thread_tiles_preshuffle.hpp | 13 + .../grid/gridwise_ab_transfer_wave_tiles.hpp | 46 ++-- ...wise_ab_transfer_wave_tiles_interleave.hpp | 43 +-- .../gridwise_gemm_wmma_cshuffle_v3_common.hpp | 27 +- ...multi_abd_bf16_i8_bf16_mk_kn_mn_common.hpp | 19 +- ..._abd_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...multi_abd_bf16_i8_bf16_mk_nk_mn_common.hpp | 15 +- ...bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...gelu_bf16_i8_bf16_mk_nk_mn_v1_instance.cpp | 8 +- ...gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...iply_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- ...gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp | 2 +- 21 files changed, 374 insertions(+), 188 deletions(-) diff --git a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_bias_fastgelu_bf16_i8.cpp b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_bias_fastgelu_bf16_i8.cpp index cf8dd31c3f..78d98e92ce 100644 --- a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_bias_fastgelu_bf16_i8.cpp +++ b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_bias_fastgelu_bf16_i8.cpp @@ -96,11 +96,11 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm 8, 8, 0, - S<8, 32, 1>, + S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, - 1, + 8, 8, 0, 1, @@ -108,7 +108,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm S<1, 32, 1, 8>, S<8, 8, 8>, ck::BlockGemmPipelineScheduler::Intrawave, - ck::BlockGemmPipelineVersion::v3>; + ck::BlockGemmPipelineVersion::v1>; int main(int argc, char* argv[]) { @@ -174,6 +174,29 @@ int main(int argc, char* argv[]) } }; + auto f_get_default_stride = + [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { + if(stride == -1 || stride == 0) + { + // give a chance if stride is -1, return a default packed stride + if constexpr(std::is_same_v) + { + return static_cast(col); + } + else + { + return static_cast(row); + } + } + else + return static_cast(stride); + }; + + StrideA = f_get_default_stride(M, K, StrideA, A0Layout{}); + StrideB = f_get_default_stride(K, N, StrideB, B0Layout{}); + StrideD = f_get_default_stride(M, N, StrideD, D0Layout{}); + StrideE = f_get_default_stride(M, N, StrideE, ELayout{}); + Tensor a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{})); Tensor b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{})); Tensor b1_k_n(f_host_tensor_descriptor(K, N, StrideB, B1Layout{})); diff --git a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fastgelu_bf16_i8.cpp b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fastgelu_bf16_i8.cpp index e4033e5bac..089404757a 100644 --- a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fastgelu_bf16_i8.cpp +++ b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fastgelu_bf16_i8.cpp @@ -94,11 +94,11 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm 8, 8, 0, - S<8, 32, 1>, + S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, - 1, + 8, 8, 0, 1, @@ -106,7 +106,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm S<1, 32, 1, 8>, S<8, 8, 8>, ck::BlockGemmPipelineScheduler::Intrawave, - ck::BlockGemmPipelineVersion::v3>; + ck::BlockGemmPipelineVersion::v1>; int main(int argc, char* argv[]) { @@ -133,7 +133,7 @@ int main(int argc, char* argv[]) init_method = std::stoi(argv[2]); time_kernel = std::stoi(argv[3]); } - else if(argc == 11) + else if(argc == 10) { do_verification = std::stoi(argv[1]); init_method = std::stoi(argv[2]); @@ -170,6 +170,28 @@ int main(int argc, char* argv[]) } }; + auto f_get_default_stride = + [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { + if(stride == -1 || stride == 0) + { + // give a chance if stride is -1, return a default packed stride + if constexpr(std::is_same_v) + { + return static_cast(col); + } + else + { + return static_cast(row); + } + } + else + return static_cast(stride); + }; + + StrideA = f_get_default_stride(M, K, StrideA, A0Layout{}); + StrideB = f_get_default_stride(K, N, StrideB, B0Layout{}); + StrideE = f_get_default_stride(M, N, StrideE, ELayout{}); + Tensor a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{})); Tensor b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{})); Tensor b1_k_n(f_host_tensor_descriptor(K, N, StrideB, B1Layout{})); diff --git a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fp16.cpp b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fp16.cpp index 5817269fdf..d5ccf7eb59 100644 --- a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fp16.cpp +++ b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_fp16.cpp @@ -141,11 +141,11 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm 8, 8, 0, - S<4, 64, 1>, + S<4, 16, 1>, S<1, 0, 2>, S<1, 0, 2>, 1, - 1, + 8, 8, 0, 1, @@ -233,6 +233,29 @@ int main(int argc, char* argv[]) } }; + auto f_get_default_stride = + [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { + if(stride == -1 || stride == 0) + { + // give a chance if stride is -1, return a default packed stride + if constexpr(std::is_same_v) + { + return static_cast(col); + } + else + { + return static_cast(row); + } + } + else + return static_cast(stride); + }; + + StrideA = f_get_default_stride(M, K, StrideA, ALayout{}); + StrideB = f_get_default_stride(K, N, StrideB, BLayout{}); + StrideD = f_get_default_stride(M, N, StrideD, DLayout{}); + StrideE = f_get_default_stride(M, N, StrideE, ELayout{}); + Tensor a0_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); Tensor a1_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{})); Tensor b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{})); diff --git a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_multiply_bias_fastgelu_bf16_i8.cpp b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_multiply_bias_fastgelu_bf16_i8.cpp index 4fb1a5ab4e..2d07bc480d 100644 --- a/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_multiply_bias_fastgelu_bf16_i8.cpp +++ b/example/60_gemm_multi_ABD/gemm_multi_ABD_wmma_multiply_bias_fastgelu_bf16_i8.cpp @@ -95,11 +95,11 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm 8, 8, 0, - S<8, 32, 1>, + S<8, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, - 1, + 8, 8, 0, 1, @@ -107,7 +107,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultipleABD_Wmm S<1, 32, 1, 8>, S<8, 8, 8>, ck::BlockGemmPipelineScheduler::Intrawave, - ck::BlockGemmPipelineVersion::v3>; + ck::BlockGemmPipelineVersion::v1>; int main(int argc, char* argv[]) { @@ -173,6 +173,29 @@ int main(int argc, char* argv[]) } }; + auto f_get_default_stride = + [](std::size_t row, std::size_t col, ck::index_t stride, auto layout) { + if(stride == -1 || stride == 0) + { + // give a chance if stride is -1, return a default packed stride + if constexpr(std::is_same_v) + { + return static_cast(col); + } + else + { + return static_cast(row); + } + } + else + return static_cast(stride); + }; + + StrideA = f_get_default_stride(M, K, StrideA, A0Layout{}); + StrideB = f_get_default_stride(K, N, StrideB, B0Layout{}); + StrideD = f_get_default_stride(M, N, StrideD, D0Layout{}); + StrideE = f_get_default_stride(M, N, StrideE, ELayout{}); + Tensor a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{})); Tensor b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{})); Tensor d0_m_n(f_host_tensor_descriptor(M, N, StrideD, D0Layout{})); diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_global.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_global.hpp index 1c322fe4a7..d1c6f30a14 100644 --- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_global.hpp +++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_global.hpp @@ -12,16 +12,17 @@ namespace ck { -template + bool DoTranspose, + index_t NumThreadScratch = 1> struct ThreadGroupTransferGlobal { static constexpr auto I0 = Number<0>{}; @@ -32,24 +33,57 @@ struct ThreadGroupTransferGlobal static constexpr auto I5 = Number<5>{}; static constexpr auto I6 = Number<6>{}; - static constexpr index_t nDim = remove_reference_t::GetNumOfDimension(); - using Index = MultiIndex; - using SrcCoord = decltype(make_tensor_coordinate(SrcDesc{}, Index{})); - using DstCoord = decltype(make_tensor_coordinate(DstDesc{}, Index{})); + // return a tuple of coordiantes for a tuple of tensor + template = false> + static constexpr auto MakeCoordinates(const Descs& descs, const Indices& indices) + { + return generate_tuple([&](auto i) { return make_tensor_coordinate(descs[i], indices[i]); }, + Number{}); + } - __device__ ThreadGroupTransferGlobal(const SrcDesc& src_desc, - const DstDesc& dst_desc, - const Index& src_block_slice_origin, - const Index& dst_block_slice_origin, - const ElementwiseOperation& element_op) - : src_coord_(make_tensor_coordinate(src_desc, src_block_slice_origin)), + static constexpr index_t nDim = + remove_cvref_t>::GetNumOfDimension(); + static constexpr index_t nSrc = SrcDescs::Size(); + using Index = MultiIndex; + using SrcCoords = decltype(MakeCoordinates(SrcDescs{}, StaticallyIndexedArray{})); + using DstCoord = decltype(make_tensor_coordinate(DstDesc{}, Index{})); + + __device__ + ThreadGroupTransferGlobal(const SrcDescs& src_descs, + const DstDesc& dst_desc, + const StaticallyIndexedArray& src_block_slice_origins, + const Index& dst_block_slice_origin, + const ElementwiseOperation& element_op) + : src_coords_(MakeCoordinates(src_descs, src_block_slice_origins)), dst_coord_(make_tensor_coordinate(dst_desc, dst_block_slice_origin)), element_op_(element_op) { } - template - __device__ void RunRead(const SrcDesc& src_desc, const GridBufferType& grid_buf) + template + __device__ static auto generate_vectors() + { + auto data_types = DataTypes_{}; + + constexpr index_t num = data_types.Size(); + + return generate_tuple( + [&](auto i) { + using DataType = remove_cvref_t; + + return vector_type_maker_t{}; + }, + Number{}); + } + + template = false> + __device__ void RunRead(SrcDescs& src_descs, + const GridBufferTypes& grid_bufs, + Number thread_scratch_id = Number{}) { constexpr auto src_access_lengths = NumberOfIterations{}; constexpr auto src_dim_access_order = IterationOrder{}; @@ -57,36 +91,6 @@ struct ThreadGroupTransferGlobal container_reorder_given_new2old(src_access_lengths, src_dim_access_order); constexpr auto ordered_fwd_step = StepsPerIteration{}; - // make forward steps - // forward step for each iteration just add 1 - const auto src_forward_steps = generate_tuple( - [&](auto i) { - Index forward_step_idx; - - static_for<0, nDim, 1>{}([&](auto j) { - forward_step_idx(j) = (i.value == j.value) ? ordered_fwd_step[i] : 0; - }); - - return make_tensor_coordinate_step(src_desc, forward_step_idx); - }, - Number{}); - - // make backward steps - // backward step at the end of the dimension iteration subtract IterationLength - 1 - const auto src_backward_steps = generate_tuple( - [&](auto i) { - Index backward_step_idx; - - static_for<0, nDim, 1>{}([&](auto j) { - backward_step_idx(j) = (i.value == j.value) - ? (-src_access_lengths[i] + 1) * ordered_fwd_step[i] - : 0; - }); - - return make_tensor_coordinate_step(src_desc, backward_step_idx); - }, - Number{}); - static_ford{}([&](auto ordered_src_access_idx) { // judge move forward or move backward constexpr auto forward_sweep = [&]() { @@ -157,10 +161,26 @@ struct ThreadGroupTransferGlobal }, Number{}); - // check if src element is valid - const bool is_src_valid = - coordinate_has_valid_offset_assuming_visible_index_is_valid(src_desc, src_coord_); - oob_thread_scratch_.template SetAsType(vgpr_data_idx_seq, is_src_valid); + auto src_vectors = generate_vectors(); + bool oob_val = true; + + static_for<0, nSrc, 1>{}([&](auto i) { + using src_vector_t = typename remove_cvref_t::type; + // check if src element is valid + const bool is_src_valid = + coordinate_has_valid_offset_assuming_visible_index_is_valid(src_descs[i], + src_coords_[i]); + + oob_val = oob_val & is_src_valid; + + // Load data from memory in src_vector first + auto index = is_src_valid || !DoTranspose ? src_coords_[i].GetOffset() : 0; + src_vectors(i).template AsType()(I0) = + grid_bufs[i].template Get(index, true); + }); + + oob_thread_scratch_(thread_scratch_id) + .template SetAsType(vgpr_data_idx_seq, oob_val); // Vector length of elementwise operation constexpr auto get_elem_op_vec_len = []() { @@ -185,57 +205,105 @@ struct ThreadGroupTransferGlobal } }; - // This is 1 for pass through because internally it's doing type conversion constexpr index_t elem_op_vec_len = get_elem_op_vec_len(); - using src_vector_container = vector_type_maker_t; - using src_vector_container_t = typename src_vector_container::type; - - using elem_op_vec_t = typename vector_type::type; - using dst_vector_type = vector_type_maker_t; using dst_vector_t = typename dst_vector_type::type; - dst_vector_type op_r_v; - // Load data from memory in src_vector first - auto index = is_src_valid || !DoTranspose ? src_coord_.GetOffset() : 0; - src_vector_container src_vector = src_vector_container{ - grid_buf.template Get(index, true)}; - // apply the src elementwise op and convert to DstData under the hood if needed static_for<0, VectorSize / elem_op_vec_len, 1>{}([&](auto idx) { - element_op_(op_r_v.template AsType()(idx), - src_vector.template AsType()[idx]); + // get reference to src data + const auto src_data_refs = generate_tie( + // return type should be lvalue + [&](auto iSrc) -> const auto& { + using SrcData = remove_cvref_t>; + + using elem_op_vec_t = typename vector_type::type; + + return src_vectors[iSrc].template AsType()[idx]; + }, + Number{}); + + // get reference to dst data + auto dst_data_refs = generate_tie( + // return type should be lvalue + [&](auto) -> auto& { + using elem_op_vec_t = typename vector_type::type; + + return op_r_v.template AsType()(idx); + }, + Number<1>{}); + + // apply pointwise function + unpack2(element_op_, dst_data_refs, src_data_refs); }); // store result in dvgpr_ (static array holding loaded data). // At this point data is already converted to DstData type and // the elementwise operation has been applied - src_dvgpr_.template SetAsType(vgpr_data_idx_seq, - op_r_v.template AsType()[I0]); + src_dvgpr_(thread_scratch_id) + .template SetAsType(vgpr_data_idx_seq, + op_r_v.template AsType()[I0]); - // For each dimension move fwd, bwd or don't move - static_for<0, nDim, 1>{}([&](auto i) { - if constexpr(move_on_dim[i]) - { - if constexpr(forward_sweep[i]) + // Move each src coordinate + static_for<0, nSrc, 1>{}([&](auto iSrc) { + // make forward steps + // forward step for each iteration just add 1 + const auto src_forward_steps = generate_tuple( + [&](auto iDim) { + Index forward_step_idx; + + static_for<0, nDim, 1>{}([&](auto j) { + forward_step_idx(j) = + (iDim.value == j.value) ? ordered_fwd_step[iDim] : 0; + }); + return make_tensor_coordinate_step(src_descs[iSrc], forward_step_idx); + }, + Number{}); + + // make backward steps + // backward step at the end of the dimension iteration subtract IterationLength - 1 + const auto src_backward_steps = generate_tuple( + [&](auto iDim) { + Index backward_step_idx; + + static_for<0, nDim, 1>{}([&](auto j) { + backward_step_idx(j) = + (iDim.value == j.value) + ? (-src_access_lengths[iDim] + 1) * ordered_fwd_step[iDim] + : 0; + }); + return make_tensor_coordinate_step(src_descs[iSrc], backward_step_idx); + }, + Number{}); + + // For each dimension move fwd, bwd or don't move + static_for<0, nDim, 1>{}([&](auto i) { + if constexpr(move_on_dim[i]) { - move_tensor_coordinate( - src_desc, src_coord_, src_forward_steps[src_dim_access_order[i]]); + if constexpr(forward_sweep[i]) + { + move_tensor_coordinate(src_descs[iSrc], + src_coords_(iSrc), + src_forward_steps[src_dim_access_order[i]]); + } + else + { + move_tensor_coordinate(src_descs[iSrc], + src_coords_(iSrc), + src_backward_steps[src_dim_access_order[i]]); + } } - else - { - move_tensor_coordinate( - src_desc, src_coord_, src_backward_steps[src_dim_access_order[i]]); - } - } + }); }); }); } - template - __device__ void RunWrite(const DstDesc& dst_desc, BlockBufferType& dst_buf) + template + __device__ void RunWrite(const DstDesc& dst_desc, + BlockBufferType& dst_buf, + Number thread_scratch_id = Number{}) { using dst_vector_type = vector_type_maker_t; using dst_vector_t = typename dst_vector_type::type; @@ -272,9 +340,10 @@ struct ThreadGroupTransferGlobal }, Number{}); - auto op_r = src_dvgpr_.template GetAsType(vgpr_data_idx_seq); + auto op_r = + src_dvgpr_(thread_scratch_id).template GetAsType(vgpr_data_idx_seq); const bool is_src_valid = - oob_thread_scratch_.template GetAsType(vgpr_data_idx_seq); + oob_thread_scratch_(thread_scratch_id).template GetAsType(vgpr_data_idx_seq); auto op_r_v = is_src_valid ? op_r : dst_vector_t(0); dst_dvgpr_.template SetAsType(vgpr_data_idx_seq, op_r_v); }); @@ -404,10 +473,12 @@ struct ThreadGroupTransferGlobal }); } - __device__ void MoveSrcSliceWindow(const SrcDesc& src_desc, const Index& step) + __device__ void MoveSrcSliceWindow(const SrcDescs& src_descs, const Index& step) { - const auto adjusted_step = make_tensor_coordinate_step(src_desc, step); - move_tensor_coordinate(src_desc, src_coord_, adjusted_step); + static_for<0, nSrc, 1>{}([&](auto iSrc) { + const auto adjusted_step = make_tensor_coordinate_step(src_descs[iSrc], step); + move_tensor_coordinate(src_descs[iSrc], src_coords_(iSrc), adjusted_step); + }); } private: @@ -443,10 +514,10 @@ struct ThreadGroupTransferGlobal decltype(src_oob_thread_scratch_desc_), true>; - ThreadScratchData src_dvgpr_; + StaticallyIndexedArray src_dvgpr_; ThreadScratchData dst_dvgpr_; - OOBThreadScratch oob_thread_scratch_; - SrcCoord src_coord_; + StaticallyIndexedArray oob_thread_scratch_; + SrcCoords src_coords_; DstCoord dst_coord_; const ElementwiseOperation element_op_; }; diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp index 96387c6f64..4d5c052e02 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp @@ -488,6 +488,19 @@ struct ABTransferThreadTiles { return make_dynamic_buffer(p_shared_AB, size); } + + template + __device__ __forceinline__ static auto get_first_element_workaround(Type& array) + { + if constexpr(numElements > 1) + { + return array; + } + else + { + return array[I0]; + } + } }; } // namespace ck diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles_preshuffle.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles_preshuffle.hpp index ad9af92ae5..fb6d1451d3 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles_preshuffle.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles_preshuffle.hpp @@ -133,6 +133,19 @@ struct ABTransferThreadTilesPreShuffle { return make_static_buffer(size); } + + template + __device__ __forceinline__ static auto get_first_element_workaround(Type& array) + { + if constexpr(numElements > 1) + { + return array; + } + else + { + return array[I0]; + } + } }; } // namespace ck diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles.hpp index caf468d6cb..63c0299750 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles.hpp @@ -318,43 +318,43 @@ struct ABTransferWaveTiles const index_t block_mn_id, const index_t) { - // Note: GlobalBufferNum is currently not used but it will be needed - // once we add other pipelines. It is currently needed only for - // consistency with the thread tiles approach - static_assert(GlobalBufferNum == 1, "single global buffer is only supported"); constexpr index_t NumABTensor = ABsDataType::Size(); - static_assert(NumABTensor == 1, "multiAB currently not supported"); - - using ABDataType = remove_cvref_t>; const auto wave_idx = GetWaveIdx(); index_t wave_idK = wave_idx[I1]; index_t wave_idMN = wave_idx[I0]; - const auto grid_lane_id = GetGridLaneIdx(); - index_t lane_group_grid = grid_lane_id[I0]; - index_t lane_local_id_grid = grid_lane_id[I1]; - const auto block_lane_id = GetBlockLaneIdx(); index_t lane_group_block = block_lane_id[I0]; index_t lane_local_id_block = block_lane_id[I1]; - return ThreadGroupTransferGlobal>; + const auto grid_lane_id = GetGridLaneIdx(); + index_t lane_group_grid = grid_lane_id[I0]; + index_t lane_local_id_grid = grid_lane_id[I1]; + return make_multi_index(block_mn_id * (MNRepeat_ * MNWaves_) + wave_idMN, + wave_idK, + lane_group_grid, + lane_local_id_grid); + }, + Number{}); + + return ThreadGroupTransferGlobal, Sequence, Sequence, ABK1Value, - ABDoTranspose>( - grid_descriptor[I0], + ABDoTranspose, + GlobalBufferNum>( + grid_descriptor, block_descriptor, - make_multi_index(block_mn_id * (MNRepeat_ * MNWaves_) + wave_idMN, - wave_idK, - lane_group_grid, - lane_local_id_grid), + idx_as_block_begin, make_multi_index(wave_idMN, wave_idK, lane_group_block, lane_local_id_block), ab_element_op); } @@ -398,6 +398,12 @@ struct ABTransferWaveTiles { return make_dynamic_buffer(p_shared_AB, size); } + + template + __device__ __forceinline__ static auto get_first_element_workaround(Type& array) + { + return array; + } }; } // namespace ck diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles_interleave.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles_interleave.hpp index bfe5b7bd08..e1ee47770b 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles_interleave.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_wave_tiles_interleave.hpp @@ -218,45 +218,46 @@ struct ABTransferWaveTilesInterleave : ABTransferWaveTiles>; const auto wave_idx = GetWaveIdx(); index_t wave_idK = wave_idx[I1]; index_t wave_idMN = wave_idx[I0]; - const auto grid_lane_id = Base::template GetGridLaneIdx(); - index_t lane_group_grid = grid_lane_id[I0]; - index_t lane_local_id_grid = grid_lane_id[I1]; - const auto block_lane_id = GetBlockLaneIdx(); index_t lane_group_block = block_lane_id[I0]; index_t lane_local_id_block = block_lane_id[I1]; constexpr index_t MNRepeatRatio = MNRepeat_Grid / MNRepeat_; - return ThreadGroupTransferGlobal>; + const auto grid_lane_id = Base::template GetGridLaneIdx(); + index_t lane_group_grid = grid_lane_id[I0]; + index_t lane_local_id_grid = grid_lane_id[I1]; + return make_multi_index(block_mn_id * MNWaves_Grid + wave_idMN / MNRepeatRatio, + wave_idK * KRepeat_Grid, + (wave_idMN % MNRepeatRatio) * MNRepeat_, + lane_group_grid, + lane_local_id_grid); + }, + Number{}); + + return ThreadGroupTransferGlobal, Sequence, Sequence, ABK1Value, - ABDoTranspose>( - grid_descriptor[I0], + ABDoTranspose, + GlobalBufferNum>( + grid_descriptor, block_descriptor, - make_multi_index(block_mn_id * MNWaves_Grid + wave_idMN / MNRepeatRatio, - wave_idK * KRepeat_Grid, - (wave_idMN % MNRepeatRatio) * MNRepeat_, - lane_group_grid, - lane_local_id_grid), + idx_as_block_begin, make_multi_index(wave_idMN / MNRepeatRatio, wave_idK * KRepeat_, (wave_idMN % MNRepeatRatio) * MNRepeat_, diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp index bcf131003c..03735bbc6a 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp @@ -364,7 +364,7 @@ struct GridwiseGemm_wmma_cshuffle_v3_base __host__ __device__ static constexpr bool AWaveTransferApplicable() { - return !ForceThreadTileTransfer && NumATensor == 1 && APackedSize == 1 && + return !ForceThreadTileTransfer && APackedSize == 1 && ABlockTransferSrcScalarPerVector == 8 && ABlockTransferDstScalarPerVector_AK1 == 8 && BlkGemmPipelineVer == BlockGemmPipelineVersion::v1 && AK1Value == 8 && !IsBPreShuffled; @@ -372,13 +372,11 @@ struct GridwiseGemm_wmma_cshuffle_v3_base __host__ __device__ static constexpr bool BWaveTransferApplicable() { - return !ForceThreadTileTransfer && NumBTensor == 1 && BPackedSize == 1 && + return !ForceThreadTileTransfer && BPackedSize == 1 && BBlockTransferSrcScalarPerVector == 8 && BBlockTransferDstScalarPerVector_BK1 == 8 && BlkGemmPipelineVer == BlockGemmPipelineVersion::v1 && BK1Value == 8; } - // Limitations of the current implementation: - // - no multiAB #ifdef __gfx12__ static constexpr bool IsAWaveTransferApplicable = AWaveTransferApplicable(); @@ -1319,19 +1317,6 @@ struct GridwiseGemm_wmma_cshuffle_v3_base } } - template - __device__ __forceinline__ static auto get_first_element_workaround(Type& array) - { - if constexpr(numElements > 1) - { - return array; - } - else - { - return array[I0]; - } - } - // Note: arguments k_batch and k_id should be set if splitk is used // with implicit gemm (no pointer shift but shift using tensor descriptors) template ( - get_first_element_workaround(as_grid_desc_ak0_m_ak1), + ATransfer::template get_first_element_workaround(as_grid_desc_ak0_m_ak1), a_block_desc_ak0_m_ak1, a_blockwise_copy, - get_first_element_workaround(as_grid_buf), + ATransfer::template get_first_element_workaround(as_grid_buf), a_block_buf, a_block_slice_copy_step, - get_first_element_workaround(bs_grid_desc_bk0_n_bk1), + BTransfer::template get_first_element_workaround(bs_grid_desc_bk0_n_bk1), b_block_desc_bk0_n_bk1, b_blockwise_copy, - get_first_element_workaround(bs_grid_buf), + BTransfer::template get_first_element_workaround(bs_grid_buf), b_block_buf, b_block_slice_copy_step, c_thread_buf, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_common.hpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_common.hpp index 4cd4403436..0dd666b3d9 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_common.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_common.hpp @@ -73,14 +73,17 @@ template using device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_comp_instances = std::tuple< // clang-format off - //###################################| ALayout| BLayout| DsLayout| ELayout| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| BlkGemmPipeSched| BlkGemmPipelineVer| - //###################################| | | | | Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Size| Block| Block| Block| | | Wmma| Wmma| | | ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MRepeat| NRepeat| _MBlock_MPerBlock_NBlock_NPerBlock| ScalarPerVector| | | - //###################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | | | Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| | | | | - //###################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 256, 256, 32, 8, 8, 16, 16, 8, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, - DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, - DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 32, 8, 8, 16, 16, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, - DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1> + //###################################| ALayout| BLayout| DsLayout| ELayout| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| BlkGemmPipeSched| BlkGemmPipelineVer| + //###################################| | | | | Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Size| Block| Block| Block| | | Wmma| Wmma| | | ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MRepeat| NRepeat| _MBlock_MPerBlock_NBlock_NPerBlock| ScalarPerVector| | | + //###################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | | | Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| | | | | + //###################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 256, 256, 32, 8, 8, 16, 16, 8, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 32, 8, 8, 16, 16, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 64, 8, 8, 16, 16, 8, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 32, 8, 8, 16, 16, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 64, 8, 8, 16, 16, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1> // clang-format on >; diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index 76a92a1971..3587c6700c 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_v1_instances( Multiply, PassThrough, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances(instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances< ck::Tuple, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_common.hpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_common.hpp index 1607b240f6..7cb50cd954 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_common.hpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_common.hpp @@ -71,12 +71,15 @@ template using device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_comp_instances = std::tuple< // clang-format off - //###################################| ALayout| BLayout| DsLayout| ELayout| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| BlkGemmPipeSched| BlkGemmPipelineVer| - //###################################| | | | | Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Size| Block| Block| Block| | | Wmma| Wmma| | | ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MRepeat| NRepeat| _MBlock_MPerBlock_NBlock_NPerBlock| ScalarPerVector| | | - //###################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | | | Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| | | | | - //###################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | - DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, - DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1> + //###################################| ALayout| BLayout| DsLayout| ELayout| AData| BData| AccData| CShuffle| DsData| EData| A| B| CDE| GEMM| Block| MPer| NPer| KPer| AK1| BK1| MPer| NPer| MRepeat| NRepeat| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| BBlockTransfer| BBlockTransfer| BBlockTransfer| BlockTransfer| BBlockTransfer| BBlockTransfer| BBlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| BlkGemmPipeSched| BlkGemmPipelineVer| + //###################################| | | | | Type| Type| Type| DataType| Type| Type| Elementwise| Elementwise| Elementwise| Spacialization| Size| Block| Block| Block| | | Wmma| Wmma| | | ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MRepeat| NRepeat| _MBlock_MPerBlock_NBlock_NPerBlock| ScalarPerVector| | | + //###################################| | | | | | | | | | | Operation| Operation| Operation| | | | | | | | | | | | Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| | | | | + //###################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlockGemmPipelineScheduler::Intrawave, BlockGemmPipelineVersion::v3>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 32, 8, 8, 16, 16, 2, 8, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 64, 1, 4>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 128, 64, 8, 8, 16, 16, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1>, + DeviceGemmMultipleABD_Wmma_CShuffleV3< AsLayout, BsLayout, DsLayout, ELayout, AsDataType, BsDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, 128, 256, 64, 8, 8, 16, 16, 8, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 0, 1, 1, S<1, 16, 1, 16>, S<8, 8, 8>, BlkGemmPipeSched, BlockGemmPipelineVersion::v1> // clang-format on >; } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index 2a4aae98a5..731518257b 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_bias_v1_instances( Multiply, Add, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances(instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances< ck::Tuple, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index 477d6811d2..0a67f2357e 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_bias_gelu_v1_instances Multiply, AddFastGelu, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances(instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances< ck::Tuple, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_nk_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_nk_mn_v1_instance.cpp index 71c04b3485..c0b4cf7b9a 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_nk_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_bias_gelu_bf16_i8_bf16_mk_nk_mn_v1_instance.cpp @@ -36,7 +36,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_bias_gelu_v1_instances ck::Tuple, AddFastGelu, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); } void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_bias_v1_instances( @@ -58,7 +58,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_bias_v1_instances( ck::Tuple, Add, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); } void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_v1_instances( @@ -80,7 +80,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_v1_instances( ck::Tuple<>, PassThrough, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); } void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_gelu_v1_instances( @@ -102,7 +102,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_nk_mn_gelu_v1_instances( ck::Tuple<>, FastGelu, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); } } // namespace instance diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index 33422fc6db..9176910cea 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_gelu_v1_instances( Multiply, FastGelu, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances(instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances< diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index 639bda6017..669eb4144a 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_v1_instances( PassThrough, Multiply, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances( instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index 7f8fea44c5..c6a812645b 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_bias_v1_insta PassThrough, MultiplyAdd, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances(instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances< ck::Tuple, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index b2bf995507..2d7ffd120d 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_bias_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_bias_gelu_v1_ PassThrough, MultiplyAddFastGelu, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances(instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances< ck::Tuple, diff --git a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp index d2adc36dc3..ab49d2f1c9 100644 --- a/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp +++ b/library/src/tensor_operation_instance/gpu/gemm_multi_abd/device_gemm_wmma_multi_abd_multiply_gelu_bf16_i8_bf16_mk_kn_mn_v1_instance.cpp @@ -39,7 +39,7 @@ void add_device_gemm_wmma_multi_abd_multiply_bf16_i8_bf16_mk_kn_mn_gelu_v1_insta PassThrough, MultiplyFastGelu, GemmMNKPadding, - Interwave>{}); + Intrawave>{}); add_device_operation_instances( instances, device_gemm_wmma_multi_abd_bf16_i8_bf16_mk_kn_mn_mem_instances,