blockwise_gemm_pipeline_wmma_selector.hpp Source File

blockwise_gemm_pipeline_wmma_selector.hpp Source File#

Composable Kernel: blockwise_gemm_pipeline_wmma_selector.hpp Source File
blockwise_gemm_pipeline_wmma_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_wmmaops_v1.hpp"
7#include "ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp"
8
9namespace ck {
10
11template <BlockGemmPipelineVersion BlkGemmPipelineVer,
12 BlockGemmPipelineScheduler BlkGemmPipeSche,
13 index_t BlockSize,
14 typename ADataType,
15 typename BDataType,
16 typename ComputeTypeA,
17 typename ComputeTypeB,
18 typename AccDataType,
19 typename AWmmaTileDesc,
20 typename BWmmaTileDesc,
21 index_t ABlockTransferSrcScalarPerVector,
22 index_t BBlockTransferSrcScalarPerVector,
23 index_t MPerBlock,
24 index_t NPerBlock,
25 index_t KPerBlock,
26 index_t MPerWmma,
27 index_t NPerWmma,
28 index_t MRepeat,
29 index_t NRepeat,
30 index_t KPack,
31 bool TransposeC = false>
32constexpr auto BlockGemmPipeline_Selector()
33{
34 if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
35 {
36 return BlockwiseGemmWmmaops_pipeline_v1<BlkGemmPipeSche,
37 BlockSize,
38 ADataType,
39 BDataType,
40 ComputeTypeA,
41 ComputeTypeB,
42 AccDataType,
43 AWmmaTileDesc,
44 BWmmaTileDesc,
45 ABlockTransferSrcScalarPerVector,
46 BBlockTransferSrcScalarPerVector,
47 MPerBlock,
48 NPerBlock,
49 KPerBlock,
50 MPerWmma,
51 NPerWmma,
52 MRepeat,
53 NRepeat,
54 KPack,
55 TransposeC>{};
56 }
57 else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
58 {
59 return BlockwiseGemmWmmaops_pipeline_v3<BlkGemmPipeSche,
60 BlockSize,
61 ADataType,
62 BDataType,
63 ComputeTypeA,
64 ComputeTypeB,
65 AccDataType,
66 AWmmaTileDesc,
67 BWmmaTileDesc,
68 ABlockTransferSrcScalarPerVector,
69 BBlockTransferSrcScalarPerVector,
70 MPerBlock,
71 NPerBlock,
72 KPerBlock,
73 MPerWmma,
74 NPerWmma,
75 MRepeat,
76 NRepeat,
77 KPack,
78 TransposeC>{};
79 }
80 else
81 {
82 static_assert(false, "BlockGemmPipeline configuration is not available");
83 }
84}
85
86} // namespace ck
blockwise_gemm_pipeline_wmmaops_v1.hpp
blockwise_gemm_pipeline_wmmaops_v3.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::v3
@ v3
Definition blkgemmpipe_scheduler.hpp:16
ck::BlockGemmPipelineVersion::v1
@ v1
Definition blkgemmpipe_scheduler.hpp:14
ck::BlockGemmPipeline_Selector
constexpr auto BlockGemmPipeline_Selector()
Definition blockwise_gemm_pipeline_wmma_selector.hpp:32
ck::BlockGemmPipelineScheduler
BlockGemmPipelineScheduler
Definition blkgemmpipe_scheduler.hpp:25
ck::BlockwiseGemmWmmaops_pipeline_v1
Definition blockwise_gemm_pipeline_wmmaops_v1.hpp:37
ck::BlockwiseGemmWmmaops_pipeline_v3
Definition blockwise_gemm_pipeline_wmmaops_v3.hpp:37