lowering.cpp 30.7 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>
10
#include <migraphx/cpu/quant_gemm.hpp>
Paul's avatar
Paul committed
11
#include <unordered_map>
Paul's avatar
Paul committed
12
#include <utility>
Paul's avatar
Paul committed
13

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

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

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

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

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

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

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

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

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

72
        if(op.bn_mode == op::batch_norm_inference::spatial)
Scott Thornton's avatar
Scott Thornton committed
73
74
75
76
        {
            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
77
                    par_dfor(num_batch, num_channels, image_height, image_width)(
Scott Thornton's avatar
Scott Thornton committed
78
                        [&](std::size_t n, std::size_t c, std::size_t h, std::size_t w) {
79
80
81
82
                            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
83
84
                        });
                });
85
86
        }

87
        if(op.bn_mode == op::batch_norm_inference::per_activation)
Scott Thornton's avatar
Scott Thornton committed
88
        {
89
90
91
            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
92
                    par_dfor(num_batch, num_channels, image_height, image_width)(
93
                        [&](std::size_t n, std::size_t c, std::size_t h, std::size_t w) {
Paul's avatar
Paul committed
94
                            assert((variance(c, h, w) + epsilon) > 0);
Scott Thornton's avatar
Scott Thornton committed
95
96
97
98
99
100
                            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);
                        });
                });
101
        }
102
103
104
105
106

        return output;
    }
};

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

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

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

Paul's avatar
Paul committed
144
145
struct cpu_convolution
{
146
    op::convolution op;
Paul's avatar
Paul committed
147
148

    std::string name() const { return "cpu::convolution"; }
Paul's avatar
Paul committed
149
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
150
151
152
153
    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
154
            auto in   = input.get_shape().lens();
Khalique's avatar
Khalique committed
155
156
            auto in_h = in[2];
            auto in_w = in[3];
Paul's avatar
Paul committed
157

Khalique's avatar
Khalique committed
158
            auto wei   = weights.get_shape().lens();
Khalique's avatar
Khalique committed
159
160
161
162
            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
163

Paul's avatar
Paul committed
164
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
165
166
167
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
168
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
Khalique's avatar
Khalique committed
169
170
                    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
171
                    const int group_id = w / (wei_n / op.group);
Paul's avatar
Paul committed
172
173
174

                    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
175
176
                        const int in_x  = start_x + x;
                        const int in_y  = start_y + y;
Khalique's avatar
Khalique committed
177
                        const int in_ch = group_id * wei_c + k;
Paul's avatar
Paul committed
178
179
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
Khalique's avatar
Khalique committed
180
                            acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
Paul's avatar
Paul committed
181
182
183
184
185
186
187
188
189
                        }
                    });
                    output(o, w, i, j) = acc;
                });
        });
        return result;
    }
};

Scott Thornton's avatar
Scott Thornton committed
190
191
struct cpu_im2col
{
192
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
193

Scott Thornton's avatar
Scott Thornton committed
194
195
    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
196

wsttiger's avatar
wsttiger committed
197
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
198
    {
Scott Thornton's avatar
Scott Thornton committed
199
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
200
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
201
202
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
203
204
            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
205
206
207
            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
208
209
            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
210
211
212
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

Paul's avatar
Paul committed
213
214
            int kdiv2_h;
            int kdiv2_w;
Scott Thornton's avatar
Scott Thornton committed
215
216
            kdiv2_h = kernel_h / 2;
            kdiv2_w = kernel_w / 2;
Scott Thornton's avatar
Scott Thornton committed
217
            // calculate output sizes
Scott Thornton's avatar
Scott Thornton committed
218
219
            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
220
            // account for padding for the starting position of the input pixels
Scott Thornton's avatar
Scott Thornton committed
221
            std::size_t iinput = kdiv2_h - pad_h;
wsttiger's avatar
wsttiger committed
222
            // loop over output pixels (ioutput, joutput)
Scott Thornton's avatar
Scott Thornton committed
223
224
225
226
227
228
229
230
            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
231
232
233
234
235
236
237
238
239
240
                    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
241
242
                }
            }
Scott Thornton's avatar
Scott Thornton committed
243
        });
Scott Thornton's avatar
Scott Thornton committed
244
245
246
247
        return result;
    }
};

Paul's avatar
Paul committed
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
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
{
275
    op::pooling op;
Paul's avatar
Paul committed
276
277

    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
278
279
    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
280
281
282
283
284
285
286
    {
        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
287
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
288
289
290
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
                [&](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
{
323
    op::contiguous op;
Paul's avatar
Paul committed
324
    std::string name() const { return "cpu::contiguous"; }
Paul's avatar
Paul committed
325
326
    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
327
    {
Paul's avatar
Paul committed
328
        return op.compute(output_shape, std::move(args));
Paul's avatar
Paul committed
329
330
331
    }
};

Khalique's avatar
Khalique committed
332
struct cpu_pad
333
{
Khalique's avatar
Khalique committed
334
335
    op::pad op;
    std::string name() const { return "cpu::contiguous"; }
336
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
    {
Khalique's avatar
Khalique committed
339
        assert(output_shape.standard());
340
        argument result{output_shape};
Khalique's avatar
Khalique committed
341
        result.visit([&](auto output) { std::fill(output.begin(), output.end(), op.value); });
Khalique's avatar
Khalique committed
342
343

        visit_all(result, args[0])([&](auto output, auto input) {
344
            shape_for_each(input.get_shape(), [&](const auto& idx) {
Khalique's avatar
Khalique committed
345
                std::vector<std::size_t> new_idx(idx.size());
Khalique's avatar
Khalique committed
346
347
348
349
                std::transform(
                    idx.begin(), idx.end(), op.pads.begin(), new_idx.begin(), [](auto i, auto j) {
                        return i + j;
                    });
Khalique's avatar
Khalique committed
350
                output(new_idx.begin(), new_idx.end()) = input(idx.begin(), idx.end());
351
            });
Khalique's avatar
Khalique committed
352
353
        });

354
355
356
357
358
359
360
361
362
363
364
        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
365
        return op.compute(output_shape, std::move(args));
366
367
368
    }
};

Paul's avatar
Paul committed
369
370
struct cpu_gemm
{
Shucai Xiao's avatar
Shucai Xiao committed
371
372
    op::dot op;
    std::string name() const { return "cpu::dot"; }
Shucai Xiao's avatar
Shucai Xiao committed
373
374
    shape compute_shape(const std::vector<shape>& inputs) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
375
376
377
378
379
        if(inputs.size() == 3)
        {
            auto c_shape = inputs.at(2);
            check_shapes{{c_shape}}.not_broadcasted();
        }
Shucai Xiao's avatar
Shucai Xiao committed
380
        return op.compute_shape(inputs);
Shucai Xiao's avatar
Shucai Xiao committed
381
    }
Paul's avatar
Paul committed
382

Paul's avatar
Paul committed
383
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
384
385
    {
        argument result{output_shape};
Shucai Xiao's avatar
Shucai Xiao committed
386
        // 3 inputs, it is alpha * A * B + beta * C, then
387
        // A and B are matrices, and C is of the same shape as A * B
Shucai Xiao's avatar
Shucai Xiao committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
        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
410
411
412
413
        return result;
    }
};

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
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
struct cpu_quant_gemm
{
    op::quant_dot op;
    std::string name() const { return "cpu::quant_dot"; }
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        if(inputs.size() == 3)
        {
            auto c_shape = inputs.at(2);
            check_shapes{{c_shape}}.not_broadcasted();
        }
        return op.compute_shape(inputs);
    }

    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        // 3 inputs, it is alpha * A * B + beta * C, then
        // A and B are matrices, and C is of the same shape to A * B

        // first, convert the args[0] and args[1] from int8_t to int32_t
        argument arg_0{{shape::int32_type, {args.at(0).get_shape().lens()}}};
        argument arg_1{{shape::int32_type, {args.at(1).get_shape().lens()}}};
        arg_0.visit([&](auto output) {
            args.at(0).visit([&](auto input) {
                std::copy(input.begin(), input.end(), output.begin());
            });
        });

        arg_1.visit([&](auto output) {
            args.at(1).visit([&](auto input) {
                std::copy(input.begin(), input.end(), output.begin());
            });
        });

        if(args.size() == 3)
        {
            // no need to consider the value of args[2]
            if(op.beta == 0)
            {
                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, arg_0, arg_1, op.alpha, op.beta);

            return result;
        }

        // 2 input arguments
        int8_t beta = 0;
        migemm(result, arg_0, arg_1, op.alpha, beta);

        return result;
    }
};

476
477
478
479
480
481
482
483
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
484
        return op.compute(output_shape, std::move(args));
485
486
487
    }
};

Paul's avatar
Paul committed
488
489
490
491
492
493
494
495
496
497
498
499
500
501
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
502
        return [](auto x) { return std::abs(make_signed(x)); };
Paul's avatar
Paul committed
503
504
505
506
507
508
509
510
511
512
513
514
    }
};

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
515
516
517
518
519
520
521
522
523
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
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
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 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); };
    }
};

578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
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
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
622
623
624
625
626
627
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
628
        return [](auto x) { return std::max(decltype(x){0}, x); };
Paul's avatar
Paul committed
629
630
631
    }
};

Khalique's avatar
Khalique committed
632
633
634
635
636
637
638
639
640
641
642
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
643
644
645
646
647
648
649
650
651
652
653
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
654
655
656
657
658
template <typename Op>
struct cpu_unary
{
    Op op;
    std::string name() const { return op.name(); }
Paul's avatar
Paul committed
659
660
    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
661
662
663
664
665
666
667
668
669
670
671
672
673
674
    {
        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
675
676
    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
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
    {
        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;
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
710
711
712
713
714
715
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
716
    template <typename T>
Shucai Xiao's avatar
Shucai Xiao committed
717
718
    std::size_t compute_batch_index(const T& idx, shape& batch_shape, int axis) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
719
        if(axis == 0)
720
721
722
723
724
725
        {
            return 0;
        }
        else
        {
            std::vector<std::size_t> batch_idx(idx.begin(), idx.begin() + axis);
Shucai Xiao's avatar
Shucai Xiao committed
726
            return batch_shape.index(batch_idx.begin(), batch_idx.end());
727
        }
Shucai Xiao's avatar
Shucai Xiao committed
728
729
730
731
732
733
    }

    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        auto lens = output_shape.lens();
734
        std::vector<std::size_t> batch_lens{};
Shucai Xiao's avatar
Shucai Xiao committed
735
        if(op.axis == 0)
736
737
738
        {
            batch_lens.push_back(1);
        }
Shucai Xiao's avatar
Shucai Xiao committed
739
        else
740
741
742
        {
            batch_lens.insert(batch_lens.begin(), lens.begin(), lens.begin() + op.axis);
        }
Shucai Xiao's avatar
Shucai Xiao committed
743
744
745
        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
746
747
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
Shucai Xiao's avatar
Shucai Xiao committed
748
            shape_for_each(output_shape, [&](auto idx) {
749
                auto index       = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
750
751
752
753
                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
754
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
755
756
757
758
759
                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) {
760
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
761
762
763
                batch_sum[index] += std::exp(output(idx.begin(), idx.end()));
            });

Shucai Xiao's avatar
Shucai Xiao committed
764
            for(std::size_t i = 0; i < batch_sum.size(); ++i)
Shucai Xiao's avatar
Shucai Xiao committed
765
766
767
768
769
            {
                batch_sum[i] = std::log(batch_sum[i]);
            }

            shape_for_each(output_shape, [&](auto idx) {
770
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
771
                output(idx.begin(), idx.end()) -= batch_sum[index];
Shucai Xiao's avatar
Shucai Xiao committed
772
773
774
775
776
777
778
            });
        });

        return result;
    }
};

Paul's avatar
Paul committed
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
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
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
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
833
834
835
836
837
template <typename Op>
struct cpu_binary
{
    Op op;
    std::string name() const { return op.name(); }
Paul's avatar
Paul committed
838
839
    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
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
    {
        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()
    {
879
880
        apply_map["im2col"]      = extend_op<cpu_im2col, op::im2col>();
        apply_map["convolution"] = extend_op<cpu_convolution, op::convolution>();
881
        apply_map["dot"]         = extend_op<cpu_gemm, op::dot>();
882
        apply_map["quant_dot"]         = extend_op<cpu_quant_gemm, op::quant_dot>();
Aditya Atluri's avatar
Aditya Atluri committed
883
        apply_map["batch_norm_inference"] =
884
            extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
Khalique's avatar
Khalique committed
885
        apply_map["lrn"]        = extend_op<cpu_lrn, op::lrn>();
886
        apply_map["contiguous"] = extend_op<cpu_contiguous, op::contiguous>();
Khalique's avatar
Khalique committed
887
        apply_map["pad"]        = extend_op<cpu_pad, op::pad>();
Scott Thornton's avatar
Scott Thornton committed
888
        apply_map["concat"]     = extend_op<cpu_concat, op::concat>();
889
        apply_map["gather"]     = extend_op<cpu_gather, op::gather>();
Shucai Xiao's avatar
Shucai Xiao committed
890
        apply_map["logsoftmax"] = extend_op<cpu_logsoftmax, op::logsoftmax>();
Khalique's avatar
Khalique committed
891
        apply_map["leaky_relu"] = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();
Khalique's avatar
Khalique committed
892
        apply_map["elu"]        = extend_op<cpu_unary<elu_op>, op::elu>();
wsttiger's avatar
wsttiger committed
893
        apply_map["identity"]   = simple_op<cpu_unary<identity_op>>();
Khalique's avatar
Khalique committed
894
        apply_map["abs"]        = simple_op<cpu_unary<abs_op>>();
895
896
        apply_map["sinh"]       = simple_op<cpu_unary<sinh_op>>();
        apply_map["cosh"]       = simple_op<cpu_unary<cosh_op>>();
wsttiger's avatar
wsttiger committed
897
898
899
        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
900
        apply_map["log"]        = simple_op<cpu_unary<log_op>>();
wsttiger's avatar
wsttiger committed
901
902
903
904
        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>>();
905
906
907
        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
908
        apply_map["relu"]       = simple_op<cpu_unary<relu_op>>();
wsttiger's avatar
wsttiger committed
909
910
911
912
        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
913
914
        apply_map["max"]        = simple_op<cpu_binary<max_op>>();
        apply_map["min"]        = simple_op<cpu_binary<min_op>>();
Paul's avatar
Paul committed
915
916
917
918
919
920
921
922
923

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

    void apply()
    {
        init();
        for(auto it : iterator_for(*prog))
        {
Khalique's avatar
Khalique committed
924
            if(it->name() == "pooling")
Paul's avatar
Paul committed
925
926
927
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
928
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
929
            {
Paul's avatar
Paul committed
930
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
931
932
933
934
935
936
937
            }
        }
    }

    template <class T>
    void apply_simple_op(instruction_ref ins)
    {
Paul's avatar
Paul committed
938
        prog->replace_instruction(ins, T{}, ins->inputs());
Paul's avatar
Paul committed
939
940
941
942
943
    }

    template <class T, class Op>
    void apply_extend_op(instruction_ref ins)
    {
944
        auto&& op = any_cast<Op>(ins->get_operator());
Paul's avatar
Paul committed
945
        prog->replace_instruction(ins, T{op}, ins->inputs());
Paul's avatar
Paul committed
946
947
948
949
    }

    void apply_pooling(instruction_ref ins)
    {
950
        auto&& op = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
951
        if(op.mode == "max")
Paul's avatar
Paul committed
952
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
953
        else if(op.mode == "average")
Paul's avatar
Paul committed
954
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
955
956
957
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
958
void lowering::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
959
960

} // namespace cpu
Paul's avatar
Paul committed
961
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
962
} // namespace migraphx