parse_resize.cpp 12.9 KB
Newer Older
Paul Fultz II's avatar
Paul Fultz II committed
1
2
3
4
5
6
7
8
9
10
11
12
13
#include <migraphx/onnx/op_parser.hpp>
#include <migraphx/onnx/checks.hpp>
#include <migraphx/ranges.hpp>
#include <migraphx/shape_for_each.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/make_op.hpp>

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace onnx {

const auto& get_nearest_op(const std::string& mode)
{
Shucai Xiao's avatar
Shucai Xiao committed
14
    using nearest_op = std::function<int(int, double)>;
Paul Fultz II's avatar
Paul Fultz II committed
15
16
    static std::unordered_map<std::string, nearest_op> const nearest_ops = {
        {"round_prefer_floor",
Shucai Xiao's avatar
Shucai Xiao committed
17
         [=](int d_in, double val) {
Paul Fultz II's avatar
Paul Fultz II committed
18
             val = std::max(0.0, std::min(d_in - 1.0, val));
Shucai Xiao's avatar
Shucai Xiao committed
19
             return static_cast<int>(std::ceil((val - 0.5)));
Paul Fultz II's avatar
Paul Fultz II committed
20
21
         }},
        {"round_prefer_ceil",
Shucai Xiao's avatar
Shucai Xiao committed
22
         [=](int d_in, double val) {
Paul Fultz II's avatar
Paul Fultz II committed
23
             val = std::max(0.0, std::min(d_in - 1.0, val));
Shucai Xiao's avatar
Shucai Xiao committed
24
             return static_cast<int>(std::round((val)));
Paul Fultz II's avatar
Paul Fultz II committed
25
26
         }},
        {"floor",
Shucai Xiao's avatar
Shucai Xiao committed
27
         [=](int d_in, double val) {
Paul Fultz II's avatar
Paul Fultz II committed
28
             val = std::max(0.0, std::min(d_in - 1.0, val));
Shucai Xiao's avatar
Shucai Xiao committed
29
             return static_cast<int>(std::floor((val)));
Paul Fultz II's avatar
Paul Fultz II committed
30
         }},
Shucai Xiao's avatar
Shucai Xiao committed
31
        {"ceil", [=](int d_in, double val) {
Paul Fultz II's avatar
Paul Fultz II committed
32
             val = std::max(0.0, std::min(d_in - 1.0, val));
Shucai Xiao's avatar
Shucai Xiao committed
33
             return static_cast<int>(std::ceil((val)));
Paul Fultz II's avatar
Paul Fultz II committed
34
35
36
37
38
39
40
41
42
43
44
45
         }}};

    if(!contains(nearest_ops, mode))
    {
        MIGRAPHX_THROW("PARSE_RESIZE: nearest_mode " + mode + " not supported!");
    }

    return nearest_ops.at(mode);
}

const auto& get_original_idx_op(const std::string& mode)
{
Shucai Xiao's avatar
Shucai Xiao committed
46
    using original_idx_op = std::function<double(int, int, int, double)>;
Paul Fultz II's avatar
Paul Fultz II committed
47
    static std::unordered_map<std::string, original_idx_op> const idx_ops = {
Shucai Xiao's avatar
Shucai Xiao committed
48
        {"half_pixel", [=](int, int, int idx, double scale) { return (idx + 0.5) / scale - 0.5; }},
Paul Fultz II's avatar
Paul Fultz II committed
49
        {"pytorch_half_pixel",
Shucai Xiao's avatar
Shucai Xiao committed
50
         [=](int, int l_out, int idx, double scale) {
Paul Fultz II's avatar
Paul Fultz II committed
51
52
53
             return l_out > 1 ? (idx + 0.5) / scale - 0.5 : 0.0;
         }},
        {"align_corners",
Shucai Xiao's avatar
Shucai Xiao committed
54
         [=](int l_in, int l_out, int idx, double) {
55
             return (l_out == 1) ? 0.0 : (1.0 * idx * (l_in - 1.0) / (l_out - 1.0));
Paul Fultz II's avatar
Paul Fultz II committed
56
         }},
Shucai Xiao's avatar
Shucai Xiao committed
57
58
59
        {"asymmetric", [=](int, int, int idx, double scale) { return idx / scale; }},
        {"tf_half_pixel_for_nn",
         [=](int, int, int idx, double scale) { return (idx + 0.5) / scale; }}};
Paul Fultz II's avatar
Paul Fultz II committed
60
61
62
63
64
65
66
67
68

    if(!contains(idx_ops, mode))
    {
        MIGRAPHX_THROW("PARSE_RESIZE: coordinate_transformation_mode " + mode + " not supported!");
    }

    return idx_ops.at(mode);
}

69
static std::vector<int>
Shucai Xiao's avatar
Shucai Xiao committed
70
calc_neighbor_points(const std::vector<std::vector<std::vector<int>>>& vvv_ind,
71
                     int i_dim,
Shucai Xiao's avatar
Shucai Xiao committed
72
                     const std::vector<std::vector<int>>& vec_dims,
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
                     const shape& in_s)
{
    if(i_dim == vvv_ind.size())
    {
        std::vector<int> vec_ind;
        vec_ind.resize(vec_dims.size());
        std::transform(vec_dims.begin(), vec_dims.end(), vec_ind.begin(), [&](auto idx) {
            return static_cast<int>(in_s.index(idx));
        });

        return vec_ind;
    }

    const auto& vv_ind = vvv_ind[i_dim];
    const auto& vv_lo  = vv_ind.at(0);
Shucai Xiao's avatar
Shucai Xiao committed
88
89
    std::vector<std::vector<int>> vec_dims1;
    for(int start = 0; start < vec_dims.size(); start += vv_lo.size())
90
91
92
93
94
95
96
97
98
99
100
101
    {
        std::transform(vv_lo.begin(),
                       vv_lo.end(),
                       vec_dims.begin() + start,
                       std::back_inserter(vec_dims1),
                       [](auto i, auto dim) {
                           dim.push_back(i);
                           return dim;
                       });
    }

    const auto& vv_hi = vv_ind.at(1);
Shucai Xiao's avatar
Shucai Xiao committed
102
    for(int start = 0; start < vec_dims.size(); start += vv_lo.size())
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
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
152
153
154
155
156
157
158
    {
        std::transform(vv_hi.begin(),
                       vv_hi.end(),
                       vec_dims.begin() + start,
                       std::back_inserter(vec_dims1),
                       [](auto i, auto dim) {
                           dim.push_back(i);
                           return dim;
                       });
    }

    return calc_neighbor_points(vvv_ind, i_dim + 1, vec_dims1, in_s);
}

static std::string get_coord_trans_mode(const onnx_parser::attribute_map& attr)
{
    std::string coord_trans_mode = "half_pixel";
    if(contains(attr, "coordinate_transformation_mode"))
    {
        coord_trans_mode = attr.at("coordinate_transformation_mode").s();
        // does not support transformation mode "tf_crop_and_resize"
        if(coord_trans_mode == "tf_crop_and_resize")
        {
            MIGRAPHX_THROW("PARSE_RESIZE: \"tf_crop_and_resize\" mode is not supported!");
        }
    }

    return coord_trans_mode;
}

static std::string get_mode(const onnx_parser::attribute_map& attr)
{
    std::string mode = "nearest";
    if(contains(attr, "mode"))
    {
        mode = attr.at("mode").s();
        if(mode != "nearest" and mode != "linear")
        {
            MIGRAPHX_THROW("PARSE_RESIZE: only nearest and linear modes are supported!");
        }
    }

    return mode;
}

static std::string get_nearest_mode(const onnx_parser::attribute_map& attr)
{
    std::string nearest_mode = "round_prefer_floor";
    if(contains(attr, "nearest_mode"))
    {
        nearest_mode = attr.at("nearest_mode").s();
    }

    return nearest_mode;
}

Paul Fultz II's avatar
Paul Fultz II committed
159
160
struct parse_resize : op_parser<parse_resize>
{
Shucai Xiao's avatar
Shucai Xiao committed
161
    std::vector<op_desc> operators() const { return {{"Resize"}, {"Upsample"}}; }
Paul Fultz II's avatar
Paul Fultz II committed
162

Shucai Xiao's avatar
Shucai Xiao committed
163
    instruction_ref parse(const op_desc& opd,
Paul Fultz II's avatar
Paul Fultz II committed
164
165
166
167
                          const onnx_parser& /*parser*/,
                          onnx_parser::node_info info,
                          std::vector<instruction_ref> args) const
    {
168
169
        // coord transform mode
        std::string coord_trans_mode = get_coord_trans_mode(info.attributes);
Paul Fultz II's avatar
Paul Fultz II committed
170

171
172
        // mode: only nearest and linear modes are supported for now
        std::string mode = get_mode(info.attributes);
Paul Fultz II's avatar
Paul Fultz II committed
173
174

        // nearest mode
175
        std::string nearest_mode = get_nearest_mode(info.attributes);
Paul Fultz II's avatar
Paul Fultz II committed
176
177

        // check exclude_outside, only support 0
178
179
        if(contains(info.attributes, "exclude_outside") and
           info.attributes.at("exclude_outside").i() == 1)
Paul Fultz II's avatar
Paul Fultz II committed
180
        {
Shucai Xiao's avatar
Shucai Xiao committed
181
            MIGRAPHX_THROW("PARSE_" + opd.op_name + ": exclude_outside 1 is not supported!");
Paul Fultz II's avatar
Paul Fultz II committed
182
183
184
185
186
187
188
        }

        // input data shape info
        auto in_s    = args[0]->get_shape();
        auto in_lens = in_s.lens();

        // output shape is explicitly specified
Shucai Xiao's avatar
Shucai Xiao committed
189
        std::vector<int> out_lens(in_lens.size());
Paul Fultz II's avatar
Paul Fultz II committed
190
191
192
193

        // scale
        std::vector<double> vec_scale;

194
        for(const auto& arg : args)
Paul Fultz II's avatar
Paul Fultz II committed
195
        {
196
            if(arg->name() == "undefined" or arg == args.front())
Paul Fultz II's avatar
Paul Fultz II committed
197
            {
198
                continue;
Paul Fultz II's avatar
Paul Fultz II committed
199
200
            }

201
202
203
204
205
206
            // skipped empty input
            auto lens = arg->get_shape().lens();
            if(lens.empty())
            {
                continue;
            }
Paul Fultz II's avatar
Paul Fultz II committed
207

208
209
210
            auto type = arg->get_shape().type();
            // output size
            if(type == shape::int64_type)
Paul Fultz II's avatar
Paul Fultz II committed
211
            {
212
                auto arg_out_s = arg->eval();
Shucai Xiao's avatar
Shucai Xiao committed
213
214
                check_arg_empty(arg_out_s,
                                "PARSE_" + opd.op_name + ": dynamic output size is not supported!");
215
216
217
218
                arg_out_s.visit([&](auto ol) { out_lens.assign(ol.begin(), ol.end()); });

                if(out_lens.size() != in_lens.size())
                {
Shucai Xiao's avatar
Shucai Xiao committed
219
220
                    MIGRAPHX_THROW("PARSE_" + opd.op_name +
                                   ": specified output size does not match input size");
221
222
223
224
225
226
227
228
229
                }

                // compute the scale
                vec_scale.resize(in_lens.size());
                std::transform(in_lens.begin(),
                               in_lens.end(),
                               out_lens.begin(),
                               vec_scale.begin(),
                               [](auto iss, auto oss) { return 1.0 * oss / iss; });
Paul Fultz II's avatar
Paul Fultz II committed
230
            }
231
232
233
234
235
236
237
238
            else
            {

                // scale input
                if(lens[0] == in_lens.size())
                {
                    auto arg_scale = arg->eval();
                    check_arg_empty(arg_scale,
Shucai Xiao's avatar
Shucai Xiao committed
239
240
                                    "PARSE_" + opd.op_name +
                                        ": dynamic input scale is not supported!");
Paul Fultz II's avatar
Paul Fultz II committed
241

242
243
244
                    arg_scale.visit([&](auto v) { vec_scale.assign(v.begin(), v.end()); });
                    if(in_lens.size() != vec_scale.size())
                    {
Shucai Xiao's avatar
Shucai Xiao committed
245
246
                        MIGRAPHX_THROW("PARSE_" + opd.op_name +
                                       ": ranks of input and scale are different!");
247
248
                    }

Shucai Xiao's avatar
Shucai Xiao committed
249
250
251
252
253
254
                    std::transform(
                        in_lens.begin(),
                        in_lens.end(),
                        vec_scale.begin(),
                        out_lens.begin(),
                        [&](auto idx, auto scale) { return static_cast<int>(idx * scale); });
255
256
                }
            }
Paul Fultz II's avatar
Paul Fultz II committed
257
258
259
        }

        shape out_s{in_s.type(), out_lens};
Shucai Xiao's avatar
Shucai Xiao committed
260
        int out_elements = out_s.elements();
Shucai Xiao's avatar
Shucai Xiao committed
261
        auto idx_op      = get_original_idx_op(coord_trans_mode);
Paul Fultz II's avatar
Paul Fultz II committed
262

263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
        // reshape input to one-dimension
        std::vector<int64_t> rsp_lens = {static_cast<int64_t>(in_s.elements())};
        args[0]                       = info.make_contiguous(args[0]);
        auto rsp = info.add_instruction(make_op("reshape", {{"dims", rsp_lens}}), args[0]);

        if(mode == "nearest")
        {
            std::vector<int> ind(out_elements);

            // map out_idx to in_idx
            auto nearest_op = get_nearest_op(nearest_mode);
            shape_for_each(out_s, [&](auto idx) {
                auto in_idx = idx;
                for(auto ii = 0; ii < in_lens.size(); ++ii)
                {
                    auto idx_val = idx_op(in_lens[ii], out_lens[ii], idx[ii], vec_scale[ii]);
                    in_idx[ii]   = nearest_op(in_lens[ii], idx_val);
                }

                ind[out_s.index(idx)] = static_cast<int64_t>(in_s.index(in_idx));
            });

            shape ind_s{shape::int32_type, out_lens};
            auto ins_ind = info.add_literal(literal(ind_s, ind));
            return info.add_instruction(make_op("gather", {{"axis", 0}}), rsp, ins_ind);
        }
        // linear mode
        else
        {
            auto nearest_floor = get_nearest_op("floor");
            auto nearest_ceil  = get_nearest_op("ceil");
Paul Fultz II's avatar
Paul Fultz II committed
294

295
            // get the number of dimensions
Shucai Xiao's avatar
Shucai Xiao committed
296
297
298
            int n_dim = out_lens.size();
            std::vector<std::vector<int>> vv_ind(2, std::vector<int>(out_elements));
            std::vector<std::vector<std::vector<int>>> vvv_ind(n_dim, vv_ind);
299
300
301
302
303
304
305
306
307
308
309
310
311
312
            std::vector<std::vector<float>> delta(n_dim, std::vector<float>(out_elements));

            shape_for_each(out_s, [&](auto idx) {
                auto in_idx  = idx;
                auto out_idx = out_s.index(idx);
                for(auto ii = 0; ii < in_lens.size(); ++ii)
                {
                    auto idx_val = idx_op(in_lens[ii], out_lens[ii], idx[ii], vec_scale[ii]);
                    vvv_ind[ii][0][out_idx] = nearest_floor(in_lens[ii], idx_val);
                    vvv_ind[ii][1][out_idx] = nearest_ceil(in_lens[ii], idx_val);
                    delta[ii][out_idx]      = idx_val - vvv_ind[ii][0][out_idx];
                }
            });

Shucai Xiao's avatar
Shucai Xiao committed
313
            std::vector<std::vector<int>> vec_dims(out_elements);
314
315
            auto ind      = calc_neighbor_points(vvv_ind, 0, vec_dims, in_s);
            auto ind_lens = out_lens;
Shucai Xiao's avatar
Shucai Xiao committed
316
            ind_lens[0] *= (int{1} << n_dim);
317
318
319
320
321
            shape ind_s{shape::int32_type, ind_lens};
            auto ins_ind = info.add_literal(literal(ind_s, ind));
            auto data    = info.add_instruction(make_op("gather", {{"axis", 0}}), rsp, ins_ind);

            auto dim_lens = out_lens;
Shucai Xiao's avatar
Shucai Xiao committed
322
323
            dim_lens[0] *= (int{1} << (n_dim - 1));
            for(int i = 0; i < n_dim; ++i)
Paul Fultz II's avatar
Paul Fultz II committed
324
            {
325
326
327
                shape dim_s{shape::float_type, dim_lens};
                const auto& dim_delta = delta[n_dim - i - 1];
                std::vector<float> delta_data;
Shucai Xiao's avatar
Shucai Xiao committed
328
                for(int j = 0; j < dim_lens[0] / out_lens[0]; ++j)
329
330
331
332
                {
                    delta_data.insert(delta_data.begin(), dim_delta.begin(), dim_delta.end());
                }
                auto ins_delta = info.add_literal(dim_s, delta_data);
Paul Fultz II's avatar
Paul Fultz II committed
333

334
                // slice the data
Paul Fultz II's avatar
Paul Fultz II committed
335
                int64_t slc_stride = dim_lens[0];
336
337
338
339
340
341
342
343
344
345
346
347
                auto low           = info.add_instruction(
                    make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {slc_stride}}}),
                    data);
                auto hi = info.add_instruction(
                    make_op("slice",
                            {{"axes", {0}}, {"starts", {slc_stride}}, {"ends", {2 * slc_stride}}}),
                    data);
                auto diff = info.add_instruction(make_op("sub"), hi, low);
                auto ddf  = info.add_instruction(make_op("mul"), diff, ins_delta);
                data      = info.add_instruction(make_op("add"), ddf, low);
                dim_lens[0] /= 2;
            }
Paul Fultz II's avatar
Paul Fultz II committed
348

349
350
            return data;
        }
Paul Fultz II's avatar
Paul Fultz II committed
351
352
353
354
    }
};

} // namespace onnx
355

Paul Fultz II's avatar
Paul Fultz II committed
356
357
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx