lowering.cpp 30 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) {
78
79
80
81
                            assert((variance[c] + epsilon) > 0);
                            result(n, c, h, w) = gamma[c] * (buffer(n, c, h, w) - mean[c]) /
                                                     std::sqrt(variance[c] + epsilon) +
                                                 bias[c];
Scott Thornton's avatar
Scott Thornton committed
82
83
                        });
                });
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;
    }
};

Khalique's avatar
Khalique committed
106
struct cpu_lrn
Khalique's avatar
Khalique committed
107
{
Khalique's avatar
Khalique committed
108
    op::lrn op;
Khalique's avatar
Khalique committed
109

Khalique's avatar
Khalique committed
110
    std::string name() const { return "cpu::lrn"; }
Khalique's avatar
Khalique committed
111
112
113
114
115
    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
116
117
118
119
            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
120
            float alphaoverarea = op.alpha / float(op.size);
Khalique's avatar
Khalique committed
121
            int radius          = (op.size - 1) / 2;
Khalique's avatar
Khalique committed
122

123
            par_dfor(n_batch, height, width)([&](int b, int h, int w) {
Khalique's avatar
Khalique committed
124
                float scale = 0;
Khalique's avatar
Khalique committed
125
126
                dfor(channels)([&](int c) {
                    auto start = (c - radius) < 0 ? 0 : (c - radius);
Khalique's avatar
Khalique committed
127
                    auto end   = (c + radius) > channels ? channels : (c + radius);
Khalique's avatar
Khalique committed
128
129
                    for(auto k = start; k < end; ++k)
                    {
Khalique's avatar
Khalique committed
130
                        scale += std::pow(input(b, k, h, w), 2);
Khalique's avatar
Khalique committed
131
132
133
                    }
                    scale *= alphaoverarea;
                    scale += op.bias;
Khalique's avatar
Khalique committed
134
                    scale              = std::pow(scale, -op.beta);
Khalique's avatar
Khalique committed
135
136
137
138
139
140
141
142
                    output(b, c, h, w) = input(b, c, h, w) * scale;
                });
            });
        });
        return result;
    }
};

Khalique's avatar
Khalique committed
143
144
145
146
147
148
struct clip_op
{
    op::clip op;
    std::string name() const { return "cpu::clip"; }
    auto fcn() const
    {
Khalique's avatar
Khalique committed
149
150
        auto max = op.max_val;
        auto min = op.min_val;
Khalique's avatar
Khalique committed
151
        return [max, min](auto x) {
Khalique's avatar
Khalique committed
152
            using type = decltype(x);
Khalique's avatar
Khalique committed
153
154
            return std::min(std::max(type(min), x), type(max));
        };
Khalique's avatar
Khalique committed
155
156
157
    }
};

Paul's avatar
Paul committed
158
159
struct cpu_convolution
{
160
    op::convolution op;
Paul's avatar
Paul committed
161
162

    std::string name() const { return "cpu::convolution"; }
Paul's avatar
Paul committed
163
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
164
165
166
167
    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
168
            auto in   = input.get_shape().lens();
Khalique's avatar
Khalique committed
169
170
            auto in_h = in[2];
            auto in_w = in[3];
Paul's avatar
Paul committed
171

Khalique's avatar
Khalique committed
172
            auto wei   = weights.get_shape().lens();
Khalique's avatar
Khalique committed
173
174
175
176
            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
177

Paul's avatar
Paul committed
178
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
179
180
181
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
182
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
Paul's avatar
Paul committed
183
184
185
                    const auto start_x  = i * op.stride[0] - op.padding[0];
                    const auto start_y  = j * op.stride[1] - op.padding[1];
                    const auto group_id = w / (wei_n / op.group);
Paul's avatar
Paul committed
186
187
188

                    double acc = 0;
                    dfor(wei_c, wei_h, wei_w)([&](std::size_t k, std::size_t x, std::size_t y) {
Paul's avatar
Paul committed
189
190
191
                        const auto in_x  = start_x + x;
                        const auto in_y  = start_y + y;
                        const auto in_ch = group_id * wei_c + k;
Paul's avatar
Paul committed
192
193
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
Khalique's avatar
Khalique committed
194
                            acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
Paul's avatar
Paul committed
195
196
197
198
199
200
201
202
203
                        }
                    });
                    output(o, w, i, j) = acc;
                });
        });
        return result;
    }
};

Scott Thornton's avatar
Scott Thornton committed
204
205
struct cpu_im2col
{
206
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
207

Scott Thornton's avatar
Scott Thornton committed
208
209
    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
210

wsttiger's avatar
wsttiger committed
211
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
212
    {
Scott Thornton's avatar
Scott Thornton committed
213
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
214
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
215
216
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
217
218
            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
219
220
221
            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
222
223
            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
224
225
226
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

Paul's avatar
Paul committed
227
228
            auto kdiv2_h = kernel_h / 2;
            auto kdiv2_w = kernel_w / 2;
Scott Thornton's avatar
Scott Thornton committed
229
            // calculate output sizes
Scott Thornton's avatar
Scott Thornton committed
230
231
            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
232
            // account for padding for the starting position of the input pixels
Scott Thornton's avatar
Scott Thornton committed
233
            std::size_t iinput = kdiv2_h - pad_h;
wsttiger's avatar
wsttiger committed
234
            // loop over output pixels (ioutput, joutput)
Scott Thornton's avatar
Scott Thornton committed
235
236
237
238
239
240
241
242
            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
243
244
245
                    dfor(channels,
                         kernel_h,
                         kernel_w)([&](std::size_t c, std::size_t koffset, std::size_t loffset) {
Paul's avatar
Paul committed
246
247
                        auto idx    = iinput + koffset - kdiv2_h;
                        auto jdx    = jinput + loffset - kdiv2_w;
wsttiger's avatar
wsttiger committed
248
249
250
251
252
                        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
253
254
                }
            }
Scott Thornton's avatar
Scott Thornton committed
255
        });
Scott Thornton's avatar
Scott Thornton committed
256
257
258
259
        return result;
    }
};

Paul's avatar
Paul committed
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
285
286
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
{
287
    op::pooling op;
Paul's avatar
Paul committed
288
289

    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
290
291
    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
292
293
294
295
296
297
298
    {
        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
299
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
300
301
302
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
                [&](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
{
335
    op::contiguous op;
Paul's avatar
Paul committed
336
    std::string name() const { return "cpu::contiguous"; }
Paul's avatar
Paul committed
337
338
    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
339
    {
Paul's avatar
Paul committed
340
        return op.compute(output_shape, std::move(args));
Paul's avatar
Paul committed
341
342
343
    }
};

Khalique's avatar
Khalique committed
344
struct cpu_pad
345
{
Khalique's avatar
Khalique committed
346
347
    op::pad op;
    std::string name() const { return "cpu::contiguous"; }
348
349
350
    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
351
        assert(output_shape.standard());
352
        argument result{output_shape};
Khalique's avatar
Khalique committed
353
        result.visit([&](auto output) { std::fill(output.begin(), output.end(), op.value); });
Khalique's avatar
Khalique committed
354
355

        visit_all(result, args[0])([&](auto output, auto input) {
356
            shape_for_each(input.get_shape(), [&](const auto& idx) {
Khalique's avatar
Khalique committed
357
                std::vector<std::size_t> new_idx(idx.size());
Khalique's avatar
Khalique committed
358
359
360
361
                std::transform(
                    idx.begin(), idx.end(), op.pads.begin(), new_idx.begin(), [](auto i, auto j) {
                        return i + j;
                    });
Khalique's avatar
Khalique committed
362
                output(new_idx.begin(), new_idx.end()) = input(idx.begin(), idx.end());
363
            });
Khalique's avatar
Khalique committed
364
365
        });

366
367
368
369
370
371
372
373
374
375
376
        return result;
    }
};

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
377
        return op.compute(output_shape, std::move(args));
378
379
380
    }
};

Paul's avatar
Paul committed
381
382
struct cpu_gemm
{
Shucai Xiao's avatar
Shucai Xiao committed
383
384
    op::dot op;
    std::string name() const { return "cpu::dot"; }
Shucai Xiao's avatar
Shucai Xiao committed
385
386
    shape compute_shape(const std::vector<shape>& inputs) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
387
388
389
390
391
        if(inputs.size() == 3)
        {
            auto c_shape = inputs.at(2);
            check_shapes{{c_shape}}.not_broadcasted();
        }
Shucai Xiao's avatar
Shucai Xiao committed
392
        return op.compute_shape(inputs);
Shucai Xiao's avatar
Shucai Xiao committed
393
    }
Paul's avatar
Paul committed
394

Paul's avatar
Paul committed
395
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
396
397
    {
        argument result{output_shape};
Shucai Xiao's avatar
Shucai Xiao committed
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
        // 3 inputs, it is alpha * A * B + beta * C, then
        // A and B are matrics, and C is broadcastable to A * B
        if(args.size() == 3)
        {
            // no need to consider the value of args[2]
            if(op.beta == 0.0f)
            {
                result.visit([&](auto output) { std::fill(output.begin(), output.end(), 0); });
            }
            else
            {
                visit_all(result, args[2])([&](auto output, auto input) {
                    std::copy(input.begin(), input.end(), output.begin());
                });
            }

            migemm(result, args[0], args[1], op.alpha, op.beta);

            return result;
        }

        // 2 input arguments
        migemm(result, args[0], args[1], op.alpha, 0.0f);

Paul's avatar
Paul committed
422
423
424
425
        return result;
    }
};

426
427
428
429
430
431
432
433
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
434
        return op.compute(output_shape, std::move(args));
435
436
437
    }
};

Paul's avatar
Paul committed
438
439
440
441
442
443
444
445
446
447
448
449
450
451
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
452
        return [](auto x) { return std::abs(make_signed(x)); };
Paul's avatar
Paul committed
453
454
455
456
457
458
459
460
461
462
463
464
    }
};

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
465
466
467
468
469
470
471
472
473
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
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
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 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); };
    }
};

528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
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
546
547
548
549
550
551
552
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
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
578
        return [](auto x) { return std::max(decltype(x){0}, x); };
Paul's avatar
Paul committed
579
580
581
    }
};

Khalique's avatar
Khalique committed
582
583
584
585
586
587
588
589
590
591
592
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
593
594
595
596
597
598
599
600
601
602
603
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
604
605
606
607
608
template <typename Op>
struct cpu_unary
{
    Op op;
    std::string name() const { return op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
609
    shape compute_shape(const std::vector<shape>& inputs) const
610
    {
Shucai Xiao's avatar
Shucai Xiao committed
611
612
613
        check_shapes{inputs}.has(1);
        auto s = inputs.at(0);
        if(s.packed())
614
        {
Shucai Xiao's avatar
Shucai Xiao committed
615
            return s;
616
617
618
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
619
            return {s.type(), s.lens()};
620
621
622
        }
    }

Paul's avatar
Paul committed
623
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
624
625
626
627
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
Shucai Xiao's avatar
Shucai Xiao committed
628
                if(input.get_shape().standard())
629
630
631
632
633
634
635
636
637
                {
                    std::transform(input.begin(), input.end(), output.begin(), op.fcn());
                }
                else
                {
                    shape_for_each(output.get_shape(), [&](const auto& idx) {
                        output(idx.begin(), idx.end()) = op.fcn()(input(idx.begin(), idx.end()));
                    });
                }
Paul's avatar
Paul committed
638
639
            });
        });
640

Paul's avatar
Paul committed
641
642
643
644
645
646
647
        return result;
    }
};

struct softmax2d
{
    std::string name() const { return "cpu::softmax2d"; }
Paul's avatar
Paul committed
648
649
    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
650
651
652
653
654
655
656
657
658
659
    {
        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();
Paul's avatar
Paul committed
660
                for(std::size_t c = 0; c < nc; c++)
Paul's avatar
Paul committed
661
662
663
                {
                    cmax = std::max(cmax, input(b, c, i, j));
                }
Paul's avatar
Paul committed
664
                for(std::size_t c = 0; c < nc; c++)
Paul's avatar
Paul committed
665
666
667
668
                {
                    output(b, c, i, j) = std::exp(input(b, c, i, j) - cmax);
                }
                value_type sum = value_type(0);
Paul's avatar
Paul committed
669
                for(std::size_t c = 0; c < nc; c++)
Paul's avatar
Paul committed
670
671
672
                {
                    sum += output(b, c, i, j);
                }
Paul's avatar
Paul committed
673
                for(std::size_t c = 0; c < nc; c++)
Paul's avatar
Paul committed
674
675
676
677
678
679
680
681
682
                {
                    output(b, c, i, j) = output(b, c, i, j) / sum;
                }
            });
        });
        return result;
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
683
684
685
686
687
688
struct cpu_logsoftmax
{
    op::logsoftmax op;
    std::string name() const { return "cpu::logsoftmax"; }
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }

Shucai Xiao's avatar
Shucai Xiao committed
689
    template <typename T>
Shucai Xiao's avatar
Shucai Xiao committed
690
691
    std::size_t compute_batch_index(const T& idx, shape& batch_shape, int axis) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
692
        if(axis == 0)
693
694
695
696
697
698
        {
            return 0;
        }
        else
        {
            std::vector<std::size_t> batch_idx(idx.begin(), idx.begin() + axis);
Shucai Xiao's avatar
Shucai Xiao committed
699
            return batch_shape.index(batch_idx.begin(), batch_idx.end());
700
        }
Shucai Xiao's avatar
Shucai Xiao committed
701
702
703
704
705
706
    }

    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        auto lens = output_shape.lens();
707
        std::vector<std::size_t> batch_lens{};
Shucai Xiao's avatar
Shucai Xiao committed
708
        if(op.axis == 0)
709
710
711
        {
            batch_lens.push_back(1);
        }
Shucai Xiao's avatar
Shucai Xiao committed
712
        else
713
714
715
        {
            batch_lens.insert(batch_lens.begin(), lens.begin(), lens.begin() + op.axis);
        }
Shucai Xiao's avatar
Shucai Xiao committed
716
717
718
        shape batch_shape{migraphx::shape::uint32_type, batch_lens};
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
Shucai Xiao's avatar
Shucai Xiao committed
719
720
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
Shucai Xiao's avatar
Shucai Xiao committed
721
            shape_for_each(output_shape, [&](auto idx) {
722
                auto index       = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
723
724
725
726
                batch_max[index] = std::max(batch_max[index], input(idx.begin(), idx.end()));
            });

            shape_for_each(output_shape, [&](auto idx) {
Shucai Xiao's avatar
Shucai Xiao committed
727
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
728
729
730
731
732
                output(idx.begin(), idx.end()) = input(idx.begin(), idx.end()) - batch_max[index];
            });

            std::vector<value_type> batch_sum(batch_shape.elements(), value_type(0));
            shape_for_each(output_shape, [&](auto idx) {
733
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
734
735
736
                batch_sum[index] += std::exp(output(idx.begin(), idx.end()));
            });

Shucai Xiao's avatar
Shucai Xiao committed
737
            for(std::size_t i = 0; i < batch_sum.size(); ++i)
Shucai Xiao's avatar
Shucai Xiao committed
738
739
740
741
742
            {
                batch_sum[i] = std::log(batch_sum[i]);
            }

            shape_for_each(output_shape, [&](auto idx) {
743
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
744
                output(idx.begin(), idx.end()) -= batch_sum[index];
Shucai Xiao's avatar
Shucai Xiao committed
745
746
747
748
749
750
751
            });
        });

        return result;
    }
};

Paul's avatar
Paul committed
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
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
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
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
806
807
808
809
template <typename Op>
struct cpu_binary
{
    Op op;
Shucai Xiao's avatar
Shucai Xiao committed
810
    std::string name() const { return "cpu::" + op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
811
    shape compute_shape(const std::vector<shape>& inputs) const
812
    {
Shucai Xiao's avatar
Shucai Xiao committed
813
814
815
816
        check_shapes{inputs}.has(2).same_type().same_dims();
        auto s0 = inputs.at(0);
        auto s1 = inputs.at(1);
        if(s0 == s1 and s0.packed())
817
        {
Shucai Xiao's avatar
Shucai Xiao committed
818
            return s0;
819
820
821
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
822
            return {s0.type(), s0.lens()};
823
824
        }
    }
Shucai Xiao's avatar
Shucai Xiao committed
825

Paul's avatar
Paul committed
826
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
827
828
829
    {
        argument result{output_shape};
        visit_all(result, args[0], args[1])([&](auto output, auto input1, auto input2) {
830
831
            auto s1 = input1.get_shape();
            auto s2 = input2.get_shape();
Shucai Xiao's avatar
Shucai Xiao committed
832
            if(s1 == s2 and s1.standard())
Paul's avatar
Paul committed
833
834
835
836
837
838
839
840
841
842
843
844
            {
                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()));
                });
            }
        });
845

Paul's avatar
Paul committed
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
        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()
    {
869
870
        apply_map["im2col"]      = extend_op<cpu_im2col, op::im2col>();
        apply_map["convolution"] = extend_op<cpu_convolution, op::convolution>();
871
        apply_map["dot"]         = extend_op<cpu_gemm, op::dot>();
Aditya Atluri's avatar
Aditya Atluri committed
872
        apply_map["batch_norm_inference"] =
873
            extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
Khalique's avatar
Khalique committed
874
        apply_map["lrn"]        = extend_op<cpu_lrn, op::lrn>();
Khalique's avatar
Khalique committed
875
        apply_map["clip"]       = extend_op<cpu_unary<clip_op>, op::clip>();
876
        apply_map["contiguous"] = extend_op<cpu_contiguous, op::contiguous>();
Khalique's avatar
Khalique committed
877
        apply_map["pad"]        = extend_op<cpu_pad, op::pad>();
Scott Thornton's avatar
Scott Thornton committed
878
        apply_map["concat"]     = extend_op<cpu_concat, op::concat>();
879
        apply_map["gather"]     = extend_op<cpu_gather, op::gather>();
Shucai Xiao's avatar
Shucai Xiao committed
880
        apply_map["logsoftmax"] = extend_op<cpu_logsoftmax, op::logsoftmax>();
Khalique's avatar
Khalique committed
881
        apply_map["leaky_relu"] = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();
Khalique's avatar
Khalique committed
882
        apply_map["elu"]        = extend_op<cpu_unary<elu_op>, op::elu>();
wsttiger's avatar
wsttiger committed
883
        apply_map["identity"]   = simple_op<cpu_unary<identity_op>>();
Khalique's avatar
Khalique committed
884
        apply_map["abs"]        = simple_op<cpu_unary<abs_op>>();
885
886
        apply_map["sinh"]       = simple_op<cpu_unary<sinh_op>>();
        apply_map["cosh"]       = simple_op<cpu_unary<cosh_op>>();
wsttiger's avatar
wsttiger committed
887
888
889
        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
890
        apply_map["log"]        = simple_op<cpu_unary<log_op>>();
wsttiger's avatar
wsttiger committed
891
892
893
894
        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>>();
895
896
897
        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
898
        apply_map["relu"]       = simple_op<cpu_unary<relu_op>>();
wsttiger's avatar
wsttiger committed
899
900
901
902
        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
903
904
        apply_map["max"]        = simple_op<cpu_binary<max_op>>();
        apply_map["min"]        = simple_op<cpu_binary<min_op>>();
Paul's avatar
Paul committed
905
906
907
908
909
910
911
912
913

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

    void apply()
    {
        init();
        for(auto it : iterator_for(*prog))
        {
Khalique's avatar
Khalique committed
914
            if(it->name() == "pooling")
Paul's avatar
Paul committed
915
916
917
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
918
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
919
            {
Paul's avatar
Paul committed
920
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
921
922
923
924
925
926
927
            }
        }
    }

    template <class T>
    void apply_simple_op(instruction_ref ins)
    {
Paul's avatar
Paul committed
928
        prog->replace_instruction(ins, T{}, ins->inputs());
Paul's avatar
Paul committed
929
930
931
932
933
    }

    template <class T, class Op>
    void apply_extend_op(instruction_ref ins)
    {
934
        auto&& op = any_cast<Op>(ins->get_operator());
Paul's avatar
Paul committed
935
        prog->replace_instruction(ins, T{op}, ins->inputs());
Paul's avatar
Paul committed
936
937
938
939
    }

    void apply_pooling(instruction_ref ins)
    {
940
        auto&& op = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
941
        if(op.mode == "max")
Paul's avatar
Paul committed
942
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
943
        else if(op.mode == "average")
Paul's avatar
Paul committed
944
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
945
946
947
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
948
void lowering::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
949
950

} // namespace cpu
Paul's avatar
Paul committed
951
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
952
} // namespace migraphx