reduction_operator.hpp 8.64 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
5
#pragma once

#include "ck/ck.hpp"
#include "ck/utility/data_type.hpp"
#include "ck/utility/type.hpp"
6
7
8
9
10
11
12
13

namespace ck {

namespace reduce {

// Every binary operator used in reduction is represented by a templated functor class. Each functor
// class must provide at least
// three members:
14
// 1) GetIdentityValue() -- the interface to return the "identity element" for the binary
15
// operator, "identity element" is the unique
16
//                    element in the algebraic space that doesn't affect the value of other elements
17
18
19
//                    when operated against them, and the concept is similar to zero vector in
//                    vector space
//                    (http://pages.cs.wisc.edu/~matthewb/pages/notes/pdf/linearalgebra/VectorSpaces.pdf).
20
21
22
23
// 2) IsCompatibleInMemoryDataOperation() -- return true if the reduction task corresponding to this
// operator can use the InMemoryDataOperation to finalize, or else it return false 3) operator() --
// the first argument of the operator must be both an input & output, and the corresponding variable
// usually stores
24
25
26
27
28
29
30
31
32
33
//                  the accumulated result of many operator() calls; the second argument is only an
//                  input. For indexable binary
//                  operator, the second version of operator() has third argument (which is an
//                  output) to indicate whether the
//                  accumulated value (the first argument) has changed, in which case the recorded
//                  accumulated index also need be
//                  changed.

struct Add
{
34
35
36
37
38
    template <typename T>
    __host__ __device__ static constexpr T GetIdentityValue()
    {
        return type_convert<T>(0.0f);
    };
39

40
    __host__ __device__ static constexpr bool
41
42
43
44
45
46
    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
    {
        return operation == InMemoryDataOperationEnum::AtomicAdd ||
               operation == InMemoryDataOperationEnum::Set;
    };

47
48
49
50
51
52
53
54
55
    template <typename T>
    __host__ __device__ inline constexpr void operator()(T& a, T b) const
    {
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, int32_t>::value,
                      "The data type is not supported by the Add accumulator!");

        a = a + b;
    }
56
57
58
59
};

struct Mul
{
60
61
62
63
64
    template <typename T>
    __host__ __device__ static constexpr T GetIdentityValue()
    {
        return type_convert<T>(1.0f);
    };
65

66
    __host__ __device__ static constexpr bool
67
68
69
70
71
    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
    {
        return operation == InMemoryDataOperationEnum::Set;
    };

72
73
74
75
76
77
78
79
80
    template <typename T>
    __host__ __device__ inline constexpr void operator()(T& a, T b) const
    {
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, int32_t>::value,
                      "The data type is not supported by the Mul accumulator!");

        a = a * b;
    }
81
82
83
84
};

struct Max
{
85
    template <typename T>
86
    __host__ __device__ static constexpr T GetIdentityValue()
87
88
89
    {
        return NumericLimits<T>::Lowest();
    };
90

91
    __host__ __device__ static constexpr bool
92
93
94
95
96
97
    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
    {
        // ToChange: atomic_max to be added
        return operation == InMemoryDataOperationEnum::Set;
    };

98
    template <typename T>
99
    __host__ __device__ inline constexpr void operator()(T& a, T b) const
100
    {
101
102
103
104
105
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                          is_same<T, int8_t>::value,
                      "The data type is not supported by the Max accumulator!");

106
107
108
109
        if(a < b)
            a = b;
    }

110
    template <typename T>
111
    __host__ __device__ inline constexpr void operator()(T& a, T b, bool& changed) const
112
    {
113
114
115
116
117
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                          is_same<T, int8_t>::value,
                      "The data type is not supported by the Max accumulator!");

118
119
120
121
122
123
124
125
126
127
        if(a < b)
        {
            a       = b;
            changed = true;
        }
    }
};

struct Min
{
128
129
130
131
132
    template <typename T>
    __host__ __device__ static constexpr T GetIdentityValue()
    {
        return NumericLimits<T>::Max();
    };
133

134
    __host__ __device__ static constexpr bool
135
136
137
138
139
140
    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
    {
        // ToChange: atomic_min to be added
        return operation == InMemoryDataOperationEnum::Set;
    };

141
    template <typename T>
142
    __host__ __device__ inline constexpr void operator()(T& a, T b) const
143
    {
144
145
146
147
148
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                          is_same<T, int8_t>::value,
                      "The data type is not supported by the Min accumulator!");

149
150
151
152
        if(a > b)
            a = b;
    }

153
    template <typename T>
154
    __host__ __device__ inline constexpr void operator()(T& a, T b, bool& changed) const
155
    {
156
157
158
159
160
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                          is_same<T, int8_t>::value,
                      "The data type is not supported by the Min accumulator!");

161
162
163
164
165
166
167
168
        if(a > b)
        {
            a       = b;
            changed = true;
        }
    }
};

169
struct AMax
170
{
171
172
173
174
175
    template <typename T>
    __host__ __device__ static constexpr T GetIdentityValue()
    {
        return type_convert<T>(0.0f);
    };
176

177
    __host__ __device__ static constexpr bool
178
179
180
181
182
183
    IsCompatibleInMemoryDataOperation(InMemoryDataOperationEnum operation)
    {
        // ToChange: atomic_max to be added
        return operation == InMemoryDataOperationEnum::Set;
    };

184
    template <typename T>
185
    __host__ __device__ inline constexpr void operator()(T& a, T b) const
186
    {
187
188
189
190
191
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                          is_same<T, int8_t>::value,
                      "The data type is not supported by the AMax accumulator!");

192
193
194
195
        if(a < b)
            a = b;
    }

196
    template <typename T>
197
    __host__ __device__ inline constexpr void operator()(T& a, T b, bool& changed) const
198
    {
199
200
201
202
203
        static_assert(is_same<T, float>::value || is_same<T, double>::value ||
                          is_same<T, half_t>::value || is_same<T, int32_t>::value ||
                          is_same<T, int8_t>::value,
                      "The data type is not supported by the AMax accumulator!");

204
205
206
207
208
209
        if(a < b)
        {
            a       = b;
            changed = true;
        }
    }
210
211
};

212
template <typename T>
213
constexpr T GetIdentityValueForInMemoryDataOperation(InMemoryDataOperationEnum operation)
214
215
216
217
218
219
220
221
222
{
    T result = ck::type_convert<T>(0.0f);

    if(operation == InMemoryDataOperationEnum::AtomicMax)
        result = ck::NumericLimits<T>::Lowest();

    return (result);
};

223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
template <InMemoryDataOperationEnum Operation, typename DataType>
struct InMemoryDataOperatonSupportedOnDataType
{
    static constexpr bool value = false;
};

template <typename DataType>
struct InMemoryDataOperatonSupportedOnDataType<InMemoryDataOperationEnum::AtomicAdd, DataType>
{
    static constexpr bool value =
        is_same<DataType, float>::value || is_same<DataType, double>::value;
};

template <typename DataType>
struct InMemoryDataOperatonSupportedOnDataType<InMemoryDataOperationEnum::AtomicMax, DataType>
{
    static constexpr bool value =
        is_same<DataType, float>::value || is_same<DataType, double>::value;
};

template <typename DataType>
struct InMemoryDataOperatonSupportedOnDataType<InMemoryDataOperationEnum::Set, DataType>
{
    static constexpr bool value =
        is_same<DataType, float>::value || is_same<DataType, double>::value ||
        is_same<DataType, half_t>::value || is_same<DataType, bhalf_t>::value ||
        is_same<DataType, int8_t>::value || is_same<DataType, int32_t>::value;
};

template <typename DataType>
struct InMemoryDataOperatonSupportedOnDataType<InMemoryDataOperationEnum::Add, DataType>
{
    static constexpr bool value =
        is_same<DataType, float>::value || is_same<DataType, double>::value ||
        is_same<DataType, half_t>::value || is_same<DataType, int8_t>::value ||
        is_same<DataType, int32_t>::value;
};

Chao Liu's avatar
Chao Liu committed
261
262
} // namespace reduce
} // namespace ck