ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-04-19 22:39:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix a_wrap preload issue for large MPerBlock.

Browse Source
This commit is contained in:
mtgu0705
2025-09-18 01:19:03 -05:00
parent f2db44710f
commit 92ad6fcc0a
1 changed files with 35 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
include/ck_tile/ops/flatmm/pipeline/mx_flatmm_pipeline_agmem_bgmem_creg_v1.hpp
View File

@@ -876,13 +876,13 @@ struct MXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Probl
(mIter_pack * MXdlPack + imxdl)) <
(KIterPerWarp * MIterPerWarp - m_preload))
{
constexpr auto AmIter =
(mIter_pack * MXdlPack + imxdl + m_preload) % MIterPerWarp;
constexpr auto AkIter =
(kIter_pack * KXdlPack + ikxdl +
(mIter_pack * MXdlPack + imxdl + m_preload) /
MIterPerWarp);
ua_ping.mxfp4 = load_tile(
constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 +
(kIter_pack * KXdlPack + ikxdl) * 2 +
(mIter_pack * MXdlPack + imxdl) / 2 * 4 +
m_preload;
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
ua_ping.mxfp4 = load_tile(
a_warp_windows_ping(number<AmIter>{})(number<AkIter>{}));
a_warp_tensor(number<AwarpIter>{}) = ua_ping.u;
}
@@ -1016,13 +1016,13 @@ struct MXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Probl
(mIter_pack * MXdlPack + imxdl)) <
(KIterPerWarp * MIterPerWarp - m_preload))
{
constexpr auto AmIter =
(mIter_pack * MXdlPack + imxdl + m_preload) % MIterPerWarp;
constexpr auto AkIter =
(kIter_pack * KXdlPack + ikxdl +
(mIter_pack * MXdlPack + imxdl + m_preload) /
MIterPerWarp);
ua_pong.mxfp4 = load_tile(
constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 +
(kIter_pack * KXdlPack + ikxdl) * 2 +
(mIter_pack * MXdlPack + imxdl) / 2 * 4 +
m_preload;
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
ua_pong.mxfp4 = load_tile(
a_warp_windows_pong(number<AmIter>{})(number<AkIter>{}));
a_warp_tensor(number<AwarpIter>{}) = ua_pong.u;
}
@@ -1156,13 +1156,13 @@ struct MXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Probl
(mIter_pack * MXdlPack + imxdl)) <
(KIterPerWarp * MIterPerWarp - m_preload))
{
constexpr auto AmIter =
(mIter_pack * MXdlPack + imxdl + m_preload) % MIterPerWarp;
constexpr auto AkIter =
(kIter_pack * KXdlPack + ikxdl +
(mIter_pack * MXdlPack + imxdl + m_preload) /
MIterPerWarp);
ua_ping.mxfp4 = load_tile(
constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 +
(kIter_pack * KXdlPack + ikxdl) * 2 +
(mIter_pack * MXdlPack + imxdl) / 2 * 4 +
m_preload;
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
ua_ping.mxfp4 = load_tile(
a_warp_windows_ping(number<AmIter>{})(number<AkIter>{}));
a_warp_tensor(number<AwarpIter>{}) = ua_ping.u;
}
@@ -1240,13 +1240,13 @@ struct MXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Probl
(mIter_pack * MXdlPack + imxdl)) <
(KIterPerWarp * MIterPerWarp - m_preload))
{
constexpr auto AmIter =
(mIter_pack * MXdlPack + imxdl + m_preload) % MIterPerWarp;
constexpr auto AkIter =
(kIter_pack * KXdlPack + ikxdl +
(mIter_pack * MXdlPack + imxdl + m_preload) /
MIterPerWarp);
ua_pong.mxfp4 = load_tile(
constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 +
(kIter_pack * KXdlPack + ikxdl) * 2 +
(mIter_pack * MXdlPack + imxdl) / 2 * 4 +
m_preload;
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
ua_pong.mxfp4 = load_tile(
a_warp_windows_pong(number<AmIter>{})(number<AkIter>{}));
a_warp_tensor(number<AwarpIter>{}) = ua_pong.u;
}
@@ -1319,13 +1319,13 @@ struct MXF4FlatmmPipelineAGmemBGmemCRegV1 : FlatmmPipelineAGmemBGmemCRegV1<Probl
(mIter_pack * MXdlPack + imxdl)) <
(KIterPerWarp * MIterPerWarp - m_preload))
{
constexpr auto AmIter =
(mIter_pack * MXdlPack + imxdl + m_preload) % MIterPerWarp;
constexpr auto AkIter =
(kIter_pack * KXdlPack + ikxdl +
(mIter_pack * MXdlPack + imxdl + m_preload) /
MIterPerWarp);
ua_ping.mxfp4 = load_tile(
constexpr auto addr = (mIter_pack * MXdlPack + imxdl) % 2 +
(kIter_pack * KXdlPack + ikxdl) * 2 +
(mIter_pack * MXdlPack + imxdl) / 2 * 4 +
m_preload;
constexpr auto AmIter = addr % 2 + addr / 4 * 2;
constexpr auto AkIter = addr / 2 % 2;
ua_ping.mxfp4 = load_tile(
a_warp_windows_ping(number<AmIter>{})(number<AkIter>{}));
a_warp_tensor(number<AwarpIter>{}) = ua_ping.u;
}