dynamic_buffer.hpp 19.8 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
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
57
58
                                         typename scalar_type<remove_cvref_t<T>>::type>::value ||
                                     !is_native_type<X>(),
Chao Liu's avatar
Chao Liu committed
59
                                 bool>::type = false>
60
61
62
    __host__ __device__ constexpr auto Get(index_t i, bool is_valid_element) const
    {
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
63
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
64

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

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

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

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

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

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

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

123
    template <InMemoryDataOperationEnum Op,
124
125
126
127
128
129
              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)
    {
130
        if constexpr(Op == InMemoryDataOperationEnum::Set)
131
132
133
        {
            this->template Set<X>(i, is_valid_element, x);
        }
134
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicAdd)
135
136
137
        {
            this->template AtomicAdd<X>(i, is_valid_element, x);
        }
rocking5566's avatar
rocking5566 committed
138
139
140
141
        else if constexpr(Op == InMemoryDataOperationEnum::AtomicMax)
        {
            this->template AtomicMax<X>(i, is_valid_element, x);
        }
142
        else if constexpr(Op == InMemoryDataOperationEnum::Add)
143
        {
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
            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);
            }
174
175
176
        }
    }

177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
    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,
                                                            element_space_size_);
    }

197
    template <typename X,
Chao Liu's avatar
Chao Liu committed
198
              typename enable_if<is_same<typename scalar_type<remove_cvref_t<X>>::type,
199
200
                                         typename scalar_type<remove_cvref_t<T>>::type>::value ||
                                     !is_native_type<X>(),
Chao Liu's avatar
Chao Liu committed
201
                                 bool>::type = false>
202
203
204
    __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
205
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
206

Chao Liu's avatar
Chao Liu committed
207
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
208
209

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

Jianfeng Yan's avatar
Jianfeng Yan committed
212
#if CK_USE_AMD_BUFFER_STORE
213
        bool constexpr use_amd_buffer_addressing = true;
214
#else
215
216
        bool constexpr use_amd_buffer_addressing      = false;
#endif
217

218
219
#if CK_WORKAROUND_SWDEV_XXXXXX_INT8_DS_WRITE_ISSUE
        bool constexpr workaround_int8_ds_write_issue = true;
220
#else
221
        bool constexpr workaround_int8_ds_write_issue = false;
222
#endif
223
224
225
226
227

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

228
            amd_buffer_store<remove_cvref_t<T>, t_per_x, coherence>(
229
                x, p_data_, i, is_valid_element, element_space_size_);
230
        }
231
232
233
        else if constexpr(GetAddressSpace() == AddressSpaceEnum::Lds &&
                          is_same<typename scalar_type<remove_cvref_t<T>>::type, int8_t>::value &&
                          workaround_int8_ds_write_issue)
234
235
236
        {
            if(is_valid_element)
            {
237
                // HACK: compiler would lower IR "store<i8, 16> address_space(3)" into inefficient
238
239
240
                // 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
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
                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)
262
                {
263
264
265
266
                    // 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);
267
                }
268
269
                else if constexpr(is_same<remove_cvref_t<T>, int8_t>::value &&
                                  is_same<remove_cvref_t<X>, int8x2_t>::value)
270
                {
271
272
273
274
275
276
277
278
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
                    // 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);
323
324
325
326
327
328
329
                }
            }
        }
        else
        {
            if(is_valid_element)
            {
330
331
332
333
334
#if CK_EXPERIMENTAL_USE_MEMCPY_FOR_VECTOR_ACCESS
                X tmp = x;

                __builtin_memcpy(&(p_data_[i]), &tmp, sizeof(X));
#else
335
                *c_style_pointer_cast<X*>(&p_data_[i]) = x;
336
#endif
337
338
339
340
            }
        }
    }

zjing14's avatar
zjing14 committed
341
    template <typename X,
Chao Liu's avatar
Chao Liu committed
342
343
344
              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
345
346
    __host__ __device__ void AtomicAdd(index_t i, bool is_valid_element, const X& x)
    {
347
348
        using scalar_t = typename scalar_type<remove_cvref_t<T>>::type;

zjing14's avatar
zjing14 committed
349
        // X contains multiple T
Chao Liu's avatar
Chao Liu committed
350
        constexpr index_t scalar_per_t_vector = scalar_type<remove_cvref_t<T>>::vector_size;
zjing14's avatar
zjing14 committed
351

Chao Liu's avatar
Chao Liu committed
352
        constexpr index_t scalar_per_x_vector = scalar_type<remove_cvref_t<X>>::vector_size;
zjing14's avatar
zjing14 committed
353
354

        static_assert(scalar_per_x_vector % scalar_per_t_vector == 0,
355
356
357
358
359
360
361
362
                      "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> ||
363
364
            (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);
365
366
367
368
369
#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> ||
370
371
            (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);
372
373
374
#else
        bool constexpr use_amd_buffer_addressing = false;
#endif
zjing14's avatar
zjing14 committed
375

376
377
378
        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
379

380
381
382
383
            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
384
        {
385
386
            if(is_valid_element)
            {
Chao Liu's avatar
Chao Liu committed
387
                atomic_add<X>(c_style_pointer_cast<X*>(&p_data_[i]), x);
388
            }
zjing14's avatar
zjing14 committed
389
390
391
        }
    }

rocking5566's avatar
rocking5566 committed
392
393
394
395
396
397
398
399
400
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
    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);
        }
    }

428
429
430
431
432
    __host__ __device__ static constexpr bool IsStaticBuffer() { return false; }

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

433
434
435
436
template <AddressSpaceEnum BufferAddressSpace,
          AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence,
          typename T,
          typename ElementSpaceSize>
437
438
__host__ __device__ constexpr auto make_dynamic_buffer(T* p, ElementSpaceSize element_space_size)
{
439
440
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize, true, coherence>{
        p, element_space_size};
441
442
}

443
template <
444
    AddressSpaceEnum BufferAddressSpace,
445
    AmdBufferCoherenceEnum coherence = AmdBufferCoherenceEnum::DefaultCoherence,
446
447
448
449
    typename T,
    typename ElementSpaceSize,
    typename X,
    typename enable_if<is_same<remove_cvref_t<T>, remove_cvref_t<X>>::value, bool>::type = false>
450
__host__ __device__ constexpr auto
451
make_dynamic_buffer(T* p, ElementSpaceSize element_space_size, X invalid_element_value)
452
{
453
    return DynamicBuffer<BufferAddressSpace, T, ElementSpaceSize, false, coherence>{
454
455
456
457
        p, element_space_size, invalid_element_value};
}

} // namespace ck