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

terrible agpr spill observed

Browse Source
This commit is contained in:
aska-0096
2025-05-23 07:51:26 +00:00
parent 5f0feabb90
commit 7151ef86f8
5 changed files with 44 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp
View File

@@ -404,17 +404,19 @@ struct BlockwiseGemmXdlops_mx_pipeline_base
protected:
// M1, N1 as double buffer index
// Read buffer + Compute buffer
// A[M0, M1, M2, KPack]
// A[MRepeat/MXdlPack, MWave, MXdlPack, KRepeat, KPack]
// 1 0 3 2 4
//
static constexpr auto a_thread_desc_ =
make_naive_tensor_descriptor(make_tuple(Number<MRepeat / MXdlPack>{},
I1,
Number<MXdlPack>{},
Number<KRepeat>{},
Number<KPack>{}),
make_tuple(Number<KPack * MXdlPack>{},
Number<KRepeat * MRepeat * KPack>{},
Number<MRepeat * KPack>{},
make_tuple(Number<KRepeat * MXdlPack* KPack>{},
Number<MRepeat * KRepeat * KPack>{},
Number<KPack>{},
Number<MXdlPack*KPack>{},
I1));
// B[N0, N1, N2, KPack]
@@ -424,10 +426,10 @@ struct BlockwiseGemmXdlops_mx_pipeline_base
Number<KRepeat>{},
Number<NXdlPack>{},
Number<KPack>{}),
make_tuple(Number<KPack * NXdlPack>{},
Number<KRepeat * NRepeat * KPack>{},
Number<NRepeat * KPack>{},
make_tuple(Number<KRepeat * NXdlPack* KPack>{},
Number<NRepeat * KRepeat * KPack>{},
Number<KPack>{},
Number<NXdlPack*KPack>{},
I1));
// C[M, N, NumRegXdlops]

7
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp
View File

@@ -246,8 +246,9 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx<BlockGemmPipelineScheduler::Intrawave,
constexpr auto num_ds_read_a_scale = MRepeat / MXdlPack * KRepeat / KXdlPack;
constexpr auto num_ds_read_b_scale = NRepeat / NXdlPack * KRepeat / KXdlPack;
constexpr auto num_dsread_a_scale_mfma = num_ds_read_a_scale;
constexpr auto num_dsread_b_scale_mfma = num_ds_read_b_scale;
// seems like compiler will use ds_read2_b32 for it.
constexpr auto num_dsread_a_scale_mfma = num_ds_read_a_scale/2;
constexpr auto num_dsread_b_scale_mfma = num_ds_read_b_scale/2;
// To relieve TA backport pressure:
// give small interval for buffer_load_dword
// give large interval for buffer_load_dwordx4
@@ -513,7 +514,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx<BlockGemmPipelineScheduler::Intrawave,
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
a_scale_blockwise_copy.Run(
a_scale_grid_desc, a_scale_grid_buf, a_scale_block_desc, a_scale_block_bufs(I0));
b_scale_blockwise_copy.Run(

9
include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp
View File

@@ -41,6 +41,7 @@ namespace ck {
*/
template <typename ThreadGroup,
typename BlockSliceLengths,
typename WaveClusterLengths, // Ugly, FIX me.
typename ThreadClusterLengths,
typename ThreadClusterArrangeOrder,
typename SrcData,
@@ -68,12 +69,8 @@ struct ThreadGroupTensorSliceTransfer_DirectLoad
static constexpr auto block_slice_lengths = BlockSliceLengths{};
static constexpr auto thread_cluster_lengths = ThreadClusterLengths{};
static constexpr auto wave_thread_cluster_lengths =
Sequence<ThreadClusterLengths{}.At(I0),
ThreadClusterLengths{}.At(I1) * 64 / ThreadGroup::GetNumOfThread(),
1>{};
static constexpr auto wave_cluster_lengths =
Sequence<1, ThreadGroup::GetNumOfThread() / 64, 1>{};
static constexpr auto wave_cluster_lengths = WaveClusterLengths{};
static constexpr auto wave_thread_cluster_lengths = thread_cluster_lengths/wave_cluster_lengths;
static constexpr auto thread_single_load_size = generate_sequence(
detail::lambda_scalar_per_access<DstVectorDim, ScalarPerVector>{}, Number<nDim>{});

22
include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp
View File

@@ -416,6 +416,28 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX<ALayout,
Run(kernel);
}
}
else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3){
if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
{
const auto kernel =
kernel_gemm_xdl_cshuffle_v3_2lds<GridwiseGemm,
false,
InMemoryDataOperationEnum::Set,
minimum_occupancy,
TailNumber::Odd>;
Run(kernel);
}
else
{
const auto kernel =
kernel_gemm_xdl_cshuffle_v3_2lds<GridwiseGemm,
false,
InMemoryDataOperationEnum::Set,
minimum_occupancy,
TailNumber::Even>;
Run(kernel);
}
}
}
return ave_time;

6
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp
View File

@@ -1253,6 +1253,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
auto a_blockwise_copy =
ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
Sequence<1, BlockSize/64, 1>,
Sequence<AK0Number, MPerBlock, AK1Number>,
ABlockTransferThreadClusterLengths_AK0_M_AK1,
ABlockTransferThreadClusterArrangeOrder,
@@ -1272,6 +1273,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
// B matrix blockwise copy
auto b_blockwise_copy =
ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
Sequence<1, BlockSize/64, 1>,
Sequence<BK0Number, NPerBlock, BK1Number>,
BBlockTransferThreadClusterLengths_BK0_N_BK1,
BBlockTransferThreadClusterArrangeOrder,
@@ -1755,6 +1757,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
auto a_blockwise_copy =
ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
Sequence<AK0Number, MPerBlock, AK1Number>,
Sequence<1, BlockSize/64, 1>,
ABlockTransferThreadClusterLengths_AK0_M_AK1,
ABlockTransferThreadClusterArrangeOrder,
ADataType,
@@ -1774,6 +1777,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
auto b_blockwise_copy =
ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
Sequence<BK0Number, NPerBlock, BK1Number>,
Sequence<1, BlockSize/64, 1>,
BBlockTransferThreadClusterLengths_BK0_N_BK1,
BBlockTransferThreadClusterArrangeOrder,
BDataType,
@@ -1828,6 +1832,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
Sequence<MPerBlock/MPerXdl / MXdlPack,
KRepeat / KXdlPack,
64* KXdlPack * MXdlPack / scale_pack_size_b>,
Sequence<BlockSize / 64, 1, 1>,
Sequence<BlockSize / 64, 1, 64>,
Sequence<0, 1, 2>,
AScaleDataType,
@@ -1847,6 +1852,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
Sequence<NPerBlock/NPerXdl / NXdlPack,
KRepeat / KXdlPack,
64* KXdlPack * NXdlPack / scale_pack_size_b>,
Sequence<BlockSize / 64, 1, 1>,
Sequence<BlockSize / 64, 1, 64>,
Sequence<0, 1, 2>,
BScaleDataType,