blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp Source File

blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp Source File#

Composable Kernel: blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp Source File
blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp
Go to the documentation of this file.
1// SPDX-License-Identifier: MIT
2// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
3
4#pragma once
5
6#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp"
7#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp"
8#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp"
9
10namespace ck {
11
12template <BlockGemmPipelineVersion BlkGemmPipelineVer,
13 BlockGemmPipelineScheduler BlkGemmPipeSche,
14 index_t BlockSize,
15 typename ADataType,
16 typename BDataType,
17 typename ComputeDataType,
18 typename AccDataType,
19 typename ATileDesc,
20 typename BTileDesc,
21 typename AMmaTileDesc,
22 typename BMmaTileDesc,
23 index_t ABlockTransferSrcScalarPerVector,
24 index_t BBlockTransferSrcScalarPerVector,
25 index_t MPerBlock,
26 index_t NPerBlock,
27 index_t KPerBlock,
28 index_t MPerXDL,
29 index_t NPerXDL,
30 index_t MRepeat,
31 index_t NRepeat,
32 index_t KPack>
33constexpr auto BlockGemmABScalePipeline_Selector()
34{
35 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
36 {
37 return BlockwiseGemmXdlops_pipeline_v1_ab_scale<BlkGemmPipeSche,
38 BlockSize,
39 ADataType,
40 BDataType,
41 ComputeDataType,
42 AccDataType,
43 ATileDesc,
44 BTileDesc,
45 AMmaTileDesc,
46 BMmaTileDesc,
47 ABlockTransferSrcScalarPerVector,
48 BBlockTransferSrcScalarPerVector,
49 MPerBlock,
50 NPerBlock,
51 KPerBlock,
52 MPerXDL,
53 NPerXDL,
54 MRepeat,
55 NRepeat,
56 KPack>{};
57 }
58 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2)
59 {
60 return BlockwiseGemmXdlops_pipeline_v2_ab_scale<BlkGemmPipeSche,
61 BlockSize,
62 ADataType,
63 BDataType,
64 ComputeDataType,
65 AccDataType,
66 ATileDesc,
67 BTileDesc,
68 AMmaTileDesc,
69 BMmaTileDesc,
70 ABlockTransferSrcScalarPerVector,
71 BBlockTransferSrcScalarPerVector,
72 MPerBlock,
73 NPerBlock,
74 KPerBlock,
75 MPerXDL,
76 NPerXDL,
77 MRepeat,
78 NRepeat,
79 KPack>{};
80 }
81 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
82 {
83 return BlockwiseGemmXdlops_pipeline_v3_ab_scale<BlkGemmPipeSche,
84 BlockSize,
85 ADataType,
86 BDataType,
87 ComputeDataType,
88 AccDataType,
89 ATileDesc,
90 BTileDesc,
91 AMmaTileDesc,
92 BMmaTileDesc,
93 ABlockTransferSrcScalarPerVector,
94 BBlockTransferSrcScalarPerVector,
95 MPerBlock,
96 NPerBlock,
97 KPerBlock,
98 MPerXDL,
99 NPerXDL,
100 MRepeat,
101 NRepeat,
102 KPack>{};
103 }
104 else
105 {
106 std::cerr << "BlockGemmPipeline configuration is not available" << std::endl;
107 }
108}
109
110} // namespace ck
blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp
blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp
blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp
ck
Definition ck.hpp:268
ck::index_t
int32_t index_t
Definition ck.hpp:299
ck::BlockGemmPipelineVersion
BlockGemmPipelineVersion
Definition blkgemmpipe_scheduler.hpp:12
ck::BlockGemmPipelineVersion::v2
@ v2
Definition blkgemmpipe_scheduler.hpp:15
ck::BlockGemmPipelineVersion::v3
@ v3
Definition blkgemmpipe_scheduler.hpp:16
ck::BlockGemmPipelineVersion::v1
@ v1
Definition blkgemmpipe_scheduler.hpp:14
ck::BlockGemmABScalePipeline_Selector
constexpr auto BlockGemmABScalePipeline_Selector()
Definition blockwise_gemm_pipeline_xdlops_ab_scale_selector.hpp:33
ck::BlockGemmPipelineScheduler
BlockGemmPipelineScheduler
Definition blkgemmpipe_scheduler.hpp:25
ck::BlockwiseGemmXdlops_pipeline_v1_ab_scale
Definition blockwise_gemm_pipeline_xdlops_v1_ab_scale.hpp:37
ck::BlockwiseGemmXdlops_pipeline_v2_ab_scale
Definition blockwise_gemm_pipeline_xdlops_v2_ab_scale.hpp:37
ck::BlockwiseGemmXdlops_pipeline_v3_ab_scale
Definition blockwise_gemm_pipeline_xdlops_v3_ab_scale.hpp:37