/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_selector.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_selector.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_selector.hpp Source File
gridwise_gemm_pipeline_selector.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: MIT
2 // Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
3 
4 #pragma once
5 
6 #if !defined(__HIPCC_RTC__) || !defined(CK_CODE_GEN_RTC)
7 #include <iostream>
8 #include <ostream>
9 #endif
10 
11 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v1.hpp"
12 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v2.hpp"
13 #include "ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v4_direct_load.hpp"
14 
15 namespace ck {
16 
17 enum struct PipelineVersion
18 {
19  v1,
20  v2,
21  // v3 is only used in the Stream-K implementation.
22  v4,
23  weight_only,
24 };
25 
26 template <PipelineVersion PipelineVer,
27  index_t NumPrefetch = 1,
28  LoopScheduler LoopSched = LoopScheduler::Default,
29  bool AEnableLds = true,
30  bool BEnableLds = true>
31 constexpr auto GridwiseGemmPipeline_Selector()
32 {
33  if constexpr(PipelineVer == PipelineVersion::v1)
34  {
35  if constexpr(LoopSched == LoopScheduler::Default)
36  {
37  return GridwiseGemmPipeline_v1<NumPrefetch, AEnableLds, BEnableLds>{};
38  }
39  else if constexpr(LoopSched == LoopScheduler::Interwave)
40  {
41  return GridwiseGemmPipelineInterwave_v1<NumPrefetch>{};
42  }
43  }
44  else if constexpr(PipelineVer == PipelineVersion::v2)
45  {
46  return GridwiseGemmPipeline_v2{};
47  }
48  else if constexpr(PipelineVer == PipelineVersion::v4)
49  {
50  return GridwiseGemmPipeline_v4<NumPrefetch>{};
51  }
52  else if constexpr(PipelineVer == PipelineVersion::weight_only)
53  {
54  return GridwiseGemmPipeline_v1_WeightOnly<NumPrefetch, AEnableLds, BEnableLds>{};
55  }
56  else
57  {
58 #if !defined(__HIPCC_RTC__) || !defined(CK_CODE_GEN_RTC)
59  std::cerr << "GridwiseGemmPipeline configuration is not available" << std::endl;
60 #endif
61  }
62 }
63 
64 } // namespace ck
65 
66 #if !defined(__HIPCC_RTC__) || !defined(CK_CODE_GEN_RTC)
67 inline std::ostream& operator<<(std::ostream& os, const ck::PipelineVersion& p)
68 {
69  switch(p)
70  {
71  case ck::PipelineVersion::v1: os << "PipelineVersion::v1"; break;
72  case ck::PipelineVersion::v2: os << "PipelineVersion::v2"; break;
73  case ck::PipelineVersion::v4: os << "PipelineVersion::v4"; break;
74  case ck::PipelineVersion::weight_only: os << "PipelineVersion::weight_only"; break;
75  default: os << "";
76  }
77  return os;
78 }
79 #endif
operator<<
std::ostream & operator<<(std::ostream &os, const ck::PipelineVersion &p)
Definition: gridwise_gemm_pipeline_selector.hpp:67
gridwise_gemm_pipeline_v1.hpp
gridwise_gemm_pipeline_v2.hpp
gridwise_gemm_pipeline_v4_direct_load.hpp
ck
Definition: ck.hpp:266
ck::GridwiseGemmPipeline_Selector
constexpr auto GridwiseGemmPipeline_Selector()
Definition: gridwise_gemm_pipeline_selector.hpp:31
ck::LoopScheduler
LoopScheduler
Definition: loop_scheduler.hpp:15
ck::index_t
int32_t index_t
Definition: ck.hpp:297
ck::PipelineVersion
PipelineVersion
Definition: gridwise_gemm_pipeline_selector.hpp:18
ck::GridwiseGemmPipeline_v1_WeightOnly
Definition: gridwise_gemm_pipeline_v1.hpp:555
ck::GridwiseGemmPipeline_v1
Definition: gridwise_gemm_pipeline_v1.hpp:13
ck::GridwiseGemmPipeline_v2
Definition: gridwise_gemm_pipeline_v2.hpp:11
ck::GridwiseGemmPipeline_v4
Definition: gridwise_gemm_pipeline_v4_direct_load.hpp:27
ck::GridwiseGemmPipelineInterwave_v1
Definition: gridwise_gemm_pipeline_v1.hpp:658