lowering.cpp 27.2 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;
    }
};

Scott Thornton's avatar
Scott Thornton committed
207
208
struct cpu_im2col
{
209
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
210

211
212
213
214
215
216
    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
217
218
    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
219

wsttiger's avatar
wsttiger committed
220
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
221
    {
Scott Thornton's avatar
Scott Thornton committed
222
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
223
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
224
225
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
226
227
            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
228
229
230
            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
231
232
            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
233
234
235
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

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

Paul's avatar
Paul committed
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
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
{
296
    op::pooling op;
Paul's avatar
Paul committed
297

298
299
300
301
302
303
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Paul's avatar
Paul committed
304
    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
305
306
    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
307
308
309
310
311
312
313
    {
        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
314
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
315
316
317
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
                [&](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;
    }
};

348
struct cpu_op
Paul's avatar
Paul committed
349
{
350
351
    operation op;
    std::string name() const { return "cpu::" + op.name(); }
Paul's avatar
Paul committed
352
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
353
    argument compute(context&, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
354
    {
Paul's avatar
Paul committed
355
        return op.compute(output_shape, args);
Paul's avatar
Paul committed
356
    }
Paul's avatar
Paul committed
357
    friend bool operator==(const cpu_op& x, const cpu_op& y) { return x.op == y.op; }
358
359
360
361
362
363
    friend bool operator==(const cpu_op& x, const operation& y)
    {
        if(x.name() != y.name())
            return false;
        return x == any_cast<cpu_op>(y);
    }
Paul's avatar
Paul committed
364
    friend bool operator==(const operation& x, const cpu_op& y) { return y == x; }
Paul's avatar
Paul committed
365
366
};

Khalique's avatar
Khalique committed
367
struct cpu_pad
368
{
Khalique's avatar
Khalique committed
369
    op::pad op;
370
371
372
373
374
375
376

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

Khalique's avatar
Khalique committed
377
    std::string name() const { return "cpu::contiguous"; }
378
379
380
    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
381
        assert(output_shape.standard());
382
        argument result{output_shape};
Khalique's avatar
Khalique committed
383
        result.visit([&](auto output) { std::fill(output.begin(), output.end(), op.value); });
Khalique's avatar
Khalique committed
384
385

        visit_all(result, args[0])([&](auto output, auto input) {
386
            shape_for_each(input.get_shape(), [&](const auto& idx) {
Khalique's avatar
Khalique committed
387
                std::vector<std::size_t> new_idx(idx.size());
Khalique's avatar
Khalique committed
388
389
390
391
                std::transform(
                    idx.begin(), idx.end(), op.pads.begin(), new_idx.begin(), [](auto i, auto j) {
                        return i + j;
                    });
Khalique's avatar
Khalique committed
392
                output(new_idx.begin(), new_idx.end()) = input(idx.begin(), idx.end());
393
            });
Khalique's avatar
Khalique committed
394
395
        });

396
397
398
399
        return result;
    }
};

Paul's avatar
Paul committed
400
401
struct cpu_gemm
{
Shucai Xiao's avatar
Shucai Xiao committed
402
    op::dot op;
403
404
405
406
407
408

    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
409
    std::string name() const { return "cpu::dot"; }
Shucai Xiao's avatar
Shucai Xiao committed
410
411
    shape compute_shape(const std::vector<shape>& inputs) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
412
413
414
415
416
        if(inputs.size() == 3)
        {
            auto c_shape = inputs.at(2);
            check_shapes{{c_shape}}.not_broadcasted();
        }
Shucai Xiao's avatar
Shucai Xiao committed
417
        return op.compute_shape(inputs);
Shucai Xiao's avatar
Shucai Xiao committed
418
    }
Paul's avatar
Paul committed
419

Paul's avatar
Paul committed
420
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
421
422
    {
        argument result{output_shape};
Shucai Xiao's avatar
Shucai Xiao committed
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
        // 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
447
448
449
450
        return result;
    }
};

Khalique's avatar
Khalique committed
451
452
453
454
455
456
457
458
459
460
461
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
462
463
464
465
466
467
468
469
470
471
472
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
473
474
475
476
template <typename Op>
struct cpu_unary
{
    Op op;
477
478
479
480
481
482

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op.op, f);
    }
Paul's avatar
Paul committed
483
    std::string name() const { return op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
484
    shape compute_shape(const std::vector<shape>& inputs) const
485
    {
Shucai Xiao's avatar
Shucai Xiao committed
486
487
488
        check_shapes{inputs}.has(1);
        auto s = inputs.at(0);
        if(s.packed())
489
        {
Shucai Xiao's avatar
Shucai Xiao committed
490
            return s;
491
492
493
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
494
            return {s.type(), s.lens()};
495
496
497
        }
    }

Paul's avatar
Paul committed
498
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
499
500
501
502
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
Shucai Xiao's avatar
Shucai Xiao committed
503
                if(input.get_shape().standard())
504
505
506
507
508
509
510
511
512
                {
                    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
513
514
            });
        });
515

Paul's avatar
Paul committed
516
517
518
519
        return result;
    }
};

Khalique's avatar
Khalique committed
520
struct cpu_softmax
Paul's avatar
Paul committed
521
{
Khalique's avatar
Khalique committed
522
523
524
525
526
527
528
529
530
531
532
533
    op::softmax op;

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

    std::string name() const { return "cpu::softmax"; }
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }

    template <typename T>
534
    std::size_t compute_batch_index(T idx, shape& batch_shape, int axis) const
Khalique's avatar
Khalique committed
535
    {
Khalique's avatar
Khalique committed
536
        idx[axis] = 0;
537
        return batch_shape.index(idx);
Khalique's avatar
Khalique committed
538
539
    }

Paul's avatar
Paul committed
540
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
541
542
    {
        argument result{output_shape};
Khalique's avatar
Khalique committed
543
        auto batch_lens     = output_shape.lens();
Khalique's avatar
Khalique committed
544
        batch_lens[op.axis] = 1;
545
546
        shape batch_shape{shape::int32_type, batch_lens};

Paul's avatar
Paul committed
547
548
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
Shucai Xiao's avatar
Shucai Xiao committed
549
550
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
Khalique's avatar
Khalique committed
551
552
553
554
555
556
557
            shape_for_each(output_shape, [&](auto idx) {
                auto index       = this->compute_batch_index(idx, batch_shape, op.axis);
                batch_max[index] = std::max(batch_max[index], input(idx.begin(), idx.end()));
            });

            shape_for_each(output_shape, [&](auto idx) {
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
558
559
                output(idx.begin(), idx.end()) =
                    std::exp(input(idx.begin(), idx.end()) - batch_max[index]);
Khalique's avatar
Khalique committed
560
561
562
563
            });

            std::vector<value_type> batch_sum(batch_shape.elements(), value_type(0));
            shape_for_each(output_shape, [&](auto idx) {
Shucai Xiao's avatar
Shucai Xiao committed
564
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Khalique's avatar
Khalique committed
565
566
567
568
569
570
                batch_sum[index] += output(idx.begin(), idx.end());
            });

            shape_for_each(output_shape, [&](auto idx) {
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
                output(idx.begin(), idx.end()) /= batch_sum[index];
Paul's avatar
Paul committed
571
572
            });
        });
Khalique's avatar
Khalique committed
573

Paul's avatar
Paul committed
574
575
576
577
        return result;
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
578
579
580
struct cpu_logsoftmax
{
    op::logsoftmax op;
581
582
583
584
585
586
587

    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
588
589
590
    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
591
    template <typename T>
592
    std::size_t compute_batch_index(T idx, const shape& batch_shape, int axis) const
Shucai Xiao's avatar
Shucai Xiao committed
593
    {
594
595
        idx[axis] = 0;
        return batch_shape.index(idx);
Shucai Xiao's avatar
Shucai Xiao committed
596
597
598
599
600
    }

    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
Shucai Xiao's avatar
Shucai Xiao committed
601
        auto batch_lens     = output_shape.lens();
602
603
604
        batch_lens[op.axis] = 1;
        shape batch_shape{shape::int32_type, batch_lens};

Shucai Xiao's avatar
Shucai Xiao committed
605
606
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
Shucai Xiao's avatar
Shucai Xiao committed
607
608
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
Shucai Xiao's avatar
Shucai Xiao committed
609
            shape_for_each(output_shape, [&](auto idx) {
610
                auto index       = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
611
612
613
614
                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
615
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
616
617
618
619
620
                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) {
621
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
Shucai Xiao's avatar
Shucai Xiao committed
622
623
624
                batch_sum[index] += std::exp(output(idx.begin(), idx.end()));
            });

Shucai Xiao's avatar
Shucai Xiao committed
625
            for(std::size_t i = 0; i < batch_sum.size(); ++i)
Shucai Xiao's avatar
Shucai Xiao committed
626
627
628
629
630
            {
                batch_sum[i] = std::log(batch_sum[i]);
            }

            shape_for_each(output_shape, [&](auto idx) {
631
                auto index = this->compute_batch_index(idx, batch_shape, op.axis);
632
                output(idx.begin(), idx.end()) -= batch_sum[index];
Shucai Xiao's avatar
Shucai Xiao committed
633
634
635
636
637
638
639
            });
        });

        return result;
    }
};

640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
struct cpu_argmax
{
    op::argmax op;

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

    std::string name() const { return "cpu::argmax"; }
    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
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
Shucai Xiao's avatar
Shucai Xiao committed
658
659
                using value_type =
                    batch_max(output_shape.elements(), std::numeric_limits<value_type>::lowest());
660
661
662
                auto data_shape = args[0].get_shape();
                shape_for_each(data_shape, [&](auto idx) {
                    auto data_index = data_shape.index(idx);
Shucai Xiao's avatar
Shucai Xiao committed
663
664
665
                    idx[axis]       = 0;
                    auto out_index  = data_shape.index(idx);
                    if(batch_max[index] < input[data_index])
666
667
                    {
                        batch_max[index] = input[data_index];
Shucai Xiao's avatar
Shucai Xiao committed
668
                        output[index]    = static_cast<int64_t>(data_index);
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
                    }
                });
            });
        });

        return result;
    }
};

struct cpu_argmin
{
    op::argmin op;

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

    std::string name() const { return "cpu::argmin"; }
    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
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
Shucai Xiao's avatar
Shucai Xiao committed
696
697
                using value_type =
                    batch_min(output_shape.elements(), std::numeric_limits<value_type>::max());
698
699
700
                auto data_shape = args[0].get_shape();
                shape_for_each(data_shape, [&](auto idx) {
                    auto data_index = data_shape.index(idx);
Shucai Xiao's avatar
Shucai Xiao committed
701
702
703
                    idx[axis]       = 0;
                    auto out_index  = data_shape.index(idx);
                    if(batch_min[index] > input[data_index])
704
705
                    {
                        batch_min[index] = input[data_index];
Shucai Xiao's avatar
Shucai Xiao committed
706
                        output[index]    = static_cast<int64_t>(data_index);
707
708
709
710
711
712
713
714
715
                    }
                });
            });
        });

        return result;
    }
};

Paul's avatar
Paul committed
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
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()
    {
Shucai Xiao's avatar
Shucai Xiao committed
735
736
        apply_map["argmax"] = extend_op<cpu_argmax, op::argmax>();
        apply_map["argmin"] = extend_op<cpu_argmin, op::argmin>();
Paul's avatar
Paul committed
737
738
        apply_map["batch_norm_inference"] =
            extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
739
        apply_map["convolution"] = extend_op<cpu_convolution, op::convolution>();
740
        apply_map["dot"]         = extend_op<cpu_gemm, op::dot>();
Paul's avatar
Paul committed
741
        apply_map["elu"]         = extend_op<cpu_unary<elu_op>, op::elu>();
Paul's avatar
Paul committed
742
        apply_map["im2col"]      = extend_op<cpu_im2col, op::im2col>();
Paul's avatar
Paul committed
743
744
745
746
        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>();
Khalique's avatar
Khalique committed
747
        apply_map["softmax"]     = extend_op<cpu_softmax, op::softmax>();
Paul's avatar
Paul committed
748
749
750
751
752
753
754
    }

    void apply()
    {
        init();
        for(auto it : iterator_for(*prog))
        {
Khalique's avatar
Khalique committed
755
            if(it->name() == "pooling")
Paul's avatar
Paul committed
756
757
758
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
759
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
760
            {
Paul's avatar
Paul committed
761
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
762
            }
Paul's avatar
Paul committed
763
            else if(is_context_free(it->get_operator()))
764
765
766
            {
                apply_cpu_op(it);
            }
Paul's avatar
Paul committed
767
768
769
        }
    }

770
771
772
773
774
    void apply_cpu_op(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_op{ins->get_operator()}, ins->inputs());
    }

Paul's avatar
Paul committed
775
776
777
    template <class T>
    void apply_simple_op(instruction_ref ins)
    {
Paul's avatar
Paul committed
778
        prog->replace_instruction(ins, T{}, ins->inputs());
Paul's avatar
Paul committed
779
780
781
782
783
    }

    template <class T, class Op>
    void apply_extend_op(instruction_ref ins)
    {
784
        auto&& op = any_cast<Op>(ins->get_operator());
Paul's avatar
Paul committed
785
        prog->replace_instruction(ins, T{op}, ins->inputs());
Paul's avatar
Paul committed
786
787
788
789
    }

    void apply_pooling(instruction_ref ins)
    {
790
        auto&& op = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
791
        if(op.mode == "max")
Paul's avatar
Paul committed
792
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
793
        else if(op.mode == "average")
Paul's avatar
Paul committed
794
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
795
796
797
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
798
void lowering::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
799
800

} // namespace cpu
Paul's avatar
Paul committed
801
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
802
} // namespace migraphx