softmax.cpp 5.32 KB
Newer Older
Khalique's avatar
Khalique committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
#include <migraphx/shape.hpp>
#include <migraphx/argument.hpp>
#include <migraphx/dfor.hpp>
#include <migraphx/gpu/device/softmax.hpp>
#include <migraphx/gpu/device/tensor.hpp>
#include <migraphx/gpu/device/launch.hpp>
#include <migraphx/gpu/device/types.hpp>
#include <migraphx/gpu/hip.hpp>

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

Shucai Xiao's avatar
Shucai Xiao committed
15
16
template <class T>
__device__ void
Shucai Xiao's avatar
Shucai Xiao committed
17
reduce_max(T* data_ptr, size_t block_size, size_t thr_idx, size_t item_num)
18
19
{
    auto stride = (item_num + 1) / 2;
Shucai Xiao's avatar
Shucai Xiao committed
20
    while(true)
21
22
23
    {
        if(thr_idx + stride < item_num)
        {
Shucai Xiao's avatar
Shucai Xiao committed
24
25
            data_ptr[thr_idx] =
                ::max(to_hip_type(data_ptr[thr_idx]), to_hip_type(data_ptr[thr_idx + stride]));
26
27
        }
        __syncthreads();
Shucai Xiao's avatar
Shucai Xiao committed
28
29
        item_num = stride;
        stride   = (stride + 1) / 2;
30
31
32
33
34
35
36

        if(item_num == 1)
            break;
    }

    if(thr_idx == 0)
    {
Shucai Xiao's avatar
Shucai Xiao committed
37
38
        data_ptr[block_size] =
            (data_ptr[0] < data_ptr[block_size]) ? data_ptr[block_size] : data_ptr[0];
39
40
41
42
43
    }

    __syncthreads();
}

Shucai Xiao's avatar
Shucai Xiao committed
44
45
template <class T>
__device__ void
Shucai Xiao's avatar
Shucai Xiao committed
46
reduce_sum(T* data_ptr, size_t block_size, size_t thr_idx, size_t item_num)
47
48
{
    auto stride = (item_num + 1) / 2;
Shucai Xiao's avatar
Shucai Xiao committed
49
    while(true)
50
51
52
53
54
55
    {
        if(thr_idx + stride < item_num)
        {
            data_ptr[thr_idx] += data_ptr[thr_idx + stride];
        }
        __syncthreads();
Shucai Xiao's avatar
Shucai Xiao committed
56
57
        item_num = stride;
        stride   = (stride + 1) / 2;
58
59
60
61
62
63
64

        if(item_num == 1)
            break;
    }

    if(thr_idx == 0)
    {
Shucai Xiao's avatar
Shucai Xiao committed
65
        data_ptr[block_size] += data_ptr[0];
66
67
68
69
70
    }

    __syncthreads();
}

Shucai Xiao's avatar
Shucai Xiao committed
71
void softmax(hipStream_t stream, const argument& result, const argument& arg, int axis)
Khalique's avatar
Khalique committed
72
{
73
    auto lens        = result.get_shape().lens();
Shucai Xiao's avatar
Shucai Xiao committed
74
    auto batch_lens  = lens;
Shucai Xiao's avatar
Shucai Xiao committed
75
    size_t batch_item_num    = lens[axis];
76
    batch_lens[axis] = 1;
77
    migraphx::shape batch_shape{result.get_shape().type(), batch_lens};
Khalique's avatar
Khalique committed
78

79
    visit_all(result, arg)([&](auto output, auto input) {
Khalique's avatar
Khalique committed
80
81
        const auto* input_ptr = device_cast(input.data());
        auto* output_ptr      = device_cast(output.data());
82
83
        visit_tensor_size(batch_shape.lens().size(), [&](auto n_dim) {
            hip_tensor_descriptor<n_dim> desc_batch(batch_shape);
84
            hip_tensor_descriptor<n_dim> desc_data(result.get_shape());
Khalique's avatar
Khalique committed
85

86
            // use one block for items in one batch.
87
            const size_t max_block_size = 1024;
Shucai Xiao's avatar
Shucai Xiao committed
88
            size_t block_size           = 1;
Shucai Xiao's avatar
Shucai Xiao committed
89
            while(block_size < max_block_size and block_size < batch_item_num)
90
91
92
93
            {
                block_size *= 2;
            }

94
95
96
97
98
99
            launch(
                stream, batch_shape.elements() * block_size, block_size)([=](auto idx) __device__ {
                size_t thr_idx = idx.local;
                size_t blk_idx = idx.group;
                using type = device_type<std::remove_cv_t<typename decltype(output)::value_type>>;

Shucai Xiao's avatar
Shucai Xiao committed
100
                MIGRAPHX_DEVICE_SHARED type lds_data[max_block_size + 1];
101
                auto batch_idx = desc_batch.multi(blk_idx);
Shucai Xiao's avatar
Shucai Xiao committed
102
                auto data_idx  = batch_idx;
103
                // load data to lds and compute the batch max
Shucai Xiao's avatar
Shucai Xiao committed
104
105
                size_t remaining_item_num          = batch_item_num;
                size_t round_item_num        = (batch_item_num + block_size - 1) / block_size * block_size;
Shucai Xiao's avatar
Shucai Xiao committed
106
                lds_data[block_size]     = input_ptr[0];
Shucai Xiao's avatar
Shucai Xiao committed
107
                for(size_t i = thr_idx; i < round_item_num; i += block_size)
108
                {
Shucai Xiao's avatar
Shucai Xiao committed
109
                    if(i < batch_item_num)
110
                    {
Shucai Xiao's avatar
Shucai Xiao committed
111
112
                        data_idx[axis]    = i;
                        lds_data[thr_idx] = input_ptr[desc_data.linear(data_idx)];
113
                    }
Khalique's avatar
Khalique committed
114

115
                    __syncthreads();
Khalique's avatar
Khalique committed
116

Shucai Xiao's avatar
Shucai Xiao committed
117
                    auto size = (remaining_item_num > block_size) ? block_size : remaining_item_num;
118
                    reduce_max<type>(lds_data, block_size, thr_idx, size);
119

Shucai Xiao's avatar
Shucai Xiao committed
120
                    remaining_item_num -= block_size;
121
                }
Khalique's avatar
Khalique committed
122

Shucai Xiao's avatar
Shucai Xiao committed
123
124
125
126
127
128
                auto batch_max = lds_data[block_size];
                __syncthreads();

                lds_data[block_size] = 0;
                remaining_item_num = batch_item_num;
                for(size_t i = thr_idx; i < round_item_num; i += block_size)
129
                {
Shucai Xiao's avatar
Shucai Xiao committed
130
                    if(i < batch_item_num)
131
132
                    {
                        data_idx[axis] = i;
Shucai Xiao's avatar
Shucai Xiao committed
133
                        lds_data[thr_idx] =
Shucai Xiao's avatar
Shucai Xiao committed
134
                            input_ptr[desc_data.linear(data_idx)] - batch_max;
Shucai Xiao's avatar
Shucai Xiao committed
135
                        lds_data[thr_idx] = ::exp(to_hip_type(lds_data[thr_idx]));
136
                    }
137
138
139

                    __syncthreads();

Shucai Xiao's avatar
Shucai Xiao committed
140
                    auto size = (remaining_item_num > block_size) ? block_size : remaining_item_num;
141
                    reduce_sum<type>(lds_data, block_size, thr_idx, size);
142

Shucai Xiao's avatar
Shucai Xiao committed
143
                    remaining_item_num -= block_size;
144
                }
Shucai Xiao's avatar
Shucai Xiao committed
145
                auto batch_sum = lds_data[block_size];
Khalique's avatar
Khalique committed
146

Shucai Xiao's avatar
Shucai Xiao committed
147
                for(size_t i = thr_idx; i < batch_item_num; i += block_size)
148
                {
149
150
                    data_idx[axis]    = i;
                    size_t index      = desc_data.linear(data_idx);
Shucai Xiao's avatar
Shucai Xiao committed
151
152
                    auto val          = input_ptr[index] - batch_max;
                    output_ptr[index] = ::exp(to_hip_type(val)) / batch_sum;
153
154
                }
            });
Khalique's avatar
Khalique committed
155
156
157
158
159
160
161
162
        });
    });
}

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