ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-17 09:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

tmp

Browse Source
This commit is contained in:
Bartlomiej Kocot
2025-09-05 18:24:59 -04:00
parent e688dcfcda
commit 98dd392f53
3 changed files with 259 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
include/ck/tensor_operation/gpu/block/blockwise_gemm_xdlops_skip_b_lds.hpp
View File

@@ -13,8 +13,8 @@ namespace ck {
template <index_t BlockSize,
typename FloatAB,
typename FloatAcc,
typename AK0MK1BlockDesc,
typename BK0K0BN0N1N2N3K1BlockDesc,
typename AK0K0BM0M1M2M3K1BlockDesc,
typename BK0NK1BlockDesc,
index_t MPerBlock,
index_t NPerBlock,
index_t K0PerBlock,
@@ -34,8 +34,8 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1
static constexpr index_t KPerBlock = K0PerBlock * KPack;
static constexpr index_t A_K0 = AK0MK1BlockDesc{}.GetLength(I0);
static constexpr index_t A_K1 = AK0MK1BlockDesc{}.GetLength(I2);
static constexpr index_t B_K0 = BK0NK1BlockDesc{}.GetLength(I0);
static constexpr index_t B_K1 = BK0NK1BlockDesc{}.GetLength(I2);
static constexpr auto xdlops_gemm = XdlopsGemm<FloatAB, MPerXDL, NPerXDL, KPack>{};
@@ -119,8 +119,8 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1
__host__ __device__ BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1()
{
static_assert(AK0MK1BlockDesc::IsKnownAtCompileTime() &&
BK0K0BN0N1N2N3K1BlockDesc::IsKnownAtCompileTime(),
static_assert(BK0NK1BlockDesc::IsKnownAtCompileTime() &&
AK0K0BM0M1M2M3K1BlockDesc::IsKnownAtCompileTime(),
"wrong! Desc should be known at compile-time");
static_assert(BlockSize == MWaves * NWaves * WaveSize,
@@ -221,58 +221,58 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1
c_grid_desc_g_m0_n0_m1_n1_m2_n2);
}
__host__ __device__ static constexpr auto MakeABlockDescriptor_M0_M1_M2_K()
__host__ __device__ static constexpr auto MakeBBlockDescriptor_N0_N1_N2_K()
{
return transform_tensor_descriptor(
AK0MK1BlockDesc{},
BK0NK1BlockDesc{},
make_tuple(
make_merge_transform_v3_division_mod(make_tuple(Number<A_K0>{}, Number<A_K1>{})),
make_merge_transform_v3_division_mod(make_tuple(Number<B_K0>{}, Number<B_K1>{})),
make_unmerge_transform(
make_tuple(Number<MRepeat>{}, Number<MWaves>{}, Number<MPerXDL>{}))),
make_tuple(Number<NRepeat>{}, Number<NWaves>{}, Number<NPerXDL>{}))),
make_tuple(Sequence<0, 2>{}, Sequence<1>{}),
make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{}));
}
__device__ void MoveABlockSliceWindow()
__device__ void MoveBBlockSliceWindow()
{
a_thread_copy_.MoveSrcSliceWindow(a_block_desc_m0_m1_m2_k,
b_thread_copy_.MoveSrcSliceWindow(b_block_desc_n0_n1_n2_k,
make_multi_index(0, 0, 0, K0PerBlock * KPack));
}
__device__ void ResetABlockStartWindow()
__device__ void ResetBBlockStartWindow()
{
a_thread_copy_.SetSrcCoord(CalculateAThreadOriginDataIndex());
b_thread_copy_.SetSrcCoord(CalculateBThreadOriginDataIndex());
}
static constexpr auto a_block_desc_m0_m1_m2_k = MakeABlockDescriptor_M0_M1_M2_K();
static constexpr auto b_block_desc_n0_n1_n2_k = MakeBBlockDescriptor_N0_N1_N2_K();
template <typename ABlockBuffer, typename BBlockBuffer, typename CThreadBuffer>
__device__ void Run(const ABlockBuffer& a_block_buf,
const BBlockBuffer& b_thread_buf,
__device__ void Run(const ABlockBuffer& a_thread_buf,
const BBlockBuffer& b_block_buf,
CThreadBuffer& c_thread_buf) const
{
auto a_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, FloatAB>(
a_thread_desc_.GetElementSpaceSize());
auto b_thread_buf = make_static_buffer<AddressSpaceEnum::Vgpr, FloatAB>(
b_thread_desc_.GetElementSpaceSize());
static_for<0, MRepeat, 1>{}([&](auto m0) {
static_for<0, NRepeat, 1>{}([&](auto n0) {
// read A
a_thread_copy_.Run(a_block_desc_m0_m1_m2_k,
make_tuple(m0, I0, I0, I0),
a_block_buf,
a_thread_desc_,
b_thread_copy_.Run(b_block_desc_n0_n1_n2_k,
make_tuple(n0, I0, I0, I0),
b_block_buf,
b_thread_desc_,
make_tuple(I0, I0, I0, I0),
a_thread_buf);
b_thread_buf);
static_for<0, NRepeat, 1>{}([&](auto n0) {
static_for<0, MRepeat, 1>{}([&](auto m0) {
// read B
static_for<0, KPerThread, KPack>{}([&](auto k) {
vector_type<FloatAB, KPack> a_thread_vec;
vector_type<FloatAB, KPack> b_thread_vec;
constexpr index_t k0 = k / KPack;
static_for<0, KPack, 1>{}([&](auto i) {
a_thread_vec.template AsType<FloatAB>()(i) = a_thread_buf
[Number<a_thread_desc_.CalculateOffset(make_tuple(0, 0, 0, k + i))>{}];
b_thread_vec.template AsType<FloatAB>()(i) = b_thread_buf
[Number<b_thread_desc_.CalculateOffset(make_tuple(k0, n0, i))>{}];
[Number<b_thread_desc_.CalculateOffset(make_tuple(0, 0, 0, k + i))>{}];
a_thread_vec.template AsType<FloatAB>()(i) = a_thread_buf
[Number<a_thread_desc_.CalculateOffset(make_tuple(k0, m0, i))>{}];
});
using mfma_input_type =
@@ -291,30 +291,30 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1
private:
// A[M0, M1, M2, KPerThread]
static constexpr auto a_thread_desc_ =
static constexpr auto b_thread_desc_ =
make_naive_tensor_descriptor_packed(make_tuple(I1, I1, I1, Number<KPerThread>{}));
// B[N0, N1, N2, KPerThread]
static constexpr auto b_thread_desc_ =
static constexpr auto a_thread_desc_ =
make_naive_tensor_descriptor_packed(make_tuple(Number<K0PerThread>{}, // KPerThread
Number<NRepeat>{}, // repeat
Number<MRepeat>{}, // repeat
Number<KPack>{}));
// C[M, N, NumRegXdlops]
static constexpr auto c_thread_desc_ = make_naive_tensor_descriptor_packed(
make_tuple(Number<MRepeat>{}, Number<NRepeat>{}, xdlops_gemm.GetRegSizePerXdlops()));
using AThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
using BThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
FloatAB,
decltype(a_block_desc_m0_m1_m2_k),
decltype(a_thread_desc_),
decltype(b_block_desc_n0_n1_n2_k),
decltype(b_thread_desc_),
Sequence<1, 1, 1, KPerThread>,
Sequence<0, 1, 2, 3>,
3,
A_K1,
A_K1>;
B_K1,
B_K1>;
AThreadCopy a_thread_copy_{CalculateAThreadOriginDataIndex()};
BThreadCopy b_thread_copy_{CalculateBThreadOriginDataIndex()};
};
} // namespace ck

139
include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp
View File

@@ -79,7 +79,7 @@ template <typename GridwiseGemm,
bool HasMainKBlockLoopInAllGemm,
bool NoMainKBlockLoopInAllGemm,
bool CTranspose,
bool SkipBLds>
bool SkipALds>
__global__ void
#if CK_USE_LAUNCH_BOUNDS
__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
@@ -151,7 +151,7 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
}
// If constexpr to be compatible with skip LDS gridwise gemm
if constexpr(SkipBLds)
if constexpr(SkipALds)
{
if constexpr(HasMainKBlockLoopInAllGemm || NoMainKBlockLoopInAllGemm)
{
@@ -348,7 +348,8 @@ template <index_t NDimSpatial,
typename BComputeType = AComputeType,
index_t MaxTransposeTransferInScalarPerVector = 1,
index_t MaxTransposeTransferOutScalarPerVector = 1,
bool SkipBLds = false>
bool SkipALds = false,
index_t ABlockBufferSize = 1>
struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
: public DeviceGroupedConvBwdDataMultipleD<NDimSpatial,
ALayout, // output image
@@ -369,7 +370,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
static_assert(NDimSpatial == 2 || NDimSpatial == 3,
"wrong! only implemented for 2D and 3D now");
static_assert(!SkipBLds || AK1 == BK1);
static_assert(!SkipALds || AK1 == BK1);
// MaxGroupedGemmGroupsNum is used to specify number of gemm args in compile time. With this
// implementation we can avoid copy data to workspace before kernel launch since number of
@@ -387,7 +388,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
static constexpr GemmSpecialization GemmSpec = GemmSpecialization::MNKPadding;
static constexpr bool IsSplitKSupported =
(CDEBlockTransferScalarPerVector_NPerBlock % 2 == 0 || sizeof(EDataType) % 4 == 0) &&
std::is_same_v<remove_cvref_t<CDEElementwiseOp>, element_wise::PassThrough> && !SkipBLds;
std::is_same_v<remove_cvref_t<CDEElementwiseOp>, element_wise::PassThrough> && !SkipALds;
// TODO: Add support for different A and B data types.
using ABDataType = ADataType;
@@ -411,7 +412,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
static constexpr bool CTranspose = (NeedTransposeKernel == false) &&
(is_same_v<ELayout, tensor_layout::convolution::NGCHW> ||
is_same_v<ELayout, tensor_layout::convolution::NGCDHW>) &&
!SkipBLds;
!SkipALds;
using ALayoutAfterTranspose = std::conditional_t<
is_NGCHW_NGKHW<ELayout, BLayout, ALayout>() && NeedTransposeKernel,
@@ -530,26 +531,53 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, \
CDEBlockTransferScalarPerVector_NPerBlock, LoopSched, PipelineVersion::v1, BComputeType
static constexpr index_t BBlockBufferSize = 1;
// static constexpr index_t BBlockBufferSize = 1;
// Force to 1, due to KN layout for GKYXC
static constexpr index_t BScalarPerVectorSkipLds = 1;
// static constexpr index_t BScalarPerVectorSkipLds = 1;
#define GridwiseGemmMultiDSkipBLdsTemplateParams \
static constexpr index_t SharedMemoryABlockBufferRequiredSize()
{
constexpr auto max_lds_align = AK1;
// A matrix in LDS memory, dst of blockwise copy
constexpr auto b_block_desc_k0_n_k1 = [&]() {
if constexpr(BBlockLdsExtraN)
{
return make_naive_tensor_descriptor(
make_tuple(Number<(KPerBlock / AK1) * ABlockBufferSize>{}, Number<NPerBlock>{}, AK1),
make_tuple(Number<NPerBlock + 1>{} * AK1, AK1, I1));
}
else
{
return make_naive_tensor_descriptor_packed(
make_tuple(Number<(KPerBlock / AK1) * ABlockBufferSize>{}, Number<NPerBlock>{}, AK1));
}
}();
constexpr auto b_block_space_size_aligned =
math::integer_least_multiple(b_block_desc_k0_n_k1.GetElementSpaceSize(), max_lds_align);
return b_block_space_size_aligned * sizeof(ABDataType);
}
static constexpr index_t ABlockBufferSizeAligned = std::max(1, std::min(65536 / SharedMemoryABlockBufferRequiredSize(), ABlockBufferSize));
#define GridwiseGemmMultiDSkipALdsTemplateParams \
BlockSize, ABDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, \
InMemoryDataOperationEnum::Set, element_wise::PassThrough, element_wise::PassThrough, \
element_wise::PassThrough, MPerBlock, NPerBlock, KPerBlock / AK1, MPerXDL, NPerXDL, AK1, \
MXdlPerWave, NXdlPerWave, ABlockTransferThreadClusterLengths_AK0_M_AK1, \
ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, \
ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, \
ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BScalarPerVectorSkipLds, \
false, BBlockBufferSize, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, \
MXdlPerWave, NXdlPerWave, \
ABlockTransferSrcScalarPerVector, false, ABlockBufferSizeAligned, \
BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, \
CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, \
CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, \
CDEBlockTransferScalarPerVector_NPerBlock
using GridwiseGemm =
std::conditional_t<SkipBLds,
std::conditional_t<SkipALds,
GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle<
GridwiseGemmMultiDSkipBLdsTemplateParams>,
GridwiseGemmMultiDSkipALdsTemplateParams>,
GridwiseGemmMultipleD_xdl_cshuffle<GridwiseGemmMultiDTemplateParams>>;
using GridwiseGemmCTranspose = std::conditional_t<
CTranspose,
@@ -1005,7 +1033,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
const auto GemmK = a_grid_desc_m_k.GetLength(I1);
const auto GemmK0 = a_grid_desc_ak0_m_ak1.GetLength(I0);
bool HasMainKBlockLoop = true;
if constexpr(SkipBLds)
if constexpr(SkipALds)
{
HasMainKBlockLoop =
GridwiseGemmCTranspose::CalculateHasMainK0BlockLoop(GemmK0);
@@ -1091,6 +1119,33 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
compute_ptr_offset_of_workspace_n_.BatchStrideE_ =
e_g_n_c_wis_strides[1] * conv_N_per_block_;
}
has_loop_in_all_gemm_.resize(gemm_kernel_args_.size());
no_loop_in_all_gemm_.resize(gemm_kernel_args_.size());
for(std::size_t gemm_set_id = 0; gemm_set_id < gemm_kernel_args_.size();
gemm_set_id++)
{
const std::array<GemmArgs, MaxGroupedGemmGroupsNum>& gemm_kernel_args =
gemm_kernel_args_[gemm_set_id];
const index_t gemms_count_for_set =
gemm_set_id == gemm_kernel_args_.size() - 1
? gemms_count_ - MaxGroupedGemmGroupsNum * gemm_set_id
: MaxGroupedGemmGroupsNum;
bool has_loop_in_all_gemm = true;
bool no_loop_in_all_gemm = true;
for(auto i = 0; i < gemms_count_for_set; i++)
{
has_loop_in_all_gemm &= gemm_kernel_args[i].HasMainKBlockLoop_;
no_loop_in_all_gemm &= !gemm_kernel_args[i].HasMainKBlockLoop_;
}
has_loop_in_all_gemm_[gemm_set_id] =has_loop_in_all_gemm;
no_loop_in_all_gemm_[gemm_set_id] =no_loop_in_all_gemm;
}
gdy_ = num_group_;
gdz_ = num_workgroups_per_Conv_N_ * k_batch_;
}
std::size_t GetWorkspaceATensorSizeBytes() const
@@ -1209,9 +1264,12 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
std::vector<index_t> gemms_grid_size_;
index_t gemms_count_ = 0;
std::vector<std::array<GemmArgs, MaxGroupedGemmGroupsNum>> gemm_kernel_args_;
std::vector<bool> has_loop_in_all_gemm_;
std::vector<bool> no_loop_in_all_gemm_;
bool bwd_needs_zero_out;
long_index_t e_space_size_bytes;
index_t gdy_, gdz_;
};
// Invoker
@@ -1224,9 +1282,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
{
float ave_time = 0;
const index_t gdy = arg.num_group_;
const index_t gdz = arg.num_workgroups_per_Conv_N_ * arg.k_batch_;
const ADataType* p_a_grid = arg.p_a_grid_;
const BDataType* p_b_grid = arg.p_b_grid_;
EDataType* p_e_grid = arg.p_e_grid_;
@@ -1257,8 +1312,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
gemm_set_id == arg.gemm_kernel_args_.size() - 1
? arg.gemms_count_ - MaxGroupedGemmGroupsNum * gemm_set_id
: MaxGroupedGemmGroupsNum;
const std::array<GemmArgs, MaxGroupedGemmGroupsNum>& gemm_kernel_args =
arg.gemm_kernel_args_[gemm_set_id];
const auto clear_workspace = [&]() {
if(arg.bwd_needs_zero_out && gemm_set_id == 0)
@@ -1268,14 +1321,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
}
};
bool has_loop_in_all_gemm = true;
bool no_loop_in_all_gemm = true;
for(auto i = 0; i < gemms_count_for_set; i++)
{
has_loop_in_all_gemm &= gemm_kernel_args[i].HasMainKBlockLoop_;
no_loop_in_all_gemm &= !gemm_kernel_args[i].HasMainKBlockLoop_;
}
auto launch_kernel = [&](auto has_main_k_block_loop, auto no_main_k_block_loop) {
constexpr bool has_main_loop = has_main_k_block_loop.value;
constexpr bool no_main_loop = no_main_k_block_loop.value;
@@ -1297,20 +1342,20 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
has_main_loop,
no_main_loop,
CTranspose,
SkipBLds>;
SkipALds>;
return launch_and_time_kernel_with_preprocess(
stream_config,
clear_workspace,
kernel,
dim3(gdx, gdy, gdz),
dim3(gdx, arg.gdy_, arg.gdz_),
dim3(BlockSize),
0,
p_b_grid,
p_a_grid,
arg.p_ds_grid_,
p_e_grid,
gemm_kernel_args,
arg.gemm_kernel_args_[gemm_set_id],
gemms_count_for_set,
arg.b_element_op_,
arg.a_element_op_,
@@ -1337,20 +1382,20 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
has_main_loop,
no_main_loop,
CTranspose,
SkipBLds>;
SkipALds>;
return launch_and_time_kernel_with_preprocess(
stream_config,
clear_workspace,
kernel,
dim3(gdx, gdy, gdz),
dim3(gdx, arg.gdy_, arg.gdz_),
dim3(BlockSize),
0,
p_a_grid,
p_b_grid,
arg.p_ds_grid_,
p_e_grid,
gemm_kernel_args,
arg.gemm_kernel_args_[gemm_set_id],
gemms_count_for_set,
arg.a_element_op_,
arg.b_element_op_,
@@ -1360,12 +1405,12 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
arg.k_batch_);
}
};
if(has_loop_in_all_gemm)
if(arg.has_loop_in_all_gemm_[gemm_set_id])
{
ave_time += launch_kernel(integral_constant<bool, true>{},
integral_constant<bool, false>{});
}
else if(no_loop_in_all_gemm)
else if(arg.no_loop_in_all_gemm_[gemm_set_id])
{
ave_time += launch_kernel(integral_constant<bool, false>{},
integral_constant<bool, true>{});
@@ -1384,10 +1429,10 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
{
float ave_time = 0;
if(stream_config.log_level_ > 0)
{
arg.Print();
}
// if(stream_config.log_level_ > 0)
// {
// arg.Print();
// }
// Transpose from NGKHW to NHWGK
if constexpr(NeedTransposeKernel)
@@ -1699,7 +1744,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
// Gridwise GEMM size
for(std::size_t i = 0; i < arg.a_grid_desc_m_k_container_.size(); i++)
{
if constexpr(SkipBLds)
if constexpr(SkipALds)
{
if(!GridwiseGemmCTranspose::CheckValidity(
arg.gemm_kernel_args_[i / MaxGroupedGemmGroupsNum]
@@ -1884,9 +1929,9 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
auto str = std::stringstream();
str << "DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1";
if constexpr(SkipBLds)
if constexpr(SkipALds)
{
str << "_SkipBLds";
str << "_SkipALds";
}
// clang-format off
@@ -1914,6 +1959,10 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
<< "TransposeTransferOutScalarPerVectorAligned: " << TransposeTransferOutScalarPerVectorAligned;
}
if constexpr(SkipALds)
{
str << ", " << ABlockBufferSize;
}
str << ">";

251
include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_skip_b_lds_multiple_d_cshuffle.hpp
View File

@@ -35,17 +35,20 @@ template <index_t BlockSize,
index_t K1Value,
index_t MXdlPerWave,
index_t NXdlPerWave,
typename ABlockTransferThreadClusterLengths_K0_M_K1,
typename ABlockTransferThreadClusterArrangeOrder,
typename ABlockTransferSrcAccessOrder,
index_t ABlockTransferSrcVectorDim,
index_t ABlockTransferSrcScalarPerVector,
index_t ABlockTransferDstScalarPerVector_K1,
bool AThreadTransferSrcResetCoordinateAfterRun,
bool ABlockLdsExtraM,
index_t ABlockBufferSize,
typename BBlockTransferThreadClusterLengths_K0_N_K1,
typename BBlockTransferThreadClusterArrangeOrder,
typename BBlockTransferSrcAccessOrder,
index_t BBlockTransferSrcVectorDim,
index_t BBlockTransferSrcScalarPerVector,
index_t BBlockTransferDstScalarPerVector_K1,
bool BThreadTransferSrcResetCoordinateAfterRun,
index_t BBlockBufferSize,
bool BBlockLdsExtraN,
index_t CShuffleMXdlPerWavePerShuffle,
index_t CShuffleNXdlPerWavePerShuffle,
typename CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
@@ -87,27 +90,27 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
using DsGridPointer = decltype(MakeDsGridPointer());
__host__ __device__ static constexpr auto GetABlockDescriptor_K0PerBlock_MPerBlock_K1()
__host__ __device__ static constexpr auto GetBBlockDescriptor_K0PerBlock_NPerBlock_K1()
{
constexpr auto max_lds_align = K1;
// A matrix in LDS memory, dst of blockwise copy
constexpr auto a_block_desc_k0_m_k1 = [&]() {
if constexpr(ABlockLdsExtraM)
constexpr auto b_block_desc_k0_n_k1 = [&]() {
if constexpr(BBlockLdsExtraN)
{
return make_naive_tensor_descriptor(
make_tuple(Number<K0PerBlock * BBlockBufferSize>{}, Number<MPerBlock>{}, K1),
make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
make_tuple(Number<K0PerBlock * ABlockBufferSize>{}, Number<NPerBlock>{}, K1),
make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
}
else
{
return make_naive_tensor_descriptor_aligned(
make_tuple(Number<K0PerBlock * BBlockBufferSize>{}, Number<MPerBlock>{}, K1),
make_tuple(Number<K0PerBlock * ABlockBufferSize>{}, Number<NPerBlock>{}, K1),
max_lds_align);
}
}();
return a_block_desc_k0_m_k1;
return b_block_desc_k0_n_k1;
}
template <typename EGridDesc_M_N>
@@ -155,7 +158,7 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
// 2-stage prefetch currently only support even number of K0 loop
// TODO: add support for odd number of K0 loop
if(!((K0 / K0PerBlock) % BBlockBufferSize == 0))
if(!((K0 / K0PerBlock) % ABlockBufferSize == 0))
{
return false;
}
@@ -167,28 +170,28 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
// TODO move this function into GEMM-pipeline class
__host__ __device__ static constexpr bool CalculateHasMainK0BlockLoop(index_t K0)
{
const bool has_main_k0_block_loop = (K0 / (BBlockBufferSize * K0PerBlock)) > 1;
const bool has_main_k0_block_loop = (K0 / (ABlockBufferSize * K0PerBlock)) > 1;
return has_main_k0_block_loop;
}
template <typename BGridDesc_K0_N_K1>
template <typename AGridDesc_K0_M_K1>
__host__ __device__ static constexpr auto
MakeBGridDescriptor_K0_K1_K2_N0_N1_N2_N3_K3(const BGridDesc_K0_N_K1& b_grid_desc_k0_n_k1)
MakeAGridDescriptor_K0_K1_K2_M0_M1_M2_M3_K3(const AGridDesc_K0_M_K1& a_grid_desc_k0_m_k1)
{
const auto K0 = b_grid_desc_k0_n_k1.GetLength(I0);
const auto N = b_grid_desc_k0_n_k1.GetLength(I1);
const auto K0 = a_grid_desc_k0_m_k1.GetLength(I0);
const auto M = a_grid_desc_k0_m_k1.GetLength(I1);
const auto b_griddesc_k0_nblockid_nrepeat_waves_nperxdlops_k1 = transform_tensor_descriptor(
b_grid_desc_k0_n_k1,
const auto a_griddesc_k0_mblockid_mrepeat_waves_mperxdlops_k1 = transform_tensor_descriptor(
a_grid_desc_k0_m_k1,
make_tuple(make_unmerge_transform(
make_tuple(K0 / K0PerBlock, xdlops_gemm.K0PerXdlops, K0PerThread)),
make_unmerge_transform(make_tuple(
N / (NXdlPerWave * NWaves * NPerXdl), NXdlPerWave, NWaves, NPerXdl)),
M / (MXdlPerWave * MWaves * MPerXdl), MXdlPerWave, MWaves, MPerXdl)),
make_pass_through_transform(K1)),
make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
make_tuple(Sequence<0, 1, 2>{}, Sequence<3, 4, 5, 6>{}, Sequence<7>{}));
return b_griddesc_k0_nblockid_nrepeat_waves_nperxdlops_k1;
return a_griddesc_k0_mblockid_mrepeat_waves_mperxdlops_k1;
}
template <typename EGridDesc_M_N>
@@ -235,14 +238,14 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
return threadid_to_wave_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
}
__device__ static auto GetWaveKNIdx(const index_t thread_id)
__device__ static auto GetWaveKMIdx(const index_t thread_id)
{
constexpr auto wave_threadid_to_nk_idx_adaptor = make_single_stage_tensor_adaptor(
make_tuple(make_merge_transform(make_tuple(xdlops_gemm.K0PerXdlops, NPerXdl))),
constexpr auto wave_threadid_to_mk_idx_adaptor = make_single_stage_tensor_adaptor(
make_tuple(make_merge_transform(make_tuple(xdlops_gemm.K0PerXdlops, MPerXdl))),
make_tuple(Sequence<0, 1>{}),
make_tuple(Sequence<0>{}));
return wave_threadid_to_nk_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
return wave_threadid_to_mk_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
}
__device__ static constexpr auto GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock()
@@ -263,12 +266,12 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
__host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
{
// LDS allocation for A and B: be careful of alignment
constexpr auto a_block_desc_k0_m_k1 = GetABlockDescriptor_K0PerBlock_MPerBlock_K1();
constexpr auto b_block_desc_k0_n_k1 = GetBBlockDescriptor_K0PerBlock_NPerBlock_K1();
constexpr auto max_lds_align = K1;
constexpr auto a_block_space_size_aligned =
math::integer_least_multiple(a_block_desc_k0_m_k1.GetElementSpaceSize(), max_lds_align);
constexpr auto b_block_space_size_aligned =
math::integer_least_multiple(b_block_desc_k0_n_k1.GetElementSpaceSize(), max_lds_align);
constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock =
GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock();
@@ -276,7 +279,7 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
constexpr auto c_block_size =
c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize();
return math::max((a_block_space_size_aligned) * sizeof(ABDataType),
return math::max((b_block_space_size_aligned) * sizeof(ABDataType),
c_block_size * sizeof(CShuffleDataType));
}
@@ -303,13 +306,13 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
const Block2CTileMap& block_2_ctile_map)
{
constexpr index_t NumDTensor = DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock::Size();
const auto b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3 =
MakeBGridDescriptor_K0_K1_K2_N0_N1_N2_N3_K3(b_grid_desc_k0_n_k1);
const auto a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3 =
MakeAGridDescriptor_K0_K1_K2_M0_M1_M2_M3_K3(a_grid_desc_k0_m_k1);
const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
p_a_grid, a_grid_desc_k0_m_k1.GetElementSpaceSize());
p_a_grid, a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3.GetElementSpaceSize());
const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
p_b_grid, b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetElementSpaceSize());
p_b_grid, b_grid_desc_k0_n_k1.GetElementSpaceSize());
const auto ds_grid_buf = generate_tuple(
[&](auto i) {
return make_dynamic_buffer<AddressSpaceEnum::Global>(
@@ -320,113 +323,113 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum::Global>(
p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize());
const auto K0 = a_grid_desc_k0_m_k1.GetLength(I0);
const auto K0 = b_grid_desc_k0_n_k1.GetLength(I0);
// divide block work by [M, N]
const auto block_work_idx =
block_2_ctile_map.CalculateBottomIndex(make_multi_index(get_block_1d_id()));
// HACK: this force m/n_block_data_idx_on_grid into SGPR
const index_t m_block_data_idx_on_grid =
__builtin_amdgcn_readfirstlane(block_work_idx[I0] * MPerBlock);
const index_t n_block_data_idx_on_grid =
__builtin_amdgcn_readfirstlane(block_work_idx[I1] * NPerBlock);
// A matrix in LDS memory, dst of blockwise copy
constexpr auto a_block_desc_k0_m_k1 = GetABlockDescriptor_K0PerBlock_MPerBlock_K1();
constexpr auto b_block_desc_k0_n_k1 = GetBBlockDescriptor_K0PerBlock_NPerBlock_K1();
// A matrix blockwise copy
auto a_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1<
auto b_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1<
ThisThreadBlock,
AElementwiseOperation,
BElementwiseOperation,
ck::tensor_operation::element_wise::PassThrough,
InMemoryDataOperationEnum::Set,
Sequence<K0PerBlock * BBlockBufferSize, MPerBlock, K1>,
ABlockTransferThreadClusterLengths_K0_M_K1,
ABlockTransferThreadClusterArrangeOrder,
Sequence<K0PerBlock * ABlockBufferSize, NPerBlock, K1>,
BBlockTransferThreadClusterLengths_K0_N_K1,
BBlockTransferThreadClusterArrangeOrder,
ABDataType,
ABDataType,
decltype(a_grid_desc_k0_m_k1),
decltype(a_block_desc_k0_m_k1),
ABlockTransferSrcAccessOrder,
decltype(b_grid_desc_k0_n_k1),
decltype(b_block_desc_k0_n_k1),
BBlockTransferSrcAccessOrder,
Sequence<1, 0, 2>,
ABlockTransferSrcVectorDim,
BBlockTransferSrcVectorDim,
2,
ABlockTransferSrcScalarPerVector,
ABlockTransferDstScalarPerVector_K1,
BBlockTransferSrcScalarPerVector,
BBlockTransferDstScalarPerVector_K1,
1,
1,
AThreadTransferSrcResetCoordinateAfterRun,
BThreadTransferSrcResetCoordinateAfterRun,
true,
1>(a_grid_desc_k0_m_k1,
make_multi_index(0, m_block_data_idx_on_grid, 0),
a_element_op,
a_block_desc_k0_m_k1,
1>(b_grid_desc_k0_n_k1,
make_multi_index(0, n_block_data_idx_on_grid, 0),
b_element_op,
b_block_desc_k0_n_k1,
make_multi_index(0, 0, 0),
ck::tensor_operation::element_wise::PassThrough{});
ignore = b_element_op;
ignore = a_element_op;
// B matrix threadwise copy
constexpr auto b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3 =
constexpr auto a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3 =
make_naive_tensor_descriptor_packed(make_tuple(I1,
I1,
Number<K0PerThread>{}, // K0PerThread
I1, // NBlockId
Number<NXdlPerWave>{}, // repeat
Number<MXdlPerWave>{}, // repeat
I1, // waves
I1, // NPerXdlops
Number<K1>{}));
auto b_thread_buf = generate_tuple(
auto a_thread_buf = generate_tuple(
[&](auto i) {
ignore = i;
return StaticBuffer<AddressSpaceEnum::Vgpr,
ABDataType,
b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetElementSpaceSize(),
a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3.GetElementSpaceSize(),
true>{};
},
Number<BBlockBufferSize>{});
Number<ABlockBufferSize>{});
const auto wave_id = GetWaveIdx();
const auto wave_k_n_id = GetWaveKNIdx(wave_id[I2]);
const auto wave_k_m_id = GetWaveKMIdx(wave_id[I2]);
#if 0
const index_t block_id = get_block_1d_id();
const index_t thread_id = get_thread_local_1d_id();
printf("block id: %d m blockid: %d n block id: %d ,thread id: %d, wave id :{%d %d %d} "
"kn id: {%d %d}\n",
block_id,
block_work_idx[I0],
block_work_idx[I1],
thread_id,
wave_id[I0],
wave_id[I1],
wave_id[I2],
wave_k_n_id[I0],
wave_k_n_id[I1]);
printf("mfma thread k per xdlops: %d K0PerThread: %d HasMainK0BlockLoop: %d K0: %d \t",
xdlops_gemm.K0PerXdlops, K0PerThread, HasMainK0BlockLoop, b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetLength(I0));
#endif
// #if 0
// const index_t block_id = get_block_1d_id();
// const index_t thread_id = get_thread_local_1d_id();
// printf("block id: %d m blockid: %d n block id: %d ,thread id: %d, wave id :{%d %d %d} "
// "kn id: {%d %d}\n",
// block_id,
// block_work_idx[I0],
// block_work_idx[I1],
// thread_id,
// wave_id[I0],
// wave_id[I1],
// wave_id[I2],
// wave_k_n_id[I0],
// wave_k_n_id[I1]);
// printf("mfma thread k per xdlops: %d K0PerThread: %d HasMainK0BlockLoop: %d K0: %d \t",
// xdlops_gemm.K0PerXdlops, K0PerThread, HasMainK0BlockLoop, b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetLength(I0));
// #endif
auto b_threadwise_copy =
auto a_threadwise_copy =
ThreadwiseTensorSliceTransfer_v2<ABDataType,
ABDataType,
decltype(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3),
decltype(b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3),
decltype(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3),
decltype(a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3),
Sequence<I1,
I1,
Number<K0PerThread>{},
I1,
Number<NXdlPerWave>{},
Number<MXdlPerWave>{},
I1,
I1,
Number<K1>{}>,
Sequence<0, 1, 2, 3, 4, 5, 6, 7>,
7,
BBlockTransferSrcScalarPerVector,
BThreadTransferSrcResetCoordinateAfterRun,
ABlockTransferSrcScalarPerVector,
AThreadTransferSrcResetCoordinateAfterRun,
true>(
b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3,
a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3,
make_multi_index(
0, wave_k_n_id[I0], 0, block_work_idx[I1], 0, wave_id[I1], wave_k_n_id[I1], 0));
0, wave_k_m_id[I0], 0, block_work_idx[I0], 0, wave_id[I0], wave_k_m_id[I1], 0));
// GEMM definition
// c_mtx += transpose(a_mtx) * b_mtx
@@ -439,8 +442,8 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
BlockSize,
ABDataType,
AccDataType,
decltype(a_block_desc_k0_m_k1),
decltype(b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3),
decltype(a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3),
decltype(b_block_desc_k0_n_k1),
MPerBlock,
NPerBlock,
K0PerBlock,
@@ -453,64 +456,64 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
auto c_thread_buf = blockwise_gemm.GetCThreadBuffer();
// LDS allocation for A
auto a_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
static_cast<ABDataType*>(p_shared), a_block_desc_k0_m_k1.GetElementSpaceSize());
auto b_block_buf = make_dynamic_buffer<AddressSpaceEnum::Lds>(
static_cast<ABDataType*>(p_shared), b_block_desc_k0_n_k1.GetElementSpaceSize());
// gridwise GEMM pipeline
constexpr auto a_block_slice_copy_step =
make_multi_index(K0PerBlock * BBlockBufferSize, 0, 0);
constexpr auto b_thread_slice_copy_step = make_multi_index(1, 0, 0, 0, 0, 0, 0, 0);
constexpr auto a_thread_slice_copy_step = make_multi_index(1, 0, 0, 0, 0, 0, 0, 0);
constexpr auto b_block_slice_copy_step =
make_multi_index(K0PerBlock * ABlockBufferSize, 0, 0);
// preload data to regiester and LDS
{
// Read
a_blockwise_copy.RunRead(a_grid_desc_k0_m_k1, a_grid_buf);
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_m_k1, a_block_slice_copy_step);
b_blockwise_copy.RunRead(b_grid_desc_k0_n_k1, b_grid_buf);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_n_k1, b_block_slice_copy_step);
static_for<0, BBlockBufferSize, 1>{}([&](auto ii) {
b_threadwise_copy.Run(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3,
b_grid_buf,
b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3,
static_for<0, ABlockBufferSize, 1>{}([&](auto ii) {
a_threadwise_copy.Run(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3,
a_grid_buf,
a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3,
make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
b_thread_buf(Number<ii>{}));
b_threadwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3,
b_thread_slice_copy_step);
a_thread_buf(Number<ii>{}));
a_threadwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3,
a_thread_slice_copy_step);
});
// Initialize C
c_thread_buf.Clear();
// a data write to lds
a_blockwise_copy.RunWrite(a_block_desc_k0_m_k1, a_block_buf);
b_blockwise_copy.RunWrite(b_block_desc_k0_n_k1, b_block_buf);
// main body
if constexpr(HasMainK0BlockLoop)
{
index_t K0BlockMainLoop =
__builtin_amdgcn_readfirstlane(K0 / (BBlockBufferSize * K0PerBlock));
__builtin_amdgcn_readfirstlane(K0 / (ABlockBufferSize * K0PerBlock));
index_t i = 0;
do
{
a_blockwise_copy.RunRead(a_grid_desc_k0_m_k1, a_grid_buf);
blockwise_gemm.ResetABlockStartWindow();
b_blockwise_copy.RunRead(b_grid_desc_k0_n_k1, b_grid_buf);
blockwise_gemm.ResetBBlockStartWindow();
block_sync_lds();
static_for<0, BBlockBufferSize, 1>{}([&](auto ii) {
blockwise_gemm.Run(a_block_buf, b_thread_buf(Number<ii>{}), c_thread_buf);
blockwise_gemm.MoveABlockSliceWindow();
static_for<0, ABlockBufferSize, 1>{}([&](auto ii) {
blockwise_gemm.Run(a_thread_buf(Number<ii>{}), b_block_buf, c_thread_buf);
blockwise_gemm.MoveBBlockSliceWindow();
s_nop();
b_threadwise_copy.Run(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3,
b_grid_buf,
b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3,
a_threadwise_copy.Run(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3,
a_grid_buf,
a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3,
make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
b_thread_buf(Number<ii>{}));
b_threadwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3,
b_thread_slice_copy_step);
a_thread_buf(Number<ii>{}));
a_threadwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3,
a_thread_slice_copy_step);
});
block_sync_lds();
a_blockwise_copy.RunWrite(a_block_desc_k0_m_k1, a_block_buf);
b_blockwise_copy.RunWrite(b_block_desc_k0_n_k1, b_block_buf);
// move a and b window
a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_m_k1,
a_block_slice_copy_step);
b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_n_k1,
b_block_slice_copy_step);
i += 1;
} while(i < (K0BlockMainLoop - 1));
@@ -520,11 +523,11 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle
{
block_sync_lds();
blockwise_gemm.ResetABlockStartWindow();
blockwise_gemm.ResetBBlockStartWindow();
static_for<0, BBlockBufferSize, 1>{}([&](auto ii) {
blockwise_gemm.Run(a_block_buf, b_thread_buf(Number<ii>{}), c_thread_buf);
blockwise_gemm.MoveABlockSliceWindow();
static_for<0, ABlockBufferSize, 1>{}([&](auto ii) {
blockwise_gemm.Run(a_thread_buf(Number<ii>{}), b_block_buf, c_thread_buf);
blockwise_gemm.MoveBBlockSliceWindow();
});
}
}