/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.0/include/ck_tile/ops/batched_transpose/pipeline/batched_transpose_policy.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.0/include/ck_tile/ops/batched_transpose/pipeline/batched_transpose_policy.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/ops/batched_transpose/pipeline/batched_transpose_policy.hpp Source File
batched_transpose_policy.hpp
Go to the documentation of this file.
1 // SPDX-License-Identifier: MIT
2 // Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
3 
4 #pragma once
5 
6 #include "ck_tile/core.hpp"
7 #include "ck_tile/ops/softmax.hpp"
8 #include "ck_tile/ops/topk.hpp"
9 
10 namespace ck_tile {
11 
12 struct BatchedTransposePolicy
13 {
14  template <typename Problem>
15  CK_TILE_HOST_DEVICE static constexpr auto MakeInputDistribution()
16  {
17  constexpr index_t BlockSize = Problem::kBlockSize;
18  constexpr index_t MPerBlock = Problem::kMPerBlock;
19  constexpr index_t NPerBlock = Problem::kNPerBlock;
20  constexpr index_t VecLoadSize = Problem::VectorSizeInput;
21  using TileEncodingPattern =
22  TileDistributionEncodingPattern2D<BlockSize,
23  MPerBlock,
24  NPerBlock,
25  VecLoadSize,
26  tile_distribution_pattern::thread_raked>;
27  return TileEncodingPattern::Make2DStaticTileDistribution();
28  }
29 
30  template <typename Problem>
31  CK_TILE_HOST_DEVICE static constexpr auto MakeOutputDistribution()
32  {
33  constexpr index_t BlockSize = Problem::kBlockSize;
34  constexpr index_t MPerBlock = Problem::kMPerBlock;
35  constexpr index_t NPerBlock = Problem::kNPerBlock;
36  constexpr index_t VecLoadSize = Problem::VectorSizeOutput;
37 
38  using TileEncodingPattern =
39  TileDistributionEncodingPattern2D<BlockSize,
40  NPerBlock,
41  MPerBlock,
42  VecLoadSize,
43  tile_distribution_pattern::thread_raked>;
44  return TileEncodingPattern::MakeShuffled2DStaticTileDistribution();
45  }
46 };
47 } // 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::index_t
int32_t index_t
Definition: integer.hpp:9
ck_tile::tile_distribution_pattern::thread_raked
@ thread_raked
Thread raked pattern.
ck_tile::BatchedTransposePolicy
Definition: batched_transpose_policy.hpp:13
ck_tile::BatchedTransposePolicy::MakeInputDistribution
static constexpr CK_TILE_HOST_DEVICE auto MakeInputDistribution()
Definition: batched_transpose_policy.hpp:15
ck_tile::BatchedTransposePolicy::MakeOutputDistribution
static constexpr CK_TILE_HOST_DEVICE auto MakeOutputDistribution()
Definition: batched_transpose_policy.hpp:31
ck_tile::TileDistributionEncodingPattern2D
Class creating 2D static tile distribution with different load/store patterns.
Definition: static_encoding_pattern.hpp:129