cpu_target.cpp 16.9 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
5
6

#include <rtg/cpu/cpu_target.hpp>
#include <rtg/instruction.hpp>
#include <rtg/dfor.hpp>
#include <rtg/operators.hpp>

Paul's avatar
Paul committed
7
8
namespace rtg {
namespace cpu {
Paul's avatar
Paul committed
9

10
template <typename T>
11
12
13
14
T zero(const T&)
{
    return T(0);
}
15

Paul's avatar
Paul committed
16
17
18
19
struct cpu_convolution
{
    convolution op;

Paul's avatar
Paul committed
20
21
    std::string name() const { return "cpu::convolution"; }
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
Paul's avatar
Paul committed
22
    argument compute(shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
23
    {
Paul's avatar
Paul committed
24
        argument result{output_shape};
Paul's avatar
Paul committed
25
26
27
        visit_all(result, args[0], args[1])([&](auto output, auto input, auto weights) {
            auto in_h = input.get_shape().lens()[2];
            auto in_w = input.get_shape().lens()[3];
Paul's avatar
Paul committed
28

Paul's avatar
Paul committed
29
30
31
            auto wei_c = weights.get_shape().lens()[1];
            auto wei_h = weights.get_shape().lens()[2];
            auto wei_w = weights.get_shape().lens()[3];
Paul's avatar
Paul committed
32

Paul's avatar
Paul committed
33
34
35
36
            dfor(output_shape.lens()[0],
                 output_shape.lens()[1],
                 output_shape.lens()[2],
                 output_shape.lens()[3])(
Paul's avatar
Paul committed
37
38
39
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
                    const int start_x = i * op.stride[0] - op.padding[0];
                    const int start_y = j * op.stride[1] - op.padding[1];
Paul's avatar
Paul committed
40

Paul's avatar
Paul committed
41
42
43
44
45
46
47
48
49
50
                    double acc = 0;
                    dfor(wei_c, wei_h, wei_w)([&](std::size_t k, std::size_t x, std::size_t y) {
                        const int in_x = start_x + x;
                        const int in_y = start_y + y;
                        if(in_x >= 0 && in_x < in_h && in_y >= 0 && in_y < in_w)
                        {
                            acc += input(o, k, in_x, in_y) * weights(w, k, x, y);
                        }
                    });
                    output(o, w, i, j) = acc;
Paul's avatar
Paul committed
51
52
53
54
55
56
                });
        });
        return result;
    }
};

Paul's avatar
Paul committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
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)
    {
        return x + y;
    }

    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)
    {
        double m = std::max(x, y);
        return (m);
    }

    static double final(double x, double y)
    {
        return x / y;
    }
};


template<class Op>
struct cpu_pooling
{
    pooling op;

    std::string name() const { return "cpu::pooling_" + Op::name(); }
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        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];

            dfor(output_shape.lens()[0],
                 output_shape.lens()[1],
                 output_shape.lens()[2],
                 output_shape.lens()[3])(
                [&](std::size_t o, std::size_t w, std::size_t i, std::size_t j) {
                    const int start_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;
    }
};

147
148
149
struct cpu_transpose
{
    transpose op;
Paul's avatar
Paul committed
150
151

    std::string name() const { return "cpu::transpose"; }
152
153
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
    argument compute(shape output_shape, std::vector<argument> args) const
154
155
156
157
158
159
160
161
162
    {
        return {output_shape, std::move(args.front().data)};
    }
};

struct cpu_contiguous
{
    contiguous op;
    std::string name() const { return "cpu::contiguous"; }
Paul's avatar
Paul committed
163
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
164
    argument compute(shape output_shape, std::vector<argument> args) const
165
166
167
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
168
            auto input_shape = args[0].get_shape();
Paul's avatar
Paul committed
169
            auto ndim        = output_shape.lens().size();
170
            using value_type = typename decltype(input)::value_type;
Paul's avatar
Paul committed
171
172
173
174
175
            value_type* ptr  = static_cast<value_type*>(output.data());
            if(ndim == 2)
            {
                dfor(input_shape.lens()[0], input_shape.lens()[1])(
                    [&](std::size_t i0, std::size_t i1) { *ptr++ = input(i0, i1); });
176
            }
Paul's avatar
Paul committed
177
178
179
            else if(ndim == 3)
            {
                dfor(input_shape.lens()[0], input_shape.lens()[1], input_shape.lens()[2])(
180
                    [&](std::size_t i0, std::size_t i1, std::size_t i2) {
Paul's avatar
Paul committed
181
                        *ptr++ = input(i0, i1, i2);
182
183
                    });
            }
Paul's avatar
Paul committed
184
185
            else if(ndim == 4)
            {
186
187
188
189
190
                dfor(input_shape.lens()[0],
                     input_shape.lens()[1],
                     input_shape.lens()[2],
                     input_shape.lens()[3])(
                    [&](std::size_t i0, std::size_t i1, std::size_t i2, std::size_t i3) {
Paul's avatar
Paul committed
191
                        *ptr++ = input(i0, i1, i2, i3);
192
193
                    });
            }
Paul's avatar
Paul committed
194
195
            else if(ndim == 5)
            {
196
197
198
199
200
                dfor(input_shape.lens()[0],
                     input_shape.lens()[1],
                     input_shape.lens()[2],
                     input_shape.lens()[3],
                     input_shape.lens()[4])(
Paul's avatar
Paul committed
201
202
203
204
205
                    [&](std::size_t i0,
                        std::size_t i1,
                        std::size_t i2,
                        std::size_t i3,
                        std::size_t i4) { *ptr++ = input(i0, i1, i2, i3, i4); });
206
            }
Paul's avatar
Paul committed
207
208
            else if(ndim == 6)
            {
209
210
211
212
213
214
                dfor(input_shape.lens()[0],
                     input_shape.lens()[1],
                     input_shape.lens()[2],
                     input_shape.lens()[3],
                     input_shape.lens()[4],
                     input_shape.lens()[5])(
Paul's avatar
Paul committed
215
216
217
218
219
220
                    [&](std::size_t i0,
                        std::size_t i1,
                        std::size_t i2,
                        std::size_t i3,
                        std::size_t i4,
                        std::size_t i5) { *ptr++ = input(i0, i1, i2, i3, i4, i5); });
221
            }
222
        });
223
        return result;
224
    }
225
};
226

227
228
struct cpu_reshape
{
229
    reshape op;
230
    std::string name() const { return "cpu::reshape"; }
231
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
232

233
    argument compute(shape output_shape, std::vector<argument> args) const
234
235
236
237
238
    {
        return {output_shape, std::move(args.front().data)};
    }
};

239
240
241
242
struct cpu_gemm
{
    gemm op;
    std::string name() const { return "cpu::gemm"; }
243
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
244

245
    argument compute(shape output_shape, std::vector<argument> args) const
246
    {
247
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
248
249
250
251
252
253
254
255
        visit_all(result, args[0], args[1])([&](auto cmat, auto amat, auto bmat) {
            auto m = amat.get_shape().lens()[0];
            auto n = bmat.get_shape().lens()[1];
            auto k = bmat.get_shape().lens()[0];

            auto a = amat.data();
            auto b = bmat.data();
            auto c = cmat.data();
256
257
258
259
260
261
            for(int ii = 0; ii < m; ii++)
            {
                for(int jj = 0; jj < n; jj++)
                {
                    c[ii * n + jj] = 0;
                }
262
            }
263
264
265
266
267
268
269
270
271
272
273
            for(int ii = 0; ii < m; ii++)
            {
                for(int kk = 0; kk < k; kk++)
                {
                    auto aik  = a[ii * k + kk];
                    auto* bkj = &b[kk * n];
                    auto* cij = &c[ii * n];
                    for(int jj = 0; jj < n; jj++, cij++, bkj++)
                    {
                        *cij += aik * (*bkj);
                    }
274
275
276
                }
            }
        });
277
        return result;
278
279
280
    }
};

281
struct identity_op
Paul's avatar
Paul committed
282
{
283
284
285
286
287
    std::string name() const { return "cpu::identity"; }
    auto fcn() const
    {
        return [](auto x) { return x; };
    }
288
};
Paul's avatar
Paul committed
289

290
struct abs_op
291
{
292
293
294
295
296
    std::string name() const { return "cpu::abs"; }
    auto fcn() const
    {
        return [](auto x) { return std::abs(x); };
    }
297
298
};

299
struct exp_op
300
{
301
302
303
304
305
    std::string name() const { return "cpu::exp"; }
    auto fcn() const
    {
        return [](auto x) { return std::exp(x); };
    }
306
307
};

308
struct sin_op
309
{
310
311
312
313
314
    std::string name() const { return "cpu::sin"; }
    auto fcn() const
    {
        return [](auto x) { return std::sin(x); };
    }
315
316
};

317
struct cos_op
318
{
319
320
321
322
323
    std::string name() const { return "cpu::cos"; }
    auto fcn() const
    {
        return [](auto x) { return std::cos(x); };
    }
324
325
};

326
struct tan_op
327
{
328
329
330
331
332
    std::string name() const { return "cpu::tan"; }
    auto fcn() const
    {
        return [](auto x) { return std::tan(x); };
    }
333
334
};

335
struct asin_op
336
{
337
338
339
340
341
    std::string name() const { return "cpu::asin"; }
    auto fcn() const
    {
        return [](auto x) { return std::asin(x); };
    }
342
343
};

344
struct acos_op
345
{
346
347
348
349
350
    std::string name() const { return "cpu::acos"; }
    auto fcn() const
    {
        return [](auto x) { return std::acos(x); };
    }
351
352
};

353
struct atan_op
354
{
355
356
357
358
359
    std::string name() const { return "cpu::atan"; }
    auto fcn() const
    {
        return [](auto x) { return std::atan(x); };
    }
360
361
362
363
};

struct tanh_op
{
364
365
366
367
368
    std::string name() const { return "cpu::tanh"; }
    auto fcn() const
    {
        return [](auto x) { return std::tanh(x); };
    }
369
370
371
372
};

struct sigmoid_op
{
373
374
375
376
377
    std::string name() const { return "cpu::sigmoid"; }
    auto fcn() const
    {
        return [](auto x) { return 1.f / (1.f + std::exp(-x)); };
    }
378
379
380
381
};

struct neg_op
{
382
383
384
385
386
    std::string name() const { return "cpu::neg"; }
    auto fcn() const
    {
        return [](auto x) { return -x; };
    }
387
388
389
390
};

struct relu_op
{
391
392
393
394
395
    std::string name() const { return "cpu::relu"; }
    auto fcn() const
    {
        return [](auto x) { return x > 0 ? x : 0; };
    }
396
397
398
399
400
};

template <typename Op>
struct cpu_unary
{
401
402
403
404
405
406
407
408
409
410
411
412
413
    Op op;
    std::string name() const { return op.name(); }
    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        result.visit([&](auto output) {
            args[0].visit([&](auto input) {
                std::transform(input.begin(), input.end(), output.begin(), op.fcn());
            });
        });
        return result;
    }
414
415
};

416
struct softmax2d
417
{
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
    std::string name() const { return "cpu::softmax2d"; }
    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
            using value_type = typename decltype(input)::value_type;
            auto nb          = input.get_shape().lens()[0];
            auto nc          = input.get_shape().lens()[1];
            auto nh          = input.get_shape().lens()[2];
            auto nw          = input.get_shape().lens()[3];
            dfor(nb, nh, nw)([&](std::size_t b, std::size_t i, std::size_t j) {
                value_type cmax = std::numeric_limits<value_type>::lowest();
                for(int c = 0; c < nc; c++)
                {
                    cmax = std::max(cmax, input(b, c, i, j));
                }
                for(int c = 0; c < nc; c++)
                {
                    output(b, c, i, j) = std::exp(input(b, c, i, j) - cmax);
                }
                value_type sum = value_type(0);
                for(int c = 0; c < nc; c++)
                {
                    sum += output(b, c, i, j);
                }
                for(int c = 0; c < nc; c++)
                {
                    output(b, c, i, j) = output(b, c, i, j) / sum;
                }
            });
        });
        return result;
    }
452
453
454
455
456
};

struct add_op
{
    std::string name() const { return "add"; }
457
458
459
460
    auto fcn() const
    {
        return [](auto x, auto y) { return x + y; };
    }
461
462
463
464
465
};

struct sub_op
{
    std::string name() const { return "sub"; }
466
467
468
469
    auto fcn() const
    {
        return [](auto x, auto y) { return x - y; };
    }
470
471
472
473
474
};

struct mul_op
{
    std::string name() const { return "mul"; }
475
476
477
478
    auto fcn() const
    {
        return [](auto x, auto y) { return x * y; };
    }
479
480
481
482
483
};

struct div_op
{
    std::string name() const { return "div"; }
484
485
486
487
    auto fcn() const
    {
        return [](auto x, auto y) { return x / y; };
    }
488
489
490
491
492
};

template <typename Op>
struct cpu_binary
{
493
494
495
496
497
498
499
500
501
502
503
    Op op;
    std::string name() const { return op.name(); }
    shape compute_shape(std::vector<shape> inputs) const { return inputs.front(); }
    argument compute(shape output_shape, std::vector<argument> args) const
    {
        argument result{output_shape};
        visit_all(result, args[0], args[1])([&](auto output, auto input1, auto input2) {
            std::transform(input1.begin(), input1.end(), input2.begin(), output.begin(), op.fcn());
        });
        return result;
    }
Paul's avatar
Paul committed
504
505
506
507
};

struct cpu_apply
{
Paul's avatar
Paul committed
508
    program* prog;
Paul's avatar
Paul committed
509
    std::unordered_map<std::string, std::function<void(instruction_ref)>> apply_map{};
Paul's avatar
Paul committed
510

Paul's avatar
Paul committed
511
    template <class T>
Paul's avatar
Paul committed
512
    auto simple_op()
Paul's avatar
Paul committed
513
    {
Paul's avatar
Paul committed
514
        return [this](instruction_ref ins) { apply_simple_op<T>(ins); };
Paul's avatar
Paul committed
515
516
    }

Paul's avatar
Paul committed
517
    template <class T, class Op>
Paul's avatar
Paul committed
518
    auto extend_op()
Paul's avatar
Paul committed
519
    {
Paul's avatar
Paul committed
520
        return [this](instruction_ref ins) { apply_extend_op<T, Op>(ins); };
Paul's avatar
Paul committed
521
522
    }

Paul's avatar
Paul committed
523
    void init()
524
    {
Paul's avatar
Paul committed
525
        apply_map["convolution"] = extend_op<cpu_convolution, convolution>();
Paul's avatar
Paul committed
526
527
528
529
530
        apply_map["gemm"]        = extend_op<cpu_gemm, gemm>();
        apply_map["reshape"]     = extend_op<cpu_reshape, reshape>();
        apply_map["contiguous"]  = extend_op<cpu_contiguous, contiguous>();
        apply_map["transpose"]   = extend_op<cpu_transpose, transpose>();

Paul's avatar
Paul committed
531
        apply_map["identity"] = simple_op<cpu_unary<identity_op>>();
Paul's avatar
Paul committed
532
533
534
535
536
537
538
        apply_map["tanh"]     = simple_op<cpu_unary<tanh_op>>();
        apply_map["sigmoid"]  = simple_op<cpu_unary<sigmoid_op>>();
        apply_map["exp"]      = simple_op<cpu_unary<exp_op>>();
        apply_map["neg"]      = simple_op<cpu_unary<neg_op>>();
        apply_map["sin"]      = simple_op<cpu_unary<sin_op>>();
        apply_map["cos"]      = simple_op<cpu_unary<cos_op>>();
        apply_map["tan"]      = simple_op<cpu_unary<tan_op>>();
Paul's avatar
Paul committed
539
540

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

Paul's avatar
Paul committed
543
    void apply()
544
    {
Paul's avatar
Paul committed
545
546
547
548
549
550
        init();
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "activation")
            {
                apply_activation(it);
Paul's avatar
Paul committed
551
            }
Paul's avatar
Paul committed
552
553
554
555
556
            else if(apply_map.count(it->op.name()) > 0)
            {
                apply_map.at(it->op.name())(it);
            }
        }
557
558
    }

Paul's avatar
Paul committed
559
    template <class T>
Paul's avatar
Paul committed
560
    void apply_simple_op(instruction_ref ins)
561
    {
Paul's avatar
Paul committed
562
        prog->replace_instruction(ins, T{}, ins->arguments);
563
564
    }

Paul's avatar
Paul committed
565
    template <class T, class Op>
Paul's avatar
Paul committed
566
    void apply_extend_op(instruction_ref ins)
567
    {
Paul's avatar
Paul committed
568
569
        auto&& op = any_cast<Op>(ins->op);
        prog->replace_instruction(ins, T{op}, ins->arguments);
570
571
    }

Paul's avatar
Paul committed
572
573
574
575
    void apply_activation(instruction_ref ins)
    {
        auto&& op = any_cast<activation>(ins->op);
        if(op.mode == "relu")
576
            prog->replace_instruction(ins, cpu_unary<relu_op>{}, ins->arguments);
Paul's avatar
Paul committed
577
    }
Paul's avatar
Paul committed
578
579
580
581
582
583
584
585
586
587

    void apply_pooling(instruction_ref ins)
    {
        auto&& op = any_cast<pooling>(ins->op);
        if(op.mode == "max")
            prog->replace_instruction(ins, cpu_pooling<max_pool>{op}, ins->arguments);
        else if(op.mode == "average")
            prog->replace_instruction(ins, cpu_pooling<avg_pool>{op}, ins->arguments);

    }
Paul's avatar
Paul committed
588
589
};

Paul's avatar
Paul committed
590
std::string cpu_target::name() const { return "cpu"; }
Paul's avatar
Paul committed
591

Paul's avatar
Paul committed
592
void cpu_target::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
593
594
595
596

} // namespace cpu

} // namespace rtg