"test/vscode:/vscode.git/clone" did not exist on "56a724eba301159f09a6661dc085954544705c00"
blockwise_4d_tensor_op.hip.hpp 26.9 KB
Newer Older
1
#pragma once
2
#include "ConstantTensorDescriptor.hip.hpp"
3

Chao Liu's avatar
Chao Liu committed
4
template <index_t BlockSize, class Float, class DstDesc, class F>
5
__device__ void
Chao Liu's avatar
Chao Liu committed
6
blockwise_4d_tensor_pointwise_operation_unary(DstDesc, Float* __restrict__ p_dst, F f)
Chao Liu's avatar
Chao Liu committed
7
{
Chao Liu's avatar
Chao Liu committed
8
9
10
11
    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};
Chao Liu's avatar
Chao Liu committed
12

13
14
    constexpr auto dst_desc = DstDesc{};

Chao Liu's avatar
Chao Liu committed
15
    constexpr auto desc = make_ConstantTensorDescriptor(dst_desc.GetLengths());
Chao Liu's avatar
Chao Liu committed
16

17
18
19
#if 0
    if(threadIdx.x == 0)
    {
Chao Liu's avatar
Chao Liu committed
20
21
        print_ConstantTensorDescriptor(dst_desc, "blockwise_4d_tensor_op_unary: dst_desc: ");
        print_ConstantTensorDescriptor(desc, "blockwise_4d_tensor_op_unary: desc: ");
22
23
24
    }
#endif

Chao Liu's avatar
Chao Liu committed
25
    constexpr index_t NLoop = desc.GetElementSize() / BlockSize;
Chao Liu's avatar
Chao Liu committed
26

Chao Liu's avatar
Chao Liu committed
27
    for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
28
    {
Chao Liu's avatar
Chao Liu committed
29
        index_t is = threadIdx.x + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
30

Chao Liu's avatar
Chao Liu committed
31
        const index_t did0 = is / desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
32
33
34

        is -= did0 * desc.GetStride(I0);

Chao Liu's avatar
Chao Liu committed
35
        const index_t did1 = is / desc.GetStride(I1);
Chao Liu's avatar
Chao Liu committed
36
37
38

        is -= did1 * desc.GetStride(I1);

Chao Liu's avatar
Chao Liu committed
39
        const index_t did2 = is / desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
40
41
42

        is -= did2 * desc.GetStride(I2);

Chao Liu's avatar
Chao Liu committed
43
        const index_t did3 = is / desc.GetStride(I3);
Chao Liu's avatar
Chao Liu committed
44

Chao Liu's avatar
Chao Liu committed
45
        const index_t dindex = dst_desc.Get1dIndex(did0, did1, did2, did3);
Chao Liu's avatar
Chao Liu committed
46

Chao Liu's avatar
Chao Liu committed
47
        f(p_dst[dindex]);
Chao Liu's avatar
Chao Liu committed
48
49
50
51
52
53
    }

    constexpr bool has_tail = (desc.GetElementSize() > NLoop * BlockSize);

    if(has_tail)
    {
Chao Liu's avatar
Chao Liu committed
54
        index_t is = threadIdx.x + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
55
56
57

        if(is < desc.GetElementSize())
        {
Chao Liu's avatar
Chao Liu committed
58
            const index_t did0 = is / desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
59
60
61

            is -= did0 * desc.GetStride(I0);

Chao Liu's avatar
Chao Liu committed
62
            const index_t did1 = is / desc.GetStride(I1);
Chao Liu's avatar
Chao Liu committed
63
64
65

            is -= did1 * desc.GetStride(I1);

Chao Liu's avatar
Chao Liu committed
66
            const index_t did2 = is / desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
67
68
69

            is -= did2 * desc.GetStride(I2);

Chao Liu's avatar
Chao Liu committed
70
            const index_t did3 = is / desc.GetStride(I3);
Chao Liu's avatar
Chao Liu committed
71

Chao Liu's avatar
Chao Liu committed
72
            const index_t dindex = dst_desc.Get1dIndex(did0, did1, did2, did3);
Chao Liu's avatar
Chao Liu committed
73

Chao Liu's avatar
Chao Liu committed
74
            f(p_dst[dindex]);
Chao Liu's avatar
Chao Liu committed
75
76
77
        }
    }
}
Chao Liu's avatar
Chao Liu committed
78

Chao Liu's avatar
Chao Liu committed
79
// Function: p_dst[reorder[i0], reorder[i1], reorder[i2], reorder[i3]] = p_src[i0,i1,i2,i3]
80
81
// TODO: in order to optimize mem access for different mem type,
// need to write specialized version
Chao Liu's avatar
Chao Liu committed
82
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
83
          class Float,
84
85
          class SrcDesc,
          class DstDesc,
Chao Liu's avatar
Chao Liu committed
86
87
          class SrcOpLengths,
          class DstFromSrcReorder,
Chao Liu's avatar
Chao Liu committed
88
          class F>
Chao Liu's avatar
Chao Liu committed
89
90
__device__ void blockwise_4d_tensor_pointwise_operation_binary_reorder_by_get_dst_from_src(
    SrcDesc,
Chao Liu's avatar
Chao Liu committed
91
    const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
92
93
94
95
96
    DstDesc,
    Float* __restrict__ p_dst,
    SrcOpLengths,
    DstFromSrcReorder,
    F f)
Chao Liu's avatar
Chao Liu committed
97
{
Chao Liu's avatar
Chao Liu committed
98
99
100
101
    constexpr auto I0 = Number<0>{};
    constexpr auto I1 = Number<1>{};
    constexpr auto I2 = Number<2>{};
    constexpr auto I3 = Number<3>{};
Chao Liu's avatar
Chao Liu committed
102

Chao Liu's avatar
Chao Liu committed
103
104
105
106
    constexpr index_t IR0 = DstFromSrcReorder{}.Get(I0);
    constexpr index_t IR1 = DstFromSrcReorder{}.Get(I1);
    constexpr index_t IR2 = DstFromSrcReorder{}.Get(I2);
    constexpr index_t IR3 = DstFromSrcReorder{}.Get(I3);
Chao Liu's avatar
Chao Liu committed
107

108
109
    constexpr auto src_desc = SrcDesc{};
    constexpr auto dst_desc = DstDesc{};
Chao Liu's avatar
Chao Liu committed
110
    constexpr auto ref_desc = make_ConstantTensorDescriptor(SrcOpLengths{});
Chao Liu's avatar
Chao Liu committed
111

Chao Liu's avatar
Chao Liu committed
112
    constexpr index_t NLoop = ref_desc.GetElementSize() / BlockSize;
Chao Liu's avatar
Chao Liu committed
113

Chao Liu's avatar
Chao Liu committed
114
    for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
115
    {
Chao Liu's avatar
Chao Liu committed
116
        index_t is = threadIdx.x + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
117

Chao Liu's avatar
Chao Liu committed
118
        index_t did[4];
Chao Liu's avatar
Chao Liu committed
119

120
        did[0] = is / ref_desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
121

122
        is -= did[0] * ref_desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
123

124
        did[1] = is / ref_desc.GetStride(I1);
Chao Liu's avatar
Chao Liu committed
125

126
        is -= did[1] * ref_desc.GetStride(I1);
Chao Liu's avatar
Chao Liu committed
127

128
        did[2] = is / ref_desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
129

130
        is -= did[2] * ref_desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
131

132
        did[3] = is / ref_desc.GetStride(I3);
Chao Liu's avatar
Chao Liu committed
133

Chao Liu's avatar
Chao Liu committed
134
        const index_t src_index = src_desc.Get1dIndex(did[0], did[1], did[2], did[3]);
Chao Liu's avatar
Chao Liu committed
135

Chao Liu's avatar
Chao Liu committed
136
        const index_t dst_index = dst_desc.Get1dIndex(did[IR0], did[IR1], did[IR2], did[IR3]);
137

Chao Liu's avatar
Chao Liu committed
138
        f(p_src[src_index], p_dst[dst_index]);
Chao Liu's avatar
Chao Liu committed
139
140
    }

141
    constexpr bool has_tail = (ref_desc.GetElementSize() > NLoop * BlockSize);
Chao Liu's avatar
Chao Liu committed
142
143
144

    if(has_tail)
    {
Chao Liu's avatar
Chao Liu committed
145
        index_t is = threadIdx.x + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
146

147
        if(is < ref_desc.GetElementSize())
Chao Liu's avatar
Chao Liu committed
148
        {
Chao Liu's avatar
Chao Liu committed
149
            index_t did[4];
150
151

            did[0] = is / ref_desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
152

153
            is -= did[0] * ref_desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
154

155
            did[1] = is / ref_desc.GetStride(I1);
Chao Liu's avatar
Chao Liu committed
156

157
            is -= did[1] * ref_desc.GetStride(I1);
Chao Liu's avatar
Chao Liu committed
158

159
            did[2] = is / ref_desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
160

161
            is -= did[2] * ref_desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
162

163
            did[3] = is / ref_desc.GetStride(I3);
Chao Liu's avatar
Chao Liu committed
164

Chao Liu's avatar
Chao Liu committed
165
            const index_t src_index = src_desc.Get1dIndex(did[0], did[1], did[2], did[3]);
166

Chao Liu's avatar
Chao Liu committed
167
            const index_t dst_index = dst_desc.Get1dIndex(did[IR0], did[IR1], did[IR2], did[IR3]);
168

Chao Liu's avatar
Chao Liu committed
169
            f(p_src[src_index], p_dst[dst_index]);
170
171
172
173
        }
    }
}

Chao Liu's avatar
Chao Liu committed
174
template <index_t BlockSize, class Float, class DstDesc>
Chao Liu's avatar
Chao Liu committed
175
__device__ void blockwise_4d_tensor_set_zero(DstDesc, Float* __restrict__ p_dst)
176
{
Chao Liu's avatar
Chao Liu committed
177
    auto f_set_zero = [](Float& v) { v = Float(0); };
Chao Liu's avatar
Chao Liu committed
178

Chao Liu's avatar
Chao Liu committed
179
    blockwise_4d_tensor_pointwise_operation_unary<BlockSize>(DstDesc{}, p_dst, f_set_zero);
Chao Liu's avatar
Chao Liu committed
180
}
181

Chao Liu's avatar
Chao Liu committed
182
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
183
          class Float,
184
185
          class SrcDesc,
          class DstDesc,
Chao Liu's avatar
Chao Liu committed
186
187
188
189
          class SrcOpLengths,
          class DstFromSrcReorder>
__device__ void
blockwise_4d_tensor_copy_reorder_by_get_dst_from_src(SrcDesc,
Chao Liu's avatar
Chao Liu committed
190
                                                     const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
191
192
193
194
                                                     DstDesc,
                                                     Float* __restrict__ p_dst,
                                                     SrcOpLengths,
                                                     DstFromSrcReorder)
195
{
Chao Liu's avatar
Chao Liu committed
196
    auto f_copy = [](const Float& src, Float& dst) { dst = src; };
197

Chao Liu's avatar
Chao Liu committed
198
199
    blockwise_4d_tensor_pointwise_operation_binary_reorder_by_get_dst_from_src<BlockSize>(
        SrcDesc{}, p_src, DstDesc{}, p_dst, SrcOpLengths{}, DstFromSrcReorder{}, f_copy);
200
201
}

Chao Liu's avatar
Chao Liu committed
202
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
203
204
205
206
          class Float,
          class SrcDesc,
          class DstDesc,
          class CopyLengths,
Chao Liu's avatar
Chao Liu committed
207
          index_t DataPerRead>
208
struct Blockwise4dTensorCopy1
Chao Liu's avatar
Chao Liu committed
209
{
210
    using vector_t = typename vector_type<Float, DataPerRead>::MemoryType;
Chao Liu's avatar
Chao Liu committed
211

212
    __device__ constexpr Blockwise4dTensorCopy1()
Chao Liu's avatar
Chao Liu committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
    {
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

        static_assert(DataPerRead == 1 ||
                          (SrcDesc{}.GetStride(I3) == 1 && DstDesc{}.GetStride(I3) == 1),
                      "wrong! only support stride3 == 1 if DataPerRead > 1!\n");

        static_assert(DataPerRead == 1 || DataPerRead == 2 || DataPerRead == 4,
                      "wrong! only support DataPerRead == 1, 2 or 4!\n");

        static_assert(SrcDesc{}.GetStride(I2) % DataPerRead == 0 &&
                          DstDesc{}.GetStride(I2) % DataPerRead == 0,
                      "src and dst stride2 should be multiple of DataPerRead to keep alignment");

        // we allow out-of-bound read from src in D3 dimension,
        //   but we need to make sure dst stride2 is big enough,
        //   so that the out-of-bound write won't contaminate next line in dst
Chao Liu's avatar
Chao Liu committed
233
234
        constexpr index_t L3          = CopyLengths{}.Get(I3);
        constexpr index_t read_per_d3 = integer_divide_ceil(L3, DataPerRead);
Chao Liu's avatar
Chao Liu committed
235
236
237
238
239

        static_assert(read_per_d3 * DataPerRead <= DstDesc{}.GetStride(I2),
                      "wrong! out-of-bound write will contaminate next line!\n");
    }

Chao Liu's avatar
Chao Liu committed
240
    __device__ void Run(const Float* __restrict__ p_src, Float* __restrict__ p_dst) const
Chao Liu's avatar
Chao Liu committed
241
    {
Chao Liu's avatar
Chao Liu committed
242
243
244
245
246
247
248
249
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

        constexpr auto src_desc = SrcDesc{};
        constexpr auto dst_desc = DstDesc{};

Chao Liu's avatar
Chao Liu committed
250
251
252
253
        constexpr index_t L0 = CopyLengths{}.Get(I0);
        constexpr index_t L1 = CopyLengths{}.Get(I1);
        constexpr index_t L2 = CopyLengths{}.Get(I2);
        constexpr index_t L3 = CopyLengths{}.Get(I3);
Chao Liu's avatar
Chao Liu committed
254

Chao Liu's avatar
Chao Liu committed
255
        constexpr index_t read_per_d3 = integer_divide_ceil(L3, DataPerRead);
Chao Liu's avatar
Chao Liu committed
256
257
258

        constexpr auto ref_desc =
            make_ConstantTensorDescriptor(Sequence<L0, L1, L2, read_per_d3>{});
259

Chao Liu's avatar
Chao Liu committed
260
        constexpr index_t NLoop = ref_desc.GetElementSize() / BlockSize;
Chao Liu's avatar
Chao Liu committed
261

Chao Liu's avatar
Chao Liu committed
262
263
        auto f_copy = [&](index_t is) {
            index_t did[4];
Chao Liu's avatar
Chao Liu committed
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278

            did[0] = is / ref_desc.GetStride(I0);

            is -= did[0] * ref_desc.GetStride(I0);

            did[1] = is / ref_desc.GetStride(I1);

            is -= did[1] * ref_desc.GetStride(I1);

            did[2] = is / ref_desc.GetStride(I2);

            is -= did[2] * ref_desc.GetStride(I2);

            did[3] = is / ref_desc.GetStride(I3);

Chao Liu's avatar
Chao Liu committed
279
            const index_t src_index =
Chao Liu's avatar
Chao Liu committed
280
                src_desc.Get1dIndex(did[0], did[1], did[2], did[3] * DataPerRead);
Chao Liu's avatar
Chao Liu committed
281
            const index_t dst_index =
Chao Liu's avatar
Chao Liu committed
282
283
284
285
286
287
                dst_desc.Get1dIndex(did[0], did[1], did[2], did[3] * DataPerRead);

            *(reinterpret_cast<vector_t*>(p_dst + dst_index)) =
                *(reinterpret_cast<const vector_t*>(p_src + src_index));
        };

Chao Liu's avatar
Chao Liu committed
288
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
289
        {
Chao Liu's avatar
Chao Liu committed
290
            index_t is = threadIdx.x + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
291
292
293
294
295
296
297
298

            f_copy(is);
        }

        constexpr bool has_tail = (ref_desc.GetElementSize() > NLoop * BlockSize);

        if(has_tail)
        {
Chao Liu's avatar
Chao Liu committed
299
            index_t is = threadIdx.x + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
300
301
302
303
304
305

            if(is < ref_desc.GetElementSize())
            {
                f_copy(is);
            }
        }
Chao Liu's avatar
Chao Liu committed
306
307
308
    }
};

Chao Liu's avatar
Chao Liu committed
309
template <index_t BlockSize,
310
311
312
313
314
          class Float,
          class SrcDesc,
          class DstDesc,
          class DstOpLengths,
          class GlobalLowerPads>
315
struct BlockwiseChwnTensorCopyPadded
316
{
Chao Liu's avatar
Chao Liu committed
317
    __device__ void Run(const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
318
319
320
321
                        index_t c_block_data_begin,
                        index_t ho_block_data_begin,
                        index_t wo_block_data_begin,
                        index_t n_block_data_begin,
322
                        Float* __restrict__ p_dst,
Chao Liu's avatar
Chao Liu committed
323
324
325
326
                        index_t h_block_pad_low,
                        index_t w_block_pad_low,
                        index_t h_block_pad_up,
                        index_t w_block_pad_up) const
327
328
329
330
331
332
333
334
335
336
337
338
339
    {
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

        constexpr auto src_desc = SrcDesc{};
        constexpr auto dst_desc = DstDesc{};
        constexpr auto ref_desc = make_ConstantTensorDescriptor(DstOpLengths{});

        constexpr auto h_global_pad_low = GlobalLowerPads{}.Get(I0);
        constexpr auto w_global_pad_low = GlobalLowerPads{}.Get(I1);

Chao Liu's avatar
Chao Liu committed
340
        constexpr index_t NLoop = ref_desc.GetElementSize() / BlockSize;
341

Chao Liu's avatar
Chao Liu committed
342
        const Float* p_src_tmp =
Chao Liu's avatar
tidy up  
Chao Liu committed
343
344
345
346
347
            p_src +
            src_desc.Get1dIndex(c_block_data_begin,
                                (ho_block_data_begin + h_block_pad_low) - h_global_pad_low,
                                (wo_block_data_begin + w_block_pad_low) - w_global_pad_low,
                                n_block_data_begin);
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370

#if 0
        if(get_thread_local_1d_id() == 0)
        {
            print_ConstantTensorDescriptor(src_desc, "src_desc: ");
            print_ConstantTensorDescriptor(dst_desc, "dst_desc: ");
            print_ConstantTensorDescriptor(ref_desc, "ref_desc: ");

            printf("%u %u, \t"
                   "h_global_pad_low %u w_global_pad_low %u \t"
                   "h_block_pad_low %u w_block_pad_low %u h_block_pad_up %u  w_block_pad_up %u \t"
                   "\n",
                   get_block_1d_id(),
                   get_thread_local_1d_id(),
                   h_global_pad_low,
                   w_global_pad_low,
                   h_block_pad_low,
                   w_block_pad_low,
                   h_block_pad_up,
                   w_block_pad_up);
        }
#endif

Chao Liu's avatar
Chao Liu committed
371
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
372
        {
Chao Liu's avatar
Chao Liu committed
373
            index_t is = threadIdx.x + iloop * BlockSize;
374

Chao Liu's avatar
Chao Liu committed
375
            index_t did[4];
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390

            did[0] = is / ref_desc.GetStride(I0);

            is -= did[0] * ref_desc.GetStride(I0);

            did[1] = is / ref_desc.GetStride(I1);

            is -= did[1] * ref_desc.GetStride(I1);

            did[2] = is / ref_desc.GetStride(I2);

            is -= did[2] * ref_desc.GetStride(I2);

            did[3] = is / ref_desc.GetStride(I3);

Chao Liu's avatar
Chao Liu committed
391
            const index_t bindex = dst_desc.Get1dIndex(did[0], did[1], did[2], did[3]);
392
393
394
395
396
397
398
399
400
401
402
403

            p_dst[bindex] =
                (did[1] < h_block_pad_low || did[1] + h_block_pad_up >= ref_desc.GetLength(I1) ||
                 did[2] < w_block_pad_low || did[2] + w_block_pad_up >= ref_desc.GetLength(I2))
                    ? Float(0)
                    : p_src_tmp[src_desc.Get1dIndex(did[0], did[1], did[2], did[3])];
        }

        constexpr bool has_tail = (ref_desc.GetElementSize() > NLoop * BlockSize);

        if(has_tail)
        {
Chao Liu's avatar
Chao Liu committed
404
            index_t is = threadIdx.x + NLoop * BlockSize;
405
406
407

            if(is < ref_desc.GetElementSize())
            {
Chao Liu's avatar
Chao Liu committed
408
                index_t did[4];
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423

                did[0] = is / ref_desc.GetStride(I0);

                is -= did[0] * ref_desc.GetStride(I0);

                did[1] = is / ref_desc.GetStride(I1);

                is -= did[1] * ref_desc.GetStride(I1);

                did[2] = is / ref_desc.GetStride(I2);

                is -= did[2] * ref_desc.GetStride(I2);

                did[3] = is / ref_desc.GetStride(I3);

Chao Liu's avatar
Chao Liu committed
424
                const index_t bindex = dst_desc.Get1dIndex(did[0], did[1], did[2], did[3]);
425
426
427
428
429
430
431
432
433
434

                p_dst[bindex] =
                    (did[1] < h_block_pad_low ||
                     did[1] + h_block_pad_up >= ref_desc.GetLength(I1) ||
                     did[2] < w_block_pad_low || did[2] + w_block_pad_up >= ref_desc.GetLength(I2))
                        ? Float(0)
                        : p_src_tmp[src_desc.Get1dIndex(did[0], did[1], did[2], did[3])];
            }
        }
    }
Chao Liu's avatar
Chao Liu committed
435
};
Chao Liu's avatar
Chao Liu committed
436
437
438

// starting point need to be aligned to float4 or float2 or float
// stride3 need to be 1 for both source and destination
Chao Liu's avatar
Chao Liu committed
439
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
440
441
442
443
444
          class Float,
          class SrcDesc,
          class DstDesc,
          class CopyLengths,
          class ThreadPerDims,
Chao Liu's avatar
Chao Liu committed
445
          index_t DataPerRead>
Chao Liu's avatar
Chao Liu committed
446
447
struct Blockwise4dTensorCopy3
{
448
    using vector_t = typename vector_type<Float, DataPerRead>::MemoryType;
Chao Liu's avatar
Chao Liu committed
449

Chao Liu's avatar
Chao Liu committed
450
451
    index_t mSrcMyThreadOffset;
    index_t mDstMyThreadOffset;
Chao Liu's avatar
Chao Liu committed
452
453
454
455
456
457
458
459

    __device__ Blockwise4dTensorCopy3()
    {
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

Chao Liu's avatar
Chao Liu committed
460
461
462
        static_assert(DataPerRead == 1 ||
                          (SrcDesc{}.GetStride(I3) == 1 && DstDesc{}.GetStride(I3) == 1),
                      "wrong! only support stride3 == 1 if DataPerRead > 1!\n");
Chao Liu's avatar
Chao Liu committed
463
464
465
466
467
468
469

        static_assert(DataPerRead == 1 || DataPerRead == 2 || DataPerRead == 4,
                      "wrong! only support DataPerRead == 1, 2 or 4!\n");

        static_assert(
            SrcDesc{}.GetStride(I2) % DataPerRead == 0 &&
                DstDesc{}.GetStride(I2) % DataPerRead == 0,
Chao Liu's avatar
Chao Liu committed
470
            "wrong! src and dst stride2 should be multiple of DataPerRead to keep alignment");
Chao Liu's avatar
Chao Liu committed
471

Chao Liu's avatar
Chao Liu committed
472
473
474
475
        constexpr index_t L0 = CopyLengths{}.Get(I0);
        constexpr index_t L1 = CopyLengths{}.Get(I1);
        constexpr index_t L2 = CopyLengths{}.Get(I2);
        constexpr index_t L3 = CopyLengths{}.Get(I3);
Chao Liu's avatar
Chao Liu committed
476

Chao Liu's avatar
Chao Liu committed
477
478
479
480
        constexpr index_t thread_per_d0 = ThreadPerDims{}.Get(I0);
        constexpr index_t thread_per_d1 = ThreadPerDims{}.Get(I1);
        constexpr index_t thread_per_d2 = ThreadPerDims{}.Get(I2);
        constexpr index_t thread_per_d3 = ThreadPerDims{}.Get(I3);
Chao Liu's avatar
Chao Liu committed
481
482
483
484

        // we allow out-of-bound read from src in D3 dimension,
        //   but we need to make sure dst stride is big enough,
        //   so that the out-of-bound write won't contaminate next line in dst
Chao Liu's avatar
Chao Liu committed
485
        constexpr index_t nloop_d3 = integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
Chao Liu's avatar
Chao Liu committed
486
487
488
489
490
491
492
493
494
495

        static_assert(nloop_d3 * thread_per_d3 * DataPerRead <= DstDesc{}.GetStride(I2),
                      "wrong! out-of-bound write will contaminate next line!\n");

        static_assert(L0 % thread_per_d0 == 0 && L1 % thread_per_d1 == 0 && L2 % thread_per_d2 == 0,
                      "wrong! L0, L1, L2 should be divided evenly!\n");

        static_assert(BlockSize >= thread_per_d0 * thread_per_d1 * thread_per_d2 * thread_per_d3,
                      "wrrong! BlockSize is not big enough for ThreadPerDims!");

Chao Liu's avatar
Chao Liu committed
496
        constexpr index_t num_active_thread =
Chao Liu's avatar
Chao Liu committed
497
498
499
500
501
502
503
504
505
506
            thread_per_d0 * thread_per_d1 * thread_per_d2 * thread_per_d3;

        if(BlockSize > num_active_thread)
        {
            if(get_thread_local_1d_id() >= num_active_thread)
            {
                return;
            }
        }

Chao Liu's avatar
Chao Liu committed
507
        const index_t thread_id_d0 =
Chao Liu's avatar
Chao Liu committed
508
            get_thread_local_1d_id() / (thread_per_d1 * thread_per_d2 * thread_per_d3);
Chao Liu's avatar
Chao Liu committed
509
510
511
        index_t itmp = get_thread_local_1d_id() -
                       thread_id_d0 * (thread_per_d1 * thread_per_d2 * thread_per_d3);
        const index_t thread_id_d1 = itmp / (thread_per_d2 * thread_per_d3);
Chao Liu's avatar
Chao Liu committed
512
        itmp -= thread_id_d1 * (thread_per_d2 * thread_per_d3);
Chao Liu's avatar
Chao Liu committed
513
514
        const index_t thread_id_d2 = itmp / thread_per_d3;
        const index_t thread_id_d3 = itmp - thread_id_d2 * thread_per_d3;
Chao Liu's avatar
Chao Liu committed
515
516
517
518
519
520
521
522
523
524
525
526
527
528

        mSrcMyThreadOffset = SrcDesc{}.Get1dIndex(
            thread_id_d0, thread_id_d1, thread_id_d2, thread_id_d3 * DataPerRead);
        mDstMyThreadOffset = DstDesc{}.Get1dIndex(
            thread_id_d0, thread_id_d1, thread_id_d2, thread_id_d3 * DataPerRead);
    }

    __device__ void Run(const Float* __restrict__ p_src, Float* __restrict__ p_dst) const
    {
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

Chao Liu's avatar
Chao Liu committed
529
530
531
532
        constexpr index_t L0 = CopyLengths{}.Get(I0);
        constexpr index_t L1 = CopyLengths{}.Get(I1);
        constexpr index_t L2 = CopyLengths{}.Get(I2);
        constexpr index_t L3 = CopyLengths{}.Get(I3);
Chao Liu's avatar
Chao Liu committed
533

Chao Liu's avatar
Chao Liu committed
534
535
536
537
        constexpr index_t thread_per_d0 = ThreadPerDims{}.Get(I0);
        constexpr index_t thread_per_d1 = ThreadPerDims{}.Get(I1);
        constexpr index_t thread_per_d2 = ThreadPerDims{}.Get(I2);
        constexpr index_t thread_per_d3 = ThreadPerDims{}.Get(I3);
Chao Liu's avatar
Chao Liu committed
538

Chao Liu's avatar
Chao Liu committed
539
        constexpr index_t num_active_thread =
Chao Liu's avatar
Chao Liu committed
540
541
542
543
544
545
546
547
548
549
            thread_per_d0 * thread_per_d1 * thread_per_d2 * thread_per_d3;

        if(BlockSize > num_active_thread)
        {
            if(get_thread_local_1d_id() >= num_active_thread)
            {
                return;
            }
        }

Chao Liu's avatar
Chao Liu committed
550
551
552
553
        constexpr index_t nloop_d0 = L0 / thread_per_d0;
        constexpr index_t nloop_d1 = L1 / thread_per_d1;
        constexpr index_t nloop_d2 = L2 / thread_per_d2;
        constexpr index_t nloop_d3 = integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
Chao Liu's avatar
Chao Liu committed
554
555

#pragma unroll
Chao Liu's avatar
Chao Liu committed
556
        for(index_t iloop_d0 = 0; iloop_d0 < nloop_d0; ++iloop_d0)
Chao Liu's avatar
Chao Liu committed
557
558
        {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
559
            for(index_t iloop_d1 = 0; iloop_d1 < nloop_d1; ++iloop_d1)
Chao Liu's avatar
Chao Liu committed
560
561
            {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
562
                for(index_t iloop_d2 = 0; iloop_d2 < nloop_d2; ++iloop_d2)
Chao Liu's avatar
Chao Liu committed
563
564
                {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
565
                    for(index_t iloop_d3 = 0; iloop_d3 < nloop_d3; ++iloop_d3)
Chao Liu's avatar
Chao Liu committed
566
                    {
Chao Liu's avatar
Chao Liu committed
567
                        const index_t src_offset =
Chao Liu's avatar
Chao Liu committed
568
569
570
571
572
                            SrcDesc{}.Get1dIndex(iloop_d0 * thread_per_d0,
                                                 iloop_d1 * thread_per_d1,
                                                 iloop_d2 * thread_per_d2,
                                                 iloop_d3 * thread_per_d3 * DataPerRead);

Chao Liu's avatar
Chao Liu committed
573
                        const index_t dst_offset =
Chao Liu's avatar
Chao Liu committed
574
575
576
577
578
                            DstDesc{}.Get1dIndex(iloop_d0 * thread_per_d0,
                                                 iloop_d1 * thread_per_d1,
                                                 iloop_d2 * thread_per_d2,
                                                 iloop_d3 * thread_per_d3 * DataPerRead);

579
                        *(reinterpret_cast<vector_t*>(&p_dst[dst_offset + mDstMyThreadOffset])) =
Chao Liu's avatar
tidy yp  
Chao Liu committed
580
581
                            *(reinterpret_cast<const vector_t*>(
                                &p_src[src_offset + mSrcMyThreadOffset]));
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
                    }
                }
            }
        }
    }

    __device__ constexpr index_t GetRegisterClipboardSize() const
    {
        static_assert(is_same<Float, float>::value, "wrong! only support float!\n");

        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

        constexpr index_t L0 = CopyLengths{}.Get(I0);
        constexpr index_t L1 = CopyLengths{}.Get(I1);
        constexpr index_t L2 = CopyLengths{}.Get(I2);
        constexpr index_t L3 = CopyLengths{}.Get(I3);

        constexpr index_t thread_per_d0 = ThreadPerDims{}.Get(I0);
        constexpr index_t thread_per_d1 = ThreadPerDims{}.Get(I1);
        constexpr index_t thread_per_d2 = ThreadPerDims{}.Get(I2);
        constexpr index_t thread_per_d3 = ThreadPerDims{}.Get(I3);

        constexpr index_t nloop_d0 = L0 / thread_per_d0;
        constexpr index_t nloop_d1 = L1 / thread_per_d1;
        constexpr index_t nloop_d2 = L2 / thread_per_d2;
        constexpr index_t nloop_d3 = integer_divide_ceil(L3, thread_per_d3 * DataPerRead);

        return DataPerRead * nloop_d0 * nloop_d1 * nloop_d2 * nloop_d3;
    }

Chao Liu's avatar
tidy yp  
Chao Liu committed
615
616
    __device__ void RunLoadRegisterClipboard(const Float* __restrict__ p_src,
                                             Float* __restrict__ p_clipboard) const
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
    {
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

        constexpr index_t L0 = CopyLengths{}.Get(I0);
        constexpr index_t L1 = CopyLengths{}.Get(I1);
        constexpr index_t L2 = CopyLengths{}.Get(I2);
        constexpr index_t L3 = CopyLengths{}.Get(I3);

        constexpr index_t thread_per_d0 = ThreadPerDims{}.Get(I0);
        constexpr index_t thread_per_d1 = ThreadPerDims{}.Get(I1);
        constexpr index_t thread_per_d2 = ThreadPerDims{}.Get(I2);
        constexpr index_t thread_per_d3 = ThreadPerDims{}.Get(I3);

        constexpr index_t num_active_thread =
            thread_per_d0 * thread_per_d1 * thread_per_d2 * thread_per_d3;

        if(BlockSize > num_active_thread)
        {
            if(get_thread_local_1d_id() >= num_active_thread)
            {
                return;
            }
        }

        constexpr index_t nloop_d0 = L0 / thread_per_d0;
        constexpr index_t nloop_d1 = L1 / thread_per_d1;
        constexpr index_t nloop_d2 = L2 / thread_per_d2;
        constexpr index_t nloop_d3 = integer_divide_ceil(L3, thread_per_d3 * DataPerRead);

#pragma unroll
        for(index_t iloop_d0 = 0; iloop_d0 < nloop_d0; ++iloop_d0)
        {
#pragma unroll
            for(index_t iloop_d1 = 0; iloop_d1 < nloop_d1; ++iloop_d1)
            {
#pragma unroll
                for(index_t iloop_d2 = 0; iloop_d2 < nloop_d2; ++iloop_d2)
                {
#pragma unroll
                    for(index_t iloop_d3 = 0; iloop_d3 < nloop_d3; ++iloop_d3)
                    {
                        const index_t src_offset =
                            SrcDesc{}.Get1dIndex(iloop_d0 * thread_per_d0,
                                                 iloop_d1 * thread_per_d1,
                                                 iloop_d2 * thread_per_d2,
                                                 iloop_d3 * thread_per_d3 * DataPerRead);

                        const index_t dst_offset =
                            DstDesc{}.Get1dIndex(iloop_d0 * thread_per_d0,
                                                 iloop_d1 * thread_per_d1,
                                                 iloop_d2 * thread_per_d2,
                                                 iloop_d3 * thread_per_d3 * DataPerRead);

                        *(reinterpret_cast<vector_t*>(&p_clipboard[dst_offset])) =
Chao Liu's avatar
tidy yp  
Chao Liu committed
674
675
                            *(reinterpret_cast<const vector_t*>(
                                &p_src[src_offset + mSrcMyThreadOffset]));
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
                    }
                }
            }
        }
    }

    __device__ void RunStoreRegisterClipboard(const Float* __restrict__ p_clipboard,
                                              Float* __restrict__ p_dst) const
    {
        constexpr auto I0 = Number<0>{};
        constexpr auto I1 = Number<1>{};
        constexpr auto I2 = Number<2>{};
        constexpr auto I3 = Number<3>{};

        constexpr index_t L0 = CopyLengths{}.Get(I0);
        constexpr index_t L1 = CopyLengths{}.Get(I1);
        constexpr index_t L2 = CopyLengths{}.Get(I2);
        constexpr index_t L3 = CopyLengths{}.Get(I3);

        constexpr index_t thread_per_d0 = ThreadPerDims{}.Get(I0);
        constexpr index_t thread_per_d1 = ThreadPerDims{}.Get(I1);
        constexpr index_t thread_per_d2 = ThreadPerDims{}.Get(I2);
        constexpr index_t thread_per_d3 = ThreadPerDims{}.Get(I3);

        constexpr index_t num_active_thread =
            thread_per_d0 * thread_per_d1 * thread_per_d2 * thread_per_d3;

        if(BlockSize > num_active_thread)
        {
            if(get_thread_local_1d_id() >= num_active_thread)
            {
                return;
            }
        }

        constexpr index_t nloop_d0 = L0 / thread_per_d0;
        constexpr index_t nloop_d1 = L1 / thread_per_d1;
        constexpr index_t nloop_d2 = L2 / thread_per_d2;
        constexpr index_t nloop_d3 = integer_divide_ceil(L3, thread_per_d3 * DataPerRead);

#pragma unroll
        for(index_t iloop_d0 = 0; iloop_d0 < nloop_d0; ++iloop_d0)
        {
#pragma unroll
            for(index_t iloop_d1 = 0; iloop_d1 < nloop_d1; ++iloop_d1)
            {
#pragma unroll
                for(index_t iloop_d2 = 0; iloop_d2 < nloop_d2; ++iloop_d2)
                {
#pragma unroll
                    for(index_t iloop_d3 = 0; iloop_d3 < nloop_d3; ++iloop_d3)
                    {
                        const index_t src_offset =
                            SrcDesc{}.Get1dIndex(iloop_d0 * thread_per_d0,
                                                 iloop_d1 * thread_per_d1,
                                                 iloop_d2 * thread_per_d2,
                                                 iloop_d3 * thread_per_d3 * DataPerRead);

                        const index_t dst_offset =
                            DstDesc{}.Get1dIndex(iloop_d0 * thread_per_d0,
                                                 iloop_d1 * thread_per_d1,
                                                 iloop_d2 * thread_per_d2,
                                                 iloop_d3 * thread_per_d3 * DataPerRead);

                        *(reinterpret_cast<vector_t*>(&p_dst[dst_offset + mDstMyThreadOffset])) =
                            *(reinterpret_cast<const vector_t*>(&p_clipboard[src_offset]));
Chao Liu's avatar
Chao Liu committed
742
743
744
745
746
747
                    }
                }
            }
        }
    }
};