cpu_target.cpp 14.2 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;
    }
};

57
58
59
60
61
62
63
struct cpu_transpose
{
    transpose op;
   
    std::string name() const { return "cpu::transpose"; } 
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
    argument compute(shape output_shape, std::vector<argument> args) const
64
65
66
67
68
69
70
71
72
73
74
75
76
77
    {
        return {output_shape, std::move(args.front().data)};
    }
};

struct cpu_contiguous
{
    contiguous op;
    std::string name() const { return "cpu::contiguous"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        return op.compute_shape(inputs);
    }
    argument compute(shape output_shape, std::vector<argument> args) const
78
79
80
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
81
82
            auto input_shape = args[0].get_shape();
            auto ndim = output_shape.lens().size();
83
84
            using value_type = typename decltype(input)::value_type;
            value_type* ptr = static_cast<value_type*>(output.data());
85
86
87
88
89
90
            if (ndim == 2) {
                dfor(input_shape.lens()[0],
                     input_shape.lens()[1])(
                    [&](std::size_t i0, std::size_t i1) {
                        *ptr++ = input(i0,i1);
                    });
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
            else if (ndim == 3) {
                dfor(input_shape.lens()[0],
                     input_shape.lens()[1],
                     input_shape.lens()[2])(
                    [&](std::size_t i0, std::size_t i1, std::size_t i2) {
                        *ptr++ = input(i0,i1,i2);
                    });
            }
            else if (ndim == 4) {
                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) {
                        *ptr++ = input(i0,i1,i2,i3);
                    });
            }
            else if (ndim == 5) {
                dfor(input_shape.lens()[0],
                     input_shape.lens()[1],
                     input_shape.lens()[2],
                     input_shape.lens()[3],
                     input_shape.lens()[4])(
                    [&](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);
                    });
            }
            else if (ndim == 6) {
                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])(
                    [&](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);
                    });
138
            }
139
        });
140
        return result;
141
    }
142
};
143

144
145
struct cpu_reshape
{
146
    reshape op;
147
    std::string name() const { return "cpu::reshape"; }
148
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
149

150
    argument compute(shape output_shape, std::vector<argument> args) const
151
152
153
154
155
    {
        return {output_shape, std::move(args.front().data)};
    }
};

156
157
158
159
struct cpu_gemm
{
    gemm op;
    std::string name() const { return "cpu::gemm"; }
160
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
161

162
    argument compute(shape output_shape, std::vector<argument> args) const
163
    {
164
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
165
166
167
168
169
170
171
172
        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();
173
174
175
176
177
178
            for(int ii = 0; ii < m; ii++)
            {
                for(int jj = 0; jj < n; jj++)
                {
                    c[ii * n + jj] = 0;
                }
179
            }
180
181
182
183
184
185
186
187
188
189
190
            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);
                    }
191
192
193
                }
            }
        });
194
        return result;
195
196
197
    }
};

198
struct identity_op
Paul's avatar
Paul committed
199
{
200
201
202
203
204
    std::string name() const { return "cpu::identity"; }
    auto fcn() const
    {
        return [](auto x) { return x; };
    }
205
};
Paul's avatar
Paul committed
206

207
struct abs_op
208
{
209
210
211
212
213
    std::string name() const { return "cpu::abs"; }
    auto fcn() const
    {
        return [](auto x) { return std::abs(x); };
    }
214
215
};

216
struct exp_op
217
{
218
219
220
221
222
    std::string name() const { return "cpu::exp"; }
    auto fcn() const
    {
        return [](auto x) { return std::exp(x); };
    }
223
224
};

225
struct sin_op
226
{
227
228
229
230
231
    std::string name() const { return "cpu::sin"; }
    auto fcn() const
    {
        return [](auto x) { return std::sin(x); };
    }
232
233
};

234
struct cos_op
235
{
236
237
238
239
240
    std::string name() const { return "cpu::cos"; }
    auto fcn() const
    {
        return [](auto x) { return std::cos(x); };
    }
241
242
};

243
struct tan_op
244
{
245
246
247
248
249
    std::string name() const { return "cpu::tan"; }
    auto fcn() const
    {
        return [](auto x) { return std::tan(x); };
    }
250
251
};

252
struct asin_op
253
{
254
255
256
257
258
    std::string name() const { return "cpu::asin"; }
    auto fcn() const
    {
        return [](auto x) { return std::asin(x); };
    }
259
260
};

261
struct acos_op
262
{
263
264
265
266
267
    std::string name() const { return "cpu::acos"; }
    auto fcn() const
    {
        return [](auto x) { return std::acos(x); };
    }
268
269
};

270
struct atan_op
271
{
272
273
274
275
276
    std::string name() const { return "cpu::atan"; }
    auto fcn() const
    {
        return [](auto x) { return std::atan(x); };
    }
277
278
279
280
};

struct tanh_op
{
281
282
283
284
285
    std::string name() const { return "cpu::tanh"; }
    auto fcn() const
    {
        return [](auto x) { return std::tanh(x); };
    }
286
287
288
289
};

struct sigmoid_op
{
290
291
292
293
294
    std::string name() const { return "cpu::sigmoid"; }
    auto fcn() const
    {
        return [](auto x) { return 1.f / (1.f + std::exp(-x)); };
    }
295
296
297
298
};

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

struct relu_op
{
308
309
310
311
312
    std::string name() const { return "cpu::relu"; }
    auto fcn() const
    {
        return [](auto x) { return x > 0 ? x : 0; };
    }
313
314
315
316
317
};

template <typename Op>
struct cpu_unary
{
318
319
320
321
322
323
324
325
326
327
328
329
330
    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;
    }
331
332
};

333
struct softmax2d
334
{
335
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
363
364
365
366
367
368
    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;
    }
369
370
371
372
373
};

struct add_op
{
    std::string name() const { return "add"; }
374
375
376
377
    auto fcn() const
    {
        return [](auto x, auto y) { return x + y; };
    }
378
379
380
381
382
};

struct sub_op
{
    std::string name() const { return "sub"; }
383
384
385
386
    auto fcn() const
    {
        return [](auto x, auto y) { return x - y; };
    }
387
388
389
390
391
};

struct mul_op
{
    std::string name() const { return "mul"; }
392
393
394
395
    auto fcn() const
    {
        return [](auto x, auto y) { return x * y; };
    }
396
397
398
399
400
};

struct div_op
{
    std::string name() const { return "div"; }
401
402
403
404
    auto fcn() const
    {
        return [](auto x, auto y) { return x / y; };
    }
405
406
407
408
409
};

template <typename Op>
struct cpu_binary
{
410
411
412
413
414
415
416
417
418
419
420
    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
421
422
423
424
};

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

Paul's avatar
Paul committed
428
429
    template<class T>
    auto simple_op()
Paul's avatar
Paul committed
430
    {
Paul's avatar
Paul committed
431
        return [this](instruction_ref ins)
Paul's avatar
Paul committed
432
        {
Paul's avatar
Paul committed
433
434
            apply_simple_op<T>(ins);
        };
Paul's avatar
Paul committed
435
436
    }

Paul's avatar
Paul committed
437
438
    template<class T, class Op>
    auto extend_op()
Paul's avatar
Paul committed
439
    {
Paul's avatar
Paul committed
440
441
442
443
        return [this](instruction_ref ins)
        {
            apply_extend_op<T, Op>(ins);
        };
Paul's avatar
Paul committed
444
445
    }

Paul's avatar
Paul committed
446
    void init()
447
    {
Paul's avatar
Paul committed
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
        apply_map["convolution"] = extend_op<cpu_convolution, convolution>();
        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>();
        
        apply_map["identity"] = simple_op<cpu_unary<identity_op>>();
        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>>();

        apply_map["softmax"] = simple_op<softmax2d>();
464
465
    }

Paul's avatar
Paul committed
466
    void apply()
467
    {
Paul's avatar
Paul committed
468
469
470
471
472
473
474
475
476
477
478
479
        init();
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "activation")
            {
                apply_activation(it);
            } 
            else if(apply_map.count(it->op.name()) > 0)
            {
                apply_map.at(it->op.name())(it);
            }
        }
480
481
    }

Paul's avatar
Paul committed
482
483
    template<class T>
    void apply_simple_op(instruction_ref ins)
484
    {
Paul's avatar
Paul committed
485
        prog->replace_instruction(ins, T{}, ins->arguments);
486
487
    }

Paul's avatar
Paul committed
488
489
    template<class T, class Op>
    void apply_extend_op(instruction_ref ins)
490
    {
Paul's avatar
Paul committed
491
492
        auto&& op = any_cast<Op>(ins->op);
        prog->replace_instruction(ins, T{op}, ins->arguments);
493
494
    }

Paul's avatar
Paul committed
495
496
497
498
    void apply_activation(instruction_ref ins)
    {
        auto&& op = any_cast<activation>(ins->op);
        if(op.mode == "relu")
499
            prog->replace_instruction(ins, cpu_unary<relu_op>{}, ins->arguments);
Paul's avatar
Paul committed
500
501
502
    }
};

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

Paul's avatar
Paul committed
505
void cpu_target::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
506
507
508
509

} // namespace cpu

} // namespace rtg