dynamic_buffer.hpp 20.1 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
struct DynamicBuffer
{
    using type = T;

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

32
33
34
35
36
37
38
    static constexpr index_t PackedSize = []() {
        if constexpr(is_same_v<remove_cvref_t<T>, pk_i4_t>)
            return 2;
        else
            return 1;
    }();

39
40
41
42
43
44
45
46
47
48
49
50
51
52
    __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}
    {
    }

53
    __host__ __device__ static constexpr AddressSpaceEnum GetAddressSpace()
54
55
56
57
    {
        return BufferAddressSpace;
    }

58
59
60
61
    __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]; }

62
    template <typename X,
Chao Liu's avatar
Chao Liu committed
63
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
64
65
                                         typename scalar_type<remove_cvref_t<T>>::type>::value ||
                                     !is_native_type<X>(),
Chao Liu's avatar
Chao Liu committed
66
                                 bool>::type = false>
67
68
69
    __host__ __device__ constexpr auto Get(index_t i, bool is_valid_element) const
    {
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
70
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
71

Chao Liu's avatar
Chao Liu committed
72
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
73
74

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

Jianfeng Yan's avatar
Jianfeng Yan committed
77
#if CK_USE_AMD_BUFFER_LOAD
78
79
80
81
82
        bool constexpr use_amd_buffer_addressing = true;
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif

83
        if constexpr(GetAddressSpace() == AddressSpaceEnum::Global && use_amd_buffer_addressing)
84
85
86
87
88
        {
            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;

            if constexpr(InvalidElementUseNumericalZeroValue)
            {
89
90
91
                return amd_buffer_load_invalid_element_return_zero<remove_cvref_t<T>,
                                                                   t_per_x,
                                                                   coherence>(
92
                    p_data_, i, is_valid_element, element_space_size_ / PackedSize);
93
94
95
            }
            else
            {
Chao Liu's avatar
Chao Liu committed
96
                return amd_buffer_load_invalid_element_return_customized_value<remove_cvref_t<T>,
97
98
                                                                               t_per_x,
                                                                               coherence>(
99
100
101
102
103
                    p_data_,
                    i,
                    is_valid_element,
                    element_space_size_ / PackedSize,
                    invalid_element_value_);
104
105
106
107
            }
        }
        else
        {
108
            if(is_valid_element)
109
            {
110
111
112
113
114
#if CK_EXPERIMENTAL_USE_MEMCPY_FOR_VECTOR_ACCESS
                X tmp;

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

115
                return tmp;
116
#else
117
                return *c_style_pointer_cast<const X*>(&p_data_[i]);
118
#endif
119
120
121
            }
            else
            {
122
123
124
125
126
127
128
129
                if constexpr(InvalidElementUseNumericalZeroValue)
                {
                    return X{0};
                }
                else
                {
                    return X{invalid_element_value_};
                }
130
131
132
133
            }
        }
    }

134
    template <InMemoryDataOperationEnum Op,
135
136
137
138
139
140
              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)
    {
141
        if constexpr(Op == InMemoryDataOperationEnum::Set)
142
143
144
        {
            this->template Set<X>(i, is_valid_element, x);
        }
145
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicAdd)
146
147
148
        {
            this->template AtomicAdd<X>(i, is_valid_element, x);
        }
rocking5566's avatar
rocking5566 committed
149
150
151
152
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicMax)
        {
            this->template AtomicMax<X>(i, is_valid_element, x);
        }
153
        else if constexpr(Op == InMemoryDataOperationEnum::Add)
154
        {
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
            auto tmp       = this->template Get<X>(i, is_valid_element);
            using scalar_t = typename scalar_type<remove_cvref_t<T>>::type;
            // handle bfloat addition
            if constexpr(is_same_v<scalar_t, bhalf_t>)
            {
                if constexpr(is_scalar_type<X>::value)
                {
                    // Scalar type
                    auto result =
                        type_convert<X>(type_convert<float>(x) + type_convert<float>(tmp));
                    this->template Set<X>(i, is_valid_element, result);
                }
                else
                {
                    // Vector type
                    constexpr auto vector_size = scalar_type<remove_cvref_t<X>>::vector_size;
                    const vector_type<scalar_t, vector_size> a_vector{tmp};
                    const vector_type<scalar_t, vector_size> b_vector{x};
                    static_for<0, vector_size, 1>{}([&](auto idx) {
                        auto result = type_convert<scalar_t>(
                            type_convert<float>(a_vector.template AsType<scalar_t>()[idx]) +
                            type_convert<float>(b_vector.template AsType<scalar_t>()[idx]));
                        this->template Set<scalar_t>(i + idx, is_valid_element, result);
                    });
                }
            }
            else
            {
                this->template Set<X>(i, is_valid_element, x + tmp);
            }
185
186
187
        }
    }

188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
    template <typename DstBuffer, index_t NumElemsPerThread>
    __host__ __device__ void DirectCopyToLds(DstBuffer& dst_buf,
                                             index_t src_offset,
                                             index_t dst_offset,
                                             bool is_valid_element) const
    {
        // Copy data from global to LDS memory using direct loads.
        static_assert(GetAddressSpace() == AddressSpaceEnum::Global,
                      "Source data must come from a global memory buffer.");
        static_assert(DstBuffer::GetAddressSpace() == AddressSpaceEnum::Lds,
                      "Destination data must be stored in an LDS memory buffer.");

        amd_direct_load_global_to_lds<T, NumElemsPerThread>(p_data_,
                                                            src_offset,
                                                            dst_buf.p_data_,
                                                            dst_offset,
                                                            is_valid_element,
205
                                                            element_space_size_ / PackedSize);
206
207
    }

208
    template <typename X,
Chao Liu's avatar
Chao Liu committed
209
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
210
211
                                         typename scalar_type<remove_cvref_t<T>>::type>::value ||
                                     !is_native_type<X>(),
Chao Liu's avatar
Chao Liu committed
212
                                 bool>::type = false>
213
214
215
    __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
216
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
217

Chao Liu's avatar
Chao Liu committed
218
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
219
220

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

Jianfeng Yan's avatar
Jianfeng Yan committed
223
#if CK_USE_AMD_BUFFER_STORE
224
        bool constexpr use_amd_buffer_addressing = true;
225
#else
226
227
        bool constexpr use_amd_buffer_addressing      = false;
#endif
228

229
230
#if CK_WORKAROUND_SWDEV_XXXXXX_INT8_DS_WRITE_ISSUE
        bool constexpr workaround_int8_ds_write_issue = true;
231
#else
232
        bool constexpr workaround_int8_ds_write_issue = false;
233
#endif
234
235
236
237
238

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

239
            amd_buffer_store<remove_cvref_t<T>, t_per_x, coherence>(
240
                x, p_data_, i, is_valid_element, element_space_size_ / PackedSize);
241
        }
242
243
244
        else if constexpr(GetAddressSpace() == AddressSpaceEnum::Lds &&
                          is_same<typename scalar_type<remove_cvref_t<T>>::type, int8_t>::value &&
                          workaround_int8_ds_write_issue)
245
246
247
        {
            if(is_valid_element)
            {
248
                // HACK: compiler would lower IR "store<i8, 16> address_space(3)" into inefficient
249
250
251
                // 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
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
                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)
273
                {
274
275
276
277
                    // 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);
278
                }
279
280
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x2_t>::value)
281
                {
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
                    // 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);
334
335
336
337
338
339
340
                }
            }
        }
        else
        {
            if(is_valid_element)
            {
341
342
343
344
345
#if CK_EXPERIMENTAL_USE_MEMCPY_FOR_VECTOR_ACCESS
                X tmp = x;

                __builtin_memcpy(&(p_data_[i]), &tmp, sizeof(X));
#else
346
                *c_style_pointer_cast<X*>(&p_data_[i]) = x;
347
#endif
348
349
350
351
            }
        }
    }

zjing14's avatar
zjing14 committed
352
    template <typename X,
Chao Liu's avatar
Chao Liu committed
353
354
355
              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
356
357
    __host__ __device__ void AtomicAdd(index_t i, bool is_valid_element, const X& x)
    {
358
359
        using scalar_t = typename scalar_type<remove_cvref_t<T>>::type;

zjing14's avatar
zjing14 committed
360
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
361
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
zjing14's avatar
zjing14 committed
362

Chao Liu's avatar
Chao Liu committed
363
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
zjing14's avatar
zjing14 committed
364
365

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
366
367
368
369
370
371
372
373
                      "wrong! X should contain multiple T");

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

#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> ||
374
375
            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0) ||
            (is_same_v<remove_cvref_t<scalar_t>, bhalf_t> && scalar_per_x_vector % 2 == 0);
376
377
378
379
380
#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> ||
381
382
            (is_same_v<remove_cvref_t<scalar_t>, half_t> && scalar_per_x_vector % 2 == 0) ||
            (is_same_v<remove_cvref_t<scalar_t>, bhalf_t> && scalar_per_x_vector % 2 == 0);
383
384
385
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif
zjing14's avatar
zjing14 committed
386

387
388
389
        if constexpr(use_amd_buffer_addressing)
        {
            constexpr index_t t_per_x = scalar_per_x_vector / scalar_per_t_vector;
zjing14's avatar
zjing14 committed
390

391
            amd_buffer_atomic_add<remove_cvref_t<T>, t_per_x>(
392
                x, p_data_, i, is_valid_element, element_space_size_ / PackedSize);
393
394
        }
        else
zjing14's avatar
zjing14 committed
395
        {
396
397
            if(is_valid_element)
            {
Chao Liu's avatar
Chao Liu committed
398
                atomic_add<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
399
            }
zjing14's avatar
zjing14 committed
400
401
402
        }
    }

rocking5566's avatar
rocking5566 committed
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
    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>(
431
                x, p_data_, i, is_valid_element, element_space_size_ / PackedSize);
rocking5566's avatar
rocking5566 committed
432
433
434
435
436
437
438
        }
        else if(is_valid_element)
        {
            atomic_max<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
        }
    }

439
440
441
442
443
    __host__ __device__ static constexpr bool IsStaticBuffer() { return false; }

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

444
445
446
447
template <AddressSpaceEnum BufferAddressSpace,
          AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence,
          typename T,
          typename ElementSpaceSize>
448
449
__host__ __device__ constexpr auto make_dynamic_buffer(T* p, ElementSpaceSize element_space_size)
{
450
451
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize, true, coherence>{
        p, element_space_size};
452
453
}

454
template <
455
    AddressSpaceEnum BufferAddressSpace,
456
    AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence,
457
458
459
460
    typename T,
    typename ElementSpaceSize,
    typename X,
    typename enable_if<is_same<remove_cvref_t<T>, remove_cvref_t<X>>::value, bool>::type = false>
461
__host__ __device__ constexpr auto
462
make_dynamic_buffer(T* p, ElementSpaceSize element_space_size, X invalid_element_value)
463
{
464
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize, false, coherence>{
465
466
467
468
        p, element_space_size, invalid_element_value};
}

} // namespace ck