blockwise_4d_tensor_op.hip.hpp 27.5 KB
Newer Older
1
#pragma once
2
#include "ConstantTensorDescriptor.hip.hpp"
Chao Liu's avatar
Chao Liu committed
3
#include "threadwise_tensor_slice_op.hip.hpp"
4

Chao Liu's avatar
Chao Liu committed
5
template <index_t BlockSize, class Float, class DstDesc, class F>
6
__device__ void
Chao Liu's avatar
Chao Liu committed
7
blockwise_4d_tensor_pointwise_operation_unary(DstDesc, Float* __restrict__ p_dst, F f)
Chao Liu's avatar
Chao Liu committed
8
{
Chao Liu's avatar
Chao Liu committed
9
10
11
12
    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
13

14
15
    constexpr auto dst_desc = DstDesc{};

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    if(has_tail)
    {
55
        index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
56
57
58

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

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

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

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

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

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

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

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

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

Chao Liu's avatar
Chao Liu committed
80
// Function: p_dst[reorder[i0], reorder[i1], reorder[i2], reorder[i3]] = p_src[i0,i1,i2,i3]
81
82
// TODO: in order to optimize mem access for different mem type,
// need to write specialized version
Chao Liu's avatar
Chao Liu committed
83
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
84
          class Float,
85
86
          class SrcDesc,
          class DstDesc,
Chao Liu's avatar
Chao Liu committed
87
          class SrcOpLengths,
88
          class MapDst2Src,
Chao Liu's avatar
Chao Liu committed
89
          class F>
Chao Liu's avatar
Chao Liu committed
90
91
__device__ void blockwise_4d_tensor_pointwise_operation_binary_reorder_by_get_dst_from_src(
    SrcDesc,
Chao Liu's avatar
Chao Liu committed
92
    const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
93
94
95
    DstDesc,
    Float* __restrict__ p_dst,
    SrcOpLengths,
96
    MapDst2Src,
Chao Liu's avatar
Chao Liu committed
97
    F f)
Chao Liu's avatar
Chao Liu committed
98
{
Chao Liu's avatar
Chao Liu committed
99
100
101
102
    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
103

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    if(has_tail)
    {
146
        index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
147

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

213
    __device__ constexpr Blockwise4dTensorCopy1()
Chao Liu's avatar
Chao Liu committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
    {
        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
234
        constexpr index_t L3          = CopyLengths{}.Get(I3);
235
        constexpr index_t read_per_d3 = mod_conv::integer_divide_ceil(L3, DataPerRead);
Chao Liu's avatar
Chao Liu committed
236
237
238
239
240

        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
241
    __device__ void Run(const Float* __restrict__ p_src, Float* __restrict__ p_dst) const
Chao Liu's avatar
Chao Liu committed
242
    {
Chao Liu's avatar
Chao Liu committed
243
244
245
246
247
248
249
250
        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
251
252
253
254
        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
255

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

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

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

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

            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
280
            const index_t src_index =
Chao Liu's avatar
Chao Liu committed
281
                src_desc.Get1dIndex(did[0], did[1], did[2], did[3] * DataPerRead);
Chao Liu's avatar
Chao Liu committed
282
            const index_t dst_index =
Chao Liu's avatar
Chao Liu committed
283
284
285
286
287
288
                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
289
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
290
        {
291
            index_t is = get_thread_local_1d_id() + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
292
293
294
295
296
297
298
299

            f_copy(is);
        }

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

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

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

Chao Liu's avatar
Chao Liu committed
310
template <index_t BlockSize,
311
312
313
314
315
          class Float,
          class SrcDesc,
          class DstDesc,
          class DstOpLengths,
          class GlobalLowerPads>
316
struct BlockwiseChwnTensorCopyPadded
317
{
Chao Liu's avatar
Chao Liu committed
318
    __device__ void Run(const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
319
320
321
322
                        index_t c_block_data_begin,
                        index_t ho_block_data_begin,
                        index_t wo_block_data_begin,
                        index_t n_block_data_begin,
323
                        Float* __restrict__ p_dst,
Chao Liu's avatar
Chao Liu committed
324
325
326
327
                        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
328
329
330
331
332
333
334
335
336
337
338
339
340
    {
        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
341
        constexpr index_t NLoop = ref_desc.GetElementSize() / BlockSize;
342

Chao Liu's avatar
Chao Liu committed
343
        const Float* p_src_tmp =
Chao Liu's avatar
Chao Liu committed
344
345
346
347
348
            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);
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371

#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
372
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
373
        {
374
            index_t is = get_thread_local_1d_id() + iloop * BlockSize;
375

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

            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
392
            const index_t bindex = dst_desc.Get1dIndex(did[0], did[1], did[2], did[3]);
393
394
395
396
397
398
399
400
401
402
403
404

            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)
        {
405
            index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
406
407
408

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

                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
425
                const index_t bindex = dst_desc.Get1dIndex(did[0], did[1], did[2], did[3]);
426
427
428
429
430
431
432
433
434
435

                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
436
};
Chao Liu's avatar
Chao Liu committed
437
438
439

// 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
440
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
441
442
443
444
445
          class Float,
          class SrcDesc,
          class DstDesc,
          class CopyLengths,
          class ThreadPerDims,
Chao Liu's avatar
Chao Liu committed
446
          index_t DataPerRead>
Chao Liu's avatar
Chao Liu committed
447
448
struct Blockwise4dTensorCopy3
{
449
    using vector_t = typename vector_type<Float, DataPerRead>::MemoryType;
Chao Liu's avatar
Chao Liu committed
450

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

    __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
461
462
463
        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
464
465
466
467
468
469
470

        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
471
            "wrong! src and dst stride2 should be multiple of DataPerRead to keep alignment");
Chao Liu's avatar
Chao Liu committed
472

Chao Liu's avatar
Chao Liu committed
473
474
475
476
        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
477

Chao Liu's avatar
Chao Liu committed
478
479
480
481
        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
482
483
484
485

        // 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
486
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
Chao Liu's avatar
Chao Liu committed
487
488
489
490
491
492
493
494
495
496

        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
497
        constexpr index_t num_active_thread =
498
            accumulate_on_sequence(ThreadPerDims{}, mod_conv::multiplies<index_t>{}, Number<1>{});
Chao Liu's avatar
Chao Liu committed
499
500
501
502
503
504
505
506
507

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

508
509
        constexpr auto thread_cluster_desc = make_ConstantTensorDescriptor(ThreadPerDims{});
        const auto thread_multi_id = thread_cluster_desc.GetMultiIndex(get_thread_local_1d_id());
Chao Liu's avatar
Chao Liu committed
510

511
512
513
514
515
516
517
518
519
        mSrcMyThreadOffset = SrcDesc{}.Get1dIndex(thread_multi_id[0],
                                                  thread_multi_id[1],
                                                  thread_multi_id[2],
                                                  thread_multi_id[3] * DataPerRead);

        mDstMyThreadOffset = DstDesc{}.Get1dIndex(thread_multi_id[0],
                                                  thread_multi_id[1],
                                                  thread_multi_id[2],
                                                  thread_multi_id[3] * DataPerRead);
Chao Liu's avatar
Chao Liu committed
520
521
522
523
524
525
526
527
528
    }

    __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
        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;
553
        constexpr index_t nloop_d3 = mod_conv::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
                    }
                }
            }
        }
    }

    __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;
610
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
611
612
613
614

        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
    {
        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;
647
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
648

649
650
651
        constexpr auto clipboard_desc = make_ConstantTensorDescriptor(
            Sequence<nloop_d0, nloop_d1, nloop_d2, nloop_d3 * DataPerRead>{});

652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
#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);

670
671
                        const index_t clipboard_offset = clipboard_desc.Get1dIndex(
                            iloop_d0, iloop_d1, iloop_d2, iloop_d3 * DataPerRead);
672

673
                        *(reinterpret_cast<vector_t*>(&p_clipboard[clipboard_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
                    }
                }
            }
        }
    }

    __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;
714
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
715

716
717
718
        constexpr auto clipboard_desc = make_ConstantTensorDescriptor(
            Sequence<nloop_d0, nloop_d1, nloop_d2, nloop_d3 * DataPerRead>{});

719
720
721
722
723
724
725
726
727
728
729
730
#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)
                    {
731
732
                        const index_t clipboard_offset = clipboard_desc.Get1dIndex(
                            iloop_d0, iloop_d1, iloop_d2, iloop_d3 * DataPerRead);
733
734
735
736
737
738
739
740

                        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])) =
741
                            *(reinterpret_cast<const vector_t*>(&p_clipboard[clipboard_offset]));
Chao Liu's avatar
Chao Liu committed
742
743
744
745
746
747
                    }
                }
            }
        }
    }
};
748
749
750
751
752
753

template <index_t BlockSize,
          class Float,
          class SrcDesc,
          class DstDesc,
          class SrcOpLengths,
754
          class MapDst2Src>
755
756
757
758
759
760
761
struct Blockwise4dTensorCopyReorder1
{
    __device__ void Run(const Float* __restrict__ p_src, Float* __restrict__ p_dst) const
    {
        auto f_copy = [](const Float& src, Float& dst) { dst = src; };

        blockwise_4d_tensor_pointwise_operation_binary_reorder_by_get_dst_from_src<BlockSize>(
762
            SrcDesc{}, p_src, DstDesc{}, p_dst, SrcOpLengths{}, MapDst2Src{}, f_copy);
763
764
    }
};