gemm_tile_partitioner.hpp Source File

gemm_tile_partitioner.hpp Source File#

Composable 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"
#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 = amd_wave_read_first_lane(blockIdx);
        const index_t iN = amd_wave_read_first_lane(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 = amd_wave_read_first_lane(blockIdx / NBlocks);
        const index_t iN = amd_wave_read_first_lane(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) const noexcept
    {
        // A WG's iter_end is either in the current C macro tile or not.
        // If it is not, then the macro tile boundary is where the WG must stop.
        uint32_t distance_to_tile_boundary =
            k_iters_per_tile.get() - (iter_start % k_iters_per_tile.get());
        return min(iter_start + distance_to_tile_boundary, iter_end) - iter_start;
    }
    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
    {
        k_iters_per_tile.divmod(iter, tile_idx, iter_offset);
    }
    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