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

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

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_waveletmodel_cshuffle.hpp Source File
Go to the documentation of this file.
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include <iostream>
 #include <sstream>
  
 #include "ck/utility/common_header.hpp"
 #include "ck/tensor_description/tensor_descriptor.hpp"
 #include "ck/tensor_description/tensor_descriptor_helper.hpp"
 #include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
 #include "ck/tensor_operation/gpu/device/device_gemm.hpp"
 #include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
 #include "ck/tensor_operation/gpu/device/matrix_padder.hpp"
 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_waveletmodel_cshuffle.hpp"
 #include "ck/host_utility/device_prop.hpp"
 #include "ck/host_utility/kernel_launch.hpp"
  
 namespace ck {
  
 template <typename GridwiseGemm,
           typename ABDataType,
           typename EDataType,
           typename AElementwiseOperation,
           typename BElementwiseOperation,
           typename EElementwiseOperation,
           typename AGridDesc_AK0_M_AK1,
           typename BGridDesc_BK0_N_BK1,
           typename EGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
           typename Block2ETileMap,
           bool HasMainKBlockLoop>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
 __launch_bounds__(CK_WAVELET_MAX_THREAD_PER_BLOCK, CK_WAVELET_MIN_BLOCK_PER_CU)
 #endif
     kernel_gemm_xdl_waveletmodel_cshuffle(const ABDataType* __restrict__ p_a_grid,
                                           const ABDataType* __restrict__ p_b_grid,
                                           EDataType* __restrict__ p_e_grid,
                                           const AElementwiseOperation a_element_op,
                                           const BElementwiseOperation b_element_op,
                                           const EElementwiseOperation e_element_op,
                                           const AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1,
                                           const BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1,
                                           const EGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock
                                               e_grid_desc_mblock_mperblock_nblock_nperblock,
                                           const Block2ETileMap block_2_etile_map)
 {
 #if defined(__gfx9__) || defined(__gfx11__) || defined(__gfx12__)
     if constexpr(GridwiseGemm::template IsValidCompilationParameter<>())
     {
         __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
  
         GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid,
                                                       p_b_grid,
                                                       p_e_grid,
                                                       p_shared,
                                                       a_element_op,
                                                       b_element_op,
                                                       e_element_op,
                                                       a_grid_desc_ak0_m_ak1,
                                                       b_grid_desc_bk0_n_bk1,
                                                       e_grid_desc_mblock_mperblock_nblock_nperblock,
                                                       block_2_etile_map);
     }
 #else
     ignore = p_a_grid;
     ignore = p_b_grid;
     ignore = p_e_grid;
     ignore = a_element_op;
     ignore = b_element_op;
     ignore = e_element_op;
     ignore = a_grid_desc_ak0_m_ak1;
     ignore = b_grid_desc_bk0_n_bk1;
     ignore = e_grid_desc_mblock_mperblock_nblock_nperblock;
     ignore = block_2_etile_map;
 #endif
 }
  
 } // namespace ck
  
 namespace ck {
 namespace tensor_operation {
 namespace device {
  
 template <typename ALayout,
           typename BLayout,
           typename ELayout,
           typename ADataType,
           typename BDataType,
           typename GemmAcEDataType,
           typename CShuffleDataType,
           typename EDataType,
           typename AElementwiseOperation,
           typename BElementwiseOperation,
           typename CDEElementwiseOperation,
           GemmSpecialization GemmSpec,
           index_t NumGemmKPrefetchStage,
           index_t TileLoadThreadGroupSize,
           index_t TileMathThreadGroupSize,
           index_t MPerBlock,
           index_t NPerBlock,
           index_t KPerBlock,
           index_t AK1,
           index_t BK1,
           index_t MPerXDL,
           index_t NPerXDL,
           index_t MXdlPerWave,
           index_t NXdlPerWave,
           typename ABlockTransferThreadClusterLengths_AK0_M_AK1,
           typename ABlockTransferThreadClusterArrangeOrder,
           typename ABlockTransferSrcAccessOrder,
           index_t ABlockTransferSrcVectorDim,
           index_t ABlockTransferSrcScalarPerVector,
           index_t ABlockTransferDstScalarPerVector_AK1,
           bool ABlockLdsExtraM,
           typename BBlockTransferThreadClusterLengths_BK0_N_BK1,
           typename BBlockTransferThreadClusterArrangeOrder,
           typename BBlockTransferSrcAccessOrder,
           index_t BBlockTransferSrcVectorDim,
           index_t BBlockTransferSrcScalarPerVector,
           index_t BBlockTransferDstScalarPerVector_BK1,
           bool BBlockLdsExtraN,
           index_t CShuffleMXdlPerWavePerShuffle,
           index_t CShuffleNXdlPerWavePerShuffle,
           typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
           index_t CShuffleBlockTransferScalarPerVector_NPerBlock>
 struct DeviceGemm_Xdl_WaveletModel_CShuffle : public DeviceGemm<ALayout,
                                                                 BLayout,
                                                                 ELayout,
                                                                 ADataType,
                                                                 BDataType,
                                                                 EDataType,
                                                                 AElementwiseOperation,
                                                                 BElementwiseOperation,
                                                                 CDEElementwiseOperation>
 {
     static constexpr auto BlockSize = math::max(TileLoadThreadGroupSize, TileMathThreadGroupSize);
     GET_NXDL_PER_WAVE_IMPL
     static constexpr auto NXdlPerWave64 = GetNXdlPerWave<true>();
     static constexpr auto NXdlPerWave32 = GetNXdlPerWave<false>();
  
     using DeviceOp = DeviceGemm_Xdl_WaveletModel_CShuffle;
  
     static constexpr auto I0 = Number<0>{};
     static constexpr auto I1 = Number<1>{};
     static constexpr auto I2 = Number<2>{};
  
     static constexpr auto matrix_padder =
         MatrixPadder<GemmSpec, index_t, index_t, index_t>{MPerBlock, NPerBlock, KPerBlock};
  
     static auto MakeAGridDescriptor_M_K(index_t MRaw, index_t KRaw, index_t StrideA)
     {
         const auto a_grid_desc_mraw_kraw = [&]() {
             if constexpr(is_same_v<tensor_layout::gemm::RowMajor, ALayout>)
             {
                 return make_naive_tensor_descriptor(make_tuple(MRaw, KRaw),
                                                     make_tuple(StrideA, I1));
             }
             else if constexpr(is_same_v<tensor_layout::gemm::ColumnMajor, ALayout>)
             {
                 return make_naive_tensor_descriptor(make_tuple(MRaw, KRaw),
                                                     make_tuple(I1, StrideA));
             }
         }();
  
         return matrix_padder.PadADescriptor_M_K(a_grid_desc_mraw_kraw);
     }
  
     static auto MakeBGridDescriptor_N_K(index_t KRaw, index_t NRaw, index_t StrideB)
     {
         const auto b_grid_desc_nraw_kraw = [&]() {
             if constexpr(is_same<tensor_layout::gemm::RowMajor, BLayout>::value)
             {
                 return make_naive_tensor_descriptor(make_tuple(NRaw, KRaw),
                                                     make_tuple(I1, StrideB));
             }
             else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, BLayout>::value)
             {
                 return make_naive_tensor_descriptor(make_tuple(NRaw, KRaw),
                                                     make_tuple(StrideB, I1));
             }
         }();
  
         return matrix_padder.PadBDescriptor_N_K(b_grid_desc_nraw_kraw);
     }
  
     template <typename ELay>
     static auto MakeEGridDescriptor_M_N(index_t MRaw, index_t NRaw, index_t StrideE)
     {
         const auto e_grid_desc_mraw_nraw = [&]() {
             if constexpr(is_same<tensor_layout::gemm::RowMajor, ELay>::value)
             {
                 return make_naive_tensor_descriptor(make_tuple(MRaw, NRaw),
                                                     make_tuple(StrideE, I1));
             }
             else if constexpr(is_same<tensor_layout::gemm::ColumnMajor, ELay>::value)
             {
                 return make_naive_tensor_descriptor(make_tuple(MRaw, NRaw),
                                                     make_tuple(I1, StrideE));
             }
         }();
  
         return matrix_padder.PadCDescriptor_M_N(e_grid_desc_mraw_nraw);
     }
  
     using AGridDesc_M_K = decltype(MakeAGridDescriptor_M_K(1, 1, 1));
     using BGridDesc_N_K = decltype(MakeBGridDescriptor_N_K(1, 1, 1));
     using EGridDesc_M_N = decltype(MakeEGridDescriptor_M_N<ELayout>(1, 1, 1));
  
     // GridwiseGemm
     template <index_t NXdlPerWave_>
     using GridwiseGemmBase = GridwiseGemm_k0mk1_k0nk1_mn_xdl_waveletmodel_cshuffle<
         ADataType, // TODO: distinguish A/B datatype
         GemmAcEDataType,
         CShuffleDataType,
         EDataType,
         AElementwiseOperation,
         BElementwiseOperation,
         CDEElementwiseOperation,
         InMemoryDataOperationEnum::Set,
         AGridDesc_M_K,
         BGridDesc_N_K,
         EGridDesc_M_N,
         NumGemmKPrefetchStage,
         TileLoadThreadGroupSize,
         TileMathThreadGroupSize,
         MPerBlock,
         NPerBlock,
         KPerBlock,
         AK1,
         BK1,
         MPerXDL,
         NPerXDL,
         MXdlPerWave,
         NXdlPerWave_,
         ABlockTransferThreadClusterLengths_AK0_M_AK1,
         ABlockTransferThreadClusterArrangeOrder,
         ABlockTransferSrcAccessOrder,
         ABlockTransferSrcVectorDim,
         ABlockTransferSrcScalarPerVector,
         ABlockTransferDstScalarPerVector_AK1,
         false,
         ABlockLdsExtraM,
         BBlockTransferThreadClusterLengths_BK0_N_BK1,
         BBlockTransferThreadClusterArrangeOrder,
         BBlockTransferSrcAccessOrder,
         BBlockTransferSrcVectorDim,
         BBlockTransferSrcScalarPerVector,
         BBlockTransferDstScalarPerVector_BK1,
         false,
         BBlockLdsExtraN,
         CShuffleMXdlPerWavePerShuffle,
         CShuffleNXdlPerWavePerShuffle,
         CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
         CShuffleBlockTransferScalarPerVector_NPerBlock>;
     using GridwiseGemm64 = GridwiseGemmBase<math::max(NXdlPerWave64, 1)>;
     using GridwiseGemm32 = GridwiseGemmBase<NXdlPerWave32>;
  
     using AGridDesc_AK0_M_AK1 =
         remove_cvref_t<decltype(GridwiseGemm64::MakeDefaultAGridDescriptor_AK0_M_AK1(
             AGridDesc_M_K{}))>;
     using BGridDesc_BK0_N_BK1 =
         remove_cvref_t<decltype(GridwiseGemm64::MakeDefaultBGridDescriptor_BK0_N_BK1(
             BGridDesc_N_K{}))>;
  
     using Block2ETileMap = typename GridwiseGemm64::DefaultBlock2ETileMap;
  
     // Argument
     struct Argument : public BaseArgument
     {
         Argument(const ADataType* p_a_grid,
                  const BDataType* p_b_grid,
                  EDataType* p_e_grid,
                  index_t MRaw,
                  index_t NRaw,
                  index_t KRaw,
                  index_t StrideA,
                  index_t StrideB,
                  index_t StrideE,
                  AElementwiseOperation a_element_op,
                  BElementwiseOperation b_element_op,
                  CDEElementwiseOperation cde_element_op)
             : p_a_grid_{static_cast<const ADataType*>(p_a_grid)},
               p_b_grid_{static_cast<const BDataType*>(p_b_grid)},
               p_e_grid_{static_cast<EDataType*>(p_e_grid)},
               a_grid_desc_m_k_{DeviceOp::MakeAGridDescriptor_M_K(MRaw, KRaw, StrideA)},
               b_grid_desc_n_k_{DeviceOp::MakeBGridDescriptor_N_K(KRaw, NRaw, StrideB)},
               e_grid_desc_m_n_{DeviceOp::MakeEGridDescriptor_M_N<ELayout>(MRaw, NRaw, StrideE)},
               a_grid_desc_ak0_m_ak1_{
                   GridwiseGemm64::MakeDefaultAGridDescriptor_AK0_M_AK1(a_grid_desc_m_k_)},
               b_grid_desc_bk0_n_bk1_{
                   GridwiseGemm64::MakeDefaultBGridDescriptor_BK0_N_BK1(b_grid_desc_n_k_)},
               block_2_etile_map_{GridwiseGemm64::MakeDefaultBlock2ETileMap(e_grid_desc_m_n_)},
               a_element_op_{a_element_op},
               b_element_op_{b_element_op},
               cde_element_op_{cde_element_op}
         {
         }
  
         void Print() const
         {
             std::cout << "A[M, K]: " << a_grid_desc_m_k_ << std::endl;
             std::cout << "B[N, K]: " << b_grid_desc_n_k_ << std::endl;
             std::cout << "E[M, N]: " << e_grid_desc_m_n_ << std::endl;
         }
  
         //  private:
         // pointers
         const ADataType* p_a_grid_;
         const BDataType* p_b_grid_;
         EDataType* p_e_grid_;
  
         // tensor descriptors for problem definiton
         AGridDesc_M_K a_grid_desc_m_k_;
         BGridDesc_N_K b_grid_desc_n_k_;
         EGridDesc_M_N e_grid_desc_m_n_;
  
         // tensor descriptors for block/thread-wise copy
         AGridDesc_AK0_M_AK1 a_grid_desc_ak0_m_ak1_;
         BGridDesc_BK0_N_BK1 b_grid_desc_bk0_n_bk1_;
  
         // block-to-e-tile map
         Block2ETileMap block_2_etile_map_;
  
         // element-wise op
         AElementwiseOperation a_element_op_;
         BElementwiseOperation b_element_op_;
         CDEElementwiseOperation cde_element_op_;
     };
  
     // Invoker
     struct Invoker : public BaseInvoker
     {
         using Argument = DeviceOp::Argument;
  
         template <typename GridwiseGemm>
         float RunImp(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
         {
 #if 0
             {
                 std::cout << "arg.a_grid_desc_ak0_m_ak1_{"
                           << arg.a_grid_desc_ak0_m_ak1_.GetLength(I0) << ", "
                           << arg.a_grid_desc_ak0_m_ak1_.GetLength(I1) << ", "
                           << arg.a_grid_desc_ak0_m_ak1_.GetLength(I2) << "}" << std::endl;
  
                 std::cout << "arg.b_grid_desc_bk0_n_bk1_{"
                           << arg.b_grid_desc_bk0_n_bk1_.GetLength(I0) << ", "
                           << arg.b_grid_desc_bk0_n_bk1_.GetLength(I1) << ", "
                           << arg.b_grid_desc_bk0_n_bk1_.GetLength(I2) << "}" << std::endl;
  
                 std::cout << "arg.e_grid_desc_m_n_{ " << arg.e_grid_desc_m_n_.GetLength(I0) << ", "
                           << arg.e_grid_desc_m_n_.GetLength(I1) << "}" << std::endl;
             }
 #endif
  
             if(!GridwiseGemm::CheckValidity(arg.a_grid_desc_m_k_,
                                             arg.b_grid_desc_n_k_,
                                             arg.e_grid_desc_m_n_,
                                             arg.block_2_etile_map_))
             {
                 throw std::runtime_error("wrong! GridwiseGemm has invalid setting");
             }
             auto e_grid_desc_mblock_mperblock_nblock_nperblock =
                 GridwiseGemm::MakeEGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(
                     arg.e_grid_desc_m_n_);
  
             const index_t grid_size = GridwiseGemm::CalculateGridSize(arg.e_grid_desc_m_n_);
             const auto K            = arg.a_grid_desc_m_k_.GetLength(I1);
  
             auto launch_kernel = [&](auto has_main_k_block_loop) {
                 constexpr bool has_main_loop = has_main_k_block_loop.value;
  
                 const auto kernel = kernel_gemm_xdl_waveletmodel_cshuffle<
                     GridwiseGemm,
                     ADataType, // TODO: distiguish A/B datatype
                     EDataType,
                     AElementwiseOperation,
                     BElementwiseOperation,
                     CDEElementwiseOperation,
                     DeviceOp::AGridDesc_AK0_M_AK1,
                     DeviceOp::BGridDesc_BK0_N_BK1,
                     typename GridwiseGemm::EGridDescriptor_MBlock_MPerBlock_NBlock_NPerBlock,
                     typename GridwiseGemm::DefaultBlock2ETileMap,
                     has_main_loop>;
  
                 return launch_and_time_kernel(
                     stream_config,
                     kernel,
                     dim3(grid_size),
                     dim3(TileLoadThreadGroupSize + TileMathThreadGroupSize),
                     0,
                     arg.p_a_grid_,
                     arg.p_b_grid_,
                     arg.p_e_grid_,
                     arg.a_element_op_,
                     arg.b_element_op_,
                     arg.cde_element_op_,
                     arg.a_grid_desc_ak0_m_ak1_,
                     arg.b_grid_desc_bk0_n_bk1_,
                     e_grid_desc_mblock_mperblock_nblock_nperblock,
                     arg.block_2_etile_map_);
             };
  
             if(GridwiseGemm::CalculateHasMainKBlockLoop(K))
             {
                 return launch_kernel(integral_constant<bool, true>{});
             }
             else
             {
                 return launch_kernel(integral_constant<bool, false>{});
             }
         }
  
         INVOKER_RUN_IMPL
  
         // polymorphic
         float Run(const BaseArgument* p_arg,
                   const StreamConfig& stream_config = StreamConfig{}) override
         {
             return Run(*dynamic_cast<const Argument*>(p_arg), stream_config);
         }
     };
  
     static bool IsSupportedArgument(const Argument& arg)
     {
         if(!ck::is_xdl_wmma_supported<ADataType, BDataType, MPerXDL, NPerXDL>())
         {
             return false;
         }
         if(get_warp_size() == 64)
         {
             if constexpr(NXdlPerWave64 > 0)
             {
                 return GridwiseGemm64::CheckValidity(arg.a_grid_desc_m_k_,
                                                      arg.b_grid_desc_n_k_,
                                                      arg.e_grid_desc_m_n_,
                                                      arg.block_2_etile_map_);
             }
         }
         else
         {
             if constexpr(NXdlPerWave32 > 0)
             {
                 return GridwiseGemm32::CheckValidity(arg.a_grid_desc_m_k_,
                                                      arg.b_grid_desc_n_k_,
                                                      arg.e_grid_desc_m_n_,
                                                      arg.block_2_etile_map_);
             }
         }
         return false;
     }
  
     // polymorphic
     bool IsSupportedArgument(const BaseArgument* p_arg) override
     {
         return IsSupportedArgument(*dynamic_cast<const Argument*>(p_arg));
     }
  
     static auto MakeArgument(const ADataType* p_a,
                              const BDataType* p_b,
                              EDataType* p_e,
                              index_t MRaw,
                              index_t NRaw,
                              index_t KRaw,
                              index_t StrideA,
                              index_t StrideB,
                              index_t StrideE,
                              AElementwiseOperation a_element_op,
                              BElementwiseOperation b_element_op,
                              CDEElementwiseOperation cde_element_op)
     {
         return Argument{p_a,
                         p_b,
                         p_e,
                         MRaw,
                         NRaw,
                         KRaw,
                         StrideA,
                         StrideB,
                         StrideE,
                         a_element_op,
                         b_element_op,
                         cde_element_op};
     }
  
     static auto MakeInvoker() { return Invoker{}; }
  
     // polymorphic
     std::unique_ptr<BaseArgument>
     MakeArgumentPointer(const void* p_a,
                         const void* p_b,
                         void* p_e,
                         index_t MRaw,
                         index_t NRaw,
                         index_t KRaw,
                         index_t StrideA,
                         index_t StrideB,
                         index_t StrideE,
                         AElementwiseOperation a_element_op,
                         BElementwiseOperation b_element_op,
                         CDEElementwiseOperation cde_element_op) override
     {
         return std::make_unique<Argument>(static_cast<const ADataType*>(p_a),
                                           static_cast<const BDataType*>(p_b),
                                           static_cast<EDataType*>(p_e),
                                           MRaw,
                                           NRaw,
                                           KRaw,
                                           StrideA,
                                           StrideB,
                                           StrideE,
                                           a_element_op,
                                           b_element_op,
                                           cde_element_op);
     }
  
     // polymorphic
     std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
     {
         return std::make_unique<Invoker>(Invoker{});
     }
  
     // polymorphic
     std::string GetTypeString() const override
     {
         auto str = std::stringstream();
  
         // clang-format off
         str << "DeviceGemm_Xdl_WaveletModel_CShuffle"
             << "<"
         << TileLoadThreadGroupSize << ", "
             << TileMathThreadGroupSize << ", "
             << MPerBlock << ", "
             << NPerBlock << ", "
             << KPerBlock << ", "
             << AK1 << ", "
             << BK1
             << ">";
         // clang-format on
  
         return str.str();
     }
 };
  
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck