/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/core/utility/reduce_operator.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/core/utility/reduce_operator.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/develop/include/ck_tile/core/utility/reduce_operator.hpp Source File
reduce_operator.hpp
Go to the documentation of this file.
1 // Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
2 // SPDX-License-Identifier: MIT
3 
4 #pragma once
5 
6 #include "ck_tile/core/config.hpp"
7 #include "ck_tile/core/arch/generic_memory_space_atomic.hpp"
8 #include "ck_tile/core/utility/type_traits.hpp"
9 
10 namespace ck_tile {
11 
12 namespace ReduceOp {
13 // y = ReduceOp(y, x);
14 struct Add
15 {
16  template <typename T>
17  CK_TILE_HOST_DEVICE static constexpr T GetIdentityValue()
18  {
19  return type_convert<T>(0.0f);
20  };
21 
22  template <typename T,
23  typename = std::enable_if_t<is_any_of<T, float, double, int32_t, int8_t>::value>>
24  CK_TILE_HOST_DEVICE constexpr T operator()(const T& y, const T x) const
25  {
26  return y + x;
27  }
28 
29  template <typename T,
30  typename = std::enable_if_t<is_any_of<T, half_t, bf16_t, fp8_t, bf8_t>::value>>
31  CK_TILE_HOST_DEVICE constexpr T operator()(T& y, T x) const
32  {
33  float y_ = type_convert<float>(y);
34  float x_ = type_convert<float>(x);
35 
36  return type_convert<T>(y_ + x_);
37  }
38 
39  CK_TILE_HOST_DEVICE static constexpr auto GetAtomic()
40  {
41  return memory_operation_enum::atomic_add;
42  }
43 };
44 
45 struct SquareAdd
46 {
47  template <typename T>
48  CK_TILE_HOST_DEVICE static constexpr T GetIdentityValue()
49  {
50  return type_convert<T>(0.0f);
51  };
52 
53  template <typename T,
54  typename = std::enable_if_t<is_any_of<T, float, double, int32_t, int8_t>::value>>
55  CK_TILE_HOST_DEVICE constexpr T operator()(const T& y, const T x) const
56  {
57  return y + (x * x);
58  }
59 
60  template <typename T,
61  typename = std::enable_if_t<is_any_of<T, half_t, bf16_t, fp8_t, bf8_t>::value>>
62  CK_TILE_HOST_DEVICE constexpr T operator()(T& y, T x) const
63  {
64  float y_ = type_convert<float>(y);
65  float x_ = type_convert<float>(x);
66  return type_convert<T>(y_ + (x_ * x_));
67  }
68 };
69 
70 struct Max
71 {
72  template <
73  typename T,
74  typename = std::enable_if_t<
75  is_any_of<T, float, double, int32_t, int8_t, half_t, bf16_t, fp8_t, bf8_t>::value>>
76  CK_TILE_HOST_DEVICE static constexpr T GetIdentityValue()
77  {
78  return numeric<T>::lowest();
79  };
80 
81  template <
82  typename T,
83  typename = std::enable_if_t<
84  is_any_of<T, float, double, int32_t, int8_t, half_t, bf16_t, fp8_t, bf8_t>::value>>
85  CK_TILE_HOST_DEVICE constexpr T operator()(const T& y, const T x) const
86  {
87  return max(y, x);
88  }
89 
90  // Overload with changed flag for index tracking
91  template <
92  typename T,
93  typename = std::enable_if_t<
94  is_any_of<T, float, double, int32_t, int8_t, half_t, bf16_t, fp8_t, bf8_t>::value>>
95  CK_TILE_HOST_DEVICE constexpr T operator()(const T& y, const T x, bool& changed) const
96  {
97  T new_max = max(y, x);
98  if(x > y)
99  {
100  changed = true;
101  }
102  return new_max;
103  }
104 };
105 
106 struct AbsMax
107 {
108  template <
109  typename T,
110  typename = std::enable_if_t<
111  is_any_of<T, float, double, int32_t, int8_t, half_t, bf16_t, fp8_t, bf8_t>::value>>
112  CK_TILE_HOST_DEVICE static constexpr T GetIdentityValue()
113  {
114  return numeric<T>::zero();
115  };
116 
117  template <
118  typename T,
119  typename = std::enable_if_t<
120  is_any_of<T, float, double, int32_t, int8_t, half_t, bf16_t, fp8_t, bf8_t>::value>>
121  CK_TILE_HOST_DEVICE constexpr T operator()(const T& y, const T x) const
122  {
123  return max(y, abs(x));
124  }
125 
126  // Overload with changed flag for index tracking
127  template <
128  typename T,
129  typename = std::enable_if_t<
130  is_any_of<T, float, double, int32_t, int8_t, half_t, bf16_t, fp8_t, bf8_t>::value>>
131  CK_TILE_HOST_DEVICE constexpr T operator()(const T& y, const T x, bool& changed) const
132  {
133  T new_max = max(y, abs(x));
134  if(abs(x) > y)
135  {
136  changed = true;
137  }
138  return new_max;
139  }
140 };
141 
142 } // namespace ReduceOp
143 } // namespace ck_tile
CK_TILE_HOST_DEVICE
#define CK_TILE_HOST_DEVICE
Definition: config.hpp:46
ck_tile
Definition: cluster_descriptor.hpp:13
ck_tile::abs
CK_TILE_HOST_DEVICE bfloat16_t abs(const bfloat16_t &x)
Definition: bfloat16.hpp:403
ck_tile::max
constexpr CK_TILE_HOST_DEVICE T max(T x)
Definition: math.hpp:157
ck::enable_if_t
typename std::enable_if< B, T >::type enable_if_t
Definition: enable_if.hpp:27
ck::atomic_add
__device__ X atomic_add(X *p_dst, const X &x)
value
const GenericPointer< typename T::ValueType > T2 value
Definition: pointer.h:1697
ck_tile::ReduceOp::AbsMax
Definition: reduce_operator.hpp:107
ck_tile::ReduceOp::AbsMax::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(const T &y, const T x, bool &changed) const
Definition: reduce_operator.hpp:131
ck_tile::ReduceOp::AbsMax::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(const T &y, const T x) const
Definition: reduce_operator.hpp:121
ck_tile::ReduceOp::AbsMax::GetIdentityValue
static constexpr CK_TILE_HOST_DEVICE T GetIdentityValue()
Definition: reduce_operator.hpp:112
ck_tile::ReduceOp::Add
Definition: reduce_operator.hpp:15
ck_tile::ReduceOp::Add::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(T &y, T x) const
Definition: reduce_operator.hpp:31
ck_tile::ReduceOp::Add::GetIdentityValue
static constexpr CK_TILE_HOST_DEVICE T GetIdentityValue()
Definition: reduce_operator.hpp:17
ck_tile::ReduceOp::Add::GetAtomic
static constexpr CK_TILE_HOST_DEVICE auto GetAtomic()
Definition: reduce_operator.hpp:39
ck_tile::ReduceOp::Add::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(const T &y, const T x) const
Definition: reduce_operator.hpp:24
ck_tile::ReduceOp::Max
Definition: reduce_operator.hpp:71
ck_tile::ReduceOp::Max::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(const T &y, const T x, bool &changed) const
Definition: reduce_operator.hpp:95
ck_tile::ReduceOp::Max::GetIdentityValue
static constexpr CK_TILE_HOST_DEVICE T GetIdentityValue()
Definition: reduce_operator.hpp:76
ck_tile::ReduceOp::Max::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(const T &y, const T x) const
Definition: reduce_operator.hpp:85
ck_tile::ReduceOp::SquareAdd
Definition: reduce_operator.hpp:46
ck_tile::ReduceOp::SquareAdd::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(const T &y, const T x) const
Definition: reduce_operator.hpp:55
ck_tile::ReduceOp::SquareAdd::operator()
constexpr CK_TILE_HOST_DEVICE T operator()(T &y, T x) const
Definition: reduce_operator.hpp:62
ck_tile::ReduceOp::SquareAdd::GetIdentityValue
static constexpr CK_TILE_HOST_DEVICE T GetIdentityValue()
Definition: reduce_operator.hpp:48
ck_tile::is_any_of
Definition: type_traits.hpp:115
ck_tile::numeric::lowest
static constexpr CK_TILE_HOST_DEVICE T lowest()
Definition: numeric.hpp:23
ck_tile::numeric::zero
static constexpr CK_TILE_HOST_DEVICE T zero()
Definition: numeric.hpp:58
generic_memory_space_atomic.hpp
type_traits.hpp