block_gemm_areg_breg_creg_v2_custom_policy.hpp Source File

block_gemm_areg_breg_creg_v2_custom_policy.hpp Source File#

Composable Kernel: block_gemm_areg_breg_creg_v2_custom_policy.hpp Source File

// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
#pragma once
 
#include "ck_tile/core.hpp"
 
namespace ck_tile {
 
enum struct GemmLoopOrder
{
    KMN,
    MNK,
};
 
template <typename AType_,
          typename BType_,
          typename CType_,
          typename BlockWarps_,
          typename WarpGemm_,
          GemmLoopOrder BlockGemmLoopOrder_ = GemmLoopOrder::KMN>
struct BlockGemmARegBRegCRegV2CustomPolicy
{
    using AType = remove_cvref_t<AType_>;
    using BType = remove_cvref_t<BType_>;
    using CType = remove_cvref_t<CType_>;
 
    using BlockWarps = remove_cvref_t<BlockWarps_>;
 
    static constexpr index_t kMWarps = BlockWarps::at(number<0>{});
    static constexpr index_t kNWarps = BlockWarps::at(number<1>{});
    static constexpr index_t kKWarps = BlockWarps::at(number<2>{});
 
    using WarpGemm = remove_cvref_t<WarpGemm_>;
 
    static constexpr auto BlockGemmLoopOrder = BlockGemmLoopOrder_;
 
    template <typename Problem>
    CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemmMWarpNWarp()
    {
        return make_tuple(WarpGemm{}, kMWarps, kNWarps);
    }
};
 
} // namespace ck_tile