diff --git a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp index 2571961c6b..3d0d504427 100644 --- a/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp +++ b/example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight.cpp @@ -215,9 +215,21 @@ int run(const std::string& in_layout, { return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); } + else if (num_groups_to_merge == 16) + { + return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); + } + else if (num_groups_to_merge == 32) + { + return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); + } + else if (num_groups_to_merge == 64) + { + return run_grouped_conv_bwd_weight_example_prec_type(in_layout, wei_layout, out_layout, argc, argv); + } else { - throw std::runtime_error("Unsupported number of groups to merge!"); + throw std::runtime_error("Unsupported number of groups to merge! The number of groups should be a power of two and at most 64."); } } diff --git a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp index b571a61633..3e4eec13a9 100644 --- a/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp +++ b/include/ck_tile/ops/grouped_convolution/kernel/grouped_convolution_backward_weight_kernel.hpp @@ -82,9 +82,11 @@ struct GroupedConvBwdWeightKernelArgs b_grid_desc_n_k = grid_descs.at(number<1>{}); c_grid_desc_m_n = grid_descs.at(number<2>{}); - group_stride_a = args.K_; // A: Out NWGK - group_stride_b = args.C_; // B: In NWGC - group_stride_c = args.K_ * args.C_ * // C: Wei GKXC + using Gm = GroupedConvTraitsType_::NumGroupsToMerge; + + group_stride_a = args.K_ * Gm; // A: Out NWGK + group_stride_b = args.C_ * Gm; // B: In NWGC + group_stride_c = args.K_ * args.C_ * Gm * // C: Wei GKXC std::accumulate(args.filter_spatial_lengths_.begin(), args.filter_spatial_lengths_.end(), 1, @@ -93,7 +95,7 @@ struct GroupedConvBwdWeightKernelArgs GemmM = a_grid_desc_m_k.get_length(number<0>{}); GemmN = b_grid_desc_n_k.get_length(number<0>{}); GemmK = a_grid_desc_m_k.get_length(number<1>{}); - GemmBatch = args.G_; + GemmBatch = args.G_ / Gm; } template < @@ -158,9 +160,11 @@ struct GroupedConvBwdWeightKernelArgs b_grid_desc_n_k = grid_descs.at(number<1>{}); c_grid_desc_m_n = grid_descs.at(number<2>{}); - group_stride_a = args.K_; // A: Out NHWGK - group_stride_b = args.C_; // B: In NHWGC - group_stride_c = args.K_ * args.C_ * // C: Wei GKYXC + using Gm = GroupedConvTraitsType_::NumGroupsToMerge; + + group_stride_a = args.K_ * Gm; // A: Out NHWGK + group_stride_b = args.C_ * Gm; // B: In NHWGC + group_stride_c = args.K_ * args.C_ * Gm * // C: Wei GKYXC std::accumulate(args.filter_spatial_lengths_.begin(), args.filter_spatial_lengths_.end(), 1, @@ -169,7 +173,7 @@ struct GroupedConvBwdWeightKernelArgs GemmM = a_grid_desc_m_k.get_length(number<0>{}); GemmN = b_grid_desc_n_k.get_length(number<0>{}); GemmK = a_grid_desc_m_k.get_length(number<1>{}); - GemmBatch = args.G_; + GemmBatch = args.G_ / Gm; } template < @@ -241,9 +245,11 @@ struct GroupedConvBwdWeightKernelArgs b_grid_desc_n_k = grid_descs.at(number<1>{}); c_grid_desc_m_n = grid_descs.at(number<2>{}); - group_stride_a = args.K_; // A: Out NDHWGK - group_stride_b = args.C_; // B: In NDHWGC - group_stride_c = args.K_ * args.C_ * // C: wEI GKZYXC + using Gm = GroupedConvTraitsType_::NumGroupsToMerge; + + group_stride_a = args.K_ * Gm; // A: Out NDHWGK + group_stride_b = args.C_ * Gm; // B: In NDHWGC + group_stride_c = args.K_ * args.C_ * Gm * // C: wEI GKZYXC std::accumulate(args.filter_spatial_lengths_.begin(), args.filter_spatial_lengths_.end(), 1, @@ -252,7 +258,7 @@ struct GroupedConvBwdWeightKernelArgs GemmM = a_grid_desc_m_k.get_length(number<0>{}); GemmN = b_grid_desc_n_k.get_length(number<0>{}); GemmK = a_grid_desc_m_k.get_length(number<1>{}); - GemmBatch = args.G_; + GemmBatch = args.G_ / Gm; } using ABCGridDescs = remove_cvref_t<