/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_ab_transfer_thread_tiles.hpp Source File
Go to the documentation of this file.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "ck/utility/amd_address_space.hpp"
 #include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v4r1.hpp"
 #include "ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_v7r2.hpp"
  
 namespace ck {
  
 template <typename ABLayout,
           typename ABMajorLayout,
           typename LDSTypeAB,
           index_t BlockSize,
           index_t MNPerBlock,
           index_t KPerBlock,
           index_t MNPerWmma,
           index_t ABK1Value,
           bool UseBlockPaddingAB,
           bool PermuteAB,
           typename ABBlockTransferThreadClusterLengths_ABK0_MN_ABK1,
           typename ABBlockTransferThreadClusterArrangeOrder,
           typename ABBlockTransferSrcAccessOrder,
           index_t ABBlockTransferSrcVectorDim,
           index_t ABBlockTransferSrcScalarPerVector,
           index_t ABBlockTransferDstScalarPerVector_ABK1,
           bool ABThreadTransferSrcResetCoordinateAfterRun>
 struct ABTransferThreadTiles
 {
     static constexpr auto ABK0Number = Number<KPerBlock / ABK1Value>{};
     static constexpr auto ABK1Number = Number<ABK1Value>{};
  
     static constexpr auto I0 = Number<0>{};
     static constexpr auto I1 = Number<1>{};
     static constexpr auto I2 = Number<2>{};
  
     static constexpr index_t ABPackedSize = []() {
         if constexpr(is_same_v<remove_cvref_t<LDSTypeAB>, pk_i4_t>)
             return 2;
         else
             return 1;
     }();
  
     using ThisThreadBlock = ThisThreadBlock<BlockSize>;
  
     template <bool PadMN, bool PadK, typename GridDescriptorBase>
     __host__ __device__ static auto MakeGridDescriptor(const GridDescriptorBase& ab_grid_desc,
                                                        index_t MN,
                                                        index_t MNPad,
                                                        index_t K,
                                                        index_t KPad,
                                                        index_t StrideAB,
                                                        index_t ABK0)
     {
  
         if constexpr(PadMN && PadK)
         {
             // pad both MN and K
             const auto ab_grid_desc_n_k =
                 transform_tensor_descriptor(ab_grid_desc,
                                             make_tuple(make_right_pad_transform(MN, MNPad - MN),
                                                        make_right_pad_transform(K, KPad - K)),
                                             make_tuple(Sequence<0>{}, Sequence<1>{}),
                                             make_tuple(Sequence<0>{}, Sequence<1>{}));
  
             const auto ab_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(
                 ab_grid_desc_n_k,
                 make_tuple(make_unmerge_transform(make_tuple(ABK0, ABK1Value)),
                            make_pass_through_transform(MNPad)),
                 make_tuple(Sequence<1>{}, Sequence<0>{}),
                 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
  
             return ab_grid_desc_bk0_n_bk1;
         }
         else if constexpr(PadMN && !PadK)
         {
             // pad MN, but not K
             const auto ab_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(
                 ab_grid_desc,
                 make_tuple(make_unmerge_transform(make_tuple(ABK0, ABK1Value)),
                            make_right_pad_transform(MN, MNPad - MN)),
                 make_tuple(Sequence<1>{}, Sequence<0>{}),
                 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
  
             return ab_grid_desc_bk0_n_bk1;
         }
         else if constexpr(!PadMN && PadK)
         {
             // pad K, but not MN
             const auto ab_grid_desc_n_k = transform_tensor_descriptor(
                 ab_grid_desc,
                 make_tuple(make_pass_through_transform(MN), make_right_pad_transform(K, KPad - K)),
                 make_tuple(Sequence<0>{}, Sequence<1>{}),
                 make_tuple(Sequence<0>{}, Sequence<1>{}));
  
             const auto ab_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(
                 ab_grid_desc_n_k,
                 make_tuple(make_unmerge_transform(make_tuple(ABK0, ABK1Value)),
                            make_pass_through_transform(MN)),
                 make_tuple(Sequence<1>{}, Sequence<0>{}),
                 make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
  
             return ab_grid_desc_bk0_n_bk1;
         }
         else
         {
             if constexpr(!PermuteAB)
             {
                 // not pad MN or K
                 const auto ab_grid_desc_bk0_n_bk1 = transform_tensor_descriptor(
                     ab_grid_desc,
                     make_tuple(make_unmerge_transform(make_tuple(ABK0, ABK1Value)),
                                make_pass_through_transform(MN)),
                     make_tuple(Sequence<1>{}, Sequence<0>{}),
                     make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
  
                 return ab_grid_desc_bk0_n_bk1;
             }
             else
             {
                 // Pre-shuffled Weight
                 // BGlobal[K / KPerBlock, MN, KPerBlock / K1, K1] -> BTile[K / K1, MN, K1]
                 constexpr index_t ABK01 = KPerBlock / ABK1Value;
                 const index_t ABK0_     = StrideAB / ABK1Value;
                 const index_t ABK00     = ABK0_ / ABK01;
  
                 const auto ab_grid_desc_abk00_mn_abk01_abk1_permute =
                     make_naive_tensor_descriptor_packed(make_tuple(ABK00, MN, ABK01, ABK1Value));
  
                 const auto ab_grid_desc_abk0_mn_abk1_permute = transform_tensor_descriptor(
                     ab_grid_desc_abk00_mn_abk01_abk1_permute,
                     make_tuple(make_merge_transform(make_tuple(ABK00, ABK01)),
                                make_pass_through_transform(make_tuple(MN)),
                                make_pass_through_transform(ABK1Value)),
                     make_tuple(Sequence<0, 2>{}, Sequence<1>{}, Sequence<3>{}),
                     make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
  
                 return ab_grid_desc_abk0_mn_abk1_permute;
             }
         }
     }
  
     __device__ static constexpr auto GetBlockDescriptor()
     {
         // A matrix in LDS memory, dst of blockwise copy
         if constexpr(UseBlockPaddingAB)
         {
             // bank conflict when writting the data into LDS, but don't worry, we have whole entire
             // loop to hide it in v4. it may give you some benefit from less valu in compute address
             return make_naive_tensor_descriptor(
                 make_tuple(ABK0Number, Number<MNPerBlock>{}, ABK1Number),
                 make_tuple(Number<MNPerBlock + 1>{} * ABK1Number, ABK1Number, I1));
         }
         // xor tensor transformation request more unnecessary vgpr usage, would cause register spill
         // in some cases.
         else if constexpr(is_same<ABMajorLayout, ABLayout>::value)
         {
             constexpr index_t LdsSize = 32 * 4 / KPerBlock / sizeof(LDSTypeAB) / ABPackedSize;
             constexpr auto MNLdsLayer = LdsSize < 1 ? 1 : LdsSize;
             constexpr auto ab_lds_block_desc = make_naive_tensor_descriptor(
                 make_tuple(ABK0Number * Number<MNLdsLayer>{},
                            Number<MNPerBlock / MNLdsLayer>{},
                            ABK1Number),
                 make_tuple(ABK1Number, Number<KPerBlock * MNLdsLayer>{}, I1));
  
             constexpr auto ab_lds_block_desc_permuted = transform_tensor_descriptor(
                 ab_lds_block_desc,
                 make_tuple(
                     make_xor_with_modulo_transform(make_tuple(Number<MNPerBlock / MNLdsLayer>{},
                                                               Number<ABK0Number * MNLdsLayer>{})),
                     make_pass_through_transform(ABK1Number)),
                 make_tuple(Sequence<1, 0>{}, Sequence<2>{}),
                 make_tuple(Sequence<1, 0>{}, Sequence<2>{}));
  
             constexpr auto ab_lds_block_desc_abk0_mnldslayer_mn_abk1 = transform_tensor_descriptor(
                 ab_lds_block_desc_permuted,
                 make_tuple(make_unmerge_transform(make_tuple(ABK0Number, Number<MNLdsLayer>{})),
                            make_pass_through_transform(Number<MNPerBlock / MNLdsLayer>{}),
                            make_pass_through_transform(ABK1Number)),
                 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
                 make_tuple(Sequence<0, 2>{}, Sequence<1>{}, Sequence<3>{}));
  
             constexpr auto ab_lds_block_desc_abk0_mn_abk1 = transform_tensor_descriptor(
                 ab_lds_block_desc_abk0_mnldslayer_mn_abk1,
                 make_tuple(make_pass_through_transform(ABK0Number),
                            make_merge_transform_v3_division_mod(
                                make_tuple(Number<MNPerBlock / MNLdsLayer>{}, Number<MNLdsLayer>{})),
                            make_pass_through_transform(ABK1Number)),
                 make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}),
                 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
  
             return ab_lds_block_desc_abk0_mn_abk1;
         }
         else
         {
             // kfold and mpair dimension is not always required.
             // more dimension in merge_transform increase the difficulty of generating immarg offset
             // for compiler.
             constexpr auto MN0 = ABBlockTransferThreadClusterLengths_ABK0_MN_ABK1{}.At(I1);
             constexpr auto MN1 = MNPerBlock / MN0;
  
             constexpr auto KThreadWrite = ABBlockTransferThreadClusterLengths_ABK0_MN_ABK1{}.At(I0);
             constexpr auto K0PerThreadWrite = ABK0Number / KThreadWrite;
             constexpr auto KThreadRead      = 64 / MNPerWmma;
             constexpr auto K0PerThreadRead  = ABK0Number / KThreadRead;
  
             constexpr auto kfold = (ABK1Number * MN0 * sizeof(LDSTypeAB) > 128)
                                        ? 1
                                        : 128 / (ABK1Number * MN0 * sizeof(LDSTypeAB));
             constexpr auto KThreadReadPerm =
                 (kfold * K0PerThreadWrite / K0PerThreadRead) > 1
                     ? KThreadRead / (kfold * K0PerThreadWrite / K0PerThreadRead)
                     : KThreadRead;
  
             // 1<=mpair<=n0
             constexpr auto mpair = (ABK1Number * MNPerWmma * sizeof(LDSTypeAB) > 128)
                                        ? 1
                                        : ((128 / (ABK1Number * MNPerWmma * sizeof(LDSTypeAB))) > MN0
                                               ? MN0
                                               : 128 / (ABK1Number * MNPerWmma * sizeof(LDSTypeAB)));
  
             constexpr auto ab_lds_block_desc = make_naive_tensor_descriptor_packed(
                 make_tuple(Number<KThreadWrite / kfold / KThreadReadPerm>{},
                            Number<K0PerThreadWrite>{},
                            Number<KThreadReadPerm * MN1>{},
                            Number<kfold * MN0 / mpair>{},
                            Number<mpair>{},
                            ABK1Number));
  
             constexpr auto ab_lds_block_desc_permuted = transform_tensor_descriptor(
                 ab_lds_block_desc,
                 make_tuple(
                     make_pass_through_transform(Number<KThreadWrite / kfold / KThreadReadPerm>{}),
                     make_pass_through_transform(Number<K0PerThreadWrite>{}),
                     make_xor_with_modulo_transform(
                         make_tuple(Number<KThreadReadPerm * MN1>{}, Number<kfold * MN0 / mpair>{})),
                     make_pass_through_transform(Number<mpair>{}),
                     make_pass_through_transform(ABK1Number)),
                 make_tuple(
                     Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4>{}, Sequence<5>{}),
                 make_tuple(
                     Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4>{}, Sequence<5>{}));
  
             constexpr auto ab_lds_block_desc_unmerged = transform_tensor_descriptor(
                 ab_lds_block_desc_permuted,
                 make_tuple(
                     make_pass_through_transform(Number<KThreadWrite / kfold / KThreadReadPerm>{}),
                     make_pass_through_transform(Number<K0PerThreadWrite>{}),
                     make_unmerge_transform(make_tuple(Number<KThreadReadPerm>{}, Number<MN1>{})),
                     make_unmerge_transform(make_tuple(Number<kfold>{}, Number<MN0 / mpair>{})),
                     make_pass_through_transform(Number<mpair>{}),
                     make_pass_through_transform(ABK1Number)),
                 make_tuple(Sequence<0>{},
                            Sequence<1>{},
                            Sequence<2>{},
                            Sequence<3>{},
                            Sequence<4>{},
                            Sequence<5>{}),
                 make_tuple(Sequence<1>{},
                            Sequence<2>{},
                            Sequence<0, 3>{},
                            Sequence<4, 5>{},
                            Sequence<6>{},
                            Sequence<7>{}));
  
             constexpr auto ab_lds_block_desc_abk0_mn_abk1 = transform_tensor_descriptor(
                 ab_lds_block_desc_unmerged,
                 make_tuple(make_merge_transform_v3_division_mod(
                                make_tuple(Number<KThreadReadPerm>{},
                                           Number<KThreadWrite / kfold / KThreadReadPerm>{},
                                           Number<kfold>{},
                                           Number<K0PerThreadWrite>{})),
                            make_merge_transform_v3_division_mod(
                                make_tuple(Number<MN0 / mpair>{}, Number<mpair>{}, Number<MN1>{})),
                            make_pass_through_transform(ABK1Number)),
                 make_tuple(Sequence<0, 1, 4, 2>{}, Sequence<5, 6, 3>{}, Sequence<7>{}),
                 make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}));
  
             return ab_lds_block_desc_abk0_mn_abk1;
         }
     }
  
     template <typename GridDescriptor,
               typename BlockDescriptor,
               typename ABsDataType,
               typename ABElementwiseOperation,
               index_t GlobalBufferNum>
     __device__ static auto GetBlockTransfer(GridDescriptor& grid_descriptor,
                                             BlockDescriptor& block_descriptor,
                                             ABElementwiseOperation& ab_element_op,
                                             const index_t block_mn_id)
     {
         constexpr index_t NumABTensor = ABsDataType::Size();
         const index_t mn_block_data_idx_on_grid =
             __builtin_amdgcn_readfirstlane(block_mn_id * MNPerBlock);
         // workaround because v7r2 is not as general as v4r1
         if constexpr(NumABTensor > 1)
         {
             const auto idx_as_block_begin = generate_tuple(
                 [&](auto) { return make_multi_index(0, mn_block_data_idx_on_grid, 0); },
                 Number<NumABTensor>{});
  
             return ThreadGroupTensorSliceTransfer_v7r2<
                 ThisThreadBlock,
                 ABsDataType,
                 Tuple<LDSTypeAB>,
                 GridDescriptor,
                 decltype(tie(block_descriptor)),
                 ABElementwiseOperation,
                 Sequence<static_cast<index_t>(InMemoryDataOperationEnum::Set)>,
                 Sequence<ABK0Number, MNPerBlock, ABK1Number>,
                 ABBlockTransferThreadClusterLengths_ABK0_MN_ABK1,
                 ABBlockTransferThreadClusterArrangeOrder,
                 ABBlockTransferSrcAccessOrder,
                 Sequence<1, 0, 2>,
                 ABBlockTransferSrcVectorDim,
                 2,
                 ABBlockTransferSrcScalarPerVector,
                 ABBlockTransferDstScalarPerVector_ABK1,
                 uniform_sequence_gen_t<NumABTensor, ABThreadTransferSrcResetCoordinateAfterRun>,
                 Sequence<true>,
                 GlobalBufferNum>{grid_descriptor,
                                  idx_as_block_begin,
                                  tie(block_descriptor),
                                  make_tuple(make_multi_index(0, 0, 0)),
                                  ab_element_op};
         }
         else
         {
             return ThreadGroupTensorSliceTransfer_v4r1<
                 ThisThreadBlock,
                 ABElementwiseOperation,
                 ck::tensor_operation::element_wise::PassThrough,
                 InMemoryDataOperationEnum::Set,
                 Sequence<ABK0Number, MNPerBlock, ABK1Number>,
                 ABBlockTransferThreadClusterLengths_ABK0_MN_ABK1,
                 ABBlockTransferThreadClusterArrangeOrder,
                 remove_cvref_t<tuple_element_t<0, ABsDataType>>,
                 remove_cvref_t<tuple_element_t<0, ABsDataType>>,
                 decltype(grid_descriptor[I0]),
                 decltype(block_descriptor),
                 ABBlockTransferSrcAccessOrder,
                 Sequence<0, 1, 2>,
                 ABBlockTransferSrcVectorDim,
                 2,
                 ABBlockTransferSrcScalarPerVector,
                 ABBlockTransferDstScalarPerVector_ABK1,
                 1,
                 1,
                 ABThreadTransferSrcResetCoordinateAfterRun,
                 true,
                 GlobalBufferNum>(grid_descriptor[I0],
                                  make_multi_index(0, mn_block_data_idx_on_grid, 0),
                                  ab_element_op,
                                  block_descriptor,
                                  make_multi_index(0, 0, 0),
                                  ck::tensor_operation::element_wise::PassThrough{});
         }
     }
  
     template <index_t MNRepeat, index_t MNWaves>
     __host__ __device__ static constexpr auto MakeWmmaTileDescriptor()
     {
         // This is a block descriptor used to read LDS memory into register
         // It's defined in a way consistent with the existing implementation to
         // avoid changes in the pipelines
         using BlockDesc = decltype(GetBlockDescriptor());
         // ABK0_MN_ABK1 -> ABK0_MNRepeat_MNWaves_KRow_MNPerWmma_ABK1
         constexpr auto ABK0 = BlockDesc{}.GetLength(I0);
         constexpr auto ABK1 = BlockDesc{}.GetLength(I2);
 #ifdef __gfx12__
         constexpr auto KRow = I2;
 #else
         constexpr auto KRow = I1;
 #endif
         return transform_tensor_descriptor(
             BlockDesc{},
             make_tuple(make_unmerge_transform(make_tuple(Number<ABK0 / KRow>{}, KRow)),
                        make_unmerge_transform(
                            make_tuple(Number<MNRepeat>{}, Number<MNWaves>{}, Number<MNPerWmma>{})),
                        make_pass_through_transform(Number<ABK1>{})),
             make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
             make_tuple(Sequence<0, 3>{}, Sequence<1, 2, 4>{}, Sequence<5>{}));
     }
  
     __device__ static constexpr auto GetBlockStep()
     {
         // Grid descriptor step (MoveSrcSliceWindow)
         return make_multi_index(KPerBlock / ABK1Number, 0, 0);
     }
  
     template <typename GridDescriptor>
     __device__ static constexpr index_t GetKDimension(const GridDescriptor& grid_desc)
     {
         // K dimension size. This should always be called with the A matrix grid descriptor
         // because it doesn't work for B matrix when packed int4 is used
         return grid_desc.GetLength(I0) * grid_desc.GetLength(I2);
     }
 };
  
 } // namespace ck