lowering.cpp 21.8 KB
Newer Older
Paul's avatar
Paul committed
1

Paul's avatar
Paul committed
2
3
4
#include <migraphx/cpu/lowering.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/dfor.hpp>
5
#include <migraphx/op/identity.hpp>
6
#include <migraphx/op/batch_norm_inference.hpp>
Paul's avatar
Paul committed
7
#include <migraphx/op/convolution.hpp>
kahmed10's avatar
kahmed10 committed
8
#include <migraphx/op/deconvolution.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
9
#include <migraphx/op/quant_convolution.hpp>
Paul's avatar
Paul committed
10
#include <migraphx/op/dot.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
11
#include <migraphx/op/quant_dot.hpp>
Paul's avatar
Paul committed
12
13
14
15
16
17
18
19
#include <migraphx/op/elu.hpp>
#include <migraphx/op/im2col.hpp>
#include <migraphx/op/leaky_relu.hpp>
#include <migraphx/op/logsoftmax.hpp>
#include <migraphx/op/lrn.hpp>
#include <migraphx/op/pad.hpp>
#include <migraphx/op/pooling.hpp>
#include <migraphx/op/softmax.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
20
21
#include <migraphx/op/argmax.hpp>
#include <migraphx/op/argmin.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
22
#include <migraphx/op/rnn_var_sl_last_output.hpp>
Paul's avatar
Paul committed
23
24
#include <migraphx/shape_for_each.hpp>
#include <migraphx/iterator_for.hpp>
Paul's avatar
Paul committed
25
#include <migraphx/par_dfor.hpp>
26
#include <migraphx/clamp.hpp>
27
#include <migraphx/cpu/context.hpp>
28
#include <migraphx/register_op.hpp>
29
#include <migraphx/make_op.hpp>
30
#include <migraphx/program.hpp>
31
#include <migraphx/tune_axis.hpp>
Paul's avatar
Paul committed
32
#include <unordered_map>
Paul's avatar
Paul committed
33
#include <utility>
kahmed10's avatar
kahmed10 committed
34
#include <iostream>
Paul's avatar
Paul committed
35

Paul's avatar
Paul committed
36
namespace migraphx {
Paul's avatar
Paul committed
37
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
38
39
40
41
42
43
44
45
namespace cpu {

template <typename T>
T zero(const T&)
{
    return T(0);
}

Khalique's avatar
Khalique committed
46
47
48
49
template <class T>
typename std::conditional_t<std::is_integral<T>{}, std::make_signed<T>, std::enable_if<true, T>>::
    type
    make_signed(T x)
Khalique's avatar
Khalique committed
50
51
52
53
{
    return x;
}

Khalique's avatar
Khalique committed
54
struct cpu_lrn
Khalique's avatar
Khalique committed
55
{
Khalique's avatar
Khalique committed
56
    op::lrn op;
Khalique's avatar
Khalique committed
57

58
59
60
61
62
63
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Khalique's avatar
Khalique committed
64
    std::string name() const { return "cpu::lrn"; }
Khalique's avatar
Khalique committed
65
66
67
68
69
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
Khalique's avatar
Khalique committed
70
71
72
73
            int n_batch         = output_shape.lens()[0];
            int channels        = output_shape.lens()[1];
            int height          = output_shape.lens()[2];
            int width           = output_shape.lens()[3];
Paul's avatar
Paul committed
74
            float alphaoverarea = op.alpha / float(op.size);
75
76
            int radius_lower    = (op.size - 1) / 2;
            int radius_upper    = op.size / 2 + 1;
Khalique's avatar
Khalique committed
77

78
            par_dfor(n_batch, height, width)([&](int b, int h, int w) {
Khalique's avatar
Khalique committed
79
                float scale = 0;
Khalique's avatar
Khalique committed
80
                dfor(channels)([&](int c) {
81
82
                    auto start = (c - radius_lower) < 0 ? 0 : (c - radius_lower);
                    auto end   = (c + radius_upper) > channels ? channels : (c + radius_upper);
Khalique's avatar
Khalique committed
83
84
                    for(auto k = start; k < end; ++k)
                    {
Khalique's avatar
Khalique committed
85
                        scale += std::pow(input(b, k, h, w), 2);
Khalique's avatar
Khalique committed
86
87
88
                    }
                    scale *= alphaoverarea;
                    scale += op.bias;
Khalique's avatar
Khalique committed
89
                    scale              = std::pow(scale, -op.beta);
Khalique's avatar
Khalique committed
90
91
92
93
94
95
96
                    output(b, c, h, w) = input(b, c, h, w) * scale;
                });
            });
        });
        return result;
    }
};
97
MIGRAPHX_REGISTER_OP(cpu_lrn)
Khalique's avatar
Khalique committed
98

kahmed10's avatar
kahmed10 committed
99
template <class Op>
100
struct cpu_deconvolution : auto_register_op<cpu_deconvolution<Op>>
kahmed10's avatar
kahmed10 committed
101
{
102
103
104
105
    cpu_deconvolution() = default;

    cpu_deconvolution(Op pop) : op(std::move(pop)) {}

kahmed10's avatar
kahmed10 committed
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
    Op op;

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

    std::string name() const { return "cpu::" + op.name(); }
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        visit_all(result, args[0], args[1])([&](auto output, auto input, auto weights) {
            using type = typename decltype(output)::value_type;

            std::fill(output.begin(), output.end(), type{0});

kahmed10's avatar
kahmed10 committed
124
125
126
            auto in_lens = input.get_shape().lens();
            auto in_n    = in_lens[0];
            auto in_c    = in_lens[1];
kahmed10's avatar
kahmed10 committed
127
128
129
130

            auto wei   = weights.get_shape().lens();
            auto wei_n = wei[0];
            auto wei_c = wei[1];
kahmed10's avatar
kahmed10 committed
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
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184

            auto out_lens = output_shape.lens();
            auto kdims    = op.kdims();

            std::vector<std::size_t> win_size{in_c};
            std::copy(in_lens.begin() + 2, in_lens.end(), std::back_inserter(win_size));
            std::copy(wei.begin() + 2, wei.end(), std::back_inserter(win_size));
            shape win_shape{output_shape.type(), win_size};

            par_dfor(in_n, wei_c)([&](int o, int k) {

                shape_for_each(win_shape, [&](auto idx_win) {
                    const int w = idx_win[0];

                    auto input_dims_start = idx_win.begin() + 1;
                    auto wei_dims_start   = idx_win.begin() + kdims + 1;

                    std::vector<std::ptrdiff_t> win_start;
                    for(std::size_t n = 0; n < kdims; ++n)
                    {
                        win_start.push_back(std::ptrdiff_t(*(input_dims_start + n) * op.stride[n]) -
                                            std::ptrdiff_t(op.padding[n]));
                    }

                    const int group_id = w / (wei_n / op.group);
                    const int in_ch    = group_id * wei_c + k;

                    std::vector<std::ptrdiff_t> idx_out{o, in_ch};

                    for(size_t n = 0; n < kdims; n++)
                    {
                        idx_out.push_back(win_start[n] + *(wei_dims_start + n) * op.dilation[n]);
                    }

                    std::vector<std::ptrdiff_t> idx_wei{w, k};
                    std::copy(wei_dims_start, idx_win.end(), std::back_inserter(idx_wei));

                    std::vector<std::ptrdiff_t> idx_in{o, w};
                    std::copy(input_dims_start, wei_dims_start, std::back_inserter(idx_in));

                    if(std::all_of(
                           idx_out.begin() + 2, idx_out.end(), [&](auto ii) { return ii >= 0; }) and
                       std::equal(idx_out.begin() + 2,
                                  idx_out.end(),
                                  out_lens.begin() + 2,
                                  out_lens.end(),
                                  std::less<std::ptrdiff_t>{}))
                    {
                        output(idx_out.begin(), idx_out.end()) +=
                            input(idx_in.begin(), idx_in.end()) *
                            weights(idx_wei.begin(), idx_wei.end());
                    }
                });

kahmed10's avatar
kahmed10 committed
185
            });
kahmed10's avatar
kahmed10 committed
186

kahmed10's avatar
kahmed10 committed
187
188
189
190
        });
        return result;
    }
};
191
template struct cpu_deconvolution<op::deconvolution>;
kahmed10's avatar
kahmed10 committed
192

Scott Thornton's avatar
Scott Thornton committed
193
194
struct cpu_im2col
{
195
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
196

197
198
199
200
201
202
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Scott Thornton's avatar
Scott Thornton committed
203
204
    static std::string name() { return "cpu::im2col"; }
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Scott Thornton's avatar
Scott Thornton committed
205

wsttiger's avatar
wsttiger committed
206
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
207
    {
Scott Thornton's avatar
Scott Thornton committed
208
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
209
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
210
211
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
212
213
            const std::size_t& height   = input_shape.lens()[2];
            const std::size_t& width    = input_shape.lens()[3];
Scott Thornton's avatar
Scott Thornton committed
214
215
216
            const std::size_t& channels = weights_shape.lens()[1];
            const std::size_t& kernel_h = weights_shape.lens()[2];
            const std::size_t& kernel_w = weights_shape.lens()[3];
Scott Thornton's avatar
Scott Thornton committed
217
218
            const std::size_t& pad_h    = op.padding[0];
            const std::size_t& pad_w    = op.padding[1];
Scott Thornton's avatar
Scott Thornton committed
219
220
221
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

Paul's avatar
Paul committed
222
223
            long kdiv2_h = long(kernel_h) / 2;
            long kdiv2_w = long(kernel_w) / 2;
Scott Thornton's avatar
Scott Thornton committed
224
            // calculate output sizes
Scott Thornton's avatar
Scott Thornton committed
225
226
            const std::size_t col_height = (height - kernel_h + 2 * pad_h) / stride_h + 1;
            const std::size_t col_width  = (width - kernel_w + 2 * pad_w) / stride_w + 1;
wsttiger's avatar
wsttiger committed
227
            // account for padding for the starting position of the input pixels
Paul's avatar
Paul committed
228
            long iinput = kdiv2_h - long(pad_h);
wsttiger's avatar
wsttiger committed
229
            // loop over output pixels (ioutput, joutput)
Scott Thornton's avatar
Scott Thornton committed
230
231
            for(std::size_t ioutput = 0; ioutput < col_height; ioutput++, iinput += stride_h)
            {
Paul's avatar
Paul committed
232
                long jinput = kdiv2_w - long(pad_w);
Scott Thornton's avatar
Scott Thornton committed
233
234
235
236
237
                for(std::size_t joutput = 0; joutput < col_width; joutput++, jinput += stride_w)
                {
                    // compute linear index for output
                    std::size_t ldx = ioutput * col_width + joutput;
                    std::size_t p   = 0;
wsttiger's avatar
wsttiger committed
238
239
240
                    dfor(channels,
                         kernel_h,
                         kernel_w)([&](std::size_t c, std::size_t koffset, std::size_t loffset) {
Paul's avatar
Paul committed
241
242
                        auto idx    = iinput + long(koffset) - kdiv2_h;
                        auto jdx    = jinput + long(loffset) - kdiv2_w;
wsttiger's avatar
wsttiger committed
243
244
245
246
247
                        col(ldx, p) = ((idx >= 0) && (idx < height) && (jdx >= 0) && (jdx < width))
                                          ? input(0, c, idx, jdx)
                                          : 0;
                        p++;
                    });
Scott Thornton's avatar
Scott Thornton committed
248
249
                }
            }
Scott Thornton's avatar
Scott Thornton committed
250
        });
Scott Thornton's avatar
Scott Thornton committed
251
252
253
        return result;
    }
};
254
MIGRAPHX_REGISTER_OP(cpu_im2col)
Scott Thornton's avatar
Scott Thornton committed
255

256
struct cpu_op
Paul's avatar
Paul committed
257
{
258
    operation op = op::identity{};
kahmed10's avatar
kahmed10 committed
259
260
261
262
263
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }
264
    std::string name() const { return "cpu::op"; }
Paul's avatar
Paul committed
265
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
266
    argument compute(context&, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
267
    {
Paul's avatar
Paul committed
268
        return op.compute(output_shape, args);
Paul's avatar
Paul committed
269
    }
270
271
272
273
274
275
276
277
278
279
280
    value to_value() const
    {
        value v;
        v["name"]     = op.name();
        v["operator"] = op.to_value();
        return v;
    }
    void from_value(const value& v)
    {
        op = make_op(v.at("name").to<std::string>(), v.at("operator"));
    }
281
    friend std::ostream& operator<<(std::ostream& os, const cpu_op& x)
Paul's avatar
Paul committed
282
    {
283
284
        os << "cpu::" << x.op;
        return os;
Paul's avatar
Paul committed
285
286
    }
};
287
MIGRAPHX_REGISTER_OP(cpu_op)
Paul's avatar
Paul committed
288

Khalique's avatar
Khalique committed
289
struct cpu_pad
290
{
Khalique's avatar
Khalique committed
291
    op::pad op;
292
293
294
295
296
297
298

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

kahmed10's avatar
kahmed10 committed
299
    std::string name() const { return "cpu::pad"; }
300
301
302
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
    {
Khalique's avatar
Khalique committed
303
        assert(output_shape.standard());
304
        argument result{output_shape};
305
306
307
308
        result.visit([&](auto output) {
            using type = typename decltype(output)::value_type;
            std::fill(output.begin(), output.end(), pad_clamp<type>(op.value));
        });
Khalique's avatar
Khalique committed
309
310

        visit_all(result, args[0])([&](auto output, auto input) {
311
            shape_for_each(input.get_shape(), [&](const auto& idx) {
Khalique's avatar
Khalique committed
312
                std::vector<std::size_t> new_idx(idx.size());
Khalique's avatar
Khalique committed
313
314
315
316
                std::transform(
                    idx.begin(), idx.end(), op.pads.begin(), new_idx.begin(), [](auto i, auto j) {
                        return i + j;
                    });
Khalique's avatar
Khalique committed
317
                output(new_idx.begin(), new_idx.end()) = input(idx.begin(), idx.end());
318
            });
Khalique's avatar
Khalique committed
319
320
        });

321
322
323
        return result;
    }
};
324
MIGRAPHX_REGISTER_OP(cpu_pad)
325

Khalique's avatar
Khalique committed
326
327
328
329
330
331
struct leaky_relu_op
{
    op::leaky_relu op;
    std::string name() const { return "cpu::leaky_relu"; }
    auto fcn() const
    {
Paul's avatar
Paul committed
332
        auto a = op.alpha;
Khalique's avatar
Khalique committed
333
334
335
336
        return [a](auto x) { return x > 0 ? x : x * a; };
    }
};

Khalique's avatar
Khalique committed
337
338
339
340
341
342
struct elu_op
{
    op::elu op;
    std::string name() const { return "cpu::elu"; }
    auto fcn() const
    {
Paul's avatar
Paul committed
343
        auto a = op.alpha;
Khalique's avatar
Khalique committed
344
345
346
347
        return [a](auto x) { return x > 0 ? x : a * std::expm1(x); };
    }
};

Paul's avatar
Paul committed
348
template <typename Op>
349
struct cpu_unary2 : auto_register_op<cpu_unary2<Op>>
Paul's avatar
Paul committed
350
{
351
    cpu_unary2() = default;
352
353

    template <class T>
354
    cpu_unary2(T pop) : op(Op{std::move(pop)})
355
356
357
    {
    }

Paul's avatar
Paul committed
358
    Op op;
359
360
361
362
363
364

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op.op, f);
    }
Paul's avatar
Paul committed
365
    std::string name() const { return op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
366
    shape compute_shape(const std::vector<shape>& inputs) const
367
    {
368
        check_shapes{inputs, *this}.has(1);
Shucai Xiao's avatar
Shucai Xiao committed
369
        auto s = inputs.at(0);
370
        return {s.type(), s.lens()};
371
372
    }

Paul's avatar
Paul committed
373
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
374
375
    {
        argument result{output_shape};
376
377
378
        visit_all(result, args[0])([&](auto output, auto input) {
            assert(input.get_shape().standard());
            std::transform(input.begin(), input.end(), output.begin(), op.fcn());
Paul's avatar
Paul committed
379
        });
380

Paul's avatar
Paul committed
381
382
383
        return result;
    }
};
384
385
template struct cpu_unary2<leaky_relu_op>;
template struct cpu_unary2<elu_op>;
Paul's avatar
Paul committed
386

387
template <class Op>
388
struct cpu_softmax : auto_register_op<cpu_softmax<Op>>
Paul's avatar
Paul committed
389
{
390
391
392
393
    cpu_softmax() = default;

    cpu_softmax(Op pop) : op(std::move(pop)) {}

394
    Op op;
Khalique's avatar
Khalique committed
395
396
397
398
399
400
401

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

402
    std::string name() const { return "cpu::" + op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
403
404
    shape compute_shape(const std::vector<shape>& inputs) const
    {
405
        check_shapes{inputs, *this}.has(1).standard();
Shucai Xiao's avatar
Shucai Xiao committed
406
407
        return op.normalize_compute_shape(inputs);
    }
Paul's avatar
Paul committed
408
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
409
410
    {
        argument result{output_shape};
411
412
413
        auto batch_lens        = output_shape.lens();
        int64_t tuned_axis     = tune_axis(args[0].get_shape().lens().size(), op.axis, op.name());
        std::size_t n_dims     = batch_lens[tuned_axis];
414
        batch_lens[tuned_axis] = 1;
415
416
        shape batch_shape{shape::int32_type, batch_lens};

Paul's avatar
Paul committed
417
418
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
Shucai Xiao's avatar
Shucai Xiao committed
419
420
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
421
422
            std::vector<value_type> batch_sum(batch_shape.elements(), value_type(0));
            par_for(batch_shape.elements(), [&](auto i) {
423
                auto idx = batch_shape.multi(i);
Shucai Xiao's avatar
Shucai Xiao committed
424
                for(std::size_t j = 0; j < n_dims; ++j)
425
                {
426
427
                    idx[tuned_axis] = j;
                    batch_max[i]    = std::max(batch_max[i], input(idx.begin(), idx.end()));
428
                }
Khalique's avatar
Khalique committed
429

Shucai Xiao's avatar
Shucai Xiao committed
430
                for(std::size_t j = 0; j < n_dims; ++j)
431
                {
432
                    idx[tuned_axis]   = j;
Shucai Xiao's avatar
Shucai Xiao committed
433
434
                    std::size_t index = output_shape.index(idx);
                    output[index]     = std::exp(input[index] - batch_max[i]);
435
                }
Khalique's avatar
Khalique committed
436

Shucai Xiao's avatar
Shucai Xiao committed
437
                for(std::size_t j = 0; j < n_dims; ++j)
438
                {
439
                    idx[tuned_axis] = j;
440
441
                    batch_sum[i] += output(idx.begin(), idx.end());
                }
Khalique's avatar
Khalique committed
442

Shucai Xiao's avatar
Shucai Xiao committed
443
                for(std::size_t j = 0; j < n_dims; ++j)
444
                {
445
                    idx[tuned_axis] = j;
446
447
                    output(idx.begin(), idx.end()) =
                        op.output()(output(idx.begin(), idx.end()), batch_sum[i]);
448
                }
Shucai Xiao's avatar
Shucai Xiao committed
449
450
451
452
453
454
            });
        });

        return result;
    }
};
455
456
template struct cpu_softmax<op::softmax>;
template struct cpu_softmax<op::logsoftmax>;
Shucai Xiao's avatar
Shucai Xiao committed
457

Shucai Xiao's avatar
Shucai Xiao committed
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
struct cpu_rnn_var_sl_last_output
{
    op::rnn_var_sl_last_output op;

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

    std::string name() const { return "cpu::rnn_var_sl_last_output"; }

    shape compute_shape(std::vector<shape> inputs) const
    {
        return op.compute_shape(std::move(inputs));
    }

    argument compute(const shape& output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        auto out_comp_lens = args[0].get_shape().lens();
        out_comp_lens[0]   = 1;
        shape out_comp_s{output_shape.type(), out_comp_lens};

        visit_all(result, args[0])([&](auto output, auto input) {
            args[1].visit([&](auto seq_lens) {
                par_for(output_shape.elements(), [&](auto i) {
                    auto idx = out_comp_s.multi(i);
                    auto b   = idx[2];
                    if(op.direction == op::rnn_direction::reverse or idx[1] == 1)
                    {
                        idx[0] = 0;
                    }
                    else
                    {
                        idx[0] = seq_lens[b] - 1;
                    }
                    output[i] = input(idx.begin(), idx.end());
                });
            });
        });

        return result;
    }
};
503
MIGRAPHX_REGISTER_OP(cpu_rnn_var_sl_last_output)
Shucai Xiao's avatar
Shucai Xiao committed
504

505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
struct cpu_literal
{
    argument data;

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return pack(f(self.data, "data"));
    }

    std::string name() const { return "cpu::literal"; }

    shape compute_shape(const std::vector<shape>&) const { return data.get_shape(); }

    argument compute(const shape&, const std::vector<argument>&) const { return data; }

    friend std::ostream& operator<<(std::ostream& os, const cpu_literal& x)
    {
        os << x.name();
        return os;
    }
};

Paul's avatar
Paul committed
528
529
struct cpu_apply
{
Shucai Xiao's avatar
Shucai Xiao committed
530
    module* modl;
531
532
533
    std::unordered_map<std::string, std::function<instruction_ref(instruction_ref)>> apply_map{};
    std::unordered_map<instruction_ref, std::string> prog_output_names{};
    instruction_ref last{};
Paul's avatar
Paul committed
534

535
536
    void create_output_names()
    {
Shucai Xiao's avatar
Shucai Xiao committed
537
        this->last = instruction::get_output_alias(std::prev(modl->end()));
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
        if(this->last->name() == "@return")
        {
            const auto& prog_outputs = last->inputs();
            std::vector<instruction_ref> outputs_alias(prog_outputs.size());

            std::transform(prog_outputs.begin(),
                           prog_outputs.end(),
                           outputs_alias.begin(),
                           [](const auto& i) { return instruction::get_output_alias(i); });

            std::size_t index = 0;
            for(auto ins : outputs_alias)
            {
                prog_output_names[ins] = "#output_" + std::to_string(index++);
            }
        }
    }

    void extend_op(const std::string& op_name, const std::string& cpu_name, bool allocate = false)
    {
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op = ins->get_operator();
            if(allocate)
                replace(ins, make_op(cpu_name, op.to_value()));
Shucai Xiao's avatar
Shucai Xiao committed
562
            return modl->replace_instruction(ins, make_op(cpu_name, op.to_value()), ins->inputs());
563
564
565
        });
    }

Paul's avatar
Paul committed
566
567
    void init()
    {
568
        create_output_names();
569
570
571
572
573
574
        extend_op("add", "dnnl::add", true);
        extend_op("mul", "dnnl::mul", true);
        extend_op("convolution", "dnnl::convolution", true);
        extend_op("dot", "dnnl::dot", true);
        extend_op("relu", "dnnl::relu", true);

575
576
577
578
579
580
581
582
583
584
585
586
        extend_op("contiguous", "cpu::contiguous", true);
        extend_op("deconvolution", "cpu::deconvolution");
        extend_op("elu", "cpu::elu");
        extend_op("im2col", "cpu::im2col");
        extend_op("leaky_relu", "cpu::leaky_relu");
        extend_op("logsoftmax", "cpu::logsoftmax");
        extend_op("lrn", "cpu::lrn");
        extend_op("pad", "cpu::pad");
        extend_op("quant_convolution", "cpu::quant_convolution", true);
        extend_op("quant_dot", "cpu::quant_dot", true);
        extend_op("rnn_var_sl_last_output", "cpu::rnn_var_sl_last_output");
        extend_op("softmax", "cpu::softmax");
Paul's avatar
Paul committed
587
588
589
590
591
    }

    void apply()
    {
        init();
Shucai Xiao's avatar
Shucai Xiao committed
592
        for(auto it : iterator_for(*modl))
Paul's avatar
Paul committed
593
        {
594
595
596
597
598
            if(it->name() == "@literal")
            {
                apply_literal(it);
            }
            else if(it->name() == "pooling")
Paul's avatar
Paul committed
599
600
601
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
602
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
603
            {
Paul's avatar
Paul committed
604
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
605
606
607
608
            }
        }
    }

609
    instruction_ref apply_literal(instruction_ref ins) const
610
    {
Shucai Xiao's avatar
Shucai Xiao committed
611
        return modl->replace_instruction(ins, cpu_literal{ins->get_literal().get_argument()});
612
613
    }

614
    instruction_ref apply_pooling(instruction_ref ins)
Paul's avatar
Paul committed
615
    {
616
617
618
619
620
621
622
623
624
625
626
        auto&& op = ins->get_operator();
        auto v    = op.to_value();
        if(has_op("dnnl::pooling") and ins->get_shape().type() == shape::type_t::float_type and
           not v["ceil_mode"].to<bool>())
            return replace(ins, make_op("dnnl::pooling", op.to_value()));
        std::string mode = v["mode"].to<std::string>();
        if(mode == "max")
            return replace(ins, make_op("cpu::pooling_max", v));
        else if(mode == "average")
            return replace(ins, make_op("cpu::pooling_average", v));
        return ins;
Paul's avatar
Paul committed
627
628
    }

629
    instruction_ref replace(instruction_ref ins, const operation& op)
Paul's avatar
Paul committed
630
    {
631
632
        auto inputs = ins->inputs();
        inputs.push_back(insert_allocation(ins, ins->get_shape()));
Shucai Xiao's avatar
Shucai Xiao committed
633
        return modl->replace_instruction(ins, op, inputs);
Paul's avatar
Paul committed
634
635
    }

636
    instruction_ref insert_allocation(instruction_ref ins, const shape& s)
Paul's avatar
Paul committed
637
    {
638
639
640
        auto ins_alias = instruction::get_output_alias(ins);
        if(last->name() == "@return" and prog_output_names.count(ins_alias) > 0)
        {
Shucai Xiao's avatar
Shucai Xiao committed
641
            return modl->add_parameter(prog_output_names[ins_alias], s);
642
643
644
        }
        else if(ins == last)
        {
Shucai Xiao's avatar
Shucai Xiao committed
645
            return modl->add_parameter("output", s);
646
647
        }

Shucai Xiao's avatar
Shucai Xiao committed
648
        return modl->insert_instruction(ins, make_op("cpu::allocate", {{"shape", to_value(s)}}));
Paul's avatar
Paul committed
649
650
651
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
652
void lowering::apply(module& m) const { cpu_apply{&m}.apply(); }
Paul's avatar
Paul committed
653
654

} // namespace cpu
Paul's avatar
Paul committed
655
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
656
} // namespace migraphx