lowering.cpp 29.1 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
7
8
9
10
11
12
13
14
15
#include <migraphx/op/batch_norm.hpp>
#include <migraphx/op/convolution.hpp>
#include <migraphx/op/dot.hpp>
#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
16
17
#include <migraphx/op/argmax.hpp>
#include <migraphx/op/argmin.hpp>
Paul's avatar
Paul committed
18
19
#include <migraphx/shape_for_each.hpp>
#include <migraphx/iterator_for.hpp>
Paul's avatar
Paul committed
20
#include <migraphx/par_dfor.hpp>
Paul's avatar
Paul committed
21
#include <migraphx/cpu/gemm.hpp>
Paul's avatar
Paul committed
22
#include <unordered_map>
Paul's avatar
Paul committed
23
#include <utility>
Paul's avatar
Paul committed
24

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

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

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

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

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

69
70
    std::string name() const { return "cpu::batch_norm_inference"; }

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

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

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

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

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

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

        return output;
    }
};

Khalique's avatar
Khalique committed
124
struct cpu_lrn
Khalique's avatar
Khalique committed
125
{
Khalique's avatar
Khalique committed
126
    op::lrn op;
Khalique's avatar
Khalique committed
127

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

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

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

Paul's avatar
Paul committed
167
168
struct cpu_convolution
{
169
    op::convolution op;
Paul's avatar
Paul committed
170

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

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

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

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

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

219
220
221
222
struct cpu_quant_convolution
{
    op::quant_convolution op;

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

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

255
                    int32_t acc = 0;
Shucai Xiao's avatar
Shucai Xiao committed
256
257
258
259
260
261
                    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
262
263
                            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
264
                        }
265
                    });
Shucai Xiao's avatar
Shucai Xiao committed
266
267
                    output(o, w, i, j) = acc;
                });
268
        });
Shucai Xiao's avatar
Shucai Xiao committed
269

270
271
272
273
        return result;
    }
};

Scott Thornton's avatar
Scott Thornton committed
274
275
struct cpu_im2col
{
276
    op::im2col op;
Scott Thornton's avatar
Scott Thornton committed
277

278
279
280
281
282
283
    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
284
285
    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
286

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

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

Paul's avatar
Paul committed
336
337
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
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
{
363
    op::pooling op;
Paul's avatar
Paul committed
364

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

Paul's avatar
Paul committed
371
    std::string name() const { return "cpu::pooling_" + Op::name(); }
Paul's avatar
Paul committed
372
373
    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
374
375
376
377
378
379
380
    {
        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
381
            par_dfor(output_shape.lens()[0],
Paul's avatar
Paul committed
382
383
384
                     output_shape.lens()[1],
                     output_shape.lens()[2],
                     output_shape.lens()[3])(
Paul's avatar
Paul committed
385
386
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
                [&](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;
    }
};

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

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

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

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

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

463
464
465
466
        return result;
    }
};

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

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

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

518
519
520
struct cpu_quant_gemm
{
    op::quant_dot op;
521
522
523
524
525
526
527

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

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

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

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

        return result;
    }
};

Khalique's avatar
Khalique committed
584
585
586
587
588
589
590
591
592
593
594
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
595
596
597
598
599
600
601
602
603
604
605
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
606
607
608
609
template <typename Op>
struct cpu_unary
{
    Op op;
610
611
612
613
614
615

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

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

Paul's avatar
Paul committed
649
650
651
652
        return result;
    }
};

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

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

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

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

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

Paul's avatar
Paul committed
707
708
709
710
        return result;
    }
};

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

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

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

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

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

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

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

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

        return result;
    }
};

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

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

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

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

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

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

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

} // namespace cpu
Paul's avatar
Paul committed
859
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
860
} // namespace migraphx