amd_buffer_addressing.hpp 45.5 KB
Newer Older
1
#pragma once
2
#include "data_type.hpp"
Chao Liu's avatar
Chao Liu committed
3
4
5
6

namespace ck {

template <typename T>
Chao Liu's avatar
tidy  
Chao Liu committed
7
union BufferResource
Chao Liu's avatar
Chao Liu committed
8
9
10
{
    // 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
11
    int32x4_t content;
12
13
14
    StaticallyIndexedArray<T*, 2> address;
    StaticallyIndexedArray<int32_t, 4> range;
    StaticallyIndexedArray<int32_t, 4> config;
Chao Liu's avatar
Chao Liu committed
15
16
17
};

template <typename T>
18
__device__ int32x4_t make_wave_buffer_resource(T* p_wave, index_t element_space_size)
Chao Liu's avatar
Chao Liu committed
19
{
Chao Liu's avatar
tidy  
Chao Liu committed
20
    BufferResource<T> wave_buffer_resource;
Chao Liu's avatar
Chao Liu committed
21
22

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

29
    return wave_buffer_resource.content;
Chao Liu's avatar
Chao Liu committed
30
31
}

Chao Liu's avatar
Chao Liu committed
32
// buffer load i8
Chao Liu's avatar
Chao Liu committed
33
__device__ int8_t
34
35
36
37
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");
38
39

__device__ int8x2_t
40
41
42
43
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");
44
45

__device__ int8x4_t
46
47
48
49
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");
50

Chao Liu's avatar
Chao Liu committed
51
// buffer load i16
52
__device__ bhalf_t
53
54
55
llvm_amdgcn_raw_buffer_load_i16(int32x4_t srsrc,
                                index_t voffset,
                                index_t soffset,
56
57
                                index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i16");

58
__device__ bhalf2_t
59
60
61
62
63
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");

64
__device__ bhalf4_t
65
66
67
68
69
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
70
// buffer load i32
Chao Liu's avatar
Chao Liu committed
71
__device__ int32_t
72
73
74
75
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
76
77

__device__ int32x2_t
78
79
80
81
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
82
83

__device__ int32x4_t
84
85
86
87
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");
88

Chao Liu's avatar
Chao Liu committed
89
// buffer load fp16
90
__device__ half_t
91
92
93
94
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");
95
96

__device__ half2_t
97
98
99
100
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");
101
102

__device__ half4_t
103
104
105
106
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
107

Chao Liu's avatar
Chao Liu committed
108
// buffer load fp32
Chao Liu's avatar
Chao Liu committed
109
__device__ float
110
111
112
113
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
114
115

__device__ float2_t
116
117
118
119
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
120
121

__device__ float4_t
122
123
124
125
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
126

Chao Liu's avatar
Chao Liu committed
127
// buffer store i8
Chao Liu's avatar
Chao Liu committed
128
__device__ void
129
130
131
132
133
134
135
136
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
137
138
139
                                  int32x4_t rsrc,
                                  index_t voffset,
                                  index_t soffset,
140
                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i8");
Chao Liu's avatar
Chao Liu committed
141

142
__device__ void
143
144
145
146
147
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");
148

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

157
__device__ void
158
llvm_amdgcn_raw_buffer_store_i16x2(bhalf2_t vdata,
159
160
161
162
163
164
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i16");

__device__ void
165
llvm_amdgcn_raw_buffer_store_i16x4(bhalf4_t vdata,
166
167
168
169
170
                                   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
171
// buffer store i32
172
173
174
175
176
177
__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");
178

Chao Liu's avatar
Chao Liu committed
179
__device__ void
180
llvm_amdgcn_raw_buffer_store_i32x2(int32x2_t vdata,
Chao Liu's avatar
Chao Liu committed
181
182
183
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
184
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i32");
Chao Liu's avatar
Chao Liu committed
185
186

__device__ void
187
llvm_amdgcn_raw_buffer_store_i32x4(int32x4_t vdata,
Chao Liu's avatar
Chao Liu committed
188
189
190
                                   int32x4_t rsrc,
                                   index_t voffset,
                                   index_t soffset,
191
                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
Chao Liu's avatar
Chao Liu committed
192

Chao Liu's avatar
Chao Liu committed
193
// buffer store fp16
Chao Liu's avatar
Chao Liu committed
194
__device__ void
195
196
197
198
199
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
200
201

__device__ void
202
203
204
205
206
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
207

208
__device__ void
209
llvm_amdgcn_raw_buffer_store_fp16x4(half4_t vdata,
210
211
212
                                    int32x4_t rsrc,
                                    index_t voffset,
                                    index_t soffset,
213
                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4f16");
214

Chao Liu's avatar
Chao Liu committed
215
// buffer store fp32
216
__device__ void
217
218
219
220
221
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");
222
223

__device__ void
224
llvm_amdgcn_raw_buffer_store_fp32x2(float2_t vdata,
Chao Liu's avatar
Chao Liu committed
225
226
227
                                    int32x4_t rsrc,
                                    index_t voffset,
                                    index_t soffset,
228
                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2f32");
Chao Liu's avatar
Chao Liu committed
229
230

__device__ void
231
232
233
234
235
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");
236

Chao Liu's avatar
Chao Liu committed
237
238
239
240
241
242
243
244
245
// 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
246
247
248
249
250
251
252
__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
253
// buffer atomic-add fp32
zjing14's avatar
zjing14 committed
254
255
256
257
258
259
__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
260
261

template <typename T, index_t N>
262
263
264
__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
265
{
266
    static_assert(
Chao Liu's avatar
Chao Liu committed
267
268
        (is_same<T, double>::value && (N == 1 || N == 2 || N == 4)) ||
            (is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
269
            (is_same<T, half_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
270
            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
Chao Liu's avatar
Chao Liu committed
271
272
            (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)),
273
        "wrong! not implemented");
Chao Liu's avatar
Chao Liu committed
274

Chao Liu's avatar
Chao Liu committed
275
276
277
278
279
280
281
282
    if constexpr(is_same<T, double>::value)
    {
        // use fp32 load to mimic fp64 load
        if constexpr(N == 1)
        {
            const float2_t tmp = llvm_amdgcn_raw_buffer_load_fp32x2(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

283
            return bit_cast<double>(tmp);
Chao Liu's avatar
Chao Liu committed
284
285
286
287
288
289
        }
        else if constexpr(N == 2)
        {
            const float4_t tmp = llvm_amdgcn_raw_buffer_load_fp32x4(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

290
            return bit_cast<double2_t>(tmp);
Chao Liu's avatar
Chao Liu committed
291
292
293
294
295
296
297
298
299
300
301
302
303
        }
        else if constexpr(N == 4)
        {
            const float4_t f32_0 = llvm_amdgcn_raw_buffer_load_fp32x4(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

            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),
                                                   0);
            vector_type<double, 4> tmp;

304
305
            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
306
307
308
309
310

            return tmp.AsType<double4_t>()(Number<0>{});
        }
    }
    else if constexpr(is_same<T, float>::value)
Chao Liu's avatar
Chao Liu committed
311
312
313
    {
        if constexpr(N == 1)
        {
314
            return llvm_amdgcn_raw_buffer_load_fp32(
Chao Liu's avatar
Chao Liu committed
315
316
317
318
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 2)
        {
319
            return llvm_amdgcn_raw_buffer_load_fp32x2(
Chao Liu's avatar
Chao Liu committed
320
321
322
323
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 4)
        {
324
            return llvm_amdgcn_raw_buffer_load_fp32x4(
Chao Liu's avatar
Chao Liu committed
325
326
327
328
329
330
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 8)
        {
            vector_type<float, 8> tmp;

331
            tmp.AsType<float4_t>()(Number<0>{}) = llvm_amdgcn_raw_buffer_load_fp32x4(
Chao Liu's avatar
Chao Liu committed
332
333
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

334
            tmp.AsType<float4_t>()(Number<1>{}) =
335
336
337
338
                llvm_amdgcn_raw_buffer_load_fp32x4(src_wave_buffer_resource,
                                                   src_thread_addr_offset,
                                                   src_wave_addr_offset + 4 * sizeof(float),
                                                   0);
339

340
            return tmp.AsType<float8_t>()(Number<0>{});
341
342
343
344
345
346
        }
    }
    else if constexpr(is_same<T, half_t>::value)
    {
        if constexpr(N == 1)
        {
347
            return llvm_amdgcn_raw_buffer_load_fp16(
348
349
350
351
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 2)
        {
352
            return llvm_amdgcn_raw_buffer_load_fp16x2(
353
354
355
356
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 4)
        {
357
            return llvm_amdgcn_raw_buffer_load_fp16x4(
358
359
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
360
        else if constexpr(N == 8)
361
        {
Chao Liu's avatar
Chao Liu committed
362
            // use fp32 load to mimic fp16 load
363
            float4_t tmp = llvm_amdgcn_raw_buffer_load_fp32x4(
zjing14's avatar
zjing14 committed
364
365
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

366
            return bit_cast<half8_t>(tmp);
Chao Liu's avatar
Chao Liu committed
367
368
        }
    }
369
    else if constexpr(is_same<T, bhalf_t>::value)
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
    {
        if constexpr(N == 1)
        {
            return llvm_amdgcn_raw_buffer_load_i16(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 2)
        {
            return llvm_amdgcn_raw_buffer_load_i16x2(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 4)
        {
            return llvm_amdgcn_raw_buffer_load_i16x4(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 8)
        {
            int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

391
            return bit_cast<bhalf8_t>(tmp);
392
393
        }
    }
Chao Liu's avatar
Chao Liu committed
394
395
396
397
    else if constexpr(is_same<T, int32_t>::value)
    {
        if constexpr(N == 1)
        {
398
            return llvm_amdgcn_raw_buffer_load_i32(
Chao Liu's avatar
Chao Liu committed
399
400
401
402
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 2)
        {
403
            return llvm_amdgcn_raw_buffer_load_i32x2(
Chao Liu's avatar
Chao Liu committed
404
405
406
407
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 4)
        {
408
            return llvm_amdgcn_raw_buffer_load_i32x4(
Chao Liu's avatar
Chao Liu committed
409
410
411
412
413
414
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
        else if constexpr(N == 8)
        {
            vector_type<int32_t, 8> tmp;

415
            tmp.AsType<int32x4_t>()(Number<0>{}) = llvm_amdgcn_raw_buffer_load_i32x4(
Chao Liu's avatar
Chao Liu committed
416
417
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

418
            tmp.AsType<int32x4_t>()(Number<1>{}) =
419
420
421
422
                llvm_amdgcn_raw_buffer_load_i32x4(src_wave_buffer_resource,
                                                  src_thread_addr_offset,
                                                  src_wave_addr_offset + 4 * sizeof(int32_t),
                                                  0);
423
            return tmp.AsType<int32x8_t>()(Number<0>{});
Chao Liu's avatar
Chao Liu committed
424
425
        }
    }
426
    else if constexpr(is_same<T, int8_t>::value)
427
428
429
    {
        if constexpr(N == 1)
        {
430
            return llvm_amdgcn_raw_buffer_load_i8(
431
432
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
        }
433
        else if constexpr(N == 2)
434
        {
435
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
436
            return llvm_amdgcn_raw_buffer_load_i8x2(
437
438
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
#else
439
            int16_t tmp = llvm_amdgcn_raw_buffer_load_i16(
440
441
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

442
            return bit_cast<int8x2_t>(tmp);
443
444
445
446
#endif
        }
        else if constexpr(N == 4)
        {
447
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
448
            return llvm_amdgcn_raw_buffer_load_i8x4(
449
450
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
#else
451
            int32_t tmp = llvm_amdgcn_raw_buffer_load_i32(
452
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
453

454
            return bit_cast<int8x4_t>(tmp);
455
456
457
458
#endif
        }
        else if constexpr(N == 8)
        {
459
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
460
461
            vector_type<int8_t, 8> tmp;

462
            tmp.AsType<int8x4_t>()(Number<0>{}) = llvm_amdgcn_raw_buffer_load_i8x4(
463
464
465
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

            tmp.AsType<int8x4_t>()(Number<1>{}) =
466
467
468
469
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 4 * sizeof(int8_t),
                                                 0);
470
471
472

            return tmp.AsType<int8x8_t>()(Number<0>{});
#else
473
            int32x2_t tmp = llvm_amdgcn_raw_buffer_load_i32x2(
474
475
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

476
            return bit_cast<int8x8_t>(tmp);
477
478
479
480
#endif
        }
        else if constexpr(N == 16)
        {
481
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
482
483
            vector_type<int8_t, 16> tmp;

484
            tmp.AsType<int8x4_t>()(Number<0>{}) = llvm_amdgcn_raw_buffer_load_i8x4(
485
486
487
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

            tmp.AsType<int8x4_t>()(Number<1>{}) =
488
489
490
491
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 4 * sizeof(int8_t),
                                                 0);
492
493

            tmp.AsType<int8x4_t>()(Number<2>{}) =
494
495
496
497
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 8 * sizeof(int8_t),
                                                 0);
498
499

            tmp.AsType<int8x4_t>()(Number<3>{}) =
500
501
502
503
                llvm_amdgcn_raw_buffer_load_i8x4(src_wave_buffer_resource,
                                                 src_thread_addr_offset,
                                                 src_wave_addr_offset + 12 * sizeof(int8_t),
                                                 0);
504
505
506

            return tmp.AsType<int8x16_t>()(Number<0>{});
#else
507
            int32x4_t tmp = llvm_amdgcn_raw_buffer_load_i32x4(
508
509
                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);

510
            return bit_cast<int8x16_t>(tmp);
511
#endif
512
513
        }
    }
Chao Liu's avatar
Chao Liu committed
514
515
516
}

template <typename T, index_t N>
517
518
519
520
__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
521
{
522
    static_assert(
Chao Liu's avatar
Chao Liu committed
523
524
525
        (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)) ||
526
            (is_same<T, bhalf_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
527
            (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4)) ||
Chao Liu's avatar
Chao Liu committed
528
            (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4 || N == 8 || N == 16)),
529
        "wrong! not implemented");
Chao Liu's avatar
Chao Liu committed
530

Chao Liu's avatar
Chao Liu committed
531
532
533
534
535
    if constexpr(is_same<T, double>::value)
    {
        // use fp32 store to mimic fp64 store
        if constexpr(N == 1)
        {
536
            llvm_amdgcn_raw_buffer_store_fp32x2(bit_cast<float2_t>(src_thread_data),
Chao Liu's avatar
Chao Liu committed
537
538
539
540
541
542
543
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
                                                0);
        }
        else if constexpr(N == 2)
        {
544
            llvm_amdgcn_raw_buffer_store_fp32x4(bit_cast<float4_t>(src_thread_data),
Chao Liu's avatar
Chao Liu committed
545
546
547
548
549
550
551
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
                                                0);
        }
    }
    else if constexpr(is_same<T, float>::value)
Chao Liu's avatar
Chao Liu committed
552
553
554
    {
        if constexpr(N == 1)
        {
555
556
557
558
559
560
561
562
563
            llvm_amdgcn_raw_buffer_store_fp32(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
                                              0);
        }
        else if constexpr(N == 2)
        {
            llvm_amdgcn_raw_buffer_store_fp32x2(src_thread_data,
Chao Liu's avatar
Chao Liu committed
564
565
566
567
568
569
570
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
                                                0);
        }
        else if constexpr(N == 4)
        {
571
572
573
574
575
            llvm_amdgcn_raw_buffer_store_fp32x4(src_thread_data,
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
                                                0);
Chao Liu's avatar
Chao Liu committed
576
577
        }
    }
Chao Liu's avatar
Chao Liu committed
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
    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,
                                              0);
        }
        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,
                                                0);
595
596
597
598
599
600
601
602
603
604
605
        }
        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,
                                                0);
        }
        else if constexpr(N == 8)
        {
606
#if 0
607
608
609
610
611
612
613
614
615
616
617
618
619
            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,
                                                0);

            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),
                                                0);
620
#else
621
            llvm_amdgcn_raw_buffer_store_fp32x4(bit_cast<float4_t>(src_thread_data),
622
623
624
625
626
                                                dst_wave_buffer_resource,
                                                dst_thread_addr_offset,
                                                dst_wave_addr_offset,
                                                0);
#endif
627
628
        }
    }
629
    else if constexpr(is_same<T, bhalf_t>::value)
630
631
632
633
634
635
636
637
638
639
640
    {
        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,
                                             0);
        }
        else if constexpr(N == 2)
        {
641
642
643
644
645
            llvm_amdgcn_raw_buffer_store_i16x2(src_thread_data,
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
Chao Liu's avatar
Chao Liu committed
646
647
648
        }
        else if constexpr(N == 4)
        {
649
650
651
652
653
            llvm_amdgcn_raw_buffer_store_i16x4(src_thread_data,
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
Chao Liu's avatar
Chao Liu committed
654
655
656
        }
        else if constexpr(N == 8)
        {
657
            vector_type<bhalf_t, 8> tmp{src_thread_data};
Chao Liu's avatar
Chao Liu committed
658

659
660
661
662
663
            llvm_amdgcn_raw_buffer_store_i16x4(tmp.AsType<bhalf4_t>()[Number<0>{}],
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
Chao Liu's avatar
Chao Liu committed
664

665
666
667
668
669
            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),
                                               0);
Chao Liu's avatar
Chao Liu committed
670
671
        }
    }
Chao Liu's avatar
Chao Liu committed
672
673
674
675
    else if constexpr(is_same<T, int32_t>::value)
    {
        if constexpr(N == 1)
        {
676
677
678
679
680
681
682
683
684
            llvm_amdgcn_raw_buffer_store_i32(src_thread_data,
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
                                             0);
        }
        else if constexpr(N == 2)
        {
            llvm_amdgcn_raw_buffer_store_i32x2(src_thread_data,
Chao Liu's avatar
Chao Liu committed
685
686
687
688
689
690
691
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
        }
        else if constexpr(N == 4)
        {
692
693
694
695
696
            llvm_amdgcn_raw_buffer_store_i32x4(src_thread_data,
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
Chao Liu's avatar
Chao Liu committed
697
698
699
700
701
702
        }
    }
    else if constexpr(is_same<T, int8_t>::value)
    {
        if constexpr(N == 1)
        {
703
704
705
706
707
            llvm_amdgcn_raw_buffer_store_i8(src_thread_data,
                                            dst_wave_buffer_resource,
                                            dst_thread_addr_offset,
                                            dst_wave_addr_offset,
                                            0);
Chao Liu's avatar
Chao Liu committed
708
709
710
        }
        else if constexpr(N == 2)
        {
711
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
712
713
714
715
716
            llvm_amdgcn_raw_buffer_store_i8x2(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
                                              0);
717
#else
718
            llvm_amdgcn_raw_buffer_store_i16(bit_cast<int16_t>(src_thread_data),
719
720
721
722
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
                                             0);
723
#endif
Chao Liu's avatar
Chao Liu committed
724
725
726
        }
        else if constexpr(N == 4)
        {
727
#if !CK_WORKAROUND_SWDEV_XXXXXX_INT8_BUFFER_LOAD_STORE_ISSUE
728
729
730
731
732
            llvm_amdgcn_raw_buffer_store_i8x4(src_thread_data,
                                              dst_wave_buffer_resource,
                                              dst_thread_addr_offset,
                                              dst_wave_addr_offset,
                                              0);
733
#else
734
            llvm_amdgcn_raw_buffer_store_i32(bit_cast<int32_t>(src_thread_data),
735
736
737
738
                                             dst_wave_buffer_resource,
                                             dst_thread_addr_offset,
                                             dst_wave_addr_offset,
                                             0);
739
#endif
Chao Liu's avatar
Chao Liu committed
740
        }
741
742
        else if constexpr(N == 8)
        {
743
            llvm_amdgcn_raw_buffer_store_i32x2(bit_cast<int32x2_t>(src_thread_data),
744
745
746
747
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
748
749
750
        }
        else if constexpr(N == 16)
        {
751
            llvm_amdgcn_raw_buffer_store_i32x4(bit_cast<int32x4_t>(src_thread_data),
752
753
754
755
                                               dst_wave_buffer_resource,
                                               dst_thread_addr_offset,
                                               dst_wave_addr_offset,
                                               0);
756
757
        }
    }
Chao Liu's avatar
Chao Liu committed
758
759
}

zjing14's avatar
zjing14 committed
760
761
762
763
764
765
766
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
767
                      (is_same<T, half_t>::value && (N == 2 || N == 4 || N == 8)) ||
zjing14's avatar
zjing14 committed
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
                      (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
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
    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
861
862
863
864
865
866
867
868
869
870
871
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
    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);
        }
    }
}

Chao Liu's avatar
Chao Liu committed
918
// buffer_load requires:
Chao Liu's avatar
Chao Liu committed
919
//   1) p_src_wave must point to global memory space
920
//   2) p_src_wave must be a wavewise pointer.
Chao Liu's avatar
Chao Liu committed
921
922
// It is user's responsibility to make sure that is true.
template <typename T, index_t N>
923
__device__ typename vector_type_maker<T, N>::type::type
Jianfeng Yan's avatar
Jianfeng Yan committed
924
925
926
927
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
928
929
{
    const int32x4_t src_wave_buffer_resource =
930
        make_wave_buffer_resource(p_src_wave, src_element_space_size);
Chao Liu's avatar
Chao Liu committed
931

932
933
934
935
    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
936

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

Chao Liu's avatar
Chao Liu committed
939
#if CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK
940
    uint32_t src_addr_shift = src_thread_element_valid ? 0 : 0x7fffffff;
Chao Liu's avatar
Chao Liu committed
941

942
    return amd_buffer_load_impl<scalar_t, vector_size>(
Chao Liu's avatar
Chao Liu committed
943
944
        src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0);
#else
945
    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size>(
946
        src_wave_buffer_resource, src_thread_addr_offset, 0);
Chao Liu's avatar
Chao Liu committed
947

948
    return src_thread_element_valid ? tmp : vector_t(0);
Chao Liu's avatar
Chao Liu committed
949
950
951
#endif
}

952
// buffer_load requires:
Chao Liu's avatar
Chao Liu committed
953
//   1) p_src_wave must point to global memory space
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
//   2) p_src_wave must be a wavewise pointer.
// It is user's responsibility to make sure that is true.
template <typename T, index_t N>
__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;

    vector_t tmp = amd_buffer_load_impl<scalar_t, vector_size>(
        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
980
// buffer_store requires:
Chao Liu's avatar
Chao Liu committed
981
//   1) p_dst_wave must point to global memory
zjing14's avatar
zjing14 committed
982
//   2) p_dst_wave must be a wavewise pointer.
Chao Liu's avatar
Chao Liu committed
983
984
// It is user's responsibility to make sure that is true.
template <typename T, index_t N>
985
986
987
988
989
__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
990
991
{
    const int32x4_t dst_wave_buffer_resource =
992
        make_wave_buffer_resource(p_dst_wave, dst_element_space_size);
Chao Liu's avatar
Chao Liu committed
993

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

996
997
998
999
    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
1000
#if CK_EXPERIMENTAL_USE_BUFFER_STORE_OOB_CHECK_OFFSET_TRICK
1001
    uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x7fffffff;
Chao Liu's avatar
Chao Liu committed
1002

1003
    amd_buffer_store_impl<scalar_t, vector_size>(
Chao Liu's avatar
Chao Liu committed
1004
1005
        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
#else
1006
    if(dst_thread_element_valid)
Chao Liu's avatar
Chao Liu committed
1007
    {
1008
        amd_buffer_store_impl<scalar_t, vector_size>(
Chao Liu's avatar
Chao Liu committed
1009
1010
1011
1012
1013
            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
    }
#endif
}

zjing14's avatar
zjing14 committed
1014
// buffer_atomic_add requires:
Chao Liu's avatar
Chao Liu committed
1015
//   1) p_dst_wave must point to global memory
zjing14's avatar
zjing14 committed
1016
1017
1018
1019
1020
1021
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
//   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
    uint32_t dst_addr_shift = dst_thread_element_valid ? 0 : 0x7fffffff;

    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
}

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