"awq/modules/vscode:/vscode.git/clone" did not exist on "bf76e108c3ee7426ec48fc28e1e9f1d30a0501a7"
blockwise_4d_tensor_op.hpp 27.9 KB
Newer Older
1
#pragma once
Chao Liu's avatar
Chao Liu committed
2
3
#include "ConstantTensorDescriptor.hpp"
#include "threadwise_tensor_slice_op.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_packed(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

46
        const index_t dindex = dst_desc.GetOffsetFromMultiIndex(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

73
            const index_t dindex = dst_desc.GetOffsetFromMultiIndex(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_packed(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

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

137
138
        const index_t dst_index =
            dst_desc.GetOffsetFromMultiIndex(did[IR0], did[IR1], did[IR2], did[IR3]);
139

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

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

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

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

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

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

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

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

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

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

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

167
168
            const index_t src_index =
                src_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
169

170
171
            const index_t dst_index =
                dst_desc.GetOffsetFromMultiIndex(did[IR0], did[IR1], did[IR2], did[IR3]);
172

Chao Liu's avatar
Chao Liu committed
173
            f(p_src[src_index], p_dst[dst_index]);
174
175
176
177
        }
    }
}

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

Chao Liu's avatar
Chao Liu committed
183
    blockwise_4d_tensor_pointwise_operation_unary<BlockSize>(DstDesc{}, p_dst, f_set_zero);
Chao Liu's avatar
Chao Liu committed
184
}
185

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

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

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

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

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

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

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

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

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

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

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

Chao Liu's avatar
Chao Liu committed
292
        for(index_t iloop = 0; iloop < NLoop; ++iloop)
Chao Liu's avatar
Chao Liu committed
293
        {
294
            index_t is = get_thread_local_1d_id() + iloop * BlockSize;
Chao Liu's avatar
Chao Liu committed
295
296
297
298
299
300
301
302

            f_copy(is);
        }

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

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

            if(is < ref_desc.GetElementSize())
            {
                f_copy(is);
            }
        }
Chao Liu's avatar
Chao Liu committed
310
311
312
    }
};

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

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

346
347
348
349
350
351
        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);
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374

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

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

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

395
            const index_t bindex = dst_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
396
397
398
399
400

            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)
401
                    : p_src_tmp[src_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3])];
402
403
404
405
406
407
        }

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

        if(has_tail)
        {
408
            index_t is = get_thread_local_1d_id() + NLoop * BlockSize;
409
410
411

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

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

428
429
                const index_t bindex =
                    dst_desc.GetOffsetFromMultiIndex(did[0], did[1], did[2], did[3]);
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)
436
437
                        : p_src_tmp[src_desc.GetOffsetFromMultiIndex(
                              did[0], did[1], did[2], did[3])];
438
439
440
            }
        }
    }
Chao Liu's avatar
Chao Liu committed
441
};
Chao Liu's avatar
Chao Liu committed
442
443
444

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

Chao Liu's avatar
Chao Liu committed
456
457
    index_t mSrcMyThreadOffset;
    index_t mDstMyThreadOffset;
Chao Liu's avatar
Chao Liu committed
458
459
460
461
462
463
464
465

    __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
466
467
468
        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
469
470
471
472
473
474
475

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

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

Chao Liu's avatar
Chao Liu committed
483
484
485
486
        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
487
488
489
490

        // 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
491
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
Chao Liu's avatar
Chao Liu committed
492
493
494
495
496
497
498
499
500
501

        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
502
        constexpr index_t num_active_thread =
503
            accumulate_on_sequence(ThreadPerDims{}, mod_conv::multiplies<index_t>{}, Number<1>{});
Chao Liu's avatar
Chao Liu committed
504
505
506
507
508
509
510
511
512

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

Chao Liu's avatar
Chao Liu committed
513
        constexpr auto thread_cluster_desc = make_ConstantTensorDescriptor_packed(ThreadPerDims{});
514
515
        const auto thread_multi_id =
            thread_cluster_desc.GetMultiIndexFrom1dIndex(get_thread_local_1d_id());
Chao Liu's avatar
Chao Liu committed
516

517
518
519
520
        mSrcMyThreadOffset = SrcDesc{}.GetOffsetFromMultiIndex(thread_multi_id[0],
                                                               thread_multi_id[1],
                                                               thread_multi_id[2],
                                                               thread_multi_id[3] * DataPerRead);
521

522
523
524
525
        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
526
527
528
529
530
531
532
533
534
    }

    __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
535
536
537
538
        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
539

Chao Liu's avatar
Chao Liu committed
540
541
542
543
        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
544

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

#pragma unroll
Chao Liu's avatar
Chao Liu committed
562
        for(index_t iloop_d0 = 0; iloop_d0 < nloop_d0; ++iloop_d0)
Chao Liu's avatar
Chao Liu committed
563
564
        {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
565
            for(index_t iloop_d1 = 0; iloop_d1 < nloop_d1; ++iloop_d1)
Chao Liu's avatar
Chao Liu committed
566
567
            {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
568
                for(index_t iloop_d2 = 0; iloop_d2 < nloop_d2; ++iloop_d2)
Chao Liu's avatar
Chao Liu committed
569
570
                {
#pragma unroll
Chao Liu's avatar
Chao Liu committed
571
                    for(index_t iloop_d3 = 0; iloop_d3 < nloop_d3; ++iloop_d3)
Chao Liu's avatar
Chao Liu committed
572
                    {
573
574
575
576
577
578
579
580
581
582
583
                        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
584

585
                        *(reinterpret_cast<vector_t*>(&p_dst[dst_offset + mDstMyThreadOffset])) =
Chao Liu's avatar
tidy yp  
Chao Liu committed
586
587
                            *(reinterpret_cast<const vector_t*>(
                                &p_src[src_offset + mSrcMyThreadOffset]));
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
615
                    }
                }
            }
        }
    }

    __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;
616
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
617
618
619
620

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

Chao Liu's avatar
tidy yp  
Chao Liu committed
621
622
    __device__ void RunLoadRegisterClipboard(const Float* __restrict__ p_src,
                                             Float* __restrict__ p_clipboard) const
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
    {
        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;
653
        constexpr index_t nloop_d3 = mod_conv::integer_divide_ceil(L3, thread_per_d3 * DataPerRead);
654

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

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)
                    {
670
671
672
673
674
                        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);
675

676
                        const index_t clipboard_offset = clipboard_desc.GetOffsetFromMultiIndex(
677
                            iloop_d0, iloop_d1, iloop_d2, iloop_d3 * DataPerRead);
678

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

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

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

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

740
741
742
743
744
                        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);
745
746

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

template <index_t BlockSize,
          class Float,
          class SrcDesc,
          class DstDesc,
          class SrcOpLengths,
760
          class MapDst2Src>
761
762
763
764
765
766
767
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>(
768
            SrcDesc{}, p_src, DstDesc{}, p_dst, SrcOpLengths{}, MapDst2Src{}, f_copy);
769
770
    }
};