"vscode:/vscode.git/clone" did not exist on "d5398a0d8d4ff30384058783fe5edfff24a8dcf2"
reduction_functions_accumulate.hpp 3.34 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
5
6
#pragma once

#include "ck/utility/data_type.hpp"
#include "ck/utility/math_v2.hpp"
#include "ck/utility/reduction_common.hpp"
#include "ck/utility/reduction_operator.hpp"
7
8
9
10

namespace ck {
namespace detail {

11
12
13
14
15
16
17
18
19
20
21
22
23
// Check for NaN; guarantee NaNs are NOT propagated to result (i.e., ignore NaNs)
template <typename ReduceOperation, typename AccDataType>
struct AccumulateWithNanIgnore
{
    __device__ static inline void Calculate(AccDataType& accuVal, AccDataType currVal)
    {
        if(!isnan(currVal))
        {
            ReduceOperation{}(accuVal, currVal);
        }
    };
};

24
25
template <bool PropagateNan, typename ReduceOperation, typename AccDataType>
struct AccumulateWithNanCheck;
26

27
28
29
30
31
// Does not check for NaN; does not guarantee NaNs be propagated to result
// e.g., given that max(a, b) = a > b ? a : b
// then  max(NaN, 1) returns 1
//       max(1, NaN) returns NaN
// since any comparison involving NaNs returns false
32
33
template <typename ReduceOperation, typename AccDataType>
struct AccumulateWithNanCheck<false, ReduceOperation, AccDataType>
34
35
{
    // cppcheck-suppress constParameter
36
    __host__ __device__ static inline void Calculate(AccDataType& accuVal, AccDataType currVal)
37
38
39
40
41
    {
        ReduceOperation{}(accuVal, currVal);
    };
};

42
// Check for NaN; guarantees NaNs be propagated to result
43
44
45
template <typename ReduceOperation, typename AccDataType>
struct AccumulateWithNanCheck<true, ReduceOperation, AccDataType>
{
46
    __host__ __device__ static inline void Calculate(AccDataType& accuVal, AccDataType currVal)
47
    {
48
49
50
        using ck::math::isnan;

        if(isnan(currVal))
51
52
53
54
55
56
57
        {
            accuVal = currVal;
        }
        else
        {
            ReduceOperation{}(accuVal, currVal);
        };
58
    };
59
60
61
62
};

template <bool PropagateNan, typename ReduceOperation, typename AccDataType, typename IndexDataType>
struct AccumulateWithIndexAndNanCheck;
63

64
65
66
template <typename ReduceOperation, typename AccDataType, typename IndexDataType>
struct AccumulateWithIndexAndNanCheck<false, ReduceOperation, AccDataType, IndexDataType>
{
67
    __host__ __device__ static inline void
68
    // cppcheck-suppress constParameter
69
70
71
72
    Calculate(AccDataType& accuVal,
              AccDataType currVal,
              IndexDataType& accuIndex,
              IndexDataType currIndex)
73
74
75
    {
        bool changed = false;

76
        ReduceOperation{}(accuVal, currVal, changed);
77
78
79
80
81
82

        if(changed)
            accuIndex = currIndex;
    };
};

83
84
template <typename ReduceOperation, typename AccDataType, typename IndexDataType>
struct AccumulateWithIndexAndNanCheck<true, ReduceOperation, AccDataType, IndexDataType>
85
{
86
    // The method is called when the ReduceOperation is indexable and the user asked for indices
87
88
89
90
    __host__ __device__ static inline void Calculate(AccDataType& accuVal,
                                                     AccDataType currVal,
                                                     IndexDataType& accuIndex,
                                                     IndexDataType currIndex)
91
    {
92
93
94
        using ck::math::isnan;

        if(isnan(currVal))
95
96
97
98
99
100
101
102
        {
            accuVal   = currVal;
            accuIndex = currIndex;
        }
        else
        {
            bool changed = false;

103
            ReduceOperation{}(accuVal, currVal, changed);
104
105
106
107
108
109
110

            if(changed)
                accuIndex = currIndex;
        }
    };
};

Chao Liu's avatar
Chao Liu committed
111
112
} // namespace detail
} // namespace ck