include/ck_tile/core/tensor/store_tile.hpp Source File

• include
• ck_tile
• core
• tensor

store_tile.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/config.hpp"

7 #include "ck_tile/core/numeric/integer.hpp"

8 #include "ck_tile/core/numeric/integral_constant.hpp"

9 #include "ck_tile/core/algorithm/coordinate_transform.hpp"

10 #include "ck_tile/core/container/container_helper.hpp"

11 #include "ck_tile/core/numeric/math.hpp"

12 #include "ck_tile/core/tensor/tile_window.hpp"

13 #include "ck_tile/core/tensor/tile_window_linear.hpp"

14 #include "ck_tile/core/utility/type_traits.hpp"

15 

16 namespace ck_tile {

17 

18 template <typename BottomTensorView_,

19  typename WindowLengths_,

20  typename TileDistribution_,

21  typename DataType_>

22 CK_TILE_DEVICE void

23 store_tile(tile_window_with_static_lengths<BottomTensorView_, WindowLengths_>& tile_window_tmp,

24  const static_distributed_tensor<DataType_, TileDistribution_>& dstr_tensor)

25 {

26  using DataType = remove_cvref_t<typename BottomTensorView_::DataType>;

27  using TileDstr = remove_cvref_t<TileDistribution_>;

28 

29  static_assert(std::is_same_v<remove_cvref_t<DataType_>, DataType>, "wrong!");

30 

31  constexpr auto tile_dstr = TileDstr{};

32 

33  auto tile_window = make_tile_window(tile_window_tmp.get_bottom_tensor_view(),

34  tile_window_tmp.get_window_lengths(),

35  tile_window_tmp.get_window_origin(),

36  tile_dstr);

37 

38  tile_window.store(dstr_tensor);

39 }

40 

41 template <typename BottomTensorView_,

42  typename WindowLengths_,

43  typename TileDistribution_,

44  typename DataType_>

45 CK_TILE_DEVICE void

46 store_tile_raw(tile_window_with_static_lengths<BottomTensorView_, WindowLengths_>& tile_window_tmp,

47  const static_distributed_tensor<DataType_, TileDistribution_>& dstr_tensor)

48 {

49  using DataType = remove_cvref_t<typename BottomTensorView_::DataType>;

50  using TileDstr = remove_cvref_t<TileDistribution_>;

51 

52  static_assert(std::is_same_v<remove_cvref_t<DataType_>, DataType>, "wrong!");

53 

54  constexpr auto tile_dstr = TileDstr{};

55 

56  auto tile_window = make_tile_window(tile_window_tmp.get_bottom_tensor_view(),

57  tile_window_tmp.get_window_lengths(),

58  tile_window_tmp.get_window_origin(),

59  tile_dstr);

60 

61  tile_window.store_raw(dstr_tensor);

62 }

63 

64 template <typename BottomTensorView_,

65  typename WindowLengths_,

66  typename TileDistribution_,

67  index_t NumCoord,

68  typename DataType_>

69 CK_TILE_DEVICE void

70 store_tile(tile_window_with_static_distribution<BottomTensorView_,

71  WindowLengths_,

72  TileDistribution_,

73  NumCoord>& tile_window,

74  const static_distributed_tensor<DataType_, TileDistribution_>& dstr_tensor)

75 {

76  tile_window.store(dstr_tensor, number<-1>{});

77 }

78 

79 template <typename BottomTensorView_,

80  typename WindowLengths_,

81  typename TileDistribution_,

82  index_t NumCoord,

83  typename DataType_>

84 CK_TILE_DEVICE void

85 store_tile_raw(tile_window_with_static_distribution<BottomTensorView_,

86  WindowLengths_,

87  TileDistribution_,

88  NumCoord>& tile_window,

89  const static_distributed_tensor<DataType_, TileDistribution_>& dstr_tensor)

90 {

91  tile_window.store_raw(dstr_tensor, number<-1>{});

92 }

93 

94 template <typename BottomTensorView_,

95  typename WindowLengths_,

96  typename TileDistribution_,

97  typename LinearBottomDims_,

98  typename DataType_>

99 CK_TILE_DEVICE void store_tile(

100  tile_window_linear<BottomTensorView_, WindowLengths_, TileDistribution_, LinearBottomDims_>&

101  tile_window,

102  const static_distributed_tensor<DataType_, TileDistribution_>& dstr_tensor)

103 {

104  tile_window.store(dstr_tensor, number<-1>{});

105 }

106 

107 template <typename BottomTensorView_,

108  typename WindowLengths_,

109  typename TileDistribution_,

110  typename LinearBottomDims_,

111  typename DataType_>

112 CK_TILE_DEVICE void store_tile_raw(

113  tile_window_linear<BottomTensorView_, WindowLengths_, TileDistribution_, LinearBottomDims_>&

114  tile_window,

115  const static_distributed_tensor<DataType_, TileDistribution_>& dstr_tensor)

116 {

117  tile_window.store_raw(dstr_tensor, number<-1>{});

118 }

119 

120 } // namespace ck_tile

config.hpp

CK_TILE_DEVICE

#define CK_TILE_DEVICE

Definition: config.hpp:40

coordinate_transform.hpp

integer.hpp

ck_tile

Definition: cluster_descriptor.hpp:13

ck_tile::index_t

int32_t index_t

Definition: integer.hpp:9

ck_tile::remove_cvref_t

remove_cv_t< std::remove_reference_t< T > > remove_cvref_t

Definition: type_traits.hpp:20

ck_tile::store_tile_raw

CK_TILE_DEVICE void store_tile_raw(tile_window_with_static_lengths< BottomTensorView_, WindowLengths_ > &tile_window_tmp, const static_distributed_tensor< DataType_, TileDistribution_ > &dstr_tensor)

Definition: store_tile.hpp:46

ck_tile::make_tile_window

constexpr CK_TILE_DEVICE auto make_tile_window(null_tensor_view, const WindowLengths &window_lengths, const multi_index< WindowLengths::size()> &, Ts &&...)

Definition: null_tile_window.hpp:72

ck_tile::store_tile

CK_TILE_DEVICE void store_tile(tile_window_with_static_lengths< BottomTensorView_, WindowLengths_ > &tile_window_tmp, const static_distributed_tensor< DataType_, TileDistribution_ > &dstr_tensor)

Definition: store_tile.hpp:23

ck::is_same_v

constexpr bool is_same_v

Definition: type.hpp:283

ck_tile::constant

Definition: integral_constant.hpp:13

ck_tile::static_distributed_tensor

Definition: static_distributed_tensor.hpp:21

ck_tile::tile_window_linear

Definition: tile_window_linear.hpp:46

ck_tile::tile_window_linear::store_raw

CK_TILE_DEVICE void store_raw(const static_distributed_tensor< DataType, TileDstr > &dstr_tensor, number< i_access >={}) const

Definition: tile_window_linear.hpp:833

ck_tile::tile_window_linear::store

CK_TILE_DEVICE void store(const static_distributed_tensor< DataType, TileDstr > &dstr_tensor, number< i_access >={}, bool_constant< oob_conditional_check >={}) const

Definition: tile_window_linear.hpp:784

ck_tile::tile_window_with_static_distribution

This class provides tile (windowed) view and access to the device memory.

Definition: tile_window.hpp:37

ck_tile::tile_window_with_static_lengths

This class provides description of tile windowed view on the device memory.

Definition: tile_window.hpp:1031

ck_tile::tile_window_with_static_lengths::get_window_lengths

constexpr CK_TILE_DEVICE auto get_window_lengths() const

Definition: tile_window.hpp:1058

ck_tile::tile_window_with_static_lengths::get_bottom_tensor_view

constexpr CK_TILE_DEVICE auto get_bottom_tensor_view() const

Definition: tile_window.hpp:1060

ck_tile::tile_window_with_static_lengths::get_window_origin

constexpr CK_TILE_DEVICE auto get_window_origin() const

Definition: tile_window.hpp:1062

container_helper.hpp

integral_constant.hpp

math.hpp

tile_window.hpp

tile_window_linear.hpp

type_traits.hpp