dynamic_buffer.hpp 17.6 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
// SPDX-License-Identifier: MIT
Illia Silin's avatar
Illia Silin committed
2
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
Chao Liu's avatar
Chao Liu committed
3

4
#pragma once
Chao Liu's avatar
Chao Liu committed
5
6

#include "ck/ck.hpp"
7
#include "ck/utility/data_type.hpp"
8
#include "enable_if.hpp"
9
10
#include "c_style_pointer_cast.hpp"
#include "amd_buffer_addressing.hpp"
rocking5566's avatar
rocking5566 committed
11
#include "generic_memory_space_atomic.hpp"
12
13
14

namespace ck {

15
16
17
18
// T may be scalar or vector
// X may be scalar or vector
// T and X have same scalar type
// X contains multiple T
19
template <AddressSpaceEnum BufferAddressSpace,
20
21
          typename T,
          typename ElementSpaceSize,
22
23
          bool InvalidElementUseNumericalZeroValue,
          AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence>
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
struct DynamicBuffer
{
    using type = T;

    T* p_data_;
    ElementSpaceSize element_space_size_;
    T invalid_element_value_ = T{0};

    __host__ __device__ constexpr DynamicBuffer(T* p_data, ElementSpaceSize element_space_size)
        : p_data_{p_data}, element_space_size_{element_space_size}
    {
    }

    __host__ __device__ constexpr DynamicBuffer(T* p_data,
                                                ElementSpaceSize element_space_size,
                                                T invalid_element_value)
        : p_data_{p_data},
          element_space_size_{element_space_size},
          invalid_element_value_{invalid_element_value}
    {
    }

46
    __host__ __device__ static constexpr AddressSpaceEnum GetAddressSpace()
47
48
49
50
    {
        return BufferAddressSpace;
    }

51
52
53
54
    __host__ __device__ constexpr const T& operator[](index_t i) const { return p_data_[i]; }

    __host__ __device__ constexpr T& operator()(index_t i) { return p_data_[i]; }

55
    template <typename X,
Chao Liu's avatar
Chao Liu committed
56
57
58
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                         typename scalar_type<remove_cvref_t<T>>::type>::value,
                                 bool>::type = false>
59
60
61
    __host__ __device__ constexpr auto Get(index_t i, bool is_valid_element) const
    {
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
62
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
63

Chao Liu's avatar
Chao Liu committed
64
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
65
66

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
67
                      "wrong! X should contain multiple T");
68

Jianfeng Yan's avatar
Jianfeng Yan committed
69
#if CK_USE_AMD_BUFFER_LOAD
70
71
72
73
74
        bool constexpr use_amd_buffer_addressing = true;
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif

75
        if constexpr(GetAddressSpace() == AddressSpaceEnum::Global && use_amd_buffer_addressing)
76
77
78
79
80
        {
            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;

            if constexpr(InvalidElementUseNumericalZeroValue)
            {
81
82
83
                return amd_buffer_load_invalid_element_return_zero<remove_cvref_t<T>,
                                                                   t_per_x,
                                                                   coherence>(
Chao Liu's avatar
Chao Liu committed
84
                    p_data_, i, is_valid_element, element_space_size_);
85
86
87
            }
            else
            {
Chao Liu's avatar
Chao Liu committed
88
                return amd_buffer_load_invalid_element_return_customized_value<remove_cvref_t<T>,
89
90
                                                                               t_per_x,
                                                                               coherence>(
91
92
93
94
95
                    p_data_, i, is_valid_element, element_space_size_, invalid_element_value_);
            }
        }
        else
        {
96
            if(is_valid_element)
97
            {
98
99
100
101
102
#if CK_EXPERIMENTAL_USE_MEMCPY_FOR_VECTOR_ACCESS
                X tmp;

                __builtin_memcpy(&tmp, &(p_data_[i]), sizeof(X));

103
                return tmp;
104
#else
105
                return *c_style_pointer_cast<const X*>(&p_data_[i]);
106
#endif
107
108
109
            }
            else
            {
110
111
112
113
114
115
116
117
                if constexpr(InvalidElementUseNumericalZeroValue)
                {
                    return X{0};
                }
                else
                {
                    return X{invalid_element_value_};
                }
118
119
120
121
            }
        }
    }

122
    template <InMemoryDataOperationEnum Op,
123
124
125
126
127
128
              typename X,
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                         typename scalar_type<remove_cvref_t<T>>::type>::value,
                                 bool>::type = false>
    __host__ __device__ void Update(index_t i, bool is_valid_element, const X& x)
    {
129
        if constexpr(Op == InMemoryDataOperationEnum::Set)
130
131
132
        {
            this->template Set<X>(i, is_valid_element, x);
        }
133
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicAdd)
134
135
136
        {
            this->template AtomicAdd<X>(i, is_valid_element, x);
        }
rocking5566's avatar
rocking5566 committed
137
138
139
140
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicMax)
        {
            this->template AtomicMax<X>(i, is_valid_element, x);
        }
141
        else if constexpr(Op == InMemoryDataOperationEnum::Add)
142
143
144
145
146
147
148
149
        {
            auto tmp = this->template Get<X>(i, is_valid_element);
            this->template Set<X>(i, is_valid_element, x + tmp);
            // tmp += x;
            // this->template Set<X>(i, is_valid_element, tmp);
        }
    }

Anthony Chang's avatar
Anthony Chang committed
150
151
152
153
154
155
156
157
158
159
    __host__ __device__ void Clear()
    {
        static_assert(GetAddressSpace() == AddressSpaceEnum::Lds,
                      "wrong! only local data share is supported");
        for(index_t i = get_thread_local_1d_id(); i < element_space_size_; i += get_block_size())
        {
            Set(i, true, T{0});
        }
    }

160
    template <typename X,
Chao Liu's avatar
Chao Liu committed
161
162
163
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                         typename scalar_type<remove_cvref_t<T>>::type>::value,
                                 bool>::type = false>
164
165
166
    __host__ __device__ void Set(index_t i, bool is_valid_element, const X& x)
    {
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
167
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
168

Chao Liu's avatar
Chao Liu committed
169
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
170
171

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
172
                      "wrong! X should contain multiple T");
173

Jianfeng Yan's avatar
Jianfeng Yan committed
174
#if CK_USE_AMD_BUFFER_STORE
175
        bool constexpr use_amd_buffer_addressing = true;
176
#else
177
178
        bool constexpr use_amd_buffer_addressing      = false;
#endif
179

180
181
#if CK_WORKAROUND_SWDEV_XXXXXX_INT8_DS_WRITE_ISSUE
        bool constexpr workaround_int8_ds_write_issue = true;
182
#else
183
        bool constexpr workaround_int8_ds_write_issue = false;
184
#endif
185
186
187
188
189

        if constexpr(GetAddressSpace() == AddressSpaceEnum::Global && use_amd_buffer_addressing)
        {
            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;

190
            amd_buffer_store<remove_cvref_t<T>, t_per_x, coherence>(
191
                x, p_data_, i, is_valid_element, element_space_size_);
192
        }
193
194
195
        else if constexpr(GetAddressSpace() == AddressSpaceEnum::Lds &&
                          is_same<typename scalar_type<remove_cvref_t<T>>::type, int8_t>::value &&
                          workaround_int8_ds_write_issue)
196
197
198
        {
            if(is_valid_element)
            {
199
                // HACK: compiler would lower IR "store<i8, 16> address_space(3)" into inefficient
200
201
202
                // ISA, so I try to let compiler emit IR "store<i32, 4>" which would be lower to
                // ds_write_b128
                // TODO: remove this after compiler fix
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
                static_assert((is_same<remove_cvref_t<T>, int8_t>::value &&
                               is_same<remove_cvref_t<X>, int8_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x2_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x4_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x8_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x16_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8x4_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x4_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8x8_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x8_t>::value) ||
                                  (is_same<remove_cvref_t<T>, int8x16_t>::value &&
                                   is_same<remove_cvref_t<X>, int8x16_t>::value),
                              "wrong! not implemented for this combination, please add "
                              "implementation");

                if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                             is_same<remove_cvref_t<X>, int8_t>::value)
224
                {
225
226
227
228
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int8_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int8_t*>(&x);
229
                }
230
231
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x2_t>::value)
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
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int16_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int16_t*>(&x);
                }
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x4_t>::value)
                {
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int32_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int32_t*>(&x);
                }
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x8_t>::value)
                {
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int32x2_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int32x2_t*>(&x);
                }
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x16_t>::value)
                {
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int32x4_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int32x4_t*>(&x);
                }
                else if constexpr(is_same<remove_cvref_t<T>, int8x4_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x4_t>::value)
                {
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int32_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int32_t*>(&x);
                }
                else if constexpr(is_same<remove_cvref_t<T>, int8x8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x8_t>::value)
                {
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int32x2_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int32x2_t*>(&x);
                }
                else if constexpr(is_same<remove_cvref_t<T>, int8x16_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x16_t>::value)
                {
                    // HACK: cast pointer of x is bad
                    // TODO: remove this after compiler fix
                    *c_style_pointer_cast<int32x4_t*>(&p_data_[i]) =
                        *c_style_pointer_cast<const int32x4_t*>(&x);
285
286
287
288
289
290
291
                }
            }
        }
        else
        {
            if(is_valid_element)
            {
292
293
294
295
296
#if CK_EXPERIMENTAL_USE_MEMCPY_FOR_VECTOR_ACCESS
                X tmp = x;

                __builtin_memcpy(&(p_data_[i]), &tmp, sizeof(X));
#else
297
                *c_style_pointer_cast<X*>(&p_data_[i]) = x;
298
#endif
299
300
301
302
            }
        }
    }

zjing14's avatar
zjing14 committed
303
    template <typename X,
Chao Liu's avatar
Chao Liu committed
304
305
306
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                         typename scalar_type<remove_cvref_t<T>>::type>::value,
                                 bool>::type = false>
zjing14's avatar
zjing14 committed
307
308
    __host__ __device__ void AtomicAdd(index_t i, bool is_valid_element, const X& x)
    {
309
310
        using scalar_t = typename scalar_type<remove_cvref_t<T>>::type;

zjing14's avatar
zjing14 committed
311
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
312
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
zjing14's avatar
zjing14 committed
313

Chao Liu's avatar
Chao Liu committed
314
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
zjing14's avatar
zjing14 committed
315
316

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
317
318
                      "wrong! X should contain multiple T");

Anthony Chang's avatar
Anthony Chang committed
319
320
321
        static_assert(GetAddressSpace() == AddressSpaceEnum::Global ||
                          GetAddressSpace() == AddressSpaceEnum::Lds,
                      "only support global mem or local data share");
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336

#if CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER && CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT
        bool constexpr use_amd_buffer_addressing =
            is_same_v<remove_cvref_t<scalar_t>, int32_t> ||
            is_same_v<remove_cvref_t<scalar_t>, float> ||
            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0);
#elif CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER && (!CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT)
        bool constexpr use_amd_buffer_addressing = is_same_v<remove_cvref_t<scalar_t>, int32_t>;
#elif(!CK_USE_AMD_BUFFER_ATOMIC_ADD_INTEGER) && CK_USE_AMD_BUFFER_ATOMIC_ADD_FLOAT
        bool constexpr use_amd_buffer_addressing =
            is_same_v<remove_cvref_t<scalar_t>, float> ||
            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0);
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif
zjing14's avatar
zjing14 committed
337

Anthony Chang's avatar
Anthony Chang committed
338
        if constexpr(use_amd_buffer_addressing && GetAddressSpace() == AddressSpaceEnum::Global)
339
340
        {
            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;
zjing14's avatar
zjing14 committed
341

342
343
344
345
            amd_buffer_atomic_add<remove_cvref_t<T>, t_per_x>(
                x, p_data_, i, is_valid_element, element_space_size_);
        }
        else
zjing14's avatar
zjing14 committed
346
        {
347
348
            if(is_valid_element)
            {
Chao Liu's avatar
Chao Liu committed
349
                atomic_add<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
350
            }
zjing14's avatar
zjing14 committed
351
352
353
        }
    }

rocking5566's avatar
rocking5566 committed
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
    template <typename X,
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
                                         typename scalar_type<remove_cvref_t<T>>::type>::value,
                                 bool>::type = false>
    __host__ __device__ void AtomicMax(index_t i, bool is_valid_element, const X& x)
    {
        // X contains multiple T
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;

        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
                      "wrong! X should contain multiple T");

        static_assert(GetAddressSpace() == AddressSpaceEnum::Global, "only support global mem");

#if CK_USE_AMD_BUFFER_ATOMIC_MAX_FLOAT64
        using scalar_t                           = typename scalar_type<remove_cvref_t<T>>::type;
        bool constexpr use_amd_buffer_addressing = is_same_v<remove_cvref_t<scalar_t>, double>;
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif

        if constexpr(use_amd_buffer_addressing)
        {
            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;

            amd_buffer_atomic_max<remove_cvref_t<T>, t_per_x>(
                x, p_data_, i, is_valid_element, element_space_size_);
        }
        else if(is_valid_element)
        {
            atomic_max<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
        }
    }

390
391
392
393
394
    __host__ __device__ static constexpr bool IsStaticBuffer() { return false; }

    __host__ __device__ static constexpr bool IsDynamicBuffer() { return true; }
};

395
396
397
398
template <AddressSpaceEnum BufferAddressSpace,
          AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence,
          typename T,
          typename ElementSpaceSize>
399
400
__host__ __device__ constexpr auto make_dynamic_buffer(T* p, ElementSpaceSize element_space_size)
{
401
402
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize, true, coherence>{
        p, element_space_size};
403
404
}

405
template <
406
    AddressSpaceEnum BufferAddressSpace,
407
    AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence,
408
409
410
411
    typename T,
    typename ElementSpaceSize,
    typename X,
    typename enable_if<is_same<remove_cvref_t<T>, remove_cvref_t<X>>::value, bool>::type = false>
412
__host__ __device__ constexpr auto
413
make_dynamic_buffer(T* p, ElementSpaceSize element_space_size, X invalid_element_value)
414
{
415
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize, false, coherence>{
416
417
418
419
        p, element_space_size, invalid_element_value};
}

} // namespace ck