Skip to content

fix(ck): convolution edge cases large tensors #11754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
EnricoDeg merged 12 commits into from
Sep 30, 2026
+318 −93
Merged

fix(ck): convolution edge cases large tensors #11754

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
fae8f6b
Workaround edge cases
EnricoDeg Sep 1, 2026
205d893
Fwd fix for edge cases without changing splitN value
EnricoDeg Sep 3, 2026
44c727e
Change flaky instance fwd
EnricoDeg Sep 3, 2026
7bfcb16
Fix bwd without changing splitN value
EnricoDeg Sep 3, 2026
aed078e
Check tensor size overflow with index_t strides (bwd weight)
EnricoDeg Sep 4, 2026
f6f0d5d
Fix compilation errror
EnricoDeg Sep 4, 2026
34db865
Cast explicitly
EnricoDeg Sep 6, 2026
a0aa728
Improve logging message
EnricoDeg Sep 7, 2026
000b504
Address review
EnricoDeg Sep 10, 2026
42a5534
Fix iteration range bwd data wmma v3
EnricoDeg Sep 17, 2026
f279f14
Fix types in bwd weight checks
EnricoDeg Sep 17, 2026
a15ed44
Merge branch 'develop' into users/enricodeg/ck/fix-edge-cases
EnricoDeg Sep 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 33 additions & 6 deletions ...ensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -576,7 +576,14 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
static bool IsSupportedArgument(const Argument& arg)
{
if(arg.stride_overflow)
{
if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
{
std::cout << "Stride overflow!" << " In " << __FILE__ << ":" << __LINE__
<< ", in function: " << __func__ << std::endl;
}
return false;
}

if(arg.k_batch_ != 1)
{
Expand Down Expand Up @@ -744,8 +751,17 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
const AElementwiseOp& a_element_op,
const BElementwiseOp& b_element_op,
const CDEElementwiseOp& cde_element_op,
const ck::index_t split_k = 1)
const index_t split_k = 1)
{
bool ds_ovf = false;
static_for<0, NumDTensor, 1>{}([&](auto i) {
using DDataType = remove_cvref_t<tuple_element_t<i.value, DsDataType>>;
ds_ovf |= tensor_exceeds_2gb<DDataType>(ds_g_n_c_wis_lengths[i]);
});
const bool stride_ovf = tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths) ||
tensor_exceeds_2gb<BDataType>(b_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<EDataType>(e_g_n_c_wis_lengths) || ds_ovf;

return Argument{p_a,
p_b,
p_ds,
Expand All @@ -765,7 +781,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
a_element_op,
b_element_op,
cde_element_op,
split_k};
split_k,
stride_ovf};
}

static auto MakeArgument(
Expand All @@ -790,7 +807,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
const AElementwiseOp& a_element_op,
const BElementwiseOp& b_element_op,
const CDEElementwiseOp& cde_element_op,
const ck::index_t split_k = 1)
const index_t split_k = 1)
{
bool ds_ovf = false;
static_for<0, NumDTensor, 1>{}([&](auto i) {
Expand Down Expand Up @@ -876,8 +893,17 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
const AElementwiseOp& a_element_op,
const BElementwiseOp& b_element_op,
const CDEElementwiseOp& cde_element_op,
const ck::index_t split_k = 1) override
const index_t split_k = 1) override
{
bool ds_ovf = false;
static_for<0, NumDTensor, 1>{}([&](auto i) {
using DDataType = remove_cvref_t<tuple_element_t<i.value, DsDataType>>;
ds_ovf |= tensor_exceeds_2gb<DDataType>(ds_g_n_c_wis_lengths[i]);
});
const bool stride_ovf = tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths) ||
tensor_exceeds_2gb<BDataType>(b_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<EDataType>(e_g_n_c_wis_lengths) || ds_ovf;

return std::make_unique<Argument>(p_a,
p_b,
p_ds,
Expand All @@ -897,7 +923,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
a_element_op,
b_element_op,
cde_element_op,
split_k);
split_k,
stride_ovf);
}

std::unique_ptr<BaseArgument> MakeArgumentPointer(
Expand All @@ -922,7 +949,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffle
const AElementwiseOp& a_element_op,
const BElementwiseOp& b_element_op,
const CDEElementwiseOp& cde_element_op,
const ck::index_t split_k = 1) override
const index_t split_k = 1) override
{
bool ds_ovf = false;
static_for<0, NumDTensor, 1>{}([&](auto i) {
Expand Down
31 changes: 31 additions & 0 deletions ...or_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1465,7 +1465,14 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3
static bool IsSupportedArgument(const Argument& arg)
{
if(arg.stride_overflow)
{
if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
{
std::cout << "Stride overflow!" << " In " << __FILE__ << ":" << __LINE__
<< ", in function: " << __func__ << std::endl;
}
return false;
}

if(!ck::is_gfx11_supported() && !ck::is_gfx12_supported())
{
Expand Down Expand Up @@ -1819,6 +1826,30 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3
}
}

// check descriptors sizes
bool is_size_valid = true;
for(std::size_t i = 0; i < arg.gemm_kernel_args_.size(); i++)
Comment thread
EnricoDeg marked this conversation as resolved.
Resolved
{
static_for<0, NumDTensor, 1>{}([&](auto j) {
using DDataType = remove_cvref_t<tuple_element_t<j.value, DsDataType>>;
is_size_valid &=
!descriptor_exceeds_2gb<DDataType>(arg.ds_grid_desc_m_n_container_[i][j]);
});

is_size_valid &= !descriptor_exceeds_2gb<ADataType>(arg.a_grid_desc_m_k_container_[i]);
is_size_valid &= !descriptor_exceeds_2gb<BDataType>(arg.b_grid_desc_n_k_container_[i]);
is_size_valid &= !descriptor_exceeds_2gb<EDataType>(arg.e_grid_desc_m_n_container_[i]);
}
if(!is_size_valid)
{
if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
{
std::cout << "Large tensor case!" << " In " << __FILE__ << ":" << __LINE__
<< ", in function: " << __func__ << std::endl;
}
return false;
}

// Check gridwise gemm validity
// Create dummy values for Ds pointers and strides
std::array<const void*, NumDTensor> p_ds_grid_dummy;
Expand Down
28 changes: 27 additions & 1 deletion ...sor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1507,8 +1507,10 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
bool used_flat_desc = false;
if constexpr(NDimSpatial == 2 && !CTranspose && NumDTensor == 0)
{
const index_t ConvK = arg.b_g_k_c_xs_lengths_[1];
if(arg.num_group_ == 1 && arg.k_batch_ == 1 && arg.gemms_count_ == 1 &&
!arg.flat_a_container_.empty())
!arg.flat_a_container_.empty() && arg.num_workgroups_per_Conv_N_ == 1 &&
ConvK % AK1 == 0 && ConvK % BK1 == 0)
{
used_flat_desc = true;
const index_t flat_idx = gemm_set_id;
Expand Down Expand Up @@ -1990,6 +1992,30 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1
return false;
}

// check descriptors sizes
Comment thread
chris-tsiaousis-hpc marked this conversation as resolved.
Resolved
bool is_size_valid = true;
for(std::size_t i = 0; i < arg.a_grid_desc_m_k_container_.size(); i++)
{
static_for<0, NumDTensor, 1>{}([&](auto j) {
using DDataType = remove_cvref_t<tuple_element_t<j.value, DsDataType>>;
is_size_valid &=
!descriptor_exceeds_2gb<DDataType>(arg.ds_grid_desc_m_n_container_[i][j]);
});

is_size_valid &= !descriptor_exceeds_2gb<ADataType>(arg.a_grid_desc_m_k_container_[i]);
is_size_valid &= !descriptor_exceeds_2gb<BDataType>(arg.b_grid_desc_n_k_container_[i]);
is_size_valid &= !descriptor_exceeds_2gb<EDataType>(arg.e_grid_desc_m_n_container_[i]);
}
if(!is_size_valid)
{
if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
{
std::cout << "Large tensor case!" << " In " << __FILE__ << ":" << __LINE__
<< ", in function: " << __func__ << std::endl;
}
return false;
}

// Gridwise GEMM size
bool isWave64 = get_warp_size() == 64;
for(std::size_t i = 0; i < arg.a_grid_desc_m_k_container_.size(); i++)
Expand Down
26 changes: 23 additions & 3 deletions ...sor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v3.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1265,10 +1265,19 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffleV3
const AElementwiseOp& a_element_op,
const BElementwiseOp& b_element_op,
const CDEElementwiseOp& cde_element_op,
const ck::index_t split_k = 1)
const index_t split_k = 1)
{
if constexpr(!LargeTensors)
{
bool ds_ovf = false;
static_for<0, NumDTensor, 1>{}([&](auto i) {
using DDataType = remove_cvref_t<tuple_element_t<i.value, DsDataType>>;
ds_ovf |= tensor_exceeds_2gb<DDataType>(ds_g_n_c_wis_lengths[i]);
});
const bool stride_ovf = tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths) ||
tensor_exceeds_2gb<BDataType>(b_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<EDataType>(e_g_n_c_wis_lengths) || ds_ovf;

return Argument{p_a,
p_b,
p_ds,
Expand All @@ -1288,7 +1297,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffleV3
a_element_op,
b_element_op,
cde_element_op,
split_k};
split_k,
stride_ovf};
}
else
{
Expand Down Expand Up @@ -1485,6 +1495,15 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffleV3
{
if constexpr(!LargeTensors)
{
bool ds_ovf = false;
static_for<0, NumDTensor, 1>{}([&](auto i) {
using DDataType = remove_cvref_t<tuple_element_t<i.value, DsDataType>>;
ds_ovf |= tensor_exceeds_2gb<DDataType>(ds_g_n_c_wis_lengths[i]);
});
const bool stride_ovf = tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths) ||
tensor_exceeds_2gb<BDataType>(b_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<EDataType>(e_g_n_c_wis_lengths) || ds_ovf;

return std::make_unique<Argument>(p_a,
p_b,
p_ds,
Expand All @@ -1504,7 +1523,8 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffleV3
a_element_op,
b_element_op,
cde_element_op,
split_k);
split_k,
stride_ovf);
}
else
{
Expand Down
44 changes: 30 additions & 14 deletions ...r_operation/gpu/device/impl/device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1169,7 +1169,14 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
static bool IsSupportedArgument(const Argument& arg)
{
if(arg.stride_overflow)
{
if(ck::EnvIsEnabled(CK_ENV(CK_LOGGING)))
{
std::cout << "Stride overflow!" << " In " << __FILE__ << ":" << __LINE__
<< ", in function: " << __func__ << std::endl;
}
return false;
}

const index_t GemmM = arg.a_grid_desc_k0_m_k1_.GetLength(I1);
const index_t GemmN = arg.b_grid_desc_k0_n_k1_.GetLength(I1);
Expand Down Expand Up @@ -1464,15 +1471,18 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
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<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,
const std::array<ck::index_t, NDimSpatial>& input_right_pads,
const std::array<index_t, NDimSpatial>& conv_filter_strides,
const std::array<index_t, NDimSpatial>& conv_filter_dilations,
const std::array<index_t, NDimSpatial>& input_left_pads,
const std::array<index_t, NDimSpatial>& input_right_pads,
InElementwiseOperation in_element_op,
WeiElementwiseOperation wei_element_op,
OutElementwiseOperation out_element_op,
const ck::index_t split_k)
const index_t split_k)
{
const bool stride_ovf = tensor_exceeds_2gb<BDataType>(b_g_n_c_wis_lengths) ||
tensor_exceeds_2gb<AccDataType>(e_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths);
return Argument{p_in_grid,
p_wei_grid,
p_out_grid,
Expand All @@ -1491,7 +1501,8 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
in_element_op,
wei_element_op,
out_element_op,
split_k};
split_k,
stride_ovf};
}

static auto MakeArgument(const InDataType* p_in_grid,
Expand All @@ -1513,7 +1524,7 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
const ck::index_t split_k)
{
const bool stride_ovf = tensor_exceeds_2gb<BDataType>(b_g_n_c_wis_lengths) ||
tensor_exceeds_2gb<EDataType>(e_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<AccDataType>(e_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths);
std::array<index_t, NDimSpatial + 3> b_g_n_c_wis_lengths_i32;
std::array<index_t, NDimSpatial + 3> b_g_n_c_wis_strides_i32;
Expand Down Expand Up @@ -1569,15 +1580,19 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
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<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,
const std::array<ck::index_t, NDimSpatial>& input_right_pads,
const std::array<index_t, NDimSpatial>& conv_filter_strides,
const std::array<index_t, NDimSpatial>& conv_filter_dilations,
const std::array<index_t, NDimSpatial>& input_left_pads,
const std::array<index_t, NDimSpatial>& input_right_pads,
InElementwiseOperation in_element_op,
WeiElementwiseOperation wei_element_op,
OutElementwiseOperation out_element_op,
const ck::index_t split_k) override
const index_t split_k) override
{
const bool stride_ovf = tensor_exceeds_2gb<BDataType>(b_g_n_c_wis_lengths) ||
tensor_exceeds_2gb<AccDataType>(e_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths);

return std::make_unique<Argument>(static_cast<const InDataType*>(p_in_grid),
static_cast<WeiDataType*>(p_wei_grid),
static_cast<const OutDataType*>(p_out_grid),
Expand All @@ -1596,7 +1611,8 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
in_element_op,
wei_element_op,
out_element_op,
split_k);
split_k,
stride_ovf);
}

std::unique_ptr<BaseArgument>
Expand All @@ -1619,7 +1635,7 @@ struct DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3
ck::index_t split_k) override
{
const bool stride_ovf = tensor_exceeds_2gb<BDataType>(b_g_n_c_wis_lengths) ||
tensor_exceeds_2gb<EDataType>(e_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<AccDataType>(e_g_k_c_xs_lengths) ||
tensor_exceeds_2gb<ADataType>(a_g_n_k_wos_lengths);

std::array<index_t, NDimSpatial + 3> b_g_n_c_wis_lengths_i32;
Expand Down
Loading
Loading