/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck/tensor_operation/gpu/grid/gridwise_set_multiple_buffer_value.hpp Source File

• develop
• include
• ck
• tensor_operation
• gpu
• grid

gridwise_set_multiple_buffer_value.hpp

Go to the documentation of this file.

1 // SPDX-License-Identifier: MIT

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

3 

4 #pragma once

5 

6 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

7 #include "ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp"

8 

9 namespace ck {

10 

11 template <typename Grid1dBufferDescTuple,

12  index_t NumBuffer,

13  index_t BlockSize,

14  typename DataTypePointerTuple,

15  typename DataTypeTuple>

16 __global__ void

17 kernel_multiple_buffer_set_value(const Grid1dBufferDescTuple grid_1d_buffer_desc_tuple,

18  DataTypePointerTuple p_global_tuple,

19  DataTypeTuple value_tuple)

20 

21 {

22  static_assert(NumBuffer == DataTypePointerTuple::Size() && NumBuffer == DataTypeTuple::Size(),

23  "The tuple size should be same as NumBuffer!");

24 

25  static_for<0, NumBuffer, 1>{}([&](auto iB) {

26  using DataTypePointer = remove_cvref_t<decltype(DataTypePointerTuple{}[iB])>;

27  using DataTypeFromPointer = remove_pointer_t<DataTypePointer>;

28  using DataType = remove_cvref_t<decltype(DataTypeTuple{}[iB])>;

29 

30  static_assert(is_same<DataType, DataTypeFromPointer>::value,

31  "Types in tuples does not match!");

32  });

33 

34  constexpr auto I0 = Number<0>{};

35 

36  const index_t thread_global_id = get_thread_global_1d_id();

37 

38  auto value_buf_tuple = generate_tuple(

39  [&](auto iB) {

40  using DataType = remove_cvref_t<decltype(DataTypeTuple{}[iB])>;

41 

42  return StaticBuffer<AddressSpaceEnum::Vgpr, DataType, 1, true>{};

43  },

44  Number<NumBuffer>{});

45 

46  static_for<0, NumBuffer, 1>{}([&](auto iB) {

47  static_for<0, 1, 1>{}([&](auto J) { value_buf_tuple(iB)(J) = value_tuple[iB]; });

48  });

49 

50  auto global_buf_tuple = generate_tuple(

51  [&](auto iB) {

52  return make_dynamic_buffer<AddressSpaceEnum::Global>(

53  p_global_tuple(iB), grid_1d_buffer_desc_tuple[iB].GetElementSpaceSize());

54  },

55  Number<NumBuffer>{});

56 

57  constexpr auto val_buff_desc = make_naive_tensor_descriptor_packed(make_tuple(Number<1>{}));

58 

59  static_for<0, NumBuffer, 1>{}([&](auto iB) {

60  using DataType = remove_cvref_t<decltype(DataTypeTuple{}[iB])>;

61  using PassThroughOp = tensor_operation::element_wise::PassThrough;

62 

63  auto threadwise_store =

64  ThreadwiseTensorSliceTransfer_v1r3<DataType,

65  DataType,

66  decltype(val_buff_desc),

67  decltype(Grid1dBufferDescTuple{}[iB]),

68  PassThroughOp,

69  Sequence<1>,

70  Sequence<0>,

71  0,

72  1,

73  InMemoryDataOperationEnum::Set,

74  1,

75  true>(

76  grid_1d_buffer_desc_tuple[iB], make_multi_index(thread_global_id), PassThroughOp{});

77 

78  threadwise_store.Run(val_buff_desc,

79  make_tuple(I0),

80  value_buf_tuple(iB),

81  grid_1d_buffer_desc_tuple[iB],

82  global_buf_tuple(iB));

83  });

84 };

85 

86 } // namespace ck

element_wise_operation.hpp

ck

Definition: ck.hpp:267

ck::make_multi_index

__host__ constexpr __device__ auto make_multi_index(Xs &&... xs)

Definition: array_multi_index.hpp:15

ck::generate_tuple

__host__ constexpr __device__ auto generate_tuple(F &&f, Number< N >)

Definition: tuple_helper.hpp:21

ck::InMemoryDataOperationEnum::Set

@ Set

ck::remove_pointer_t

typename remove_pointer< T >::type remove_pointer_t

Definition: type.hpp:300

ck::make_naive_tensor_descriptor_packed

__host__ constexpr __device__ auto make_naive_tensor_descriptor_packed(const Tuple< Lengths... > &lengths)

Definition: tensor_descriptor_helper.hpp:101

ck::kernel_multiple_buffer_set_value

__global__ void kernel_multiple_buffer_set_value(const Grid1dBufferDescTuple grid_1d_buffer_desc_tuple, DataTypePointerTuple p_global_tuple, DataTypeTuple value_tuple)

Definition: gridwise_set_multiple_buffer_value.hpp:17

ck::get_thread_global_1d_id

__device__ index_t get_thread_global_1d_id()

Definition: get_id.hpp:54

ck::make_tuple

__host__ constexpr __device__ auto make_tuple(Xs &&... xs)

Definition: tuple.hpp:211

ck::remove_cvref_t

remove_cv_t< remove_reference_t< T > > remove_cvref_t

Definition: type.hpp:297

ck::index_t

int32_t index_t

Definition: ck.hpp:298

ck::Sequence

Definition: sequence.hpp:43

ck::StaticBuffer

Definition: static_buffer.hpp:16

ck::ThreadwiseTensorSliceTransfer_v1r3

Definition: threadwise_tensor_slice_transfer.hpp:39

ck::integral_constant

Definition: integral_constant.hpp:20

ck::is_same

Definition: type.hpp:177

ck::static_for

Definition: functional2.hpp:33

ck::tensor_operation::element_wise::PassThrough

Definition: unary_element_wise_operation.hpp:334

threadwise_tensor_slice_transfer.hpp