Global mem write debugging.

This commit is contained in:
Ville Pietilä
2026-01-27 11:23:37 -05:00
parent 99b38bd260
commit 970869661c
6 changed files with 154 additions and 83 deletions

View File

@@ -13,6 +13,7 @@
#include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp"
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

View File

@@ -4,7 +4,7 @@
// Best instance for most difficult fwd case
template <ck::index_t NDimSpatial>
using DeviceConvFwdInstance =
ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle<
ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<
NDimSpatial,
InputLayout<NDimSpatial>,
WeightLayout<NDimSpatial>,
@@ -21,7 +21,6 @@ using DeviceConvFwdInstance =
OutElementOp,
ConvSpec, // ConvForwardSpecialization
GemmSpec, // GemmSpecialization
1, //
256, // BlockSize
64, // MPerBlock
64, // NPerBlock
@@ -46,14 +45,16 @@ using DeviceConvFwdInstance =
4, // BBlockTransferSrcScalarPerVector
8, // BBlockTransferDstScalarPerVector_BK1
1, // BBlockLdsExtraN
1,
1,
S<1, 32, 1, 2, 2>, // CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
1, // Vector load/store size for output tensor = CDEBlockTransferScalarPerVector_NPerBlock
1, // CShuffleMXdlPerWavePerShuffle
1, // CShuffleNXdlPerWavePerShuffle
S<1, 32, 1, 4>, // CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, must match with num merged groups
2, // Vector load/store size for output tensor = CDEBlockTransferScalarPerVector_NPerBlock
ck::BlockGemmPipelineScheduler::Intrawave,
ck::BlockGemmPipelineVersion::v1,
InKernelDataType,
WeiKernelDataType,
ck::LoopScheduler::Default,
2>; // Number of merged groups
WeiKernelDataType,
false, // No direct load
2>; // Num groups to merge
// Best instance for 2nd most difficult fwd case, increase the B block src vector size?
/*

View File

@@ -357,6 +357,11 @@ check_err(const Range& out,
}
res = false;
}
else
{
std::cerr << "Success: Correct results!" << std::setw(12) << std::setprecision(7) << " out[" << i
<< "] == ref[" << i << "]: " << o << " == " << r << std::endl;
}
}
if(!res)
{

View File

@@ -382,7 +382,8 @@ template <index_t NDimSpatial,
// in tuple for MultiAB), unpack if tuple was
// passed
typename BComputeDataType = AComputeDataType,
bool DirectLoad = false>
bool DirectLoad = false,
index_t NumGroupsToMerge = 1>
struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
: public DeviceGroupedConvFwdMultipleABD<NDimSpatial,
ALayout,
@@ -418,6 +419,8 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
Wave32MaxMNPerXDL,
NXdlPerWave*(NPerXDL / Wave32MaxMNPerXDL)>();
static_assert(NumGroupsToMerge >= 1);
static constexpr bool isMultiA = is_detected<is_tuple, ADataType>::value;
static constexpr bool isMultiB = is_detected<is_tuple, BDataType>::value;
static constexpr bool isMultiD = DsDataType::Size() > 0;
@@ -447,7 +450,8 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
ConvForwardSpecialization,
true /*SplitN*/,
ADataType,
EDataType>;
EDataType,
NumGroupsToMerge>;
using ComputePtrOffset = ComputePtrOffsetOfStridedBatch<I1, I1, NumDTensor>;
@@ -783,8 +787,8 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
cde_element_op_{cde_element_op}
{
// A/B/E Batch/N Stride
compute_ptr_offset_of_groups_.BatchStrideA_ = a_g_n_c_wis_strides_[0];
compute_ptr_offset_of_groups_.BatchStrideB_ = b_g_k_c_xs_strides_[0];
compute_ptr_offset_of_groups_.BatchStrideA_ = a_g_n_c_wis_strides_[0] * NumGroupsToMerge;
compute_ptr_offset_of_groups_.BatchStrideB_ = b_g_k_c_xs_strides_[0] * NumGroupsToMerge;
compute_ptr_offset_of_n_.BatchStrideA_ = a_g_n_c_wis_strides_[1] * conv_N_per_block_;
// p_as and p_bs are pointers
@@ -795,7 +799,7 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
static_for<0, NumDTensor, 1>{}([&](auto i) {
using DLayout = remove_cvref_t<tuple_element_t<i.value, DsLayout>>;
// D batch stride
compute_ptr_offset_of_groups_.BatchStrideDs_(i) = ds_g_n_k_wos_strides_[i][0];
compute_ptr_offset_of_groups_.BatchStrideDs_(i) = ds_g_n_k_wos_strides_[i][0] * NumGroupsToMerge;
compute_ptr_offset_of_n_.BatchStrideDs_(i) =
ds_g_n_k_wos_strides_[i][1] * conv_N_per_block_;
@@ -815,7 +819,7 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
DeviceOp::MakeEGridDescriptor_M_N<DLayout>(conv_to_gemm_transformer_d);
});
compute_ptr_offset_of_groups_.BatchStrideE_ = e_g_n_k_wos_strides_[0];
compute_ptr_offset_of_groups_.BatchStrideE_ = e_g_n_k_wos_strides_[0] * NumGroupsToMerge;
compute_ptr_offset_of_n_.BatchStrideE_ = e_g_n_k_wos_strides_[1] * conv_N_per_block_;
if constexpr(is_NGCHW_GKCYX_NGKHW<ALayout, BLayout, ELayout>() ||
@@ -998,7 +1002,7 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
std::tie(gdx, gdy, gdz) =
GridwiseGemm::CalculateGridSize(GemmM, GemmN, I1 /*arg.KBatch*/);
gdy = arg.num_group_;
gdy = arg.num_group_ / NumGroupsToMerge;
gdz = num_workgroups_per_Conv_N;
index_t K_split = (GemmK + KPerBlock - 1) / KPerBlock * KPerBlock;
@@ -1498,6 +1502,19 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
}
}
if constexpr(NumGroupsToMerge > 1)
{
if(G % NumGroupsToMerge != 0)
{
if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
{
std::cout << "Unsupported! G % NumGroupsToMerge != 0: G=" << G
<< ", NumGroupsToMerge=" << NumGroupsToMerge << std::endl;
}
return false;
}
}
if(get_device_name() == "gfx908")
{
// FIXME: re-enable fp64 when SWDEV-335738 is fixed
@@ -2124,7 +2141,8 @@ struct DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3
<< "BlkGemmPipelineScheduler: "
<< BlkGemmPipelineSchedulerToString[BlkGemmPipeSched] << ", "
<< "BlkGemmPipelineVersion: "
<< BlkGemmPipelineVersionToString[BlkGemmPipelineVer]
<< BlkGemmPipelineVersionToString[BlkGemmPipelineVer] << ", "
<< NumGroupsToMerge
<< ">";
// clang-format on

View File

@@ -96,6 +96,8 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
static constexpr auto AK0PerBlock = Number<KPerBlock / AK1Value>{};
static constexpr auto BK0PerBlock = Number<KPerBlock / BK1Value>{};
static constexpr ck::index_t Gm = 2;
using ThisThreadBlock = ThisThreadBlock<BlockSize>;
using GridwiseGemmPipe = remove_cvref_t<
@@ -145,24 +147,6 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
return c_shuffle_block_desc_mblock_mperblock_nblock_nperblock;
}
__host__ __device__ static constexpr auto
GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerGroup_Gm()
{
constexpr ck::index_t Gm = 2;
constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl);
constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl);
constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_npergroup_gm =
make_naive_tensor_descriptor_packed(
make_tuple(I1,
Number<CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl>{},
I1,
Number<CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl / Gm>{},
Number<Gm>{}));
return c_shuffle_block_desc_mblock_mperblock_nblock_npergroup_gm;
}
// ck::Tuple<const D0DataType*, const D1DataType*, ...>
static constexpr auto MakeDsGridPointer()
{
@@ -244,21 +228,22 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
const auto M = e_grid_desc_m_n.GetLength(I0);
const auto N = e_grid_desc_m_n.GetLength(I1);
constexpr ck::index_t Gm = 2;
const auto MPerGroup = MPerBlock / Gm;
const auto NPerGroup = NPerBlock / Gm;
const auto MBlock = M / MPerBlock;
const auto NBlock = N / NPerBlock;
std::cout << "MBlock: " << MBlock << ", NBlock: " << NBlock << std::endl;
std::cout << "MPerBlock: " << MPerBlock << ", NPerBlock: " << NPerBlock << ", NPerGroup: " << NPerGroup << std::endl;
std::cout << "MPerBlock: " << MPerBlock << ", NPerBlock: " << NPerBlock << std::endl;
std::cout << "MPerGroup: " << MPerGroup << ", NPerGroup: " << NPerGroup << std::endl;
const auto e_grid_desc_mblock_mperblock_nblock_nperblock = transform_tensor_descriptor(
e_grid_desc_m_n,
make_tuple(make_unmerge_transform(make_tuple(MBlock, Number<MPerBlock>{})),
make_unmerge_transform(make_tuple(NBlock, Number<NPerGroup>{}, Number<Gm>{}))),
make_unmerge_transform(make_tuple(NBlock, Number<NPerBlock>{}))),
make_tuple(Sequence<0>{}, Sequence<1>{}),
make_tuple(Sequence<0, 1>{}, Sequence<2, 3, 4>{}));
make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}));
return e_grid_desc_mblock_mperblock_nblock_nperblock;
}
@@ -778,9 +763,6 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =
GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock();
constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_npergroup_gm =
GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerGroup_Gm();
auto c_shuffle_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
static_cast<CShuffleDataType*>(p_shared),
c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());
@@ -888,9 +870,8 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
vpgr_to_lds_element_op()};
// tuple of reference to C/Ds tensor descriptors
// This needs to be 5D, i.e., split NPerBlock into (NPerBlock/Gm, Gm)
const auto c_ds_desc_refs = concat_tuple_of_reference(
tie(c_shuffle_block_desc_mblock_mperblock_nblock_npergroup_gm),
tie(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock),
generate_tie([&](auto i) -> const auto& // return type should be reference
{ return ds_grid_desc_mblock_mperblock_nblock_nperblock[i]; },
Number<NumDTensor>{}));
@@ -904,15 +885,13 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
// tuple of starting index of C/Ds blockwise copy
const auto idx_c_ds_block_begin = container_concat(
make_tuple(make_multi_index(0, 0, 0, 0, 0)),
make_tuple(make_multi_index(0, 0, 0, 0)),
generate_tuple(
[&](auto) {
return make_multi_index(block_work_idx[I0], 0, block_work_idx[I1], 0, 0);
return make_multi_index(block_work_idx[I0], 0, block_work_idx[I1], 0);
},
Number<NumDTensor>{}));
constexpr ck::index_t Gm = 2;
// blockwise copy C/D/E between LDS and global
auto cde_block_copy_lds_and_global = ThreadGroupTensorSliceTransfer_v7<
ThisThreadBlock,
@@ -928,11 +907,10 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
Sequence<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl/Gm,
Gm>, // BlockSliceLengths,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths,
CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
Sequence<0, 1, 2, 3, 4>, // typename ThreadClusterArrangeOrder,
Sequence<0, 1, 2, 3, 4>, // typename DimAccessOrder,
Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder,
Sequence<0, 1, 2, 3>, // typename DimAccessOrder,
3, // index_t VectorDim,
CDEShuffleBlockTransferScalarPerVector_NPerBlock, // ScalarPerVector
sequence_merge_t<
@@ -943,7 +921,7 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
{c_ds_desc_refs,
idx_c_ds_block_begin, // Src (LDS) slice origins
tie(e_grid_desc_mblock_mperblock_nblock_nperblock),
make_tuple(make_multi_index(block_work_idx[I0], 0, block_work_idx[I1], 0, 0)), // Dst (global) slice origins
make_tuple(make_multi_index(block_work_idx[I0], 0, block_work_idx[I1], 0)), // Dst (global) slice origins
lds_to_global_element_op()};
// space filling curve for threadwise C in VGPR before shuffle
@@ -961,13 +939,12 @@ struct GridwiseGemmMultipleD_xdl_cshuffle
// space filling curve for shuffled blockwise C/D/E
constexpr auto sfc_cde_block =
SpaceFillingCurve<Sequence<1, MPerBlock, 1, NPerBlock/Gm, Gm>,
Sequence<0, 2, 1, 3, 4>,
SpaceFillingCurve<Sequence<1, MPerBlock, 1, NPerBlock>,
Sequence<0, 2, 1, 3>,
Sequence<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl / Gm,
Gm>>{};
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>{};
constexpr index_t num_access = sfc_c_vgpr.GetNumOfAccess();

View File

@@ -1850,12 +1850,12 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
MakeDsGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
ds_grid_desc_m_n, problem.MBlock, problem.NBlock);
const auto ds_grid_buf = generate_tuple(
[&](auto i) {
return make_dynamic_buffer<AddressSpaceEnum::Global>(
p_ds_grid[i], ds_grid_desc_m_n[i].GetElementSpaceSize());
},
Number<NumDTensor>{});
// const auto ds_grid_buf = generate_tuple(
// [&](auto i) {
// return make_dynamic_buffer<AddressSpaceEnum::Global>(
// p_ds_grid[i], ds_grid_desc_m_n[i].GetElementSpaceSize());
// },
// Number<NumDTensor>{});
// tuple of reference to C/Ds tensor descriptors
const auto c_ds_desc_refs = concat_tuple_of_reference(
@@ -1865,11 +1865,11 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
Number<NumDTensor>{}));
// tuple of reference to C/Ds tensor descriptors
const auto c_ds_buf_refs = concat_tuple_of_reference(
tie(c_shuffle_block_buf),
generate_tie([&](auto i) -> const auto& // return type should be reference
{ return ds_grid_buf[i]; },
Number<NumDTensor>{}));
// const auto c_ds_buf_refs = concat_tuple_of_reference(
// tie(c_shuffle_block_buf),
// generate_tie([&](auto i) -> const auto& // return type should be reference
// { return ds_grid_buf[i]; },
// Number<NumDTensor>{}));
// tuple of starting index of C/Ds blockwise copy
const auto idx_c_ds_block_begin = container_concat(
@@ -1891,8 +1891,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
ThisThreadBlock,
decltype(container_concat(make_tuple(CShuffleDataType{}), DsDataType{})),
Tuple<EDataType>,
decltype(c_ds_desc_refs),
decltype(tie(e_grid_desc_mblock_mperblock_nblock_nperblock)),
decltype(c_ds_desc_refs), // SrcDesc
decltype(tie(e_grid_desc_mblock_mperblock_nblock_nperblock)), // DstDesc
conditional_t<!DoElementwiseBeforeCShuffle,
CElementwiseOperation,
tensor_operation::element_wise::PassThrough>,
@@ -1903,9 +1903,9 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>, // BlockSliceLengths,
CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
Sequence<0, 1, 2, 3>, // typename ThreadClusterArrangeOrder,
Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder,
Sequence<0, 1, 2, 3>, // typename DstDimAccessOrder,
Sequence<0, 2, 1, 3>, // typename ThreadClusterArrangeOrder,
Sequence<0, 2, 1, 3>, // typename SrcDimAccessOrder,
Sequence<0, 2, 1, 3>, // typename DstDimAccessOrder // Sequence<0, 1, 2, 3>
3, // index_t SrcVectorDim,
3, // index_t DstVectorDim,
CDEShuffleBlockTransferScalarPerVectors,
@@ -1915,10 +1915,10 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
uniform_sequence_gen_t<NumDTensor,
false>>, // ThreadTransferSrcResetCoordinateAfterRunFlags
Sequence<false>> // ThreadTransferDstResetCoordinateAfterRunFlags
{c_ds_desc_refs,
idx_c_ds_block_begin,
tie(e_grid_desc_mblock_mperblock_nblock_nperblock),
make_tuple(make_multi_index(block_m_id, 0, block_n_id, 0)),
{c_ds_desc_refs, // SrcDesc
idx_c_ds_block_begin, // SrcSliceBegin
tie(e_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc
make_tuple(make_multi_index(block_m_id, 0, block_n_id, 0)), // DstSliceBegin
lds_to_global_element_op()};
// space filling curve for threadwise C in VGPR
@@ -1939,7 +1939,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
// space filling curve for shuffled blockwise C/D/E
constexpr auto sfc_cde_block =
SpaceFillingCurve<Sequence<1, MPerBlock, 1, NPerBlock>,
Sequence<0, 2, 1, 3>,
Sequence<0, 2, 1, 3>,
Sequence<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
@@ -1962,11 +1962,80 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3
block_sync_lds();
// each block copy its data from LDS to global
cde_block_copy_lds_and_global.Run(
c_ds_desc_refs,
c_ds_buf_refs,
// cde_block_copy_lds_and_global.Run(
// c_ds_desc_refs,
// c_ds_buf_refs,
// tie(e_grid_desc_mblock_mperblock_nblock_nperblock),
// tie(c_grid_buf));
const auto& lds_to_linear_index = [](ck::index_t lds_row, ck::index_t lds_col) {
const auto nPerBlock = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl;
return lds_row * nPerBlock + lds_col;
};
const auto tid = threadIdx.x;
//const auto diagonal_block_id = tid % 32;
//const ck::index_t convK = 4;
const auto group_id = threadIdx.x % 2;
const ck::index_t row_offset = group_id;
const auto lds_col_idx = tid;
float2 buffer;
ck::half_t* buf_ptr = reinterpret_cast<ck::half_t*>(&buffer);
auto lds_ptr = static_cast<CShuffleDataType*>(p_shared);
buf_ptr[0] = lds_ptr[lds_to_linear_index(row_offset, lds_col_idx)];
buf_ptr[1] = lds_ptr[lds_to_linear_index(row_offset + 2, lds_col_idx)];
buf_ptr[2] = lds_ptr[lds_to_linear_index(row_offset + 4, lds_col_idx)];
buf_ptr[4] = lds_ptr[lds_to_linear_index(row_offset + 6, lds_col_idx)];
cde_block_copy_lds_and_global.RunWrite(
tie(e_grid_desc_mblock_mperblock_nblock_nperblock),
tie(c_grid_buf));
tie(c_grid_buf),
buffer);
// LDS Debugging Start
__syncthreads();
if (threadIdx.x == 0 && blockIdx.x == 0)
{
// Print LDS content for MBlock, NBlock, MPerBlock, NPerBlock
const auto mBlock = 1;
const auto nBlock = 1;
const auto mPerBlock = CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl;
const auto nPerBlock = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl;
printf("LDS content for %d x %d block, %d x %d per block:\n",
static_cast<int>(mBlock),
static_cast<int>(nBlock),
static_cast<int>(mPerBlock),
static_cast<int>(nPerBlock));
for (index_t mb = 0; mb < mBlock; ++mb)
{
for (index_t nb = 0; nb < nBlock; ++nb)
{
printf("Block (%d, %d):\n", static_cast<int>(mb), static_cast<int>(nb));
for (index_t npb = 0; npb < nPerBlock; ++npb)
{
for (index_t mpb = 0; mpb < mPerBlock; ++mpb)
{
// Calculate the linear index in LDS
index_t lds_index =
mb * mPerBlock * nBlock * nPerBlock +
nb * nPerBlock * mPerBlock +
mpb * nPerBlock +
npb;
// Assuming CShuffleDataType is int for printing
int value = static_cast<int>(static_cast<CShuffleDataType*>(
p_shared)[lds_index]);
printf("%4d ", value);
}
printf("\n");
}
printf("----\n");
}
printf("========\n");
}
}
// LDS Debugging End
if constexpr(access_id < num_access - 1)
{