convnd_fwd_common.hpp 14.6 KB
Newer Older
Chao Liu's avatar
Chao Liu committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.

#include <cstdlib>
#include <iostream>
#include <numeric>
#include <type_traits>

#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"

#include "ck/library/utility/check_err.hpp"
#include "ck/library/utility/device_memory.hpp"
#include "ck/library/utility/host_tensor.hpp"
#include "ck/library/utility/host_tensor_generator.hpp"
Chao Liu's avatar
Chao Liu committed
17
#include "ck/library/utility/convolution_parameter.hpp"
Chao Liu's avatar
clean  
Chao Liu committed
18
#include "ck/library/utility/convolution_host_tensor_descriptor_helper.hpp"
Chao Liu's avatar
Chao Liu committed
19
#include "ck/library/reference_tensor_operation/cpu/reference_conv_fwd.hpp"
Chao Liu's avatar
Chao Liu committed
20

Chao Liu's avatar
Chao Liu committed
21
22
23
24
25
26
27
28
29
void print_helper_msg()
{
    std::cout << "arg1: verification (0=no, 1=yes)\n"
              << "arg2: initialization (0=no init, 1=integer value, 2=decimal value)\n"
              << "arg3: time kernel (0=no, 1=yes)\n"
              << "arg4: N spatial dimensions (default 2)\n"
              << "Following arguments (depending on number of spatial dims):\n"
              << " N, K, C, \n"
              << " <filter spatial dimensions>, (ie Y, X for 2D)\n"
Chao Liu's avatar
Chao Liu committed
30
              << " <input image spatial dimensions>, (ie Hi, Wi for 2D)\n"
Chao Liu's avatar
Chao Liu committed
31
32
33
34
35
36
37
              << " <strides>, (ie Sy, Sx for 2D)\n"
              << " <dilations>, (ie Dy, Dx for 2D)\n"
              << " <left padding>, (ie LeftPy, LeftPx for 2D)\n"
              << " <right padding>, (ie RightPy, RightPx for 2D)\n"
              << std::endl;
}

Chao Liu's avatar
clean  
Chao Liu committed
38
ck::utils::conv::ConvParam parse_conv_params(int num_dim_spatial, int arg_idx, char* const argv[])
Chao Liu's avatar
Chao Liu committed
39
{
Chao Liu's avatar
Chao Liu committed
40
41
42
    const ck::index_t N = std::stoi(argv[arg_idx++]);
    const ck::index_t K = std::stoi(argv[arg_idx++]);
    const ck::index_t C = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
43

Chao Liu's avatar
Chao Liu committed
44
45
46
47
48
49
    std::vector<ck::index_t> filter_spatial_lengths(num_dim_spatial);
    std::vector<ck::index_t> input_spatial_lengths(num_dim_spatial);
    std::vector<ck::index_t> conv_filter_strides(num_dim_spatial);
    std::vector<ck::index_t> conv_filter_dilations(num_dim_spatial);
    std::vector<ck::index_t> input_left_pads(num_dim_spatial);
    std::vector<ck::index_t> input_right_pads(num_dim_spatial);
Chao Liu's avatar
Chao Liu committed
50
51
52

    for(int i = 0; i < num_dim_spatial; ++i)
    {
Chao Liu's avatar
Chao Liu committed
53
        filter_spatial_lengths[i] = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
54
    }
Chao Liu's avatar
Chao Liu committed
55

Chao Liu's avatar
Chao Liu committed
56
57
    for(int i = 0; i < num_dim_spatial; ++i)
    {
Chao Liu's avatar
Chao Liu committed
58
        input_spatial_lengths[i] = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
59
    }
Chao Liu's avatar
Chao Liu committed
60

Chao Liu's avatar
Chao Liu committed
61
62
    for(int i = 0; i < num_dim_spatial; ++i)
    {
Chao Liu's avatar
Chao Liu committed
63
        conv_filter_strides[i] = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
64
    }
Chao Liu's avatar
Chao Liu committed
65

Chao Liu's avatar
Chao Liu committed
66
67
    for(int i = 0; i < num_dim_spatial; ++i)
    {
Chao Liu's avatar
Chao Liu committed
68
        conv_filter_dilations[i] = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
69
    }
Chao Liu's avatar
Chao Liu committed
70

Chao Liu's avatar
Chao Liu committed
71
72
    for(int i = 0; i < num_dim_spatial; ++i)
    {
Chao Liu's avatar
Chao Liu committed
73
        input_left_pads[i] = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
74
    }
Chao Liu's avatar
Chao Liu committed
75

Chao Liu's avatar
Chao Liu committed
76
77
    for(int i = 0; i < num_dim_spatial; ++i)
    {
Chao Liu's avatar
Chao Liu committed
78
        input_right_pads[i] = std::stoi(argv[arg_idx++]);
Chao Liu's avatar
Chao Liu committed
79
80
    }

Chao Liu's avatar
clean  
Chao Liu committed
81
82
83
84
85
86
87
88
89
90
    return ck::utils::conv::ConvParam{num_dim_spatial,
                                      N,
                                      K,
                                      C,
                                      filter_spatial_lengths,
                                      input_spatial_lengths,
                                      conv_filter_strides,
                                      conv_filter_dilations,
                                      input_left_pads,
                                      input_right_pads};
Chao Liu's avatar
Chao Liu committed
91
92
}

Chao Liu's avatar
Chao Liu committed
93
// FIXME: current implementation only support NCHW/NHWC layout
Chao Liu's avatar
Chao Liu committed
94
template <ck::index_t NDimSpatial,
Chao Liu's avatar
Chao Liu committed
95
96
97
          typename InLayout,
          typename WeiLayout,
          typename OutLayout,
Chao Liu's avatar
Chao Liu committed
98
99
100
101
102
103
          typename InDataType,
          typename WeiDataType,
          typename OutDataType,
          typename InElementOp,
          typename WeiElementOp,
          typename OutElementOp,
Chao Liu's avatar
Chao Liu committed
104
          typename DeviceConvNDFwdInstance>
Chao Liu's avatar
Chao Liu committed
105
106
107
int run_conv_fwd(bool do_verification,
                 int init_method,
                 bool time_kernel,
Chao Liu's avatar
clean  
Chao Liu committed
108
                 const ck::utils::conv::ConvParam& conv_param,
Chao Liu's avatar
Chao Liu committed
109
110
111
                 const InElementOp& in_element_op,
                 const WeiElementOp& wei_element_op,
                 const OutElementOp& out_element_op)
Chao Liu's avatar
Chao Liu committed
112
{
Chao Liu's avatar
clean  
Chao Liu committed
113
114
115
    const auto in_desc  = ck::utils::conv::get_input_host_tensor_descriptor<InLayout>(conv_param);
    const auto wei_desc = ck::utils::conv::get_weight_host_tensor_descriptor<WeiLayout>(conv_param);
    const auto out_desc = ck::utils::conv::get_output_host_tensor_descriptor<OutLayout>(conv_param);
Chao Liu's avatar
Chao Liu committed
116

Chao Liu's avatar
Chao Liu committed
117
118
119
120
121
122
123
124
    // hacky, hardcoded for 2d NHWK
    const auto bias_desc = HostTensorDescriptor(
        std::vector<std::size_t>{static_cast<std::size_t>(conv_param.N_),
                                 static_cast<std::size_t>(conv_param.output_spatial_lengths_[0]),
                                 static_cast<std::size_t>(conv_param.output_spatial_lengths_[1]),
                                 static_cast<std::size_t>(conv_param.K_)},
        std::vector<std::size_t>{0, 0, 0, 1});

Chao Liu's avatar
Chao Liu committed
125
126
    Tensor<InDataType> in(in_desc);
    Tensor<WeiDataType> wei(wei_desc);
Chao Liu's avatar
Chao Liu committed
127
    Tensor<OutDataType> bias(bias_desc);
Chao Liu's avatar
Chao Liu committed
128
129
    Tensor<OutDataType> out_host(out_desc);
    Tensor<OutDataType> out_device(out_desc);
Chao Liu's avatar
Chao Liu committed
130

Chao Liu's avatar
Chao Liu committed
131
132
    std::cout << "in: " << in.mDesc << std::endl;
    std::cout << "wei: " << wei.mDesc << std::endl;
Chao Liu's avatar
Chao Liu committed
133
    std::cout << "bias: " << bias.mDesc << std::endl;
Chao Liu's avatar
Chao Liu committed
134
    std::cout << "out: " << out_host.mDesc << std::endl;
Chao Liu's avatar
Chao Liu committed
135
136
137
138
139

    switch(init_method)
    {
    case 0: break;
    case 1:
Chao Liu's avatar
Chao Liu committed
140
141
        in.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5});
        wei.GenerateTensorValue(GeneratorTensor_2<WeiDataType>{-5, 5});
Chao Liu's avatar
Chao Liu committed
142
        bias.GenerateTensorValue(GeneratorTensor_2<OutDataType>{-5, 5});
Chao Liu's avatar
Chao Liu committed
143
144
        break;
    default:
Chao Liu's avatar
Chao Liu committed
145
146
        in.GenerateTensorValue(GeneratorTensor_3<InDataType>{0.0, 1.0});
        wei.GenerateTensorValue(GeneratorTensor_3<WeiDataType>{-0.5, 0.5});
Chao Liu's avatar
Chao Liu committed
147
        bias.GenerateTensorValue(GeneratorTensor_3<OutDataType>{-0.5, 0.5});
Chao Liu's avatar
Chao Liu committed
148
149
    }

Chao Liu's avatar
Chao Liu committed
150
151
    DeviceMem in_device_buf(sizeof(InDataType) * in.mDesc.GetElementSpace());
    DeviceMem wei_device_buf(sizeof(WeiDataType) * wei.mDesc.GetElementSpace());
Chao Liu's avatar
Chao Liu committed
152
    DeviceMem bias_device_buf(sizeof(OutDataType) * bias.mDesc.GetElementSpace());
Chao Liu's avatar
Chao Liu committed
153
    DeviceMem out_device_buf(sizeof(OutDataType) * out_device.mDesc.GetElementSpace());
Chao Liu's avatar
Chao Liu committed
154

Chao Liu's avatar
Chao Liu committed
155
156
    in_device_buf.ToDevice(in.mData.data());
    wei_device_buf.ToDevice(wei.mData.data());
Chao Liu's avatar
Chao Liu committed
157
    bias_device_buf.ToDevice(bias.mData.data());
Chao Liu's avatar
Chao Liu committed
158

Chao Liu's avatar
Chao Liu committed
159
    // tensor descriptor in NCHW/KXYC/NKHW dimensional order
Chao Liu's avatar
Chao Liu committed
160
161
162
163
    HostTensorDescriptor in_n_c_wis_desc   = in_desc;
    HostTensorDescriptor wei_k_c_xs_desc   = wei_desc;
    HostTensorDescriptor bias_n_k_wos_desc = bias_desc;
    HostTensorDescriptor out_n_k_wos_desc  = out_desc;
Chao Liu's avatar
Chao Liu committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203

    // input
    if constexpr(ck::is_same_v<InLayout, ck::tensor_layout::convolution::NWC>)
    {
        in_n_c_wis_desc = transpose_host_tensor_descriptor_given_new2old(
            in_desc, std::vector<std::size_t>{0, 2, 1});
    }
    else if constexpr(ck::is_same_v<InLayout, ck::tensor_layout::convolution::NHWC>)
    {
        in_n_c_wis_desc = transpose_host_tensor_descriptor_given_new2old(
            in_desc, std::vector<std::size_t>{0, 3, 1, 2});
    }
    else if constexpr(ck::is_same_v<InLayout, ck::tensor_layout::convolution::NDHWC>)
    {
        in_n_c_wis_desc = transpose_host_tensor_descriptor_given_new2old(
            in_desc, std::vector<std::size_t>{0, 4, 1, 2, 3});
    }

    // weight
    if constexpr(ck::is_same_v<WeiLayout, ck::tensor_layout::convolution::KXC>)
    {
        wei_k_c_xs_desc = transpose_host_tensor_descriptor_given_new2old(
            wei_desc, std::vector<std::size_t>{0, 2, 1});
    }
    else if constexpr(ck::is_same_v<WeiLayout, ck::tensor_layout::convolution::KYXC>)
    {
        wei_k_c_xs_desc = transpose_host_tensor_descriptor_given_new2old(
            wei_desc, std::vector<std::size_t>{0, 3, 1, 2});
    }
    else if constexpr(ck::is_same_v<WeiLayout, ck::tensor_layout::convolution::KZYXC>)
    {
        wei_k_c_xs_desc = transpose_host_tensor_descriptor_given_new2old(
            wei_desc, std::vector<std::size_t>{0, 4, 1, 2, 3});
    }

    // output
    if constexpr(ck::is_same_v<OutLayout, ck::tensor_layout::convolution::NWK>)
    {
        out_n_k_wos_desc = transpose_host_tensor_descriptor_given_new2old(
            out_desc, std::vector<std::size_t>{0, 2, 1});
Chao Liu's avatar
Chao Liu committed
204
205
206

        bias_n_k_wos_desc = transpose_host_tensor_descriptor_given_new2old(
            bias_desc, std::vector<std::size_t>{0, 2, 1});
Chao Liu's avatar
Chao Liu committed
207
208
209
210
211
    }
    else if constexpr(ck::is_same_v<OutLayout, ck::tensor_layout::convolution::NHWK>)
    {
        out_n_k_wos_desc = transpose_host_tensor_descriptor_given_new2old(
            out_desc, std::vector<std::size_t>{0, 3, 1, 2});
Chao Liu's avatar
Chao Liu committed
212
213
214

        bias_n_k_wos_desc = transpose_host_tensor_descriptor_given_new2old(
            bias_desc, std::vector<std::size_t>{0, 3, 1, 2});
Chao Liu's avatar
Chao Liu committed
215
216
217
218
219
    }
    else if constexpr(ck::is_same_v<OutLayout, ck::tensor_layout::convolution::NDHWK>)
    {
        out_n_k_wos_desc = transpose_host_tensor_descriptor_given_new2old(
            out_desc, std::vector<std::size_t>{0, 4, 1, 2, 3});
Chao Liu's avatar
Chao Liu committed
220
221
222

        bias_n_k_wos_desc = transpose_host_tensor_descriptor_given_new2old(
            bias_desc, std::vector<std::size_t>{0, 4, 1, 2, 3});
Chao Liu's avatar
Chao Liu committed
223
224
225
226
227
228
    }

    std::array<ck::index_t, NDimSpatial + 2> a_n_c_wis_lengths{};
    std::array<ck::index_t, NDimSpatial + 2> a_n_c_wis_strides{};
    std::array<ck::index_t, NDimSpatial + 2> b_k_c_xs_lengths{};
    std::array<ck::index_t, NDimSpatial + 2> b_k_c_xs_strides{};
Chao Liu's avatar
Chao Liu committed
229
230
    std::array<ck::index_t, NDimSpatial + 2> d_n_k_wos_lengths{};
    std::array<ck::index_t, NDimSpatial + 2> d_n_k_wos_strides{};
Chao Liu's avatar
Chao Liu committed
231
232
233
234
235
236
237
238
239
240
241
242
243
    std::array<ck::index_t, NDimSpatial + 2> e_n_k_wos_lengths{};
    std::array<ck::index_t, NDimSpatial + 2> e_n_k_wos_strides{};
    std::array<ck::index_t, NDimSpatial> conv_filter_strides{};
    std::array<ck::index_t, NDimSpatial> conv_filter_dilations{};
    std::array<ck::index_t, NDimSpatial> input_left_pads{};
    std::array<ck::index_t, NDimSpatial> input_right_pads{};

    auto copy = [](auto& x, auto& y) { std::copy(x.begin(), x.end(), y.begin()); };

    copy(in_n_c_wis_desc.GetLengths(), a_n_c_wis_lengths);
    copy(in_n_c_wis_desc.GetStrides(), a_n_c_wis_strides);
    copy(wei_k_c_xs_desc.GetLengths(), b_k_c_xs_lengths);
    copy(wei_k_c_xs_desc.GetStrides(), b_k_c_xs_strides);
Chao Liu's avatar
Chao Liu committed
244
245
    copy(bias_n_k_wos_desc.GetLengths(), d_n_k_wos_lengths);
    copy(bias_n_k_wos_desc.GetStrides(), d_n_k_wos_strides);
Chao Liu's avatar
Chao Liu committed
246
247
248
249
250
251
252
    copy(out_n_k_wos_desc.GetLengths(), e_n_k_wos_lengths);
    copy(out_n_k_wos_desc.GetStrides(), e_n_k_wos_strides);
    copy(conv_param.conv_filter_strides_, conv_filter_strides);
    copy(conv_param.conv_filter_dilations_, conv_filter_dilations);
    copy(conv_param.input_left_pads_, input_left_pads);
    copy(conv_param.input_right_pads_, input_right_pads);

Chao Liu's avatar
Chao Liu committed
253
254
255
    // do GEMM
    auto conv     = DeviceConvNDFwdInstance{};
    auto invoker  = conv.MakeInvoker();
Chao Liu's avatar
Chao Liu committed
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
    auto argument = conv.MakeArgument(
        in_device_buf.GetDeviceBuffer(),
        wei_device_buf.GetDeviceBuffer(),
        std::array<const void*, 1>{bias_device_buf.GetDeviceBuffer()},
        out_device_buf.GetDeviceBuffer(),
        a_n_c_wis_lengths,
        a_n_c_wis_strides,
        b_k_c_xs_lengths,
        b_k_c_xs_strides,
        std::array<std::array<ck::index_t, NDimSpatial + 2>, 1>{{d_n_k_wos_lengths}},
        std::array<std::array<ck::index_t, NDimSpatial + 2>, 1>{{d_n_k_wos_strides}},
        e_n_k_wos_lengths,
        e_n_k_wos_strides,
        conv_filter_strides,
        conv_filter_dilations,
        input_left_pads,
        input_right_pads,
        in_element_op,
        wei_element_op,
        out_element_op);
Chao Liu's avatar
Chao Liu committed
276
277
278
279
280
281
282
283

    if(!conv.IsSupportedArgument(argument))
    {
        throw std::runtime_error(
            "wrong! device_conv with the specified compilation parameters does "
            "not support this Conv problem");
    }

Chao Liu's avatar
Chao Liu committed
284
    float avg_time = invoker.Run(argument, StreamConfig{nullptr, time_kernel});
Chao Liu's avatar
Chao Liu committed
285

Chao Liu's avatar
clean  
Chao Liu committed
286
287
    std::size_t flop      = conv_param.GetFlops();
    std::size_t num_btype = conv_param.GetByte<InDataType, WeiDataType, OutDataType>();
Chao Liu's avatar
Chao Liu committed
288

Chao Liu's avatar
Chao Liu committed
289
290
291
    float tflops     = static_cast<float>(flop) / 1.E9 / avg_time;
    float gb_per_sec = num_btype / 1.E6 / avg_time;
    std::cout << "Perf: " << avg_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
Chao Liu's avatar
Chao Liu committed
292
293
294
295
              << conv.GetTypeString() << std::endl;

    if(do_verification)
    {
Chao Liu's avatar
Chao Liu committed
296
297
298
299
        using PassThrough = ck::tensor_operation::element_wise::PassThrough;

        Tensor<OutDataType> c_host(out_desc);

Chao Liu's avatar
Chao Liu committed
300
301
302
303
304
305
306
307
308
        auto ref_conv = ck::tensor_operation::host::ReferenceConvFwd<NDimSpatial,
                                                                     InLayout,
                                                                     WeiLayout,
                                                                     OutLayout,
                                                                     InDataType,
                                                                     WeiDataType,
                                                                     OutDataType,
                                                                     InElementOp,
                                                                     WeiElementOp,
Chao Liu's avatar
Chao Liu committed
309
                                                                     PassThrough>();
Chao Liu's avatar
Chao Liu committed
310
311

        auto ref_invoker  = ref_conv.MakeInvoker();
Chao Liu's avatar
Chao Liu committed
312
313
        auto ref_argument = ref_conv.MakeArgument(in,
                                                  wei,
Chao Liu's avatar
Chao Liu committed
314
                                                  c_host,
Chao Liu's avatar
clean  
Chao Liu committed
315
316
317
318
                                                  conv_param.conv_filter_strides_,
                                                  conv_param.conv_filter_dilations_,
                                                  conv_param.input_left_pads_,
                                                  conv_param.input_right_pads_,
Chao Liu's avatar
Chao Liu committed
319
320
                                                  in_element_op,
                                                  wei_element_op,
Chao Liu's avatar
Chao Liu committed
321
                                                  PassThrough{});
Chao Liu's avatar
Chao Liu committed
322
323
324

        ref_invoker.Run(ref_argument);

Chao Liu's avatar
Chao Liu committed
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
        for(int n = 0; n < out_host.mDesc.GetLengths()[0]; n++)
        {
            for(int ho = 0; ho < out_host.mDesc.GetLengths()[1]; ho++)
            {
                for(int wo = 0; wo < out_host.mDesc.GetLengths()[2]; wo++)
                {
                    for(int k = 0; k < out_host.mDesc.GetLengths()[3]; k++)
                    {
                        out_element_op(
                            out_host(n, ho, wo, k), c_host(n, ho, wo, k), bias(n, ho, wo, k));
                    }
                }
            }
        }

Chao Liu's avatar
Chao Liu committed
340
        out_device_buf.FromDevice(out_device.mData.data());
Chao Liu's avatar
Chao Liu committed
341

Chao Liu's avatar
Chao Liu committed
342
        return ck::utils::check_err(
Chao Liu's avatar
Chao Liu committed
343
                   out_device.mData, out_host.mData, "Error: incorrect results!", 1e-5f, 1e-4f)
Chao Liu's avatar
Chao Liu committed
344
345
346
347
348
349
                   ? 0
                   : 1;
    }

    return 0;
}