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

• develop
• include
• ck_tile
• ops
• smoothquant
• pipeline

smoothquant_pipeline_default_policy.hpp

Go to the documentation of this file.

1 // SPDX-License-Identifier: MIT

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

3 

4 #pragma once

5 

6 #include "ck_tile/core.hpp"

7 #include "ck_tile/ops/reduce/block/block_reduce2d_problem.hpp"

8 #include "ck_tile/ops/reduce/block/block_reduce2d.hpp"

9 

10 namespace ck_tile {

11 

12 struct SmoothquantPipelineDefaultPolicy

13 {

14  template <typename Problem>

15  CK_TILE_DEVICE static constexpr auto MakeXBlockTileDistribution()

16  {

17  using S = typename Problem::BlockShape;

18 

19  return make_static_tile_distribution(

20  tile_distribution_encoding<

21  sequence<>,

22  tuple<sequence<S::Repeat_M, S::WarpPerBlock_M, S::ThreadPerWarp_M, S::Vector_M>,

23  sequence<S::Repeat_N, S::WarpPerBlock_N, S::ThreadPerWarp_N, S::Vector_N>>,

24  tuple<sequence<1, 2>, sequence<1, 2>>,

25  tuple<sequence<1, 1>, sequence<2, 2>>,

26  sequence<1, 1, 2, 2>,

27  sequence<0, 3, 0, 3>>{});

28  }

29 

30  template <typename Problem>

31  CK_TILE_DEVICE static constexpr auto MakeSmoothScaleBlockTileDistribution()

32  {

33  using S = typename Problem::BlockShape;

34 

35  return make_static_tile_distribution(

36  tile_distribution_encoding<

37  sequence<S::WarpPerBlock_M, S::ThreadPerWarp_M>,

38  tuple<sequence<S::Repeat_N, S::WarpPerBlock_N, S::ThreadPerWarp_N, S::Vector_N>>,

39  tuple<sequence<0, 1>, sequence<0, 1>>,

40  tuple<sequence<0, 1>, sequence<1, 2>>,

41  sequence<1, 1>,

42  sequence<0, 3>>{});

43  }

44 

45  template <typename Problem>

46  CK_TILE_HOST_DEVICE static constexpr auto GetBlockReduce2d()

47  {

48  using P_ = BlockReduce2dProblem<typename Problem::ComputeDataType,

49  typename Problem::ComputeDataType,

50  typename Problem::BlockShape>;

51  return BlockReduce2d<P_>{};

52  }

53 

54  template <typename Problem>

55  CK_TILE_HOST_DEVICE static constexpr auto GetBlockReduce2dSync()

56  {

57  using P_ = BlockReduce2dProblem<typename Problem::ComputeDataType,

58  typename Problem::ComputeDataType,

59  typename Problem::BlockShape>;

60  return BlockReduce2dSync<P_>{};

61  }

62 

63  template <typename Problem>

64  CK_TILE_HOST_DEVICE static constexpr auto GetBlockReduce2dCrossWarpSync()

65  {

66  using P_ = BlockReduce2dProblem<typename Problem::ComputeDataType,

67  typename Problem::ComputeDataType,

68  typename Problem::BlockShape>;

69  return BlockReduce2dCrossWarpSync<P_>{};

70  }

71 

72  template <typename Problem>

73  CK_TILE_HOST_DEVICE static constexpr index_t GetSmemSize()

74  {

75  if constexpr(Problem::kNeedCrossWarpSync)

76  {

77  using P_ = BlockReduce2dProblem<typename Problem::XDataType,

78  typename Problem::ComputeDataType,

79  typename Problem::BlockShape>;

80 

81  using block_reduce2d = BlockReduce2d<P_>;

82  using x_block_tile =

83  decltype(make_static_distributed_tensor<typename Problem::XDataType>(

84  MakeXBlockTileDistribution<Problem>()));

85  using y_block_tile = decltype(block_reduce2d::template MakeYBlockTile<x_block_tile>());

86 

87  return GetBlockReduce2dCrossWarpSync<Problem>().template GetSmemSize<y_block_tile>();

88  }

89  else

90  {

91  return 1; // zero size arrays are an extension

92  }

93  }

94 };

95 } // namespace ck_tile

block_reduce2d.hpp

block_reduce2d_problem.hpp

CK_TILE_DEVICE

#define CK_TILE_DEVICE

Definition: config.hpp:41

CK_TILE_HOST_DEVICE

#define CK_TILE_HOST_DEVICE

Definition: config.hpp:42

core.hpp

ck_tile

Definition: cluster_descriptor.hpp:13

ck_tile::index_t

int32_t index_t

Definition: integer.hpp:9

ck_tile::make_static_tile_distribution

constexpr CK_TILE_HOST_DEVICE auto make_static_tile_distribution(StaticTileDistributionEncoding_)

Definition: tile_distribution.hpp:480

ck_tile::BlockReduce2dCrossWarpSync

Definition: block_reduce2d.hpp:200

ck_tile::BlockReduce2d

Definition: block_reduce2d.hpp:45

ck_tile::BlockReduce2dProblem

Definition: block_reduce2d_problem.hpp:12

ck_tile::BlockReduce2dSync

Definition: block_reduce2d.hpp:135

ck_tile::SmoothquantPipelineDefaultPolicy

Definition: smoothquant_pipeline_default_policy.hpp:13

ck_tile::SmoothquantPipelineDefaultPolicy::GetBlockReduce2d

static constexpr CK_TILE_HOST_DEVICE auto GetBlockReduce2d()

Definition: smoothquant_pipeline_default_policy.hpp:46

ck_tile::SmoothquantPipelineDefaultPolicy::GetBlockReduce2dCrossWarpSync

static constexpr CK_TILE_HOST_DEVICE auto GetBlockReduce2dCrossWarpSync()

Definition: smoothquant_pipeline_default_policy.hpp:64

ck_tile::SmoothquantPipelineDefaultPolicy::GetSmemSize

static constexpr CK_TILE_HOST_DEVICE index_t GetSmemSize()

Definition: smoothquant_pipeline_default_policy.hpp:73

ck_tile::SmoothquantPipelineDefaultPolicy::MakeXBlockTileDistribution

static constexpr CK_TILE_DEVICE auto MakeXBlockTileDistribution()

Definition: smoothquant_pipeline_default_policy.hpp:15

ck_tile::SmoothquantPipelineDefaultPolicy::MakeSmoothScaleBlockTileDistribution

static constexpr CK_TILE_DEVICE auto MakeSmoothScaleBlockTileDistribution()

Definition: smoothquant_pipeline_default_policy.hpp:31

ck_tile::SmoothquantPipelineDefaultPolicy::GetBlockReduce2dSync

static constexpr CK_TILE_HOST_DEVICE auto GetBlockReduce2dSync()

Definition: smoothquant_pipeline_default_policy.hpp:55

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