Skip to content

Commit d559ec0

Browse files
EnricoDegassistant-librarian[bot]
EnricoDeg
authored and
assistant-librarian[bot]
committed
[rocm-libraries] ROCm/rocm-libraries#8554 (commit be9af54)
refactor(ck): mx gemm kernel unification ## Motivation CK tile currently has two separate MX GEMM kernels for gfx950 and gfx1250. This pull request refactors and modernizes the MX GEMM kernel and example to use new scale tensor handling, improved kernel argument structures, and updated pipeline and kernel APIs. The changes simplify the interface and improve type safety. JIRA ID ROCM-26313 ## Technical Details - Add support for gfx950 in MX GEMM kernel for gfx1250 and remove unused kernel - Unify comp async pipeline for GEMM and MX GEMM - Unify eight waves pipeline for GEMM and MX GEMM - Move preshuffle MX GEMM pipeline to gemm ops and remove gemm_mx ops - Unify testing framework for MX GEMM - Add gfx950 tests for grouped MX GEMM ## Test Plan - `test_mx_gemm_async.cpp` for MX GEMM on gfx950 - `test_mx_grouped_gemm_comp_async.cpp` for grouped MX GEMM on gfx950 ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
1 parent 604c56b commit d559ec0

60 files changed

Lines changed: 3703 additions & 5217 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎example/ck_tile/42_mx_gemm/mx_gemm.cpp‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,13 @@ static constexpr inline auto is_row_major(Layout layout_)
2424
template <typename GemmConfig,
2525
typename ADataType,
2626
typename BDataType,
27+
typename AScaleDataType,
28+
typename BScaleDataType,
2729
typename AccDataType,
2830
typename CDataType,
2931
typename ALayout,
3032
typename BLayout,
3133
typename CLayout,
32-
typename ScaleM,
33-
typename ScaleN,
3434
bool UsePersistentKernel = false>
3535
float invoke_mx_gemm(ck_tile::DeviceMem& a_dev_buf,
3636
ck_tile::DeviceMem& b_dev_buf,
@@ -42,36 +42,38 @@ float invoke_mx_gemm(ck_tile::DeviceMem& a_dev_buf,
4242
ck_tile::index_t stride_B,
4343
ck_tile::index_t stride_C,
4444
ck_tile::index_t kbatch,
45-
ScaleM scale_m,
46-
ScaleN scale_n,
45+
ck_tile::DeviceMem& scale_m,
46+
ck_tile::DeviceMem& scale_n,
4747
int n_warmup,
4848
int n_repeat)
4949
{
50-
MXGemmHostArgs<ScaleM, ScaleN> args(a_dev_buf.GetDeviceBuffer(),
51-
b_dev_buf.GetDeviceBuffer(),
52-
c_dev_buf.GetDeviceBuffer(),
53-
kbatch,
54-
M,
55-
N,
56-
K,
57-
stride_A,
58-
stride_B,
59-
stride_C,
60-
scale_m,
61-
scale_n);
50+
ck_tile::MxGemmHostArgs<1, 1, 0> args({static_cast<const void*>(a_dev_buf.GetDeviceBuffer())},
51+
{static_cast<const void*>(scale_m.GetDeviceBuffer())},
52+
{static_cast<const void*>(b_dev_buf.GetDeviceBuffer())},
53+
{static_cast<const void*>(scale_n.GetDeviceBuffer())},
54+
{},
55+
c_dev_buf.GetDeviceBuffer(),
56+
kbatch,
57+
M,
58+
N,
59+
K,
60+
{stride_A},
61+
{stride_B},
62+
{},
63+
stride_C);
6264

6365
// Simplified invocation - comp_async handles hot loop and tail internally
6466
auto invoke_splitk_path = [&](auto split_k_) {
6567
return mx_gemm_calc<GemmConfig,
6668
ADataType,
6769
BDataType,
70+
AScaleDataType,
71+
BScaleDataType,
6872
AccDataType,
6973
CDataType,
7074
ALayout,
7175
BLayout,
7276
CLayout,
73-
ScaleM,
74-
ScaleN,
7577
UsePersistentKernel,
7678
split_k_.value>(
7779
args, ck_tile::stream_config{nullptr, true, 1, n_warmup, n_repeat, true, true, 50});

‎example/ck_tile/42_mx_gemm/mx_gemm.hpp‎

Lines changed: 4 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -9,49 +9,8 @@
99
#include "ck_tile/host/kernel_launch.hpp"
1010
#include "ck_tile/ops/epilogue.hpp"
1111
#include "ck_tile/ops/gemm.hpp"
12-
#include "ck_tile/ops/gemm_mx.hpp"
13-
#include "ck_tile/ops/gemm_mx/kernel/scale_pointer.hpp"
14-
15-
template <typename ScaleM, typename ScaleN>
16-
struct MXGemmHostArgs : ck_tile::UniversalGemmHostArgs<1, 1, 0>
17-
{
18-
using Base = ck_tile::UniversalGemmHostArgs<1, 1, 0>;
19-
20-
MXGemmHostArgs(const void* a_ptr,
21-
const void* b_ptr,
22-
void* c_ptr_,
23-
ck_tile::index_t k_batch_,
24-
ck_tile::index_t M_,
25-
ck_tile::index_t N_,
26-
ck_tile::index_t K_,
27-
ck_tile::index_t stride_A_,
28-
ck_tile::index_t stride_B_,
29-
ck_tile::index_t stride_C_,
30-
ScaleM scale_m_,
31-
ScaleN scale_n_)
32-
: Base({a_ptr},
33-
{b_ptr},
34-
{},
35-
c_ptr_,
36-
k_batch_,
37-
M_,
38-
N_,
39-
K_,
40-
{stride_A_},
41-
{stride_B_},
42-
{},
43-
stride_C_),
44-
scale_m(scale_m_),
45-
scale_n(scale_n_)
46-
{
47-
}
48-
49-
ScaleM scale_m;
50-
ScaleN scale_n;
51-
};
5212

5313
// GEMM config with 16x16 warp tile
54-
5514
struct MxGemmConfig
5615
{
5716
static constexpr ck_tile::index_t M_Tile = 128;
@@ -78,12 +37,15 @@ struct MxGemmConfig
7837
static constexpr int TileParitionerM01 = 4;
7938
static constexpr auto Scheduler = ck_tile::GemmPipelineScheduler::Intrawave;
8039
static constexpr ck_tile::index_t NumWaveGroups = 1;
81-
static constexpr bool DoubleSmemBuffer = false; // comp_async uses double buffer
40+
static constexpr bool DoubleSmemBuffer = true; // comp_async uses double buffer
8241
static constexpr bool Preshuffle = false;
8342
static constexpr ck_tile::index_t BContiguousItemsPerAccess = 16;
8443

8544
static constexpr int N_Repeat = N_Tile / N_Warp_Tile / N_Warp;
8645
static constexpr bool TiledMMAPermuteN = false;
46+
47+
using AScaleDataType = ck_tile::e8m0_t;
48+
using BScaleDataType = ck_tile::e8m0_t;
8749
};
8850

8951
struct MX_GemmConfigEightWaves : MxGemmConfig

‎example/ck_tile/42_mx_gemm/mx_gemm_instance.hpp‎

Lines changed: 33 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,9 @@
55

66
#include "ck_tile/host.hpp"
77
#include "mx_gemm.hpp"
8-
#include "ck_tile/ops/gemm_mx/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp"
9-
#include "ck_tile/ops/gemm_mx/pipeline/gemm_pipeline_ag_bg_cr_comp_async_eight_waves.hpp"
10-
#include "ck_tile/ops/gemm_mx/kernel/gemm_mx_kernel.hpp"
11-
#include "ck_tile/ops/gemm_mx/pipeline/wp_pipeline_agmem_bgmem_creg_v1.hpp"
8+
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async.hpp"
9+
#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_async_eight_waves.hpp"
10+
#include "ck_tile/ops/gemm/pipeline/wp_mx_pipeline_agmem_bgmem_creg_v1.hpp"
1211

1312
template <typename Layout>
1413
using is_row_major_t = ck_tile::bool_constant<
@@ -17,16 +16,16 @@ using is_row_major_t = ck_tile::bool_constant<
1716
template <typename GemmConfig,
1817
typename ADataType,
1918
typename BDataType,
19+
typename AScaleDataType,
20+
typename BScaleDataType,
2021
typename AccDataType,
2122
typename CDataType,
2223
typename ALayout,
2324
typename BLayout,
2425
typename CLayout,
25-
typename ScaleM,
26-
typename ScaleN,
2726
bool persistent,
2827
bool Splitk>
29-
float mx_gemm_calc(const MXGemmHostArgs<ScaleM, ScaleN>& args, const ck_tile::stream_config& s)
28+
float mx_gemm_calc(const ck_tile::MxGemmHostArgs<1, 1, 0>& args, const ck_tile::stream_config& s)
3029
{
3130
using GemmShape = ck_tile::TileGemmShape<
3231
ck_tile::sequence<GemmConfig::M_Tile, GemmConfig::N_Tile, GemmConfig::K_Tile>,
@@ -51,29 +50,38 @@ float mx_gemm_calc(const MXGemmHostArgs<ScaleM, ScaleN>& args, const ck_tile::st
5150
static_assert(sizeof(ComputeDataType) >= sizeof(BDataType),
5251
"mixed_prec_gemm requires ADataType is a wider type than BDataType");
5352

54-
using MXPipelineProblem = ck_tile::UniversalGemmPipelineProblem<ADataType,
55-
BDataType,
56-
AccDataType,
57-
GemmShape,
58-
MXGemmTraits,
59-
GemmConfig::Scheduler>;
53+
using AComputeDataType = ADataType;
54+
using BComputeDataType = BDataType;
55+
56+
using MXPipelineProblem = ck_tile::MxGemmPipelineProblem<ADataType,
57+
BDataType,
58+
AccDataType,
59+
GemmShape,
60+
MXGemmTraits,
61+
GemmConfig::Scheduler,
62+
ck_tile::element_wise::PassThrough,
63+
ck_tile::element_wise::PassThrough,
64+
AComputeDataType,
65+
BComputeDataType,
66+
AScaleDataType,
67+
BScaleDataType>;
6068

61-
// Use the MX GEMM Preshuffle pipeline or
62-
// the new MX comp_async pipeline with MX scaling support
6369
constexpr bool IsEightWave =
6470
(GemmConfig::M_Warp * GemmConfig::N_Warp * GemmConfig::K_Warp) == 8;
6571
using MXGemmPipeline = std::conditional_t<
6672
GemmConfig::Preshuffle,
6773
ck_tile::MXGemmPreshufflePipelineAGmemBGmemCRegV1<MXPipelineProblem>,
6874
std::conditional_t<IsEightWave,
69-
ck_tile::MXGemmPipelineAgBgCrCompAsyncEightWaves<MXPipelineProblem>,
70-
ck_tile::MXGemmPipelineAgBgCrCompAsync<MXPipelineProblem>>>;
75+
ck_tile::GemmPipelineAgBgCrCompAsyncEightWaves<MXPipelineProblem>,
76+
ck_tile::GemmPipelineAgBgCrCompAsync<MXPipelineProblem>>>;
7177

7278
using TilePartitioner =
7379
ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
7480
GemmConfig::TileParitionerGroupNum,
7581
GemmConfig::TileParitionerM01>;
7682

83+
constexpr ck_tile::index_t BlockedXDLNPerWarp = GemmConfig::Preshuffle ? 2 : 1;
84+
7785
using GemmEpilogue =
7886
std::conditional_t<GemmConfig::TiledMMAPermuteN,
7987
ck_tile::PermuteNEpilogue<
@@ -115,38 +123,27 @@ float mx_gemm_calc(const MXGemmHostArgs<ScaleM, ScaleN>& args, const ck_tile::st
115123
GemmConfig::NumWaveGroups,
116124
false, // FixedVectorSize_ (Default)
117125
1, // VectorSizeC_ (Default)
118-
ck_tile::MXEpilogueTraits<GemmConfig>::BlockedXDLNPerWarp,
126+
BlockedXDLNPerWarp,
119127
false, // DoubleSmemBuffer_ (Default)
120128
ComputeDataType, // AComputeDataType
121129
ComputeDataType, // BComputeDataType
122130
!GemmConfig::Preshuffle>>>; // TilesPacked_ (because of
123131
// packed scales)
124132

125-
using Kernel = ck_tile::MXGemmKernel<TilePartitioner, MXGemmPipeline, GemmEpilogue>;
126-
127-
auto kargs = Kernel::MakeKernelArgs(std::array<const void*, 1>{args.as_ptr},
128-
std::array<const void*, 1>{args.bs_ptr},
129-
std::array<const void*, 0>{},
130-
args.e_ptr,
131-
args.k_batch,
132-
args.M,
133-
args.N,
134-
args.K,
135-
std::array<ck_tile::index_t, 1>{args.stride_As},
136-
std::array<ck_tile::index_t, 1>{args.stride_Bs},
137-
std::array<ck_tile::index_t, 0>{},
138-
args.stride_E,
139-
args.scale_m,
140-
args.scale_n);
133+
using Kernel = ck_tile::MxGemmKernel<TilePartitioner, MXGemmPipeline, GemmEpilogue>;
134+
135+
auto kargs = Kernel::MakeKernelArgs(args);
141136

142137
if(!Kernel::IsSupportedArgument(kargs))
143138
{
144139
throw std::runtime_error(
145140
"MX GEMM: unsupported shape/configuration (set CK_TILE_LOGGING=1 for details).");
146141
}
147142

148-
const auto kernel = ck_tile::make_kernel<Kernel::kBlockPerCu>(
149-
Kernel{}, Kernel::GridSize(kargs), Kernel::BlockSize(), 0, kargs);
143+
constexpr int kBlockPerCu = 1;
144+
145+
const auto kernel = ck_tile::make_kernel<kBlockPerCu>(
146+
Kernel{}, Kernel::GridSize(args.M, args.N, args.k_batch), Kernel::BlockSize(), 0, kargs);
150147

151148
// For split-K (k_batch > 1) the kernel's epilogue uses atomic_add into C, so C must be
152149
// zeroed before every kernel launch -- not just once before the first warmup iteration.

0 commit comments

Comments
 (0)