direct_convolution_2.cuh 9.42 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
#pragma once
#include "constant_tensor_descriptor.cuh"
#include "blockwise_tensor_op.cuh"
Chao Liu's avatar
Chao Liu committed
4
#include "blockwise_convolution.cuh"
Chao Liu's avatar
Chao Liu committed
5
6
7
8
9
10
11

template <class TFloat,
          class InDesc,
          class WeiDesc,
          class OutDesc,
          unsigned OutTileSizeH,
          unsigned OutTileSizeW,
Chao Liu's avatar
faster  
Chao Liu committed
12
13
14
          unsigned NPerBlock,
          unsigned KPerBlock,
          unsigned CPerBlock,
Chao Liu's avatar
Chao Liu committed
15
16
          unsigned YPerBlock,
          unsigned XPerBlock,
Chao Liu's avatar
Chao Liu committed
17
18
19
20
          unsigned NBlockOpLen0,
          unsigned NBlockOpLen1,
          unsigned NBlockOpLen2,
          unsigned NBlockOpLen3,
Chao Liu's avatar
Chao Liu committed
21
22
23
          unsigned BlockSize,
          unsigned GridSize>
__global__ void gridwise_convolution(InDesc,
Chao Liu's avatar
Chao Liu committed
24
                                     TFloat* const __restrict__ p_in_glb,
Chao Liu's avatar
Chao Liu committed
25
                                     WeiDesc,
Chao Liu's avatar
Chao Liu committed
26
                                     TFloat* const __restrict__ p_wei_glb,
Chao Liu's avatar
Chao Liu committed
27
                                     OutDesc,
Chao Liu's avatar
Chao Liu committed
28
                                     TFloat* __restrict__ p_out_glb)
Chao Liu's avatar
Chao Liu committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
{
    constexpr auto I0 = Index<0>{};
    constexpr auto I1 = Index<1>{};
    constexpr auto I2 = Index<2>{};
    constexpr auto I3 = Index<3>{};

    constexpr auto in_desc  = InDesc{};
    constexpr auto wei_desc = WeiDesc{};
    constexpr auto out_desc = OutDesc{};

    constexpr unsigned S = wei_desc.GetLength(I2);
    constexpr unsigned R = wei_desc.GetLength(I3);

    constexpr unsigned HoPerBlock = OutTileSizeH * YPerBlock;
    constexpr unsigned WoPerBlock = OutTileSizeW * XPerBlock;

    constexpr unsigned HiPerBlock = YPerBlock * OutTileSizeH + S - 1;
    constexpr unsigned WiPerBlock = XPerBlock * OutTileSizeW + R - 1;

    constexpr unsigned NBlockWork = (out_desc.GetLength(I0) + NPerBlock - 1) / NPerBlock;
    constexpr unsigned KBlockWork = (out_desc.GetLength(I1) + KPerBlock - 1) / KPerBlock;
    constexpr unsigned YBlockWork = (out_desc.GetLength(I2) + HoPerBlock - 1) / HoPerBlock;
    constexpr unsigned XBlockWork = (out_desc.GetLength(I3) + WoPerBlock - 1) / WoPerBlock;

    constexpr auto in_block_glb_desc = make_ConstantTensorDescriptor(
Chao Liu's avatar
faster  
Chao Liu committed
54
        Sequence<NPerBlock, CPerBlock, HiPerBlock, WiPerBlock>{}, in_desc.GetStrides());
Chao Liu's avatar
Chao Liu committed
55
56

    constexpr auto wei_block_glb_desc = make_ConstantTensorDescriptor(
Chao Liu's avatar
faster  
Chao Liu committed
57
        Sequence<KPerBlock, CPerBlock, S, R>{}, wei_desc.GetStrides());
Chao Liu's avatar
Chao Liu committed
58
59
60
61
62
63
64
65
66
67
68

    constexpr auto out_block_glb_desc = make_ConstantTensorDescriptor(
        Sequence<NPerBlock, KPerBlock, HoPerBlock, WoPerBlock>{}, out_desc.GetStrides());

    constexpr auto in_block_lds_desc =
        make_ConstantTensorDescriptor(in_block_glb_desc.GetLengths());
    constexpr auto wei_block_lds_desc =
        make_ConstantTensorDescriptor(wei_block_glb_desc.GetLengths());
    constexpr auto out_block_lds_desc =
        make_ConstantTensorDescriptor(out_block_glb_desc.GetLengths());

Chao Liu's avatar
Chao Liu committed
69
70
71
    constexpr unsigned in_block_size  = in_block_lds_desc.GetElementSpace();
    constexpr unsigned wei_block_size = wei_block_lds_desc.GetElementSpace();
    constexpr unsigned out_block_size = out_block_lds_desc.GetElementSpace();
Chao Liu's avatar
Chao Liu committed
72

Chao Liu's avatar
Chao Liu committed
73
74
75
    __shared__ TFloat p_in_block_lds[in_block_size];
    __shared__ TFloat p_wei_block_lds[wei_block_size];
    __shared__ TFloat p_out_block_lds[out_block_size];
Chao Liu's avatar
Chao Liu committed
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

    const unsigned block_id = blockIdx.x;

    unsigned itmp            = block_id;
    unsigned n_block_work_id = itmp / (KBlockWork * YBlockWork * XBlockWork);
    itmp -= n_block_work_id * (KBlockWork * YBlockWork * XBlockWork);
    unsigned k_block_work_id = itmp / (YBlockWork * XBlockWork);
    itmp -= k_block_work_id * (YBlockWork * XBlockWork);
    unsigned y_block_work_id = itmp / XBlockWork;
    unsigned x_block_work_id = itmp - y_block_work_id * XBlockWork;

    unsigned n_block_work_begin = n_block_work_id * NPerBlock;
    unsigned k_block_work_begin = k_block_work_id * KPerBlock;
    unsigned y_block_work_begin = y_block_work_id * YPerBlock;
    unsigned x_block_work_begin = x_block_work_id * XPerBlock;

    unsigned ho_block_work_begin = y_block_work_begin * OutTileSizeH;
    unsigned wo_block_work_begin = x_block_work_begin * OutTileSizeW;

    unsigned hi_block_work_begin = ho_block_work_begin; // minus padding
    unsigned wi_block_work_begin = wo_block_work_begin; // minus padding

#if 0
    if(threadIdx.x == 0)
    {
        print_ConstantTensorDescriptor( in_desc, "gridwise_convolution:  in_desc: ");
        print_ConstantTensorDescriptor(wei_desc, "gridwise_convolution: wei_desc: ");
        print_ConstantTensorDescriptor(out_desc, "gridwise_convolution: out_desc: ");
        print_ConstantTensorDescriptor( in_block_glb_desc, "gridwise_convolution:  in_block_glb_desc: ");
        print_ConstantTensorDescriptor(wei_block_glb_desc, "gridwise_convolution: wei_block_glb_desc: ");
        print_ConstantTensorDescriptor(out_block_glb_desc, "gridwise_convolution: out_block_glb_desc: ");
        print_ConstantTensorDescriptor( in_block_lds_desc, "gridwise_convolution:  in_block_lds_desc: ");
        print_ConstantTensorDescriptor(wei_block_lds_desc, "gridwise_convolution: wei_block_lds_desc: ");
        print_ConstantTensorDescriptor(out_block_lds_desc, "gridwise_convolution: out_block_lds_desc: ");

        printf("NBlockWork %u, KBlockWork %u, YBlockWork %u, XBlockWork %u \t"
               "block_id %u, n_block_work_id %u, k_block_work_id %u, y_block_work_id %u, "
               "x_block_work_id %u\n",
               NBlockWork,
               KBlockWork,
               YBlockWork,
               XBlockWork,
               block_id,
               n_block_work_id,
               k_block_work_id,
               y_block_work_id,
               x_block_work_id);
    }
#endif

Chao Liu's avatar
Chao Liu committed
126
    auto f_set0 = [](TFloat& v) { v = TFloat(0); };
Chao Liu's avatar
faster  
Chao Liu committed
127
128
129
    auto f_copy = [](const TFloat& src, TFloat& dst) { dst = src; };
    auto f_accu = [](const TFloat& src, TFloat& dst) { dst += src; };

Chao Liu's avatar
Chao Liu committed
130
131
132
    // set output tensor in LDS to 0
    blockwise_4d_tensor_op_unary<TFloat,
                                 decltype(out_block_lds_desc),
Chao Liu's avatar
Chao Liu committed
133
134
135
136
                                 NBlockOpLen0,
                                 NBlockOpLen1,
                                 NBlockOpLen2,
                                 NBlockOpLen3,
Chao Liu's avatar
Chao Liu committed
137
                                 decltype(f_set0),
Chao Liu's avatar
Chao Liu committed
138
                                 BlockSize>(out_block_lds_desc, p_out_block_lds, f_set0);
Chao Liu's avatar
faster  
Chao Liu committed
139

Chao Liu's avatar
Chao Liu committed
140
    for(unsigned c_block_work_begin = 0; c_block_work_begin < in_desc.GetLength(I1);
Chao Liu's avatar
faster  
Chao Liu committed
141
        c_block_work_begin += CPerBlock)
Chao Liu's avatar
Chao Liu committed
142
143
144
    {

        // copy input tensor to LDS
Chao Liu's avatar
Chao Liu committed
145
146
147
        blockwise_4d_tensor_op_binary<TFloat,
                                      decltype(in_block_glb_desc),
                                      decltype(in_block_lds_desc),
Chao Liu's avatar
Chao Liu committed
148
149
150
151
                                      NBlockOpLen0,
                                      NBlockOpLen1,
                                      NBlockOpLen2,
                                      NBlockOpLen3,
Chao Liu's avatar
Chao Liu committed
152
153
154
                                      decltype(f_copy),
                                      BlockSize>(
            in_block_glb_desc,
Chao Liu's avatar
Chao Liu committed
155
156
157
158
            p_in_glb + in_block_glb_desc.Get1dIndex(n_block_work_begin,
                                                    c_block_work_begin,
                                                    hi_block_work_begin,
                                                    wi_block_work_begin),
Chao Liu's avatar
Chao Liu committed
159
            in_block_lds_desc,
Chao Liu's avatar
Chao Liu committed
160
            p_in_block_lds,
Chao Liu's avatar
Chao Liu committed
161
            f_copy);
Chao Liu's avatar
Chao Liu committed
162
163

        // copy weight tensor to LDS
Chao Liu's avatar
Chao Liu committed
164
165
166
        blockwise_4d_tensor_op_binary<TFloat,
                                      decltype(wei_block_glb_desc),
                                      decltype(wei_block_lds_desc),
Chao Liu's avatar
Chao Liu committed
167
168
169
170
                                      NBlockOpLen0,
                                      NBlockOpLen1,
                                      NBlockOpLen2,
                                      NBlockOpLen3,
Chao Liu's avatar
Chao Liu committed
171
172
                                      decltype(f_copy),
                                      BlockSize>(
Chao Liu's avatar
Chao Liu committed
173
            wei_block_glb_desc,
Chao Liu's avatar
Chao Liu committed
174
            p_wei_glb + wei_block_glb_desc.Get1dIndex(k_block_work_begin, c_block_work_begin, 0, 0),
Chao Liu's avatar
Chao Liu committed
175
            wei_block_lds_desc,
Chao Liu's avatar
Chao Liu committed
176
            p_wei_block_lds,
Chao Liu's avatar
Chao Liu committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
            f_copy);

#if 1
        __syncthreads();
#endif

        // blockwise convolution
        blockwise_convolution<TFloat,
                              decltype(in_block_lds_desc),
                              decltype(wei_block_lds_desc),
                              decltype(out_block_lds_desc),
                              OutTileSizeH,
                              OutTileSizeW,
                              BlockSize>(in_block_lds_desc,
Chao Liu's avatar
Chao Liu committed
191
                                         p_in_block_lds,
Chao Liu's avatar
Chao Liu committed
192
                                         wei_block_lds_desc,
Chao Liu's avatar
Chao Liu committed
193
                                         p_wei_block_lds,
Chao Liu's avatar
Chao Liu committed
194
                                         out_block_lds_desc,
Chao Liu's avatar
Chao Liu committed
195
                                         p_out_block_lds);
Chao Liu's avatar
Chao Liu committed
196
197
198
199

#if 1
        __syncthreads();
#endif
Chao Liu's avatar
faster  
Chao Liu committed
200
    }
Chao Liu's avatar
Chao Liu committed
201

Chao Liu's avatar
faster  
Chao Liu committed
202
    // copy output tensor from LDS to device mem
Chao Liu's avatar
Chao Liu committed
203
204
205
    blockwise_4d_tensor_op_binary<TFloat,
                                  decltype(out_block_lds_desc),
                                  decltype(out_block_glb_desc),
Chao Liu's avatar
Chao Liu committed
206
207
208
209
                                  NBlockOpLen0,
                                  NBlockOpLen1,
                                  NBlockOpLen2,
                                  NBlockOpLen3,
Chao Liu's avatar
Chao Liu committed
210
211
212
                                  decltype(f_copy),
                                  BlockSize>(
        out_block_lds_desc,
Chao Liu's avatar
Chao Liu committed
213
        p_out_block_lds,
Chao Liu's avatar
Chao Liu committed
214
        out_block_glb_desc,
Chao Liu's avatar
Chao Liu committed
215
        p_out_glb +
Chao Liu's avatar
Chao Liu committed
216
217
218
            out_block_glb_desc.Get1dIndex(
                n_block_work_begin, k_block_work_begin, ho_block_work_begin, wo_block_work_begin),
        f_copy);
Chao Liu's avatar
Chao Liu committed
219
}