cpu_target.cpp 12.7 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
64
65
66
67
68
69
70
71
72
73
74
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
    {
        argument result{output_shape};
        visit_all(result, args[0])([&](auto output, auto input) {
            dfor(output_shape.lens()[0],
                 output_shape.lens()[0],
                 output_shape.lens()[0],
                 output_shape.lens()[0]
        });
    }
}

75
76
struct cpu_reshape
{
77
    reshape op;
78
    std::string name() const { return "cpu::reshape"; }
79
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
80

81
    argument compute(shape output_shape, std::vector<argument> args) const
82
83
84
85
86
    {
        return {output_shape, std::move(args.front().data)};
    }
};

87
88
89
90
struct cpu_gemm
{
    gemm op;
    std::string name() const { return "cpu::gemm"; }
91
    shape compute_shape(std::vector<shape> inputs) const { return op.compute_shape(inputs); }
92

93
    argument compute(shape output_shape, std::vector<argument> args) const
94
    {
95
        argument result{output_shape};
Scott Thornton's avatar
Scott Thornton committed
96
97
98
99
100
101
102
103
        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();
104
105
106
107
108
109
            for(int ii = 0; ii < m; ii++)
            {
                for(int jj = 0; jj < n; jj++)
                {
                    c[ii * n + jj] = 0;
                }
110
            }
111
112
113
114
115
116
117
118
119
120
121
            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);
                    }
122
123
124
                }
            }
        });
125
        return result;
126
127
128
    }
};

129
struct identity_op
Paul's avatar
Paul committed
130
{
131
132
133
134
135
    std::string name() const { return "cpu::identity"; }
    auto fcn() const
    {
        return [](auto x) { return x; };
    }
136
};
Paul's avatar
Paul committed
137

138
struct abs_op
139
{
140
141
142
143
144
    std::string name() const { return "cpu::abs"; }
    auto fcn() const
    {
        return [](auto x) { return std::abs(x); };
    }
145
146
};

147
struct exp_op
148
{
149
150
151
152
153
    std::string name() const { return "cpu::exp"; }
    auto fcn() const
    {
        return [](auto x) { return std::exp(x); };
    }
154
155
};

156
struct sin_op
157
{
158
159
160
161
162
    std::string name() const { return "cpu::sin"; }
    auto fcn() const
    {
        return [](auto x) { return std::sin(x); };
    }
163
164
};

165
struct cos_op
166
{
167
168
169
170
171
    std::string name() const { return "cpu::cos"; }
    auto fcn() const
    {
        return [](auto x) { return std::cos(x); };
    }
172
173
};

174
struct tan_op
175
{
176
177
178
179
180
    std::string name() const { return "cpu::tan"; }
    auto fcn() const
    {
        return [](auto x) { return std::tan(x); };
    }
181
182
};

183
struct asin_op
184
{
185
186
187
188
189
    std::string name() const { return "cpu::asin"; }
    auto fcn() const
    {
        return [](auto x) { return std::asin(x); };
    }
190
191
};

192
struct acos_op
193
{
194
195
196
197
198
    std::string name() const { return "cpu::acos"; }
    auto fcn() const
    {
        return [](auto x) { return std::acos(x); };
    }
199
200
};

201
struct atan_op
202
{
203
204
205
206
207
    std::string name() const { return "cpu::atan"; }
    auto fcn() const
    {
        return [](auto x) { return std::atan(x); };
    }
208
209
210
211
};

struct tanh_op
{
212
213
214
215
216
    std::string name() const { return "cpu::tanh"; }
    auto fcn() const
    {
        return [](auto x) { return std::tanh(x); };
    }
217
218
219
220
};

struct sigmoid_op
{
221
222
223
224
225
    std::string name() const { return "cpu::sigmoid"; }
    auto fcn() const
    {
        return [](auto x) { return 1.f / (1.f + std::exp(-x)); };
    }
226
227
228
229
};

struct neg_op
{
230
231
232
233
234
    std::string name() const { return "cpu::neg"; }
    auto fcn() const
    {
        return [](auto x) { return -x; };
    }
235
236
237
238
};

struct relu_op
{
239
240
241
242
243
    std::string name() const { return "cpu::relu"; }
    auto fcn() const
    {
        return [](auto x) { return x > 0 ? x : 0; };
    }
244
245
246
247
248
};

template <typename Op>
struct cpu_unary
{
249
250
251
252
253
254
255
256
257
258
259
260
261
    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;
    }
262
263
};

264
struct softmax2d
265
{
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
    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;
    }
300
301
302
303
304
};

struct add_op
{
    std::string name() const { return "add"; }
305
306
307
308
    auto fcn() const
    {
        return [](auto x, auto y) { return x + y; };
    }
309
310
311
312
313
};

struct sub_op
{
    std::string name() const { return "sub"; }
314
315
316
317
    auto fcn() const
    {
        return [](auto x, auto y) { return x - y; };
    }
318
319
320
321
322
};

struct mul_op
{
    std::string name() const { return "mul"; }
323
324
325
326
    auto fcn() const
    {
        return [](auto x, auto y) { return x * y; };
    }
327
328
329
330
331
};

struct div_op
{
    std::string name() const { return "div"; }
332
333
334
335
    auto fcn() const
    {
        return [](auto x, auto y) { return x / y; };
    }
336
337
338
339
340
};

template <typename Op>
struct cpu_binary
{
341
342
343
344
345
346
347
348
349
350
351
    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
352
353
354
355
};

struct cpu_apply
{
Paul's avatar
Paul committed
356
    program* prog;
Paul's avatar
Paul committed
357
358
359

    void apply()
    {
Paul's avatar
Paul committed
360
361
362
363
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
            if(it->op.name() == "convolution")
            {
Paul's avatar
Paul committed
364
                apply_convolution(it);
Paul's avatar
Paul committed
365
            }
366
367
368
369
370
371
372
373
            else if(it->op.name() == "gemm")
            {
                apply_gemm(it);
            }
            else if(it->op.name() == "reshape")
            {
                apply_reshape(it);
            }
Paul's avatar
Paul committed
374
375
            else if(it->op.name() == "activation")
            {
Paul's avatar
Paul committed
376
377
                apply_activation(it);
            }
378
379
380
381
382
383
384
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
            else if(it->op.name() == "identity")
            {
                apply_identity(it);
            }
            else if(it->op.name() == "softmax")
            {
                apply_softmax(it);
            }
            else if(it->op.name() == "tanh")
            {
                apply_tanh(it);
            }
            else if(it->op.name() == "sigmoid")
            {
                apply_sigmoid(it);
            }
            else if(it->op.name() == "exp")
            {
                apply_exp(it);
            }
            else if(it->op.name() == "neg")
            {
                apply_neg(it);
            }
            else if(it->op.name() == "sin")
            {
                apply_sin(it);
            }
            else if(it->op.name() == "cos")
            {
                apply_cos(it);
            }
            else if(it->op.name() == "tan")
            {
                apply_tan(it);
            }
Paul's avatar
Paul committed
414
415
416
417
418
419
420
421
422
        }
    }

    void apply_convolution(instruction_ref ins)
    {
        auto&& op = any_cast<convolution>(ins->op);
        prog->replace_instruction(ins, cpu_convolution{op}, ins->arguments);
    }

423
424
425
426
427
428
    void apply_gemm(instruction_ref ins)
    {
        auto&& op = any_cast<gemm>(ins->op);
        prog->replace_instruction(ins, cpu_gemm{op}, ins->arguments);
    }

429
430
431
432
433
434
    void apply_reshape(instruction_ref ins)
    {
        auto&& op = any_cast<reshape>(ins->op);
        prog->replace_instruction(ins, cpu_reshape{op}, ins->arguments);
    }

Paul's avatar
Paul committed
435
436
437
438
    void apply_activation(instruction_ref ins)
    {
        auto&& op = any_cast<activation>(ins->op);
        if(op.mode == "relu")
439
            prog->replace_instruction(ins, cpu_unary<relu_op>{}, ins->arguments);
Paul's avatar
Paul committed
440
    }
441
442
443
444
445
446
447
448

    void apply_identity(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<identity_op>{}, ins->arguments);
    }

    void apply_softmax(instruction_ref ins)
    {
449
        prog->replace_instruction(ins, softmax2d{}, ins->arguments);
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
    }

    void apply_tanh(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<tanh_op>{}, ins->arguments);
    }

    void apply_sigmoid(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<sigmoid_op>{}, ins->arguments);
    }

    void apply_exp(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<exp_op>{}, ins->arguments);
    }

    void apply_neg(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<neg_op>{}, ins->arguments);
    }

    void apply_sin(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<sin_op>{}, ins->arguments);
    }

    void apply_cos(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<cos_op>{}, ins->arguments);
    }

    void apply_tan(instruction_ref ins)
    {
        prog->replace_instruction(ins, cpu_unary<tan_op>{}, ins->arguments);
    }
Paul's avatar
Paul committed
486
487
};

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

Paul's avatar
Paul committed
490
void cpu_target::apply(program& p) const { cpu_apply{&p}.apply(); }
Paul's avatar
Paul committed
491
492
493
494

} // namespace cpu

} // namespace rtg