maxpool2d_bwd_common.hpp 10.4 KB
Newer Older
rocking's avatar
rocking committed
1
// SPDX-License-Identifier: MIT
rocking's avatar
rocking committed
2
// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
rocking's avatar
rocking committed
3
4
5
6
7
8
9
10

#pragma once

#include <iostream>

#include "ck/ck.hpp"
#include "ck/utility/reduction_enums.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_pool2d_fwd_nhwc_nhwc.hpp"
rocking's avatar
rocking committed
11
#include "ck/tensor_operation/gpu/device/impl/device_index_pool_bwd_impl.hpp"
rocking's avatar
rocking committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
#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"
#include "ck/library/utility/literals.hpp"
#include "ck/library/reference_tensor_operation/cpu/reference_pool_fwd.hpp"
#include "ck/library/reference_tensor_operation/cpu/reference_maxpool_bwd.hpp"

template <typename InDataType,
          typename OutDataType,
          typename IndexDataType,
          typename ComputeDataType,
          typename DInDataType,
          typename DOutDataType,
rocking's avatar
rocking committed
28
          bool PropagateNan>
rocking's avatar
rocking committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
bool maxpool_bwd_test(bool do_verification,
                      bool time_kernel,
                      ck::index_t N,
                      ck::index_t C,
                      ck::index_t Y,
                      ck::index_t X,
                      ck::index_t Hi,
                      ck::index_t Wi,
                      ck::index_t window_stride_h,
                      ck::index_t window_stride_w,
                      ck::index_t in_left_pad_h,
                      ck::index_t in_left_pad_w,
                      ck::index_t in_right_pad_h,
                      ck::index_t in_right_pad_w)
{
    using PassThrough = ck::tensor_operation::element_wise::PassThrough;

    using DevicePoolFwdInstance =
        ck::tensor_operation::device::DevicePool2dFwd_Input_N_Hi_Wi_C_Output_N_Ho_Wo_C<
            InDataType,      // InDataType
            OutDataType,     // OutDataType
            IndexDataType,   // IndexDataType
            ComputeDataType, // ComputeDataType
rocking's avatar
rocking committed
52
            ck::ReduceTensorOp::MAX,
rocking's avatar
rocking committed
53
54
55
56
57
58
59
60
61
            true, // OutputIndex
            64,   // BlockSize
            64,   // ReduceMThreadClusterSize
            1,    // ReduceKThreadClusterSize
            4,    // ReduceMThreadSliceSize
            1,    // ReduceKThreadSliceSize
            1>;   // InSrcOutDstVectorSize

    using DeviceMaxPoolBwdInstance = ck::tensor_operation::device::
rocking's avatar
rocking committed
62
        DeviceIndexPoolBwdImpl<DOutDataType, IndexDataType, DInDataType, 4>;
rocking's avatar
rocking committed
63
64
65
66
67
68
69
70
71
72

    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::vector<ck::index_t> window_spatial_lengths{Y, X};
    const std::vector<ck::index_t> window_strides{window_stride_h, window_stride_w};
    const std::vector<ck::index_t> input_left_pads{in_left_pad_h, in_left_pad_w};
    const std::vector<ck::index_t> input_right_pads{in_right_pad_h, in_right_pad_w};

    auto f_host_tensor_descriptor =
rocking's avatar
rocking committed
73
        [](std::size_t N_, std::size_t C_, std::size_t H, std::size_t W) {
rocking's avatar
rocking committed
74
            using namespace ck::literals;
rocking's avatar
rocking committed
75
76
            // reference need Tensor with NCHW order
            return HostTensorDescriptor({N_, C_, H, W}, {C_ * H * W, 1_uz, W * C_, C_});
rocking's avatar
rocking committed
77
78
79
        };

    // in
rocking's avatar
rocking committed
80
    Tensor<InDataType> in_n_c_hi_wi(f_host_tensor_descriptor(N, C, Hi, Wi));
rocking's avatar
rocking committed
81
82

    // out
rocking's avatar
rocking committed
83
84
    Tensor<OutDataType> out_n_c_ho_wo_host(f_host_tensor_descriptor(N, C, Ho, Wo));
    Tensor<OutDataType> out_n_c_ho_wo_device(f_host_tensor_descriptor(N, C, Ho, Wo));
rocking's avatar
rocking committed
85
86

    // indices
rocking's avatar
rocking committed
87
88
    Tensor<IndexDataType> indices_n_c_ho_wo_device(f_host_tensor_descriptor(N, C, Ho, Wo));
    Tensor<IndexDataType> indices_n_c_ho_wo_host(f_host_tensor_descriptor(N, C, Ho, Wo));
rocking's avatar
rocking committed
89
90

    // dout
rocking's avatar
rocking committed
91
    Tensor<DOutDataType> dout_n_c_ho_wo(f_host_tensor_descriptor(N, C, Ho, Wo));
rocking's avatar
rocking committed
92
93

    // din
rocking's avatar
rocking committed
94
95
    Tensor<DInDataType> din_n_c_hi_wi_host(f_host_tensor_descriptor(N, C, Hi, Wi));
    Tensor<DInDataType> din_n_c_hi_wi_device(f_host_tensor_descriptor(N, C, Hi, Wi));
rocking's avatar
rocking committed
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

    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;
    std::cout << "indices_n_c_ho_wo: " << indices_n_c_ho_wo_host.mDesc << std::endl;
    std::cout << "dout_n_c_ho_wo: " << dout_n_c_ho_wo.mDesc << std::endl;
    std::cout << "din_n_c_hi_wi: " << din_n_c_hi_wi_host.mDesc << std::endl;

    in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_3<InDataType>{-1.0, 1.0});
    dout_n_c_ho_wo.GenerateTensorValue(GeneratorTensor_3<DOutDataType>{-1.0, 1.0});

    DeviceMem in_device_buf(sizeof(InDataType) * in_n_c_hi_wi.mDesc.GetElementSpaceSize());
    DeviceMem out_device_buf(sizeof(OutDataType) *
                             out_n_c_ho_wo_device.mDesc.GetElementSpaceSize());
    DeviceMem indices_device_buf(sizeof(IndexDataType) *
                                 indices_n_c_ho_wo_device.mDesc.GetElementSpaceSize());
    DeviceMem dout_device_buf(sizeof(DOutDataType) * dout_n_c_ho_wo.mDesc.GetElementSpaceSize());
    DeviceMem din_device_buf(sizeof(DInDataType) *
                             din_n_c_hi_wi_device.mDesc.GetElementSpaceSize());

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

    auto pool_fwd              = DevicePoolFwdInstance{};
    auto pool_fwd_invoker_ptr  = pool_fwd.MakeInvokerPointer();
    auto pool_fwd_argument_ptr = pool_fwd.MakeArgumentPointer(
        static_cast<InDataType*>(in_device_buf.GetDeviceBuffer()),
        static_cast<OutDataType*>(out_device_buf.GetDeviceBuffer()),
        static_cast<IndexDataType*>(indices_device_buf.GetDeviceBuffer()),
        {N, C, Hi, Wi},
rocking's avatar
rocking committed
125
        window_spatial_lengths,
rocking's avatar
rocking committed
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
        {N, C, Ho, Wo},
        {C * Hi * Wi, 1, Wi * C, C},
        {C * Ho * Wo, 1, Wo * C, C},
        {C * Ho * Wo, 1, Wo * C, C},
        window_strides,
        input_left_pads,
        input_right_pads,
        {2, 3});

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

    float ave_time_fwd =
        pool_fwd_invoker_ptr->Run(pool_fwd_argument_ptr.get(), StreamConfig{nullptr, time_kernel});

    auto pool_bwd              = DeviceMaxPoolBwdInstance{};
    auto pool_bwd_invoker_ptr  = pool_bwd.MakeInvokerPointer();
    auto pool_bwd_argument_ptr = pool_bwd.MakeArgumentPointer(
        static_cast<DOutDataType*>(dout_device_buf.GetDeviceBuffer()),
        static_cast<IndexDataType*>(indices_device_buf.GetDeviceBuffer()),
        static_cast<DInDataType*>(din_device_buf.GetDeviceBuffer()),
        dout_n_c_ho_wo.mDesc.GetElementSpaceSize(),
        din_n_c_hi_wi_device.mDesc.GetElementSpaceSize(),
rocking's avatar
rocking committed
152
153
        window_spatial_lengths,
        window_strides);
rocking's avatar
rocking committed
154
155
156
157
158
159
160

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

rocking's avatar
rocking committed
161
162
163
164
165
    size_t pool_bwd_workspace_sz = pool_bwd.GetWorkSpaceSize(pool_bwd_argument_ptr.get());
    DeviceMem pool_bwd_workspace_device_buf(pool_bwd_workspace_sz);
    pool_bwd.SetWorkSpacePointer(pool_bwd_argument_ptr.get(),
                                 pool_bwd_workspace_device_buf.GetDeviceBuffer());

rocking's avatar
rocking committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
    float ave_time_bwd =
        pool_bwd_invoker_ptr->Run(pool_bwd_argument_ptr.get(), StreamConfig{nullptr, time_kernel});

    std::cout << "Pool fwd perf: " << ave_time_fwd << " ms" << std::endl;
    std::cout << "Pool bwd perf: " << ave_time_bwd << " ms" << std::endl;

    bool pass = true;

    if(do_verification)
    {
        using ReferencePoolingFwdInstance =
            ck::tensor_operation::host::ReferencePoolingFwd<4,
                                                            2,
                                                            InDataType,
                                                            OutDataType,
                                                            ComputeDataType,
                                                            IndexDataType,
rocking's avatar
rocking committed
183
                                                            ck::ReduceTensorOp::MAX,
rocking's avatar
rocking committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
                                                            PropagateNan,
                                                            true>;

        auto ref_pooling_fwd          = ReferencePoolingFwdInstance{};
        auto ref_pooling_fwd_invoker  = ref_pooling_fwd.MakeInvoker();
        auto ref_pooling_fwd_argument = ref_pooling_fwd.MakeArgument(in_n_c_hi_wi,
                                                                     out_n_c_ho_wo_host,
                                                                     indices_n_c_ho_wo_host,
                                                                     window_spatial_lengths,
                                                                     window_strides,
                                                                     input_left_pads,
                                                                     input_right_pads);
        ref_pooling_fwd_invoker.Run(ref_pooling_fwd_argument);

rocking's avatar
rocking committed
198
199
200
201
202
203
        using ReferencePoolingBwdInstance =
            ck::tensor_operation::host::ReferenceMaxPoolBwd<DOutDataType,
                                                            IndexDataType,
                                                            ComputeDataType,
                                                            DInDataType,
                                                            PassThrough>;
rocking's avatar
rocking committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222

        auto ref_pooling_bwd          = ReferencePoolingBwdInstance{};
        auto ref_pooling_bwd_invoker  = ref_pooling_bwd.MakeInvoker();
        auto ref_pooling_bwd_argument = ref_pooling_bwd.MakeArgument(
            dout_n_c_ho_wo, indices_n_c_ho_wo_host, din_n_c_hi_wi_host, PassThrough{});

        ref_pooling_bwd_invoker.Run(ref_pooling_bwd_argument);

        out_device_buf.FromDevice(out_n_c_ho_wo_device.mData.data());
        indices_device_buf.FromDevice(indices_n_c_ho_wo_device.mData.data());
        din_device_buf.FromDevice(din_n_c_hi_wi_device.mData.data());

        pass = pass && ck::utils::check_err(out_n_c_ho_wo_device, out_n_c_ho_wo_host);
        pass = pass && ck::utils::check_err(indices_n_c_ho_wo_device, indices_n_c_ho_wo_host);
        pass = pass && ck::utils::check_err(din_n_c_hi_wi_device, din_n_c_hi_wi_host);
    }

    return (pass);
};