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