mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-11 01:27:34 +00:00
Add a lot of GEMM template, commiting to save progress.
This commit has a lot of template code extracted from the gemm universal example, but it needs more debugging before we can compile the kernel object.
This commit is contained in:
John Shumway
@@ -336,7 +336,7 @@ find_package(Threads REQUIRED)
|
||||
link_libraries(Threads::Threads)
|
||||
|
||||
## C++
|
||||
set(CMAKE_CXX_STANDARD 17)
|
||||
set(CMAKE_CXX_STANDARD 20)
|
||||
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
||||
set(CMAKE_CXX_EXTENSIONS OFF)
|
||||
message(STATUS "CMAKE_CXX_COMPILER: ${CMAKE_CXX_COMPILER}")
|
||||
@@ -568,6 +568,7 @@ SET(BUILD_DEV ON CACHE BOOL "BUILD_DEV")
|
||||
if(BUILD_DEV)
|
||||
add_compile_options(-Werror)
|
||||
add_compile_options(-Weverything)
|
||||
add_compile_options(-Wno-c++20-compat)
|
||||
endif()
|
||||
message(STATUS "CMAKE_CXX_FLAGS: ${CMAKE_CXX_FLAGS}")
|
||||
|
||||
|
||||
@@ -2,34 +2,37 @@
|
||||
#include <iostream>
|
||||
#include <concepts>
|
||||
|
||||
namespace ck_tile::builder {
|
||||
#include "ck_tile/host.hpp"
|
||||
#include "ck_tile/ops/gemm/pipeline/tile_gemm_shape.hpp"
|
||||
#include "ck_tile/ops/gemm/kernel/gemm_tile_partitioner.hpp"
|
||||
#include "ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp"
|
||||
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp"
|
||||
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v3.hpp"
|
||||
#include "ck_tile/ops/gemm/kernel/gemm_kernel.hpp"
|
||||
#include "ck_tile/ops/epilogue/cshuffle_epilogue.hpp"
|
||||
|
||||
using index_t = std::size_t;
|
||||
namespace ck_tile::builder {
|
||||
|
||||
// Some sample host args to describe a GEMM, with defaults set to zero.
|
||||
struct GemmHostArgs
|
||||
{
|
||||
index_t m = 0;
|
||||
index_t n = 0;
|
||||
index_t k = 0;
|
||||
index_t lda = 0;
|
||||
index_t ldb = 0;
|
||||
index_t ldc = 0;
|
||||
index_t k_batch_ = 0;
|
||||
const void* a = nullptr;
|
||||
const void* b = nullptr;
|
||||
void* c = nullptr;
|
||||
ck_tile::index_t m = 0;
|
||||
ck_tile::index_t n = 0;
|
||||
ck_tile::index_t k = 0;
|
||||
ck_tile::index_t lda = 0;
|
||||
ck_tile::index_t ldb = 0;
|
||||
ck_tile::index_t ldc = 0;
|
||||
index_t k_batch_ = 0;
|
||||
const void* a = nullptr;
|
||||
const void* b = nullptr;
|
||||
void* c = nullptr;
|
||||
};
|
||||
|
||||
// Tag for column major layout.
|
||||
struct ColMajor
|
||||
{
|
||||
};
|
||||
using ColMajor = ck_tile::tensor_layout::gemm::ColumnMajor;
|
||||
|
||||
// Tag for row major layout.
|
||||
struct RowMajor
|
||||
{
|
||||
};
|
||||
using RowMajor = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
|
||||
// Requirements for struct to define the data types used in the GEMM operation.
|
||||
//
|
||||
@@ -76,14 +79,144 @@ class Gemm
|
||||
}
|
||||
};
|
||||
|
||||
// Returns the GemmUniversal GemmConfig.
|
||||
//
|
||||
// Captures all the madness in tile_example_gemm_universal!
|
||||
template <DefinesGemmTypes Types>
|
||||
struct GemmConfigForTypes
|
||||
{
|
||||
using PrecType = Types::AccDataType;
|
||||
|
||||
static constexpr int CK_TILE_PIPELINE_COMPUTE_V3 = 1;
|
||||
static consteval auto get_k_warp_tile(auto M_Warp_Tile)
|
||||
{
|
||||
return (M_Warp_Tile == 32) ? 16 : 32;
|
||||
}
|
||||
// Use GemmConfigComputeV3 from tile_example_gemm_universal.
|
||||
struct GemmConfig
|
||||
{
|
||||
// Compute V3 only support Intrawave scheduler
|
||||
static constexpr ck_tile::index_t M_Tile = 16;
|
||||
static constexpr ck_tile::index_t N_Tile = 64;
|
||||
static constexpr ck_tile::index_t K_Tile = 256 / sizeof(PrecType);
|
||||
|
||||
static constexpr ck_tile::index_t M_Warp = 1;
|
||||
static constexpr ck_tile::index_t N_Warp = 4;
|
||||
static constexpr ck_tile::index_t K_Warp = 1;
|
||||
|
||||
static constexpr ck_tile::index_t M_Warp_Tile = 16;
|
||||
static constexpr ck_tile::index_t N_Warp_Tile = 16;
|
||||
static constexpr ck_tile::index_t K_Warp_Tile = get_k_warp_tile(M_Warp_Tile);
|
||||
|
||||
static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V3;
|
||||
|
||||
static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
|
||||
|
||||
static constexpr bool kPadM = false;
|
||||
static constexpr bool kPadN = false;
|
||||
static constexpr bool kPadK = false;
|
||||
|
||||
static constexpr bool PermuteA = false;
|
||||
static constexpr bool PermuteB = false;
|
||||
static constexpr bool TransposeC = false;
|
||||
|
||||
static constexpr bool UseStructuredSparsity = false;
|
||||
|
||||
static constexpr int kBlockPerCu = 1;
|
||||
static constexpr ck_tile::index_t TileParitionerGroupNum = 8;
|
||||
static constexpr ck_tile::index_t TileParitionerM01 = 4;
|
||||
static constexpr ck_tile::index_t NumWaveGroups = 1;
|
||||
|
||||
static constexpr bool Preshuffle = false;
|
||||
static constexpr bool DoubleSmemBuffer = false;
|
||||
};
|
||||
};
|
||||
|
||||
// A minimal GEMM builder, this is where all the work will be.
|
||||
template <DefinesGemmTypes Types, DefinesGemmLayout Layout>
|
||||
class GemmBuilder
|
||||
struct GemmBuilder
|
||||
{
|
||||
public:
|
||||
using value = Gemm;
|
||||
using types_type = Types;
|
||||
using layout_type = Layout;
|
||||
|
||||
static constexpr bool PERSISTENT = false;
|
||||
|
||||
using GemmConfig = typename GemmConfigForTypes<Types>::GemmConfig;
|
||||
|
||||
using GemmShape = ck_tile::TileGemmShape<
|
||||
ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
|
||||
ck_tile::sequence<GemmConfig::M_Warp, GemmConfig::N_Warp, GemmConfig::K_Warp>,
|
||||
ck_tile::
|
||||
sequence<GemmConfig::M_Warp_Tile, GemmConfig::N_Warp_Tile, GemmConfig::K_Warp_Tile>,
|
||||
GemmConfig::PermuteA,
|
||||
GemmConfig::PermuteB>;
|
||||
|
||||
using TilePartitioner =
|
||||
ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
|
||||
GemmConfig::TileParitionerGroupNum,
|
||||
GemmConfig::TileParitionerM01>;
|
||||
|
||||
using Traits = ck_tile::TileGemmTraits<GemmConfig::kPadM,
|
||||
GemmConfig::kPadN,
|
||||
GemmConfig::kPadK,
|
||||
typename Layout::ALayout,
|
||||
typename Layout::BLayout,
|
||||
typename Layout::CLayout,
|
||||
GemmConfig::NumWaveGroups>;
|
||||
|
||||
using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<GemmConfig::kPadM,
|
||||
GemmConfig::kPadN,
|
||||
GemmConfig::kPadK,
|
||||
GemmConfig::DoubleSmemBuffer,
|
||||
typename Layout::ALayout,
|
||||
typename Layout::BLayout,
|
||||
typename Layout::CLayout,
|
||||
GemmConfig::TransposeC,
|
||||
GemmConfig::UseStructuredSparsity,
|
||||
PERSISTENT,
|
||||
GemmConfig::NumWaveGroups,
|
||||
GemmConfig::Preshuffle>;
|
||||
|
||||
using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<typename Types::ADataType,
|
||||
typename Types::BDataType,
|
||||
typename Types::AccDataType,
|
||||
GemmShape,
|
||||
GemmUniversalTraits,
|
||||
GemmConfig::Scheduler>;
|
||||
|
||||
using GemmPipelineProblem = ck_tile::GemmPipelineProblem<typename Types::ADataType,
|
||||
typename Types::BDataType,
|
||||
typename Types::AccDataType,
|
||||
GemmShape,
|
||||
Traits>;
|
||||
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
|
||||
|
||||
using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV3<GemmPipelineProblem>;
|
||||
|
||||
using GemmEpilogue = ck_tile::CShuffleEpilogue<
|
||||
ck_tile::CShuffleEpilogueProblem<typename Types::ADataType,
|
||||
typename Types::BDataType,
|
||||
ck_tile::tuple<>,
|
||||
typename Types::AccDataType,
|
||||
typename Types::CDataType,
|
||||
ck_tile::tuple<>,
|
||||
typename Layout::CLayout,
|
||||
ck_tile::element_wise::PassThrough,
|
||||
UniversalGemmProblem::kBlockSize,
|
||||
TilePartitioner::MPerBlock,
|
||||
TilePartitioner::NPerBlock,
|
||||
GemmConfig::M_Warp,
|
||||
GemmConfig::N_Warp,
|
||||
GemmConfig::M_Warp_Tile,
|
||||
GemmConfig::N_Warp_Tile,
|
||||
GemmConfig::K_Warp_Tile,
|
||||
UniversalGemmProblem::TransposeC,
|
||||
ck_tile::memory_operation_enum::set,
|
||||
GemmConfig::NumWaveGroups>>;
|
||||
|
||||
using Kernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmPipeline>;
|
||||
};
|
||||
|
||||
} // namespace ck_tile::builder
|
||||
|
||||
@@ -63,7 +63,10 @@ struct MyGemmLayout
|
||||
using CLayout = ckb::RowMajor;
|
||||
};
|
||||
|
||||
using MyGemm = ckb::GemmBuilder<MyGemmTypes, MyGemmLayout>::value;
|
||||
using Builder = ckb::GemmBuilder<MyGemmTypes, MyGemmLayout>;
|
||||
|
||||
using Gemm = Builder::value;
|
||||
using Kernel = Builder::Kernel;
|
||||
|
||||
} // namespace example
|
||||
|
||||
@@ -71,7 +74,13 @@ int main()
|
||||
{
|
||||
// Create the GEMM kernel.
|
||||
const int M = 1024, N = 2048, K = 64;
|
||||
example::MyGemm gemm;
|
||||
example::Gemm gemm;
|
||||
|
||||
// Describe the GEMM kernel:
|
||||
std::cout << "Shape: " << example::Builder::GemmShape::GetName() << std::endl;
|
||||
std::cout << "Problem: " << example::Builder::UniversalGemmProblem::GetName() << std::endl;
|
||||
// std::cout << "Pipeline: " << example::Builder::GemmPipeline::GetName() << std::endl;
|
||||
// std::cout << "Kernel name: " << Kernel::GetName() << std::endl;
|
||||
|
||||
// Try GPU execution.
|
||||
try
|
||||
|
||||
@@ -7,6 +7,9 @@
|
||||
#include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp"
|
||||
#include "ck_tile/ops/common/tensor_layout.hpp"
|
||||
|
||||
#include "ck_tile/ops/gemm/block/block_gemm_asmem_bsmem_creg_v1_custom_policy.hpp"
|
||||
#include "ck_tile/ops/gemm/block/block_universal_gemm_as_bs_cr.hpp"
|
||||
|
||||
namespace ck_tile {
|
||||
|
||||
template <typename Derived>
|
||||
|
||||
Reference in New Issue
Block a user