ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-15 03:30:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

refine universal gemm multi-config

Browse Source
This commit is contained in:
Qun Lin
2025-05-15 17:48:21 +08:00
parent 303336858f
commit feeb6f7cd8
2 changed files with 81 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
example/ck_tile/03_gemm/universal_gemm.cpp
View File

@@ -339,7 +339,7 @@ int run_gemm_example_prec_type(std::string a_layout, std::string b_layout, int a
}
}
template<typename GemmConfig>
template<template<typename PreType> typename GemmConfig>
int run_gemm_example(int argc, char* argv[])
{
auto [result, arg_parser] = create_args(argc, argv);
@@ -352,31 +352,35 @@ int run_gemm_example(int argc, char* argv[])
if(data_type == "fp16")
{
return run_gemm_example_prec_type<GemmConfig, ck_tile::half_t>(a_layout, b_layout, argc, argv);
return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, ck_tile::half_t>(a_layout, b_layout, argc, argv);
}
else if(data_type == "bf16")
{
return run_gemm_example_prec_type<GemmConfig, ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, ck_tile::bf16_t>(a_layout, b_layout, argc, argv);
}
else if(data_type == "fp8")
{
return run_gemm_example_prec_type<GemmConfig, ck_tile::fp8_t, ck_tile::fp8_t, ck_tile::half_t>(
return run_gemm_example_prec_type<GemmConfig<ck_tile::fp8_t>, ck_tile::fp8_t, ck_tile::fp8_t, ck_tile::half_t>(
a_layout, b_layout, argc, argv);
}
else if(data_type == "bf8")
{
return run_gemm_example_prec_type<GemmConfig, ck_tile::bf8_t, ck_tile::bf8_t, ck_tile::half_t>(
return run_gemm_example_prec_type<GemmConfig<ck_tile::bf8_t>, ck_tile::bf8_t, ck_tile::bf8_t, ck_tile::half_t>(
a_layout, b_layout, argc, argv);
}
#if(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V3)
else if(data_type == "pk_int4_t")
{
// TODO: Add support for bhalf_t ADataType
return run_gemm_example_prec_type<GemmConfig, ck_tile::half_t, ck_tile::pk_int4_t, ck_tile::half_t>(
a_layout, b_layout, argc, argv);
// TODO: Add support for bhalf_t ADataType
if constexpr(GemmConfig<ck_tile::half_t>::Pipeline == CK_TILE_PIPELINE_COMPUTE_V3)
{
return run_gemm_example_prec_type<GemmConfig<ck_tile::half_t>, ck_tile::half_t, ck_tile::pk_int4_t, ck_tile::half_t>(
a_layout, b_layout, argc, argv);
}
else
{
throw std::runtime_error("Unsupported pipeline for this operation !!!");
}
}
#endif
else
{
throw std::runtime_error("Unsupported data type for this operation !!!");
@@ -390,10 +394,12 @@ int main(int argc, char* argv[])
run_gemm_example<GemmConfig_1>(argc, argv);
run_gemm_example<GemmConfig_1_1>(argc, argv);
run_gemm_example<GemmConfig_2>(argc, argv);
//run_gemm_example<GemmConfig_2_1>(argc, argv);
//run_gemm_example<GemmConfig_2_1>(argc, argv); xx
run_gemm_example<GemmConfig_2_2>(argc, argv);
run_gemm_example<GemmConfig_2_3>(argc, argv);
run_gemm_example<GemmConfig_3>(argc, argv);
//run_gemm_example<GemmConfig_3_1>(argc, argv);
//run_gemm_example<GemmConfig_3_1>(argc, argv); xx
run_gemm_example<GemmConfig_3_2>(argc, argv);
}
catch(const std::runtime_error& e)
{