/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/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-6.4.3/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-6.4.3/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  using S = Problem;
18  return make_static_tile_distribution(
19  tile_distribution_encoding<
20  sequence<>,
21  tuple<sequence<S::kMWarpPerBlock, S::kMThreadPerWarp, S::kMPerThread>,
22  sequence<S::kNWarpPerBlock, S::kNThreadPerWarp, S::kNPerThread>>,
23  tuple<sequence<1, 2>, sequence<1, 2>>,
24  tuple<sequence<0, 0>, sequence<1, 1>>,
25  sequence<1, 2>,
26  sequence<2, 2>>{});
27  }
28 
29  template <typename Problem>
30  CK_TILE_HOST_DEVICE static constexpr auto MakeOutputDistribution()
31  {
32  using S = Problem;
33  return make_static_tile_distribution(
34  tile_distribution_encoding<
35  sequence<>,
36  tuple<sequence<S::kNWarpPerBlock, S::kNThreadPerWarp, S::kNPerThread>,
37  sequence<S::kMWarpPerBlock, S::kMThreadPerWarp, S::kMPerThread>>,
38  tuple<sequence<2, 1>, sequence<2, 1>>,
39  tuple<sequence<0, 0>, sequence<1, 1>>,
40  sequence<2, 1>,
41  sequence<2, 2>>{});
42  }
43 };
44 } // 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::make_static_tile_distribution
constexpr CK_TILE_HOST_DEVICE auto make_static_tile_distribution(StaticTileDistributionEncoding_)
Definition: tile_distribution.hpp:498
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:30
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