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

fix fail

Browse Source
This commit is contained in:
Jin Zhou
2025-09-10 03:46:28 +00:00
parent 91d993909d
commit b2680f3e2f
3 changed files with 299 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_b_preshuffle_v2.hpp
View File

@@ -235,7 +235,15 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v2<BlockGemmPipelineScheduler::I
typename BGridBuffer,
typename BBlockBuffer,
typename BBlockTransferStep,
typename CThreadBuffer>
typename CThreadBuffer,
typename D0GridDesc,
typename D0GridBuffer,
typename D0BlockTransfer,
typename D0ThreadBuffer,
typename D1GridDesc,
typename D1GridBuffer,
typename D1BlockTransfer,
typename D1ThreadBuffer>
__device__ void Run(const AGridDesc& a_grid_desc,
const ABlockDesc& a_block_desc,
ABlockTransfer& a_blockwise_copy,
@@ -248,7 +256,15 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v2<BlockGemmPipelineScheduler::I
BBlockBuffer& b_block_buf,
const BBlockTransferStep& b_block_copy_step,
CThreadBuffer& c_thread_buf,
index_t num_loop) const
index_t num_loop,
const D0GridDesc& d0_grid_desc,
const D0GridBuffer& d0_grid_buf,
D0BlockTransfer& d0_blockwise_copy,
D0ThreadBuffer& d0_thread_buf,
const D1GridDesc& d1_grid_desc,
const D1GridBuffer& d1_grid_buf,
D1BlockTransfer& d1_blockwise_copy,
D1ThreadBuffer& d1_thread_buf) const
{
ignore = b_block_buf;
__builtin_amdgcn_sched_barrier(0);
@@ -455,6 +471,16 @@ struct BlockwiseGemmXdlops_pipeline_bpreshuffle_v2<BlockGemmPipelineScheduler::I
b_block_origin_idx,
b_thread_bufs(local_read_reg));
d0_blockwise_copy.Run(d0_grid_desc,
d0_grid_buf,
make_tuple(I0, I0, I0, I0),
d0_thread_buf);
d1_blockwise_copy.Run(d1_grid_desc,
d1_grid_buf,
make_tuple(I0, I0, I0, I0),
d1_thread_buf);
static_for<0, MRepeat, 1>{}([&](auto m0) {
static_for<0, KRepeat, 1>{}([&](auto k0) {
static_for<0, KGroup, 1>{}([&](auto kg0) {

325
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp
View File

@@ -570,7 +570,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
using DsGridDesc_M_N = remove_cvref_t<decltype(MakeDsGridDescriptor_M_N(0, 0, 0, 0, {}))>;
template <index_t idx,
typename DGridDesc>
typename DGridDesc,
typename SfcCdeBlock>
struct DsBlockTransfer{
static constexpr auto DIdx = Number<idx>{};
static constexpr auto ScalarPerVectors = CDEShuffleBlockTransferScalarPerVectors{}[I1 + DIdx];
@@ -578,13 +579,24 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
using DDataType = remove_cvref_t<decltype(DsDataType{}[DIdx])>;
using DLayout = remove_cvref_t<tuple_element_t<DIdx.value, DsLayout>>;
using Index = MultiIndex<4>;
static constexpr auto ThreadSliceLengths =
Sequence<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>{}
/ CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{};
static constexpr auto sfc_d_block = remove_cvref_t<SfcCdeBlock>{};
static constexpr auto ThreadSliceLengths = []() constexpr {
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
return Sequence<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>{}
/
CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{};
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
return Sequence<1, 1, 1, 1>{};
}
}();
static constexpr auto DsBufferTransferScalarPerVectors =
generate_tuple(
[](auto i) {
@@ -592,18 +604,16 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
},
Number<NumDTensor>{});
static constexpr auto MRepeat = MXdlPerWave / CShuffleMXdlPerWavePerShuffle;
static constexpr auto NRepeat = NXdlPerWave / CShuffleNXdlPerWavePerShuffle;
static constexpr auto MPerCopy = CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl;
static constexpr auto NPerCopy = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl;
__device__ static constexpr auto GetBuffSize()
{
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
return d_buff_desc.GetElementSpaceSize() * Number<CShuffleNXdlPerWavePerShuffle>{};
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
return d_buff_desc.GetElementSpaceSize()
* Number<CShuffleMXdlPerWavePerShuffle>{}
* Number<ThreadSliceLengths[I3]>{};
}
return d_buff_desc.GetElementSpaceSize();
}
__device__ constexpr DsBlockTransfer(const DGridDesc& d_grid_desc,
@@ -613,18 +623,23 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
d_buffer_transfer_.SetSrcSliceOrigin(d_grid_desc,
make_multi_index(src_origin[I0],
0,
src_origin[I2],
get_thread_local_1d_id() * ThreadSliceLengths[I3] % NPerBlock));
make_multi_index(
src_origin[I0],
0,
src_origin[I2],
get_thread_local_1d_id() * ThreadSliceLengths[I3]
% Number<NPerBlock /NRepeat>{}));
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
constexpr auto slice_lenght_n = CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl
/ CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock{}[I3];
d_buffer_transfer_.SetSrcSliceOrigin(d_grid_desc,
make_multi_index(src_origin[I0],
get_thread_local_1d_id() * ThreadSliceLengths[I3] / MPerBlock,
src_origin[I2],
0));
make_multi_index(
src_origin[I0],
get_thread_local_1d_id() / Number<NPerBlock /NRepeat/ slice_lenght_n>{} ,
src_origin[I2],
0));
}
}
@@ -637,72 +652,268 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
auto dst_slice_origni = DstSliceOriginIdx{};
static_for<0, CShuffleNXdlPerWavePerShuffle, 1>{}([&](auto i) {
// static_assert(d_buff_desc.CalculateOffset(
// to_multi_index(make_tuple(I0, I1, I0, I0))) ==0);
static_for<0, NRepeat, 1>{}([&](auto i) {
d_buffer_transfer_.Run(src_desc,
src_buf,
d_buff_desc,
dst_slice_origni,// + makeindex(I0, i, I0, I0),
make_tuple(I0, i, I0, I0),
dst_buf);
#if 0
if constexpr(i < CShuffleNXdlPerWavePerShuffle - 1)
#if 1
if constexpr(i < NRepeat - 1)
{
d_buffer_transfer_.MoveSrcSliceWindow(
src_desc,
make_tuple(I0, I0, I0, ));
Index(I0, I0, I0, Number<NPerBlock / NRepeat>{}));
}
#endif
});
#if 0
const size_t M = 1;
const size_t N = BlockSize * NPerThread;
__shared__ float lds_d[M* N];
if (blockIdx.x == 0)
{
block_sync_lds();
for(size_t i = 0; i < M*N; i++)
{
lds_d[i] = 0;
}
block_sync_lds();
static_for<0, NPerThread, 1>{}([&](auto i) {
size_t lds_idx = get_thread_local_1d_id()*NPerThread + i.value;
lds_d[lds_idx] = dst_buf[i];
});
block_sync_lds();
if (threadIdx.x == 0)
{
for (size_t i = 0; i < M; i++)
{
printf("[lds %zu] ", i);
for (size_t j = 0; j < N; j++)
{
/* code */
printf(" %0.2f ", lds_d[i*M + j]);
}
printf(" \n ");
}
}
}
#endif
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
#if 1
auto dst_slice_origni = DstSliceOriginIdx{};
static_for<0, CShuffleMXdlPerWavePerShuffle, 1>{}([&](auto i) {
static_for<0, MRepeat, 1>{}([&](auto i) {
d_buffer_transfer_.Run(src_desc,
src_buf,
d_buff_desc,
make_tuple(I0, i, I0, I0),
dst_buf);
#if 1
if constexpr(i < CShuffleMXdlPerWavePerShuffle - 1)
if constexpr(i < MRepeat - 1)
{
d_buffer_transfer_.MoveSrcSliceWindow(
src_desc,
Index(I0,
Number<CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl>{},
Number<MPerBlock / MRepeat>{},
I0,
I0));
}
#endif
});
#else
constexpr index_t num_access = sfc_d_block.GetNumOfAccess();
static_for<0, num_access, 1>{}([&](auto access_id) {
constexpr auto idx_md = sfc_d_block.GetIndex(access_id);
d_buffer_transfer_.Run(src_desc,
src_buf,
d_buff_desc,
make_tuple(I0,
// Number<idx_md[I2] / 16>{},
Number<2>{},
I0,
I0),
dst_buf);
if constexpr(access_id < num_access - 1)
{
#if 0
constexpr auto d_block_step =
sfc_d_block.GetForwardStep(access_id);
d_buffer_transfer_.MoveSrcSliceWindow(
src_desc,
d_block_step);
#endif
d_buffer_transfer_.MoveSrcSliceWindow(
src_desc,
Index(I0,
//Number<CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl>{},
Number<MPerBlock / MRepeat>{},
I0,
I0));
}
});
#endif
#if 0
const size_t M = MPerThread;
const size_t N = 1;
__shared__ float lds_d[M* N];
if (blockIdx.x == 0)
{
block_sync_lds();
for(size_t i = 0; i < M*N; i++)
{
lds_d[i] = 0;
}
block_sync_lds();
static_for<0, MPerThread, 1>{}([&](auto i) {
size_t lds_idx = get_thread_local_1d_id()*MPerThread + i.value;
lds_d[lds_idx] = dst_buf[i];
});
block_sync_lds();
if (threadIdx.x == 0)
{
for (size_t i = 0; i < M; i++)
{
printf("[lds %zu] ", i);
for (size_t j = 0; j < N; j++)
{
/* code */
printf(" %0.2f ", lds_d[i*M + j]);
}
printf(" \n ");
}
}
}
#endif
}
}
template <index_t access_id>
__device__ static constexpr auto GetOrignal(const Number<access_id>&)
{
constexpr auto idx_md = sfc_d_block.GetIndex(Number<access_id>{});
//static_assert(idx_md[I3] == 128);
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
return make_tuple(I0,
Number<idx_md[I3] / NPerCopy>{},
I0,
I0);
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
return make_tuple(I0,
Number<idx_md[I1] / MPerCopy>{},
I0,
I0);
}
}
// for write
__device__ static constexpr auto MakeDBufferDescriptor()
{
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
return make_naive_tensor_descriptor_packed(
make_tuple(I1,
Number<CShuffleNXdlPerWavePerShuffle>{},
Number<NRepeat>{},
I1,
Number<ThreadSliceLengths[I3]>{}
));
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
return make_naive_tensor_descriptor_packed(
make_tuple(I1,
Number<CShuffleMXdlPerWavePerShuffle>{},
Number<MRepeat>{},
I1,
I1));
}
}
// for read
__device__ static constexpr auto MakeDBufferDescriptor_mblock_mperblock_nblock_nperblock()
{
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
#if 1
return make_naive_tensor_descriptor_packed(
make_tuple(I1,
Number<NRepeat>{},
I1,
Number<ThreadSliceLengths[I3]>{}
));
#else
return transform_tensor_descriptor(
make_naive_tensor_descriptor_packed(
#if 0
make_tuple(I1,
Number<MPerBlock>{},
I1,
Number<NPerBlock>{})
#endif
make_tuple(I1,
I5, Number<16>{},
I1,
I2, Number<128>{})
),
make_tuple(
make_pass_through_transform(Number<1>{}),
make_merge_transform(
make_tuple(I5, Number<16>{})),
make_pass_through_transform(Number<1>{}),
make_merge_transform(make_tuple(I2, Number<128>{}))
),
make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}, Sequence<4, 5>{}),
//make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{})
);
#endif
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
return make_naive_tensor_descriptor_packed(
make_tuple(I1,
Number<MRepeat>{},
I1,
I1));
}
}
static constexpr auto d_buff_desc = MakeDBufferDescriptor();
static constexpr auto d_buff_desc_mblock_mperblock_nblock_nperblock =
MakeDBufferDescriptor_mblock_mperblock_nblock_nperblock();
static_assert(d_buff_desc_mblock_mperblock_nblock_nperblock.IsKnownAtCompileTime());
private:
__device__ static constexpr auto GetVectorDim()
{
if constexpr(is_same<tensor_layout::gemm::RowMajor, DLayout>::value)
{
return I3;
}
else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, DLayout>::value)
{
return I1;
}
}
using DBufferTransfer = ThreadwiseTensorSliceTransfer_v2<
DDataType, // SrcData
DDataType, // DstData
@@ -710,7 +921,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
remove_cvref_t<decltype(d_buff_desc)>, //DstDesc
decltype(ThreadSliceLengths), // SliceLengths
Sequence<0, 1, 2, 3>, // DimAccessOrder,
3, // SrcVectorDim
GetVectorDim(), // SrcVectorDim
DsBufferTransferScalarPerVectors[DIdx], // SrcScalarPerVector
1, // SrcScalarStrideInVector
false>;
@@ -1857,11 +2068,20 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
auto blockwise_gemm_pipeline = BlockwiseGemmPipe{};
auto c_thread_buf = blockwise_gemm_pipeline.GetCThreadBuffer();
using SfcCdeBlock =
SpaceFillingCurve<Sequence<1, MPerBlock, 1, NPerBlock>,
Sequence<0, 2, 1, 3>,
Sequence<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>;
// Ds
using D0BlockTransfer = DsBlockTransfer<0,
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0])>;
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0]),
SfcCdeBlock>;
using D1BlockTransfer = DsBlockTransfer<1,
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1])>;
decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock[I1]),
SfcCdeBlock>;
D0BlockTransfer d0_block_copy_to_vgpr(ds_grid_desc_mblock_mperblock_nblock_nperblock[I0],
make_multi_index(block_m_id, 0, block_n_id, 0));
@@ -2053,8 +2273,8 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
remove_cvref_t<decltype(DsDataType{}[I0])>, // Src2Data
EDataType, // DstData
decltype(c_shuffle_block_desc_mblock_mperblock_nblock_nperblock), // Src0Desc
remove_cvref_t<decltype(D0BlockTransfer::d_buff_desc)>, // Src1Desc
remove_cvref_t<decltype(D1BlockTransfer::d_buff_desc)>, // Src2Desc
remove_cvref_t<decltype(D0BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock)>, // Src1Desc
remove_cvref_t<decltype(D1BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock)>, // Src2Desc
decltype(e_grid_desc_mblock_mperblock_nblock_nperblock), // DstDesc
Sequence<0, 1, 2, 3>, // typename SrcDimAccessOrder,
3, // index_t SrcVectorDim,
@@ -2083,15 +2303,10 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
constexpr index_t num_access = sfc_c_vgpr.GetNumOfAccess();
// 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<1,
CShuffleMXdlPerWavePerShuffle * MWave * MPerXdl,
1,
CShuffleNXdlPerWavePerShuffle * NWave * NPerXdl>>{};
constexpr SfcCdeBlock sfc_cde_block;
static_assert(num_access == sfc_cde_block.GetNumOfAccess(), "wrong!");
static_for<0, num_access, 1>{}([&](auto access_id) {
// make sure it's safe to write to LDS
block_sync_lds();
@@ -2105,16 +2320,16 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle
// make sure it's safe to read from LDS
block_sync_lds();
constexpr auto idx_md = SfcCdeBlock::GetIndex(I1);
// each block copy its data from LDS to global
cde_block_copy_lds_and_global.Run(
c_shuffle_block_desc_mblock_mperblock_nblock_nperblock,
c_shuffle_block_buf,
D0BlockTransfer::d_buff_desc,
make_tuple(I0, I0, I0, I0),
D0BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock,
D0BlockTransfer::GetOrignal(access_id),
d0_thread_buf,
D1BlockTransfer::d_buff_desc,
make_tuple(I0, access_id, I0, I0),
D1BlockTransfer::d_buff_desc_mblock_mperblock_nblock_nperblock,
D1BlockTransfer::GetOrignal(access_id),
d1_thread_buf,
e_grid_desc_mblock_mperblock_nblock_nperblock,
c_grid_buf);

6
include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer_v6r4.hpp
View File

@@ -151,11 +151,7 @@ struct ThreadwiseTensorSliceTransfer_v6r4
static_for<0, ScalarPerVector, 1>{}([&](auto i) {
element_op_(dst_vector_container.template AsType<DstData>()(i),
src0_vector_container.template AsType<Src0Data>()[i],
#if 1
src1_buf[Number<i>{}],
#else
type_convert<Src1Data>(src0_coord_.GetOffset()+i),
#endif
src1_buf[Number<src1_offset+i>{}],
src2_buf[Number<src2_offset>{}]);
});