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

#include <rtg/cpu/cpu_target.hpp>
#include <rtg/instruction.hpp>
#include <rtg/dfor.hpp>
#include <rtg/operators.hpp>
Paul's avatar
Paul committed
6
#include <rtg/shape_for_each.hpp>
Paul's avatar
Paul committed
7

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

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

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

Paul's avatar
Paul committed
21
22
    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
23
    argument compute(shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
24
    {
Paul's avatar
Paul committed
25
        argument result{output_shape};
Paul's avatar
Paul committed
26
27
28
        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
29

Paul's avatar
Paul committed
30
31
32
            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
33

Paul's avatar
Paul committed
34
35
36
37
            dfor(output_shape.lens()[0],
                 output_shape.lens()[1],
                 output_shape.lens()[2],
                 output_shape.lens()[3])(
Paul's avatar
Paul committed
38
39
40
                [&](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
41

Paul's avatar
Paul committed
42
43
44
45
46
47
48
49
50
51
                    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
52
53
54
55
56
57
                });
        });
        return result;
    }
};

Paul's avatar
Paul committed
58
59
60
struct max_pool
{
    static std::string name() { return "max"; }
Paul's avatar
Paul committed
61
    static double start() { return std::numeric_limits<double>::lowest(); }
Paul's avatar
Paul committed
62

Paul's avatar
Paul committed
63
    static double apply(double x, double y) { return x + y; }
Paul's avatar
Paul committed
64

Paul's avatar
Paul committed
65
    static double final(double x, double) { return (x); }
Paul's avatar
Paul committed
66
67
68
69
70
};

struct avg_pool
{
    static std::string name() { return "average"; }
Paul's avatar
Paul committed
71
    static double start() { return 0.0; }
Paul's avatar
Paul committed
72
73
74
75
76
77
78

    static double apply(double x, double y)
    {
        double m = std::max(x, y);
        return (m);
    }

Paul's avatar
Paul committed
79
    static double final(double x, double y) { return x / y; }
Paul's avatar
Paul committed
80
81
};

Paul's avatar
Paul committed
82
template <class Op>
Paul's avatar
Paul committed
83
84
85
86
87
88
89
90
91
92
93
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;
Paul's avatar
Paul committed
94
95
            auto in_h  = input.get_shape().lens()[2];
            auto in_w  = input.get_shape().lens()[3];
Paul's avatar
Paul committed
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

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

131
132
133
struct cpu_transpose
{
    transpose op;
Paul's avatar
Paul committed
134
135

    std::string name() const { return "cpu::transpose"; }
136
137
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
    argument compute(shape output_shape, std::vector<argument> args) const
138
139
140
141
142
143
144
145
146
    {
        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
147
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
148
    argument compute(shape output_shape, std::vector<argument> args) const
149
150
151
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
152
            auto input_shape = args[0].get_shape();
Paul's avatar
Paul committed
153
            auto ndim        = output_shape.lens().size();
154
            using value_type = typename decltype(input)::value_type;
Paul's avatar
Paul committed
155
156
157
158
159
            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); });
160
            }
Paul's avatar
Paul committed
161
162
163
            else if(ndim == 3)
            {
                dfor(input_shape.lens()[0], input_shape.lens()[1], input_shape.lens()[2])(
164
                    [&](std::size_t i0, std::size_t i1, std::size_t i2) {
Paul's avatar
Paul committed
165
                        *ptr++ = input(i0, i1, i2);
166
167
                    });
            }
Paul's avatar
Paul committed
168
169
            else if(ndim == 4)
            {
170
171
172
173
174
                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
175
                        *ptr++ = input(i0, i1, i2, i3);
176
177
                    });
            }
Paul's avatar
Paul committed
178
179
            else if(ndim == 5)
            {
180
181
182
183
184
                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
185
186
187
188
189
                    [&](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); });
190
            }
Paul's avatar
Paul committed
191
192
            else if(ndim == 6)
            {
193
194
195
196
197
198
                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
199
200
201
202
203
204
                    [&](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); });
205
            }
206
        });
207
        return result;
208
    }
209
};
210

211
212
struct cpu_reshape
{
213
    reshape op;
214
    std::string name() const { return "cpu::reshape"; }
215
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
216

217
    argument compute(shape output_shape, std::vector<argument> args) const
218
219
220
221
222
    {
        return {output_shape, std::move(args.front().data)};
    }
};

223
224
225
226
struct cpu_gemm
{
    gemm op;
    std::string name() const { return "cpu::gemm"; }
227
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
228

229
    argument compute(shape output_shape, std::vector<argument> args) const
230
    {
231
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
232
233
234
235
236
237
238
239
        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();
240
241
242
243
244
245
            for(int ii = 0; ii < m; ii++)
            {
                for(int jj = 0; jj < n; jj++)
                {
                    c[ii * n + jj] = 0;
                }
246
            }
247
248
249
250
251
252
253
254
255
256
257
            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);
                    }
258
259
260
                }
            }
        });
261
        return result;
262
263
264
    }
};

265
struct identity_op
Paul's avatar
Paul committed
266
{
267
268
269
270
271
    std::string name() const { return "cpu::identity"; }
    auto fcn() const
    {
        return [](auto x) { return x; };
    }
272
};
Paul's avatar
Paul committed
273

274
struct abs_op
275
{
276
277
278
279
280
    std::string name() const { return "cpu::abs"; }
    auto fcn() const
    {
        return [](auto x) { return std::abs(x); };
    }
281
282
};

283
struct exp_op
284
{
285
286
287
288
289
    std::string name() const { return "cpu::exp"; }
    auto fcn() const
    {
        return [](auto x) { return std::exp(x); };
    }
290
291
};

292
struct sin_op
293
{
294
295
296
297
298
    std::string name() const { return "cpu::sin"; }
    auto fcn() const
    {
        return [](auto x) { return std::sin(x); };
    }
299
300
};

301
struct cos_op
302
{
303
304
305
306
307
    std::string name() const { return "cpu::cos"; }
    auto fcn() const
    {
        return [](auto x) { return std::cos(x); };
    }
308
309
};

310
struct tan_op
311
{
312
313
314
315
316
    std::string name() const { return "cpu::tan"; }
    auto fcn() const
    {
        return [](auto x) { return std::tan(x); };
    }
317
318
};

319
struct asin_op
320
{
321
322
323
324
325
    std::string name() const { return "cpu::asin"; }
    auto fcn() const
    {
        return [](auto x) { return std::asin(x); };
    }
326
327
};

328
struct acos_op
329
{
330
331
332
333
334
    std::string name() const { return "cpu::acos"; }
    auto fcn() const
    {
        return [](auto x) { return std::acos(x); };
    }
335
336
};

337
struct atan_op
338
{
339
340
341
342
343
    std::string name() const { return "cpu::atan"; }
    auto fcn() const
    {
        return [](auto x) { return std::atan(x); };
    }
344
345
346
347
};

struct tanh_op
{
348
349
350
351
352
    std::string name() const { return "cpu::tanh"; }
    auto fcn() const
    {
        return [](auto x) { return std::tanh(x); };
    }
353
354
355
356
};

struct sigmoid_op
{
357
358
359
360
361
    std::string name() const { return "cpu::sigmoid"; }
    auto fcn() const
    {
        return [](auto x) { return 1.f / (1.f + std::exp(-x)); };
    }
362
363
364
365
};

struct neg_op
{
366
367
368
369
370
    std::string name() const { return "cpu::neg"; }
    auto fcn() const
    {
        return [](auto x) { return -x; };
    }
371
372
373
374
};

struct relu_op
{
375
376
377
378
379
    std::string name() const { return "cpu::relu"; }
    auto fcn() const
    {
        return [](auto x) { return x > 0 ? x : 0; };
    }
380
381
382
383
384
};

template <typename Op>
struct cpu_unary
{
385
386
387
388
389
390
391
392
393
394
395
396
397
    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;
    }
398
399
};

400
struct softmax2d
401
{
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
    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;
    }
436
437
438
439
440
};

struct add_op
{
    std::string name() const { return "add"; }
441
442
443
444
    auto fcn() const
    {
        return [](auto x, auto y) { return x + y; };
    }
445
446
447
448
449
};

struct sub_op
{
    std::string name() const { return "sub"; }
450
451
452
453
    auto fcn() const
    {
        return [](auto x, auto y) { return x - y; };
    }
454
455
456
457
458
};

struct mul_op
{
    std::string name() const { return "mul"; }
459
460
461
462
    auto fcn() const
    {
        return [](auto x, auto y) { return x * y; };
    }
463
464
465
466
467
};

struct div_op
{
    std::string name() const { return "div"; }
468
469
470
471
    auto fcn() const
    {
        return [](auto x, auto y) { return x / y; };
    }
472
473
474
475
476
};

template <typename Op>
struct cpu_binary
{
477
478
479
480
481
482
483
    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) {
Paul's avatar
Paul committed
484
485
486
487
488
489
490
            if(input1.get_shape().packed() and input2.get_shape().packed()) {
                std::transform(input1.begin(), input1.end(), input2.begin(), output.begin(), op.fcn());
            } else {
                shape_for_each(output.get_shape(), [&](const auto& idx) {
                    output(idx.begin(), idx.end()) = op.fcn()(input1(idx.begin(), idx.end()), input2(idx.begin(), idx.end()));
                });
            }
491
492
493
        });
        return result;
    }
Paul's avatar
Paul committed
494
495
496
497
};

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

Paul's avatar
Paul committed
501
    template <class T>
Paul's avatar
Paul committed
502
    auto simple_op()
Paul's avatar
Paul committed
503
    {
Paul's avatar
Paul committed
504
        return [this](instruction_ref ins) { apply_simple_op<T>(ins); };
Paul's avatar
Paul committed
505
506
    }

Paul's avatar
Paul committed
507
    template <class T, class Op>
Paul's avatar
Paul committed
508
    auto extend_op()
Paul's avatar
Paul committed
509
    {
Paul's avatar
Paul committed
510
        return [this](instruction_ref ins) { apply_extend_op<T, Op>(ins); };
Paul's avatar
Paul committed
511
512
    }

Paul's avatar
Paul committed
513
    void init()
514
    {
Paul's avatar
Paul committed
515
        apply_map["convolution"] = extend_op<cpu_convolution, convolution>();
Paul's avatar
Paul committed
516
517
518
519
520
        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
521
        apply_map["identity"] = simple_op<cpu_unary<identity_op>>();
Paul's avatar
Paul committed
522
523
524
525
526
527
528
        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
529
        apply_map["add"]      = simple_op<cpu_binary<add_op>>();
Paul's avatar
Paul committed
530
531
532
        apply_map["sub"]      = simple_op<cpu_binary<sub_op>>();
        apply_map["mul"]      = simple_op<cpu_binary<mul_op>>();
        apply_map["div"]      = simple_op<cpu_binary<div_op>>();
Paul's avatar
Paul committed
533
534

        apply_map["softmax"] = simple_op<softmax2d>();
535
    }
Paul's avatar
Paul committed
536
537
538

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

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

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

Paul's avatar
Paul committed
570
571
572
573
    void apply_activation(instruction_ref ins)
    {
        auto&& op = any_cast<activation>(ins->op);
        if(op.mode == "relu")
574
            prog->replace_instruction(ins, cpu_unary<relu_op>{}, ins->arguments);
Paul's avatar
Paul committed
575
    }
576

Paul's avatar
Paul committed
577
    void apply_pooling(instruction_ref ins)
578
    {
Paul's avatar
Paul committed
579
580
581
582
583
        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);
584
    }
Paul's avatar
Paul committed
585
586
};

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

Paul's avatar
Paul committed
589
void cpu_target::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
590
591
592
593

} // namespace cpu

} // namespace rtg