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

Paul's avatar
Paul committed
2
3
4
5
6
7
#include <migraphx/cpu/lowering.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/dfor.hpp>
#include <migraphx/operators.hpp>
#include <migraphx/shape_for_each.hpp>
#include <migraphx/iterator_for.hpp>
Paul's avatar
Paul committed
8
#include <migraphx/par_dfor.hpp>
Paul's avatar
Paul committed
9
#include <migraphx/cpu/gemm.hpp>
Paul's avatar
Paul committed
10
#include <unordered_map>
Paul's avatar
Paul committed
11
#include <utility>
Paul's avatar
Paul committed
12

Paul's avatar
Paul committed
13
namespace migraphx {
Paul's avatar
Paul committed
14
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
15
16
17
18
19
20
21
22
namespace cpu {

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

Khalique's avatar
Khalique committed
23
24
25
26
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
27
28
29
30
{
    return x;
}

31
32
33
34
//
// cpu implemenataion of batch norm for inference
//
// inputs are:
35
36
37
38
// args[0] -> input data buffer
// args[1] -> mini batch mean
// args[2] -> mini batch variance
// args[3] -> gamma
Aditya Atluri's avatar
Aditya Atluri committed
39
// args[4] -> bias
40
41
42
//
// The equation to compute batch norm for inference is:
//
Aditya Atluri's avatar
Aditya Atluri committed
43
// output[i] = bias + gamma * (input[i] + mean) / sqrt(variance + epsilon)
44
45
46
47
48
//
// the input data format should be nchw
//
struct cpu_batch_norm_inference
{
49
    op::batch_norm_inference op;
50

51
52
    std::string name() const { return "cpu::batch_norm_inference"; }

Paul's avatar
Paul committed
53
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
54

Paul's avatar
Paul committed
55
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
56
    {
57
58
        argument output{output_shape};

Aditya Atluri's avatar
Aditya Atluri committed
59
60
        double epsilon           = op.epsilon;
        auto input               = args[0];
Paul's avatar
Paul committed
61
62
63
64
        auto arg_gamma           = args[1];
        auto arg_bias            = args[2];
        auto mini_batch_mean     = args[3];
        auto mini_batch_variance = args[4];
65

66
        auto num_batch    = output_shape.lens()[0];
Aditya Atluri's avatar
Aditya Atluri committed
67
68
        auto num_channels = output_shape.lens()[1];
        auto image_height = output_shape.lens()[2];
69
        auto image_width  = output_shape.lens()[3];
Aditya Atluri's avatar
Aditya Atluri committed
70

71
        if(op.bn_mode == op::batch_norm_inference::spatial)
Scott Thornton's avatar
Scott Thornton committed
72
73
74
75
        {
            visit_all(output, input, mini_batch_mean, mini_batch_variance, arg_gamma, arg_bias)(
                [&](auto result, auto buffer, auto mean, auto variance, auto gamma, auto bias) {

Paul's avatar
Paul committed
76
                    par_dfor(num_batch, num_channels, image_height, image_width)(
Scott Thornton's avatar
Scott Thornton committed
77
                        [&](std::size_t n, std::size_t c, std::size_t h, std::size_t w) {
Paul's avatar
Paul committed
78
                            assert((variance(c) + epsilon) > 0);
Scott Thornton's avatar
Scott Thornton committed
79
80
81
82
83
                            result(n, c, h, w) = gamma(c) * (buffer(n, c, h, w) - mean(c)) /
                                                     std::sqrt(variance(c) + epsilon) +
                                                 bias(c);
                        });
                });
84
85
        }

86
        if(op.bn_mode == op::batch_norm_inference::per_activation)
Scott Thornton's avatar
Scott Thornton committed
87
        {
88
89
90
            visit_all(output, input, mini_batch_mean, mini_batch_mean, arg_gamma, arg_bias)(
                [&](auto result, auto buffer, auto mean, auto variance, auto gamma, auto bias) {

Paul's avatar
Paul committed
91
                    par_dfor(num_batch, num_channels, image_height, image_width)(
92
                        [&](std::size_t n, std::size_t c, std::size_t h, std::size_t w) {
Paul's avatar
Paul committed
93
                            assert((variance(c, h, w) + epsilon) > 0);
Scott Thornton's avatar
Scott Thornton committed
94
95
96
97
98
99
                            result(n, c, h, w) = gamma(c, h, w) *
                                                     (buffer(n, c, h, w) - mean(c, h, w)) /
                                                     std::sqrt(variance(c, h, w) + epsilon) +
                                                 bias(c, h, w);
                        });
                });
100
        }
101
102
103
104
105

        return output;
    }
};

Paul's avatar
Paul committed
106
107
struct cpu_convolution
{
108
    op::convolution op;
Paul's avatar
Paul committed
109
110

    std::string name() const { return "cpu::convolution"; }
Paul's avatar
Paul committed
111
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
112
113
114
115
    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) {
Khalique's avatar
Khalique committed
116
            auto in   = input.get_shape().lens();
Khalique's avatar
Khalique committed
117
118
            auto in_h = in[2];
            auto in_w = in[3];
Paul's avatar
Paul committed
119

Khalique's avatar
Khalique committed
120
            auto wei   = weights.get_shape().lens();
Khalique's avatar
Khalique committed
121
122
123
124
            auto wei_n = wei[0];
            auto wei_c = wei[1];
            auto wei_h = wei[2];
            auto wei_w = wei[3];
Paul's avatar
Paul committed
125

Paul's avatar
Paul committed
126
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
127
128
129
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
130
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
Khalique's avatar
Khalique committed
131
132
                    const int start_x  = i * op.stride[0] - op.padding[0];
                    const int start_y  = j * op.stride[1] - op.padding[1];
Khalique's avatar
Khalique committed
133
                    const int group_id = w / (wei_n / op.group);
Paul's avatar
Paul committed
134
135
136

                    double acc = 0;
                    dfor(wei_c, wei_h, wei_w)([&](std::size_t k, std::size_t x, std::size_t y) {
Khalique's avatar
Khalique committed
137
138
                        const int in_x  = start_x + x;
                        const int in_y  = start_y + y;
Khalique's avatar
Khalique committed
139
                        const int in_ch = group_id * wei_c + k;
Paul's avatar
Paul committed
140
141
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
Khalique's avatar
Khalique committed
142
                            acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
Paul's avatar
Paul committed
143
144
145
146
147
148
149
150
151
                        }
                    });
                    output(o, w, i, j) = acc;
                });
        });
        return result;
    }
};

Scott Thornton's avatar
Scott Thornton committed
152
153
struct cpu_im2col
{
154
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
155

Scott Thornton's avatar
Scott Thornton committed
156
157
    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
158

wsttiger's avatar
wsttiger committed
159
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
160
    {
Scott Thornton's avatar
Scott Thornton committed
161
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
162
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
163
164
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
165
166
            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
167
168
169
            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
170
171
            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
172
173
174
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

Paul's avatar
Paul committed
175
176
            int kdiv2_h;
            int kdiv2_w;
Scott Thornton's avatar
Scott Thornton committed
177
178
            kdiv2_h = kernel_h / 2;
            kdiv2_w = kernel_w / 2;
Scott Thornton's avatar
Scott Thornton committed
179
            // calculate output sizes
Scott Thornton's avatar
Scott Thornton committed
180
181
            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
182
            // account for padding for the starting position of the input pixels
Scott Thornton's avatar
Scott Thornton committed
183
            std::size_t iinput = kdiv2_h - pad_h;
wsttiger's avatar
wsttiger committed
184
            // loop over output pixels (ioutput, joutput)
Scott Thornton's avatar
Scott Thornton committed
185
186
187
188
189
190
191
192
            for(std::size_t ioutput = 0; ioutput < col_height; ioutput++, iinput += stride_h)
            {
                std::size_t jinput = kdiv2_w - pad_w;
                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
193
194
195
196
197
198
199
200
201
202
                    dfor(channels,
                         kernel_h,
                         kernel_w)([&](std::size_t c, std::size_t koffset, std::size_t loffset) {
                        int idx     = iinput + koffset - kdiv2_h;
                        int jdx     = jinput + loffset - kdiv2_w;
                        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
203
204
                }
            }
Scott Thornton's avatar
Scott Thornton committed
205
        });
Scott Thornton's avatar
Scott Thornton committed
206
207
208
209
        return result;
    }
};

Paul's avatar
Paul committed
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
struct max_pool
{
    static std::string name() { return "max"; }
    static double start() { return std::numeric_limits<double>::lowest(); }

    static double apply(double x, double y)
    {
        double m = std::max(x, y);
        return (m);
    }

    static double final(double x, double) { return (x); }
};

struct avg_pool
{
    static std::string name() { return "average"; }
    static double start() { return 0.0; }

    static double apply(double x, double y) { return x + y; }

    static double final(double x, double y) { return x / y; }
};

template <class Op>
struct cpu_pooling
{
237
    op::pooling op;
Paul's avatar
Paul committed
238
239

    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
240
241
    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
Paul's avatar
Paul committed
242
243
244
245
246
247
248
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
            using type = typename decltype(output)::value_type;
            auto in_h  = input.get_shape().lens()[2];
            auto in_w  = input.get_shape().lens()[3];

Paul's avatar
Paul committed
249
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
250
251
252
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
                    const int start_x0 = i * op.stride[0] - op.padding[0];
                    const int start_y0 = j * op.stride[1] - op.padding[1];

                    const int hend = std::min(start_x0 + op.lengths[0], in_h);
                    const int wend = std::min(start_y0 + op.lengths[1], in_w);

                    const int start_x = std::max(start_x0, 0);
                    const int start_y = std::max(start_y0, 0);

                    const int w_h       = (hend - start_x);
                    const int w_w       = (wend - start_y);
                    const int pool_size = std::max(w_h * w_w, 1);

                    double acc = Op::start();
                    dfor(w_h, w_w)([&](int x, int y) {
                        const int in_x = start_x + x;
                        const int in_y = start_y + y;
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
                            acc = Op::apply(acc, input(o, w, in_x, in_y));
                        }
                    });
                    output(o, w, i, j) = type(Op::final(acc, pool_size));
                });
        });
        return result;
    }
};

struct cpu_contiguous
{
285
    op::contiguous op;
Paul's avatar
Paul committed
286
    std::string name() const { return "cpu::contiguous"; }
Paul's avatar
Paul committed
287
288
    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
Paul's avatar
Paul committed
289
    {
Paul's avatar
Paul committed
290
        assert(output_shape.standard());
Paul's avatar
Paul committed
291
292
293
294
295
296
297
298
299
300
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
            shape_for_each(output.get_shape(), [&](const auto& idx) {
                output(idx.begin(), idx.end()) = input(idx.begin(), idx.end());
            });
        });
        return result;
    }
};

Khalique's avatar
Khalique committed
301
302
303
304
305
306
307
308
309
struct cpu_pad
{
    op::pad op;
    std::string name() const { return "cpu::contiguous"; }
    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
    {
        assert(output_shape.standard());
        argument result{output_shape};
Khalique's avatar
Khalique committed
310
        result.visit([&](auto output) { std::fill(output.begin(), output.end(), op.value); });
Khalique's avatar
Khalique committed
311
312

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

        return result;
    }
};

327
328
329
330
331
332
333
struct cpu_concat
{
    op::concat op;
    std::string name() const { return "cpu::concat"; }
    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
    {
Paul's avatar
Paul committed
334
        return op.compute(output_shape, std::move(args));
335
336
337
    }
};

Paul's avatar
Paul committed
338
339
struct cpu_gemm
{
Shucai Xiao's avatar
Shucai Xiao committed
340
341
    op::dot op;
    std::string name() const { return "cpu::dot"; }
Paul's avatar
Paul committed
342
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
343

Paul's avatar
Paul committed
344
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
345
346
    {
        argument result{output_shape};
Paul's avatar
Paul committed
347
        migemm(result, args[0], args[1], op.alpha, op.beta);
Paul's avatar
Paul committed
348
349
350
351
        return result;
    }
};

352
353
354
355
356
357
358
359
struct cpu_gather
{
    op::gather op;
    std::string name() const { return "cpu::gather"; }
    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
    {
Shucai Xiao's avatar
Shucai Xiao committed
360
        return op.compute(output_shape, std::move(args));
361
362
363
    }
};

Paul's avatar
Paul committed
364
365
366
367
368
369
370
371
372
373
374
375
376
377
struct identity_op
{
    std::string name() const { return "cpu::identity"; }
    auto fcn() const
    {
        return [](auto x) { return x; };
    }
};

struct abs_op
{
    std::string name() const { return "cpu::abs"; }
    auto fcn() const
    {
Khalique's avatar
Khalique committed
378
        return [](auto x) { return std::abs(make_signed(x)); };
Paul's avatar
Paul committed
379
380
381
382
383
384
385
386
387
388
389
390
    }
};

struct exp_op
{
    std::string name() const { return "cpu::exp"; }
    auto fcn() const
    {
        return [](auto x) { return std::exp(x); };
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
391
392
393
394
395
396
397
398
399
struct log_op
{
    std::string name() const { return "cpu::log"; }
    auto fcn() const
    {
        return [](auto x) { return std::log(x); };
    }
};

Paul's avatar
Paul committed
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
struct sin_op
{
    std::string name() const { return "cpu::sin"; }
    auto fcn() const
    {
        return [](auto x) { return std::sin(x); };
    }
};

struct cos_op
{
    std::string name() const { return "cpu::cos"; }
    auto fcn() const
    {
        return [](auto x) { return std::cos(x); };
    }
};

struct tan_op
{
    std::string name() const { return "cpu::tan"; }
    auto fcn() const
    {
        return [](auto x) { return std::tan(x); };
    }
};

struct asin_op
{
    std::string name() const { return "cpu::asin"; }
    auto fcn() const
    {
        return [](auto x) { return std::asin(x); };
    }
};

struct acos_op
{
    std::string name() const { return "cpu::acos"; }
    auto fcn() const
    {
        return [](auto x) { return std::acos(x); };
    }
};

struct atan_op
{
    std::string name() const { return "cpu::atan"; }
    auto fcn() const
    {
        return [](auto x) { return std::atan(x); };
    }
};

454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
struct sinh_op
{
    std::string name() const { return "cpu::sinh"; }
    auto fcn() const
    {
        return [](auto x) { return std::sinh(x); };
    }
};

struct cosh_op
{
    std::string name() const { return "cpu::cosh"; }
    auto fcn() const
    {
        return [](auto x) { return std::cosh(x); };
    }
};

Paul's avatar
Paul committed
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
503
struct tanh_op
{
    std::string name() const { return "cpu::tanh"; }
    auto fcn() const
    {
        return [](auto x) { return std::tanh(x); };
    }
};

struct sigmoid_op
{
    std::string name() const { return "cpu::sigmoid"; }
    auto fcn() const
    {
        return [](auto x) { return 1.f / (1.f + std::exp(-x)); };
    }
};

struct neg_op
{
    std::string name() const { return "cpu::neg"; }
    auto fcn() const
    {
        return [](auto x) { return -x; };
    }
};

struct relu_op
{
    std::string name() const { return "cpu::relu"; }
    auto fcn() const
    {
Paul's avatar
Paul committed
504
        return [](auto x) { return std::max(decltype(x){0}, x); };
Paul's avatar
Paul committed
505
506
507
    }
};

Khalique's avatar
Khalique committed
508
509
510
511
512
513
514
515
516
517
518
struct leaky_relu_op
{
    op::leaky_relu op;
    std::string name() const { return "cpu::leaky_relu"; }
    auto fcn() const
    {
        auto& a = op.alpha;
        return [a](auto x) { return x > 0 ? x : x * a; };
    }
};

Khalique's avatar
Khalique committed
519
520
521
522
523
524
525
526
527
528
529
struct elu_op
{
    op::elu op;
    std::string name() const { return "cpu::elu"; }
    auto fcn() const
    {
        auto& a = op.alpha;
        return [a](auto x) { return x > 0 ? x : a * std::expm1(x); };
    }
};

Paul's avatar
Paul committed
530
531
532
533
534
template <typename Op>
struct cpu_unary
{
    Op op;
    std::string name() const { return op.name(); }
Paul's avatar
Paul committed
535
536
    shape compute_shape(const std::vector<shape>& inputs) const { return inputs.front(); }
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
537
538
539
540
541
542
543
544
545
546
547
548
549
550
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
                std::transform(input.begin(), input.end(), output.begin(), op.fcn());
            });
        });
        return result;
    }
};

struct softmax2d
{
    std::string name() const { return "cpu::softmax2d"; }
Paul's avatar
Paul committed
551
552
    shape compute_shape(const std::vector<shape>& inputs) const { return inputs.front(); }
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
            auto nb          = input.get_shape().lens()[0];
            auto nc          = input.get_shape().lens()[1];
            auto nh          = input.get_shape().lens()[2];
            auto nw          = input.get_shape().lens()[3];
            dfor(nb, nh, nw)([&](std::size_t b, std::size_t i, std::size_t j) {
                value_type cmax = std::numeric_limits<value_type>::lowest();
                for(int c = 0; c < nc; c++)
                {
                    cmax = std::max(cmax, input(b, c, i, j));
                }
                for(int c = 0; c < nc; c++)
                {
                    output(b, c, i, j) = std::exp(input(b, c, i, j) - cmax);
                }
                value_type sum = value_type(0);
                for(int c = 0; c < nc; c++)
                {
                    sum += output(b, c, i, j);
                }
                for(int c = 0; c < nc; c++)
                {
                    output(b, c, i, j) = output(b, c, i, j) / sum;
                }
            });
        });
        return result;
    }
};

struct add_op
{
    std::string name() const { return "add"; }
    auto fcn() const
    {
        return [](auto x, auto y) { return x + y; };
    }
};

struct sub_op
{
    std::string name() const { return "sub"; }
    auto fcn() const
    {
        return [](auto x, auto y) { return x - y; };
    }
};

struct mul_op
{
    std::string name() const { return "mul"; }
    auto fcn() const
    {
        return [](auto x, auto y) { return x * y; };
    }
};

struct div_op
{
    std::string name() const { return "div"; }
    auto fcn() const
    {
        return [](auto x, auto y) { return x / y; };
    }
};

Khalique's avatar
Khalique committed
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
struct max_op
{
    std::string name() const { return "max"; }
    auto fcn() const
    {
        return [](auto x, auto y) { return std::max(x, y); };
    }
};

struct min_op
{
    std::string name() const { return "min"; }
    auto fcn() const
    {
        return [](auto x, auto y) { return std::min(x, y); };
    }
};

Paul's avatar
Paul committed
640
641
642
643
644
template <typename Op>
struct cpu_binary
{
    Op op;
    std::string name() const { return op.name(); }
Paul's avatar
Paul committed
645
646
    shape compute_shape(const std::vector<shape>& inputs) const { return inputs.front(); }
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
    {
        argument result{output_shape};
        visit_all(result, args[0], args[1])([&](auto output, auto input1, auto input2) {
            if(input1.get_shape().packed() and input2.get_shape().packed())
            {
                std::transform(
                    input1.begin(), input1.end(), input2.begin(), output.begin(), op.fcn());
            }
            else
            {
                shape_for_each(output.get_shape(), [&](const auto& idx) {
                    output(idx.begin(), idx.end()) =
                        op.fcn()(input1(idx.begin(), idx.end()), input2(idx.begin(), idx.end()));
                });
            }
        });
        return result;
    }
};

struct cpu_apply
{
    program* prog;
    std::unordered_map<std::string, std::function<void(instruction_ref)>> apply_map{};

    template <class T>
    auto simple_op()
    {
        return [this](instruction_ref ins) { apply_simple_op<T>(ins); };
    }

    template <class T, class Op>
    auto extend_op()
    {
        return [this](instruction_ref ins) { apply_extend_op<T, Op>(ins); };
    }

    void init()
    {
686
687
        apply_map["im2col"]      = extend_op<cpu_im2col, op::im2col>();
        apply_map["convolution"] = extend_op<cpu_convolution, op::convolution>();
688
        apply_map["dot"]         = extend_op<cpu_gemm, op::dot>();
Aditya Atluri's avatar
Aditya Atluri committed
689
        apply_map["batch_norm_inference"] =
690
691
            extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
        apply_map["contiguous"] = extend_op<cpu_contiguous, op::contiguous>();
Khalique's avatar
Khalique committed
692
        apply_map["pad"]        = extend_op<cpu_pad, op::pad>();
Scott Thornton's avatar
Scott Thornton committed
693
        apply_map["concat"]     = extend_op<cpu_concat, op::concat>();
694
        apply_map["gather"]     = extend_op<cpu_gather, op::gather>();
Khalique's avatar
Khalique committed
695
        apply_map["leaky_relu"] = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();
Khalique's avatar
Khalique committed
696
        apply_map["elu"]        = extend_op<cpu_unary<elu_op>, op::elu>();
wsttiger's avatar
wsttiger committed
697
        apply_map["identity"]   = simple_op<cpu_unary<identity_op>>();
Khalique's avatar
Khalique committed
698
        apply_map["abs"]        = simple_op<cpu_unary<abs_op>>();
699
700
        apply_map["sinh"]       = simple_op<cpu_unary<sinh_op>>();
        apply_map["cosh"]       = simple_op<cpu_unary<cosh_op>>();
wsttiger's avatar
wsttiger committed
701
702
703
        apply_map["tanh"]       = simple_op<cpu_unary<tanh_op>>();
        apply_map["sigmoid"]    = simple_op<cpu_unary<sigmoid_op>>();
        apply_map["exp"]        = simple_op<cpu_unary<exp_op>>();
Shucai Xiao's avatar
Shucai Xiao committed
704
        apply_map["log"]        = simple_op<cpu_unary<log_op>>();
wsttiger's avatar
wsttiger committed
705
706
707
708
        apply_map["neg"]        = simple_op<cpu_unary<neg_op>>();
        apply_map["sin"]        = simple_op<cpu_unary<sin_op>>();
        apply_map["cos"]        = simple_op<cpu_unary<cos_op>>();
        apply_map["tan"]        = simple_op<cpu_unary<tan_op>>();
709
710
711
        apply_map["asin"]       = simple_op<cpu_unary<asin_op>>();
        apply_map["acos"]       = simple_op<cpu_unary<acos_op>>();
        apply_map["atan"]       = simple_op<cpu_unary<atan_op>>();
Khalique's avatar
Khalique committed
712
        apply_map["relu"]       = simple_op<cpu_unary<relu_op>>();
wsttiger's avatar
wsttiger committed
713
714
715
716
        apply_map["add"]        = simple_op<cpu_binary<add_op>>();
        apply_map["sub"]        = simple_op<cpu_binary<sub_op>>();
        apply_map["mul"]        = simple_op<cpu_binary<mul_op>>();
        apply_map["div"]        = simple_op<cpu_binary<div_op>>();
Khalique's avatar
Khalique committed
717
718
        apply_map["max"]        = simple_op<cpu_binary<max_op>>();
        apply_map["min"]        = simple_op<cpu_binary<min_op>>();
Paul's avatar
Paul committed
719
720
721
722
723
724
725
726
727

        apply_map["softmax"] = simple_op<softmax2d>();
    }

    void apply()
    {
        init();
        for(auto it : iterator_for(*prog))
        {
Khalique's avatar
Khalique committed
728
            if(it->name() == "pooling")
Paul's avatar
Paul committed
729
730
731
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
732
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
733
            {
Paul's avatar
Paul committed
734
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
735
736
737
738
739
740
741
            }
        }
    }

    template <class T>
    void apply_simple_op(instruction_ref ins)
    {
Paul's avatar
Paul committed
742
        prog->replace_instruction(ins, T{}, ins->inputs());
Paul's avatar
Paul committed
743
744
745
746
747
    }

    template <class T, class Op>
    void apply_extend_op(instruction_ref ins)
    {
748
        auto&& op = any_cast<Op>(ins->get_operator());
Paul's avatar
Paul committed
749
        prog->replace_instruction(ins, T{op}, ins->inputs());
Paul's avatar
Paul committed
750
751
752
753
    }

    void apply_pooling(instruction_ref ins)
    {
754
        auto&& op = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
755
        if(op.mode == "max")
Paul's avatar
Paul committed
756
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
757
        else if(op.mode == "average")
Paul's avatar
Paul committed
758
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
759
760
761
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
762
void lowering::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
763
764

} // namespace cpu
Paul's avatar
Paul committed
765
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
766
} // namespace migraphx