Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "79c3e8ad73adb83cb861ba1e4f514663faeb1079"
Commit 1867ffa3 authored by mtgu0705's avatar mtgu0705
Browse files

Add Scale_Block_M = 1 based on ab_scale support

parent f728087c
Hide whitespace changes
Inline Side-by-side
Showing with 61 additions and 35 deletions
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp
View file @ 1867ffa3
......@@ -338,18 +338,27 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlockGemmPipelineScheduler::Intr
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
static_for<0, MRepeat, 1>{}([&](auto m0){
a_scale_thread_copy.Run(a_scale_grid_desc,
a_scale_grid_buf,
a_scale_thread_desc,
make_tuple(m0, I0),
a_scale_thread_buf);
static_for<0, MRepeat, 1>{}([&](auto m0) {
static_for<0, Number<16>{}, 1>{}([&](auto i0) {
a_scale_thread_copy.Run(a_scale_grid_desc,
a_scale_grid_buf,
a_scale_thread_desc,
make_tuple(m0, i0, I0),
a_scale_thread_buf);
a_scale_thread_copy.MoveSrcSliceWindow(
a_scale_grid_desc, a_scale_thread_copy_step.At(Number<2>{}));
if((i0 + 1) % 4 == 0)
{
a_scale_thread_copy.MoveSrcSliceWindow(
a_scale_grid_desc, a_scale_thread_copy_step.At(Number<1>{}));
}
});
a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc,
a_scale_thread_copy_step.At(Number<0>{}));
});
a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc,
a_scale_thread_copy_step.At(Number<1>{}));
a_scale_thread_copy_step.At(Number<3>{}));
b_scale_thread_copy.Run(b_scale_grid_desc,
b_scale_grid_buf,
......@@ -443,7 +452,8 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlockGemmPipelineScheduler::Intr
c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t));
c_thread_buf(Number<c_offset>{}) +=
c_thread_buf_per_scale[Number<t>{}] *
type_convert<AccDataType>(a_scale_thread_buf[m0]) *
type_convert<AccDataType>(a_scale_thread_buf(Number<m0 * xdlops_gemm.GetRegSizePerXdlops() + t>{})) *
// type_convert<AccDataType>(a_scale_thread_buf(Number<a_offset>{})) *
type_convert<AccDataType>(b_scale_thread_buf[I0]);
});
});
......@@ -468,18 +478,27 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlockGemmPipelineScheduler::Intr
b_thread_buf);
});
});
static_for<0,MRepeat,1>{}([&](auto m0){
a_scale_thread_copy.Run(a_scale_grid_desc,
a_scale_grid_buf,
a_scale_thread_desc,
make_tuple(m0, I0),
a_scale_thread_buf);
static_for<0, MRepeat, 1>{}([&](auto m0) {
static_for<0, Number<16>{}, 1>{}([&](auto i0) {
a_scale_thread_copy.Run(a_scale_grid_desc,
a_scale_grid_buf,
a_scale_thread_desc,
make_tuple(m0, i0, I0),
a_scale_thread_buf);
a_scale_thread_copy.MoveSrcSliceWindow(
a_scale_grid_desc, a_scale_thread_copy_step.At(Number<2>{}));
if((i0 + 1) % 4 == 0)
{
a_scale_thread_copy.MoveSrcSliceWindow(
a_scale_grid_desc, a_scale_thread_copy_step.At(Number<1>{}));
}
});
a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc,
a_scale_thread_copy_step.At(Number<0>{}));
a_scale_thread_copy_step.At(Number<0>{}));
});
a_scale_thread_copy.MoveSrcSliceWindow(a_scale_grid_desc,
a_scale_thread_copy_step.At(Number<1>{}));
a_scale_thread_copy_step.At(Number<3>{}));
b_scale_thread_copy.Run(b_scale_grid_desc,
b_scale_grid_buf,
......@@ -525,9 +544,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlockGemmPipelineScheduler::Intr
constexpr index_t c_offset =
c_thread_desc_.CalculateOffset(make_tuple(m0, n0, t));
c_thread_buf(Number<c_offset>{}) +=
c_thread_buf_per_scale[Number<t>{}] *
type_convert<AccDataType>(a_scale_thread_buf[m0]) *
type_convert<AccDataType>(b_scale_thread_buf[I0]);
c_thread_buf_per_scale[Number<t>{}] *
type_convert<AccDataType>(a_scale_thread_buf(Number<m0 * xdlops_gemm.GetRegSizePerXdlops() + t>{})) *
// type_convert<AccDataType>(a_scale_thread_buf(Number<a_offset>{})) *
type_convert<AccDataType>(b_scale_thread_buf[I0]);
});
});
});
......
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_ab_scale.hpp
View file @ 1867ffa3
......@@ -1213,8 +1213,9 @@ struct GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3
const auto a_scale_grid_desc_am_ak = make_naive_tensor_descriptor(
make_tuple(math::integer_divide_ceil(problem.M, ScaleBlockM),
16,
math::integer_divide_ceil(problem.K, ScaleBlockK)),
make_tuple(math::integer_divide_ceil(problem.K, ScaleBlockK), 1));
make_tuple(math::integer_divide_ceil(problem.K, ScaleBlockK),0, 1));
const auto b_scale_grid_desc_bn_ak = make_naive_tensor_descriptor(
make_tuple(math::integer_divide_ceil(problem.N, ScaleBlockN),
math::integer_divide_ceil(problem.K, ScaleBlockK)),
......@@ -1357,34 +1358,36 @@ struct GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3
(a_grid_desc_ak0_m_ak1.GetLength(I0) * a_grid_desc_ak0_m_ak1.GetLength(I2)) /
KPerBlock);
const index_t ScaleSliceSizeM = MXdlPerWave;
const index_t ScaleSliceSizeN = 1;
const index_t ScaleSliceSizeK = 1;
static constexpr auto xdlops_gemm =
XdlopsGemm<ComputeTypeA, MPerXdl, NPerXdl, KPack, ComputeTypeA, false>{};
const index_t ScaleSliceSizeM = MXdlPerWave;
const index_t RegSizePerXdlops = xdlops_gemm.GetRegSizePerXdlops();
const index_t ScaleSliceSizeN = 1;
const index_t ScaleSliceSizeK = 1;
constexpr auto a_scale_thread_desc = make_naive_tensor_descriptor_packed(
make_tuple(Number<ScaleSliceSizeM>{}, Number<ScaleSliceSizeK>{}));
make_tuple(Number<ScaleSliceSizeM>{}, Number<RegSizePerXdlops>{}, Number<ScaleSliceSizeK>{}));
constexpr auto b_scale_thread_desc = make_naive_tensor_descriptor_packed(
make_tuple(Number<ScaleSliceSizeN>{}, Number<ScaleSliceSizeK>{}));
constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXdl);
// auto a_thread_offset =
// get_thread_local_1d_id() % MPerXdl + (get_thread_local_1d_id() / 64) % MWaves * MPerXdl;
auto a_thread_offset = get_thread_local_1d_id() % MPerXdl + (get_thread_local_1d_id() / 128) * MPerXdl;
auto a_thread_offset =
get_thread_local_1d_id() / 128 * 32 + get_thread_local_1d_id() % 64 / 32 * 4;
auto a_scale_thread_copy =
ThreadwiseTensorSliceTransfer_v2<AScaleType,
AScaleType,
decltype(a_scale_grid_desc_am_ak),
decltype(a_scale_thread_desc),
Sequence<1, ScaleSliceSizeK>,
Sequence<0, 1>,
Sequence<1, 1, ScaleSliceSizeK>,
Sequence<0, 1, 2>,
1,
1,
1,
false>(
a_scale_grid_desc_am_ak, make_multi_index(block_m_id * MPerBlock / ScaleBlockM + a_thread_offset, 0));
a_scale_grid_desc_am_ak, make_multi_index(block_m_id * MPerBlock / ScaleBlockM + a_thread_offset,0, 0));
auto b_scale_thread_copy =
ThreadwiseTensorSliceTransfer_v2<BScaleType,
......@@ -1400,7 +1403,10 @@ struct GridwiseGemmMultiD_ABScale_xdl_cshuffle_v3
b_scale_grid_desc_bn_ak, make_multi_index(block_n_id * NPerBlock / ScaleBlockN, 0));
constexpr auto a_scale_thread_slice_copy_step =
make_tuple(make_multi_index(MWaves * MPerXdl, 0), make_multi_index(-MPerBlock, 1));
make_tuple(make_multi_index(MPerXdl, 0, 0),
make_multi_index(4, 0, 0),
make_multi_index(1, 0, 0),
make_multi_index(-MPerBlock, 0, 1));
constexpr auto b_scale_thread_slice_copy_step = make_multi_index(0, 1);
const index_t num_k_block_per_scale = ScaleBlockK / KPerBlock;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment