dynamic_buffer.hpp 20 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};

Jing Zhang's avatar
Jing Zhang committed
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
64
65
              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>
66
67
68
    __host__ __device__ constexpr auto Get(index_t i, bool is_valid_element) const
    {
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
69
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
70

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

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

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

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

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

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

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

133
    template <InMemoryDataOperationEnum Op,
134
135
136
137
138
139
              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)
    {
140
        if constexpr(Op == InMemoryDataOperationEnum::Set)
141
142
143
        {
            this->template Set<X>(i, is_valid_element, x);
        }
144
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicAdd)
145
146
147
        {
            this->template AtomicAdd<X>(i, is_valid_element, x);
        }
rocking5566's avatar
rocking5566 committed
148
149
150
151
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicMax)
        {
            this->template AtomicMax<X>(i, is_valid_element, x);
        }
152
        else if constexpr(Op == InMemoryDataOperationEnum::Add)
153
        {
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
            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);
            }
184
185
186
        }
    }

187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    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,
Jing Zhang's avatar
Jing Zhang committed
204
                                                            element_space_size_ / PackedSize);
205
206
    }

207
    template <typename X,
Chao Liu's avatar
Chao Liu committed
208
209
210
              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>
211
212
213
    __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
214
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
215

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

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

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

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

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

237
            amd_buffer_store<remove_cvref_t<T>, t_per_x, coherence>(
Jing Zhang's avatar
Jing Zhang committed
238
                x, p_data_, i, is_valid_element, element_space_size_ / PackedSize);
239
        }
240
241
242
        else if constexpr(GetAddressSpace() == AddressSpaceEnum::Lds &&
                          is_same<typename scalar_type<remove_cvref_t<T>>::type, int8_t>::value &&
                          workaround_int8_ds_write_issue)
243
244
245
        {
            if(is_valid_element)
            {
246
                // HACK: compiler would lower IR "store<i8, 16> address_space(3)" into inefficient
247
248
249
                // 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
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
                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)
271
                {
272
273
274
275
                    // 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);
276
                }
277
278
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x2_t>::value)
279
                {
280
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
                    // 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);
332
333
334
335
336
337
338
                }
            }
        }
        else
        {
            if(is_valid_element)
            {
339
340
341
342
343
#if CK_EXPERIMENTAL_USE_MEMCPY_FOR_VECTOR_ACCESS
                X tmp = x;

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

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

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

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

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
364
365
366
367
368
369
370
371
                      "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> ||
372
373
            (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);
374
375
376
377
378
#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> ||
379
380
            (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);
381
382
383
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif
zjing14's avatar
zjing14 committed
384

385
386
387
        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
388

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

rocking5566's avatar
rocking5566 committed
401
402
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
    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>(
Jing Zhang's avatar
Jing Zhang committed
429
                x, p_data_, i, is_valid_element, element_space_size_ / PackedSize);
rocking5566's avatar
rocking5566 committed
430
431
432
433
434
435
436
        }
        else if(is_valid_element)
        {
            atomic_max<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
        }
    }

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

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

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

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

} // namespace ck