/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.0/include/ck_tile/core/arch/amd_transpose_load_encoding.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.0/include/ck_tile/core/arch/amd_transpose_load_encoding.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.0/include/ck_tile/core/arch/amd_transpose_load_encoding.hpp Source File
amd_transpose_load_encoding.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/config.hpp"
7 #include "ck_tile/core/container/sequence.hpp"
8 #include "ck_tile/core/tensor/tile_distribution_encoding.hpp"
9 
10 namespace ck_tile {
11 
12 // this generate wave level tile distribution
13 template <typename T, typename = void>
14 struct LaneGroupTransposeTraits;
15 
16 template <typename T>
17 struct LaneGroupTransposeTraits<T, std::enable_if_t<sizeof(T) == 2>>
18 {
19  // before transpose, 4x16
20  static constexpr index_t ksecondDim = 4;
21  static constexpr index_t kleadDim = 16;
22  // after transpose, 16x4
23  static constexpr index_t ksecondDimT = 16;
24  static constexpr index_t kleadDimT = 4;
25  template <index_t kOuterDistDim0,
26  index_t kOuterDistDim1,
27  index_t kInnerDistDim0,
28  index_t kInnerDistDim1>
29  using TileDistribution =
30  tile_distribution_encoding<sequence<>,
31  tuple<sequence<kOuterDistDim0, kOuterDistDim1, 4>,
32  sequence<kInnerDistDim0, kInnerDistDim1, 4, 4>>,
33  tuple<sequence<1, 2, 1, 2>>,
34  tuple<sequence<0, 0, 2, 2>>,
35  sequence<2, 1, 2>,
36  sequence<1, 1, 3>>;
37 };
38 
39 template <typename T>
40 struct LaneGroupTransposeTraits<T, std::enable_if_t<sizeof(T) == 1>>
41 {
42  static constexpr index_t ksecondDim = 8;
43  static constexpr index_t kleadDim = 16;
44 
45  static constexpr index_t ksecondDimT = 16;
46  static constexpr index_t kleadDimT = 8;
47 
48  template <index_t kOuterDistDim0,
49  index_t kOuterDistDim1,
50  index_t kInnerDistDim0,
51  index_t kInnerDistDim1>
52  using TileDistribution =
53  tile_distribution_encoding<sequence<>,
54  tuple<sequence<kOuterDistDim0, kOuterDistDim1, 8>,
55  sequence<kInnerDistDim0, kInnerDistDim1, 2, 8>>,
56  tuple<sequence<1, 2, 1, 2>>,
57  tuple<sequence<0, 0, 2, 2>>,
58  sequence<2, 1, 2>,
59  sequence<1, 1, 3>>;
60 };
61 
62 /*
63  * @brief This function is used to generate the transposed distribution encoding
64  * for the given data type and distribution dimensions.
65  *
66  * @tparam T The data type of the elements in the tensor.
67  * @tparam kOuterDistDim0 The outer distribution dimension 0, which is outer dimension for stride.
68  * @tparam kOuterDistDim1 The outer distribution dimension 1, which is inner dimension for stride.
69  * @tparam kInnerDistDim0 The inner distribution dimension 0, which is outer dimension for
70  * consecutive.
71  * @tparam kInnerDistDim1 The inner distribution dimension 1, which is inner dimension for
72  * consecutive.
73  */
74 template <typename T,
75  index_t kOuterDistDim0,
76  index_t kOuterDistDim1,
77  index_t kInnerDistDim0,
78  index_t kInnerDistDim1>
79 CK_TILE_DEVICE constexpr auto make_transposed_distr_encode()
80 {
81  using xdllevel_dstr_encoding = typename LaneGroupTransposeTraits<T>::
82  template TileDistribution<kOuterDistDim0, kOuterDistDim1, kInnerDistDim0, kInnerDistDim1>;
83  return xdllevel_dstr_encoding{};
84 }
85 
86 } // namespace ck_tile
CK_TILE_DEVICE
#define CK_TILE_DEVICE
Definition: config.hpp:40
ck_tile
Definition: cluster_descriptor.hpp:13
ck_tile::index_t
int32_t index_t
Definition: integer.hpp:9
ck_tile::make_transposed_distr_encode
constexpr CK_TILE_DEVICE auto make_transposed_distr_encode()
Definition: amd_transpose_load_encoding.hpp:79
ck::enable_if_t
typename std::enable_if< B, T >::type enable_if_t
Definition: enable_if.hpp:27
ck_tile::LaneGroupTransposeTraits
Definition: amd_transpose_load_encoding.hpp:14
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_encoding.hpp