driver.cpp 27.7 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
#include <iostream>
Chao Liu's avatar
Chao Liu committed
2
3
#include <numeric>
#include <initializer_list>
Chao Liu's avatar
Chao Liu committed
4
#include <cstdlib>
Chao Liu's avatar
Chao Liu committed
5
#include <stdlib.h>
Chao Liu's avatar
Chao Liu committed
6
7
8
#include "config.hpp"
#include "ConstantTensorDescriptor.hpp"
#include "device.hpp"
Chao Liu's avatar
Chao Liu committed
9
#include "conv_common.hpp"
Chao Liu's avatar
Chao Liu committed
10
#include "device_convolution_direct_v2_nchw_kcyx_nkhw.hpp"
Chao Liu's avatar
Chao Liu committed
11
#include "device_convolution_implicit_gemm_v1_chwn_cyxk_khwn.hpp"
Chao Liu's avatar
Chao Liu committed
12
#include "device_convolution_implicit_gemm_v1_nchw_cyxk_nkhw.hpp"
Chao Liu's avatar
Chao Liu committed
13
#include "device_convolution_implicit_gemm_v2_chwn_cyxk_khwn.hpp"
Chao Liu's avatar
Chao Liu committed
14
#include "device_convolution_implicit_gemm_v3_nchw_cyxk_nkhw.hpp"
15
#include "device_convolution_implicit_gemm_v4_nchw_kcyx_nkhw.hpp"
Chao Liu's avatar
Chao Liu committed
16

17
18
using namespace ck;

Chao Liu's avatar
Chao Liu committed
19
struct GeneratorTensor_1
Chao Liu's avatar
Chao Liu committed
20
21
{
    template <class... Is>
Chao Liu's avatar
Chao Liu committed
22
    double operator()(Is... is)
Chao Liu's avatar
Chao Liu committed
23
    {
Chao Liu's avatar
Chao Liu committed
24
        return 1;
Chao Liu's avatar
Chao Liu committed
25
26
27
    }
};

Chao Liu's avatar
Chao Liu committed
28
29
30
31
32
33
34
35
36
37
38
39
struct GeneratorTensor_2
{
    int min_value = 0;
    int max_value = 1;

    template <class... Is>
    double operator()(Is...)
    {
        return (std::rand() % (max_value - min_value)) + min_value;
    }
};

40
41
42
43
44
45
46
struct GeneratorTensor_3
{
    template <class... Is>
    double operator()(Is... is)
    {
        std::array<index_t, sizeof...(Is)> dims = {{static_cast<index_t>(is)...}};

47
        auto f_acc = [](auto a, auto b) { return 100 * a + b; };
48

49
        return std::accumulate(dims.begin(), dims.end(), index_t(0), f_acc);
50
51
52
    }
};

Chao Liu's avatar
Chao Liu committed
53
54
55
56
57
struct GeneratorTensor_Checkboard
{
    template <class... Ts>
    double operator()(Ts... Xs) const
    {
Chao Liu's avatar
Chao Liu committed
58
        std::array<index_t, sizeof...(Ts)> dims = {{Xs...}};
Chao Liu's avatar
Chao Liu committed
59
60
61
        return std::accumulate(dims.begin(),
                               dims.end(),
                               true,
Chao Liu's avatar
Chao Liu committed
62
                               [](bool init, index_t x) -> int { return init != (x % 2); })
Chao Liu's avatar
Chao Liu committed
63
64
65
66
67
                   ? 1
                   : -1;
    }
};

Chao Liu's avatar
Chao Liu committed
68
69
70
71
72
73
// this is ugly, only for 4d
template <class TConstTensorDesc>
void ostream_ConstantTensorDescriptor(TConstTensorDesc, std::ostream& os = std::cout)
{
    static_assert(TConstTensorDesc::nDim == 4, "nDim is not 4");

Chao Liu's avatar
Chao Liu committed
74
75
76
77
    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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
    constexpr auto desc = TConstTensorDesc{};

    os << "Lengths: {" << desc.GetLength(I0) << ", " << desc.GetLength(I1) << ", "
       << desc.GetLength(I2) << ", " << desc.GetLength(I3) << "}, "
       << "Strides: {" << desc.GetStride(I0) << ", " << desc.GetStride(I1) << ", "
       << desc.GetStride(I2) << ", " << desc.GetStride(I3) << "}" << std::endl;
}

// this is ugly, only for 4d
template <class TConstTensorDesc>
auto make_TensorDescriptor(TConstTensorDesc)
{
    static_assert(TConstTensorDesc::nDim == 4, "nDim is not 4");

Chao Liu's avatar
Chao Liu committed
92
93
94
95
    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
96
97
    constexpr auto desc = TConstTensorDesc{};

Chao Liu's avatar
Chao Liu committed
98
    std::initializer_list<index_t> lengths = {
Chao Liu's avatar
Chao Liu committed
99
        desc.GetLength(I0), desc.GetLength(I1), desc.GetLength(I2), desc.GetLength(I3)};
Chao Liu's avatar
Chao Liu committed
100
    std::initializer_list<index_t> strides = {
Chao Liu's avatar
Chao Liu committed
101
102
103
104
105
        desc.GetStride(I0), desc.GetStride(I1), desc.GetStride(I2), desc.GetStride(I3)};

    return TensorDescriptor(lengths, strides);
}

106
107
108
109
110
111
112
template <class TIn,
          class TWei,
          class TOut,
          class ConvStrides,
          class ConvDilations,
          class LowerPads,
          class UpperPads>
113
114
115
void host_direct_convolution(const Tensor<TIn>& in_nchw,
                             const Tensor<TWei>& wei_kcyx,
                             Tensor<TOut>& out_nkhw,
116
117
                             ConvStrides,
                             ConvDilations,
118
119
                             LowerPads,
                             UpperPads)
Chao Liu's avatar
Chao Liu committed
120
{
Chao Liu's avatar
Chao Liu committed
121
122
    index_t h_pad_low = LowerPads{}.Get(Number<0>{});
    index_t w_pad_low = LowerPads{}.Get(Number<1>{});
123

Chao Liu's avatar
Chao Liu committed
124
125
    index_t h_pad_up = UpperPads{}.Get(Number<0>{});
    index_t w_pad_up = UpperPads{}.Get(Number<1>{});
126

Chao Liu's avatar
Chao Liu committed
127
128
    auto f = [&](auto n, auto k, auto ho, auto wo) {
        double v = 0;
Chao Liu's avatar
Chao Liu committed
129
        for(int c = 0; c < wei_kcyx.mDesc.GetLengths()[1]; ++c)
Chao Liu's avatar
Chao Liu committed
130
        {
Chao Liu's avatar
Chao Liu committed
131
            for(int y = 0; y < wei_kcyx.mDesc.GetLengths()[2]; ++y)
Chao Liu's avatar
Chao Liu committed
132
            {
133
                int hi = ho * ConvStrides{}[0] + y * ConvDilations{}[0] - h_pad_low;
Chao Liu's avatar
Chao Liu committed
134
                for(int x = 0; x < wei_kcyx.mDesc.GetLengths()[3]; ++x)
Chao Liu's avatar
Chao Liu committed
135
                {
136
                    int wi = wo * ConvStrides{}[1] + x * ConvDilations{}[1] - w_pad_low;
137
138
139
                    if(hi >= 0 && hi < in_nchw.mDesc.GetLengths()[2] && wi >= 0 &&
                       wi < in_nchw.mDesc.GetLengths()[3])
                    {
140
                        v += double(in_nchw(n, c, hi, wi)) * double(wei_kcyx(k, c, y, x));
141
                    }
Chao Liu's avatar
Chao Liu committed
142
143
144
                }
            }
        }
145
        out_nkhw(n, k, ho, wo) = v;
Chao Liu's avatar
Chao Liu committed
146
147
148
    };

    auto f_par = make_ParallelTensorFunctor(f,
149
150
151
152
                                            out_nkhw.mDesc.GetLengths()[0],
                                            out_nkhw.mDesc.GetLengths()[1],
                                            out_nkhw.mDesc.GetLengths()[2],
                                            out_nkhw.mDesc.GetLengths()[3]);
Chao Liu's avatar
Chao Liu committed
153

Chao Liu's avatar
Chao Liu committed
154
    f_par(std::thread::hardware_concurrency());
Chao Liu's avatar
Chao Liu committed
155
156
}

157
158
159
160
161
162
template <class TIn, class TWei, class TOut, class LowerPads, class UpperPads>
void host_winograd_3x3_convolution(const Tensor<TIn>& in_nchw,
                                   const Tensor<TWei>& wei_kcyx,
                                   Tensor<TOut>& out_nkhw,
                                   LowerPads,
                                   UpperPads)
Chao Liu's avatar
Chao Liu committed
163
{
Chao Liu's avatar
Chao Liu committed
164
165
    constexpr std::size_t HoPerTile = 2;
    constexpr std::size_t WoPerTile = 2;
Chao Liu's avatar
Chao Liu committed
166

Chao Liu's avatar
Chao Liu committed
167
168
169
170
    std::size_t N  = in_nchw.mDesc.GetLengths()[0];
    std::size_t C  = in_nchw.mDesc.GetLengths()[1];
    std::size_t HI = in_nchw.mDesc.GetLengths()[2];
    std::size_t WI = in_nchw.mDesc.GetLengths()[3];
Chao Liu's avatar
Chao Liu committed
171

Chao Liu's avatar
Chao Liu committed
172
173
174
    std::size_t K = wei_kcyx.mDesc.GetLengths()[0];
    std::size_t Y = wei_kcyx.mDesc.GetLengths()[2];
    std::size_t X = wei_kcyx.mDesc.GetLengths()[3];
Chao Liu's avatar
Chao Liu committed
175

176
177
    std::size_t HO = out_nkhw.mDesc.GetLengths()[2];
    std::size_t WO = out_nkhw.mDesc.GetLengths()[3];
Chao Liu's avatar
Chao Liu committed
178

Chao Liu's avatar
Chao Liu committed
179
180
    index_t h_pad_low = LowerPads{}.Get(Number<0>{});
    index_t w_pad_low = LowerPads{}.Get(Number<1>{});
181

Chao Liu's avatar
Chao Liu committed
182
183
    index_t h_pad_up = UpperPads{}.Get(Number<0>{});
    index_t w_pad_up = UpperPads{}.Get(Number<1>{});
184

Chao Liu's avatar
Chao Liu committed
185
186
    std::size_t HiPerTile = HoPerTile + Y - 1;
    std::size_t WiPerTile = WoPerTile + X - 1;
Chao Liu's avatar
Chao Liu committed
187

Chao Liu's avatar
Chao Liu committed
188
189
    std::size_t HTile = (HO + HoPerTile - 1) / HoPerTile;
    std::size_t WTile = (WO + WoPerTile - 1) / WoPerTile;
Chao Liu's avatar
Chao Liu committed
190

191
192
193
194
195
    Tensor<double> in_hold({N, C, HTile, WTile, HiPerTile, WiPerTile});
    Tensor<double> in_transform({N, C, HTile, WTile, HiPerTile, WiPerTile});
    Tensor<double> wei_transform({K, C, HiPerTile, WiPerTile});
    Tensor<double> out_transform({N, K, HTile, WTile, HiPerTile, HiPerTile});
    Tensor<double> out_hold({N, K, HTile, WTile, HoPerTile, WoPerTile});
Chao Liu's avatar
Chao Liu committed
196

Chao Liu's avatar
Chao Liu committed
197
198
    auto f_in_hold = [&](auto n, auto c, auto htile, auto wtile) {
        for(int j = 0; j < HiPerTile; ++j)
Chao Liu's avatar
Chao Liu committed
199
        {
Chao Liu's avatar
Chao Liu committed
200
201
            int hi = HoPerTile * htile + j - h_pad_low;
            for(int i = 0; i < WiPerTile; ++i)
Chao Liu's avatar
Chao Liu committed
202
            {
Chao Liu's avatar
Chao Liu committed
203
                int wi = WoPerTile * wtile + i - w_pad_low;
204
205
206
207

                if(hi >= 0 && hi < in_nchw.mDesc.GetLengths()[2] && wi >= 0 &&
                   wi < in_nchw.mDesc.GetLengths()[3])
                {
Chao Liu's avatar
Chao Liu committed
208
                    in_hold(n, c, htile, wtile, j, i) = in_nchw(n, c, hi, wi);
209
210
211
                }
                else
                {
212
                    in_hold(n, c, htile, wtile, j, i) = TIn(0);
213
                }
Chao Liu's avatar
Chao Liu committed
214
215
216
217
            }
        }
    };

Chao Liu's avatar
Chao Liu committed
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
    auto f_in_transform = [&](auto n, auto c, auto htile, auto wtile) {
        in_transform(n, c, htile, wtile, 0, 0) =
            in_hold(n, c, htile, wtile, 0, 0) - in_hold(n, c, htile, wtile, 0, 2) -
            in_hold(n, c, htile, wtile, 2, 0) + in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 0, 1) =
            in_hold(n, c, htile, wtile, 0, 1) + in_hold(n, c, htile, wtile, 0, 2) -
            in_hold(n, c, htile, wtile, 2, 1) - in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 0, 2) =
            -in_hold(n, c, htile, wtile, 0, 1) + in_hold(n, c, htile, wtile, 0, 2) +
            in_hold(n, c, htile, wtile, 2, 1) - in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 0, 3) =
            in_hold(n, c, htile, wtile, 0, 1) - in_hold(n, c, htile, wtile, 0, 3) -
            in_hold(n, c, htile, wtile, 2, 1) + in_hold(n, c, htile, wtile, 2, 3);

        in_transform(n, c, htile, wtile, 1, 0) =
            in_hold(n, c, htile, wtile, 1, 0) - in_hold(n, c, htile, wtile, 1, 2) +
            in_hold(n, c, htile, wtile, 2, 0) - in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 1, 1) =
            in_hold(n, c, htile, wtile, 1, 1) + in_hold(n, c, htile, wtile, 1, 2) +
            in_hold(n, c, htile, wtile, 2, 1) + in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 1, 2) =
            -in_hold(n, c, htile, wtile, 1, 1) + in_hold(n, c, htile, wtile, 1, 2) -
            in_hold(n, c, htile, wtile, 2, 1) + in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 1, 3) =
            in_hold(n, c, htile, wtile, 1, 1) - in_hold(n, c, htile, wtile, 1, 3) +
            in_hold(n, c, htile, wtile, 2, 1) - in_hold(n, c, htile, wtile, 2, 3);

        in_transform(n, c, htile, wtile, 2, 0) =
            -in_hold(n, c, htile, wtile, 1, 0) + in_hold(n, c, htile, wtile, 1, 2) +
            in_hold(n, c, htile, wtile, 2, 0) - in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 2, 1) =
            -in_hold(n, c, htile, wtile, 1, 1) - in_hold(n, c, htile, wtile, 1, 2) +
            in_hold(n, c, htile, wtile, 2, 1) + in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 2, 2) =
            in_hold(n, c, htile, wtile, 1, 1) - in_hold(n, c, htile, wtile, 1, 2) -
            in_hold(n, c, htile, wtile, 2, 1) + in_hold(n, c, htile, wtile, 2, 2);
        in_transform(n, c, htile, wtile, 2, 3) =
            -in_hold(n, c, htile, wtile, 1, 1) + in_hold(n, c, htile, wtile, 1, 3) +
            in_hold(n, c, htile, wtile, 2, 1) - in_hold(n, c, htile, wtile, 2, 3);

        in_transform(n, c, htile, wtile, 3, 0) =
            in_hold(n, c, htile, wtile, 1, 0) - in_hold(n, c, htile, wtile, 1, 2) -
            in_hold(n, c, htile, wtile, 3, 0) + in_hold(n, c, htile, wtile, 3, 2);
        in_transform(n, c, htile, wtile, 3, 1) =
            in_hold(n, c, htile, wtile, 1, 1) + in_hold(n, c, htile, wtile, 1, 2) -
            in_hold(n, c, htile, wtile, 3, 1) - in_hold(n, c, htile, wtile, 3, 2);
        in_transform(n, c, htile, wtile, 3, 2) =
            -in_hold(n, c, htile, wtile, 1, 1) + in_hold(n, c, htile, wtile, 1, 2) +
            in_hold(n, c, htile, wtile, 3, 1) - in_hold(n, c, htile, wtile, 3, 2);
        in_transform(n, c, htile, wtile, 3, 3) =
            in_hold(n, c, htile, wtile, 1, 1) - in_hold(n, c, htile, wtile, 1, 3) -
            in_hold(n, c, htile, wtile, 3, 1) + in_hold(n, c, htile, wtile, 3, 3);
Chao Liu's avatar
Chao Liu committed
270
271
272
    };

    auto f_wei_transform = [&](auto k, auto c) {
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
        wei_transform(k, c, 0, 0) = double(wei_kcyx(k, c, 0, 0));
        wei_transform(k, c, 0, 1) = 0.5 * double(wei_kcyx(k, c, 0, 0)) +
                                    0.5 * double(wei_kcyx(k, c, 0, 1)) +
                                    0.5 * double(wei_kcyx(k, c, 0, 2));
        wei_transform(k, c, 0, 2) = 0.5 * double(wei_kcyx(k, c, 0, 0)) -
                                    0.5 * double(wei_kcyx(k, c, 0, 1)) +
                                    0.5 * double(wei_kcyx(k, c, 0, 2));
        wei_transform(k, c, 0, 3) = double(wei_kcyx(k, c, 0, 2));

        wei_transform(k, c, 1, 0) = 0.5 * double(wei_kcyx(k, c, 0, 0)) +
                                    0.5 * double(wei_kcyx(k, c, 1, 0)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 0));
        wei_transform(k, c, 1, 1) =
            0.25 * double(wei_kcyx(k, c, 0, 0)) + 0.25 * double(wei_kcyx(k, c, 0, 1)) +
            0.25 * double(wei_kcyx(k, c, 0, 2)) + 0.25 * double(wei_kcyx(k, c, 1, 0)) +
            0.25 * double(wei_kcyx(k, c, 1, 1)) + 0.25 * double(wei_kcyx(k, c, 1, 2)) +
            0.25 * double(wei_kcyx(k, c, 2, 0)) + 0.25 * double(wei_kcyx(k, c, 2, 1)) +
            0.25 * double(wei_kcyx(k, c, 2, 2));
        wei_transform(k, c, 1, 2) =
            0.25 * double(wei_kcyx(k, c, 0, 0)) - 0.25 * double(wei_kcyx(k, c, 0, 1)) +
            0.25 * double(wei_kcyx(k, c, 0, 2)) + 0.25 * double(wei_kcyx(k, c, 1, 0)) -
            0.25 * double(wei_kcyx(k, c, 1, 1)) + 0.25 * double(wei_kcyx(k, c, 1, 2)) +
            0.25 * double(wei_kcyx(k, c, 2, 0)) - 0.25 * double(wei_kcyx(k, c, 2, 1)) +
            0.25 * double(wei_kcyx(k, c, 2, 2));
        wei_transform(k, c, 1, 3) = 0.5 * double(wei_kcyx(k, c, 0, 2)) +
                                    0.5 * double(wei_kcyx(k, c, 1, 2)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 2));

        wei_transform(k, c, 2, 0) = 0.5 * double(wei_kcyx(k, c, 0, 0)) -
                                    0.5 * double(wei_kcyx(k, c, 1, 0)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 0));
        wei_transform(k, c, 2, 1) =
            0.25 * double(wei_kcyx(k, c, 0, 0)) + 0.25 * double(wei_kcyx(k, c, 0, 1)) +
            0.25 * double(wei_kcyx(k, c, 0, 2)) - 0.25 * double(wei_kcyx(k, c, 1, 0)) -
            0.25 * double(wei_kcyx(k, c, 1, 1)) - 0.25 * double(wei_kcyx(k, c, 1, 2)) +
            0.25 * double(wei_kcyx(k, c, 2, 0)) + 0.25 * double(wei_kcyx(k, c, 2, 1)) +
            0.25 * double(wei_kcyx(k, c, 2, 2));
        wei_transform(k, c, 2, 2) =
            0.25 * double(wei_kcyx(k, c, 0, 0)) - 0.25 * double(wei_kcyx(k, c, 0, 1)) +
            0.25 * double(wei_kcyx(k, c, 0, 2)) - 0.25 * double(wei_kcyx(k, c, 1, 0)) +
            0.25 * double(wei_kcyx(k, c, 1, 1)) - 0.25 * double(wei_kcyx(k, c, 1, 2)) +
            0.25 * double(wei_kcyx(k, c, 2, 0)) - 0.25 * double(wei_kcyx(k, c, 2, 1)) +
            0.25 * double(wei_kcyx(k, c, 2, 2));
        wei_transform(k, c, 2, 3) = 0.5 * double(wei_kcyx(k, c, 0, 2)) -
                                    0.5 * double(wei_kcyx(k, c, 1, 2)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 2));

        wei_transform(k, c, 3, 0) = double(wei_kcyx(k, c, 2, 0));
        wei_transform(k, c, 3, 1) = 0.5 * double(wei_kcyx(k, c, 2, 0)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 1)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 2));
        wei_transform(k, c, 3, 2) = 0.5 * double(wei_kcyx(k, c, 2, 0)) -
                                    0.5 * double(wei_kcyx(k, c, 2, 1)) +
                                    0.5 * double(wei_kcyx(k, c, 2, 2));
        wei_transform(k, c, 3, 3) = double(wei_kcyx(k, c, 2, 2));
Chao Liu's avatar
Chao Liu committed
328
329
    };

Chao Liu's avatar
Chao Liu committed
330
331
    auto f_out_transform = [&](auto n, auto k, auto htile, auto wtile) {
        for(int j = 0; j < HiPerTile; ++j)
Chao Liu's avatar
Chao Liu committed
332
        {
Chao Liu's avatar
Chao Liu committed
333
            for(int i = 0; i < WiPerTile; ++i)
Chao Liu's avatar
Chao Liu committed
334
335
336
337
            {
                double v = 0;
                for(int c = 0; c < C; ++c)
                {
Chao Liu's avatar
Chao Liu committed
338
                    v += in_transform(n, c, htile, wtile, j, i) * wei_transform(k, c, j, i);
Chao Liu's avatar
Chao Liu committed
339
340
                }

Chao Liu's avatar
Chao Liu committed
341
                out_transform(n, k, htile, wtile, j, i) = v;
Chao Liu's avatar
Chao Liu committed
342
343
344
345
            }
        }
    };

Chao Liu's avatar
Chao Liu committed
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
    auto f_out_hold = [&](auto n, auto k, auto htile, auto wtile) {
        out_hold(n, k, htile, wtile, 0, 0) =
            out_transform(n, k, htile, wtile, 0, 0) + out_transform(n, k, htile, wtile, 0, 1) +
            out_transform(n, k, htile, wtile, 0, 2) + out_transform(n, k, htile, wtile, 1, 0) +
            out_transform(n, k, htile, wtile, 1, 1) + out_transform(n, k, htile, wtile, 1, 2) +
            out_transform(n, k, htile, wtile, 2, 0) + out_transform(n, k, htile, wtile, 2, 1) +
            out_transform(n, k, htile, wtile, 2, 2);
        out_hold(n, k, htile, wtile, 0, 1) =
            out_transform(n, k, htile, wtile, 0, 1) - out_transform(n, k, htile, wtile, 0, 2) -
            out_transform(n, k, htile, wtile, 0, 3) + out_transform(n, k, htile, wtile, 1, 1) -
            out_transform(n, k, htile, wtile, 1, 2) - out_transform(n, k, htile, wtile, 1, 3) +
            out_transform(n, k, htile, wtile, 2, 1) - out_transform(n, k, htile, wtile, 2, 2) -
            out_transform(n, k, htile, wtile, 2, 3);
        out_hold(n, k, htile, wtile, 1, 0) =
            out_transform(n, k, htile, wtile, 1, 0) + out_transform(n, k, htile, wtile, 1, 1) +
            out_transform(n, k, htile, wtile, 1, 2) - out_transform(n, k, htile, wtile, 2, 0) -
            out_transform(n, k, htile, wtile, 2, 1) - out_transform(n, k, htile, wtile, 2, 2) -
            out_transform(n, k, htile, wtile, 3, 0) - out_transform(n, k, htile, wtile, 3, 1) -
            out_transform(n, k, htile, wtile, 3, 2);
        out_hold(n, k, htile, wtile, 1, 1) =
            out_transform(n, k, htile, wtile, 1, 1) - out_transform(n, k, htile, wtile, 1, 2) -
            out_transform(n, k, htile, wtile, 1, 3) - out_transform(n, k, htile, wtile, 2, 1) +
            out_transform(n, k, htile, wtile, 2, 2) + out_transform(n, k, htile, wtile, 2, 3) -
            out_transform(n, k, htile, wtile, 3, 1) + out_transform(n, k, htile, wtile, 3, 2) +
            out_transform(n, k, htile, wtile, 3, 3);
Chao Liu's avatar
Chao Liu committed
371
372
    };

Chao Liu's avatar
Chao Liu committed
373
374
    auto f_out = [&](auto n, auto k, auto htile, auto wtile) {
        for(int j = 0; j < HoPerTile; ++j)
Chao Liu's avatar
Chao Liu committed
375
        {
Chao Liu's avatar
Chao Liu committed
376
377
            std::size_t ho = HoPerTile * htile + j;
            for(int i = 0; i < WoPerTile; ++i)
Chao Liu's avatar
Chao Liu committed
378
            {
379
                std::size_t wo = WoPerTile * wtile + i;
380
                out_nkhw(n, k, ho, wo) = out_hold(n, k, htile, wtile, j, i);
Chao Liu's avatar
Chao Liu committed
381
382
383
384
385
386
            }
        }
    };

    std::size_t num_thread = std::thread::hardware_concurrency();

Chao Liu's avatar
Chao Liu committed
387
388
    make_ParallelTensorFunctor(f_in_hold, N, C, HTile, WTile)(num_thread);
    make_ParallelTensorFunctor(f_in_transform, N, C, HTile, WTile)(num_thread);
Chao Liu's avatar
Chao Liu committed
389
    make_ParallelTensorFunctor(f_wei_transform, K, C)(num_thread);
Chao Liu's avatar
Chao Liu committed
390
391
392
    make_ParallelTensorFunctor(f_out_transform, N, K, HTile, WTile)(num_thread);
    make_ParallelTensorFunctor(f_out_hold, N, K, HTile, WTile)(num_thread);
    make_ParallelTensorFunctor(f_out, N, K, HTile, WTile)(num_thread);
Chao Liu's avatar
Chao Liu committed
393
394
395
396
397
398
}

template <class T>
void check_error(const Tensor<T>& ref, const Tensor<T>& result)
{
    float error     = 0;
Chao Liu's avatar
Chao Liu committed
399
    float max_diff  = -1;
Chao Liu's avatar
Chao Liu committed
400
401
402
    float ref_value = 0, result_value = 0;
    for(int i = 0; i < ref.mData.size(); ++i)
    {
403
404
        error += std::abs(double(ref.mData[i]) - double(result.mData[i]));
        float diff = std::abs(double(ref.mData[i]) - double(result.mData[i]));
Chao Liu's avatar
Chao Liu committed
405
406
407
408
409
410
411
412
413
414
415
416
        if(max_diff < diff)
        {
            max_diff     = diff;
            ref_value    = ref.mData[i];
            result_value = result.mData[i];
        }
    }

    std::cout << "error: " << error << std::endl;
    std::cout << "max_diff: " << max_diff << ", " << ref_value << ", " << result_value << std::endl;
}

Chao Liu's avatar
Chao Liu committed
417
int main(int argc, char* argv[])
Chao Liu's avatar
Chao Liu committed
418
{
Chao Liu's avatar
Chao Liu committed
419
420
#if 0
    constexpr index_t N  = 8;
Chao Liu's avatar
Chao Liu committed
421
    constexpr index_t C  = 16;
Chao Liu's avatar
Chao Liu committed
422
423
424
425
426
427
428
429
    constexpr index_t HI = 3;
    constexpr index_t WI = 18;
    constexpr index_t K  = 128;
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;

    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
430
#elif 0
431
    // 3x3, 34x34
432
    constexpr index_t N  = 128;
433
    constexpr index_t C  = 256;
434
435
    constexpr index_t HI = 34;
    constexpr index_t WI = 34;
436
437
438
    constexpr index_t K  = 128;
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;
Chao Liu's avatar
Chao Liu committed
439

440
441
442
    using ConvStrides   = Sequence<2, 2>;
    using ConvDilations = Sequence<1, 1>;

Chao Liu's avatar
Chao Liu committed
443
444
    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
445
#elif 0
446
    // 3x3, 56x56
Chao Liu's avatar
Chao Liu committed
447
448
    constexpr index_t N  = 64;
    constexpr index_t C  = 64;
449
450
    constexpr index_t HI = 56;
    constexpr index_t WI = 56;
Chao Liu's avatar
Chao Liu committed
451
452
453
    constexpr index_t K  = 128;
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;
Chao Liu's avatar
Chao Liu committed
454
455
456

    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
457
#elif 1
Chao Liu's avatar
Chao Liu committed
458
459
460
461
462
    // 3x3 filter, 28x28 image
    constexpr index_t N  = 128;
    constexpr index_t C  = 256;
    constexpr index_t HI = 28;
    constexpr index_t WI = 28;
463
    constexpr index_t K  = 128;
Chao Liu's avatar
Chao Liu committed
464
465
466
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;

467
468
469
    using ConvStrides   = Sequence<2, 2>;
    using ConvDilations = Sequence<1, 1>;

Chao Liu's avatar
Chao Liu committed
470
471
    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
472
#elif 0
Chao Liu's avatar
Chao Liu committed
473
    // 1x1 filter, 28x28 image
474
475
    constexpr index_t N  = 128;
    constexpr index_t C  = 512;
Chao Liu's avatar
Chao Liu committed
476
477
478
479
480
481
482
483
    constexpr index_t HI = 28;
    constexpr index_t WI = 28;
    constexpr index_t K  = 512;
    constexpr index_t Y  = 1;
    constexpr index_t X  = 1;

    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
484
485
#elif 0
    // 3x3 filter, 20x84 image, 1x1 padding
Chao Liu's avatar
Chao Liu committed
486
487
488
489
490
491
492
493
494
495
    constexpr index_t N  = 16;
    constexpr index_t C  = 256;
    constexpr index_t HI = 20;
    constexpr index_t WI = 84;
    constexpr index_t K  = 256;
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;

    constexpr index_t HPad = 1;
    constexpr index_t WPad = 1;
Chao Liu's avatar
Chao Liu committed
496
497
#elif 0
    // 3x3 filter, 112x112 image, 1x1 padding
Chao Liu's avatar
Chao Liu committed
498
499
500
501
502
503
504
505
506
507
    constexpr index_t N  = 16;
    constexpr index_t C  = 64;
    constexpr index_t HI = 112;
    constexpr index_t WI = 112;
    constexpr index_t K  = 128;
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;

    constexpr index_t HPad = 1;
    constexpr index_t WPad = 1;
508
#elif 0
509
510
511
512
513
514
515
516
517
518
519
    // 5x5 filter, 20x86 image
    constexpr index_t N  = 16;
    constexpr index_t C  = 256;
    constexpr index_t HI = 20;
    constexpr index_t WI = 86;
    constexpr index_t K  = 512;
    constexpr index_t Y  = 5;
    constexpr index_t X  = 5;

    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
520
521
#elif 0
    // 5x5 filter, 20x86 image, 1x1 padding
Chao Liu's avatar
Chao Liu committed
522
523
524
525
526
527
528
529
530
531
    constexpr index_t N  = 16;
    constexpr index_t C  = 256;
    constexpr index_t HI = 20;
    constexpr index_t WI = 86;
    constexpr index_t K  = 512;
    constexpr index_t Y  = 5;
    constexpr index_t X  = 5;

    constexpr index_t HPad = 1;
    constexpr index_t WPad = 1;
Chao Liu's avatar
Chao Liu committed
532
533
#elif 0
    // 5x5 filter, 28x28 image, 2x2 padding
Chao Liu's avatar
Chao Liu committed
534
535
536
537
538
539
540
541
542
543
    constexpr index_t N  = 16;
    constexpr index_t C  = 192;
    constexpr index_t HI = 28;
    constexpr index_t WI = 28;
    constexpr index_t K  = 32;
    constexpr index_t Y  = 5;
    constexpr index_t X  = 5;

    constexpr index_t HPad = 2;
    constexpr index_t WPad = 2;
Chao Liu's avatar
Chao Liu committed
544
#elif 0
545
    // 3x3 filter, 14x14 image
Chao Liu's avatar
Chao Liu committed
546
    constexpr index_t N  = 128;
547
    constexpr index_t C  = 256;
Chao Liu's avatar
Chao Liu committed
548
549
    constexpr index_t HI = 14;
    constexpr index_t WI = 14;
550
551
552
    constexpr index_t K  = 128;
    constexpr index_t Y  = 3;
    constexpr index_t X  = 3;
Chao Liu's avatar
Chao Liu committed
553
554
555

    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
556
#elif 0
557
    // 1x1 filter, 14x14 image
Chao Liu's avatar
Chao Liu committed
558
559
560
561
562
563
564
565
    constexpr index_t N  = 128;
    constexpr index_t C  = 512;
    constexpr index_t HI = 14;
    constexpr index_t WI = 14;
    constexpr index_t K  = 512;
    constexpr index_t Y  = 1;
    constexpr index_t X  = 1;

Chao Liu's avatar
Chao Liu committed
566
567
568
569
570
571
572
573
574
575
576
577
    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
#elif 0
    // 1x1 filter, 7x7 image
    constexpr index_t N  = 128;
    constexpr index_t C  = 512;
    constexpr index_t HI = 7;
    constexpr index_t WI = 7;
    constexpr index_t K  = 2048;
    constexpr index_t Y  = 1;
    constexpr index_t X  = 1;

578
579
    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
580
#elif 0
581
582
    // 1x1 filter, 73x73 image
    constexpr index_t N  = 128;
Chao Liu's avatar
Chao Liu committed
583
    constexpr index_t C  = 512;
584
585
586
587
588
589
    constexpr index_t HI = 73;
    constexpr index_t WI = 73;
    constexpr index_t K  = 128;
    constexpr index_t Y  = 1;
    constexpr index_t X  = 1;

Chao Liu's avatar
Chao Liu committed
590
591
    constexpr index_t HPad = 0;
    constexpr index_t WPad = 0;
Chao Liu's avatar
Chao Liu committed
592
#endif
Chao Liu's avatar
Chao Liu committed
593

594
595
596
    auto lower_pads = Sequence<HPad, WPad>{};
    auto upper_pads = Sequence<HPad, WPad>{};

Chao Liu's avatar
Chao Liu committed
597
598
    auto in_nchw_desc  = make_ConstantTensorDescriptor_packed(Sequence<N, C, HI, WI>{});
    auto wei_kcyx_desc = make_ConstantTensorDescriptor_packed(Sequence<K, C, Y, X>{});
599
    auto out_nkhw_desc = get_convolution_with_padding_output_default_4d_tensor_descriptor(
600
        in_nchw_desc, wei_kcyx_desc, ConvStrides{}, ConvDilations{}, lower_pads, upper_pads);
Chao Liu's avatar
Chao Liu committed
601

Chao Liu's avatar
Chao Liu committed
602
    ostream_ConstantTensorDescriptor(in_nchw_desc, std::cout << "in_nchw_desc: ");
Chao Liu's avatar
Chao Liu committed
603
    ostream_ConstantTensorDescriptor(wei_kcyx_desc, std::cout << "wei_kcyx_desc: ");
Chao Liu's avatar
Chao Liu committed
604
    ostream_ConstantTensorDescriptor(out_nkhw_desc, std::cout << "out_nkhw_desc: ");
Chao Liu's avatar
Chao Liu committed
605

Chao Liu's avatar
Chao Liu committed
606
607
    using in_data_t  = float;
    using out_data_t = float;
608
609
610
611
    Tensor<in_data_t> in_nchw(make_TensorDescriptor(in_nchw_desc));
    Tensor<in_data_t> wei_kcyx(make_TensorDescriptor(wei_kcyx_desc));
    Tensor<out_data_t> out_nkhw_host(make_TensorDescriptor(out_nkhw_desc));
    Tensor<out_data_t> out_nkhw_device(make_TensorDescriptor(out_nkhw_desc));
Chao Liu's avatar
Chao Liu committed
612

Chao Liu's avatar
Chao Liu committed
613
    std::size_t num_thread = std::thread::hardware_concurrency();
Chao Liu's avatar
Chao Liu committed
614

Chao Liu's avatar
Chao Liu committed
615
616
617
618
619
620
621
    if(argc != 3)
    {
        printf("arg1: do_verification, arg2: nrepeat\n");
        exit(1);
    }

    bool do_verification = atoi(argv[1]);
Chao Liu's avatar
Chao Liu committed
622
    index_t nrepeat      = atoi(argv[2]);
623
624
625

    if(do_verification)
    {
Chao Liu's avatar
Chao Liu committed
626
#if 0
627
        in_nchw.GenerateTensorValue(GeneratorTensor_1{}, num_thread);
Chao Liu's avatar
Chao Liu committed
628
        wei_kcyx.GenerateTensorValue(GeneratorTensor_1{}, num_thread);
Chao Liu's avatar
Chao Liu committed
629
630
631
#elif 0
        in_nchw.GenerateTensorValue(GeneratorTensor_1{}, num_thread);
        wei_kcyx.GenerateTensorValue(GeneratorTensor_2{-5, 5}, num_thread);
632
633
634
#elif 0
        in_nchw.GenerateTensorValue(GeneratorTensor_3{}, num_thread);
        wei_kcyx.GenerateTensorValue(GeneratorTensor_1{}, num_thread);
Chao Liu's avatar
Chao Liu committed
635
#elif 1
636
        in_nchw.GenerateTensorValue(GeneratorTensor_2{-5, 5}, num_thread);
Chao Liu's avatar
Chao Liu committed
637
        wei_kcyx.GenerateTensorValue(GeneratorTensor_2{-5, 5}, num_thread);
Chao Liu's avatar
Chao Liu committed
638
#elif 0
639
640
641
642
643
644
        in_nchw.GenerateTensorValue(GeneratorTensor_2{1, 5}, num_thread);

        auto gen_wei = [](auto... is) {
            return GeneratorTensor_2{1, 5}(is...) * GeneratorTensor_Checkboard{}(is...);
        };
        wei_kcyx.GenerateTensorValue(gen_wei, num_thread);
Chao Liu's avatar
Chao Liu committed
645
#endif
646
    }
Chao Liu's avatar
Chao Liu committed
647

Chao Liu's avatar
Chao Liu committed
648
#if 1
Chao Liu's avatar
Chao Liu committed
649
#if 0
Chao Liu's avatar
Chao Liu committed
650
    device_convolution_direct_v2_nchw_kcyx_nkhw
Chao Liu's avatar
Chao Liu committed
651
#elif 0
652
    device_convolution_implicit_gemm_v1_chwn_cyxk_khwn
Chao Liu's avatar
Chao Liu committed
653
#elif 0
Chao Liu's avatar
Chao Liu committed
654
    device_convolution_implicit_gemm_v1_nchw_cyxk_nkhw
655
#elif 0
Chao Liu's avatar
Chao Liu committed
656
    device_convolution_implicit_gemm_v2_chwn_cyxk_khwn
Chao Liu's avatar
Chao Liu committed
657
#elif 0
Chao Liu's avatar
Chao Liu committed
658
659
    device_convolution_implicit_gemm_v3_nchw_cyxk_nkhw
#elif 1
660
    device_convolution_implicit_gemm_v4_nchw_kcyx_nkhw
661
#endif
662
663
664
665
666
667
668
669
670
    (in_nchw_desc,
     in_nchw,
     wei_kcyx_desc,
     wei_kcyx,
     out_nkhw_desc,
     out_nkhw_device,
     ConvStrides{},
     ConvDilations{},
     nrepeat);
671

672
#elif 0
Chao Liu's avatar
Chao Liu committed
673
    device_implicit_gemm_convolution_1_chwn_cyxk_khwn_padded(in_nchw_desc,
Chao Liu's avatar
Chao Liu committed
674
                                                             in_nchw,
Chao Liu's avatar
Chao Liu committed
675
676
                                                             wei_kcyx_desc,
                                                             wei_kcyx,
Chao Liu's avatar
Chao Liu committed
677
678
679
680
681
                                                             out_nkhw_desc,
                                                             out_nkhw_device,
                                                             lower_pads,
                                                             upper_pads,
                                                             nrepeat);
682
#endif
Chao Liu's avatar
Chao Liu committed
683

684
    if(do_verification)
685
    {
Chao Liu's avatar
Chao Liu committed
686
#if 1
687
688
        if(Y == 3 && X == 3 && ConvStrides{}[0] == 1 && ConvStrides{}[1] == 1 &&
           ConvDilations{}[0] == 1 && ConvDilations{}[1] == 1)
689
        {
Chao Liu's avatar
Chao Liu committed
690
            host_winograd_3x3_convolution(in_nchw, wei_kcyx, out_nkhw_host, lower_pads, upper_pads);
691
692
        }
        else
Chao Liu's avatar
Chao Liu committed
693
#endif
694
        {
695
696
697
698
699
700
701
            host_direct_convolution(in_nchw,
                                    wei_kcyx,
                                    out_nkhw_host,
                                    ConvStrides{},
                                    ConvDilations{},
                                    lower_pads,
                                    upper_pads);
702
703
        }
        check_error(out_nkhw_host, out_nkhw_device);
Chao Liu's avatar
Chao Liu committed
704

Chao Liu's avatar
Chao Liu committed
705
#if 0
706
        LogRange(std::cout << "in_nchw : ", in_nchw.mData, ",") << std::endl;
Chao Liu's avatar
Chao Liu committed
707
        LogRange(std::cout << "wei_kcyx: ", wei_kcyx.mData, ",") << std::endl;
708
709
        LogRange(std::cout << "out_nkhw_host  : ", out_nkhw_host.mData, ",") << std::endl;
        LogRange(std::cout << "out_nkhw_device: ", out_nkhw_device.mData, ",") << std::endl;
Chao Liu's avatar
Chao Liu committed
710
#endif
711
    }
712
}