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

• develop
• include
• ck_tile
• ops
• batched_transpose
• pipeline

batched_transpose_lds_policy.hpp

Go to the documentation of this file.

1 // SPDX-License-Identifier: MIT

2 // Copyright (c) Advanced Micro Devices, Inc. All rights reserved.

3 

4 #pragma once

5 

6 #include "ck_tile/core.hpp"

7 #include "batched_transpose_common_policy.hpp"

8 

9 namespace ck_tile {

10 

11 struct BatchedTransposeLdsPolicy : public BatchedTransposeCommonPolicy

12 {

13  template <typename Problem>

14  CK_TILE_DEVICE static constexpr index_t GetSmemSize()

15  {

16  return integer_least_multiple(

17  sizeof(typename Problem::DataType) *

18  MakeLdsStoreBlockDescriptor<Problem>().get_element_space_size(),

19  16);

20  }

21 

22  template <typename Problem>

23  CK_TILE_DEVICE static constexpr auto MakeOutputDistribution()

24  {

25  constexpr auto input_dstr = MakeLdsLoadTileDistribution<Problem>();

26 

27  using OutTileDstrEncode =

28  typename OutputTileDistributionTraits<typename decltype(input_dstr)::DstrEncode,

29  typename Problem::DataType>::TransposedDstrEncode;

30  constexpr auto block_dstr = make_static_tile_distribution(OutTileDstrEncode{});

31 

32  return block_dstr;

33  }

34 

35  template <typename Problem>

36  CK_TILE_DEVICE static constexpr auto MakeLdsStoreBlockDescriptor()

37  {

38  constexpr index_t kLeadDimPerBlock = Problem::kLeadSizePerBlock;

39  constexpr index_t kSecondDimPerBlock = Problem::kSecondSizePerBlock;

40  constexpr index_t kVectorSize = Problem::LDSVectorSize;

41 

42  constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(

43  make_tuple(number<kSecondDimPerBlock>{},

44  number<kLeadDimPerBlock / kVectorSize>{},

45  number<kVectorSize>{}),

46  make_tuple(number<kLeadDimPerBlock>{}, number<kVectorSize>{}, number<1>{}),

47  number<kVectorSize>{},

48  number<1>{});

49 

50  constexpr auto lds_block_desc = transform_tensor_descriptor(

51  lds_block_desc_0,

52  make_tuple(make_pass_through_transform(number<kSecondDimPerBlock>{}),

53  make_merge_transform(make_tuple(number<kLeadDimPerBlock / kVectorSize>{},

54  number<kVectorSize>{}))),

55  make_tuple(sequence<0>{}, sequence<1, 2>{}),

56  make_tuple(sequence<0>{}, sequence<1>{}));

57 

58  return lds_block_desc;

59  }

60 

61  template <typename Problem>

62  CK_TILE_DEVICE static constexpr auto MakeLdsLoadBlockDescriptor()

63  {

64  constexpr index_t kLeadDimPerBlock = Problem::kLeadSizePerBlock;

65  constexpr index_t kSecondDimPerBlock = Problem::kSecondSizePerBlock;

66  constexpr index_t kVectorSize = Problem::LDSVectorSize;

67 

68  constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(

69  make_tuple(number<kSecondDimPerBlock>{},

70  number<kLeadDimPerBlock / kVectorSize>{},

71  number<kVectorSize>{}),

72  make_tuple(number<kLeadDimPerBlock>{}, number<kVectorSize>{}, number<1>{}),

73  number<kVectorSize>{},

74  number<1>{});

75 

76  constexpr auto lds_block_desc = transform_tensor_descriptor(

77  lds_block_desc_0,

78  make_tuple(make_pass_through_transform(number<kSecondDimPerBlock>{}),

79  make_merge_transform(make_tuple(number<kLeadDimPerBlock / kVectorSize>{},

80  number<kVectorSize>{}))),

81  make_tuple(sequence<0>{}, sequence<1, 2>{}),

82  make_tuple(sequence<0>{}, sequence<1>{}));

83 

84  return lds_block_desc;

85  }

86 

87  template <typename Problem>

88  CK_TILE_DEVICE static constexpr auto MakeLdsLoadTileDistribution()

89  {

90  using DataType = typename Problem::DataType;

91 

92  // Calculate block-level dimensions

93  constexpr index_t kLeadIterPerWarp = 1;

94  constexpr index_t kSecondIterPerWarp = 1;

95  constexpr index_t kLeadNumWarps = Problem::kLeadNumWarps;

96  constexpr index_t kSecondNumWarps = Problem::kSecondNumWarps;

97 

98  // Calculate repetitions of base pattern

99  constexpr index_t kLeadRepetitions = Problem::kQuadNumPerLeadDim;

100  constexpr index_t kSecondRepetitions = Problem::kQuadNumPerSecondDim;

101  constexpr index_t kSecondDimIterations = Problem::kIterationsInSecondDim;

102  constexpr index_t kSecondDimStrSub = kSecondRepetitions / kSecondDimIterations;

103 

104  constexpr index_t kLaneGroupSize = 16;

105  constexpr auto xdllevel_dstr_encoding = make_transposed_distr_encode<DataType,

106  kLaneGroupSize,

107  kSecondDimStrSub,

108  kSecondDimIterations,

109  kLeadRepetitions,

110  1>();

111 

112  constexpr auto input_tile_encode =

113  InputTileDistributionEncoding<decltype(xdllevel_dstr_encoding),

114  kLeadIterPerWarp,

115  kSecondIterPerWarp,

116  kLeadNumWarps,

117  kSecondNumWarps>();

118  constexpr auto block_dstr = make_static_tile_distribution(input_tile_encode);

119  return block_dstr;

120  }

121 };

122 

123 } // namespace ck_tile

batched_transpose_common_policy.hpp

CK_TILE_DEVICE

#define CK_TILE_DEVICE

Definition: config.hpp:41

core.hpp

ck_tile

Definition: cluster_descriptor.hpp:13

ck_tile::make_naive_tensor_descriptor

constexpr CK_TILE_HOST_DEVICE auto make_naive_tensor_descriptor(const tuple< Lengths... > &lengths, const tuple< Strides... > &strides, number< GuaranteedLastDimensionVectorLength >=number<-1 >{}, number< GuaranteedLastDimensionVectorStride >=number<-1 >{})

Definition: tensor_descriptor.hpp:268

ck_tile::integer_least_multiple

constexpr CK_TILE_HOST_DEVICE auto integer_least_multiple(X x, Y y)

Definition: math.hpp:155

ck_tile::make_merge_transform

constexpr CK_TILE_HOST_DEVICE auto make_merge_transform(const LowLengths &low_lengths)

Definition: coordinate_transform.hpp:1615

ck_tile::index_t

int32_t index_t

Definition: integer.hpp:9

ck_tile::make_pass_through_transform

constexpr CK_TILE_HOST_DEVICE auto make_pass_through_transform(const LowLength &low_length)

Definition: coordinate_transform.hpp:1558

ck_tile::transform_tensor_descriptor

constexpr CK_TILE_HOST_DEVICE auto transform_tensor_descriptor(const OldTensorDescriptor &old_tensor_desc, const NewTransforms &new_transforms, NewLowerDimensionOldTopIdss, NewUpperDimensionNewTopIdss)

Definition: tensor_descriptor.hpp:197

ck_tile::InputTileDistributionEncoding

constexpr CK_TILE_HOST_DEVICE auto InputTileDistributionEncoding()

Definition: load_tile_transpose.hpp:351

ck_tile::make_tuple

constexpr CK_TILE_HOST_DEVICE auto make_tuple(Xs &&... xs)

Definition: tuple.hpp:360

ck_tile::make_transposed_distr_encode

constexpr CK_TILE_DEVICE auto make_transposed_distr_encode()

Definition: amd_transpose_load_encoding.hpp:82

ck_tile::make_static_tile_distribution

constexpr CK_TILE_HOST_DEVICE auto make_static_tile_distribution(StaticTileDistributionEncoding_)

Definition: tile_distribution.hpp:480

ck_tile::BatchedTransposeCommonPolicy

Definition: batched_transpose_common_policy.hpp:11

ck_tile::BatchedTransposeLdsPolicy

Definition: batched_transpose_lds_policy.hpp:12

ck_tile::BatchedTransposeLdsPolicy::MakeLdsStoreBlockDescriptor

static constexpr CK_TILE_DEVICE auto MakeLdsStoreBlockDescriptor()

Definition: batched_transpose_lds_policy.hpp:36

ck_tile::BatchedTransposeLdsPolicy::MakeLdsLoadTileDistribution

static constexpr CK_TILE_DEVICE auto MakeLdsLoadTileDistribution()

Definition: batched_transpose_lds_policy.hpp:88

ck_tile::BatchedTransposeLdsPolicy::MakeOutputDistribution

static constexpr CK_TILE_DEVICE auto MakeOutputDistribution()

Definition: batched_transpose_lds_policy.hpp:23

ck_tile::BatchedTransposeLdsPolicy::GetSmemSize

static constexpr CK_TILE_DEVICE index_t GetSmemSize()

Definition: batched_transpose_lds_policy.hpp:14

ck_tile::BatchedTransposeLdsPolicy::MakeLdsLoadBlockDescriptor

static constexpr CK_TILE_DEVICE auto MakeLdsLoadBlockDescriptor()

Definition: batched_transpose_lds_policy.hpp:62

ck_tile::TransposeTileDistributionTraits

Definition: load_tile_transpose.hpp:207

ck_tile::constant

Definition: integral_constant.hpp:13

ck_tile::sequence

Definition: sequence.hpp:49