From fc82ebc174ec719f445889699a89db428a973b6a Mon Sep 17 00:00:00 2001 From: Sami Aario Date: Thu, 9 Oct 2025 08:07:04 +0000 Subject: [PATCH] Introduce DetermineWarpPrecType for determining warp GEMM precision types --- include/ck_tile/ops/common.hpp | 1 + .../ops/common/determine_warp_prec_type.hpp | 41 +++++++++++++++++++ .../ops/epilogue/cshuffle_epilogue.hpp | 14 ++----- .../block/block_universal_gemm_as_bs_cr.hpp | 8 +--- ...emm_universal_pipeline_ag_bg_cr_policy.hpp | 12 ++---- 5 files changed, 51 insertions(+), 25 deletions(-) create mode 100644 include/ck_tile/ops/common/determine_warp_prec_type.hpp diff --git a/include/ck_tile/ops/common.hpp b/include/ck_tile/ops/common.hpp index ad7da5c183..0113d8c9a2 100644 --- a/include/ck_tile/ops/common.hpp +++ b/include/ck_tile/ops/common.hpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: MIT #pragma once +#include "ck_tile/ops/common/determine_warp_prec_type.hpp" #include "ck_tile/ops/common/generic_2d_block_shape.hpp" #include "ck_tile/ops/common/load_and_convert_tile.hpp" #include "ck_tile/ops/common/streamk_common.hpp" diff --git a/include/ck_tile/ops/common/determine_warp_prec_type.hpp b/include/ck_tile/ops/common/determine_warp_prec_type.hpp new file mode 100644 index 0000000000..05ad47b526 --- /dev/null +++ b/include/ck_tile/ops/common/determine_warp_prec_type.hpp @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#pragma once + +#include "ck_tile/core.hpp" + +// DetermineWarpPrecType is a set of rules to determine the right precision type to use +// for the warp GEMM, given the other precision type. This gives rise to a type conversion: +// type conversions are sometimes needed to obtain a pair of types that are compatible with +// the hardware matrix operations available. A typical use case is mixed precision GEMMs. + +namespace ck_tile { +// For the most general case, we default to no conversion. +template +struct DetermineWarpPrecType +{ + using prec_type = PrecType; +}; + +// For pk_int4_t, we convert to the other precision type. +template +struct DetermineWarpPrecType +{ + using prec_type = OtherPrecType; +}; + +// For pk_fp4_t, we convert to the other precision type. +template +struct DetermineWarpPrecType +{ + using prec_type = OtherPrecType; +}; + +// For pk_fp4_raw_t, we convert to the other precision type. +template +struct DetermineWarpPrecType +{ + using prec_type = OtherPrecType; +}; +}; // namespace ck_tile diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp index 97f936fde9..d8844cb27e 100644 --- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp +++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp @@ -7,7 +7,7 @@ #include "ck_tile/core.hpp" #include "ck_tile/ops/common/utils.hpp" #include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp" -#include "ck_tile/ops/common/tensor_layout.hpp" +#include "ck_tile/ops/common.hpp" #include "ck_tile/ops/elementwise/unary_element_wise_operation.hpp" #include @@ -92,16 +92,8 @@ struct CShuffleEpilogue using ADataType = remove_cvref_t{}, AsDataTypeTuple>>; using BDataType = remove_cvref_t{}, BsDataTypeTuple>>; - using ATypeToUse = std::conditional_t || - std::is_same_v, - BDataType, - ADataType>; - // Used for weight-only quantization kernel, B would be dequantized to the same data type as A - using BTypeToUse = std::conditional_t || - std::is_same_v || - std::is_same_v, - ADataType, - BDataType>; + using ATypeToUse = typename DetermineWarpPrecType::prec_type; + using BTypeToUse = typename DetermineWarpPrecType::prec_type; using ELayout = remove_cvref_t; using CDElementwise = remove_cvref_t; diff --git a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp index 381a5513ea..fffcf8ba4c 100644 --- a/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp +++ b/include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp @@ -94,12 +94,8 @@ struct BlockUniversalGemmAsBsCr using ComputeDataType = remove_cvref_t; using CDataType = remove_cvref_t; - using ATypeToUse = - std::conditional_t, BDataType, ADataType>; - using BTypeToUse = std::conditional_t || - std::is_same_v, - ADataType, - BDataType>; + using ATypeToUse = typename DetermineWarpPrecType::prec_type; + using BTypeToUse = typename DetermineWarpPrecType::prec_type; using WarpGemm = remove_cvref_t; 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 4bb7829fb6..ae76b22c75 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 @@ -895,14 +895,10 @@ struct UniversalGemmPipelineAgBgCrPolicy : vector_size * 4 == thread_elements ? WGAttrNumAccessEnum::Quad : WGAttrNumAccessEnum::Invalid; - using ADataType = remove_cvref_t; - using BDataType = remove_cvref_t; - using ATypeToUse = - std::conditional_t, BDataType, ADataType>; - using BTypeToUse = std::conditional_t || - std::is_same_v, - ADataType, - BDataType>; + using ADataType = remove_cvref_t; + using BDataType = remove_cvref_t; + using ATypeToUse = typename DetermineWarpPrecType::prec_type; + using BTypeToUse = typename DetermineWarpPrecType::prec_type; using WarpGemm = WarpGemmDispatcher