Skip to content

feat(ck): convolution LDS Transpose and TransposeC support #11905

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 23, 2026
Merged

feat(ck): convolution LDS Transpose and TransposeC support #11905

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
864208a
Initial support for LDS load transpose on gfx1250 (old ck)
EnricoDeg Aug 24, 2026
e35fe72
Fix bug tensor descriptors
EnricoDeg Aug 25, 2026
f8191fe
Fix LDS layout
EnricoDeg Aug 25, 2026
bc590ba
Add bf16 instruction for ds load transpose
EnricoDeg Aug 26, 2026
1ef1188
TransposeC support in cshuffle epilogue Wmma implementation
EnricoDeg Aug 26, 2026
6e58eeb
Expose TransposeC and LdsTranspose as template parameters in Gridwise…
EnricoDeg Aug 26, 2026
31a6491
Integrate TransposeC and LdsTranspose in bwd data and add instances
EnricoDeg Aug 26, 2026
f0955f5
Enable LdsTranspose and TransposeC in bwd weight
EnricoDeg Sep 8, 2026
ef92e52
Add instances for bwd weight
EnricoDeg Sep 8, 2026
0b98c2b
Fix rebase
EnricoDeg Sep 9, 2026
e8bedd1
Clean up and generalize
EnricoDeg Sep 9, 2026
f31f966
Fix definition
EnricoDeg Sep 21, 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
12 changes: 9 additions & 3 deletions ...lder/reflect/instance_traits_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 @@ -60,7 +60,9 @@ template <index_t NDimSpatial,
typename ComputeTypeA,
typename ComputeTypeB,
ck::index_t max_transpose_transfer_src_scalar_per_vector,
ck::index_t max_transpose_transfer_dst_scalar_per_vector>
ck::index_t max_transpose_transfer_dst_scalar_per_vector,
bool UseLdsTranspose,
bool TransposeC>
struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3;

} // namespace ck::tensor_operation::device
Expand Down Expand Up @@ -120,7 +122,9 @@ template <index_t NDimSpatial,
typename ComputeTypeA_,
typename ComputeTypeB_,
ck::index_t max_transpose_transfer_src_scalar_per_vector,
ck::index_t max_transpose_transfer_dst_scalar_per_vector>
ck::index_t max_transpose_transfer_dst_scalar_per_vector,
bool UseLdsTranspose,
bool TransposeC>
struct InstanceTraits<
ck::tensor_operation::device::DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3<
NDimSpatial,
Expand Down Expand Up @@ -173,7 +177,9 @@ struct InstanceTraits<
ComputeTypeA_,
ComputeTypeB_,
max_transpose_transfer_src_scalar_per_vector,
max_transpose_transfer_dst_scalar_per_vector>>
max_transpose_transfer_dst_scalar_per_vector,
UseLdsTranspose,
TransposeC>>
{
static constexpr auto kTensorOpName = "DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3";

Expand Down
12 changes: 9 additions & 3 deletions ..._tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,9 @@ template <ck::index_t NDimSpatial,
typename ComputeTypeA,
typename ComputeTypeB,
ck::index_t MaxTransposeTransferSrcScalarPerVector,
ck::index_t MaxTransposeTransferDstScalarPerVector>
ck::index_t MaxTransposeTransferDstScalarPerVector,
bool UseLdsTranspose,
bool TransposeC>
struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3;

} // namespace ck::tensor_operation::device
Expand Down Expand Up @@ -112,7 +114,9 @@ template <ck::index_t NDimSpatial,
typename ComputeTypeA_,
typename ComputeTypeB_,
ck::index_t MaxTransposeTransferSrcScalarPerVector,
ck::index_t MaxTransposeTransferDstScalarPerVector>
ck::index_t MaxTransposeTransferDstScalarPerVector,
bool UseLdsTranspose,
bool TransposeC>
struct InstanceTraits<ck::tensor_operation::device::DeviceGroupedConvBwdWeight_Wmma_CShuffleV3<
NDimSpatial,
InLayout_,
Expand Down Expand Up @@ -158,7 +162,9 @@ struct InstanceTraits<ck::tensor_operation::device::DeviceGroupedConvBwdWeight_W
ComputeTypeA_,
ComputeTypeB_,
MaxTransposeTransferSrcScalarPerVector,
MaxTransposeTransferDstScalarPerVector>>
MaxTransposeTransferDstScalarPerVector,
UseLdsTranspose,
TransposeC>>
{
static constexpr auto kTensorOpName = "DeviceGroupedConvBwdWeight_Wmma_CShuffleV3";
using device_kernel_tag = DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3_Tag;
Expand Down
14 changes: 10 additions & 4 deletions ...blekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmma_selector.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,10 @@ template <BlockGemmPipelineVersion BlkGemmPipelineVer,
index_t NRepeat,
index_t KPack,
index_t KInner,
bool TransposeC = false,
bool BSkipLDS = false>
bool TransposeC = false,
bool BSkipLDS = false,
bool UseLdsTransposeA = false,
bool UseLdsTransposeB = false>
constexpr auto BlockGemmPipeline_Selector()
{
if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
Expand All @@ -56,7 +58,9 @@ constexpr auto BlockGemmPipeline_Selector()
KPack,
KInner,
TransposeC,
BSkipLDS>{};
BSkipLDS,
UseLdsTransposeA,
UseLdsTransposeB>{};
}
else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
{
Expand All @@ -81,7 +85,9 @@ constexpr auto BlockGemmPipeline_Selector()
KPack,
KInner,
TransposeC,
BSkipLDS>{};
BSkipLDS,
UseLdsTransposeA,
UseLdsTransposeB>{};
}
else
{
Expand Down
69 changes: 61 additions & 8 deletions ...ablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,9 @@ template <index_t BlockSize,
index_t NRepeat,
index_t KPack,
index_t KInner,
bool TransposeC = false>
bool TransposeC = false,
bool UseLdsTransposeA = false,
bool UseLdsTransposeB = false>
struct BlockwiseGemmWmmaops_pipeline_base
{
static constexpr auto I0 = Number<0>{};
Expand All @@ -48,6 +50,8 @@ struct BlockwiseGemmWmmaops_pipeline_base

using ThisThreadBlock = ThisThreadBlock<BlockSize>;

static constexpr bool TransposeC_ = TransposeC;

static constexpr index_t WaveSize = 32;

static constexpr index_t MWaves = MPerBlock / (MRepeat * MPerWmma);
Expand All @@ -70,8 +74,12 @@ struct BlockwiseGemmWmmaops_pipeline_base
TransposeC>{};

static constexpr index_t KPerThread = wmma_gemm.wmma_instr.k_per_blk * KInner;
static constexpr index_t A_K1 = ck::math::min(AWmmaTileDesc{}.GetLength(I6), KPerThread);
static constexpr index_t B_K1 = ck::math::min(BWmmaTileDesc{}.GetLength(I6), KPerThread);
static constexpr index_t A_K1 = UseLdsTransposeA
? ck::math::min(AWmmaTileDesc{}.GetLength(I5), KPerThread)
: ck::math::min(AWmmaTileDesc{}.GetLength(I6), KPerThread);
static constexpr index_t B_K1 = UseLdsTransposeB
? ck::math::min(BWmmaTileDesc{}.GetLength(I5), KPerThread)
: ck::math::min(BWmmaTileDesc{}.GetLength(I6), KPerThread);

static_assert(KPack % (A_K1 * A_KRow) == 0, "wrong!");
static_assert(KPack % (B_K1 * B_KRow) == 0, "wrong!");
Expand Down Expand Up @@ -282,8 +290,20 @@ struct BlockwiseGemmWmmaops_pipeline_base
#else
const auto wmma_krow = 0;
#endif

return make_tuple(0, 0, 0, waveId_m, wmma_krow, wmma_a_idx, 0);
if constexpr(UseLdsTransposeA)
{
return make_tuple(0,
0,
0,
waveId_m,
wmma_krow,
wmma_a_idx % A_K1,
(wmma_a_idx / A_K1) * (MPerWmma / 2));
}
else
{
return make_tuple(0, 0, 0, waveId_m, wmma_krow, wmma_a_idx, 0);
}
}

__device__ static auto CalculateBThreadOriginDataIndex()
Expand All @@ -300,7 +320,20 @@ struct BlockwiseGemmWmmaops_pipeline_base
const auto wmma_krow = 0;
#endif

return make_tuple(0, 0, 0, waveId_n, wmma_krow, wmma_b_idx, 0);
if constexpr(UseLdsTransposeB)
{
return make_tuple(0,
0,
0,
waveId_n,
wmma_krow,
wmma_b_idx % B_K1,
(wmma_b_idx / B_K1) * (NPerWmma / 2));
}
else
{
return make_tuple(0, 0, 0, waveId_n, wmma_krow, wmma_b_idx, 0);
}
}

template <index_t m0, index_t n0>
Expand Down Expand Up @@ -405,6 +438,7 @@ struct BlockwiseGemmWmmaops_pipeline_base
AccStride));
}

// C = A * B
__host__ __device__ static constexpr auto
GetCBlockDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs()
{
Expand All @@ -421,6 +455,23 @@ struct BlockwiseGemmWmmaops_pipeline_base
c_block_desc_mrepeat_mwave_mperwmma_nrepeat_nwave_nperwmma);
}

// C' = B' * A'
__host__ __device__ static constexpr auto
GetCBlockDescriptor_MRepeat_MWave_MThreadPerSubGroup_NRepeat_NWave_NSubGroup_MAccVgprs()
{
constexpr auto c_block_desc_mrepeat_mwave_mperwmma_nrepeat_nwave_nperwmma =
make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{},
Number<MWaves>{},
Number<MPerWmma>{},
Number<NRepeat>{},
Number<NWaves>{},
Number<NPerWmma>{}));

return wmma_gemm
.MakeCDesc_MBlockxRepeat_MWave_MThreadPerSubGroup_NBlockxRepeat_NWave_NSubGroup_NAccVgprs(
c_block_desc_mrepeat_mwave_mperwmma_nrepeat_nwave_nperwmma);
}

// Describe how data allocated in thread copy src buffer
// M0_M1_M2 = MRepeat_MWave_MPerWmma, N0_N1_N2 = NRepeat_NWave_NPerWmma
static constexpr AWmmaTileDesc a_block_desc_k0_m0_m1_m2_k1;
Expand Down Expand Up @@ -472,7 +523,8 @@ struct BlockwiseGemmWmmaops_pipeline_base
Sequence<0, 1, 2, 3, 4, 5, 6>,
6,
A_K1,
A_K1>;
A_K1,
UseLdsTransposeA>;

using BThreadCopy =
ThreadwiseTensorSliceTransfer_v4<BDataType,
Expand All @@ -483,7 +535,8 @@ struct BlockwiseGemmWmmaops_pipeline_base
Sequence<0, 1, 2, 3, 4, 5, 6>,
6,
B_K1,
B_K1>;
B_K1,
UseLdsTransposeB>;

AThreadCopy a_thread_copy_;
BThreadCopy b_thread_copy_;
Expand Down
Loading
Loading