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

Paul's avatar
Paul committed
2
3
4
#include <migraphx/cpu/lowering.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/dfor.hpp>
Paul's avatar
Paul committed
5
6
#include <migraphx/op/batch_norm.hpp>
#include <migraphx/op/convolution.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
7
#include <migraphx/op/quant_convolution.hpp>
Paul's avatar
Paul committed
8
#include <migraphx/op/dot.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
9
#include <migraphx/op/quant_dot.hpp>
Paul's avatar
Paul committed
10
11
12
13
14
15
16
17
#include <migraphx/op/elu.hpp>
#include <migraphx/op/im2col.hpp>
#include <migraphx/op/leaky_relu.hpp>
#include <migraphx/op/logsoftmax.hpp>
#include <migraphx/op/lrn.hpp>
#include <migraphx/op/pad.hpp>
#include <migraphx/op/pooling.hpp>
#include <migraphx/op/softmax.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
18
19
#include <migraphx/op/argmax.hpp>
#include <migraphx/op/argmin.hpp>
Paul's avatar
Paul committed
20
21
#include <migraphx/shape_for_each.hpp>
#include <migraphx/iterator_for.hpp>
Paul's avatar
Paul committed
22
#include <migraphx/par_dfor.hpp>
Paul's avatar
Paul committed
23
#include <migraphx/cpu/gemm.hpp>
Paul's avatar
Paul committed
24
#include <unordered_map>
Paul's avatar
Paul committed
25
#include <utility>
Paul's avatar
Paul committed
26

Paul's avatar
Paul committed
27
namespace migraphx {
Paul's avatar
Paul committed
28
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
29
30
31
32
33
34
35
36
namespace cpu {

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

Khalique's avatar
Khalique committed
37
38
39
40
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
41
42
43
44
{
    return x;
}

45
46
47
48
//
// cpu implemenataion of batch norm for inference
//
// inputs are:
49
50
51
52
// 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
53
// args[4] -> bias
54
55
56
//
// The equation to compute batch norm for inference is:
//
Aditya Atluri's avatar
Aditya Atluri committed
57
// output[i] = bias + gamma * (input[i] + mean) / sqrt(variance + epsilon)
58
59
60
61
62
//
// the input data format should be nchw
//
struct cpu_batch_norm_inference
{
63
    op::batch_norm_inference op;
64

65
66
67
68
69
70
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

71
72
    std::string name() const { return "cpu::batch_norm_inference"; }

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

Paul's avatar
Paul committed
75
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
76
    {
77
78
        argument output{output_shape};

Aditya Atluri's avatar
Aditya Atluri committed
79
80
        double epsilon           = op.epsilon;
        auto input               = args[0];
Paul's avatar
Paul committed
81
82
83
84
        auto arg_gamma           = args[1];
        auto arg_bias            = args[2];
        auto mini_batch_mean     = args[3];
        auto mini_batch_variance = args[4];
85

86
        auto num_batch    = output_shape.lens()[0];
Aditya Atluri's avatar
Aditya Atluri committed
87
88
        auto num_channels = output_shape.lens()[1];
        auto image_height = output_shape.lens()[2];
89
        auto image_width  = output_shape.lens()[3];
Aditya Atluri's avatar
Aditya Atluri committed
90

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

106
        if(op.bn_mode == op::batch_norm_inference::per_activation)
Scott Thornton's avatar
Scott Thornton committed
107
        {
108
109
110
            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
111
                    par_dfor(num_batch, num_channels, image_height, image_width)(
112
                        [&](std::size_t n, std::size_t c, std::size_t h, std::size_t w) {
Paul's avatar
Paul committed
113
                            assert((variance(c, h, w) + epsilon) > 0);
Scott Thornton's avatar
Scott Thornton committed
114
115
116
117
118
119
                            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);
                        });
                });
120
        }
121
122
123
124
125

        return output;
    }
};

Khalique's avatar
Khalique committed
126
struct cpu_lrn
Khalique's avatar
Khalique committed
127
{
Khalique's avatar
Khalique committed
128
    op::lrn op;
Khalique's avatar
Khalique committed
129

130
131
132
133
134
135
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Khalique's avatar
Khalique committed
136
    std::string name() const { return "cpu::lrn"; }
Khalique's avatar
Khalique committed
137
138
139
140
141
    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
142
143
144
145
            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
146
            float alphaoverarea = op.alpha / float(op.size);
Khalique's avatar
Khalique committed
147
            int radius          = (op.size - 1) / 2;
Khalique's avatar
Khalique committed
148

149
            par_dfor(n_batch, height, width)([&](int b, int h, int w) {
Khalique's avatar
Khalique committed
150
                float scale = 0;
Khalique's avatar
Khalique committed
151
152
                dfor(channels)([&](int c) {
                    auto start = (c - radius) < 0 ? 0 : (c - radius);
Khalique's avatar
Khalique committed
153
                    auto end   = (c + radius) > channels ? channels : (c + radius);
Khalique's avatar
Khalique committed
154
155
                    for(auto k = start; k < end; ++k)
                    {
Khalique's avatar
Khalique committed
156
                        scale += std::pow(input(b, k, h, w), 2);
Khalique's avatar
Khalique committed
157
158
159
                    }
                    scale *= alphaoverarea;
                    scale += op.bias;
Khalique's avatar
Khalique committed
160
                    scale              = std::pow(scale, -op.beta);
Khalique's avatar
Khalique committed
161
162
163
164
165
166
167
168
                    output(b, c, h, w) = input(b, c, h, w) * scale;
                });
            });
        });
        return result;
    }
};

Paul's avatar
Paul committed
169
170
struct cpu_convolution
{
171
    op::convolution op;
Paul's avatar
Paul committed
172

173
174
175
176
177
178
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Paul's avatar
Paul committed
179
    std::string name() const { return "cpu::convolution"; }
Paul's avatar
Paul committed
180
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
181
182
183
184
    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
185
            auto in   = input.get_shape().lens();
Khalique's avatar
Khalique committed
186
187
            auto in_h = in[2];
            auto in_w = in[3];
Paul's avatar
Paul committed
188

Khalique's avatar
Khalique committed
189
            auto wei   = weights.get_shape().lens();
Khalique's avatar
Khalique committed
190
191
192
193
            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
194

Paul's avatar
Paul committed
195
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
196
197
198
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
199
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
Paul's avatar
Paul committed
200
201
202
                    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
203
204
205

                    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
206
207
208
                        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
209
210
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
Khalique's avatar
Khalique committed
211
                            acc += input(o, in_ch, in_x, in_y) * weights(w, k, x, y);
Paul's avatar
Paul committed
212
213
214
215
216
217
218
219
220
                        }
                    });
                    output(o, w, i, j) = acc;
                });
        });
        return result;
    }
};

221
222
223
224
struct cpu_quant_convolution
{
    op::quant_convolution op;

225
226
227
228
229
230
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

231
232
233
234
235
    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};
236
        auto output = result.get<int32_t>();
Shucai Xiao's avatar
Shucai Xiao committed
237
238
239
240
241
242
243
244
245
246
247
248
        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],
Shucai Xiao's avatar
Shucai Xiao committed
249
250
251
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Shucai Xiao's avatar
Shucai Xiao committed
252
253
254
255
256
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
                    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);

257
                    int32_t acc = 0;
Shucai Xiao's avatar
Shucai Xiao committed
258
259
260
261
262
263
                    dfor(wei_c, wei_h, wei_w)([&](std::size_t k, std::size_t x, std::size_t y) {
                        const auto in_x  = start_x + x;
                        const auto in_y  = start_y + y;
                        const auto in_ch = group_id * wei_c + k;
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
Shucai Xiao's avatar
Shucai Xiao committed
264
265
                            acc += static_cast<int32_t>(input(o, in_ch, in_x, in_y)) *
                                   weights(w, k, x, y);
Shucai Xiao's avatar
Shucai Xiao committed
266
                        }
267
                    });
Shucai Xiao's avatar
Shucai Xiao committed
268
269
                    output(o, w, i, j) = acc;
                });
270
        });
Shucai Xiao's avatar
Shucai Xiao committed
271

272
273
274
275
        return result;
    }
};

Scott Thornton's avatar
Scott Thornton committed
276
277
struct cpu_im2col
{
278
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
279

280
281
282
283
284
285
    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
286
287
    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
288

wsttiger's avatar
wsttiger committed
289
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
290
    {
Scott Thornton's avatar
Scott Thornton committed
291
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
292
        auto input_shape   = args[0].get_shape();
Scott Thornton's avatar
Scott Thornton committed
293
294
        auto weights_shape = args[1].get_shape();
        visit_all(result, args[0])([&](auto col, auto input) {
Scott Thornton's avatar
Scott Thornton committed
295
296
            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
297
298
299
            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
300
301
            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
302
303
304
            const std::size_t& stride_h = op.stride[0];
            const std::size_t& stride_w = op.stride[1];

Paul's avatar
Paul committed
305
306
            long kdiv2_h = long(kernel_h) / 2;
            long kdiv2_w = long(kernel_w) / 2;
Scott Thornton's avatar
Scott Thornton committed
307
            // calculate output sizes
Scott Thornton's avatar
Scott Thornton committed
308
309
            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
310
            // account for padding for the starting position of the input pixels
Paul's avatar
Paul committed
311
            long iinput = kdiv2_h - long(pad_h);
wsttiger's avatar
wsttiger committed
312
            // loop over output pixels (ioutput, joutput)
Scott Thornton's avatar
Scott Thornton committed
313
314
            for(std::size_t ioutput = 0; ioutput < col_height; ioutput++, iinput += stride_h)
            {
Paul's avatar
Paul committed
315
                long jinput = kdiv2_w - long(pad_w);
Scott Thornton's avatar
Scott Thornton committed
316
317
318
319
320
                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
321
322
323
                    dfor(channels,
                         kernel_h,
                         kernel_w)([&](std::size_t c, std::size_t koffset, std::size_t loffset) {
Paul's avatar
Paul committed
324
325
                        auto idx    = iinput + long(koffset) - kdiv2_h;
                        auto jdx    = jinput + long(loffset) - kdiv2_w;
wsttiger's avatar
wsttiger committed
326
327
328
329
330
                        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
331
332
                }
            }
Scott Thornton's avatar
Scott Thornton committed
333
        });
Scott Thornton's avatar
Scott Thornton committed
334
335
336
337
        return result;
    }
};

Paul's avatar
Paul committed
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
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
{
365
    op::pooling op;
Paul's avatar
Paul committed
366

367
368
369
370
371
372
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op, f);
    }

Paul's avatar
Paul committed
373
    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
374
375
    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
376
377
378
379
380
381
382
    {
        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
383
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
384
385
386
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
                [&](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;
    }
};

417
struct cpu_op
Paul's avatar
Paul committed
418
{
419
420
    operation op;
    std::string name() const { return "cpu::" + op.name(); }
Paul's avatar
Paul committed
421
    shape compute_shape(const std::vector<shape>& inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
422
    argument compute(context&, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
423
    {
Paul's avatar
Paul committed
424
        return op.compute(output_shape, args);
Paul's avatar
Paul committed
425
    }
Paul's avatar
Paul committed
426
    friend bool operator==(const cpu_op& x, const cpu_op& y) { return x.op == y.op; }
427
    friend bool operator==(const cpu_op& x, const operation& y)
Paul's avatar
Paul committed
428
    {
429
430
431
        if(x.name() != y.name())
            return false;
        return x == any_cast<cpu_op>(y);
Paul's avatar
Paul committed
432
    }
Paul's avatar
Paul committed
433
    friend bool operator==(const operation& x, const cpu_op& y) { return y == x; }
Paul's avatar
Paul committed
434
435
};

Khalique's avatar
Khalique committed
436
struct cpu_pad
437
{
Khalique's avatar
Khalique committed
438
    op::pad op;
439
440
441
442
443
444
445

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

Khalique's avatar
Khalique committed
446
    std::string name() const { return "cpu::contiguous"; }
447
448
449
    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
450
        assert(output_shape.standard());
451
        argument result{output_shape};
Khalique's avatar
Khalique committed
452
        result.visit([&](auto output) { std::fill(output.begin(), output.end(), op.value); });
Khalique's avatar
Khalique committed
453
454

        visit_all(result, args[0])([&](auto output, auto input) {
455
            shape_for_each(input.get_shape(), [&](const auto& idx) {
Khalique's avatar
Khalique committed
456
                std::vector<std::size_t> new_idx(idx.size());
Khalique's avatar
Khalique committed
457
458
459
460
                std::transform(
                    idx.begin(), idx.end(), op.pads.begin(), new_idx.begin(), [](auto i, auto j) {
                        return i + j;
                    });
Khalique's avatar
Khalique committed
461
                output(new_idx.begin(), new_idx.end()) = input(idx.begin(), idx.end());
462
            });
Khalique's avatar
Khalique committed
463
464
        });

465
466
467
468
        return result;
    }
};

Paul's avatar
Paul committed
469
470
struct cpu_gemm
{
Shucai Xiao's avatar
Shucai Xiao committed
471
    op::dot op;
472
473
474
475
476
477

    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
478
    std::string name() const { return "cpu::dot"; }
Shucai Xiao's avatar
Shucai Xiao committed
479
480
    shape compute_shape(const std::vector<shape>& inputs) const
    {
Shucai Xiao's avatar
Shucai Xiao committed
481
482
483
484
485
        if(inputs.size() == 3)
        {
            auto c_shape = inputs.at(2);
            check_shapes{{c_shape}}.not_broadcasted();
        }
Shucai Xiao's avatar
Shucai Xiao committed
486
        return op.compute_shape(inputs);
Shucai Xiao's avatar
Shucai Xiao committed
487
    }
Paul's avatar
Paul committed
488

Paul's avatar
Paul committed
489
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
490
491
    {
        argument result{output_shape};
Shucai Xiao's avatar
Shucai Xiao committed
492
        // 3 inputs, it is alpha * A * B + beta * C, then
493
        // A and B are matrices, and C is of the same shape as A * B
Shucai Xiao's avatar
Shucai Xiao committed
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
        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
516
517
518
519
        return result;
    }
};

520
521
522
struct cpu_quant_gemm
{
    op::quant_dot op;
523
524
525
526
527
528
529

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

530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
    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
551
552
            args.at(0).visit(
                [&](auto input) { std::copy(input.begin(), input.end(), output.begin()); });
553
554
555
        });

        arg_1.visit([&](auto output) {
Shucai Xiao's avatar
Shucai Xiao committed
556
557
            args.at(1).visit(
                [&](auto input) { std::copy(input.begin(), input.end(), output.begin()); });
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
        });

        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
580
        migemm(result, arg_0, arg_1, op.alpha, int32_t{0});
581
582
583
584
585

        return result;
    }
};

Khalique's avatar
Khalique committed
586
587
588
589
590
591
struct leaky_relu_op
{
    op::leaky_relu op;
    std::string name() const { return "cpu::leaky_relu"; }
    auto fcn() const
    {
Paul's avatar
Paul committed
592
        auto a = op.alpha;
Khalique's avatar
Khalique committed
593
594
595
596
        return [a](auto x) { return x > 0 ? x : x * a; };
    }
};

Khalique's avatar
Khalique committed
597
598
599
600
601
602
struct elu_op
{
    op::elu op;
    std::string name() const { return "cpu::elu"; }
    auto fcn() const
    {
Paul's avatar
Paul committed
603
        auto a = op.alpha;
Khalique's avatar
Khalique committed
604
605
606
607
        return [a](auto x) { return x > 0 ? x : a * std::expm1(x); };
    }
};

Paul's avatar
Paul committed
608
609
610
611
template <typename Op>
struct cpu_unary
{
    Op op;
612
613
614
615
616
617

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return migraphx::reflect(self.op.op, f);
    }
Paul's avatar
Paul committed
618
    std::string name() const { return op.name(); }
Shucai Xiao's avatar
Shucai Xiao committed
619
    shape compute_shape(const std::vector<shape>& inputs) const
620
    {
Shucai Xiao's avatar
Shucai Xiao committed
621
622
623
        check_shapes{inputs}.has(1);
        auto s = inputs.at(0);
        if(s.packed())
624
        {
Shucai Xiao's avatar
Shucai Xiao committed
625
            return s;
626
627
628
        }
        else
        {
Shucai Xiao's avatar
Shucai Xiao committed
629
            return {s.type(), s.lens()};
630
631
632
        }
    }

Paul's avatar
Paul committed
633
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
634
635
636
637
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
Shucai Xiao's avatar
Shucai Xiao committed
638
                if(input.get_shape().standard())
639
640
641
642
643
644
645
646
647
                {
                    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
648
649
            });
        });
650

Paul's avatar
Paul committed
651
652
653
654
        return result;
    }
};

Khalique's avatar
Khalique committed
655
struct cpu_softmax
Paul's avatar
Paul committed
656
{
Khalique's avatar
Khalique committed
657
658
659
660
661
662
663
664
665
666
    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); }
Paul's avatar
Paul committed
667
    argument compute(context&, const shape& output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
668
669
    {
        argument result{output_shape};
Khalique's avatar
Khalique committed
670
        auto batch_lens     = output_shape.lens();
Shucai Xiao's avatar
Shucai Xiao committed
671
        std::size_t n_dims  = batch_lens[op.axis];
Khalique's avatar
Khalique committed
672
        batch_lens[op.axis] = 1;
673
674
        shape batch_shape{shape::int32_type, batch_lens};

Paul's avatar
Paul committed
675
676
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
Shucai Xiao's avatar
Shucai Xiao committed
677
678
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
679
680
            std::vector<value_type> batch_sum(batch_shape.elements(), value_type(0));
            par_for(batch_shape.elements(), [&](auto i) {
681
                auto idx = batch_shape.multi(i);
Shucai Xiao's avatar
Shucai Xiao committed
682
                for(std::size_t j = 0; j < n_dims; ++j)
683
684
685
686
                {
                    idx[op.axis] = j;
                    batch_max[i] = std::max(batch_max[i], input(idx.begin(), idx.end()));
                }
Khalique's avatar
Khalique committed
687

Shucai Xiao's avatar
Shucai Xiao committed
688
                for(std::size_t j = 0; j < n_dims; ++j)
689
                {
Shucai Xiao's avatar
Shucai Xiao committed
690
691
692
                    idx[op.axis]      = j;
                    std::size_t index = output_shape.index(idx);
                    output[index]     = std::exp(input[index] - batch_max[i]);
693
                }
Khalique's avatar
Khalique committed
694

Shucai Xiao's avatar
Shucai Xiao committed
695
                for(std::size_t j = 0; j < n_dims; ++j)
696
697
698
699
                {
                    idx[op.axis] = j;
                    batch_sum[i] += output(idx.begin(), idx.end());
                }
Khalique's avatar
Khalique committed
700

Shucai Xiao's avatar
Shucai Xiao committed
701
                for(std::size_t j = 0; j < n_dims; ++j)
702
703
704
705
                {
                    idx[op.axis] = j;
                    output(idx.begin(), idx.end()) /= batch_sum[i];
                }
Paul's avatar
Paul committed
706
707
            });
        });
Khalique's avatar
Khalique committed
708

Paul's avatar
Paul committed
709
710
711
712
        return result;
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
713
714
715
struct cpu_logsoftmax
{
    op::logsoftmax op;
716
717
718
719
720
721
722

    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
723
724
725
726
727
    std::string name() const { return "cpu::logsoftmax"; }
    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};
Shucai Xiao's avatar
Shucai Xiao committed
728
        auto batch_lens     = output_shape.lens();
Shucai Xiao's avatar
Shucai Xiao committed
729
        std::size_t n_dims  = batch_lens[op.axis];
730
731
732
        batch_lens[op.axis] = 1;
        shape batch_shape{shape::int32_type, batch_lens};

733
734
        // use a parallel implementation to acheive better performance
        // one thread for one batch
Shucai Xiao's avatar
Shucai Xiao committed
735
736
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
Shucai Xiao's avatar
Shucai Xiao committed
737
738
            std::vector<value_type> batch_max(batch_shape.elements(),
                                              std::numeric_limits<value_type>::lowest());
739
            std::vector<value_type> batch_sum(batch_shape.elements(), value_type(0));
Shucai Xiao's avatar
Shucai Xiao committed
740

741
            par_for(batch_shape.elements(), [&](auto i) {
742
                auto idx = batch_shape.multi(i);
Shucai Xiao's avatar
Shucai Xiao committed
743
                for(std::size_t j = 0; j < n_dims; ++j)
744
745
746
747
                {
                    idx[op.axis] = j;
                    batch_max[i] = std::max(batch_max[i], input(idx.begin(), idx.end()));
                }
Shucai Xiao's avatar
Shucai Xiao committed
748

Shucai Xiao's avatar
Shucai Xiao committed
749
                for(std::size_t j = 0; j < n_dims; ++j)
750
                {
Shucai Xiao's avatar
Shucai Xiao committed
751
752
753
                    idx[op.axis]      = j;
                    std::size_t index = output_shape.index(idx);
                    output[index]     = input[index] - batch_max[i];
754
755
                }

Shucai Xiao's avatar
Shucai Xiao committed
756
                for(std::size_t j = 0; j < n_dims; ++j)
757
758
759
760
                {
                    idx[op.axis] = j;
                    batch_sum[i] += std::exp(output(idx.begin(), idx.end()));
                }
Shucai Xiao's avatar
Shucai Xiao committed
761
762
763

                batch_sum[i] = std::log(batch_sum[i]);

Shucai Xiao's avatar
Shucai Xiao committed
764
                for(std::size_t j = 0; j < n_dims; ++j)
765
766
767
768
                {
                    idx[op.axis] = j;
                    output(idx.begin(), idx.end()) -= batch_sum[i];
                }
Shucai Xiao's avatar
Shucai Xiao committed
769
770
771
772
773
774
775
            });
        });

        return result;
    }
};

Paul's avatar
Paul committed
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
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
795
        apply_map["batch_norm_inference"] =
796
            extend_op<cpu_batch_norm_inference, op::batch_norm_inference>();
Shucai Xiao's avatar
Shucai Xiao committed
797
798
        apply_map["convolution"]       = extend_op<cpu_convolution, op::convolution>();
        apply_map["dot"]               = extend_op<cpu_gemm, op::dot>();
799
800
        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
801
802
803
804
805
806
        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>();
807
        apply_map["softmax"]           = extend_op<cpu_softmax, op::softmax>();
Paul's avatar
Paul committed
808
809
810
811
812
813
814
    }

    void apply()
    {
        init();
        for(auto it : iterator_for(*prog))
        {
Khalique's avatar
Khalique committed
815
            if(it->name() == "pooling")
Paul's avatar
Paul committed
816
817
818
            {
                apply_pooling(it);
            }
Paul's avatar
Paul committed
819
            else if(apply_map.count(it->name()) > 0)
Paul's avatar
Paul committed
820
            {
Paul's avatar
Paul committed
821
                apply_map.at(it->name())(it);
Paul's avatar
Paul committed
822
            }
Paul's avatar
Paul committed
823
            else if(is_context_free(it->get_operator()))
824
825
826
            {
                apply_cpu_op(it);
            }
Paul's avatar
Paul committed
827
828
829
        }
    }

830
831
832
833
834
    void apply_cpu_op(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_op{ins->get_operator()}, ins->inputs());
    }

Paul's avatar
Paul committed
835
836
837
    template <class T>
    void apply_simple_op(instruction_ref ins)
    {
Paul's avatar
Paul committed
838
        prog->replace_instruction(ins, T{}, ins->inputs());
Paul's avatar
Paul committed
839
840
841
842
843
    }

    template <class T, class Op>
    void apply_extend_op(instruction_ref ins)
    {
844
        auto&& op = any_cast<Op>(ins->get_operator());
Paul's avatar
Paul committed
845
        prog->replace_instruction(ins, T{op}, ins->inputs());
Paul's avatar
Paul committed
846
847
848
849
    }

    void apply_pooling(instruction_ref ins)
    {
850
        auto&& op = any_cast<op::pooling>(ins->get_operator());
Paul's avatar
Paul committed
851
        if(op.mode == "max")
Paul's avatar
Paul committed
852
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
853
        else if(op.mode == "average")
Paul's avatar
Paul committed
854
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->inputs());
Paul's avatar
Paul committed
855
856
857
    }
};

Shucai Xiao's avatar
Shucai Xiao committed
858
void lowering::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
859
860

} // namespace cpu
Paul's avatar
Paul committed
861
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
862
} // namespace migraphx