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

Add number of groups to merge to ck tile grouped gemm example.

Browse Source
This commit is contained in:
Ville Pietilä
2025-09-04 14:24:23 +00:00
parent e8d0c04a1b
commit 61b3c96273
8 changed files with 63 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp
View File

@@ -20,6 +20,7 @@ template <ck_tile::index_t NDimSpatial,
typename InLayout,
typename WeiLayout,
typename OutLayout,
ck_tile::index_t NumGroupsToMerge,
typename DsDataType = ck_tile::tuple<>,
typename DsLayout = ck_tile::tuple<>,
typename CDEElementWise = ck_tile::element_wise::PassThrough>
@@ -53,7 +54,7 @@ float grouped_conv_bwd_weight(const ck_tile::GroupedConvBwdWeightHostArgs& args,
constexpr auto ConvSpec = ck_tile::ConvolutionSpecialization::Default;
using TilePartitioner = ck_tile::GemmTile1DPartitioner<CodegenShape>;
using GroupedConvTraitsType =
ck_tile::GroupedConvTraits<NDimSpatial, ConvSpec, InLayout, WeiLayout, DsLayout, OutLayout>;
ck_tile::GroupedConvTraits<NDimSpatial, ConvSpec, InLayout, WeiLayout, DsLayout, OutLayout, NumGroupsToMerge>;
using CodegenPipelineProblem =
ck_tile::GemmPipelineProblem<InDataType,
WeiDataType,
@@ -141,7 +142,7 @@ float grouped_conv_bwd_weight(const ck_tile::GroupedConvBwdWeightHostArgs& args,
#include "run_grouped_convolution_bwd_weight_example.inc"
template <typename InPrecType, typename WeiPrecType = InPrecType, typename OutPrecType = InPrecType>
template <typename InPrecType, ck_tile::index_t NumGroupsToMerge = 1, typename WeiPrecType = InPrecType, typename OutPrecType = InPrecType>
int run_grouped_conv_bwd_weight_example_prec_type(
std::string in_layout, std::string wei_layout, std::string out_layout, int argc, char* argv[])
{
@@ -162,7 +163,8 @@ int run_grouped_conv_bwd_weight_example_prec_type(
return run_grouped_conv_bwd_weight_example_with_layouts<ck_tile::number<1>{},
InPrecType,
WeiPrecType,
OutPrecType>(
OutPrecType,
NumGroupsToMerge>(
argc, argv, NWGC{}, GKXC{}, NWGK{});
}
else if(in_layout == "NHWGC" && wei_layout == "GKYXC" && out_layout == "NHWGK")
@@ -170,7 +172,8 @@ int run_grouped_conv_bwd_weight_example_prec_type(
return run_grouped_conv_bwd_weight_example_with_layouts<ck_tile::number<2>{},
InPrecType,
WeiPrecType,
OutPrecType>(
OutPrecType,
NumGroupsToMerge>(
argc, argv, NHWGC{}, GKYXC{}, NHWGK{});
}
else if(in_layout == "NDHWGC" && wei_layout == "GKZYXC" && out_layout == "NDHWGK")
@@ -178,7 +181,8 @@ int run_grouped_conv_bwd_weight_example_prec_type(
return run_grouped_conv_bwd_weight_example_with_layouts<ck_tile::number<3>{},
InPrecType,
WeiPrecType,
OutPrecType>(
OutPrecType,
NumGroupsToMerge>(
argc, argv, NDHWGC{}, GKZYXC{}, NDHWGK{});
}
else
@@ -187,6 +191,36 @@ int run_grouped_conv_bwd_weight_example_prec_type(
}
}
template <typename InPrecType>
int run(const std::string& in_layout,
const std::string& wei_layout,
const std::string& out_layout,
int num_groups_to_merge,
int argc,
char* argv[])
{
if (num_groups_to_merge == 1)
{
return run_grouped_conv_bwd_weight_example_prec_type<InPrecType, 1>(in_layout, wei_layout, out_layout, argc, argv);
}
else if (num_groups_to_merge == 2)
{
return run_grouped_conv_bwd_weight_example_prec_type<InPrecType, 2>(in_layout, wei_layout, out_layout, argc, argv);
}
else if (num_groups_to_merge == 4)
{
return run_grouped_conv_bwd_weight_example_prec_type<InPrecType, 4>(in_layout, wei_layout, out_layout, argc, argv);
}
else if (num_groups_to_merge == 8)
{
return run_grouped_conv_bwd_weight_example_prec_type<InPrecType, 8>(in_layout, wei_layout, out_layout, argc, argv);
}
else
{
throw std::runtime_error("Unsupported number of groups to merge!");
}
}
int run_grouped_conv_bwd_weight_example(int argc, char* argv[])
{
auto [result, arg_parser] = create_args(argc, argv);
@@ -197,16 +231,15 @@ int run_grouped_conv_bwd_weight_example(int argc, char* argv[])
std::string in_layout = arg_parser.get_str("in_layout");
std::string wei_layout = arg_parser.get_str("wei_layout");
std::string out_layout = arg_parser.get_str("out_layout");
ck_tile::index_t num_groups_to_merge = arg_parser.get_int("num_groups_to_merge");
if(data_type == "fp16")
{
return run_grouped_conv_bwd_weight_example_prec_type<ck_tile::half_t>(
in_layout, wei_layout, out_layout, argc, argv);
return run<ck_tile::half_t>(in_layout, wei_layout, out_layout, num_groups_to_merge, argc, argv);
}
else if(data_type == "bf16")
{
return run_grouped_conv_bwd_weight_example_prec_type<ck_tile::bf16_t>(
in_layout, wei_layout, out_layout, argc, argv);
return run<ck_tile::bf16_t>(in_layout, wei_layout, out_layout, num_groups_to_merge, argc, argv);
}
else
{

3
example/ck_tile/20_grouped_convolution/grouped_convolution_utils.hpp
View File

@@ -121,7 +121,8 @@ auto create_args(int argc, char* argv[])
.insert("timer", "gpu", "gpu:gpu timer, cpu:cpu timer")
.insert("split_k", "1", "splitK value")
.insert("init", "0", "0:random, 1:linear, 2:constant(1)")
.insert("json", "0", "0: No Json, 1: Dump Results in Json format");
.insert("json", "0", "0: No Json, 1: Dump Results in Json format")
.insert("num_groups_to_merge", "1", "Number of groups to merge for grouped convolution");
bool result = arg_parser.parse(argc, argv);
return std::make_tuple(result, arg_parser);

11
example/ck_tile/20_grouped_convolution/run_grouped_convolution_bwd_weight_example.inc
View File

@@ -9,7 +9,8 @@ template <ck_tile::index_t NDimSpatial,
typename OutDataType,
typename InLayout,
typename WeiLayout,
typename OutLayout>
typename OutLayout,
ck_tile::index_t NumGroupsToMerge>
float invoke_grouped_conv_bwd_weight(ck_tile::GroupedConvBwdWeightHostArgs& args,
int n_warmup,
int n_repeat)
@@ -21,7 +22,8 @@ float invoke_grouped_conv_bwd_weight(ck_tile::GroupedConvBwdWeightHostArgs& args
OutDataType,
InLayout,
WeiLayout,
OutLayout>(
OutLayout,
NumGroupsToMerge>(
args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat});
std::size_t flop = args.GetFlops();
@@ -39,6 +41,7 @@ template <ck_tile::index_t NDimSpatial,
typename InDataType,
typename WeiDataType = InDataType,
typename OutDataType = InDataType,
ck_tile::index_t NumGroupsToMerge,
typename InLayout,
typename WeiLayout,
typename OutLayout>
@@ -134,6 +137,7 @@ int run_grouped_conv_bwd_weight_example_with_layouts(
std::cout << "input: " << input.mDesc << std::endl;
std::cout << "weight: " << weight.mDesc << std::endl;
std::cout << "output: " << output.mDesc << std::endl;
std::cout << "number of groups to merge: " << NumGroupsToMerge << std::endl;
invoke_grouped_conv_bwd_weight<NDimSpatial,
InDataType,
@@ -142,7 +146,8 @@ int run_grouped_conv_bwd_weight_example_with_layouts(
OutDataType,
InLayout,
WeiLayout,
OutLayout>(args, n_warmup, n_repeat);
OutLayout,
NumGroupsToMerge>(args, n_warmup, n_repeat);
weight_dev_buf.FromDevice(weight.data());
bool pass = true;

8
include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp
View File

@@ -23,7 +23,8 @@ struct GroupedConvBwdWeightKernelArgs
using ConvToGemmTransformer =
TransformConvBwdWeightToGemm<GroupedConvTraitsType_::NDimSpatial,
GroupedConvTraitsType_::ConvSpecialization>;
GroupedConvTraitsType_::ConvSpecialization,
GroupedConvTraitsType_::NumGroupsToMerge>;
static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor;
template <
@@ -314,10 +315,9 @@ struct GroupedConvBwdWeightKernelArgs
/// the policy is responsible for definition of all necessary data layouts and thread's
/// work distribution.
///
/// @tparam GroupedConvTraitsType_ The type of class providing traits for grouped convolution.
/// @tparam GroupedConvTraitsType_ The type of class providing traits for grouped convolution.
/// @tparam TilePartitioner_ The type of class providing mapping of workgroup index into
/// the
/// output data tile to be calculated. It determines the
/// the output data tile to be calculated. It determines the
/// workgroup to data relationship (or in other words - which
/// data would be processed and calculated by which workgroup).
/// @tparam GemmPipeline_ The type of class which provides the core part of matrix

3
include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_forward_kernel.hpp
View File

@@ -23,7 +23,8 @@ struct GroupedConvFwdKernelArgs
using ConvToGemmFwdTransformer =
TransformConvFwdToGemm<GroupedConvTraitsType_::NDimSpatial,
GroupedConvTraitsType_::ConvSpecialization>;
GroupedConvTraitsType_::ConvSpecialization,
GroupedConvTraitsType_::NumGroupsToMerge>;
static constexpr index_t NumDTensor = GroupedConvTraitsType_::NumDTensor;
template <

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

@@ -49,7 +49,8 @@ template <index_t NDimSpatial_,
typename InLayout_,
typename WeiLayout_,
typename DsLayout_,
typename OutLayout_>
typename OutLayout_,
index_t NumGroupsToMerge_ = 1>
struct GroupedConvTraits
{
private:
@@ -60,7 +61,7 @@ struct GroupedConvTraits
}
public:
static constexpr index_t NumGroupsToMerge = 1;
static constexpr index_t NumGroupsToMerge = NumGroupsToMerge_;
static constexpr index_t NDimSpatial = NDimSpatial_;
static constexpr ConvolutionSpecialization ConvSpecialization = ConvSpecialization_;
using InLayout = InLayout_;

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

@@ -10,10 +10,10 @@ namespace ck_tile {
template <index_t NDimSpatial,
ConvolutionSpecialization ConvolutionSpecialization,
index_t NumGroupsToMerge = 1,
bool SplitN = false,
typename ADataType = float,
typename CDataType = float,
index_t NumGroupsToMerge = 1,
typename IndexType = index_t>
struct TransformConvBwdWeightToGemm
{

2
include/ck_tile/ops/grouped_convolution/utils/transform_conv_fwd_to_gemm.hpp
View File

@@ -10,10 +10,10 @@ namespace ck_tile {
template <index_t NDimSpatial,
ConvolutionSpecialization ConvSpecialization,
index_t NumGroupsToMerge = 1,
bool SplitN = false,
typename ADataType = float,
typename CDataType = float,
index_t NumGroupsToMerge = 1,
typename IndexType = index_t>
struct TransformConvFwdToGemm
{