This commit is contained in:
Bartlomiej Kocot
2025-09-29 22:54:04 +00:00
parent bebf0e9d15
commit dbc623f455
6 changed files with 41 additions and 50 deletions

View File

@@ -36,8 +36,8 @@ struct GroupedConvolutionBackwardDataInvoker
constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
constexpr ck_tile::index_t VectorSizeA = 1;
constexpr ck_tile::index_t VectorSizeB = 1;
constexpr ck_tile::index_t VectorSizeA = 8;
constexpr ck_tile::index_t VectorSizeB = 8;
constexpr ck_tile::index_t VectorSizeC = 8;
// Implicit GEMM Traits

View File

@@ -35,8 +35,8 @@ struct GroupedConvolutionBackwardWeightInvoker
constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
constexpr ck_tile::index_t VectorSizeA = 1;
constexpr ck_tile::index_t VectorSizeB = 1;
constexpr ck_tile::index_t VectorSizeA = 8;
constexpr ck_tile::index_t VectorSizeB = 8;
constexpr ck_tile::index_t VectorSizeC = 8;
// Implicit GEMM Traits

View File

@@ -37,8 +37,8 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
constexpr ck_tile::index_t N_Warp_Tile = GemmWarpConfig::N_Warp_Tile;
constexpr ck_tile::index_t K_Warp_Tile = GemmWarpConfig::K_Warp_Tile;
constexpr ck_tile::index_t VectorSizeA = 1;
constexpr ck_tile::index_t VectorSizeB = 1;
constexpr ck_tile::index_t VectorSizeA = 8;
constexpr ck_tile::index_t VectorSizeB = 8;
constexpr ck_tile::index_t VectorSizeC = 1;
// Implicit GEMM Traits

View File

@@ -109,7 +109,7 @@ struct GroupedConvBwdDataKernelArgs
GroupedConvTraitsType_::NDimSpatial>(1);
a_grid_descs_m_k[gemm_count] = grid_descs.at(number<0>{});
b_grid_descs_n_k[gemm_count] = grid_descs.at(number<1>{});
b_grid_descs_k_n[gemm_count] = grid_descs.at(number<1>{});
c_grid_descs_m_n[gemm_count] = grid_descs.at(number<2>{});
const index_t grid_size_grp =
@@ -225,7 +225,7 @@ struct GroupedConvBwdDataKernelArgs
GroupedConvTraitsType_::NDimSpatial>(1);
a_grid_descs_m_k[gemm_count] = grid_descs.at(number<0>{});
b_grid_descs_n_k[gemm_count] = grid_descs.at(number<1>{});
b_grid_descs_k_n[gemm_count] = grid_descs.at(number<1>{});
c_grid_descs_m_n[gemm_count] = grid_descs.at(number<2>{});
const index_t grid_size_grp =
@@ -357,7 +357,7 @@ struct GroupedConvBwdDataKernelArgs
GroupedConvTraitsType_::NDimSpatial>(1);
a_grid_descs_m_k[gemm_count] = grid_descs.at(number<0>{});
b_grid_descs_n_k[gemm_count] = grid_descs.at(number<1>{});
b_grid_descs_k_n[gemm_count] = grid_descs.at(number<1>{});
c_grid_descs_m_n[gemm_count] = grid_descs.at(number<2>{});
const index_t grid_size_grp =
@@ -416,7 +416,7 @@ struct GroupedConvBwdDataKernelArgs
const void* wei_ptr;
array<AGridDescMK, MaxGroupedGemmGroupsNum> a_grid_descs_m_k;
array<BGridDescNK, MaxGroupedGemmGroupsNum> b_grid_descs_n_k;
array<BGridDescNK, MaxGroupedGemmGroupsNum> b_grid_descs_k_n;
array<CGridDescMN, MaxGroupedGemmGroupsNum> c_grid_descs_m_n;
array<index_t, MaxGroupedGemmGroupsNum> block_starts;
@@ -471,10 +471,6 @@ template <typename GroupedConvTraitsType_,
typename EpiloguePipeline_>
struct GroupedConvolutionBackwardDataKernel
{
// Todo: Enable Vector Load Size > 1
static_assert(GroupedConvTraitsType_::VectorSizeA == 1 &&
GroupedConvTraitsType_::VectorSizeB == 1);
static constexpr index_t NDimSpatial = GroupedConvTraitsType_::NDimSpatial_;
static constexpr ConvolutionSpecialization ConvSpecialization =
GroupedConvTraitsType_::ConvSpecialization;
@@ -516,13 +512,10 @@ struct GroupedConvolutionBackwardDataKernel
static_assert(GemmPipeline::kPadM && GemmPipeline::kPadN && GemmPipeline::kPadK,
"Not supported!");
static_assert(std::is_same_v<GemmALayout, tensor_layout::gemm::RowMajor>, "Not supported!");
static_assert(std::is_same_v<GemmBLayout, tensor_layout::gemm::ColumnMajor>, "Not supported!");
// TODO: Change to and enable vector load
// static_assert(std::is_same_v<GemmALayout, tensor_layout::gemm::RowMajor>,
// "Not supported A GEMM layout!");
// static_assert(std::is_same_v<GemmBLayout, tensor_layout::gemm::RowMajor>,
// "Not supported B GEMM layout!");
static_assert(std::is_same_v<GemmALayout, tensor_layout::gemm::RowMajor>,
"Not supported A GEMM layout!");
static_assert(std::is_same_v<GemmBLayout, tensor_layout::gemm::RowMajor>,
"Not supported B GEMM layout!");
static_assert(std::is_same_v<GemmCLayout, tensor_layout::gemm::RowMajor>,
"Not supported C GEMM layout!");
@@ -703,7 +696,7 @@ struct GroupedConvolutionBackwardDataKernel
const auto& b_tensor_view = [&]() {
return make_tensor_view<address_space_enum::global>(
b_ptr,
kargs.b_grid_descs_n_k[group_id]); // B: weight
kargs.b_grid_descs_k_n[group_id]); // B: weight
}();
const auto& c_tensor_view = [&]() {
@@ -742,8 +735,10 @@ struct GroupedConvolutionBackwardDataKernel
const auto& b_pad_view = [&]() {
const auto& b_tensor_view = views.at(I1);
return pad_tensor_view(b_tensor_view,
make_tuple(number<TilePartitioner::NPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
make_tuple(
number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::NPerBlock>{}
),
sequence<true, true>{});
}();
@@ -788,9 +783,11 @@ struct GroupedConvolutionBackwardDataKernel
const auto& b_block_window = [&]() {
return make_tile_window(b_pad_view,
make_tuple(number<TilePartitioner::NPerBlock>{},
number<TilePartitioner::KPerBlock>{}),
{i_n, i_k});
make_tuple(
number<TilePartitioner::KPerBlock>{},
number<TilePartitioner::NPerBlock>{}
),
{i_k, i_n});
}();
const auto ds_block_window = generate_tuple(

View File

@@ -81,21 +81,15 @@ struct GroupedConvTraits
TileGemmTraits<true,
true,
true,
ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::ColumnMajor,
// TODO: Change to and enable vector load
// ck_tile::tensor_layout::gemm::RowMajor,
// ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::RowMajor>;
using GroupedConvImplicitGemmTraitsBwdWeight =
TileGemmTraits<true,
true,
true,
ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::ColumnMajor,
// TODO: Change to and enable vector load
// ck_tile::tensor_layout::gemm::ColumnMajor,
// ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::ColumnMajor,
ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::tensor_layout::gemm::RowMajor>;
static constexpr ck_tile::index_t VectorSizeA = VectorSizeA_;
static constexpr ck_tile::index_t VectorSizeB = VectorSizeB_;

View File

@@ -502,7 +502,7 @@ struct TransformConvBwdDataToGemm
// TODO Add support for NumGroupsToMerge > 1
return make_naive_tensor_descriptor(make_tuple(N_, Hi_, Wi_, C_),
make_tuple(NStride, HiStride, WiStride, CStride),
number<VectorSizeB>{},
number<VectorSizeC>{},
I1);
}
@@ -512,7 +512,7 @@ struct TransformConvBwdDataToGemm
// GKYXC
return make_naive_tensor_descriptor(make_tuple(K_, Y_, X_, C_),
make_tuple(C_ * X_ * Y_, C_ * X_, C_, I1),
number<VectorSizeC>{},
number<VectorSizeB>{},
I1);
}
@@ -547,7 +547,7 @@ struct TransformConvBwdDataToGemm
return make_naive_tensor_descriptor(
make_tuple(N_, Di_, Hi_, Wi_, C_),
make_tuple(NStride, DiStride, HiStride, WiStride, CStride),
number<VectorSizeB>{},
number<VectorSizeC>{},
I1);
}
@@ -558,7 +558,7 @@ struct TransformConvBwdDataToGemm
return make_naive_tensor_descriptor(
make_tuple(K_, Z_, Y_, X_, C_),
make_tuple(C_ * X_ * Y_ * Z_, C_ * X_ * Y_, C_ * X_, C_, I1),
number<VectorSizeC>{},
number<VectorSizeB>{},
I1);
}
// TODO: implement ck_tile::tensor_layout::convolution that describe packed/strided dimemsion as
@@ -637,12 +637,12 @@ struct TransformConvBwdDataToGemm
make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}, sequence<3>{}),
make_tuple(sequence<0>{}, sequence<1>{}, sequence<>{}, sequence<2>{}));
const auto wei_gemmn_gemmkraw_grid_desc =
const auto wei_gemmkraw_gemmn_grid_desc =
transform_tensor_descriptor(wei_k_xdotslice_c_grid_desc,
make_tuple(make_merge_transform(make_tuple(XDotSlice, K_)),
make_pass_through_transform(C_)),
make_tuple(sequence<1, 0>{}, sequence<2>{}),
make_tuple(sequence<1>{}, sequence<0>{}));
make_tuple(sequence<0>{}, sequence<1>{}));
// c: input
const auto in_n_wip_c_grid_desc = transform_tensor_descriptor(
@@ -679,7 +679,7 @@ struct TransformConvBwdDataToGemm
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tuple(out_gemmm_gemmkraw_grid_desc,
wei_gemmn_gemmkraw_grid_desc,
wei_gemmkraw_gemmn_grid_desc,
in_gemmmraw_gemmnraw_grid_desc);
}
@@ -792,12 +792,12 @@ struct TransformConvBwdDataToGemm
sequence<>{},
sequence<3>{}));
const auto wei_gemmn_gemmkraw_grid_desc = transform_tensor_descriptor(
const auto wei_gemmkraw_gemmn_grid_desc = transform_tensor_descriptor(
wei_k_ydotslice_xdotslice_c_grid_desc,
make_tuple(make_merge_transform(make_tuple(YDotSlice, XDotSlice, K_)),
make_pass_through_transform(C_)),
make_tuple(sequence<1, 2, 0>{}, sequence<3>{}),
make_tuple(sequence<1>{}, sequence<0>{}));
make_tuple(sequence<0>{}, sequence<1>{}));
// c: input
const auto in_n_hip_wip_c_grid_desc = transform_tensor_descriptor(
@@ -849,7 +849,7 @@ struct TransformConvBwdDataToGemm
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tuple(out_gemmm_gemmkraw_grid_desc,
wei_gemmn_gemmkraw_grid_desc,
wei_gemmkraw_gemmn_grid_desc,
in_gemmmraw_gemmnraw_grid_desc);
}
@@ -994,12 +994,12 @@ struct TransformConvBwdDataToGemm
sequence<>{},
sequence<4>{}));
const auto wei_gemmn_gemmkraw_grid_desc = transform_tensor_descriptor(
const auto wei_gemmkraw_gemmn_grid_desc = transform_tensor_descriptor(
wei_k_ydotslice_xdotslice_c_grid_desc,
make_tuple(make_merge_transform(make_tuple(ZDotSlice, YDotSlice, XDotSlice, K_)),
make_pass_through_transform(C_)),
make_tuple(sequence<1, 2, 3, 0>{}, sequence<4>{}),
make_tuple(sequence<1>{}, sequence<0>{}));
make_tuple(sequence<0>{}, sequence<1>{}));
// c: input
const auto in_n_hip_wip_c_grid_desc = transform_tensor_descriptor(
@@ -1064,7 +1064,7 @@ struct TransformConvBwdDataToGemm
make_tuple(sequence<0>{}, sequence<1>{}));
return make_tuple(out_gemmm_gemmkraw_grid_desc,
wei_gemmn_gemmkraw_grid_desc,
wei_gemmkraw_gemmn_grid_desc,
in_gemmmraw_gemmnraw_grid_desc);
}