include/ck_tile/core/algorithm/static_encoding_pattern.hpp Source File

include/ck_tile/core/algorithm/static_encoding_pattern.hpp Source File#

Composable Kernel: include/ck_tile/core/algorithm/static_encoding_pattern.hpp Source File
static_encoding_pattern.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_tile/core/arch/arch.hpp"
7 #include "ck_tile/core/config.hpp"
8 #include "ck_tile/core/container/sequence.hpp"
9 #include "ck_tile/core/container/tuple.hpp"
10 #include "ck_tile/core/numeric/integer.hpp"
11 #include "ck_tile/core/tensor/tile_distribution.hpp"
12 #include "ck_tile/core/tensor/tile_distribution_encoding.hpp"
13 
14 namespace ck_tile {
15 
20 enum struct tile_distribution_pattern
21 {
26  thread_raked,
31  warp_raked,
36  block_raked,
37 };
38 
39 struct TileDistributionEncodingPattern
40 {
41 };
42 
55 template <index_t BlockSize,
56  index_t YPerTile,
57  index_t XPerTile,
58  index_t VecSize,
59  tile_distribution_pattern DistributionPattern>
60 struct TileDistributionEncodingPattern2D : public TileDistributionEncodingPattern
61 {
62 };
63 
64 // Thread raked
65 template <index_t BlockSize, index_t YPerTile, index_t XPerTile, index_t VecSize>
66 struct TileDistributionEncodingPattern2D<BlockSize,
67  YPerTile,
68  XPerTile,
69  VecSize,
70  tile_distribution_pattern::thread_raked>
71  : public TileDistributionEncodingPattern
72 {
73 
74  // TODO: make pattern where below condition does not need to hold - GGemmMultiDSplitk!
75  static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!");
76  static constexpr index_t warp_size = get_warp_size();
77  static constexpr index_t num_warps = BlockSize / get_warp_size();
78  static constexpr index_t X1 = VecSize;
79  static constexpr index_t X0 = XPerTile / X1; // # of threads in X dim
80 
81  // # of rows in Y dim accessed by single wavefront in one iteration
82  static constexpr index_t Y1 = warp_size / X0;
83  static_assert(X0 * Y1 == warp_size, "X0 * Y1 must cover whole wavefront!");
84 
85  static constexpr index_t Y0 = num_warps;
86  // YPerWarp = YPerTile / Y0;
87  // Y2 = YPerWarp / Y1;
88  static constexpr index_t Y2 = YPerTile / (Y1 * Y0); // # of iters within wavefront
89 
90  static_assert(X0 * Y1 * Y0 == BlockSize, "X0 * warp_ys * Y0 must cover whole workgroup!");
91  static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover whole YPerTile");
92 
93  CK_TILE_HOST_DEVICE static constexpr auto Make2DStaticTileDistribution()
94  {
95  return make_static_tile_distribution(
96  tile_distribution_encoding<sequence<1>,
97  tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
98  tuple<sequence<1>, sequence<1, 2>>,
99  tuple<sequence<0>, sequence<1, 0>>,
100  sequence<1, 2>,
101  sequence<2, 1>>{});
102  }
103 
104  CK_TILE_HOST_DEVICE static constexpr auto MakeShuffled2DStaticTileDistribution()
105  {
106  return make_static_tile_distribution(
107  tile_distribution_encoding<sequence<1>,
108  tuple<sequence<X0, X1>, sequence<Y0, Y1, Y2>>,
109  tuple<sequence<2>, sequence<2, 1>>,
110  tuple<sequence<0>, sequence<1, 0>>,
111  sequence<1, 2>,
112  sequence<1, 2>>{});
113  }
114 };
115 
116 // Warp raked
117 template <index_t BlockSize, index_t YPerTile, index_t XPerTile, index_t VecSize>
118 struct TileDistributionEncodingPattern2D<BlockSize,
119  YPerTile,
120  XPerTile,
121  VecSize,
122  tile_distribution_pattern::warp_raked>
123  : public TileDistributionEncodingPattern
124 {
125 
126  static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!");
127  static constexpr index_t warp_size = get_warp_size();
128  static constexpr index_t num_warps = BlockSize / get_warp_size();
129  static constexpr index_t X1 = VecSize;
130  static constexpr index_t X0 = XPerTile / X1; // # of threads in X dim
131 
132  static constexpr index_t Y2 = warp_size / X0; // # of rows in Y dim to cover whole wavefront
133  static_assert(X0 * Y2 == warp_size, "X0 * Y2 must cover whole wavefront!");
134 
135  static constexpr index_t Y0 = num_warps;
136  static_assert(X0 * Y2 * Y0 == BlockSize, "X0 * Y2 * Y1 must cover whole workgroup!");
137 
138  static constexpr index_t Y1 = YPerTile / (Y2 * Y0); // # of iters within wavefront
139  static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover whole YPerTile");
140 
141  CK_TILE_HOST_DEVICE static constexpr auto Make2DStaticTileDistribution()
142  {
143  return make_static_tile_distribution(
144  tile_distribution_encoding<sequence<1>,
145  tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
146  tuple<sequence<1>, sequence<1, 2>>,
147  tuple<sequence<0>, sequence<2, 0>>,
148  sequence<1, 2>,
149  sequence<1, 1>>{});
150  }
151 
152  CK_TILE_HOST_DEVICE static constexpr auto MakeShuffled2DStaticTileDistribution()
153  {
154  return make_static_tile_distribution(
155  tile_distribution_encoding<sequence<1>,
156  tuple<sequence<X0, X1>, sequence<Y0, Y1, Y2>>,
157  tuple<sequence<2>, sequence<2, 1>>,
158  tuple<sequence<0>, sequence<2, 0>>,
159  sequence<1, 2>,
160  sequence<1, 1>>{});
161  }
162 };
163 
164 // Block raked
165 template <index_t BlockSize, index_t YPerTile, index_t XPerTile, index_t VecSize>
166 struct TileDistributionEncodingPattern2D<BlockSize,
167  YPerTile,
168  XPerTile,
169  VecSize,
170  tile_distribution_pattern::block_raked>
171  : public TileDistributionEncodingPattern
172 {
173 
174  // TODO: make pattern where below condition does not need to hold - GGemmMultiDSplitk!
175  static_assert(XPerTile % VecSize == 0, "XPerTile must be a multiple of VecSize!");
176  static constexpr index_t warp_size = get_warp_size();
177  static constexpr index_t num_warps = BlockSize / get_warp_size();
178  static constexpr index_t X1 = VecSize;
179  static constexpr index_t X0 = XPerTile / X1; // # of threads in X dim
180  static constexpr index_t Y2 = warp_size / X0; // # of rows in Y dim to cover whole wavefront
181  static_assert(X0 * Y2 == warp_size, "X0 * Y2 must cover whole wavefront!");
182  static constexpr index_t Y1 = num_warps;
183  static_assert(X0 * Y2 * Y1 == BlockSize, "X0 * Y2 * Y1 must cover whole workgroup!");
184  static constexpr index_t Y0 = YPerTile / (Y2 * Y1); // # of iters
185  static_assert(Y0 * Y1 * Y2 == YPerTile, "Y0, Y1, Y2 must cover whole YPerTile");
186 
187  CK_TILE_HOST_DEVICE static constexpr auto Make2DStaticTileDistribution()
188  {
189  return make_static_tile_distribution(
190  tile_distribution_encoding<sequence<1>,
191  tuple<sequence<Y0, Y1, Y2>, sequence<X0, X1>>,
192  tuple<sequence<1>, sequence<1, 2>>,
193  tuple<sequence<1>, sequence<2, 0>>,
194  sequence<1, 2>,
195  sequence<0, 1>>{});
196  }
197 
198  CK_TILE_HOST_DEVICE static constexpr auto MakeShuffled2DStaticTileDistribution()
199  {
200  return make_static_tile_distribution(
201  tile_distribution_encoding<sequence<1>,
202  tuple<sequence<X0, X1>, sequence<Y0, Y1, Y2>>,
203  tuple<sequence<2>, sequence<2, 1>>,
204  tuple<sequence<1>, sequence<2, 0>>,
205  sequence<1, 2>,
206  sequence<1, 0>>{});
207  }
208 };
209 
210 } // namespace ck_tile
CK_TILE_HOST_DEVICE
#define CK_TILE_HOST_DEVICE
Definition: config.hpp:41
ck_tile
Definition: cluster_descriptor.hpp:13
ck_tile::get_warp_size
constexpr CK_TILE_HOST_DEVICE index_t get_warp_size()
Definition: arch.hpp:51
ck_tile::index_t
int32_t index_t
Definition: integer.hpp:9
ck_tile::tile_distribution_pattern
tile_distribution_pattern
Enumeration describing static tile distribution patterns.
Definition: static_encoding_pattern.hpp:21
ck_tile::tile_distribution_pattern::block_raked
@ block_raked
Block raked pattern - aka linear.
ck_tile::tile_distribution_pattern::thread_raked
@ thread_raked
Thread raked pattern.
ck_tile::tile_distribution_pattern::warp_raked
@ warp_raked
Warp raked pattern.
ck_tile::make_static_tile_distribution
constexpr CK_TILE_HOST_DEVICE auto make_static_tile_distribution(StaticTileDistributionEncoding_)
Definition: tile_distribution.hpp:498
ck_tile::TileDistributionEncodingPattern2D< BlockSize, YPerTile, XPerTile, VecSize, tile_distribution_pattern::block_raked >::MakeShuffled2DStaticTileDistribution
static constexpr CK_TILE_HOST_DEVICE auto MakeShuffled2DStaticTileDistribution()
Definition: static_encoding_pattern.hpp:198
ck_tile::TileDistributionEncodingPattern2D< BlockSize, YPerTile, XPerTile, VecSize, tile_distribution_pattern::block_raked >::Make2DStaticTileDistribution
static constexpr CK_TILE_HOST_DEVICE auto Make2DStaticTileDistribution()
Definition: static_encoding_pattern.hpp:187
ck_tile::TileDistributionEncodingPattern2D< BlockSize, YPerTile, XPerTile, VecSize, tile_distribution_pattern::warp_raked >::MakeShuffled2DStaticTileDistribution
static constexpr CK_TILE_HOST_DEVICE auto MakeShuffled2DStaticTileDistribution()
Definition: static_encoding_pattern.hpp:152
ck_tile::TileDistributionEncodingPattern2D< BlockSize, YPerTile, XPerTile, VecSize, tile_distribution_pattern::warp_raked >::Make2DStaticTileDistribution
static constexpr CK_TILE_HOST_DEVICE auto Make2DStaticTileDistribution()
Definition: static_encoding_pattern.hpp:141
ck_tile::TileDistributionEncodingPattern2D< BlockSize, YPerTile, XPerTile, VecSize, tile_distribution_pattern::thread_raked >::MakeShuffled2DStaticTileDistribution
static constexpr CK_TILE_HOST_DEVICE auto MakeShuffled2DStaticTileDistribution()
Definition: static_encoding_pattern.hpp:104
ck_tile::TileDistributionEncodingPattern2D< BlockSize, YPerTile, XPerTile, VecSize, tile_distribution_pattern::thread_raked >::Make2DStaticTileDistribution
static constexpr CK_TILE_HOST_DEVICE auto Make2DStaticTileDistribution()
Definition: static_encoding_pattern.hpp:93
ck_tile::TileDistributionEncodingPattern2D
Class creating 2D static tile distribution with different load/store patterns.
Definition: static_encoding_pattern.hpp:61
ck_tile::TileDistributionEncodingPattern
Definition: static_encoding_pattern.hpp:40
ck_tile::sequence
Definition: sequence.hpp:52
ck_tile::tile_distribution_encoding
Definition: tile_distribution_encoding.hpp:26
ck_tile::tuple
Definition: tuple.hpp:192
tile_distribution.hpp
tile_distribution_encoding.hpp