ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 01:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

tmp

Browse Source
This commit is contained in:
Jakub Piasecki
2025-06-23 14:02:03 +00:00
parent 1ac8f1c744
commit f5345174e4
7 changed files with 89 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
include/ck_tile/ops/grouped_convolution/utils/grouped_convolution_utils.hpp
View File

@@ -14,14 +14,15 @@ namespace ck_tile {
/// This structure is passed to Grouped Convolution Kernels when creating kernel
/// arguments object. It contain all necessary information required to
/// build proper kernel argument and launch kernel on GPU.
template <typename InPtr, typename WeiPtr, typename OutPtr>
struct GroupedConvHostArgs : public conv::ConvParam
{
CK_TILE_HOST GroupedConvHostArgs() = delete;
CK_TILE_HOST GroupedConvHostArgs(ConvParam conv_param,
const void* in_ptr_,
const void* wei_ptr_,
InPtr in_ptr_,
WeiPtr wei_ptr_,
const std::vector<const void*> ds_ptr_,
void* out_ptr_,
OutPtr out_ptr_,
index_t k_batch_)
: conv::ConvParam(conv_param),
in_ptr(in_ptr_),
@@ -32,37 +33,15 @@ struct GroupedConvHostArgs : public conv::ConvParam
{
}
const void* in_ptr;
const void* wei_ptr;
InPtr in_ptr;
WeiPtr wei_ptr;
const std::vector<const void*> ds_ptr;
void* out_ptr;
OutPtr out_ptr;
index_t k_batch;
};
struct GroupedConvBwdWeightHostArgs : public conv::ConvParam
{
CK_TILE_HOST GroupedConvBwdWeightHostArgs() = delete;
CK_TILE_HOST GroupedConvBwdWeightHostArgs(ConvParam conv_param,
const void* in_ptr_,
void* wei_ptr_,
const std::vector<const void*> ds_ptr_,
const void* out_ptr_,
index_t k_batch_)
: conv::ConvParam(conv_param),
in_ptr(in_ptr_),
wei_ptr(wei_ptr_),
ds_ptr(ds_ptr_),
out_ptr(out_ptr_),
k_batch(k_batch_)
{
}
const void* in_ptr;
void* wei_ptr;
const std::vector<const void*> ds_ptr;
const void* out_ptr;
index_t k_batch;
};
using GroupedConvFwdHostArgs = GroupedConvHostArgs<const void*, const void*, void*>;
using GroupedConvBwdWeightHostArgs = GroupedConvHostArgs<const void*, void*, const void*>;
template <index_t NDimSpatial_,
ConvolutionSpecialization ConvSpecialization_,

22
include/ck_tile/ops/grouped_convolution/utils/transform_conv_bwd_weight_to_gemm.hpp
View File

@@ -415,7 +415,7 @@ struct TransformConvBwdWeightToGemm
#endif
template <index_t NDim = NDimSpatial, typename std::enable_if<NDim == 1, bool>::type = false>
CK_TILE_HOST auto make_out_grid_desc() const
CK_TILE_HOST auto make_out_grid_desc(const index_t GemmKBatch) const
{
// NWGK
const index_t NDoHoWoStride = G_ * K_;
@@ -423,7 +423,7 @@ struct TransformConvBwdWeightToGemm
// TODO Add support for NumGroupsToMerge > 1
return make_naive_tensor_descriptor(make_tuple(N_ * Wo_, K_),
return make_naive_tensor_descriptor(make_tuple(N_ * Wo / GemmKBatch_, K_),
make_tuple(NDoHoWoStride, KStride));
}
@@ -538,23 +538,22 @@ struct TransformConvBwdWeightToGemm
// properties
template <index_t NDim = NDimSpatial, typename std::enable_if<NDim == 1, bool>::type = false>
CK_TILE_HOST auto MakeABCGridDescriptor_A_K0_M_K1_B_K0_N_K1_C_M_N() const
CK_TILE_HOST auto MakeABCGridDescriptor_A_K0_M_K1_B_K0_N_K1_C_M_N(const index_t GemmKBatch) const
{
// Assume NumGroupsToMerge == 1 for now
const index_t GemmKTotal = N_ * Wo_;
const index_t GemmKTotal = N_ * Wo_ / KBatch; // tmp
const index_t GemmM = K_ * NumGroupsToMerge;
const index_t GemmN = C_ * X_ * NumGroupsToMerge;
const auto PadGemmM = MPerBlock - GemmM % MPerBlock;
const auto PadGemmN = NPerBlock - GemmN % NPerBlock;
const index_t GemmKBatch = 1;
const index_t GemmK0 =
integer_divide_ceil(GemmKTotal, GemmK1Number * K0PerBlock * GemmKBatch) * K0PerBlock;
const index_t GemmKPad = GemmKBatch * GemmK0 * GemmK1Number;
const auto out_grid_desc = make_out_grid_desc<NDimSpatial>();
const auto in_grid_desc = make_in_grid_desc<NDimSpatial>();
const auto out_grid_desc = make_out_grid_desc<NDimSpatial>(GemmKBatch);
const auto in_grid_desc = make_in_grid_desc<NDimSpatial>(GemmKBatch);
const auto wei_grid_desc = make_wei_grid_desc<NDimSpatial>();
// A: output tensor comes in K_M
@@ -597,6 +596,13 @@ struct TransformConvBwdWeightToGemm
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<1>{}, sequence<0>{}));
const auto in_gemmkpad_gemmn_grid_desc = transform_tensor_descriptor(
in_gemmktotal_gemmn_grid_desc,
make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
make_right_pad_transform(GemmN, PadGemmN)),
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<1>{}, sequence<0>{}));
const auto wei_gemmm_gemmn_pad_grid_desc =
transform_tensor_descriptor(wei_grid_desc,
make_tuple(make_right_pad_transform(GemmM, PadGemmM),
@@ -604,7 +610,7 @@ struct TransformConvBwdWeightToGemm
make_tuple(sequence<0>{}, sequence<1>{}),
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tuple(out_gemmkpad_gemmm_grid_desc,
return make_tuple(out_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc,
in_gemmkpad_gemmn_grid_desc,
wei_gemmm_gemmn_pad_grid_desc);
}