pool2d_fwd.cpp 12.5 KB
Newer Older
1
2
3
4
5
#include <iostream>
#include <numeric>
#include <initializer_list>
#include <cstdlib>
#include <stdlib.h>
6
7

#include "check_err.hpp"
8
9
10
11
12
13
14
15
16
#include "config.hpp"
#include "print.hpp"
#include "device.hpp"
#include "host_tensor.hpp"
#include "host_tensor_generator.hpp"
#include "host_reduce_util.hpp"
#include "device_tensor.hpp"
#include "tensor_layout.hpp"
#include "reduction_operator.hpp"
Chao Liu's avatar
Chao Liu committed
17
#include "device_pool2d_fwd_nhwc_nhwc.hpp"
18
19
20
21
22

using InDataType  = ck::half_t;
using OutDataType = ck::half_t;
using AccDataType = float;

23
24
using IndexDataType = int32_t;

25
26
27
28
using InLayout  = ck::tensor_layout::convolution::NHWC;
using OutLayout = ck::tensor_layout::convolution::NHWC;

#if 1
29
static constexpr auto ReduceOpId = ck::ReduceTensorOp::MAX;
30
#else
31
static constexpr auto ReduceOpId = ck::ReduceTensorOp::AVG;
32
33
#endif

34
static constexpr bool OutputIndex  = false;
35
36
37
38
39
40
41
42
static constexpr bool PropagateNan = false;

using DevicePoolFwdInstance =
    ck::tensor_operation::device::DevicePool2dFwd_Input_N_Hi_Wi_C_Output_N_Ho_Wo_C<
        InDataType,  // InDataType
        OutDataType, // OutDataType
        AccDataType, // AccDataType
        ReduceOpId,
43
        OutputIndex,
44
45
46
47
48
49
50
51
52
53
        64, // BlockSize
        64, // ReduceMThreadClusterSize
        1,  // ReduceKThreadClusterSize
        4,  // ReduceMThreadSliceSize
        1,  // ReduceKThreadSliceSize
        4>; // InSrcOutDstVectorSize

template <typename InDataType,
          typename OutDataType,
          typename AccDataType,
54
          ck::ReduceTensorOp ReduceOpId,
55
          bool PropagateNan,
56
          bool OutputIndex>
57
58
static void pool_host_verify(const Tensor<InDataType>& in,
                             Tensor<OutDataType>& out,
59
                             Tensor<IndexDataType>& out_indices,
60
61
62
63
64
65
66
                             const std::array<ck::index_t, 2>& window_spatial_lengths,
                             const std::array<ck::index_t, 2>& window_strides,
                             const std::array<ck::index_t, 2>& in_left_pads,
                             const std::array<ck::index_t, 2>& /*in_right_pads*/)
{
    using namespace ck::host_reduce;

67
    const int32_t divider = window_spatial_lengths[0] * window_spatial_lengths[1];
68
69
70
71

    const auto PreUnaryOp = PreUnaryOpFn<AccDataType, ReduceOpId>(divider);
    const auto PosUnaryOp = PosUnaryOpFn<AccDataType, ReduceOpId>(divider);

72
    if constexpr(!OutputIndex)
73
74
75
76
77
78
    {
        auto opReduce = ReduceOpFn<AccDataType, ReduceOpId>();

        auto f_nchw = [&](auto n, auto c, auto ho, auto wo) {
            auto accuVal = ReduceOpZeroVal<AccDataType, ReduceOpId>();

79
            for(ck::index_t y = 0; y < window_spatial_lengths[0]; ++y)
80
            {
81
82
                ck::index_t hi = ho * window_strides[0] + y - in_left_pads[0];
                for(ck::index_t x = 0; x < window_spatial_lengths[1]; ++x)
83
                {
84
85
86
                    ck::index_t wi = wo * window_strides[1] + x - in_left_pads[1];
                    if(hi >= 0 && hi < static_cast<ck::index_t>(in.mDesc.GetLengths()[2]) &&
                       wi >= 0 && wi < static_cast<ck::index_t>(in.mDesc.GetLengths()[3]))
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
                    {
                        AccDataType currVal = static_cast<AccDataType>(in(n, c, hi, wi));

                        PreUnaryOp(currVal);

                        binop_with_nan_check<AccDataType, PropagateNan>(opReduce, accuVal, currVal);
                    }
                }
            }

            PosUnaryOp(accuVal);

            out(n, c, ho, wo) = accuVal;
        };

        make_ParallelTensorFunctor(f_nchw,
                                   out.mDesc.GetLengths()[0],
                                   out.mDesc.GetLengths()[1],
                                   out.mDesc.GetLengths()[2],
                                   out.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
    }
    else
    {
        auto opReduce = ReduceOpFn2<AccDataType, ReduceOpId>();

        auto f_nchw = [&](auto n, auto c, auto ho, auto wo) {
113
114
            auto accuVal            = ReduceOpZeroVal<AccDataType, ReduceOpId>();
            IndexDataType accuIndex = 0;
115

116
            for(ck::index_t y = 0; y < window_spatial_lengths[0]; ++y)
117
            {
118
119
                ck::index_t hi = ho * window_strides[0] + y - in_left_pads[0];
                for(ck::index_t x = 0; x < window_spatial_lengths[1]; ++x)
120
                {
121
                    ck::index_t wi = wo * window_strides[1] + x - in_left_pads[1];
122
123
124
                    if(hi >= 0 && hi < in.mDesc.GetLengths()[2] && wi >= 0 &&
                       wi < in.mDesc.GetLengths()[3])
                    {
125
126
                        AccDataType currVal     = static_cast<AccDataType>(in(n, c, hi, wi));
                        IndexDataType currIndex = y * window_spatial_lengths[1] + x;
127
128
129

                        PreUnaryOp(currVal);

130
                        binop_with_index_and_nan_check<AccDataType, IndexDataType, PropagateNan>(
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
                            opReduce, accuVal, currVal, accuIndex, currIndex);
                    }
                }
            }

            PosUnaryOp(accuVal);

            out(n, c, ho, wo)         = accuVal;
            out_indices(n, c, ho, wo) = accuIndex;
        };

        make_ParallelTensorFunctor(f_nchw,
                                   out.mDesc.GetLengths()[0],
                                   out.mDesc.GetLengths()[1],
                                   out.mDesc.GetLengths()[2],
                                   out.mDesc.GetLengths()[3])(std::thread::hardware_concurrency());
    };
}

int main(int argc, char* argv[])
{
    using namespace ck::host_reduce;

154
155
156
    bool do_verification;
    int init_method;
    bool time_kernel;
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171

    // Pool shape
    ck::index_t N               = 128;
    ck::index_t C               = 192;
    ck::index_t Y               = 3;
    ck::index_t X               = 3;
    ck::index_t Hi              = 71;
    ck::index_t Wi              = 71;
    ck::index_t window_stride_h = 2;
    ck::index_t window_stride_w = 2;
    ck::index_t in_left_pad_h   = 1;
    ck::index_t in_left_pad_w   = 1;
    ck::index_t in_right_pad_h  = 1;
    ck::index_t in_right_pad_w  = 1;

172
173
174
175
176
177
178
    if(argc == 1)
    {
        do_verification = true;
        init_method     = 1;
        time_kernel     = true;
    }
    else if(argc == 4)
179
180
181
    {
        do_verification = std::stoi(argv[1]);
        init_method     = std::stoi(argv[2]);
182
        time_kernel     = static_cast<bool>(std::stoi(argv[3]));
183
184
185
186
187
    }
    else if(argc == 16)
    {
        do_verification = std::stoi(argv[1]);
        init_method     = std::stoi(argv[2]);
188
        time_kernel     = static_cast<bool>(std::stoi(argv[3]));
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206

        N               = std::stoi(argv[4]);
        C               = std::stoi(argv[5]);
        Y               = std::stoi(argv[6]);
        X               = std::stoi(argv[7]);
        Hi              = std::stoi(argv[8]);
        Wi              = std::stoi(argv[9]);
        window_stride_h = std::stoi(argv[10]);
        window_stride_w = std::stoi(argv[11]);
        in_left_pad_h   = std::stoi(argv[12]);
        in_left_pad_w   = std::stoi(argv[13]);
        in_right_pad_h  = std::stoi(argv[14]);
        in_right_pad_w  = std::stoi(argv[15]);
    }
    else
    {
        printf("arg1: verification (0=no, 1=yes)\n");
        printf("arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n");
207
        printf("arg3: time kernel (0=no, 1=yes)\n");
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
        printf("arg4 to 15: N, C, Y, X, Hi, Wi, Sy, Sx, LeftPy, LeftPx, RightPy, "
               "RightPx\n");
        exit(0);
    }

    const ck::index_t Ho = (Hi + in_left_pad_h + in_right_pad_h - Y) / window_stride_h + 1;
    const ck::index_t Wo = (Wi + in_left_pad_w + in_right_pad_w - X) / window_stride_w + 1;

    const std::array<ck::index_t, 2> window_spatial_lengths{{Y, X}};
    const std::array<ck::index_t, 2> window_strides{{window_stride_h, window_stride_w}};
    const std::array<ck::index_t, 2> input_left_pads{{in_left_pad_h, in_left_pad_w}};
    const std::array<ck::index_t, 2> input_right_pads{{in_right_pad_h, in_right_pad_w}};

    // tensor layout
    auto f_host_tensor_descriptor =
        [](std::size_t N_, std::size_t C_, std::size_t H, std::size_t W, auto layout) {
            if constexpr(ck::is_same<decltype(layout), ck::tensor_layout::convolution::NCHW>::value)
            {
                return HostTensorDescriptor(std::vector<std::size_t>({N_, C_, H, W}),
                                            std::vector<std::size_t>({C_ * H * W, H * W, W, 1}));
            }
            else if constexpr(ck::is_same<decltype(layout),
                                          ck::tensor_layout::convolution::NHWC>::value)
            {
                return HostTensorDescriptor(std::vector<std::size_t>({N_, C_, H, W}),
                                            std::vector<std::size_t>({C_ * H * W, 1, W * C_, C_}));
            }
        };

    Tensor<InDataType> in_n_c_hi_wi(f_host_tensor_descriptor(N, C, Hi, Wi, InLayout{}));
    Tensor<OutDataType> out_n_c_ho_wo_host(f_host_tensor_descriptor(N, C, Ho, Wo, OutLayout{}));
239
240
    Tensor<IndexDataType> out_indices_n_c_ho_wo_host(
        f_host_tensor_descriptor(N, C, Ho, Wo, OutLayout{}));
241
    Tensor<OutDataType> out_n_c_ho_wo_device(f_host_tensor_descriptor(N, C, Ho, Wo, OutLayout{}));
242
243
    Tensor<IndexDataType> out_indices_n_c_ho_wo_device(
        f_host_tensor_descriptor(N, C, Ho, Wo, OutLayout{}));
244
245
246
247
248
249
250

    std::cout << "in_n_c_hi_wi: " << in_n_c_hi_wi.mDesc << std::endl;
    std::cout << "out_n_c_ho_wo: " << out_n_c_ho_wo_host.mDesc << std::endl;

    switch(init_method)
    {
    case 0: break;
251
252
253
    case 1: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_1<InDataType>{1}); break;
    case 2: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5}); break;
    default: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_3<InDataType>{-5.0, 5.0});
254
255
256
257
    }

    DeviceMem in_device_buf(sizeof(InDataType) * in_n_c_hi_wi.mDesc.GetElementSpace());
    DeviceMem out_device_buf(sizeof(OutDataType) * out_n_c_ho_wo_device.mDesc.GetElementSpace());
258
    DeviceMem out_indices_device_buf(sizeof(IndexDataType) *
259
260
261
262
                                     out_indices_n_c_ho_wo_device.mDesc.GetElementSpace());

    in_device_buf.ToDevice(in_n_c_hi_wi.mData.data());

263
264
265
266
267
268
269
270
271
272
273
274
275
276
    auto pool         = DevicePoolFwdInstance{};
    auto invoker_ptr  = pool.MakeInvokerPointer();
    auto argument_ptr = pool.MakeArgumentPointer(
        static_cast<InDataType*>(in_device_buf.GetDeviceBuffer()),
        static_cast<OutDataType*>(out_device_buf.GetDeviceBuffer()),
        static_cast<IndexDataType*>(out_indices_device_buf.GetDeviceBuffer()),
        N,
        C,
        std::array<ck::index_t, 2>{{Hi, Wi}},
        std::array<ck::index_t, 2>{{Y, X}},
        std::array<ck::index_t, 2>{{Ho, Wo}},
        window_strides,
        input_left_pads,
        input_right_pads);
277
278
279
280
281
282
283

    if(!pool.IsSupportedArgument(argument_ptr.get()))
    {
        throw std::runtime_error("wrong! device_op with the specified compilation parameters does "
                                 "not support this problem");
    }

JD's avatar
JD committed
284
    float ave_time = invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, time_kernel});
285
286
287
288
289
290
291
292
293
294
295
296
297

    std::size_t flop = std::size_t(2) * N * C * Ho * Wo * Y * X;

    std::size_t num_btype =
        sizeof(InDataType) * (N * C * Hi * Wi) + sizeof(OutDataType) * (N * C * Ho * Wo);

    float tflops = static_cast<float>(flop) / 1.E9 / ave_time;

    float gb_per_sec = num_btype / 1.E6 / ave_time;

    std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s"
              << std::endl;

Anthony Chang's avatar
Anthony Chang committed
298
    bool pass = true;
299

300
301
302
303
304
305
306
    if(do_verification)
    {
        pool_host_verify<InDataType,
                         OutDataType,
                         AccDataType,
                         ReduceOpId,
                         PropagateNan,
307
                         OutputIndex>(in_n_c_hi_wi,
308
309
310
311
312
313
314
315
316
                                      out_n_c_ho_wo_host,
                                      out_indices_n_c_ho_wo_host,
                                      window_spatial_lengths,
                                      window_strides,
                                      input_left_pads,
                                      input_right_pads);

        out_device_buf.FromDevice(out_n_c_ho_wo_device.mData.data());

317
        pass = pass && ck::utils::check_err(out_n_c_ho_wo_device.mData, out_n_c_ho_wo_host.mData);
318

319
        if constexpr(OutputIndex)
320
321
322
        {
            out_indices_device_buf.FromDevice(out_indices_n_c_ho_wo_device.mData.data());

323
324
            pass = pass && ck::utils::check_err(out_indices_n_c_ho_wo_device.mData,
                                                out_indices_n_c_ho_wo_host.mData);
325
326
        };
    }
327
328

    return (pass ? 0 : 1);
329
}