ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 18:17:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Introduce DetermineWarpPrecType for determining warp GEMM precision types

Browse Source
This commit is contained in:
Sami Aario
2025-10-09 08:07:04 +00:00
parent e62c96f1dd
commit fc82ebc174
5 changed files with 51 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
include/ck_tile/ops/common.hpp
View File

@@ -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"

41
include/ck_tile/ops/common/determine_warp_prec_type.hpp Normal file
View File

@@ -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 <typename PrecType, typename OtherPrecType>
struct DetermineWarpPrecType
{
using prec_type = PrecType;
};
// For pk_int4_t, we convert to the other precision type.
template <typename OtherPrecType>
struct DetermineWarpPrecType<ck_tile::pk_int4_t, OtherPrecType>
{
using prec_type = OtherPrecType;
};
// For pk_fp4_t, we convert to the other precision type.
template <typename OtherPrecType>
struct DetermineWarpPrecType<ck_tile::pk_fp4_t, OtherPrecType>
{
using prec_type = OtherPrecType;
};
// For pk_fp4_raw_t, we convert to the other precision type.
template <typename OtherPrecType>
struct DetermineWarpPrecType<ck_tile::pk_fp4_raw_t, OtherPrecType>
{
using prec_type = OtherPrecType;
};
}; // namespace ck_tile

14
include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
View File

@@ -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 <type_traits>
@@ -92,16 +92,8 @@ struct CShuffleEpilogue
using ADataType = remove_cvref_t<std::tuple_element_t<number<0>{}, AsDataTypeTuple>>;
using BDataType = remove_cvref_t<std::tuple_element_t<number<0>{}, BsDataTypeTuple>>;
using ATypeToUse = std::conditional_t<std::is_same_v<ADataType, pk_int4_t> ||
std::is_same_v<ADataType, pk_fp4_t>,
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<BDataType, pk_int4_t> ||
std::is_same_v<BDataType, pk_fp4_t> ||
std::is_same_v<BDataType, pk_fp4_raw_t>,
ADataType,
BDataType>;
using ATypeToUse = typename DetermineWarpPrecType<ADataType, BDataType>::prec_type;
using BTypeToUse = typename DetermineWarpPrecType<BDataType, ADataType>::prec_type;
using ELayout = remove_cvref_t<typename Problem::ELayout>;
using CDElementwise = remove_cvref_t<typename Problem::CDElementwise>;

8
include/ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp
View File

@@ -94,12 +94,8 @@ struct BlockUniversalGemmAsBsCr
using ComputeDataType = remove_cvref_t<typename Traits::ComputeDataType>;
using CDataType = remove_cvref_t<typename Traits::CDataType>;
using ATypeToUse =
std::conditional_t<std::is_same_v<ADataType, pk_int4_t>, BDataType, ADataType>;
using BTypeToUse = std::conditional_t<std::is_same_v<BDataType, pk_int4_t> ||
std::is_same_v<BDataType, pk_fp4_raw_t>,
ADataType,
BDataType>;
using ATypeToUse = typename DetermineWarpPrecType<ADataType, BDataType>::prec_type;
using BTypeToUse = typename DetermineWarpPrecType<BDataType, ADataType>::prec_type;
using WarpGemm = remove_cvref_t<typename Traits::WarpGemm>;

12
include/ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp
View File

@@ -895,14 +895,10 @@ struct UniversalGemmPipelineAgBgCrPolicy
: vector_size * 4 == thread_elements ? WGAttrNumAccessEnum::Quad
: WGAttrNumAccessEnum::Invalid;
using ADataType = remove_cvref_t<typename Problem::ADataType>;
using BDataType = remove_cvref_t<typename Problem::BDataType>;
using ATypeToUse =
std::conditional_t<std::is_same_v<ADataType, pk_int4_t>, BDataType, ADataType>;
using BTypeToUse = std::conditional_t<std::is_same_v<BDataType, pk_int4_t> ||
std::is_same_v<BDataType, pk_fp4_raw_t>,
ADataType,
BDataType>;
using ADataType = remove_cvref_t<typename Problem::ADataType>;
using BDataType = remove_cvref_t<typename Problem::BDataType>;
using ATypeToUse = typename DetermineWarpPrecType<ADataType, BDataType>::prec_type;
using BTypeToUse = typename DetermineWarpPrecType<BDataType, ADataType>::prec_type;
using WarpGemm = WarpGemmDispatcher<ATypeToUse,
BTypeToUse,