Skip to content

[CK_TILE] Blockwise GEMM Pipeline V5 #2360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 16 commits into
base: develop
Choose a base branch
from
+942 −205
Open

[CK_TILE] Blockwise GEMM Pipeline V5 #2360

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
7c63567
First checkpoint
aledudek May 8, 2025
06deffc
Second checkpoint - hot loop scheduler
aledudek May 8, 2025
9d63430
Third checkpoint - init main operator
aledudek May 9, 2025
71bb790
Fourth checkpoint - main loop ready
aledudek May 9, 2025
a8d47fb
Fifth checkpoint - main loop fix
aledudek May 9, 2025
5b3eff9
Sixth checkpoint - ReadWritecompFunc
aledudek May 9, 2025
b72a894
Seventh checkpoint - Tail finished
aledudek May 9, 2025
6f3441f
[CK_TILE] Blockwise gemm pipeline v5 complete
aledudek May 9, 2025
5b3388d
Merge branch 'develop' into _bgp_v5_a
aledudek May 9, 2025
a1e7960
Working
aledudek Jun 17, 2025
0ba66f0
Working fixes 2
aledudek Jun 17, 2025
35e372e
Change name to avoid conflict
aledudek Jun 20, 2025
4c0a5d9
Merge branch 'develop' into _bgp_v5_b
aledudek Jun 20, 2025
2c0d5b4
Merge branch 'develop' into _bgp_v5_b
aledudek Jun 20, 2025
573b676
Remove maybe unused has hot loop in tailhandler
aledudek Jun 20, 2025
654ff5a
Change naming to more accurate
aledudek Jun 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 30 additions & 1 deletion example/ck_tile/03_gemm/gemm_utils.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#define CK_TILE_PIPELINE_MEMORY 2
#define CK_TILE_PIPELINE_COMPUTE_V4 3
#define CK_TILE_PIPELINE_COMPUTE_V5 4
#define CK_TILE_PIPELINE_COMPUTE_V6 5

struct GemmConfigBase
{
Expand Down Expand Up @@ -178,6 +179,25 @@ struct GemmConfigComputeV4_1 : public GemmConfigBase

template <typename PrecType>
struct GemmConfigComputeV5 : public GemmConfigBase
{
static constexpr ck_tile::index_t M_Tile = 256;
static constexpr ck_tile::index_t N_Tile = 256;
static constexpr ck_tile::index_t K_Tile = 32;

static constexpr ck_tile::index_t M_Warp = 2;
static constexpr ck_tile::index_t N_Warp = 2;
static constexpr ck_tile::index_t K_Warp = 1;

static constexpr ck_tile::index_t M_Warp_Tile = 32;
static constexpr ck_tile::index_t N_Warp_Tile = 32;
static constexpr ck_tile::index_t K_Warp_Tile = 16;

static constexpr bool DoubleSmemBuffer = false;
static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V5;
};

template <typename PrecType>
struct GemmConfigComputeV6 : public GemmConfigBase
{
static constexpr ck_tile::index_t M_Tile = 128;
static constexpr ck_tile::index_t N_Tile = 128;
Expand All @@ -192,7 +212,7 @@ struct GemmConfigComputeV5 : public GemmConfigBase
static constexpr ck_tile::index_t K_Warp_Tile = sizeof(PrecType) == 2 ? 16 : 64;

static constexpr bool DoubleSmemBuffer = false;
static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V5;
static constexpr ck_tile::index_t Pipeline = CK_TILE_PIPELINE_COMPUTE_V6;
static constexpr ck_tile::index_t NumWaNumWaveGroups = 2;
};

Expand Down Expand Up @@ -329,6 +349,15 @@ struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V5>
using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV5<PipelineProblem>;
};

template <>
struct PipelineTypeTraits<CK_TILE_PIPELINE_COMPUTE_V6>
{
template <typename PipelineProblem>
using GemmPipeline = ck_tile::GemmPipelineAgBgCrCompV6<PipelineProblem>;
template <typename PipelineProblem>
using UniversalGemmPipeline = ck_tile::BaseGemmPipelineAgBgCrCompV6<PipelineProblem>;
};

auto create_args(int argc, char* argv[])
{
ck_tile::ArgParser arg_parser;
Expand Down
3 changes: 1 addition & 2 deletions example/ck_tile/03_gemm/universal_gemm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -324,13 +324,12 @@ int main(int argc, char* argv[])
{
try
{
return !run_gemm_example<GemmConfigComputeV3>(argc, argv);
return !run_gemm_example<GemmConfigComputeV5>(argc, argv);
}
catch(const std::runtime_error& e)
{
std::cerr << "Caught runtime error: " << e.what() << '\n';
// Return a non-zero code to indicate failure
return EXIT_FAILURE;
}
return EXIT_SUCCESS;
}
2 changes: 2 additions & 0 deletions include/ck_tile/ops/gemm.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v4_default_policy.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v6_default_policy.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_mem.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp"
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp"
Expand Down
Loading