include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp Source File

include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp Source File#

Composable Kernel: include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp Source File
 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include "ck_tile/core.hpp"
  
 namespace ck_tile {
  
 template <typename BlockGemmShapeType>
 struct GemmTile2DPartitioner
 {
     using BlockGemmShape = remove_cvref_t<BlockGemmShapeType>;
  
     static constexpr index_t MPerBlock = BlockGemmShape::kM;
     static constexpr index_t NPerBlock = BlockGemmShape::kN;
     static constexpr index_t KPerBlock = BlockGemmShape::kK;
  
     CK_TILE_HOST_DEVICE GemmTile2DPartitioner() noexcept = delete;
     CK_TILE_HOST_DEVICE GemmTile2DPartitioner([[maybe_unused]] index_t M,
                                               [[maybe_unused]] index_t N) noexcept;
  
     CK_TILE_HOST static auto
     GridSize(index_t M, index_t N) noexcept(noexcept(MPerBlock != 0 && NPerBlock != 0)) -> dim3
     {
         const index_t GridDimX = (M + MPerBlock - 1) / MPerBlock;
         const index_t GridDimY = (N + NPerBlock - 1) / NPerBlock;
         return dim3(GridDimX, GridDimY, 1);
     }
  
     CK_TILE_HOST_DEVICE static auto GetLoopNum(index_t K) noexcept -> index_t
     {
         return integer_divide_ceil(K, KPerBlock);
     }
  
     CK_TILE_DEVICE static auto GetOutputTileIndex(index_t blockIdx, index_t blockIdy) noexcept
         -> const tuple<index_t, index_t>
     {
         const index_t iM = __builtin_amdgcn_readfirstlane(blockIdx);
         const index_t iN = __builtin_amdgcn_readfirstlane(blockIdy);
         return make_tuple(iM, iN);
     }
 };
  
 template <typename BlockGemmShape_>
 struct GemmTile1DPartitioner
 {
     using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
  
     static constexpr index_t MPerBlock = BlockGemmShape::kM;
     static constexpr index_t NPerBlock = BlockGemmShape::kN;
     static constexpr index_t KPerBlock = BlockGemmShape::kK;
  
     CK_TILE_HOST_DEVICE GemmTile1DPartitioner() noexcept = delete;
  
     CK_TILE_HOST_DEVICE GemmTile1DPartitioner([[maybe_unused]] index_t M, index_t N) noexcept
     {
         N_ = N;
     }
  
     CK_TILE_HOST static auto
     GridSize(index_t M, index_t N) noexcept(noexcept(MPerBlock != 0 && NPerBlock != 0)) -> index_t
     {
         const index_t GridDimX = (M + MPerBlock - 1) / MPerBlock;
         const index_t GridDimY = (N + NPerBlock - 1) / NPerBlock;
         return GridDimX * GridDimY;
     }
  
     CK_TILE_HOST_DEVICE static auto GetLoopNum(index_t K) noexcept -> index_t
     {
         return integer_divide_ceil(K, KPerBlock);
     }
  
     CK_TILE_DEVICE static auto GetOutputTileIndex(index_t blockIdx) noexcept
         -> const tuple<index_t, index_t>
     {
         const index_t NBlocks = integer_divide_ceil(N_, NPerBlock);
  
         const index_t iM = __builtin_amdgcn_readfirstlane(blockIdx / NBlocks);
         const index_t iN = __builtin_amdgcn_readfirstlane(blockIdx - iM * NBlocks);
         return make_tuple(iM, iN);
     }
  
     private:
     CK_TILE_DEVICE static index_t N_;
 };
  
 template <typename, typename = void>
 struct HasFnOneArgImpl : std::false_type
 {
 };
  
 template <typename T>
 struct HasFnOneArgImpl<T, std::void_t<decltype(std::declval<T>().GetOutputTileIndex(1))>>
     : std::true_type
 {
 };
  
 template <typename TilePartitioner,
           typename = typename std::enable_if_t<HasFnOneArgImpl<TilePartitioner>{}>>
 struct OffsettedTile1DPartitioner
 {
     [[nodiscard]] CK_TILE_DEVICE static auto
     GetOffsetedTileIndex(index_t block_start, index_t M, index_t N) noexcept
         -> const tuple<index_t, index_t>
     {
         const auto [iM, iN] = TilePartitioner{M, N}.GetOutputTileIndex(blockIdx.x - block_start);
         return make_tuple(iM, iN);
     }
 };
  
 template <typename BlockGemmShapeType, index_t GroupNum, index_t M01>
 struct GemmSpatiallyLocalTilePartitioner
 {
     using BlockGemmShape = remove_cvref_t<BlockGemmShapeType>;
  
     static constexpr index_t MPerBlock = BlockGemmShape::kM;
     static constexpr index_t NPerBlock = BlockGemmShape::kN;
     static constexpr index_t KPerBlock = BlockGemmShape::kK;
  
     CK_TILE_HOST_DEVICE GemmSpatiallyLocalTilePartitioner() noexcept = delete;
     CK_TILE_HOST_DEVICE GemmSpatiallyLocalTilePartitioner(index_t M_, index_t N_) noexcept
         : M(M_), N(N_)
     {
     }
  
     CK_TILE_HOST static auto
     GridSize(index_t M, index_t N) noexcept(noexcept(MPerBlock != 0 && NPerBlock != 0)) -> index_t
     {
         const index_t GridDimX = integer_divide_ceil(M, MPerBlock);
         const index_t GridDimY = integer_divide_ceil(N, NPerBlock);
         return GridDimX * GridDimY;
     }
  
     CK_TILE_HOST_DEVICE static auto GetLoopNum(index_t K) noexcept -> index_t
     {
         return integer_divide_ceil(K, KPerBlock);
     }
  
     CK_TILE_DEVICE auto GetOutputTileIndex(index_t block_1d_id) noexcept
         -> const tuple<index_t, index_t>
     {
         const auto M0 = integer_divide_ceil(M, MPerBlock);
         const auto N0 = integer_divide_ceil(N, NPerBlock);
  
         if(M0 == 1)
         {
             return make_tuple(0, block_1d_id);
         }
         else if(N0 == 1)
         {
             return make_tuple(block_1d_id, 0);
         }
         // block_1d_id = block_1d_id % (M0 * N0); // swallow batch index
         else
         {
             const auto group_size    = integer_divide_ceil(M0 * N0, GroupNum);
             const auto big_group_num = GroupNum - (group_size * GroupNum - M0 * N0);
             const auto group_id_y    = block_1d_id / GroupNum;
             const auto group_id_x    = block_1d_id - group_id_y * GroupNum;
             const auto remap_block_1d_id =
                 group_id_x <= big_group_num
                     ? group_id_x * group_size + group_id_y
                     : group_id_x * group_size + big_group_num - group_id_x + group_id_y;
  
             const index_t idx_M0 = remap_block_1d_id / N0;
             const index_t idx_N0 = remap_block_1d_id - idx_M0 * N0;
  
             const index_t M0_tmp     = M0 / M01;
             const index_t M0_mod_M01 = M0 - M0_tmp * M01;
  
             const auto M01_adapt = (idx_M0 < M0 - M0_mod_M01) ? M01 : M0_mod_M01;
  
             const index_t idx_M00          = idx_M0 / M01;
             const index_t idx_M01          = idx_M0 - idx_M00 * M01;
             const index_t idx_N0_M01_local = idx_N0 + idx_M01 * N0;
  
             const index_t N_out           = idx_N0_M01_local / M01_adapt;
             const index_t idx_loc_mod_M01 = idx_N0_M01_local - N_out * M01_adapt;
  
             return make_tuple(idx_loc_mod_M01 + idx_M00 * M01, N_out);
         }
     }
  
     private:
     index_t M;
     index_t N;
 };
  
 } // namespace ck_tile