reduce_sum.cpp 5.55 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
#include <migraphx/gpu/device/reduce_sum.hpp>
#include <migraphx/gpu/device/launch.hpp>
#include <migraphx/gpu/device/visit.hpp>
Paul's avatar
Paul committed
4
#include <migraphx/requires.hpp>
Paul's avatar
Paul committed
5
6
7
8
9
10
11
12

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {
namespace device {

struct sum
{
Paul's avatar
Paul committed
13
14
    template <class T, class U>
    MIGRAPHX_DEVICE_CONSTEXPR auto operator()(T x, U y) const
Paul's avatar
Paul committed
15
16
17
18
19
    {
        return x + y;
    }
};

Paul's avatar
Paul committed
20
#ifdef MIGRAPHX_NO_DPP
Paul's avatar
Paul committed
21
template <std::size_t N, class Op, class T, class F>
Paul's avatar
Paul committed
22
23
24
25
26
__device__ auto block_reduce(index idx, Op op, T init, std::size_t n, F f)
{
    using type = decltype(f(idx.local));
    MIGRAPHX_DEVICE_SHARED type buffer[N];
    type x = init;
Paul's avatar
Paul committed
27
    idx.local_stride(n, [&](auto i) { x = op(x, f(i)); });
Paul's avatar
Paul committed
28
29
30
    buffer[idx.local] = x;
    __syncthreads();

Paul's avatar
Paul committed
31
    for(std::size_t s = 1; s < idx.nlocal(); s *= 2)
Paul's avatar
Paul committed
32
33
    {
        const std::size_t index = 2 * s * idx.local;
Paul's avatar
Paul committed
34
        if(index < idx.nlocal())
Paul's avatar
Paul committed
35
36
37
38
39
40
41
        {
            buffer[index] = op(buffer[index], buffer[index + s]);
        }
        __syncthreads();
    }
    return buffer[0];
}
Paul's avatar
Paul committed
42
#else
Paul's avatar
Paul committed
43
constexpr unsigned int dpp_row_shr(unsigned int x) { return 0x110 | x; }
Paul's avatar
Paul committed
44
45
46
47
48
49

constexpr unsigned int dpp_row_bcast(unsigned int x)
{
    unsigned int y = 0;
    switch(x)
    {
Paul's avatar
Paul committed
50
51
52
    case 15: y = 0x142; break;
    case 31: y = 0x143; break;
    default: throw std::runtime_error("Unknown bcast");
Paul's avatar
Paul committed
53
54
55
56
    }
    return y;
}

Paul's avatar
Paul committed
57
58
59
60
61
template <unsigned int DppCtrl,
          unsigned int RowMask  = 0xf,
          unsigned int BankMask = 0xf,
          bool BoundCtrl        = false,
          class T>
Paul's avatar
Paul committed
62
63
64
65
66
67
68
69
70
71
72
__device__ T dpp_mov(T& x)
{
    static const std::size_t n = sizeof(T) < 4 ? 1 : sizeof(T) / 4;
    union type
    {
        uint32_t reg[n];
        T data;
    };
    type output;
    type input;
    input.data = x;
Paul's avatar
Paul committed
73
    for(std::size_t i = 0; i < n; i++)
Paul's avatar
Paul committed
74
75
76
77
78
79
    {
        output.reg[i] = __llvm_amdgcn_move_dpp(input.reg[i], DppCtrl, RowMask, BankMask, BoundCtrl);
    }
    return output.data;
}

Paul's avatar
Paul committed
80
template <class T, class Op>
Paul's avatar
Paul committed
81
82
83
84
__device__ void dpp_reduce(T& in, Op op)
{
    T out;
    out = dpp_mov<dpp_row_shr(1)>(in);
Paul's avatar
Paul committed
85
    in  = op(in, out);
Paul's avatar
Paul committed
86
    out = dpp_mov<dpp_row_shr(2)>(in);
Paul's avatar
Paul committed
87
    in  = op(in, out);
Paul's avatar
Paul committed
88
    out = dpp_mov<dpp_row_shr(4), 0xf, 0xe>(in);
Paul's avatar
Paul committed
89
    in  = op(in, out);
Paul's avatar
Paul committed
90
    out = dpp_mov<dpp_row_shr(8), 0xf, 0xc>(in);
Paul's avatar
Paul committed
91
    in  = op(in, out);
Paul's avatar
Paul committed
92
    out = dpp_mov<dpp_row_bcast(15), 0xa>(in);
Paul's avatar
Paul committed
93
    in  = op(in, out);
Paul's avatar
Paul committed
94
    out = dpp_mov<dpp_row_bcast(31), 0xc>(in);
Paul's avatar
Paul committed
95
    in  = op(in, out);
Paul's avatar
Paul committed
96
}
Paul's avatar
Paul committed
97

Paul's avatar
Paul committed
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
__device__ void dpp_reduce(float& x, sum)
{
    __asm__ volatile("s_nop 4\n"
                     "v_add_f32 %0 %0 %0 row_shr:1\n"
                     "s_nop 1\n"
                     "v_add_f32 %0 %0 %0 row_shr:2\n"
                     "s_nop 1\n"
                     "v_add_f32 %0 %0 %0 row_shr:4 bank_mask:0xe\n"
                     "s_nop 1\n"
                     "v_add_f32 %0 %0 %0 row_shr:8 bank_mask:0xc\n"
                     "s_nop 1\n"
                     "v_add_f32 %0 %0 %0 row_bcast:15 row_mask:0xa\n"
                     "s_nop 1\n"
                     "v_add_f32 %0 %0 %0 row_bcast:31 row_mask:0xc\n"
                     "s_nop 1\n"
                     : "=v"(x)
                     : "0"(x));
}

template <std::size_t N, class Op, class T, class F>
__device__ auto block_reduce(index idx, Op op, T init, std::size_t n, F f)
{
Paul's avatar
Paul committed
120
    using type                  = decltype(f(idx.local));
Paul's avatar
Paul committed
121
    const auto std::size_t wave = 64;
Paul's avatar
Paul committed
122
    MIGRAPHX_DEVICE_SHARED type buffer[N / 64];
Paul's avatar
Paul committed
123
124
125
    type x = init;
    idx.local_stride(n, [&](auto i) { x = op(x, f(i)); });
    dpp_reduce(x, op);
Paul's avatar
Paul committed
126

Paul's avatar
Paul committed
127
128
129
130
131
132
133
134
    const auto ldsidx = idx.local / 64;
    if((idx.local % 64) == 63)
    {
        buffer[ldsidx] = x;
    }
    __syncthreads();

    type y = 0;
Paul's avatar
Paul committed
135
    for(std::size_t i = 0; i < idx.nlocal() / 64; i++)
Paul's avatar
Paul committed
136
137
138
139
140
141
    {
        y += buffer[i];
    }
    return y;
}
#endif
Paul's avatar
Paul committed
142
143
constexpr std::size_t compute_block_size(std::size_t n, std::size_t max_block_size)
{
Paul's avatar
Paul committed
144
    size_t block_size = 64;
Paul's avatar
Paul committed
145
146
147
148
149
    while(block_size < max_block_size and block_size < n)
        block_size *= 2;
    return block_size;
}

Paul's avatar
Paul committed
150
151
152
void reduce_sum(hipStream_t stream, const argument& result, const argument& arg)
{
    auto&& output_shape = result.get_shape();
Paul's avatar
Paul committed
153
    auto&& input_shape  = arg.get_shape();
Paul's avatar
Paul committed
154
    std::vector<std::size_t> reduce_lens;
Paul's avatar
Paul committed
155
156
157
158
159
160
161
162
163
164
    std::transform(output_shape.lens().begin(),
                   output_shape.lens().end(),
                   input_shape.lens().begin(),
                   std::back_inserter(reduce_lens),
                   [](auto x, auto y) -> std::size_t {
                       if(x == y)
                           return 1;
                       else
                           return y;
                   });
Paul's avatar
Paul committed
165
    shape reduce_slice{output_shape.type(), reduce_lens};
Paul's avatar
Paul committed
166
167
168
169
    hip_visit_all(result, arg, reduce_slice)([&](auto output, auto input, auto reduce_shape) {
        auto nelements = result.get_shape().elements();
        auto relements = reduce_slice.elements();

Paul's avatar
Paul committed
170
        const std::size_t max_block_size = 1024;
Paul's avatar
Paul committed
171
        const std::size_t block_size     = compute_block_size(relements, max_block_size);
Paul's avatar
Paul committed
172
        gs_launch(stream, nelements * block_size, block_size)([=](auto i, auto idx) __device__ {
Paul's avatar
Paul committed
173
            const auto out_idx = i / block_size;
Paul's avatar
Paul committed
174
            auto base_idx      = output.get_shape().multi(out_idx);
Paul's avatar
Paul committed
175
            auto r = block_reduce<max_block_size>(idx, sum{}, 0, relements, [&](auto j) __device__ {
Paul's avatar
Paul committed
176
177
                auto reduce_idx = reduce_shape.multi(j);
                return input[reduce_idx + base_idx];
Paul's avatar
Paul committed
178
            });
Paul's avatar
Paul committed
179
            if(idx.local == 0)
Paul's avatar
Paul committed
180
                output.data()[out_idx] = r;
Paul's avatar
Paul committed
181
182
183
184
185
186
187
188
        });
    });
}

} // namespace device
} // namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx