blockwise_4d_tensor_op.hpp 27.9 KB
Newer Older
1
2
3
#ifndef CK_BLOCKWISE_4D_TENSOR_OP_HPP
#define CK_BLOCKWISE_4D_TENSOR_OP_HPP

Chao Liu's avatar
Chao Liu committed
4
#include "ConstantTensorDescriptor.hpp"
5
6
7
#include "threadwise_tensor_slice_copy.hpp"

namespace ck {
8

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

18
19
    constexpr auto dst_desc = DstDesc{};

Chao Liu's avatar
Chao Liu committed
20
    constexpr auto desc = make_ConstantTensorDescriptor_packed(dst_desc.GetLengths());
Chao Liu's avatar
Chao Liu committed
21

22
#if 0
23
    if(get_thread_local_1d_id() == 0)
24
    {
Chao Liu's avatar
Chao Liu committed
25
26
        print_ConstantTensorDescriptor(dst_desc, "blockwise_4d_tensor_op_unary: dst_desc: ");
        print_ConstantTensorDescriptor(desc, "blockwise_4d_tensor_op_unary: desc: ");
27
28
29
    }
#endif

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

Chao Liu's avatar
Chao Liu committed
32
    for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
33
    {
34
        index_t is = get_thread_local_1d_id() + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
35

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

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

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

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

Chao Liu's avatar
Chao Liu committed
44
        const index_t did2 = is / desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
45
46
47

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

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

50
        const index_t dindex = dst_desc.GetOffsetFromMultiIndex(did0, did1, did2, did3);
Chao Liu's avatar
Chao Liu committed
51

Chao Liu's avatar
Chao Liu committed
52
        f(p_dst[dindex]);
Chao Liu's avatar
Chao Liu committed
53
54
55
56
57
58
    }

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

    if(has_tail)
    {
59
        index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
60
61
62

        if(is < desc.GetElementSize())
        {
Chao Liu's avatar
Chao Liu committed
63
            const index_t did0 = is / desc.GetStride(I0);
Chao Liu's avatar
Chao Liu committed
64
65
66

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

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

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

Chao Liu's avatar
Chao Liu committed
71
            const index_t did2 = is / desc.GetStride(I2);
Chao Liu's avatar
Chao Liu committed
72
73
74

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

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

77
            const index_t dindex = dst_desc.GetOffsetFromMultiIndex(did0, did1, did2, did3);
Chao Liu's avatar
Chao Liu committed
78

Chao Liu's avatar
Chao Liu committed
79
            f(p_dst[dindex]);
Chao Liu's avatar
Chao Liu committed
80
81
82
        }
    }
}
Chao Liu's avatar
Chao Liu committed
83

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

108
109
110
111
    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
112

113
114
    constexpr auto src_desc = SrcDesc{};
    constexpr auto dst_desc = DstDesc{};
Chao Liu's avatar
Chao Liu committed
115
    constexpr auto ref_desc = make_ConstantTensorDescriptor_packed(SrcOpLengths{});
Chao Liu's avatar
Chao Liu committed
116

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

Chao Liu's avatar
Chao Liu committed
119
    for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
120
    {
121
        index_t is = get_thread_local_1d_id() + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
122

Chao Liu's avatar
Chao Liu committed
123
        index_t did[4];
Chao Liu's avatar
Chao Liu committed
124

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

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

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

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

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

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

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

139
        const index_t src_index = src_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
Chao Liu's avatar
Chao Liu committed
140

141
142
        const index_t dst_index =
            dst_desc.GetOffsetFromMultiIndex(did[IR0], did[IR1], did[IR2], did[IR3]);
143

Chao Liu's avatar
Chao Liu committed
144
        f(p_src[src_index], p_dst[dst_index]);
Chao Liu's avatar
Chao Liu committed
145
146
    }

147
    constexpr bool has_tail = (ref_desc.GetElementSize() > NLoop * BlockSize);
Chao Liu's avatar
Chao Liu committed
148
149
150

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

153
        if(is < ref_desc.GetElementSize())
Chao Liu's avatar
Chao Liu committed
154
        {
Chao Liu's avatar
Chao Liu committed
155
            index_t did[4];
156
157

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

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

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

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

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

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

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

171
172
            const index_t src_index =
                src_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
173

174
175
            const index_t dst_index =
                dst_desc.GetOffsetFromMultiIndex(did[IR0], did[IR1], did[IR2], did[IR3]);
176

Chao Liu's avatar
Chao Liu committed
177
            f(p_src[src_index], p_dst[dst_index]);
178
179
180
181
        }
    }
}

Chao Liu's avatar
Chao Liu committed
182
template <index_t BlockSize, class Float, class DstDesc>
Chao Liu's avatar
Chao Liu committed
183
__device__ void blockwise_4d_tensor_set_zero(DstDesc, Float* __restrict__ p_dst)
184
{
Chao Liu's avatar
Chao Liu committed
185
    auto f_set_zero = [](Float& v) { v = Float(0); };
Chao Liu's avatar
Chao Liu committed
186

Chao Liu's avatar
Chao Liu committed
187
    blockwise_4d_tensor_pointwise_operation_unary<BlockSize>(DstDesc{}, p_dst, f_set_zero);
Chao Liu's avatar
Chao Liu committed
188
}
189

Chao Liu's avatar
Chao Liu committed
190
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
191
          class Float,
192
193
          class SrcDesc,
          class DstDesc,
Chao Liu's avatar
Chao Liu committed
194
          class SrcOpLengths,
195
          class MapDst2Src>
Chao Liu's avatar
Chao Liu committed
196
197
__device__ void
blockwise_4d_tensor_copy_reorder_by_get_dst_from_src(SrcDesc,
Chao Liu's avatar
Chao Liu committed
198
                                                     const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
199
200
201
                                                     DstDesc,
                                                     Float* __restrict__ p_dst,
                                                     SrcOpLengths,
202
                                                     MapDst2Src)
203
{
Chao Liu's avatar
Chao Liu committed
204
    auto f_copy = [](const Float& src, Float& dst) { dst = src; };
205

Chao Liu's avatar
Chao Liu committed
206
    blockwise_4d_tensor_pointwise_operation_binary_reorder_by_get_dst_from_src<BlockSize>(
207
        SrcDesc{}, p_src, DstDesc{}, p_dst, SrcOpLengths{}, MapDst2Src{}, f_copy);
208
209
}

Chao Liu's avatar
Chao Liu committed
210
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
211
212
213
214
          class Float,
          class SrcDesc,
          class DstDesc,
          class CopyLengths,
Chao Liu's avatar
Chao Liu committed
215
          index_t DataPerRead>
216
struct Blockwise4dTensorCopy1
Chao Liu's avatar
Chao Liu committed
217
{
218
    using vector_t = typename vector_type<Float, DataPerRead>::MemoryType;
Chao Liu's avatar
Chao Liu committed
219

220
    __device__ constexpr Blockwise4dTensorCopy1()
Chao Liu's avatar
Chao Liu committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
    {
        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
241
        constexpr index_t L3          = CopyLengths{}.Get(I3);
242
        constexpr index_t read_per_d3 = math::integer_divide_ceil(L3, DataPerRead);
Chao Liu's avatar
Chao Liu committed
243
244
245
246
247

        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
248
    __device__ void Run(const Float* __restrict__ p_src, Float* __restrict__ p_dst) const
Chao Liu's avatar
Chao Liu committed
249
    {
Chao Liu's avatar
Chao Liu committed
250
251
252
253
254
255
256
257
        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
258
259
260
261
        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
262

263
        constexpr index_t read_per_d3 = math::integer_divide_ceil(L3, DataPerRead);
Chao Liu's avatar
Chao Liu committed
264
265

        constexpr auto ref_desc =
Chao Liu's avatar
Chao Liu committed
266
            make_ConstantTensorDescriptor_packed(Sequence<L0, L1, L2, read_per_d3>{});
267

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

Chao Liu's avatar
Chao Liu committed
270
271
        auto f_copy = [&](index_t is) {
            index_t did[4];
Chao Liu's avatar
Chao Liu committed
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286

            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
287
            const index_t src_index =
288
                src_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3] * DataPerRead);
Chao Liu's avatar
Chao Liu committed
289
            const index_t dst_index =
290
                dst_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3] * DataPerRead);
Chao Liu's avatar
Chao Liu committed
291
292
293
294
295

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

Chao Liu's avatar
Chao Liu committed
296
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
297
        {
298
            index_t is = get_thread_local_1d_id() + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
299
300
301
302
303
304
305
306

            f_copy(is);
        }

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

        if(has_tail)
        {
307
            index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
Chao Liu's avatar
Chao Liu committed
308
309
310
311
312
313

            if(is < ref_desc.GetElementSize())
            {
                f_copy(is);
            }
        }
Chao Liu's avatar
Chao Liu committed
314
315
316
    }
};

Chao Liu's avatar
Chao Liu committed
317
template <index_t BlockSize,
318
319
320
321
322
          class Float,
          class SrcDesc,
          class DstDesc,
          class DstOpLengths,
          class GlobalLowerPads>
323
struct BlockwiseChwnTensorCopyPadded
324
{
Chao Liu's avatar
Chao Liu committed
325
    __device__ void Run(const Float* __restrict__ p_src,
Chao Liu's avatar
Chao Liu committed
326
327
328
329
                        index_t c_block_data_begin,
                        index_t ho_block_data_begin,
                        index_t wo_block_data_begin,
                        index_t n_block_data_begin,
330
                        Float* __restrict__ p_dst,
Chao Liu's avatar
Chao Liu committed
331
332
333
334
                        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
335
336
337
338
339
340
341
342
    {
        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
343
        constexpr auto ref_desc = make_ConstantTensorDescriptor_packed(DstOpLengths{});
344
345
346
347

        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
348
        constexpr index_t NLoop = ref_desc.GetElementSize() / BlockSize;
349

350
351
352
353
354
355
        const Float* p_src_tmp = p_src +
                                 src_desc.GetOffsetFromMultiIndex(
                                     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);
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378

#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
379
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
380
        {
381
            index_t is = get_thread_local_1d_id() + iloop * BlockSize;
382

Chao Liu's avatar
Chao Liu committed
383
            index_t did[4];
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398

            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);

399
            const index_t bindex = dst_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
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)
405
                    : p_src_tmp[src_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3])];
406
407
408
409
410
411
        }

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

        if(has_tail)
        {
412
            index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
413
414
415

            if(is < ref_desc.GetElementSize())
            {
Chao Liu's avatar
Chao Liu committed
416
                index_t did[4];
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431

                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);

432
433
                const index_t bindex =
                    dst_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
434
435
436
437
438
439

                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)
440
441
                        : p_src_tmp[src_desc.GetOffsetFromMultiIndex(
                              did[0], did[1], did[2], did[3])];
442
443
444
            }
        }
    }
Chao Liu's avatar
Chao Liu committed
445
};
Chao Liu's avatar
Chao Liu committed
446
447
448

// 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
449
template <index_t BlockSize,
Chao Liu's avatar
Chao Liu committed
450
451
452
453
454
          class Float,
          class SrcDesc,
          class DstDesc,
          class CopyLengths,
          class ThreadPerDims,
Chao Liu's avatar
Chao Liu committed
455
          index_t DataPerRead>
Chao Liu's avatar
Chao Liu committed
456
457
struct Blockwise4dTensorCopy3
{
458
    using vector_t = typename vector_type<Float, DataPerRead>::MemoryType;
Chao Liu's avatar
Chao Liu committed
459

Chao Liu's avatar
Chao Liu committed
460
461
    index_t mSrcMyThreadOffset;
    index_t mDstMyThreadOffset;
Chao Liu's avatar
Chao Liu committed
462
463
464
465
466
467
468
469

    __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
470
471
472
        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
473
474
475
476
477
478
479

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

Chao Liu's avatar
Chao Liu committed
482
483
484
485
        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
486

Chao Liu's avatar
Chao Liu committed
487
488
489
490
        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
491
492
493
494

        // 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
495
        constexpr index_t nloop_d3 = math::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
Chao Liu's avatar
Chao Liu committed
496
497
498
499
500
501
502
503
504
505

        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
506
        constexpr index_t num_active_thread =
507
            accumulate_on_sequence(ThreadPerDims{}, math::multiplies<index_t>{}, Number<1>{});
Chao Liu's avatar
Chao Liu committed
508
509
510
511
512
513
514
515
516

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

Chao Liu's avatar
Chao Liu committed
517
        constexpr auto thread_cluster_desc = make_ConstantTensorDescriptor_packed(ThreadPerDims{});
518
519
        const auto thread_multi_id =
            thread_cluster_desc.GetMultiIndexFrom1dIndex(get_thread_local_1d_id());
Chao Liu's avatar
Chao Liu committed
520

521
522
523
524
        mSrcMyThreadOffset = SrcDesc{}.GetOffsetFromMultiIndex(thread_multi_id[0],
                                                               thread_multi_id[1],
                                                               thread_multi_id[2],
                                                               thread_multi_id[3] * DataPerRead);
525

526
527
528
529
        mDstMyThreadOffset = DstDesc{}.GetOffsetFromMultiIndex(thread_multi_id[0],
                                                               thread_multi_id[1],
                                                               thread_multi_id[2],
                                                               thread_multi_id[3] * DataPerRead);
Chao Liu's avatar
Chao Liu committed
530
531
532
533
534
535
536
537
538
    }

    __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
539
540
541
542
        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
543

Chao Liu's avatar
Chao Liu committed
544
545
546
547
        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
548

Chao Liu's avatar
Chao Liu committed
549
        constexpr index_t num_active_thread =
Chao Liu's avatar
Chao Liu committed
550
551
552
553
554
555
556
557
558
559
            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
560
561
562
        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;
563
        constexpr index_t nloop_d3 = math::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
Chao Liu's avatar
Chao Liu committed
564
565

#pragma unroll
Chao Liu's avatar
Chao Liu committed
566
        for(index_t iloop_d0 = 0; iloop_d0 < nloop_d0; ++iloop_d0)
Chao Liu's avatar
Chao Liu committed
567
568
        {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
569
            for(index_t iloop_d1 = 0; iloop_d1 < nloop_d1; ++iloop_d1)
Chao Liu's avatar
Chao Liu committed
570
571
            {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
572
                for(index_t iloop_d2 = 0; iloop_d2 < nloop_d2; ++iloop_d2)
Chao Liu's avatar
Chao Liu committed
573
574
                {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
575
                    for(index_t iloop_d3 = 0; iloop_d3 < nloop_d3; ++iloop_d3)
Chao Liu's avatar
Chao Liu committed
576
                    {
577
578
579
580
581
582
583
584
585
586
587
                        const index_t src_offset = SrcDesc{}.GetOffsetFromMultiIndex(
                            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{}.GetOffsetFromMultiIndex(
                            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
588

589
                        *(reinterpret_cast<vector_t*>(&p_dst[dst_offset + mDstMyThreadOffset])) =
Chao Liu's avatar
tidy yp  
Chao Liu committed
590
591
                            *(reinterpret_cast<const vector_t*>(
                                &p_src[src_offset + mSrcMyThreadOffset]));
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
                    }
                }
            }
        }
    }

    __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;
620
        constexpr index_t nloop_d3 = math::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
621
622
623
624

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

Chao Liu's avatar
tidy yp  
Chao Liu committed
625
626
    __device__ void RunLoadRegisterClipboard(const Float* __restrict__ p_src,
                                             Float* __restrict__ p_clipboard) const
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
    {
        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;
657
        constexpr index_t nloop_d3 = math::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
658

Chao Liu's avatar
Chao Liu committed
659
        constexpr auto clipboard_desc = make_ConstantTensorDescriptor_packed(
660
661
            Sequence<nloop_d0, nloop_d1, nloop_d2, nloop_d3 * DataPerRead>{});

662
663
664
665
666
667
668
669
670
671
672
673
#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)
                    {
674
675
676
677
678
                        const index_t src_offset = SrcDesc{}.GetOffsetFromMultiIndex(
                            iloop_d0 * thread_per_d0,
                            iloop_d1 * thread_per_d1,
                            iloop_d2 * thread_per_d2,
                            iloop_d3 * thread_per_d3 * DataPerRead);
679

680
                        const index_t clipboard_offset = clipboard_desc.GetOffsetFromMultiIndex(
681
                            iloop_d0, iloop_d1, iloop_d2, iloop_d3 * DataPerRead);
682

683
                        *(reinterpret_cast<vector_t*>(&p_clipboard[clipboard_offset])) =
Chao Liu's avatar
tidy yp  
Chao Liu committed
684
685
                            *(reinterpret_cast<const vector_t*>(
                                &p_src[src_offset + mSrcMyThreadOffset]));
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
                    }
                }
            }
        }
    }

    __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;
724
        constexpr index_t nloop_d3 = math::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
725

Chao Liu's avatar
Chao Liu committed
726
        constexpr auto clipboard_desc = make_ConstantTensorDescriptor_packed(
727
728
            Sequence<nloop_d0, nloop_d1, nloop_d2, nloop_d3 * DataPerRead>{});

729
730
731
732
733
734
735
736
737
738
739
740
#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)
                    {
741
                        const index_t clipboard_offset = clipboard_desc.GetOffsetFromMultiIndex(
742
                            iloop_d0, iloop_d1, iloop_d2, iloop_d3 * DataPerRead);
743

744
745
746
747
748
                        const index_t dst_offset = DstDesc{}.GetOffsetFromMultiIndex(
                            iloop_d0 * thread_per_d0,
                            iloop_d1 * thread_per_d1,
                            iloop_d2 * thread_per_d2,
                            iloop_d3 * thread_per_d3 * DataPerRead);
749
750

                        *(reinterpret_cast<vector_t*>(&p_dst[dst_offset + mDstMyThreadOffset])) =
751
                            *(reinterpret_cast<const vector_t*>(&p_clipboard[clipboard_offset]));
Chao Liu's avatar
Chao Liu committed
752
753
754
755
756
757
                    }
                }
            }
        }
    }
};
758
759
760
761
762
763

template <index_t BlockSize,
          class Float,
          class SrcDesc,
          class DstDesc,
          class SrcOpLengths,
764
          class MapDst2Src>
765
766
767
768
769
770
771
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>(
772
            SrcDesc{}, p_src, DstDesc{}, p_dst, SrcOpLengths{}, MapDst2Src{}, f_copy);
773
774
    }
};
775
776
777
778

} // namespace

#endif