/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/ops/gemm/kernel/gemm_tile_partitioner.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 "ck_tile/core.hpp"
 #include "ck_tile/ops/common.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_DEVICE 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);
     }
  
     [[nodiscard]] CK_TILE_DEVICE static auto
     GetOffsetedTileIndex(index_t block_start, index_t M, index_t N, index_t block_idx) noexcept
         -> const tuple<index_t, index_t>
     {
         const auto [iM, iN] = TilePartitioner{M, N}.GetOutputTileIndex(block_idx - 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_DEVICE 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;
 };
  
 template <typename BlockGemmShapeType,
           StreamKReductionStrategy ReductionStrategy = ck_tile::StreamKReductionStrategy::Atomic,
           uint32_t TileSwizzleSubM                   = 8>
 struct StreamKTilePartitioner
 {
     using BlockGemmShape = BlockGemmShapeType;
  
     static constexpr uint32_t MPerBlock = BlockGemmShape::kM;
     static constexpr uint32_t NPerBlock = BlockGemmShape::kN;
     static constexpr uint32_t KPerBlock = BlockGemmShape::kK;
  
     CK_TILE_HOST_DEVICE StreamKTilePartitioner() noexcept = delete;
  
     CK_TILE_HOST_DEVICE StreamKTilePartitioner(uint32_t M,
                                                uint32_t N,
                                                uint32_t K,
                                                uint32_t num_cu,
                                                uint32_t occupancy,
                                                uint32_t sk_blocks = 0xffffffff) noexcept
         : M_(M), N_(N), K_(K)
     {
         num_tile_m_ = integer_divide_ceil(M, MPerBlock);
         num_tile_n_ = integer_divide_ceil(N, NPerBlock);
         num_tile_k_ = integer_divide_ceil(K, KPerBlock);
  
         constexpr uint32_t min_k_iters_per_sk_block = 2;
         uint32_t num_tiles                          = num_tile_m_ * num_tile_n_;
         k_iters_per_tile                            = mdiv(num_tile_k_);
  
         // one cu can hold one wg at one time, from the whole cZ's point of view
         // if number of wg is same as num_cu, we call it 1 dispatch
         // if number of wg is 2x num_cu, we call it 2 dispatches.
         // one dispatch can deliver wg same as num_cu (full dispatch), or less than num_cu (partial
         // dispatch)
         //
         const uint32_t full_dispatches        = num_tiles / num_cu;
         const uint32_t full_dispatch_tiles    = full_dispatches * num_cu;
         const uint32_t partial_dispatch_tiles = num_tiles - full_dispatch_tiles;
  
         uint32_t sk_occupancy = occupancy;
         uint32_t dp_tiles     = full_dispatch_tiles;
         uint32_t sk_tiles     = partial_dispatch_tiles;
  
         if(full_dispatches < occupancy)
         {
             // in this case, we allocate all blocks as sk blocks
             // sk_occupancy = occupancy - full_dispatches;
             sk_occupancy = 1;
             dp_tiles     = full_dispatch_tiles;
             sk_tiles     = partial_dispatch_tiles;
         }
         else if((occupancy > 1) && (full_dispatches % occupancy == occupancy - 1))
         {
             // e.g. occupancy = 2, full_dispatches = 3, 5, 7 ...
             //      occupancy = 3, full_dispatches = 5, 8, 11 ...
             //      occupancy = 4, full_dispatches = 7, 11 ...
             sk_occupancy = 1; // left 1 slot for sk occupancy
             dp_tiles     = full_dispatch_tiles;
             sk_tiles     = partial_dispatch_tiles;
         }
         else
         {
             // otherwise, we reduce 1 dispatch from dp, together with partial dispatch,
             // to construct sk dispatch
             sk_occupancy = occupancy - ((full_dispatches - 1) % occupancy);
             dp_tiles     = full_dispatch_tiles - num_cu;
             sk_tiles     = partial_dispatch_tiles + num_cu;
         }
  
         // uint32_t dp_iters_per_block = k_iters_per_tile.get();
         uint32_t sk_total_iters = k_iters_per_tile.get() * sk_tiles;
         uint32_t dp_num_blocks  = 0;
  
         {
             const uint32_t min_sk_tiles = (sk_tiles >= num_cu) ? num_cu : (sk_tiles + 1);
             const uint32_t max_sk_tiles =
                 (sk_tiles >= num_cu) ? num_cu * sk_occupancy
                                      : min(num_cu, sk_total_iters / min_k_iters_per_sk_block);
  
             // if use dp for sk-block, how many iters do we need
             const uint32_t dp_for_sk_iters = k_iters_per_tile.get();
  
             uint32_t best_sk_score =
                 std::numeric_limits<int>::max(); // we need to find the smallest sk iters
             for(uint32_t tentative_sk_blocks = min_sk_tiles; tentative_sk_blocks < max_sk_tiles;
                 tentative_sk_blocks++)
             {
                 const uint32_t tentative_sk_iters_per_block =
                     (sk_total_iters + tentative_sk_blocks - 1) / tentative_sk_blocks;
                 const uint32_t tentative_sk_iters = tentative_sk_iters_per_block;
                 const uint32_t sk_blocks_per_tile = (tentative_sk_blocks + sk_tiles - 1) / sk_tiles;
  
                 //       the more sk_blocks_per_tile, the worse the overhead
                 uint32_t cross_sk_blocks_overhead = sk_blocks_per_tile;
                 if(tentative_sk_blocks % sk_tiles != 0)
                 {
                     // penalty for uneven divide
                     cross_sk_blocks_overhead +=
                         sk_blocks_per_tile * tentative_sk_iters_per_block / 50;
                 }
  
                 const uint32_t tentative_sk_score = tentative_sk_iters + cross_sk_blocks_overhead;
  
                 if(tentative_sk_score < best_sk_score)
                 {
                     best_sk_score = tentative_sk_score;
                     sk_num_blocks = tentative_sk_blocks;
                 }
             }
  
             if(best_sk_score >= dp_for_sk_iters)
             {
                 sk_num_blocks = 0;
             }
  
             // give a chance to control num of sk blocks
             sk_num_blocks = sk_blocks != 0xffffffff ? sk_blocks : sk_num_blocks;
  
             if(sk_num_blocks == 0)
             {
                 sk_num_big_blocks     = 0;
                 k_iters_per_big_block = 0;
  
                 dp_num_blocks      = num_tiles; // all tile to be dp block
                 dp_start_block_idx = 0;
                 sk_total_iters     = 0; // clear this tiles
             }
             else
             {
                 // k_iters_per_sk_block is the floor of avg each ck block loop over tiles.
                 // we need to decide how many iters for each sk block
                 // let m = k_iters_per_sk_block
                 // some of the sk block (little) will cover m iters, some (big) will cover m+1
                 // we have
                 // 1) l + b = sk_blocks
                 // 2) l * m + b * (m + 1) = sk_total_iters
                 //      => (l + b) * m + b = sk_total_iters
                 //      => sk_blocks * m + b = sk_total_iters
                 //      => b = sk_total_iters - m * sk_blocks
                 //      NOTE: big could be zero
                 const uint32_t k_iters_per_sk_block = sk_total_iters / sk_num_blocks;
                 sk_num_big_blocks     = sk_total_iters - k_iters_per_sk_block * sk_num_blocks;
                 k_iters_per_big_block = k_iters_per_sk_block + 1;
  
                 dp_num_blocks      = dp_tiles;
                 dp_start_block_idx = (sk_num_blocks + num_cu - 1) / num_cu * num_cu;
             }
         }
         n_tiles                   = mdiv2(num_tile_n_);
         reduction_start_block_idx = dp_start_block_idx + dp_num_blocks;
  
         if constexpr(ReductionStrategy == ck_tile::StreamKReductionStrategy::Reduction)
         {
             const uint32_t upper_big    = lcm(k_iters_per_big_block, k_iters_per_tile.get());
             const uint32_t upper_little = lcm(k_iters_per_big_block - 1, k_iters_per_tile.get());
             equiv_tiles_big             = mdiv(upper_big / k_iters_per_tile.get());
             equiv_tiles_little          = mdiv(upper_little / k_iters_per_tile.get());
         }
     }
  
     CK_TILE_HOST auto GridSize() const noexcept -> dim3
     {
         if constexpr(ReductionStrategy == ck_tile::StreamKReductionStrategy::Reduction)
         {
             return dim3(reduction_start_block_idx + GetSkTiles(), 1, 1);
         }
         else
             return dim3(reduction_start_block_idx, 1, 1);
     }
  
     CK_TILE_HOST_DEVICE static auto GetLoopNum(uint32_t K) noexcept -> uint32_t
     {
         return integer_divide_ceil(K, KPerBlock); // Stream-K processes one K-slice at a time
     }
  
     CK_TILE_DEVICE auto
     GetOutputTileIndex(uint32_t tile_idx) const noexcept -> tuple<uint32_t, uint32_t>
     {
         uint32_t m_tile_idx, n_tile_idx;
         n_tiles.divmod(tile_idx, num_tile_n_, m_tile_idx, n_tile_idx);
  
         // swizzle tile
  
         uint32_t tile_swizzle_sub_m_rem = num_tile_m_ % TileSwizzleSubM;
  
         const auto sub_m_adapt = (m_tile_idx < (num_tile_m_ - tile_swizzle_sub_m_rem))
                                      ? TileSwizzleSubM
                                      : tile_swizzle_sub_m_rem;
  
         uint32_t m_tile_idx_sub0, m_tile_idx_sub1;
         m_tile_idx_sub0 = m_tile_idx / TileSwizzleSubM;
         m_tile_idx_sub1 = m_tile_idx % TileSwizzleSubM;
  
         uint32_t tile_idx_local = n_tile_idx + m_tile_idx_sub1 * num_tile_n_;
  
         uint32_t m_tile_idx_with_adapt, n_tile_idx_with_adapt;
  
         n_tile_idx_with_adapt = tile_idx_local / sub_m_adapt;
         m_tile_idx_with_adapt = tile_idx_local % sub_m_adapt;
         return make_tuple(m_tile_idx_with_adapt + m_tile_idx_sub0 * TileSwizzleSubM,
                           n_tile_idx_with_adapt);
     }
  
     CK_TILE_DEVICE void
     GetBlockItr(uint32_t block_idx, uint32_t& iter_start, uint32_t& iter_end) const noexcept
     {
         if(block_idx < sk_num_big_blocks)
         {
             iter_start = block_idx * k_iters_per_big_block;
             iter_end   = iter_start + k_iters_per_big_block;
         }
         else if(block_idx < sk_num_blocks)
         {
             iter_start = (sk_num_big_blocks * k_iters_per_big_block) +
                          (block_idx - sk_num_big_blocks) * (k_iters_per_big_block - 1);
             iter_end = iter_start + (k_iters_per_big_block - 1);
         }
         else if(block_idx >= dp_start_block_idx)
         {
             uint32_t sk_total_iters     = GetSkTotalIters();
             uint32_t dp_iters_per_block = k_iters_per_tile.get();
             iter_start = sk_total_iters + (block_idx - dp_start_block_idx) * dp_iters_per_block;
             iter_end   = iter_start + dp_iters_per_block;
         }
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetSkTotalIters() const noexcept
     {
         uint32_t sk_total_iters = sk_num_big_blocks * k_iters_per_big_block +
                                   (sk_num_blocks - sk_num_big_blocks) * (k_iters_per_big_block - 1);
         return sk_total_iters;
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetSkTiles() const noexcept
     {
         // tiles for sk
         uint32_t sk_total_iters = GetSkTotalIters();
         return k_iters_per_tile.div(sk_total_iters);
     }
  
     CK_TILE_DEVICE uint32_t GetCurrentIterLength(uint32_t iter_start,
                                                  uint32_t iter_end,
                                                  uint32_t total_iter_length) const noexcept
     {
         uint32_t iter_length_mod, iter_length_quo /*unused*/;
         k_iters_per_tile.divmod(iter_end, iter_length_quo, iter_length_mod);
         uint32_t total_iter_length_val = static_cast<uint32_t>(total_iter_length);
         uint32_t current_iter_length =
             min(iter_length_mod == 0 ? (iter_end - iter_start) : iter_length_mod,
                 total_iter_length_val);
         return current_iter_length;
     }
  
     CK_TILE_DEVICE uint32_t GetTileIdx(uint32_t iter) const noexcept
     {
         return k_iters_per_tile.div(iter);
     }
  
     CK_TILE_DEVICE void
     GetTileIdxWithOffset(uint32_t iter, uint32_t& tile_idx, uint32_t& iter_offset) const noexcept
     {
         uint32_t tile_idx_val    = static_cast<uint32_t>(tile_idx);
         uint32_t iter_offset_val = static_cast<uint32_t>(iter_offset);
         k_iters_per_tile.divmod(iter, tile_idx_val, iter_offset_val);
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetWorkSpaceSizeForAcc(uint32_t acc_element_bytes) const noexcept
     {
         static constexpr uint32_t alignment = 128;
         uint32_t acc_buffer_bytes =
             MPerBlock * NPerBlock * GetTotalAccBuffers() * acc_element_bytes;
         return (acc_buffer_bytes + alignment - 1) / alignment * alignment;
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetWorkSpaceSizeForSemaphore() const noexcept
     {
         return GetSkTiles() * sizeof(uint32_t);
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetWorkSpaceSize(uint32_t acc_element_bytes) const noexcept
     {
         return GetWorkSpaceSizeForAcc(acc_element_bytes) + GetWorkSpaceSizeForSemaphore();
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetTileIntersections(uint32_t tiles_,
                                                       const mdiv& equiv_tiles_) const noexcept
     {
         uint32_t tile_idx_        = tiles_ == 0 ? 0 : (tiles_ - 1);
         uint32_t max_equiv_tiles_ = equiv_tiles_.get() - 1;
         uint32_t quo_, rem_;
         equiv_tiles_.divmod(tile_idx_, quo_, rem_);
         return quo_ * max_equiv_tiles_ + rem_;
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetTilesCoverSkBlock(uint32_t num_sk_blocks_,
                                                       uint32_t iters_per_sk_block_) const noexcept
     {
         return k_iters_per_tile.div(num_sk_blocks_ * iters_per_sk_block_ + k_iters_per_tile.get() -
                                     1);
     }
  
     CK_TILE_HOST_DEVICE uint32_t GetTotalAccBuffers() const noexcept
     {
         uint32_t tiles_cover_big_blocks =
             GetTilesCoverSkBlock(sk_num_big_blocks, k_iters_per_big_block);
         uint32_t tiles_cover_little_blocks =
             GetTilesCoverSkBlock(sk_num_blocks - sk_num_big_blocks, k_iters_per_big_block - 1);
  
         uint32_t total_intersec_big = GetTileIntersections(tiles_cover_big_blocks, equiv_tiles_big);
         uint32_t total_intersec_little =
             GetTileIntersections(tiles_cover_little_blocks, equiv_tiles_little);
  
         return sk_num_blocks + total_intersec_big + total_intersec_little;
     }
  
     CK_TILE_DEVICE uint32_t GetAccBufferOffsetFromTile(uint32_t tile_idx_) const noexcept
     {
         uint32_t tiles_cover_big_blocks =
             GetTilesCoverSkBlock(sk_num_big_blocks, k_iters_per_big_block);
         if(tile_idx_ < tiles_cover_big_blocks)
         {
             uint32_t touched_sk_blocks =
                 (tile_idx_ * k_iters_per_tile.get() + k_iters_per_big_block - 1) /
                 k_iters_per_big_block;
             uint32_t current_intersec = GetTileIntersections(tile_idx_, equiv_tiles_big);
             return touched_sk_blocks + current_intersec;
         }
         else
         {
             uint32_t iters_per_little_sk_block = k_iters_per_big_block - 1;
             uint32_t tile_idx_little_reverse   = GetSkTiles() - tile_idx_;
             uint32_t touched_sk_blocks =
                 (tile_idx_little_reverse * k_iters_per_tile.get() + iters_per_little_sk_block - 1) /
                 iters_per_little_sk_block;
             uint32_t current_intersec =
                 GetTileIntersections(tile_idx_little_reverse, equiv_tiles_little);
             return GetTotalAccBuffers() - (touched_sk_blocks + current_intersec);
         }
     }
  
     CK_TILE_DEVICE uint32_t GetAccBufferOffsetFromBlock(uint32_t block_idx_) const noexcept
     {
         uint32_t iters_per_big_sk_block    = k_iters_per_big_block;
         uint32_t iters_per_little_sk_block = k_iters_per_big_block - 1;
         if(block_idx_ < sk_num_big_blocks)
         {
             uint32_t touched_tiles    = k_iters_per_tile.div(block_idx_ * iters_per_big_sk_block +
                                                           k_iters_per_tile.get() - 1);
             uint32_t current_intersec = GetTileIntersections(touched_tiles, equiv_tiles_big);
             return block_idx_ + current_intersec;
         }
         else
         {
             uint32_t block_idx_little_reverse = sk_num_blocks - block_idx_;
             uint32_t touched_tiles            = k_iters_per_tile.div(
                 block_idx_little_reverse * iters_per_little_sk_block + k_iters_per_tile.get() - 1);
             uint32_t current_intersec = GetTileIntersections(touched_tiles, equiv_tiles_little);
             return GetTotalAccBuffers() - (block_idx_little_reverse + current_intersec);
         }
     }
  
     // Getters for problem dimensions
     CK_TILE_HOST_DEVICE uint32_t GetNumTileM() const noexcept { return num_tile_m_; }
     CK_TILE_HOST_DEVICE uint32_t GetNumTileN() const noexcept { return num_tile_n_; }
     CK_TILE_HOST_DEVICE uint32_t GetNumTileK() const noexcept { return num_tile_k_; }
  
     uint32_t sk_num_blocks;
     uint32_t sk_num_big_blocks;
     uint32_t dp_start_block_idx;
     uint32_t reduction_start_block_idx;
     uint32_t k_iters_per_big_block;
     mdiv2 n_tiles;
     mdiv k_iters_per_tile;
     mdiv equiv_tiles_big;    // for reduction
     mdiv equiv_tiles_little; // for reduction
  
     private:
     uint32_t M_, N_, K_;
     uint32_t num_tile_m_, num_tile_n_, num_tile_k_;
 };
 } // namespace ck_tile