blockwise_gemm_pipeline_xdlops_mx_bpreshuffle_selector.hpp Source File

blockwise_gemm_pipeline_xdlops_mx_bpreshuffle_selector.hpp Source File#

Composable Kernel: blockwise_gemm_pipeline_xdlops_mx_bpreshuffle_selector.hpp Source File
blockwise_gemm_pipeline_xdlops_mx_bpreshuffle_selector.hpp
Go to the documentation of this file.
1// SPDX-License-Identifier: MIT
2// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
3
4#pragma once
5
6#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx_bpreshuffle.hpp"
7
8namespace ck {
9template <BlockGemmPipelineVersion BlkGemmPipelineVer,
10 BlockGemmPipelineScheduler BlkGemmPipeSche,
11 index_t ThreadBlockSize,
12 index_t ScaleBlockSize,
13 typename ADataType,
14 typename AScaleDataType,
15 typename BDataType,
16 typename BScaleDataType,
17 typename ComputeDataType, // TODO: remove this as in this pipeline ADataType and BDataType
18 // must be used for compute
19 typename AccDataType,
20 typename ATileDesc,
21 typename BTileDesc,
22 typename AMmaTileDesc,
23 typename BMmaTileDesc,
24 index_t ABlockTransferSrcScalarPerVector,
25 index_t BBlockTransferSrcScalarPerVector,
26 index_t MPerBlock,
27 index_t NPerBlock,
28 index_t KPerBlock,
29 index_t MPerXDL,
30 index_t NPerXDL,
31 index_t MRepeat,
32 index_t NRepeat,
33 index_t KPack>
34constexpr auto BlockGemmMXBPreshufflePipeline_Selector()
35{
36
37 // Hardware MX GEMM pipeline
38 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
39 {
40 return BlockwiseGemmXdlops_pipeline_v3_mx_bprehuffle<BlkGemmPipeSche,
41 ThreadBlockSize,
42 ScaleBlockSize,
43 ADataType,
44 AScaleDataType,
45 BDataType,
46 BScaleDataType,
47 ATileDesc,
48 BTileDesc,
49 AMmaTileDesc,
50 BMmaTileDesc,
51 ABlockTransferSrcScalarPerVector,
52 BBlockTransferSrcScalarPerVector,
53 MPerBlock,
54 NPerBlock,
55 KPerBlock,
56 MPerXDL,
57 NPerXDL,
58 MRepeat,
59 NRepeat,
60 KPack>{};
61 }
62 else
63 {
64 std::cerr << "MX GEMM Pipeline configuration is not available" << std::endl;
65 }
66}
67
68} // namespace ck
blockwise_gemm_pipeline_xdlops_v3_mx_bpreshuffle.hpp
ck
Definition ck.hpp:268
ck::index_t
int32_t index_t
Definition ck.hpp:299
ck::BlockGemmMXBPreshufflePipeline_Selector
constexpr auto BlockGemmMXBPreshufflePipeline_Selector()
Definition blockwise_gemm_pipeline_xdlops_b_preshuffle_mx_moe_selector.hpp:37
ck::BlockGemmPipelineVersion
BlockGemmPipelineVersion
Definition blkgemmpipe_scheduler.hpp:12
ck::BlockGemmPipelineVersion::v3
@ v3
Definition blkgemmpipe_scheduler.hpp:16
ck::BlockGemmPipelineScheduler
BlockGemmPipelineScheduler
Definition blkgemmpipe_scheduler.hpp:25
ck::BlockwiseGemmXdlops_pipeline_v3_mx_bprehuffle
Definition blockwise_gemm_pipeline_xdlops_v3_mx_bpreshuffle.hpp:38