lowering.cpp 28.6 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
53
54
55
56
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

57
58
    std::string name() const { return "cpu::batch_norm_inference"; }

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

Paul's avatar
Paul committed
61
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
62
    {
63
64
        argument output{output_shape};

Aditya Atluri's avatar
Aditya Atluri committed
65
66
        double epsilon           = op.epsilon;
        auto input               = args[0];
Paul's avatar
Paul committed
67
68
69
70
        auto arg_gamma           = args[1];
        auto arg_bias            = args[2];
        auto mini_batch_mean     = args[3];
        auto mini_batch_variance = args[4];
71

72
        auto num_batch    = output_shape.lens()[0];
Aditya Atluri's avatar
Aditya Atluri committed
73
74
        auto num_channels = output_shape.lens()[1];
        auto image_height = output_shape.lens()[2];
75
        auto image_width  = output_shape.lens()[3];
Aditya Atluri's avatar
Aditya Atluri committed
76

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

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

        return output;
    }
};

Khalique's avatar
Khalique committed
112
struct cpu_lrn
Khalique's avatar
Khalique committed
113
{
Khalique's avatar
Khalique committed
114
    op::lrn op;
Khalique's avatar
Khalique committed
115

116
117
118
119
120
121
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Khalique's avatar
Khalique committed
122
    std::string name() const { return "cpu::lrn"; }
Khalique's avatar
Khalique committed
123
124
125
126
127
    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
128
129
130
131
            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
132
            float alphaoverarea = op.alpha / float(op.size);
Khalique's avatar
Khalique committed
133
            int radius          = (op.size - 1) / 2;
Khalique's avatar
Khalique committed
134

135
            par_dfor(n_batch, height, width)([&](int b, int h, int w) {
Khalique's avatar
Khalique committed
136
                float scale = 0;
Khalique's avatar
Khalique committed
137
138
                dfor(channels)([&](int c) {
                    auto start = (c - radius) < 0 ? 0 : (c - radius);
Khalique's avatar
Khalique committed
139
                    auto end   = (c + radius) > channels ? channels : (c + radius);
Khalique's avatar
Khalique committed
140
141
                    for(auto k = start; k < end; ++k)
                    {
Khalique's avatar
Khalique committed
142
                        scale += std::pow(input(b, k, h, w), 2);
Khalique's avatar
Khalique committed
143
144
145
                    }
                    scale *= alphaoverarea;
                    scale += op.bias;
Khalique's avatar
Khalique committed
146
                    scale              = std::pow(scale, -op.beta);
Khalique's avatar
Khalique committed
147
148
149
150
151
152
153
154
                    output(b, c, h, w) = input(b, c, h, w) * scale;
                });
            });
        });
        return result;
    }
};

Paul's avatar
Paul committed
155
156
struct cpu_convolution
{
157
    op::convolution op;
Paul's avatar
Paul committed
158

159
160
161
162
163
164
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

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

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

Paul's avatar
Paul committed
181
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
182
183
184
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
185
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
Paul's avatar
Paul committed
186
187
188
                    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
189
190
191

                    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
192
193
194
                        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
195
196
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
Khalique's avatar
Khalique committed
197
                            acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
Paul's avatar
Paul committed
198
199
200
201
202
203
204
205
206
                        }
                    });
                    output(o, w, i, j) = acc;
                });
        });
        return result;
    }
};

207
208
209
210
struct cpu_quant_convolution
{
    op::quant_convolution op;

211
212
213
214
215
216
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

217
218
219
220
221
    std::string name() const { return "cpu::quant_convolution"; }
    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};
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
        result.visit([&](auto output) {
            visit_all(args[0], args[1])([&](auto input, auto weights) {
                auto in   = input.get_shape().lens();
                auto in_h = in[2];
                auto in_w = in[3];

                auto wei   = weights.get_shape().lens();
                auto wei_n = wei[0];
                auto wei_c = wei[1];
                auto wei_h = wei[2];
                auto wei_w = wei[3];

                par_dfor(output_shape.lens()[0],
                        output_shape.lens()[1],
                        output_shape.lens()[2],
                        output_shape.lens()[3])(
                    [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
                        const int start_x  = i * op.stride[0] - op.padding[0];
                        const int start_y  = j * op.stride[1] - op.padding[1];
                        const int group_id = w / (wei_n / op.group);

                        float acc = 0;
                        dfor(wei_c, wei_h, wei_w)([&](std::size_t k, std::size_t x, std::size_t y) {
                            const int in_x  = start_x + x;
                            const int in_y  = start_y + y;
                            const int in_ch = group_id * wei_c + k;
                            if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                            {
                                acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
                            }
                        });
                        output(o, w, i, j) = acc;
254
                    });
255
            });
256
        });
257
        
258
259
260
261
        return result;
    }
};

Scott Thornton's avatar
Scott Thornton committed
262
263
struct cpu_im2col
{
264
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
265

266
267
268
269
270
271
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Scott Thornton's avatar
Scott Thornton committed
272
273
    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
274

wsttiger's avatar
wsttiger committed
275
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
276
    {
Scott Thornton's avatar
Scott Thornton committed
277
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
278
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
279
280
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
281
282
            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
283
284
285
            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
286
287
            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
288
289
290
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

Paul's avatar
Paul committed
291
292
            auto kdiv2_h = kernel_h / 2;
            auto kdiv2_w = kernel_w / 2;
Scott Thornton's avatar
Scott Thornton committed
293
            // calculate output sizes
Scott Thornton's avatar
Scott Thornton committed
294
295
            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
296
            // account for padding for the starting position of the input pixels
Scott Thornton's avatar
Scott Thornton committed
297
            std::size_t iinput = kdiv2_h - pad_h;
wsttiger's avatar
wsttiger committed
298
            // loop over output pixels (ioutput, joutput)
Scott Thornton's avatar
Scott Thornton committed
299
300
301
302
303
304
305
306
            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
307
308
309
                    dfor(channels,
                         kernel_h,
                         kernel_w)([&](std::size_t c, std::size_t koffset, std::size_t loffset) {
Paul's avatar
Paul committed
310
311
                        auto idx    = iinput + koffset - kdiv2_h;
                        auto jdx    = jinput + loffset - kdiv2_w;
wsttiger's avatar
wsttiger committed
312
313
314
315
316
                        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
317
318
                }
            }
Scott Thornton's avatar
Scott Thornton committed
319
        });
Scott Thornton's avatar
Scott Thornton committed
320
321
322
323
        return result;
    }
};

Paul's avatar
Paul committed
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
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
{
351
    op::pooling op;
Paul's avatar
Paul committed
352

353
354
355
356
357
358
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Paul's avatar
Paul committed
359
    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
360
361
    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
362
363
364
365
366
367
368
    {
        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
369
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
370
371
372
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
                [&](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;
    }
};

403
struct cpu_op
Paul's avatar
Paul committed
404
{
405
406
    operation op;
    std::string name() const { return "cpu::" + op.name(); }
Paul's avatar
Paul committed
407
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
408
    argument compute(context&, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
409
    {
Paul's avatar
Paul committed
410
        return op.compute(output_shape, args);
Paul's avatar
Paul committed
411
    }
Paul's avatar
Paul committed
412
    friend bool operator==(const cpu_op& x, const cpu_op& y) { return x.op == y.op; }
413
    friend bool operator==(const cpu_op& x, const operation& y)
Paul's avatar
Paul committed
414
    {
415
416
417
        if(x.name() != y.name())
            return false;
        return x == any_cast<cpu_op>(y);
Paul's avatar
Paul committed
418
    }
Paul's avatar
Paul committed
419
    friend bool operator==(const operation& x, const cpu_op& y) { return y == x; }
Paul's avatar
Paul committed
420
421
};

Khalique's avatar
Khalique committed
422
struct cpu_pad
423
{
Khalique's avatar
Khalique committed
424
    op::pad op;
425
426
427
428
429
430
431

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

Khalique's avatar
Khalique committed
432
    std::string name() const { return "cpu::contiguous"; }
433
434
435
    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
436
        assert(output_shape.standard());
437
        argument result{output_shape};
Khalique's avatar
Khalique committed
438
        result.visit([&](auto output) { std::fill(output.begin(), output.end(), op.value); });
Khalique's avatar
Khalique committed
439
440

        visit_all(result, args[0])([&](auto output, auto input) {
441
            shape_for_each(input.get_shape(), [&](const auto& idx) {
Khalique's avatar
Khalique committed
442
                std::vector<std::size_t> new_idx(idx.size());
Khalique's avatar
Khalique committed
443
444
445
446
                std::transform(
                    idx.begin(), idx.end(), op.pads.begin(), new_idx.begin(), [](auto i, auto j) {
                        return i + j;
                    });
Khalique's avatar
Khalique committed
447
                output(new_idx.begin(), new_idx.end()) = input(idx.begin(), idx.end());
448
            });
Khalique's avatar
Khalique committed
449
450
        });

451
452
453
454
        return result;
    }
};

Paul's avatar
Paul committed
455
456
struct cpu_gemm
{
Shucai Xiao's avatar
Shucai Xiao committed
457
    op::dot op;
458
459
460
461
462
463

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }
Shucai Xiao's avatar
Shucai Xiao committed
464
    std::string name() const { return "cpu::dot"; }
Shucai Xiao's avatar
Shucai Xiao committed
465
466
    shape compute_shape(const std::vector<shape>& inputs) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
467
468
469
470
471
        if(inputs.size() == 3)
        {
            auto c_shape = inputs.at(2);
            check_shapes{{c_shape}}.not_broadcasted();
        }
Shucai Xiao's avatar
Shucai Xiao committed
472
        return op.compute_shape(inputs);
Shucai Xiao's avatar
Shucai Xiao committed
473
    }
Paul's avatar
Paul committed
474

Paul's avatar
Paul committed
475
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
476
477
    {
        argument result{output_shape};
Shucai Xiao's avatar
Shucai Xiao committed
478
        // 3 inputs, it is alpha * A * B + beta * C, then
479
        // A and B are matrices, and C is of the same shape as A * B
Shucai Xiao's avatar
Shucai Xiao committed
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
        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
502
503
504
505
        return result;
    }
};

506
507
508
struct cpu_quant_gemm
{
    op::quant_dot op;
509
510
511
512
513
514
515

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

516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
    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) {
Shucai Xiao's avatar
Shucai Xiao committed
537
538
            args.at(0).visit(
                [&](auto input) { std::copy(input.begin(), input.end(), output.begin()); });
539
540
541
        });

        arg_1.visit([&](auto output) {
Shucai Xiao's avatar
Shucai Xiao committed
542
543
            args.at(1).visit(
                [&](auto input) { std::copy(input.begin(), input.end(), output.begin()); });
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
        });

        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;
    }
};

Khalique's avatar
Khalique committed
573
574
575
576
577
578
579
580
581
582
583
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
584
585
586
587
588
589
590
591
592
593
594
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
595
596
597
598
template <typename Op>
struct cpu_unary
{
    Op op;
599
600
601
602
603
604

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op.op, f);
    }
Paul's avatar
Paul committed
605
    std::string name() const { return op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
606
    shape compute_shape(const std::vector<shape>& inputs) const
607
    {
Shucai Xiao's avatar
Shucai Xiao committed
608
609
610
        check_shapes{inputs}.has(1);
        auto s = inputs.at(0);
        if(s.packed())
611
        {
Shucai Xiao's avatar
Shucai Xiao committed
612
            return s;
613
614
615
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
616
            return {s.type(), s.lens()};
617
618
619
        }
    }

Paul's avatar
Paul committed
620
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
621
622
623
624
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
Shucai Xiao's avatar
Shucai Xiao committed
625
                if(input.get_shape().standard())
626
627
628
629
630
631
632
633
634
                {
                    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
635
636
            });
        });
637

Paul's avatar
Paul committed
638
639
640
641
642
643
644
        return result;
    }
};

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

Shucai Xiao's avatar
Shucai Xiao committed
680
681
682
struct cpu_logsoftmax
{
    op::logsoftmax op;
683
684
685
686
687
688
689

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

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

    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        auto lens = output_shape.lens();
711
        std::vector<std::size_t> batch_lens{};
Shucai Xiao's avatar
Shucai Xiao committed
712
        if(op.axis == 0)
713
714
715
        {
            batch_lens.push_back(1);
        }
Shucai Xiao's avatar
Shucai Xiao committed
716
        else
717
718
719
        {
            batch_lens.insert(batch_lens.begin(), lens.begin(), lens.begin() + op.axis);
        }
Shucai Xiao's avatar
Shucai Xiao committed
720
721
722
        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
723
724
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
Shucai Xiao's avatar
Shucai Xiao committed
725
            shape_for_each(output_shape, [&](auto idx) {
726
                auto index       = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
727
728
729
730
                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
731
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
732
733
734
735
736
                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) {
737
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
738
739
740
                batch_sum[index] += std::exp(output(idx.begin(), idx.end()));
            });

Shucai Xiao's avatar
Shucai Xiao committed
741
            for(std::size_t i = 0; i < batch_sum.size(); ++i)
Shucai Xiao's avatar
Shucai Xiao committed
742
743
744
745
746
            {
                batch_sum[i] = std::log(batch_sum[i]);
            }

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

        return result;
    }
};

Paul's avatar
Paul committed
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
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()
    {
Aditya Atluri's avatar
Aditya Atluri committed
775
        apply_map["batch_norm_inference"] =
776
            extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
Shucai Xiao's avatar
Shucai Xiao committed
777
778
        apply_map["convolution"]       = extend_op<cpu_convolution, op::convolution>();
        apply_map["dot"]               = extend_op<cpu_gemm, op::dot>();
779
780
        apply_map["quant_dot"]         = extend_op<cpu_quant_gemm, op::quant_dot>();
        apply_map["quant_convolution"] = extend_op<cpu_quant_convolution, op::quant_convolution>();
Shucai Xiao's avatar
Shucai Xiao committed
781
782
783
784
785
786
787
        apply_map["elu"]               = extend_op<cpu_unary<elu_op>, op::elu>();
        apply_map["im2col"]            = extend_op<cpu_im2col, op::im2col>();
        apply_map["leaky_relu"]        = extend_op<cpu_unary<leaky_relu_op>, op::leaky_relu>();
        apply_map["logsoftmax"]        = extend_op<cpu_logsoftmax, op::logsoftmax>();
        apply_map["lrn"]               = extend_op<cpu_lrn, op::lrn>();
        apply_map["pad"]               = extend_op<cpu_pad, op::pad>();
        apply_map["softmax"]           = simple_op<softmax2d>();
Paul's avatar
Paul committed
788
789
790
791
792
793
794
    }

    void apply()
    {
        init();
        for(auto it : iterator_for(*prog))
        {
Khalique's avatar
Khalique committed
795
            if(it->name() == "pooling")
Paul's avatar
Paul committed
796
797
798
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
799
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
800
            {
Paul's avatar
Paul committed
801
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
802
            }
Paul's avatar
Paul committed
803
            else if(is_context_free(it->get_operator()))
804
805
806
            {
                apply_cpu_op(it);
            }
Paul's avatar
Paul committed
807
808
809
        }
    }

810
811
812
813
814
    void apply_cpu_op(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_op{ins->get_operator()}, ins->inputs());
    }

Paul's avatar
Paul committed
815
816
817
    template <class T>
    void apply_simple_op(instruction_ref ins)
    {
Paul's avatar
Paul committed
818
        prog->replace_instruction(ins, T{}, ins->inputs());
Paul's avatar
Paul committed
819
820
821
822
823
    }

    template <class T, class Op>
    void apply_extend_op(instruction_ref ins)
    {
824
        auto&& op = any_cast<Op>(ins->get_operator());
Paul's avatar
Paul committed
825
        prog->replace_instruction(ins, T{op}, ins->inputs());
Paul's avatar
Paul committed
826
827
828
829
    }

    void apply_pooling(instruction_ref ins)
    {
830
        auto&& op = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
831
        if(op.mode == "max")
Paul's avatar
Paul committed
832
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
833
        else if(op.mode == "average")
Paul's avatar
Paul committed
834
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
835
836
837
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
838
void lowering::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
839
840

} // namespace cpu
Paul's avatar
Paul committed
841
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
842
} // namespace migraphx