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

tmp code 2

Browse Source
This commit is contained in:
Qun Lin
2025-05-28 22:44:24 +08:00
parent 45b4c48000
commit b9466aff4d
1 changed files with 248 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files

526
example/20_grouped_conv_bwd_weight/grouped_conv_bwd_weight_xdl_fp16.cpp
View File

@@ -16,266 +16,251 @@ using InElementOp = PassThrough;
using WeiElementOp = PassThrough;
using OutElementOp = PassThrough;
#if 0
template <ck::index_t NDimSpatial>
using DeviceConvBwdWeightInstance =
ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Xdl_CShuffle<
NDimSpatial,
ck::tuple_element_t<NDimSpatial - 1,
ck::Tuple<ck::tensor_layout::convolution::GNWC,
ck::tensor_layout::convolution::GNHWC,
ck::tensor_layout::convolution::GNDHWC>>,
ck::tuple_element_t<NDimSpatial - 1,
ck::Tuple<ck::tensor_layout::convolution::GKXC,
ck::tensor_layout::convolution::GKYXC,
ck::tensor_layout::convolution::GKZYXC>>,
ck::tuple_element_t<NDimSpatial - 1,
ck::Tuple<ck::tensor_layout::convolution::GNWK,
ck::tensor_layout::convolution::GNHWK,
ck::tensor_layout::convolution::GNDHWK>>,
InDataType, // InDataType
WeiDataType, // WeiDataType
OutDataType, // OutDataType
AccDataType, // AccDataType
InElementOp, // InElementwiseOperation
WeiElementOp, // WeiElementwiseOperation
OutElementOp, // OutElementwiseOperation
ConvBwdWeightDefault, // ConvolutionBackwardWeightSpecialization
256, // BlockSize
128, // MPerBlock
128, // NPerBlock
4, // K0PerBlock
8, // K1
32, // MPerXdl
32, // NPerXdl
2, // MXdlPerWave
2, // NXdlPerWave
S<1, 4, 16, 4>, // ABlockTransferThreadClusterLengths_K0_M_K1
S<0, 3, 1, 2>, // ABlockTransferThreadClusterArrangeOrder
S<0, 2, 1, 3>, // ABlockTransferSrcAccessOrder
2, // ABlockTransferSrcVectorDim
1, // ABlockTransferSrcScalarPerVector
1, // ABlockTransferDstScalarPerVector_K1
true, // ABlockLdsAddExtraM
S<1, 4, 16, 4>, // BBlockTransferThreadClusterLengths_K0_N_K1
S<0, 3, 1, 2>, // BBlockTransferThreadClusterArrangeOrder
S<0, 2, 1, 3>, // BBlockTransferSrcAccessOrder
2, // BBlockTransferSrcVectorDim
1, // BBlockTransferSrcScalarPerVector
1, // BBlockTransferDstScalarPerVector_K1
true, // BBlockLdsAddExtraN
1, // CShuffleMXdlPerWavePerShuffle
1, // CShuffleNXdlPerWavePerShuffle
S<1, 32, 1, 4>, // CBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock
1>; // CBlockTransferScalarPerVector_NWaveNPerXdl
#endif
// DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 64, 32, 32, 32, 8, 32, 32, 1, 1, S<4, 8, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 2, 2, false, S<4, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 2, 2, false, 1, 1, S<1, 8, 1, 8>, 1, Scheduler, PipelineVersion, 2>,
//ConvBwdWeightDefault,
//is_NHWGC_GKYXC_NHWGK
using ALayout = ck::tensor_layout::convolution::NHWGC;
using BLayout = ck::tensor_layout::convolution::GKYXC;
using ELayout = ck::tensor_layout::convolution::NHWGK;
//using Scheduler =ck::BlockGemmPipelineScheduler::Intrawave;
//using PipelineVersion =ck::BlockGemmPipelineVersion::v1;
template <ck::index_t NDimSpatial>
using DeviceConvBwdWeightInstance =
ck::tensor_operation::device::DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle<
NDimSpatial,
ALayout,
BLayout,
ELayout,
F16,
F16,
F16,
F32,
PassThrough,
PassThrough,
PassThrough,
ConvBwdWeightDefault,
64,
32,//16,
64,
32,//64,
8,
32, //16,
32, //16,
1,
2, //4,
S<4, 8, 1>,// S<8, 8, 1>
S<2, 0, 1>,
S<1, 0, 2>,
1,
2,
2,
false,
S<4, 16, 1>,
S<2, 0, 1>,
S<1, 0, 2>,
1,
2,
2,
false,
1,
1,
S<1, 8, 1, 8>,
1,
ck::BlockGemmPipelineScheduler::Intrawave,
ck::BlockGemmPipelineVersion::v3,
2>;
#if 0
DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle< NDimSpatial, ALayout, BLayout, ELayout, F16, F16, F16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 64, 32, 64, 32, 8, 32, 32, 1, 2, S<4, 8, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 2, 2, false, S<4, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 2, 2, false, 1, 1, S<1, 8, 1, 8>, 1, Scheduler, PipelineVersion, 2>,
64, 16, 16, 32, 8, 16, 16, 1, 1, S<4, 8, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 1, 4, false, S<4, 8, 1>,
S<2, 0, 1>, S<1, 0, 2>, 1, 1, 4, false, 1, 1, S<1, 8, 1, 8>, 1,
ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, 8 > ;
#endif
namespace ck {
namespace tensor_operation {
namespace device {
static constexpr index_t WaveSize = 64;
template<typename Argument,
typename InDataType>
void __device__ load_input_from_global(const Argument* arg, InDataType* p_in, index_t n, uint32_t* p_scratch)
static constexpr index_t W_PACK = 2; // WaveSize / arg->input_spatial_lengths_[1];
static constexpr index_t Tile_H = 32;
static constexpr index_t Tile_W = 32;
static constexpr index_t N_Pack = 2;
static constexpr index_t Pad_H = 2;
static constexpr index_t Pad_W = 2;
static constexpr index_t Filter_X = 5;
static constexpr index_t Filter_Y = 5;
static constexpr index_t SizeOfType = 2;
static constexpr index_t ShareMemSize = Tile_H * Tile_W * N_Pack * SizeOfType;
static constexpr index_t ScratchSize = ShareMemSize / 64 / 4;
template <typename DataType>
void __device__ load_data_from_global(DataType* p,
index_t n_stride,
index_t h,
index_t w,
index_t h_stride,
index_t w_stride,
uint32_t* p_scratch)
{
InDataType* p_in_n = p_in + arg->a_g_n_k_wos_strides[1] * n;
InDataType* p_in_n_1 = p_in + arg->a_g_n_k_wos_strides[1] * (n + 1);
ignore = h;
ignore = w;
DataType* p_1 = p + n_stride;
static_assert(sizeof(DataType) == 2);
static_assert(Pad_H % W_PACK == 0);
const uint32_t W_PACK = 2; //WaveSize / arg->input_spatial_lengths_[1];
static_assert(sizeof(InDataType) == 2);
const index_t x = threadIdx.x % (WaveSize/W_PACK);
const index_t y_base = threadIdx.x / (WaveSize/W_PACK);
auto get_offset = [&](index_t y, index_t x)
auto get_offset = [&](index_t y_, index_t x_)
{
return y * arg->input_spatial_stride_[0] + x * arg->input_spatial_stride_[1];
}
for (uint32_t i = 0; i < arg->input_spatial_lengths_[1] / W_PACK; i++)
return y_ * h_stride + x_ * w_stride;
};
if(x >= Pad_W && x < w + Pad_W)
{
const index_t offset = get_offset(i * W_PACK + threadIdx.x / (64/W_PACK), threadIdx.x % (64/W_PACK));
auto tmp0 = p_in_n[offset];
auto tmp1 = p_in_n_1[offset];
InDataType* p_scratch_offset = reinterpret_cast<InDataType*>(&p_scratch[i]);
p_scratch_offset[0] = tmp1;
p_scratch_offset[1] = tmp1;
static_for<0, Tile_H / W_PACK, 1>{}([&](auto i) {
const index_t y = y_base + i * W_PACK;
if constexpr (i * W_PACK >= Pad_H && i * W_PACK < Tile_H / W_PACK - Pad_H)
{
const index_t offset = get_offset(y, x);
half2_t tmp = {};
tmp[0] = p[offset];
tmp[1] = p_1[offset];
p_scratch[i] = bit_cast<uint32_t>(tmp);
}
});
}
}
template<typename Argument,
typename OutDataType>
void __device__ load_output_from_global(const Argument* arg, OutDataType* p_out, index_t n, uint32_t* p_scratch)
void __device__ write_data_to_lds(const uint32_t* p_scratch, uint32_t* p_sharemem)
{
OutDataType* p_out_n = p_out + arg->a_g_n_k_wos_strides[1] * n;
OutDataType* p_out_n_1 = p_out + arg->a_g_n_k_wos_strides[1] * (n + 1);
const index_t x = threadIdx.x % (WaveSize/W_PACK);
const index_t y_base = threadIdx.x / (WaveSize/W_PACK);
//static_assert(N_Pack * sizeof(InDataType)/ sizeof(uint32_t) == 1);
const uint32_t W_PACK = 2; //WaveSize / arg->input_spatial_lengths_[1];
static_assert(sizeof(OutDataType) == 2);
auto get_offset = [&](index_t y, index_t x)
{
return y * arg->output_spatial_stride_[0] + x * arg->output_spatial_stride_[1];
}
for (uint32_t i = 0; i < arg->output_spatial_lengths_[1] / W_PACK; i++)
{
const index_t offset = get_offset(i * W_PACK + threadIdx.x / (64/W_PACK), threadIdx.x % (64/W_PACK));
auto tmp0 = p_out_n[offset];
auto tmp1 = p_out_n_1[offset];
InDataType* p_scratch_offset = reinterpret_cast<InDataType*>(&p_scratch[i]);
p_scratch_offset[0] = tmp1;
p_scratch_offset[1] = tmp1;
}
auto get_offset = [&](index_t y_, index_t x_) { return (y_ * Tile_W + x_); };
static_for<0, Tile_H / W_PACK, 1>{}([&](auto i) {
const index_t y = y_base + i * W_PACK;
const index_t offset = get_offset(y, x);
p_sharemem[offset] = p_scratch[i];
});
}
write_input_to_lds
void __device__ run_conv_bwd_weight(index_t x, index_t y, index_t H, index_t W, uint32_t* p_share_in, uint32_t* p_share_out,float& acc)
{
ignore = H;
ignore = W;
auto get_in = [&](int h_, int w_)
{
return p_share_in[(h_ + y) * Tile_W + w_ + x];
};
auto get_out = [&](int h_, int w_)
{
return p_share_out[(h_ + Pad_H) * Tile_W + w_ + Pad_W];
};
if (x < Filter_X && y < Filter_Y)
{
//for (int ho = 0; ho < H; ho++)
static_for<0, Tile_H - Pad_H - Pad_H, 1>{}([&](auto ho)
{
//for (int wo = 0; wo < W; wo++)
static_for<0, Tile_W - Pad_W - Pad_W, 1>{}([&](auto wo)
{
uint32_t v_in = get_in(ho, wo);
uint32_t v_out = get_out(ho, wo);
acc = __builtin_amdgcn_fdot2(bit_cast<half2_t>(v_in), bit_cast<half2_t>(v_out), acc, false);
});
});
}
}
template <typename Argument, typename WeiDataType>
void __device__ write_output(const Argument* arg, int g, int y, int x, WeiDataType acc)
{
const index_t Wei_G_Stride = arg->wei_g_k_c_xs_strides_[0];
const index_t Y_Stride = arg->wei_g_k_c_xs_strides_[3];
const index_t X_Stride = arg->wei_g_k_c_xs_strides_[4];
if (y < Filter_Y && x < Filter_X)
{
auto p_wei = arg->p_wei_grid_ + Wei_G_Stride * g + y * Y_Stride + x * X_Stride;
*p_wei = acc;
}
}
template <typename Argument, index_t MinimumOccupancy = 1>
__global__ void
#if CK_USE_LAUNCH_BOUNDS
__launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
__launch_bounds__(64, MinimumOccupancy)
#endif
kernel_grouped_conv_bwd_weight_naive(Argument* arg)
kernel_grouped_conv_bwd_weight_naive(const Argument* arg)
{
#if(!defined(__HIP_DEVICE_COMPILE__) || defined(__gfx9__))
const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.x);
index_t n_idx = 0;
constexpr index_t Tile_H = 32;
constexpr index_t Tile_W = 32;
constexpr index_t N_Pack = 2;
constexpr index_t SizeOfType = 2;
constexpr index_t ShareMemSize = Tile_H * Tile_W * N_Pack * SizeOfType;
__shared__ char p_input_0[ShareMemSize];
__shared__ char p_input_1[ShareMemSize];
__shared__ char p_output_0[ShareMemSize];
__shared__ char p_output_1[ShareMemSize];
__shared__ uint32_t p_input_0[ShareMemSize/sizeof(uint32_t)];
//__shared__ char p_input_1[ShareMemSize];
__shared__ uint32_t p_output_0[ShareMemSize/sizeof(uint32_t)];
//__shared__ char p_output_1[ShareMemSize];
constexpr index_t ScratchSize = ShareMemSize / 64 / 4;
uint32_t p_input_0_scratch[ScratchSize];
uint32_t p_input_1_scratch[ScratchSize];
uint32_t p_output_0_scratch[ScratchSize];
uint32_t p_output_1_scratch[ScratchSize];
InDataType* p_in = arg->p_in_grid + g_idx * arg->a_g_n_k_wos_strides[0];
OutDataType* p_out = arg->p_out_grid + g_idx * arg->e_g_k_c_xs_strides[0];
static constexpr index_t spatial_offset = 3;
//const index_t G = arg->in_g_n_c_wis_lengths[0];
const index_t N = arg->in_g_n_c_wis_lengths_[1];
// In
const index_t Hi = arg->in_g_n_c_wis_lengths_[spatial_offset + 0];
const index_t Wi = arg->in_g_n_c_wis_lengths_[spatial_offset + 1];
const index_t Hi_Stride = arg->in_g_n_c_wis_strides_[spatial_offset + 0];
const index_t Wi_Stride = arg->in_g_n_c_wis_strides_[spatial_offset + 0];
const index_t In_G_Stride = arg->in_g_n_c_wis_strides_[0];
const index_t In_N_Stride = arg->in_g_n_c_wis_strides_[1];
// Out
const index_t Ho = arg->out_g_n_k_wos_lengths_[spatial_offset + 0];
const index_t Wo = arg->out_g_n_k_wos_lengths_[spatial_offset + 1];
const index_t Ho_Stride = arg->out_g_n_k_wos_strides_[spatial_offset + 0];
const index_t Wo_Stride = arg->out_g_n_k_wos_strides_[spatial_offset + 0];
const index_t Out_G_Stride = arg->out_g_n_k_wos_strides_[0];
const index_t Out_N_Stride = arg->out_g_n_k_wos_strides_[1];
static_for<0,ScratchSize, 1>{}([&](auto i)
{
p_input_0_scratch[i] = 0;
p_output_0_scratch[i] = 0;
p_input_1_scratch[i] = 0;
p_output_1_scratch[i] = 0;
});
// Wei
//static_assert(sizeof(InDataType) == 2);
//static_assert(sizeof(OutputDataType) == 2);
//
auto* p_in = arg->p_in_grid_ + g_idx * In_G_Stride + n_idx * In_N_Stride;
auto* p_out = arg->p_out_grid_ + g_idx * Out_G_Stride + n_idx * Out_N_Stride;
// prefetch 0
load_input_from_global(arg, p_in, n_idx, p_input_0_scratch);
load_output_from_global(arg, p_out, n_idx, p_output_0_scratch);
load_data_from_global(p_in, In_N_Stride, Hi, Wi, Hi_Stride, Wi_Stride, p_input_0_scratch);
load_data_from_global(p_out, In_N_Stride, Ho, Wo, Ho_Stride, Wo_Stride, p_output_0_scratch);
//load_data_from_global(p_in + In_N_Stride, Hi, Wi, Hi_Stride, Wi_Stride, p_input_1_scratch);
//load_data_from_global(p_out + Out_N_Stride, Ho, Wo, Ho_Stride, Wo_Stride, p_output_1_scratch);
p_in += 2 * In_N_Stride;
p_out += 2 * Out_N_Stride;
#if 0
static_for<0,ScratchSize, 1>{}([&](auto i)
{
p_input_0_scratch[i] = (p_input_0_scratch[i] << 16) | (p_input_1_scratch[i] & 0xffff);
p_output_0_scratch[i] = (p_output_0_scratch[i] << 16) | (p_output_1_scratch[i] & 0xffff);
});
#endif
// prefetch 1
load_input_from_global(arg, p_in, n_idx + 1, p_input_1_scratch);
load_output_from_global(arg, p_out, n_idx + 1, p_output_0_scratch);
//load_input_from_global(arg, p_in, n_idx + 1, p_input_1_scratch);
//load_output_from_global(arg, p_out, n_idx + 1, p_output_0_scratch);
// write 0
write_input_to_lds(arg, p_input_0_scratch);
write_output_to_lds(arg, p_output_0_scratch);
write_data_to_lds(p_input_0_scratch, p_input_0);
write_data_to_lds(p_output_0_scratch, p_output_0);
index_t x = threadIdx.x % Filter_X;
index_t y = threadIdx.x / Filter_Y;
float acc = 0;
index_t num_loop = N / 2 - 1;
while(num_loop > 0)
{
// prefetch 0
load_input_from_global();
load_output_from_global();
load_data_from_global(p_in, In_N_Stride, Hi, Wi, Hi_Stride, Wi_Stride, p_input_0_scratch);
load_data_from_global(p_out, In_N_Stride, Ho, Wo, Ho_Stride, Wo_Stride, p_output_0_scratch);
//load_data_from_global(p_in + In_N_Stride, Hi, Wi, Hi_Stride, Wi_Stride, p_input_1_scratch);
//load_data_from_global(p_out + Out_N_Stride, Ho, Wo, Ho_Stride, Wo_Stride, p_output_1_scratch);
p_in += 2 * In_N_Stride;
p_out += 2 * Out_N_Stride;
#if 0
static_for<0,ScratchSize, 1>{}([&](auto i)
{
p_input_0_scratch[i] = (p_input_0_scratch[i] << 16) | (p_input_1_scratch[i] & 0xffff);
p_output_0_scratch[i] = (p_output_0_scratch[i] << 16) | (p_output_1_scratch[i] & 0xffff);
});
#endif
// do conv_bwd on 0
run_conv_bwd_weight();
run_conv_bwd_weight(x, y, Ho, Wo, p_input_0, p_output_0, acc);
// write 1
write_input_to_lds();
write_output_to_lds();
//write_input_to_lds();
//write_output_to_lds();
// prefetch 1
load_input_from_global();
load_output_from_global();
//load_input_from_global();
//load_output_from_global();
// do conv_bwd on 1
run_conv_bwd_weight();
//run_conv_bwd_weight();
// write 0
write_input_to_lds();
write_output_to_lds();
write_data_to_lds(p_input_0_scratch, p_input_0);
write_data_to_lds(p_output_0_scratch, p_output_0);
num_loop --;
};
if (tail_num == 1)
// tail
{
run_conv_bwd_weight(x, y, Ho, Wo, p_input_0, p_output_0, acc);
}
if (tail_num == 2)
{
}
write_output();
write_output(arg, g_idx, y, x, acc);
#else
ignore = karg;
#endif // end of if (defined(__gfx9__))
}
template <ck::index_t NDimSpatial
template <ck::index_t NDimSpatial,
typename InLayout,
typename WeiLayout,
typename OutLayout,
@@ -302,6 +287,7 @@ struct DeviceGroupedConvBwdWeightNaive
ComputeTypeA,
ComputeTypeB>
{
using DeviceOp = DeviceGroupedConvBwdWeightNaive;
static_assert(is_same_v<InElementwiseOperation, element_wise::PassThrough>);
static_assert(is_same_v<WeiElementwiseOperation, element_wise::PassThrough>);
static_assert(is_same_v<OutElementwiseOperation, element_wise::PassThrough>);
@@ -311,12 +297,12 @@ struct DeviceGroupedConvBwdWeightNaive
Argument(const InDataType* p_in_grid,
WeiDataType* p_wei_grid,
const OutDataType* p_out_grid,
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_lengths, // input
const std::array<index_t, NDimSpatial + 3>& b_g_n_c_wis_strides,
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_lengths, // weight
const std::array<index_t, NDimSpatial + 3>& e_g_k_c_xs_strides,
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_lengths, // output
const std::array<index_t, NDimSpatial + 3>& a_g_n_k_wos_strides,
const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_lengths, // input
const std::array<index_t, NDimSpatial + 3>& in_g_n_c_wis_strides,
const std::array<index_t, NDimSpatial + 3>& wei_g_k_c_xs_lengths, // weight
const std::array<index_t, NDimSpatial + 3>& wei_g_k_c_xs_strides,
const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_lengths, // output
const std::array<index_t, NDimSpatial + 3>& out_g_n_k_wos_strides,
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides,
const std::array<ck::index_t, NDimSpatial>& conv_filter_dilations,
const std::array<ck::index_t, NDimSpatial>& input_left_pads,
@@ -325,39 +311,24 @@ struct DeviceGroupedConvBwdWeightNaive
WeiElementwiseOperation wei_element_op,
OutElementwiseOperation out_element_op,
ck::index_t split_k)
: p_a_grid_{p_out_grid},
p_b_grid_{p_in_grid},
p_e_grid_{p_wei_grid},
a_grid_desc_k0_m_k1_{},
b_grid_desc_k0_n_k1_{},
ce_grid_desc_m_n_{},
c_grid_desc_mblock_mperblock_nblock_nperblock_{},
compute_ptr_offset_of_batch_{},
a_element_op_{out_element_op},
b_element_op_{in_element_op},
cde_element_op_{wei_element_op},
Conv_G_{b_g_n_c_wis_lengths[0]},
Conv_N_{b_g_n_c_wis_lengths[1]},
Conv_K_{e_g_k_c_xs_lengths[1]},
Conv_C_{b_g_n_c_wis_lengths[2]},
input_spatial_lengths_{},
filter_spatial_lengths_{},
output_spatial_lengths_{},
conv_filter_strides_{conv_filter_strides},
input_left_pads_{input_left_pads},
input_right_pads_{input_right_pads},
: p_in_grid_{p_in_grid},
p_wei_grid_{p_wei_grid},
p_out_grid_{p_out_grid},
out_element_op_{out_element_op},
in_element_op_{in_element_op},
wei_element_op_{wei_element_op},
in_g_n_c_wis_lengths_(in_g_n_c_wis_lengths),
in_g_n_c_wis_strides_(in_g_n_c_wis_strides),
wei_g_k_c_xs_lengths_(wei_g_k_c_xs_lengths),
wei_g_k_c_xs_strides_(wei_g_k_c_xs_strides),
out_g_n_k_wos_lengths_(out_g_n_k_wos_lengths),
out_g_n_k_wos_strides_(out_g_n_k_wos_strides),
conv_filter_strides_(conv_filter_strides),
conv_filter_dilations_(conv_filter_dilations),
input_left_pads_(input_left_pads),
input_right_pads_(input_right_pads),
k_batch_{split_k}
{
constexpr index_t spatial_offset = 3;
std::copy(begin(b_g_n_c_wis_lengths) + spatial_offset,
end(b_g_n_c_wis_lengths),
begin(input_spatial_lengths_));
std::copy(begin(e_g_k_c_xs_lengths) + spatial_offset,
end(e_g_k_c_xs_lengths),
begin(filter_spatial_lengths_));
std::copy(begin(a_g_n_k_wos_lengths) + spatial_offset,
end(a_g_n_k_wos_lengths),
begin(output_spatial_lengths_));
}
std::size_t GetWorkspaceSizeBytes() const
@@ -365,28 +336,24 @@ struct DeviceGroupedConvBwdWeightNaive
return 0;
}
const ADataType* p_a_grid_;
const BDataType* p_b_grid_;
EDataType* p_e_grid_;
const InDataType* p_in_grid_;
WeiDataType* p_wei_grid_;
const OutDataType* p_out_grid_;
index_t M01_;
index_t N01_;
OutElementwiseOperation out_element_op_;
InElementwiseOperation in_element_op_;
WeiElementwiseOperation wei_element_op_;
OutElementwiseOperation a_element_op_;
InElementwiseOperation b_element_op_;
WeiElementwiseOperation cde_element_op_;
// for checking IsSupportedArgument()
const index_t Conv_G_;
const index_t Conv_N_;
const index_t Conv_K_;
const index_t Conv_C_;
std::array<ck::index_t, NDimSpatial> input_spatial_lengths_;
std::array<ck::index_t, NDimSpatial> filter_spatial_lengths_;
std::array<ck::index_t, NDimSpatial> output_spatial_lengths_;
const std::array<ck::index_t, NDimSpatial>& conv_filter_strides_;
const std::array<ck::index_t, NDimSpatial>& input_left_pads_;
const std::array<ck::index_t, NDimSpatial>& input_right_pads_;
std::array<index_t, NDimSpatial + 3> in_g_n_c_wis_lengths_;
std::array<index_t, NDimSpatial + 3> in_g_n_c_wis_strides_;
std::array<index_t, NDimSpatial + 3> wei_g_k_c_xs_lengths_;
std::array<index_t, NDimSpatial + 3> wei_g_k_c_xs_strides_;
std::array<index_t, NDimSpatial + 3> out_g_n_k_wos_lengths_;
std::array<index_t, NDimSpatial + 3> out_g_n_k_wos_strides_;
std::array<ck::index_t, NDimSpatial> conv_filter_strides_;
std::array<ck::index_t, NDimSpatial> conv_filter_dilations_;
std::array<ck::index_t, NDimSpatial> input_left_pads_;
std::array<ck::index_t, NDimSpatial> input_right_pads_;
const index_t k_batch_;
};
@@ -395,13 +362,13 @@ struct DeviceGroupedConvBwdWeightNaive
{
using Argument = DeviceOp::Argument;
void ShowInfo(const Argument& arg)
void ShowInfo(const Argument&)
{
}
index_t CalculateGridSize(const Argument& arg)
{
return arg.Conv_G_;
return arg.in_g_n_c_wis_lengths_[0];;
}
float RunGemmV3(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
@@ -411,33 +378,19 @@ struct DeviceGroupedConvBwdWeightNaive
float ave_time = 0;
constexpr index_t minimum_occupancy =
BlkGemmPipeSched == BlockGemmPipelineScheduler::Intrawave ? 1 : 2;
constexpr index_t minimum_occupancy = 2;
// BlkGemmPipeSched == BlockGemmPipelineScheduler::Intrawave ? 1 : 2;
constexpr index_t BlockSize = 64;
const auto kernel = kernel_grouped_conv_bwd_weight_naive<
GridwiseGemm,
remove_reference_t<DeviceOp::AGridDesc_K0_M_K1>,
remove_reference_t<DeviceOp::BGridDesc_K0_N_K1>,
remove_reference_t<
DeviceOp::CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock>,
ComputePtrOffsetOfStridedBatch<I1, I1, I0>,
NumGroupsToMerge,
true,
InMemoryDataOperationEnum::Set,
minimum_occupancy>;
Argument, minimum_occupancy>;
ave_time += launch_and_time_kernel(
stream_config,
kernel,
dim3(gdx, gdy, gdz),
dim3(gdx),
dim3(BlockSize),
0,
gemm_arg,
arg.a_grid_desc_k0_m_k1_,
arg.b_grid_desc_k0_n_k1_,
arg.c_grid_desc_mblock_mperblock_nblock_nperblock_,
arg.compute_ptr_offset_of_batch_,
num_k_per_block);
&arg);
return ave_time;
}
@@ -462,7 +415,7 @@ struct DeviceGroupedConvBwdWeightNaive
return true;
}
static bool IsSupportedArgument(const Argument& arg)
static bool IsSupportedArgument(const Argument&)
{
return true;
}
@@ -594,6 +547,23 @@ struct DeviceGroupedConvBwdWeightNaive
}
}
using ALayout = ck::tensor_layout::convolution::NHWGC;
using BLayout = ck::tensor_layout::convolution::GKYXC;
using ELayout = ck::tensor_layout::convolution::NHWGK;
template <ck::index_t NDimSpatial>
using DeviceConvBwdWeightInstance =
ck::tensor_operation::device::DeviceGroupedConvBwdWeightNaive<NDimSpatial,
ALayout,
BLayout,
ELayout,
InDataType,
WeiDataType,
OutDataType,
InElementOp,
WeiElementOp,
OutElementOp>;
template <ck::index_t NDimSpatial>
using HostConvBwdWeightInstance = ck::tensor_operation::host::ReferenceConvBwdWeight<NDimSpatial,
InDataType,