operators.hpp 13.7 KB
Newer Older
Paul's avatar
Paul committed
1
2
#ifndef MIGRAPH_GUARD_OPERATORS_HPP
#define MIGRAPH_GUARD_OPERATORS_HPP
Paul's avatar
Paul committed
3

4
#include <array>
Paul's avatar
Paul committed
5
#include <migraph/operation.hpp>
Paul's avatar
Paul committed
6
#include <migraph/check_shapes.hpp>
Paul's avatar
Paul committed
7
8
#include <migraph/stringutils.hpp>
#include <migraph/streamutils.hpp>
Paul's avatar
Paul committed
9
#include <cmath>
Paul's avatar
Paul committed
10

Paul's avatar
Paul committed
11
namespace migraph {
Paul's avatar
Paul committed
12

Paul's avatar
Paul committed
13
14
struct not_computable
{
Paul's avatar
Paul committed
15
16
17
18
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Paul's avatar
Paul committed
19
20
};

21
22
struct batch_norm_inference
{
23
24
25
26
    float epsilon  = 1.0e-6f;
    float momentum = 0.9f;
    bool spatial   = true;
    bool is_test   = false;
27
28
29

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

30
31
32
33
34
35
36
37
    enum bn_infer_mode_t
    {
        per_activation,
        spatial,
    };

    bn_infer_mode_t bn_mode = spatial;

38
39
40
41
42
43
44
45
46
47
48
49
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs, *this}.has(5);
        return inputs.front();
    }

    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
};

Paul's avatar
Paul committed
50
struct convolution
Paul's avatar
Paul committed
51
{
Paul's avatar
Paul committed
52
53
54
    std::array<std::size_t, 2> padding  = {{0, 0}};
    std::array<std::size_t, 2> stride   = {{1, 1}};
    std::array<std::size_t, 2> dilation = {{1, 1}};
Paul's avatar
Paul committed
55
56
57
58
59
60
61
    enum padding_mode_t
    {
        default_, // NOLINT
        same,
        valid
    };
    padding_mode_t padding_mode = default_;
Paul's avatar
Paul committed
62
    std::string name() const { return "convolution"; }
Paul's avatar
Paul committed
63
64
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
65
        check_shapes{inputs, *this}.has(2).same_type().same_ndims().only_dims(4);
Paul's avatar
Paul committed
66

Paul's avatar
Paul committed
67
        const shape& input   = inputs.at(0);
Paul's avatar
Paul committed
68
        const shape& weights = inputs.at(1);
Paul's avatar
Paul committed
69
        auto t               = input.type();
Paul's avatar
Paul committed
70
71
        if(padding_mode == default_)
        {
Paul's avatar
Paul committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
            return {t,
                    {
                        input.lens()[0],
                        weights.lens()[0],
                        std::size_t(std::max<std::ptrdiff_t>(
                            1,
                            (input.lens()[2] - (1 + dilation[0] * (weights.lens()[2] - 1)) +
                             2 * padding[0]) /
                                    stride[0] +
                                1)),
                        std::size_t(std::max<std::ptrdiff_t>(
                            1,
                            (input.lens()[3] - (1 + dilation[1] * (weights.lens()[3] - 1)) +
                             2 * padding[1]) /
                                    stride[1] +
                                1)),
                    }};
Paul's avatar
Paul committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
        }
        else if(padding_mode == same)
        {
            return {t,
                    {input.lens()[0],
                     weights.lens()[0],
                     static_cast<std::size_t>(
                         std::ceil(static_cast<double>(input.lens()[2]) / stride[0])),
                     static_cast<std::size_t>(
                         std::ceil(static_cast<double>(input.lens()[3]) / stride[1]))}};
        }
        else if(padding_mode == valid)
        {
            return {
                t,
                {input.lens()[0],
                 weights.lens()[0],
                 static_cast<std::size_t>(std::ceil(
                     static_cast<double>(input.lens()[2] - weights.lens()[2] + 1) / stride[0])),
                 static_cast<std::size_t>(std::ceil(
                     static_cast<double>(input.lens()[3] - weights.lens()[3] + 1) / stride[1]))}};
        }
        else
        {
Paul's avatar
Paul committed
113
            MIGRAPH_THROW("Invalid padding mode");
Paul's avatar
Paul committed
114
        }
Paul's avatar
Paul committed
115
    }
Paul's avatar
Paul committed
116

Paul's avatar
Paul committed
117
118
119
120
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Paul's avatar
Paul committed
121

Paul's avatar
Paul committed
122
    friend std::ostream& operator<<(std::ostream& os, const convolution& op)
Paul's avatar
Paul committed
123
    {
Paul's avatar
Paul committed
124
125
126
127
128
        os << op.name() << "[";
        os << "padding={" << stream_range(op.padding) << "}, ";
        os << "stride={" << stream_range(op.stride) << "}, ";
        os << "dilation={" << stream_range(op.dilation) << "}";
        os << "]";
Paul's avatar
Paul committed
129
130
        return os;
    }
Paul's avatar
Paul committed
131
132
};

Paul's avatar
Paul committed
133
struct pooling
Paul's avatar
Paul committed
134
135
{
    std::string mode;
Paul's avatar
Paul committed
136
137
138
    std::array<std::size_t, 2> padding = {{0, 0}};
    std::array<std::size_t, 2> stride  = {{1, 1}};
    std::array<std::size_t, 2> lengths = {{1, 1}};
Paul's avatar
Paul committed
139
    std::string name() const { return "pooling"; }
Scott Thornton's avatar
Scott Thornton committed
140

Paul's avatar
Paul committed
141
142
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
143
        check_shapes{inputs, *this}.has(1).only_dims(4);
Paul's avatar
Paul committed
144

Paul's avatar
Paul committed
145
        const shape& input = inputs.at(0);
Paul's avatar
Paul committed
146
        auto t             = input.type();
Paul's avatar
Paul committed
147

Paul's avatar
Paul committed
148
149
        assert(lengths[0] < (input.lens()[2] + 2 * padding[0]));
        assert(lengths[1] < (input.lens()[3] + 2 * padding[1]));
Paul's avatar
Paul committed
150

Scott Thornton's avatar
Scott Thornton committed
151
152
153
154
155
156
        return {t,
                {
                    input.lens()[0],
                    input.lens()[1],
                    std::size_t(std::max<std::ptrdiff_t>(
                        1,
Paul's avatar
Paul committed
157
                        std::ptrdiff_t(std::ceil((input.lens()[2] + 2 * padding[0] - lengths[0]) /
Paul's avatar
Paul committed
158
                                                 static_cast<float>(stride[0]))) +
Scott Thornton's avatar
Scott Thornton committed
159
160
161
                            1)),
                    std::size_t(std::max<std::ptrdiff_t>(
                        1,
Paul's avatar
Paul committed
162
                        std::ptrdiff_t(std::ceil((input.lens()[3] + 2 * padding[1] - lengths[1]) /
Paul's avatar
Paul committed
163
                                                 static_cast<float>(stride[1]))) +
Scott Thornton's avatar
Scott Thornton committed
164
165
                            1)),
                }};
Paul's avatar
Paul committed
166
    }
Paul's avatar
Paul committed
167

Paul's avatar
Paul committed
168
169
170
171
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Paul's avatar
Paul committed
172

Paul's avatar
Paul committed
173
    friend std::ostream& operator<<(std::ostream& os, const pooling& op)
Paul's avatar
Paul committed
174
    {
Paul's avatar
Paul committed
175
176
177
178
179
        os << op.name() << "[";
        os << "padding={" << stream_range(op.padding) << "}, ";
        os << "stride={" << stream_range(op.stride) << "}, ";
        os << "lengths={" << stream_range(op.lengths) << "}";
        os << "]";
Paul's avatar
Paul committed
180
181
        return os;
    }
Paul's avatar
Paul committed
182
183
};

Paul's avatar
Paul committed
184
struct activation
Paul's avatar
Paul committed
185
186
{
    std::string mode;
Paul's avatar
Paul committed
187
    std::string name() const { return "activation"; }
Paul's avatar
Paul committed
188
189
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
190
        check_shapes{inputs, *this}.has(1);
Paul's avatar
Paul committed
191
192
        return inputs.front();
    }
Paul's avatar
Paul committed
193

Paul's avatar
Paul committed
194
195
196
197
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Paul's avatar
Paul committed
198
    friend std::ostream& operator<<(std::ostream& os, const activation& op)
Paul's avatar
Paul committed
199
    {
Paul's avatar
Paul committed
200
        os << op.name() << ":" << op.mode;
Paul's avatar
Paul committed
201
202
        return os;
    }
Paul's avatar
Paul committed
203
204
};

205
206
207
208
209
210
struct transpose
{
    std::vector<int64_t> dims;
    std::string name() const { return "transpose"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
211
        check_shapes{inputs, *this}.has(1);
212
        auto input         = inputs.at(0);
213
        auto input_lens    = input.lens();
214
215
        auto input_strides = input.strides();
        auto t             = input.type();
Paul's avatar
Paul committed
216
217
        if(dims.size() != input_lens.size())
        {
Paul's avatar
Paul committed
218
            MIGRAPH_THROW("Permutation has wrong number of axes");
219
220
221
        }
        std::vector<int64_t> axes(dims.size());
        std::iota(axes.begin(), axes.end(), 0);
Paul's avatar
Paul committed
222
223
        if(!std::is_permutation(axes.begin(), axes.end(), dims.begin()))
        {
Paul's avatar
Paul committed
224
            MIGRAPH_THROW("Invalid permutation");
225
        }
226
227
        std::vector<size_t> output_lens(input_lens.size());
        std::vector<size_t> output_strides(input_lens.size());
Paul's avatar
Paul committed
228
229
230
        for(int i = 0; i < output_lens.size(); i++)
        {
            output_lens[i]    = input_lens[dims[i]];
231
232
            output_strides[i] = input_strides[dims[i]];
        }
233
        return {t, output_lens, output_strides};
234
    }
Paul's avatar
Paul committed
235
236
237
238
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
239
240
};

Paul's avatar
Paul committed
241
struct contiguous
242
243
244
245
{
    std::string name() const { return "contiguous"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
246
        check_shapes{inputs, *this}.has(1);
Paul's avatar
Paul committed
247
248
249
250
        auto lens = inputs.at(0).lens();
        auto t    = inputs.at(0).type();
        if(lens.size() < 2)
        {
Paul's avatar
Paul committed
251
            MIGRAPH_THROW("Number of dimensions should exceed 1");
252
253
254
        }
        return {t, lens};
    }
Paul's avatar
Paul committed
255
256
257
258
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
259
260
};

Paul's avatar
Paul committed
261
262
263
struct reshape
{
    std::vector<int64_t> dims;
Paul's avatar
Paul committed
264
    std::string name() const { return "reshape"; }
Paul's avatar
Paul committed
265
266
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
267
        check_shapes{inputs, *this}.has(1);
Paul's avatar
Paul committed
268
269
        auto&& idims = inputs.front().lens();
        std::vector<std::size_t> rdims(dims.begin(), dims.end());
270
271
272
        auto n_neg_dims = std::count(dims.begin(), dims.end(), -1);
        if(n_neg_dims > 1)
            MIGRAPH_THROW("Dimensions for reshape can only have one -1 dim");
Paul's avatar
Paul committed
273
        for(std::size_t i = 0; i < dims.size(); i++)
Paul's avatar
Paul committed
274
275
276
277
        {
            if(dims[i] == 0)
                rdims[i] = idims[i];
        }
278
279
280
281
282
283
284
285
286
287
288
        if(n_neg_dims > 0)
        {
            size_t missing_dim =
                -inputs.front().elements() /
                std::accumulate(rdims.begin(), rdims.end(), 1, std::multiplies<int64_t>());
            for(std::size_t i = 0; i < rdims.size(); i++)
            {
                if(dims[i] == -1)
                    rdims[i] = missing_dim;
            }
        }
Paul's avatar
Paul committed
289
290
291
        if(dims.back() == -1)
        {
            rdims.pop_back();
Paul's avatar
Paul committed
292
            std::copy(idims.begin() + rdims.size(), idims.end(), std::back_inserter(rdims));
Paul's avatar
Paul committed
293
        }
Scott Thornton's avatar
Scott Thornton committed
294
        shape s{inputs.front().type(), rdims};
Paul's avatar
Paul committed
295
        if(s.elements() != inputs.front().elements())
Paul's avatar
Paul committed
296
            MIGRAPH_THROW("Wrong number of elements for reshape");
Scott Thornton's avatar
Scott Thornton committed
297
        return s;
Paul's avatar
Paul committed
298
299
    }

Paul's avatar
Paul committed
300
301
302
303
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Paul's avatar
Paul committed
304

Paul's avatar
Paul committed
305
    friend std::ostream& operator<<(std::ostream& os, const reshape& op)
Paul's avatar
Paul committed
306
    {
Paul's avatar
Paul committed
307
308
309
        os << op.name() << "[";
        os << "dims={" << stream_range(op.dims) << "}, ";
        os << "]";
Paul's avatar
Paul committed
310
311
        return os;
    }
Paul's avatar
Paul committed
312
313
};

314
315
struct gemm
{
316
    std::string name() const { return "gemm"; }
317
318
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
319
        check_shapes{inputs, *this}.has(2).same_type();
320
321
        const shape& a = inputs.at(0);
        const shape& b = inputs.at(1);
Scott Thornton's avatar
Scott Thornton committed
322
        auto t         = a.type();
323

324
        if(a.lens()[1] != b.lens()[0])
Paul's avatar
Paul committed
325
            MIGRAPH_THROW("Inner dimensions do not match");
Scott Thornton's avatar
Scott Thornton committed
326
        return {t, {a.lens()[0], b.lens()[1]}};
327
    }
328

Paul's avatar
Paul committed
329
330
331
332
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
333
334

    friend std::ostream& operator<<(std::ostream& os, const gemm& op)
335
336
    {
        os << op.name() << "[";
337
        os << "]";
Scott Thornton's avatar
Scott Thornton committed
338
        return os;
339
340
341
    }
};

342
struct unary
Scott Thornton's avatar
Scott Thornton committed
343
{
344
345
    shape compute_shape(std::vector<shape> inputs) const
    {
346
347
        check_shapes{inputs}.has(1);
        return inputs.at(0);
348
    }
Paul's avatar
Paul committed
349
350
351
352
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Scott Thornton's avatar
Scott Thornton committed
353
354
};

355
356
struct identity : unary
{
357
    std::string name() const { return "identity"; }
358
359
360
};

struct abs : unary
Scott Thornton's avatar
Scott Thornton committed
361
{
362
    std::string name() const { return "abs"; }
Scott Thornton's avatar
Scott Thornton committed
363
364
};

365
struct exp : unary
Scott Thornton's avatar
Scott Thornton committed
366
{
367
    std::string name() const { return "exp"; }
Scott Thornton's avatar
Scott Thornton committed
368
369
};

370
struct sin : unary
Scott Thornton's avatar
Scott Thornton committed
371
{
372
    std::string name() const { return "sin"; }
Scott Thornton's avatar
Scott Thornton committed
373
374
};

375
struct cos : unary
Scott Thornton's avatar
Scott Thornton committed
376
{
377
    std::string name() const { return "cos"; }
Scott Thornton's avatar
Scott Thornton committed
378
379
};

380
struct tan : unary
Scott Thornton's avatar
Scott Thornton committed
381
{
382
    std::string name() const { return "tan"; }
Scott Thornton's avatar
Scott Thornton committed
383
384
};

385
struct asin : unary
Scott Thornton's avatar
Scott Thornton committed
386
{
387
    std::string name() const { return "asin"; }
Scott Thornton's avatar
Scott Thornton committed
388
389
};

390
struct acos : unary
Scott Thornton's avatar
Scott Thornton committed
391
{
392
    std::string name() const { return "acos"; }
Scott Thornton's avatar
Scott Thornton committed
393
394
};

395
struct atan : unary
Scott Thornton's avatar
Scott Thornton committed
396
{
397
    std::string name() const { return "atan"; }
Scott Thornton's avatar
Scott Thornton committed
398
399
};

400
struct softmax : unary
Scott Thornton's avatar
Scott Thornton committed
401
{
402
    std::string name() const { return "softmax"; }
Scott Thornton's avatar
Scott Thornton committed
403
404
};

405
struct tanh : unary
Scott Thornton's avatar
Scott Thornton committed
406
{
407
    std::string name() const { return "tanh"; }
Scott Thornton's avatar
Scott Thornton committed
408
409
};

410
struct sigmoid : unary
Scott Thornton's avatar
Scott Thornton committed
411
{
412
    std::string name() const { return "sigmoid"; }
Scott Thornton's avatar
Scott Thornton committed
413
414
};

415
struct neg : unary
Scott Thornton's avatar
Scott Thornton committed
416
{
417
    std::string name() const { return "neg"; }
Scott Thornton's avatar
Scott Thornton committed
418
419
};

420
struct flatten
Scott Thornton's avatar
Scott Thornton committed
421
422
423
424
{
    std::string name() const { return "flatten"; }
};

425
426
427
428
429
430
struct broadcast
{
    uint64_t axis = 0;
    std::string name() const { return "broadcast"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
431
432
433
434
        auto t      = inputs.at(0).type();
        auto result = inputs.at(0);
        auto input  = inputs.at(1);

Paul's avatar
Paul committed
435
        std::vector<size_t> bcast_strides(result.lens().size(), 0);
436

Paul's avatar
Paul committed
437
438
        if(std::all_of(
               result.lens().cbegin(), result.lens().cend(), [&](auto x) { return x == 1; }))
439
        {
Scott Thornton's avatar
Scott Thornton committed
440
            if(axis != 0)
Paul's avatar
Paul committed
441
                MIGRAPH_THROW("when broadcasting tensor of size 1, axis should be 0");
Paul's avatar
Paul committed
442
            return {t, result.lens(), std::move(bcast_strides)};
443
444
445
        }
        else
        {
Paul's avatar
Paul committed
446
447
            assert(result.lens().size() - axis >= input.lens().size());
            if(!std::equal(input.lens().begin(), input.lens().end(), result.lens().begin() + axis))
Paul's avatar
Paul committed
448
                MIGRAPH_THROW("when broadcasting success sizes must match");
Paul's avatar
Paul committed
449
            std::copy(input.strides().begin(), input.strides().end(), bcast_strides.begin() + axis);
Paul's avatar
Paul committed
450
            return {t, result.lens(), std::move(bcast_strides)};
451
452
        }
    }
Paul's avatar
Paul committed
453
    argument compute(context&, shape output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
454
    {
455
        return {output_shape, std::move(args.at(1).data)};
Scott Thornton's avatar
Scott Thornton committed
456
    }
457
458
};

459
struct binary
Scott Thornton's avatar
Scott Thornton committed
460
{
461
    uint64_t broadcast = 0;
462
463
    shape compute_shape(std::vector<shape> inputs) const
    {
464
465
        check_shapes{inputs}.has(2).same_type().same_dims();
        return inputs.at(0);
466
    }
Paul's avatar
Paul committed
467
468
469
470
    argument compute(context&, shape, std::vector<argument>) const
    {
        MIGRAPH_THROW("not computable");
    }
Scott Thornton's avatar
Scott Thornton committed
471
472
};

473
474
475
476
477
478
struct add : binary
{
    std::string name() const { return "add"; }
};

struct sub : binary
Scott Thornton's avatar
Scott Thornton committed
479
480
481
482
{
    std::string name() const { return "sub"; }
};

483
struct mul : binary
Scott Thornton's avatar
Scott Thornton committed
484
485
486
487
{
    std::string name() const { return "mul"; }
};

488
struct div : binary
Scott Thornton's avatar
Scott Thornton committed
489
490
491
492
{
    std::string name() const { return "div"; }
};

Paul's avatar
Paul committed
493
struct outline
Scott Thornton's avatar
Scott Thornton committed
494
{
Paul's avatar
Paul committed
495
496
497
498
    shape s;
    std::string name() const { return "outline"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
499
        check_shapes{inputs, *this}.has(0);
Paul's avatar
Paul committed
500
501
        return s;
    }
Paul's avatar
Paul committed
502
    argument compute(context&, shape, std::vector<argument>) const { return {s, nullptr}; }
Scott Thornton's avatar
Scott Thornton committed
503
504
};

Paul's avatar
Paul committed
505
} // namespace migraph
Paul's avatar
Paul committed
506
507

#endif