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

4
#pragma once
5
#include "data_type.hpp"
Chao Liu's avatar
Chao Liu committed
6
7
8
9

namespace ck {

template <typename T>
Chao Liu's avatar
tidy  
Chao Liu committed
10
union BufferResource
Chao Liu's avatar
Chao Liu committed
11
{
12
13
    __device__ constexpr BufferResource() : content{} {}

Chao Liu's avatar
Chao Liu committed
14
15
    // 128 bit SGPRs to supply buffer resource in buffer instructions
    // https://rocm-documentation.readthedocs.io/en/latest/GCN_ISA_Manuals/testdocbook.html#vector-memory-buffer-instructions
16
    int32x4_t content;
17
18
19
    StaticallyIndexedArray<T*, 2> address;
    StaticallyIndexedArray<int32_t, 4> range;
    StaticallyIndexedArray<int32_t, 4> config;
Chao Liu's avatar
Chao Liu committed
20
21
22
};

template <typename T>
23
__device__ int32x4_t make_wave_buffer_resource(T* p_wave, index_t element_space_size)
Chao Liu's avatar
Chao Liu committed
24
{
Chao Liu's avatar
tidy  
Chao Liu committed
25
    BufferResource<T> wave_buffer_resource;
Chao Liu's avatar
Chao Liu committed
26
27

    // wavewise base address (64 bit)
28
    wave_buffer_resource.address(Number<0>{}) = const_cast<remove_cv_t<T>*>(p_wave);
Chao Liu's avatar
Chao Liu committed
29
    // wavewise range (32 bit)
30
    wave_buffer_resource.range(Number<2>{}) = element_space_size * sizeof(T);
Chao Liu's avatar
Chao Liu committed
31
    // wavewise setting (32 bit)
32
    wave_buffer_resource.config(Number<3>{}) = CK_BUFFER_RESOURCE_3RD_DWORD;
Chao Liu's avatar
Chao Liu committed
33

34
    return wave_buffer_resource.content;
Chao Liu's avatar
Chao Liu committed
35
36
}

carlushuang's avatar
carlushuang committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
template <typename T>
__device__ int32x4_t make_wave_buffer_resource_with_default_range(T* p_wave)
{
    BufferResource<T> wave_buffer_resource;

    // wavewise base address (64 bit)
    wave_buffer_resource.address(Number<0>{}) = const_cast<remove_cv_t<T>*>(p_wave);
    // wavewise range (32 bit)
    wave_buffer_resource.range(Number<2>{}) = 0xffffffff; // max possible range
    // wavewise setting (32 bit)
    wave_buffer_resource.config(Number<3>{}) = CK_BUFFER_RESOURCE_3RD_DWORD;

    return wave_buffer_resource.content;
}

Chao Liu's avatar
Chao Liu committed
52
// buffer load i8
Chao Liu's avatar
Chao Liu committed
53
__device__ int8_t
54
55
56
57
llvm_amdgcn_raw_buffer_load_i8(int32x4_t srsrc,
                               index_t voffset,
                               index_t soffset,
                               index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i8");
58
59

__device__ int8x2_t
60
61
62
63
llvm_amdgcn_raw_buffer_load_i8x2(int32x4_t srsrc,
                                 index_t voffset,
                                 index_t soffset,
                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2i8");
64
65

__device__ int8x4_t
66
67
68
69
llvm_amdgcn_raw_buffer_load_i8x4(int32x4_t srsrc,
                                 index_t voffset,
                                 index_t soffset,
                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4i8");
70

Chao Liu's avatar
Chao Liu committed
71
// buffer load i16
72
__device__ bhalf_t
73
74
75
llvm_amdgcn_raw_buffer_load_i16(int32x4_t srsrc,
                                index_t voffset,
                                index_t soffset,
76
77
                                index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i16");

78
__device__ bhalf2_t
79
80
81
82
83
llvm_amdgcn_raw_buffer_load_i16x2(int32x4_t srsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2i16");

84
__device__ bhalf4_t
85
86
87
88
89
llvm_amdgcn_raw_buffer_load_i16x4(int32x4_t srsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4i16");

Chao Liu's avatar
Chao Liu committed
90
// buffer load i32
Chao Liu's avatar
Chao Liu committed
91
__device__ int32_t
92
93
94
95
llvm_amdgcn_raw_buffer_load_i32(int32x4_t srsrc,
                                index_t voffset,
                                index_t soffset,
                                index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i32");
Chao Liu's avatar
Chao Liu committed
96
97

__device__ int32x2_t
98
99
100
101
llvm_amdgcn_raw_buffer_load_i32x2(int32x4_t srsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2i32");
Chao Liu's avatar
Chao Liu committed
102
103

__device__ int32x4_t
104
105
106
107
llvm_amdgcn_raw_buffer_load_i32x4(int32x4_t srsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
108

Chao Liu's avatar
Chao Liu committed
109
// buffer load fp16
110
__device__ half_t
111
112
113
114
llvm_amdgcn_raw_buffer_load_fp16(int32x4_t srsrc,
                                 index_t voffset,
                                 index_t soffset,
                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.f16");
115
116

__device__ half2_t
117
118
119
120
llvm_amdgcn_raw_buffer_load_fp16x2(int32x4_t srsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2f16");
121
122

__device__ half4_t
123
124
125
126
llvm_amdgcn_raw_buffer_load_fp16x4(int32x4_t srsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4f16");
Chao Liu's avatar
Chao Liu committed
127

Chao Liu's avatar
Chao Liu committed
128
// buffer load fp32
Chao Liu's avatar
Chao Liu committed
129
__device__ float
130
131
132
133
llvm_amdgcn_raw_buffer_load_fp32(int32x4_t srsrc,
                                 index_t voffset,
                                 index_t soffset,
                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.f32");
Chao Liu's avatar
Chao Liu committed
134
135

__device__ float2_t
136
137
138
139
llvm_amdgcn_raw_buffer_load_fp32x2(int32x4_t srsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2f32");
Chao Liu's avatar
Chao Liu committed
140
141

__device__ float4_t
142
143
144
145
llvm_amdgcn_raw_buffer_load_fp32x4(int32x4_t srsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4f32");
Chao Liu's avatar
Chao Liu committed
146

Chao Liu's avatar
Chao Liu committed
147
// buffer store i8
Chao Liu's avatar
Chao Liu committed
148
__device__ void
149
150
151
152
153
154
155
156
llvm_amdgcn_raw_buffer_store_i8(int8_t vdata,
                                int32x4_t rsrc,
                                index_t voffset,
                                index_t soffset,
                                index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i8");

__device__ void
llvm_amdgcn_raw_buffer_store_i8x2(int8x2_t vdata,
Chao Liu's avatar
Chao Liu committed
157
158
159
                                  int32x4_t rsrc,
                                  index_t voffset,
                                  index_t soffset,
160
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i8");
Chao Liu's avatar
Chao Liu committed
161

162
__device__ void
163
164
165
166
167
llvm_amdgcn_raw_buffer_store_i8x4(int8x4_t vdata,
                                  int32x4_t rsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i8");
168

Chao Liu's avatar
Chao Liu committed
169
// buffer store i16
170
__device__ void
171
llvm_amdgcn_raw_buffer_store_i16(bhalf_t vdata,
172
173
174
175
176
                                 int32x4_t rsrc,
                                 index_t voffset,
                                 index_t soffset,
                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i16");

177
__device__ void
178
llvm_amdgcn_raw_buffer_store_i16x2(bhalf2_t vdata,
179
180
181
182
183
184
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i16");

__device__ void
185
llvm_amdgcn_raw_buffer_store_i16x4(bhalf4_t vdata,
186
187
188
189
190
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i16");

Chao Liu's avatar
Chao Liu committed
191
// buffer store i32
192
193
194
195
196
197
__device__ void
llvm_amdgcn_raw_buffer_store_i32(int32_t vdata,
                                 int32x4_t rsrc,
                                 index_t voffset,
                                 index_t soffset,
                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i32");
198

Chao Liu's avatar
Chao Liu committed
199
__device__ void
200
llvm_amdgcn_raw_buffer_store_i32x2(int32x2_t vdata,
Chao Liu's avatar
Chao Liu committed
201
202
203
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
204
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i32");
Chao Liu's avatar
Chao Liu committed
205
206

__device__ void
207
llvm_amdgcn_raw_buffer_store_i32x4(int32x4_t vdata,
Chao Liu's avatar
Chao Liu committed
208
209
210
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
211
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
Chao Liu's avatar
Chao Liu committed
212

Chao Liu's avatar
Chao Liu committed
213
// buffer store fp16
Chao Liu's avatar
Chao Liu committed
214
__device__ void
215
216
217
218
219
llvm_amdgcn_raw_buffer_store_fp16(half_t vdata,
                                  int32x4_t rsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.f16");
Chao Liu's avatar
Chao Liu committed
220
221

__device__ void
222
223
224
225
226
llvm_amdgcn_raw_buffer_store_fp16x2(half2_t vdata,
                                    int32x4_t rsrc,
                                    index_t voffset,
                                    index_t soffset,
                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2f16");
Chao Liu's avatar
Chao Liu committed
227

228
__device__ void
229
llvm_amdgcn_raw_buffer_store_fp16x4(half4_t vdata,
230
231
232
                                    int32x4_t rsrc,
                                    index_t voffset,
                                    index_t soffset,
233
                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4f16");
234

Chao Liu's avatar
Chao Liu committed
235
// buffer store fp32
236
__device__ void
237
238
239
240
241
llvm_amdgcn_raw_buffer_store_fp32(float vdata,
                                  int32x4_t rsrc,
                                  index_t voffset,
                                  index_t soffset,
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.f32");
242
243

__device__ void
244
llvm_amdgcn_raw_buffer_store_fp32x2(float2_t vdata,
Chao Liu's avatar
Chao Liu committed
245
246
247
                                    int32x4_t rsrc,
                                    index_t voffset,
                                    index_t soffset,
248
                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2f32");
Chao Liu's avatar
Chao Liu committed
249
250

__device__ void
251
252
253
254
255
llvm_amdgcn_raw_buffer_store_fp32x4(float4_t vdata,
                                    int32x4_t rsrc,
                                    index_t voffset,
                                    index_t soffset,
                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4f32");
256

Chao Liu's avatar
Chao Liu committed
257
258
259
260
261
262
263
264
265
// buffer atomic-add fp16
__device__ half2_t llvm_amdgcn_raw_buffer_atomic_add_fp16x2(
    half2_t vdata,
    int32x4_t rsrc,
    index_t voffset,
    index_t soffset,
    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fadd.v2f16");

// buffer atomic-add i32
zjing14's avatar
zjing14 committed
266
267
268
269
270
271
272
__device__ int32_t llvm_amdgcn_raw_buffer_atomic_add_i32(
    int32_t vdata,
    int32x4_t rsrc,
    index_t voffset,
    index_t soffset,
    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.add.i32");

Chao Liu's avatar
Chao Liu committed
273
// buffer atomic-add fp32
zjing14's avatar
zjing14 committed
274
275
276
277
278
279
__device__ float llvm_amdgcn_raw_buffer_atomic_add_fp32(
    float vdata,
    int32x4_t rsrc,
    index_t voffset,
    index_t soffset,
    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fadd.f32");
Chao Liu's avatar
Chao Liu committed
280

rocking5566's avatar
rocking5566 committed
281
282
283
284
285
286
287
288
// buffer atomic-add fp32
__device__ double
llvm_amdgcn_raw_buffer_atomic_max_fp64(double vdata,
                                       int32x4_t rsrc, // dst_wave_buffer_resource
                                       int voffset,    // dst_thread_addr_offset
                                       int soffset,    // dst_wave_addr_offset
                                       int glc_slc) __asm("llvm.amdgcn.raw.buffer.atomic.fmax.f64");

289
290
291
292
293
294
295
296
297
298
299
300
// memory coherency bit for buffer store/load instruction
enum struct amd_buffer_coherence_bits
{
    default_coherence = 0, // default value
    glc               = 1,
    slc               = 2,
    glc_slc           = 3,
};

template <typename T,
          index_t N,
          amd_buffer_coherence_bits coherence = amd_buffer_coherence_bits::default_coherence>
301
302
303
__device__ typename vector_type<T, N>::type amd_buffer_load_impl(int32x4_t src_wave_buffer_resource,
                                                                 index_t src_thread_addr_offset,
                                                                 index_t src_wave_addr_offset)
Chao Liu's avatar
Chao Liu committed
304
{
305
    static_assert(
Chao Liu's avatar
Chao Liu committed
306
307
        (is_same<T, double>::value && (N == 1 || N == 2 || N == 4)) ||
            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
308
            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
309
            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
Chao Liu's avatar
Chao Liu committed
310
311
            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
            (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
312
        "wrong! not implemented");
Chao Liu's avatar
Chao Liu committed
313

Chao Liu's avatar
Chao Liu committed
314
315
316
317
318
    if constexpr(is_same<T, double>::value)
    {
        // use fp32 load to mimic fp64 load
        if constexpr(N == 1)
        {
319
320
321
322
323
            const float2_t tmp =
                llvm_amdgcn_raw_buffer_load_fp32x2(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset,
                                                   static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
324

325
            return bit_cast<double>(tmp);
Chao Liu's avatar
Chao Liu committed
326
327
328
        }
        else if constexpr(N == 2)
        {
329
330
331
332
333
            const float4_t tmp =
                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset,
                                                   static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
334

335
            return bit_cast<double2_t>(tmp);
Chao Liu's avatar
Chao Liu committed
336
337
338
        }
        else if constexpr(N == 4)
        {
339
340
341
342
343
            const float4_t f32_0 =
                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset,
                                                   static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
344
345
346
347
348

            const float4_t f32_1 =
                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset + 4 * sizeof(float),
349
                                                   static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
350
351
            vector_type<double, 4> tmp;

352
353
            tmp.AsType<double2_t>()(Number<0>{}) = bit_cast<double2_t>(f32_0);
            tmp.AsType<double2_t>()(Number<1>{}) = bit_cast<double2_t>(f32_1);
Chao Liu's avatar
Chao Liu committed
354
355
356
357
358

            return tmp.AsType<double4_t>()(Number<0>{});
        }
    }
    else if constexpr(is_same<T, float>::value)
Chao Liu's avatar
Chao Liu committed
359
360
361
    {
        if constexpr(N == 1)
        {
362
363
364
365
            return llvm_amdgcn_raw_buffer_load_fp32(src_wave_buffer_resource,
                                                    src_thread_addr_offset,
                                                    src_wave_addr_offset,
                                                    static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
366
367
368
        }
        else if constexpr(N == 2)
        {
369
370
371
372
            return llvm_amdgcn_raw_buffer_load_fp32x2(src_wave_buffer_resource,
                                                      src_thread_addr_offset,
                                                      src_wave_addr_offset,
                                                      static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
373
374
375
        }
        else if constexpr(N == 4)
        {
376
377
378
379
            return llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                      src_thread_addr_offset,
                                                      src_wave_addr_offset,
                                                      static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
380
381
382
383
384
        }
        else if constexpr(N == 8)
        {
            vector_type<float, 8> tmp;

385
386
387
388
389
            tmp.AsType<float4_t>()(Number<0>{}) =
                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset,
                                                   static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
390

391
            tmp.AsType<float4_t>()(Number<1>{}) =
392
393
394
                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset + 4 * sizeof(float),
395
                                                   static_cast<index_t>(coherence));
396

397
            return tmp.AsType<float8_t>()(Number<0>{});
398
399
400
401
402
403
        }
    }
    else if constexpr(is_same<T, half_t>::value)
    {
        if constexpr(N == 1)
        {
404
405
406
407
            return llvm_amdgcn_raw_buffer_load_fp16(src_wave_buffer_resource,
                                                    src_thread_addr_offset,
                                                    src_wave_addr_offset,
                                                    static_cast<index_t>(coherence));
408
409
410
        }
        else if constexpr(N == 2)
        {
411
412
413
414
            return llvm_amdgcn_raw_buffer_load_fp16x2(src_wave_buffer_resource,
                                                      src_thread_addr_offset,
                                                      src_wave_addr_offset,
                                                      static_cast<index_t>(coherence));
415
416
417
        }
        else if constexpr(N == 4)
        {
418
419
420
421
            return llvm_amdgcn_raw_buffer_load_fp16x4(src_wave_buffer_resource,
                                                      src_thread_addr_offset,
                                                      src_wave_addr_offset,
                                                      static_cast<index_t>(coherence));
422
        }
423
        else if constexpr(N == 8)
424
        {
Chao Liu's avatar
Chao Liu committed
425
            // use fp32 load to mimic fp16 load
426
427
428
429
            float4_t tmp = llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                              src_thread_addr_offset,
                                                              src_wave_addr_offset,
                                                              static_cast<index_t>(coherence));
zjing14's avatar
zjing14 committed
430

431
            return bit_cast<half8_t>(tmp);
Chao Liu's avatar
Chao Liu committed
432
433
        }
    }
434
    else if constexpr(is_same<T, bhalf_t>::value)
435
436
437
    {
        if constexpr(N == 1)
        {
438
439
440
441
            return llvm_amdgcn_raw_buffer_load_i16(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset,
                                                   static_cast<index_t>(coherence));
442
443
444
        }
        else if constexpr(N == 2)
        {
445
446
447
448
            return llvm_amdgcn_raw_buffer_load_i16x2(src_wave_buffer_resource,
                                                     src_thread_addr_offset,
                                                     src_wave_addr_offset,
                                                     static_cast<index_t>(coherence));
449
450
451
        }
        else if constexpr(N == 4)
        {
452
453
454
455
            return llvm_amdgcn_raw_buffer_load_i16x4(src_wave_buffer_resource,
                                                     src_thread_addr_offset,
                                                     src_wave_addr_offset,
                                                     static_cast<index_t>(coherence));
456
457
458
        }
        else if constexpr(N == 8)
        {
459
460
461
462
            int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                              src_thread_addr_offset,
                                                              src_wave_addr_offset,
                                                              static_cast<index_t>(coherence));
463

464
            return bit_cast<bhalf8_t>(tmp);
465
466
        }
    }
Chao Liu's avatar
Chao Liu committed
467
468
469
470
    else if constexpr(is_same<T, int32_t>::value)
    {
        if constexpr(N == 1)
        {
471
472
473
474
            return llvm_amdgcn_raw_buffer_load_i32(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset,
                                                   static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
475
476
477
        }
        else if constexpr(N == 2)
        {
478
479
480
481
            return llvm_amdgcn_raw_buffer_load_i32x2(src_wave_buffer_resource,
                                                     src_thread_addr_offset,
                                                     src_wave_addr_offset,
                                                     static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
482
483
484
        }
        else if constexpr(N == 4)
        {
485
486
487
488
            return llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                     src_thread_addr_offset,
                                                     src_wave_addr_offset,
                                                     static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
489
490
491
492
493
        }
        else if constexpr(N == 8)
        {
            vector_type<int32_t, 8> tmp;

494
495
496
497
498
            tmp.AsType<int32x4_t>()(Number<0>{}) =
                llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                  src_thread_addr_offset,
                                                  src_wave_addr_offset,
                                                  static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
499

500
            tmp.AsType<int32x4_t>()(Number<1>{}) =
501
502
503
                llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                  src_thread_addr_offset,
                                                  src_wave_addr_offset + 4 * sizeof(int32_t),
504
                                                  static_cast<index_t>(coherence));
505
            return tmp.AsType<int32x8_t>()(Number<0>{});
Chao Liu's avatar
Chao Liu committed
506
507
        }
    }
508
    else if constexpr(is_same<T, int8_t>::value)
509
510
511
    {
        if constexpr(N == 1)
        {
512
513
514
515
            return llvm_amdgcn_raw_buffer_load_i8(src_wave_buffer_resource,
                                                  src_thread_addr_offset,
                                                  src_wave_addr_offset,
                                                  static_cast<index_t>(coherence));
516
        }
517
        else if constexpr(N == 2)
518
        {
519
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
520
521
522
523
            return llvm_amdgcn_raw_buffer_load_i8x2(src_wave_buffer_resource,
                                                    src_thread_addr_offset,
                                                    src_wave_addr_offset,
                                                    static_cast<index_t>(coherence));
524
#else
525
526
527
528
            int16_t tmp = llvm_amdgcn_raw_buffer_load_i16(src_wave_buffer_resource,
                                                          src_thread_addr_offset,
                                                          src_wave_addr_offset,
                                                          static_cast<index_t>(coherence));
529

530
            return bit_cast<int8x2_t>(tmp);
531
532
533
534
#endif
        }
        else if constexpr(N == 4)
        {
535
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
536
537
538
539
            return llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                    src_thread_addr_offset,
                                                    src_wave_addr_offset,
                                                    static_cast<index_t>(coherence));
540
#else
541
542
543
544
            int32_t tmp = llvm_amdgcn_raw_buffer_load_i32(src_wave_buffer_resource,
                                                          src_thread_addr_offset,
                                                          src_wave_addr_offset,
                                                          static_cast<index_t>(coherence));
545

546
            return bit_cast<int8x4_t>(tmp);
547
548
549
550
#endif
        }
        else if constexpr(N == 8)
        {
551
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
552
553
            vector_type<int8_t, 8> tmp;

554
555
556
557
558
            tmp.AsType<int8x4_t>()(Number<0>{}) =
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset,
                                                 static_cast<index_t>(coherence));
559
560

            tmp.AsType<int8x4_t>()(Number<1>{}) =
561
562
563
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 4 * sizeof(int8_t),
564
                                                 static_cast<index_t>(coherence));
565
566
567

            return tmp.AsType<int8x8_t>()(Number<0>{});
#else
568
569
570
571
            int32x2_t tmp = llvm_amdgcn_raw_buffer_load_i32x2(src_wave_buffer_resource,
                                                              src_thread_addr_offset,
                                                              src_wave_addr_offset,
                                                              static_cast<index_t>(coherence));
572

573
            return bit_cast<int8x8_t>(tmp);
574
575
576
577
#endif
        }
        else if constexpr(N == 16)
        {
578
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
579
580
            vector_type<int8_t, 16> tmp;

581
582
583
584
585
            tmp.AsType<int8x4_t>()(Number<0>{}) =
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset,
                                                 static_cast<index_t>(coherence));
586
587

            tmp.AsType<int8x4_t>()(Number<1>{}) =
588
589
590
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 4 * sizeof(int8_t),
591
                                                 static_cast<index_t>(coherence));
592
593

            tmp.AsType<int8x4_t>()(Number<2>{}) =
594
595
596
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 8 * sizeof(int8_t),
597
                                                 static_cast<index_t>(coherence));
598
599

            tmp.AsType<int8x4_t>()(Number<3>{}) =
600
601
602
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 12 * sizeof(int8_t),
603
                                                 static_cast<index_t>(coherence));
604
605
606

            return tmp.AsType<int8x16_t>()(Number<0>{});
#else
607
608
609
610
            int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                              src_thread_addr_offset,
                                                              src_wave_addr_offset,
                                                              static_cast<index_t>(coherence));
611

612
            return bit_cast<int8x16_t>(tmp);
613
#endif
614
615
        }
    }
Chao Liu's avatar
Chao Liu committed
616
617
}

618
619
620
template <typename T,
          index_t N,
          amd_buffer_coherence_bits coherence = amd_buffer_coherence_bits::default_coherence>
621
622
623
624
__device__ void amd_buffer_store_impl(const typename vector_type<T, N>::type src_thread_data,
                                      int32x4_t dst_wave_buffer_resource,
                                      index_t dst_thread_addr_offset,
                                      index_t dst_wave_addr_offset)
Chao Liu's avatar
Chao Liu committed
625
{
626
    static_assert(
Chao Liu's avatar
Chao Liu committed
627
628
629
        (is_same<T, double>::value && (N == 1 || N == 2)) ||
            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4)) ||
            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
630
            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
631
            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4)) ||
Chao Liu's avatar
Chao Liu committed
632
            (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
633
        "wrong! not implemented");
Chao Liu's avatar
Chao Liu committed
634

Chao Liu's avatar
Chao Liu committed
635
636
637
638
639
    if constexpr(is_same<T, double>::value)
    {
        // use fp32 store to mimic fp64 store
        if constexpr(N == 1)
        {
640
            llvm_amdgcn_raw_buffer_store_fp32x2(bit_cast<float2_t>(src_thread_data),
Chao Liu's avatar
Chao Liu committed
641
642
643
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
644
                                                static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
645
646
647
        }
        else if constexpr(N == 2)
        {
648
            llvm_amdgcn_raw_buffer_store_fp32x4(bit_cast<float4_t>(src_thread_data),
Chao Liu's avatar
Chao Liu committed
649
650
651
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
652
                                                static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
653
654
655
        }
    }
    else if constexpr(is_same<T, float>::value)
Chao Liu's avatar
Chao Liu committed
656
657
658
    {
        if constexpr(N == 1)
        {
659
660
661
662
            llvm_amdgcn_raw_buffer_store_fp32(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
663
                                              static_cast<index_t>(coherence));
664
665
666
667
        }
        else if constexpr(N == 2)
        {
            llvm_amdgcn_raw_buffer_store_fp32x2(src_thread_data,
Chao Liu's avatar
Chao Liu committed
668
669
670
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
671
                                                static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
672
673
674
        }
        else if constexpr(N == 4)
        {
675
676
677
678
            llvm_amdgcn_raw_buffer_store_fp32x4(src_thread_data,
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
679
                                                static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
680
681
        }
    }
Chao Liu's avatar
Chao Liu committed
682
683
684
685
686
687
688
689
    else if constexpr(is_same<T, half_t>::value)
    {
        if constexpr(N == 1)
        {
            llvm_amdgcn_raw_buffer_store_fp16(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
690
                                              static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
691
692
693
694
695
696
697
        }
        else if constexpr(N == 2)
        {
            llvm_amdgcn_raw_buffer_store_fp16x2(src_thread_data,
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
698
                                                static_cast<index_t>(coherence));
699
700
701
702
703
704
705
        }
        else if constexpr(N == 4)
        {
            llvm_amdgcn_raw_buffer_store_fp16x4(src_thread_data,
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
706
                                                static_cast<index_t>(coherence));
707
708
709
        }
        else if constexpr(N == 8)
        {
710
#if 0
711
712
713
714
715
716
            vector_type<half_t, 8> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_store_fp16x4(tmp.AsType<half4_t>()[Number<0>{}],
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
717
                                                static_cast<index_t>(coherence));
718
719
720
721
722

            llvm_amdgcn_raw_buffer_store_fp16x4(tmp.AsType<half4_t>()[Number<1>{}],
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset + 4 * sizeof(half_t),
723
                                                static_cast<index_t>(coherence));
724
#else
725
            llvm_amdgcn_raw_buffer_store_fp32x4(bit_cast<float4_t>(src_thread_data),
726
727
728
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
729
                                                static_cast<index_t>(coherence));
730
#endif
731
732
        }
    }
733
    else if constexpr(is_same<T, bhalf_t>::value)
734
735
736
737
738
739
740
    {
        if constexpr(N == 1)
        {
            llvm_amdgcn_raw_buffer_store_i16(src_thread_data,
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
741
                                             static_cast<index_t>(coherence));
742
743
744
        }
        else if constexpr(N == 2)
        {
745
746
747
748
            llvm_amdgcn_raw_buffer_store_i16x2(src_thread_data,
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
749
                                               static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
750
751
752
        }
        else if constexpr(N == 4)
        {
753
754
755
756
            llvm_amdgcn_raw_buffer_store_i16x4(src_thread_data,
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
757
                                               static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
758
759
760
        }
        else if constexpr(N == 8)
        {
761
            vector_type<bhalf_t, 8> tmp{src_thread_data};
Chao Liu's avatar
Chao Liu committed
762

763
764
765
766
            llvm_amdgcn_raw_buffer_store_i16x4(tmp.AsType<bhalf4_t>()[Number<0>{}],
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
767
                                               static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
768

769
770
771
772
            llvm_amdgcn_raw_buffer_store_i16x4(tmp.AsType<bhalf4_t>()[Number<1>{}],
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset + 4 * sizeof(bhalf_t),
773
                                               static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
774
775
        }
    }
Chao Liu's avatar
Chao Liu committed
776
777
778
779
    else if constexpr(is_same<T, int32_t>::value)
    {
        if constexpr(N == 1)
        {
780
781
782
783
            llvm_amdgcn_raw_buffer_store_i32(src_thread_data,
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
784
                                             static_cast<index_t>(coherence));
785
786
787
788
        }
        else if constexpr(N == 2)
        {
            llvm_amdgcn_raw_buffer_store_i32x2(src_thread_data,
Chao Liu's avatar
Chao Liu committed
789
790
791
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
792
                                               static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
793
794
795
        }
        else if constexpr(N == 4)
        {
796
797
798
799
            llvm_amdgcn_raw_buffer_store_i32x4(src_thread_data,
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
800
                                               static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
801
802
803
804
805
806
        }
    }
    else if constexpr(is_same<T, int8_t>::value)
    {
        if constexpr(N == 1)
        {
807
808
809
810
            llvm_amdgcn_raw_buffer_store_i8(src_thread_data,
                                            dst_wave_buffer_resource,
                                            dst_thread_addr_offset,
                                            dst_wave_addr_offset,
811
                                            static_cast<index_t>(coherence));
Chao Liu's avatar
Chao Liu committed
812
813
814
        }
        else if constexpr(N == 2)
        {
815
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
816
817
818
819
            llvm_amdgcn_raw_buffer_store_i8x2(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
820
                                              static_cast<index_t>(coherence));
821
#else
822
            llvm_amdgcn_raw_buffer_store_i16(bit_cast<int16_t>(src_thread_data),
823
824
825
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
826
                                             static_cast<index_t>(coherence));
827
#endif
Chao Liu's avatar
Chao Liu committed
828
829
830
        }
        else if constexpr(N == 4)
        {
831
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
832
833
834
835
            llvm_amdgcn_raw_buffer_store_i8x4(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
836
                                              static_cast<index_t>(coherence));
837
#else
838
            llvm_amdgcn_raw_buffer_store_i32(bit_cast<int32_t>(src_thread_data),
839
840
841
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
842
                                             static_cast<index_t>(coherence));
843
#endif
Chao Liu's avatar
Chao Liu committed
844
        }
845
846
        else if constexpr(N == 8)
        {
847
            llvm_amdgcn_raw_buffer_store_i32x2(bit_cast<int32x2_t>(src_thread_data),
848
849
850
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
851
                                               static_cast<index_t>(coherence));
852
853
854
        }
        else if constexpr(N == 16)
        {
855
            llvm_amdgcn_raw_buffer_store_i32x4(bit_cast<int32x4_t>(src_thread_data),
856
857
858
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
859
                                               static_cast<index_t>(coherence));
860
861
        }
    }
Chao Liu's avatar
Chao Liu committed
862
863
}

zjing14's avatar
zjing14 committed
864
865
866
867
868
869
870
template <typename T, index_t N>
__device__ void amd_buffer_atomic_add_impl(const typename vector_type<T, N>::type src_thread_data,
                                           int32x4_t dst_wave_buffer_resource,
                                           index_t dst_thread_addr_offset,
                                           index_t dst_wave_addr_offset)
{
    static_assert((is_same<T, float>::value && (N == 1 || N == 2 || N == 4)) ||
Chao Liu's avatar
Chao Liu committed
871
                      (is_same<T, half_t>::value && (N == 2 || N == 4 || N == 8)) ||
zjing14's avatar
zjing14 committed
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
                      (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4)),
                  "wrong! not implemented");

    if constexpr(is_same<T, float>::value)
    {
        if constexpr(N == 1)
        {
            llvm_amdgcn_raw_buffer_atomic_add_fp32(src_thread_data,
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset,
                                                   0);
        }
        else if constexpr(N == 2)
        {
            vector_type<float, 2> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_atomic_add_fp32(tmp.AsType<float>()[Number<0>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset,
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_add_fp32(tmp.AsType<float>()[Number<1>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + sizeof(float),
                                                   0);
        }
        else if constexpr(N == 4)
        {
            vector_type<float, 4> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_atomic_add_fp32(tmp.AsType<float>()[Number<0>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset,
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_add_fp32(tmp.AsType<float>()[Number<1>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + sizeof(float),
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_add_fp32(tmp.AsType<float>()[Number<2>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + 2 * sizeof(float),
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_add_fp32(tmp.AsType<float>()[Number<3>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + 3 * sizeof(float),
                                                   0);
        }
    }
Chao Liu's avatar
Chao Liu committed
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
    else if constexpr(is_same<T, half_t>::value)
    {
        if constexpr(N == 2)
        {
            llvm_amdgcn_raw_buffer_atomic_add_fp16x2(src_thread_data,
                                                     dst_wave_buffer_resource,
                                                     dst_thread_addr_offset,
                                                     dst_wave_addr_offset,
                                                     0);
        }
        else if constexpr(N == 4)
        {
            vector_type<half_t, 4> tmp{src_thread_data};

            static_for<0, 2, 1>{}([&](auto i) {
                llvm_amdgcn_raw_buffer_atomic_add_fp16x2(tmp.AsType<half2_t>()[i],
                                                         dst_wave_buffer_resource,
                                                         dst_thread_addr_offset,
                                                         dst_wave_addr_offset + i * sizeof(half2_t),
                                                         0);
            });
        }
        else if constexpr(N == 8)
        {
            vector_type<half_t, 8> tmp{src_thread_data};

            static_for<0, 4, 1>{}([&](auto i) {
                llvm_amdgcn_raw_buffer_atomic_add_fp16x2(tmp.AsType<half2_t>()[i],
                                                         dst_wave_buffer_resource,
                                                         dst_thread_addr_offset,
                                                         dst_wave_addr_offset + i * sizeof(half2_t),
                                                         0);
            });
        }
    }
zjing14's avatar
zjing14 committed
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
    else if constexpr(is_same<T, int32_t>::value)
    {
        if constexpr(N == 1)
        {
            llvm_amdgcn_raw_buffer_atomic_add_i32(src_thread_data,
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset,
                                                  0);
        }
        else if constexpr(N == 2)
        {
            vector_type<int32_t, 2> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_atomic_add_i32(tmp.AsType<int32_t>()[Number<0>{}],
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset,
                                                  0);

            llvm_amdgcn_raw_buffer_atomic_add_i32(tmp.AsType<int32_t>()[Number<1>{}],
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset + sizeof(int32_t),
                                                  0);
        }
        else if constexpr(N == 4)
        {
            vector_type<int32_t, 4> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_atomic_add_i32(tmp.AsType<int32_t>()[Number<0>{}],
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset,
                                                  0);

            llvm_amdgcn_raw_buffer_atomic_add_i32(tmp.AsType<int32_t>()[Number<1>{}],
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset + sizeof(int32_t),
                                                  0);

            llvm_amdgcn_raw_buffer_atomic_add_i32(tmp.AsType<int32_t>()[Number<2>{}],
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset + 2 * sizeof(int32_t),
                                                  0);

            llvm_amdgcn_raw_buffer_atomic_add_i32(tmp.AsType<int32_t>()[Number<3>{}],
                                                  dst_wave_buffer_resource,
                                                  dst_thread_addr_offset,
                                                  dst_wave_addr_offset + 3 * sizeof(int32_t),
                                                  0);
        }
    }
}

rocking5566's avatar
rocking5566 committed
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
template <typename T, index_t N>
__device__ void amd_buffer_atomic_max_impl(const typename vector_type<T, N>::type src_thread_data,
                                           int32x4_t dst_wave_buffer_resource,
                                           index_t dst_thread_addr_offset,
                                           index_t dst_wave_addr_offset)
{
    static_assert((is_same<T, double>::value && (N == 1 || N == 2 || N == 4)),
                  "wrong! not implemented");
    if constexpr(is_same<T, double>::value)
    {
        if constexpr(N == 1)
        {
            llvm_amdgcn_raw_buffer_atomic_max_fp64(src_thread_data,
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset,
                                                   0);
        }
        else if constexpr(N == 2)
        {
            vector_type<double, 2> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_atomic_max_fp64(tmp.AsType<double>()[Number<0>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset,
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_max_fp64(tmp.AsType<double>()[Number<1>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + sizeof(double),
                                                   0);
        }
        else if constexpr(N == 4)
        {
            vector_type<double, 4> tmp{src_thread_data};

            llvm_amdgcn_raw_buffer_atomic_max_fp64(tmp.AsType<double>()[Number<0>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset,
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_max_fp64(tmp.AsType<double>()[Number<1>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + sizeof(double),
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_max_fp64(tmp.AsType<double>()[Number<2>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + 2 * sizeof(double),
                                                   0);

            llvm_amdgcn_raw_buffer_atomic_max_fp64(tmp.AsType<double>()[Number<3>{}],
                                                   dst_wave_buffer_resource,
                                                   dst_thread_addr_offset,
                                                   dst_wave_addr_offset + 3 * sizeof(double),
                                                   0);
        }
    }
}

Chao Liu's avatar
Chao Liu committed
1087
// buffer_load requires:
Chao Liu's avatar
Chao Liu committed
1088
//   1) p_src_wave must point to global memory space
1089
//   2) p_src_wave must be a wavewise pointer.
Chao Liu's avatar
Chao Liu committed
1090
// It is user's responsibility to make sure that is true.
1091
1092
1093
template <typename T,
          index_t N,
          amd_buffer_coherence_bits coherence = amd_buffer_coherence_bits::default_coherence>
1094
__device__ typename vector_type_maker<T, N>::type::type
Jianfeng Yan's avatar
Jianfeng Yan committed
1095
1096
1097
1098
amd_buffer_load_invalid_element_return_zero(const T* p_src_wave,
                                            index_t src_thread_element_offset,
                                            bool src_thread_element_valid,
                                            index_t src_element_space_size)
Chao Liu's avatar
Chao Liu committed
1099
1100
{
    const int32x4_t src_wave_buffer_resource =
1101
        make_wave_buffer_resource(p_src_wave, src_element_space_size);
Chao Liu's avatar
Chao Liu committed
1102

1103
1104
1105
1106
    index_t src_thread_addr_offset = src_thread_element_offset * sizeof(T);

    using vector_t = typename vector_type_maker<T, N>::type::type;
    using scalar_t = typename scalar_type<vector_t>::type;
Chao Liu's avatar
Chao Liu committed
1107

1108
1109
    constexpr index_t vector_size = scalar_type<vector_t>::vector_size;

Chao Liu's avatar
Chao Liu committed
1110
#if CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK
1111
    uint32_t src_addr_shift = src_thread_element_valid ? 0 : 0x80000000;
Chao Liu's avatar
Chao Liu committed
1112

1113
    return amd_buffer_load_impl<scalar_t, vector_size, coherence>(
Chao Liu's avatar
Chao Liu committed
1114
1115
        src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0);
#else
1116
    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size, coherence>(
1117
        src_wave_buffer_resource, src_thread_addr_offset, 0);
Chao Liu's avatar
Chao Liu committed
1118

1119
    return src_thread_element_valid ? tmp : vector_t(0);
Chao Liu's avatar
Chao Liu committed
1120
1121
1122
#endif
}

1123
// buffer_load requires:
Chao Liu's avatar
Chao Liu committed
1124
//   1) p_src_wave must point to global memory space
1125
1126
//   2) p_src_wave must be a wavewise pointer.
// It is user's responsibility to make sure that is true.
1127
1128
1129
template <typename T,
          index_t N,
          amd_buffer_coherence_bits coherence = amd_buffer_coherence_bits::default_coherence>
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
__device__ typename vector_type_maker<T, N>::type::type
amd_buffer_load_invalid_element_return_customized_value(const T* p_src_wave,
                                                        index_t src_thread_element_offset,
                                                        bool src_thread_element_valid,
                                                        index_t src_element_space_size,
                                                        T customized_value)
{
    const int32x4_t src_wave_buffer_resource =
        make_wave_buffer_resource(p_src_wave, src_element_space_size);

    index_t src_thread_addr_offset = src_thread_element_offset * sizeof(T);

    using vector_t = typename vector_type_maker<T, N>::type::type;
    using scalar_t = typename scalar_type<vector_t>::type;

    constexpr index_t vector_size = scalar_type<vector_t>::vector_size;

1147
    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size, coherence>(
1148
1149
1150
1151
1152
        src_wave_buffer_resource, src_thread_addr_offset, 0);

    return src_thread_element_valid ? tmp : vector_t(customized_value);
}

Chao Liu's avatar
Chao Liu committed
1153
// buffer_store requires:
Chao Liu's avatar
Chao Liu committed
1154
//   1) p_dst_wave must point to global memory
zjing14's avatar
zjing14 committed
1155
//   2) p_dst_wave must be a wavewise pointer.
Chao Liu's avatar
Chao Liu committed
1156
// It is user's responsibility to make sure that is true.
1157
1158
1159
template <typename T,
          index_t N,
          amd_buffer_coherence_bits coherence = amd_buffer_coherence_bits::default_coherence>
1160
1161
1162
1163
1164
__device__ void amd_buffer_store(const typename vector_type_maker<T, N>::type::type src_thread_data,
                                 T* p_dst_wave,
                                 const index_t dst_thread_element_offset,
                                 const bool dst_thread_element_valid,
                                 const index_t dst_element_space_size)
Chao Liu's avatar
Chao Liu committed
1165
1166
{
    const int32x4_t dst_wave_buffer_resource =
1167
        make_wave_buffer_resource(p_dst_wave, dst_element_space_size);
Chao Liu's avatar
Chao Liu committed
1168

1169
    index_t dst_thread_addr_offset = dst_thread_element_offset * sizeof(T);
Chao Liu's avatar
Chao Liu committed
1170

1171
1172
1173
1174
    using vector_t                = typename vector_type_maker<T, N>::type::type;
    using scalar_t                = typename scalar_type<vector_t>::type;
    constexpr index_t vector_size = scalar_type<vector_t>::vector_size;

Chao Liu's avatar
Chao Liu committed
1175
#if CK_EXPERIMENTAL_USE_BUFFER_STORE_OOB_CHECK_OFFSET_TRICK
1176
    uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000;
Chao Liu's avatar
Chao Liu committed
1177

1178
    amd_buffer_store_impl<scalar_t, vector_size, coherence>(
Chao Liu's avatar
Chao Liu committed
1179
1180
        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
#else
1181
    if(dst_thread_element_valid)
Chao Liu's avatar
Chao Liu committed
1182
    {
1183
        amd_buffer_store_impl<scalar_t, vector_size, coherence>(
Chao Liu's avatar
Chao Liu committed
1184
1185
1186
1187
1188
            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
    }
#endif
}

zjing14's avatar
zjing14 committed
1189
// buffer_atomic_add requires:
Chao Liu's avatar
Chao Liu committed
1190
//   1) p_dst_wave must point to global memory
zjing14's avatar
zjing14 committed
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
//   2) p_dst_wave must be a wavewise pointer.
// It is user's responsibility to make sure that is true.
template <typename T, index_t N>
__device__ void
amd_buffer_atomic_add(const typename vector_type_maker<T, N>::type::type src_thread_data,
                      T* p_dst_wave,
                      const index_t dst_thread_element_offset,
                      const bool dst_thread_element_valid,
                      const index_t dst_element_space_size)
{
    const int32x4_t dst_wave_buffer_resource =
        make_wave_buffer_resource(p_dst_wave, dst_element_space_size);

    index_t dst_thread_addr_offset = dst_thread_element_offset * sizeof(T);

    using vector_t                = typename vector_type_maker<T, N>::type::type;
    using scalar_t                = typename scalar_type<vector_t>::type;
    constexpr index_t vector_size = scalar_type<vector_t>::vector_size;

#if CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_ADD_OOB_CHECK_OFFSET_TRICK
1211
    uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000;
zjing14's avatar
zjing14 committed
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223

    amd_buffer_atomic_add_impl<scalar_t, vector_size>(
        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
#else
    if(dst_thread_element_valid)
    {
        amd_buffer_atomic_add_impl<scalar_t, vector_size>(
            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
    }
#endif
}

rocking5566's avatar
rocking5566 committed
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
// buffer_atomic_max requires:
//   1) p_dst_wave must point to global memory
//   2) p_dst_wave must be a wavewise pointer.
// It is user's responsibility to make sure that is true.
template <typename T, index_t N>
__device__ void
amd_buffer_atomic_max(const typename vector_type_maker<T, N>::type::type src_thread_data,
                      T* p_dst_wave,
                      const index_t dst_thread_element_offset,
                      const bool dst_thread_element_valid,
                      const index_t dst_element_space_size)
{
    const int32x4_t dst_wave_buffer_resource =
        make_wave_buffer_resource(p_dst_wave, dst_element_space_size);

    index_t dst_thread_addr_offset = dst_thread_element_offset * sizeof(T);

    using vector_t                = typename vector_type_maker<T, N>::type::type;
    using scalar_t                = typename scalar_type<vector_t>::type;
    constexpr index_t vector_size = scalar_type<vector_t>::vector_size;

#if CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_MAX_OOB_CHECK_OFFSET_TRICK
1246
    uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x80000000;
rocking5566's avatar
rocking5566 committed
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258

    amd_buffer_atomic_max_impl<scalar_t, vector_size>(
        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
#else
    if(dst_thread_element_valid)
    {
        amd_buffer_atomic_max_impl<scalar_t, vector_size>(
            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
    }
#endif
}

Chao Liu's avatar
Chao Liu committed
1259
} // namespace ck