mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-19 20:40:07 +00:00
Add support for mixed precision in contraction scale and bilinear (#973)
* Add support for mixed precision in contraction scale and bilinear (#936)
* Extract common functionality to separate files
* Reference contraction: Remove incorrect consts from type_converts
* Reference contraction: Add missing type_convert for dst value
* Reference contraction: Fix incorrect order of B matrix dimensions
* Add support for mixed precision in contraction scale and bilinear
* Move using statements from instances to a common file
* Move using statements from examples to a common file
* Fix the order of B matrix dimensions across examples and profiler
* Fix the computation of error threshold
* Make ComputeDataType an optional argument
* Include possible DataType -> ComputeDataType casting error in the threshold
* Remove commented code
* Make the ComputeDataType an optional argument in instance
---------
Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
[ROCm/composable_kernel commit: 4ef704d8a6]
This commit is contained in:
Bartlomiej Wroblewski
committed by
GitHub
GitHub
parent
3eaee9196f
commit
382a513acb
@@ -17,8 +17,9 @@
|
||||
static void print_helper_msg()
|
||||
{
|
||||
std::cout << "arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n"
|
||||
<< "arg2: data type (0: fp32; 1: f64)\n"
|
||||
<< "arg3: matrix layout (0: A[m0, m1, k0, k1] * B[k0, k1, n0, n1] + "
|
||||
<< "arg2: data type (0: fp32; 1: f64; 2: f16; 3: bf16)\n"
|
||||
<< "arg3: compute data type (0: fp32; 1: f64; 2: f16; 3: bf16)\n"
|
||||
<< "arg4: matrix layout (0: A[m0, m1, k0, k1] * B[k0, k1, n0, n1] + "
|
||||
"D[m0, m1, n0, n1] = E[m0, m1, n0, n1];\n"
|
||||
<< " 1: A[m0, m1, k0, k1] * B[n0, n1, k0, k1] + "
|
||||
"D[m0, m1, n0, n1] = E[m0, m1, n0, n1];\n"
|
||||
@@ -26,39 +27,40 @@ static void print_helper_msg()
|
||||
"D[m0, m1, n0, n1] = E[m0, m1, n0, n1];\n"
|
||||
<< " 3: A[k0, k1, m0, m1] * B[n0, n1, k0, k1] + "
|
||||
"D[m0, m1, n0, n1] = E[m0, m1, n0, n1])\n"
|
||||
<< "arg4: verification (0: no; 1: yes)\n"
|
||||
<< "arg5: initialization (0: no init; 1: integer value; 2: decimal "
|
||||
<< "arg5: verification (0: no; 1: yes)\n"
|
||||
<< "arg6: initialization (0: no init; 1: integer value; 2: decimal "
|
||||
<< "value)\n"
|
||||
<< "arg6: print tensor value (0: no; 1: yes)\n"
|
||||
<< "arg7: time kernel (0: no, 1: yes)\n"
|
||||
<< "arg8: alpha\n"
|
||||
<< "arg9 to 14: M0, M1, N0, N1, K0, K1\n"
|
||||
<< "arg15 to 30: Strides for A, B, D and E (skip for default)\n"
|
||||
<< "arg7: print tensor value (0: no; 1: yes)\n"
|
||||
<< "arg8: time kernel (0: no, 1: yes)\n"
|
||||
<< "arg9: alpha\n"
|
||||
<< "arg10 to 15: M0, M1, N0, N1, K0, K1\n"
|
||||
<< "arg16 to 31: Strides for A, B, D and E (skip for default)\n"
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
int profile_contraction_scale(int argc, char* argv[])
|
||||
{
|
||||
const bool default_strides = argc == 15;
|
||||
const bool default_strides = argc == 16;
|
||||
|
||||
if(argc != 31 && argc != 15)
|
||||
if(argc != 32 && argc != 16)
|
||||
{
|
||||
print_helper_msg();
|
||||
exit(1);
|
||||
}
|
||||
|
||||
const auto data_type = static_cast<ContractionDataType>(std::stoi(argv[2]));
|
||||
const auto layout = static_cast<ContractionMatrixLayout>(std::stoi(argv[3]));
|
||||
const bool do_verification = std::stoi(argv[4]);
|
||||
const ck::index_t init_method = std::stoi(argv[5]);
|
||||
const bool do_log = std::stoi(argv[6]);
|
||||
const bool time_kernel = std::stoi(argv[7]);
|
||||
const float alpha = std::stof(argv[8]);
|
||||
const auto compute_data_type = static_cast<ContractionComputeDataType>(std::stoi(argv[3]));
|
||||
const auto layout = static_cast<ContractionMatrixLayout>(std::stoi(argv[4]));
|
||||
const bool do_verification = std::stoi(argv[5]);
|
||||
const ck::index_t init_method = std::stoi(argv[6]);
|
||||
const bool do_log = std::stoi(argv[7]);
|
||||
const bool time_kernel = std::stoi(argv[8]);
|
||||
const float alpha = std::stof(argv[9]);
|
||||
|
||||
std::vector<ck::index_t> M;
|
||||
std::vector<ck::index_t> N;
|
||||
std::vector<ck::index_t> K;
|
||||
const ck::index_t dims_arg_num = 9;
|
||||
const ck::index_t dims_arg_num = 10;
|
||||
collect_index_params(argv, M, dims_arg_num, 2);
|
||||
collect_index_params(argv, N, dims_arg_num + 2, 2);
|
||||
collect_index_params(argv, K, dims_arg_num + 4, 2);
|
||||
@@ -75,88 +77,131 @@ int profile_contraction_scale(int argc, char* argv[])
|
||||
collect_index_params(argv, StridesD, dims_arg_num + 18, 4);
|
||||
}
|
||||
|
||||
using F32 = float;
|
||||
using F64 = double;
|
||||
using F16 = ck::half_t;
|
||||
using BF16 = ck::bhalf_t;
|
||||
using F32 = float;
|
||||
using F64 = double;
|
||||
|
||||
auto profile = [&](auto a_layout, auto b_layout, auto cde_layout, auto type) {
|
||||
using ALayout = decltype(a_layout);
|
||||
using BLayout = decltype(b_layout);
|
||||
using CDELayout = decltype(cde_layout);
|
||||
auto profile =
|
||||
[&](auto a_layout, auto b_layout, auto cde_layout, auto type, auto compute_type) {
|
||||
using ALayout = decltype(a_layout);
|
||||
using BLayout = decltype(b_layout);
|
||||
using CDELayout = decltype(cde_layout);
|
||||
|
||||
using DataType = decltype(type);
|
||||
using DataType = decltype(type);
|
||||
using ComputeDataType = decltype(compute_type);
|
||||
|
||||
if(default_strides)
|
||||
if(default_strides)
|
||||
{
|
||||
assign_default_strides(a_layout, StridesA, {M[0], M[1], K[0], K[1]});
|
||||
assign_default_strides(b_layout, StridesB, {N[0], N[1], K[0], K[1]});
|
||||
assign_default_strides(cde_layout, StridesE, {M[0], M[1], N[0], N[1]});
|
||||
assign_default_strides(cde_layout, StridesD, {M[0], M[1], N[0], N[1]});
|
||||
}
|
||||
|
||||
bool pass = ck::profiler::profile_contraction_impl<ALayout,
|
||||
BLayout,
|
||||
CDELayout,
|
||||
DataType,
|
||||
ComputeDataType,
|
||||
ck::Tuple<>,
|
||||
Scale>(do_verification,
|
||||
init_method,
|
||||
do_log,
|
||||
time_kernel,
|
||||
Scale{alpha},
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
StridesA,
|
||||
StridesB,
|
||||
StridesE,
|
||||
StridesD);
|
||||
|
||||
return pass;
|
||||
};
|
||||
|
||||
auto run_profile_for_datatype = [&](auto type, auto compute_type) {
|
||||
if(layout == ContractionMatrixLayout::MK_KN_MN_MN)
|
||||
{
|
||||
assign_default_strides(a_layout, StridesA, {M[0], M[1], K[0], K[1]});
|
||||
assign_default_strides(b_layout, StridesB, {K[0], K[1], N[0], N[1]});
|
||||
assign_default_strides(cde_layout, StridesE, {M[0], M[1], N[0], N[1]});
|
||||
assign_default_strides(cde_layout, StridesD, {M[0], M[1], N[0], N[1]});
|
||||
return profile(Row{}, Row{}, Row{}, type, compute_type);
|
||||
}
|
||||
|
||||
bool pass = ck::profiler::
|
||||
profile_contraction_impl<ALayout, BLayout, CDELayout, DataType, ck::Tuple<>, Scale>(
|
||||
do_verification,
|
||||
init_method,
|
||||
do_log,
|
||||
time_kernel,
|
||||
Scale{alpha},
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
StridesA,
|
||||
StridesB,
|
||||
StridesE,
|
||||
StridesD);
|
||||
|
||||
return pass;
|
||||
else if(layout == ContractionMatrixLayout::MK_NK_MN_MN)
|
||||
{
|
||||
return profile(Row{}, Col{}, Row{}, type, compute_type);
|
||||
}
|
||||
else if(layout == ContractionMatrixLayout::KM_KN_MN_MN)
|
||||
{
|
||||
return profile(Col{}, Row{}, Row{}, type, compute_type);
|
||||
}
|
||||
else if(layout == ContractionMatrixLayout::KM_NK_MN_MN)
|
||||
{
|
||||
return profile(Col{}, Col{}, Row{}, type, compute_type);
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
if(data_type == ContractionDataType::F32_F32_F32_F32 &&
|
||||
layout == ContractionMatrixLayout::MK_KN_MN_MN)
|
||||
if(data_type == ContractionDataType::F32_F32_F32_F32)
|
||||
{
|
||||
return profile(Row{}, Row{}, Row{}, F32{});
|
||||
if(compute_data_type == ContractionComputeDataType::F32)
|
||||
{
|
||||
return run_profile_for_datatype(F32{}, F32{});
|
||||
}
|
||||
else if(compute_data_type == ContractionComputeDataType::F16)
|
||||
{
|
||||
return run_profile_for_datatype(F32{}, F16{});
|
||||
}
|
||||
else if(compute_data_type == ContractionComputeDataType::BF16)
|
||||
{
|
||||
return run_profile_for_datatype(F32{}, BF16{});
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Incorrect combination of data type and compute data type." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
else if(data_type == ContractionDataType::F32_F32_F32_F32 &&
|
||||
layout == ContractionMatrixLayout::MK_NK_MN_MN)
|
||||
else if(data_type == ContractionDataType::F64_F64_F64_F64)
|
||||
{
|
||||
return profile(Row{}, Col{}, Row{}, F32{});
|
||||
if(compute_data_type == ContractionComputeDataType::F64)
|
||||
{
|
||||
return run_profile_for_datatype(F64{}, F64{});
|
||||
}
|
||||
else if(compute_data_type == ContractionComputeDataType::F32)
|
||||
{
|
||||
return run_profile_for_datatype(F64{}, F32{});
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Incorrect combination of data type and compute data type." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
else if(data_type == ContractionDataType::F32_F32_F32_F32 &&
|
||||
layout == ContractionMatrixLayout::KM_KN_MN_MN)
|
||||
else if(data_type == ContractionDataType::F16_F16_F16_F16)
|
||||
{
|
||||
return profile(Col{}, Row{}, Row{}, F32{});
|
||||
if(compute_data_type == ContractionComputeDataType::F32)
|
||||
{
|
||||
return run_profile_for_datatype(F16{}, F32{});
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Incorrect combination of data type and compute data type." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
else if(data_type == ContractionDataType::F32_F32_F32_F32 &&
|
||||
layout == ContractionMatrixLayout::KM_NK_MN_MN)
|
||||
else if(data_type == ContractionDataType::BF16_BF16_BF16_BF16)
|
||||
{
|
||||
return profile(Col{}, Col{}, Row{}, F32{});
|
||||
}
|
||||
else if(data_type == ContractionDataType::F64_F64_F64_F64 &&
|
||||
layout == ContractionMatrixLayout::MK_KN_MN_MN)
|
||||
{
|
||||
return profile(Row{}, Row{}, Row{}, F64{});
|
||||
}
|
||||
else if(data_type == ContractionDataType::F64_F64_F64_F64 &&
|
||||
layout == ContractionMatrixLayout::MK_NK_MN_MN)
|
||||
{
|
||||
return profile(Row{}, Col{}, Row{}, F64{});
|
||||
}
|
||||
else if(data_type == ContractionDataType::F64_F64_F64_F64 &&
|
||||
layout == ContractionMatrixLayout::KM_KN_MN_MN)
|
||||
{
|
||||
return profile(Col{}, Row{}, Row{}, F64{});
|
||||
}
|
||||
else if(data_type == ContractionDataType::F64_F64_F64_F64 &&
|
||||
layout == ContractionMatrixLayout::KM_NK_MN_MN)
|
||||
{
|
||||
return profile(Col{}, Col{}, Row{}, F64{});
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "this data_type & layout is not implemented" << std::endl;
|
||||
|
||||
return 1;
|
||||
if(compute_data_type == ContractionComputeDataType::F32)
|
||||
{
|
||||
return run_profile_for_datatype(BF16{}, F32{});
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "Incorrect combination of data type and compute data type." << std::endl;
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
REGISTER_PROFILER_OPERATION(OP_NAME, OP_DESC, profile_contraction_scale);
|
||||
|
||||
Reference in New Issue
Block a user