operators.hpp 27 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
6
7
8
9
#include <migraphx/operation.hpp>
#include <migraphx/check_shapes.hpp>
#include <migraphx/stringutils.hpp>
#include <migraphx/streamutils.hpp>
#include <migraphx/config.hpp>
Paul's avatar
Paul committed
10
#include <cmath>
Paul's avatar
Paul committed
11
#include <utility>
Paul's avatar
Paul committed
12

Paul's avatar
Paul committed
13
namespace migraphx {
14
inline namespace MIGRAPH_INLINE_NS {
15
namespace op {
Paul's avatar
Paul committed
16

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

25
26
struct batch_norm_inference
{
27
28
    float epsilon  = 1.0e-6f;
    float momentum = 0.9f;
29
30
31

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

32
33
34
35
36
37
38
39
    enum bn_infer_mode_t
    {
        per_activation,
        spatial,
    };

    bn_infer_mode_t bn_mode = spatial;

Paul's avatar
Paul committed
40
41
42
43
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return pack(
Paul's avatar
Paul committed
44
            f(self.epsilon, "epsilon"), f(self.momentum, "momentum"), f(self.bn_mode, "bn_mode"));
Paul's avatar
Paul committed
45
    }
46

47
48
49
50
51
52
53
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs, *this}.has(5);
        return inputs.front();
    }
};

Paul's avatar
Paul committed
54
struct convolution
Paul's avatar
Paul committed
55
{
Paul's avatar
Paul committed
56
57
58
    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
59
60
61
62
63
64
65
    enum padding_mode_t
    {
        default_, // NOLINT
        same,
        valid
    };
    padding_mode_t padding_mode = default_;
Paul's avatar
Paul committed
66
67
68
69

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
70
71
72
73
        return pack(f(self.padding, "padding"),
                    f(self.stride, "stride"),
                    f(self.dilation, "dilation"),
                    f(self.padding_mode, "padding_mode"));
Paul's avatar
Paul committed
74
75
    }

Paul's avatar
Paul committed
76
    std::string name() const { return "convolution"; }
Paul's avatar
Paul committed
77
78
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
79
        check_shapes{inputs, *this}.has(2).same_type().same_ndims().only_dims(4);
Paul's avatar
Paul committed
80

Paul's avatar
Paul committed
81
        const shape& input   = inputs.at(0);
Paul's avatar
Paul committed
82
        const shape& weights = inputs.at(1);
Paul's avatar
Paul committed
83
        auto t               = input.type();
Paul's avatar
Paul committed
84
85
        if(padding_mode == default_)
        {
Paul's avatar
Paul committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
            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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
        }
        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
127
            MIGRAPH_THROW("Invalid padding mode");
Paul's avatar
Paul committed
128
        }
Paul's avatar
Paul committed
129
130
131
    }
};

Scott Thornton's avatar
Scott Thornton committed
132
133
struct im2col
{
Scott Thornton's avatar
Scott Thornton committed
134
135
136
137
138
139
140
141
142
    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}};
    enum padding_mode_t
    {
        default_, // NOLINT
        same,
        valid
    };
Paul's avatar
Paul committed
143
144
145
146
147
    padding_mode_t padding_mode = default_;

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
148
149
150
151
        return pack(f(self.padding, "padding"),
                    f(self.stride, "stride"),
                    f(self.dilation, "dilation"),
                    f(self.padding_mode, "padding_mode"));
Paul's avatar
Paul committed
152
    }
Scott Thornton's avatar
Scott Thornton committed
153
154
155
156
157

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

    shape compute_shape(std::vector<shape> inputs) const
    {
Scott Thornton's avatar
Scott Thornton committed
158
159
160
        auto input          = inputs[0];
        auto weights        = inputs[1];
        auto batch_size     = input.lens()[0];
Scott Thornton's avatar
Scott Thornton committed
161
        auto input_channels = weights.lens()[1];
Scott Thornton's avatar
Scott Thornton committed
162
163
        auto kernel_height  = weights.lens()[2];
        auto kernel_width   = weights.lens()[3];
Scott Thornton's avatar
Scott Thornton committed
164
        check_shapes{inputs, *this}.has(2);
Scott Thornton's avatar
Scott Thornton committed
165
166
        if(batch_size != 1)
            MIGRAPH_THROW("im2col only support batch_size 1");
Scott Thornton's avatar
Scott Thornton committed
167
        auto output_height = std::size_t(std::max<std::ptrdiff_t>(
Scott Thornton's avatar
Scott Thornton committed
168
169
170
            1,
            (input.lens()[2] - (1 + dilation[0] * (kernel_height - 1)) + 2 * padding[0]) /
                    stride[0] +
Scott Thornton's avatar
Scott Thornton committed
171
                1));
Scott Thornton's avatar
Scott Thornton committed
172
173
174
175
        auto output_width  = std::size_t(std::max<std::ptrdiff_t>(
            1,
            (input.lens()[3] - (1 + dilation[1] * (kernel_width - 1)) + 2 * padding[1]) /
                    stride[1] +
Scott Thornton's avatar
Scott Thornton committed
176
                1));
Scott Thornton's avatar
Scott Thornton committed
177
178
        auto channels_col  = kernel_height * kernel_width * input_channels;
        return {input.type(), {output_height * output_width, channels_col}};
Scott Thornton's avatar
Scott Thornton committed
179
180
181
    }
};

Paul's avatar
Paul committed
182
struct pooling
Paul's avatar
Paul committed
183
{
Paul's avatar
Paul committed
184
    std::string mode                   = "average";
Paul's avatar
Paul committed
185
186
187
    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
188
189
190
191

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
192
193
194
195
        return pack(f(self.mode, "mode"),
                    f(self.padding, "padding"),
                    f(self.stride, "stride"),
                    f(self.lengths, "lengths"));
Paul's avatar
Paul committed
196
197
    }

Paul's avatar
Paul committed
198
    std::string name() const { return "pooling"; }
Scott Thornton's avatar
Scott Thornton committed
199

Paul's avatar
Paul committed
200
201
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
202
        check_shapes{inputs, *this}.has(1).only_dims(4);
Paul's avatar
Paul committed
203

Paul's avatar
Paul committed
204
        const shape& input = inputs.at(0);
Paul's avatar
Paul committed
205
        auto t             = input.type();
Paul's avatar
Paul committed
206

Paul's avatar
Paul committed
207
208
        assert(lengths[0] <= (input.lens()[2] + 2 * padding[0]));
        assert(lengths[1] <= (input.lens()[3] + 2 * padding[1]));
Paul's avatar
Paul committed
209

Scott Thornton's avatar
Scott Thornton committed
210
211
212
213
214
215
        return {t,
                {
                    input.lens()[0],
                    input.lens()[1],
                    std::size_t(std::max<std::ptrdiff_t>(
                        1,
Paul's avatar
Paul committed
216
                        std::ptrdiff_t(std::floor((input.lens()[2] + 2 * padding[0] - lengths[0]) /
Paul's avatar
Paul committed
217
                                                  static_cast<float>(stride[0]))) +
Scott Thornton's avatar
Scott Thornton committed
218
219
220
                            1)),
                    std::size_t(std::max<std::ptrdiff_t>(
                        1,
Paul's avatar
Paul committed
221
                        std::ptrdiff_t(std::floor((input.lens()[3] + 2 * padding[1] - lengths[1]) /
Paul's avatar
Paul committed
222
                                                  static_cast<float>(stride[1]))) +
Scott Thornton's avatar
Scott Thornton committed
223
224
                            1)),
                }};
Paul's avatar
Paul committed
225
226
227
    }
};

Khalique's avatar
Khalique committed
228
229
230
231
232
233
234
235
236
struct leaky_relu
{
    std::string name() const { return "leaky_relu"; }
    float alpha;
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs, *this}.has(1);
        return inputs.front();
    }
Khalique's avatar
Khalique committed
237
238
239

    template <class Self, class F>
    static auto reflect(Self& self, F f)
Khalique's avatar
Khalique committed
240
    {
Khalique's avatar
Khalique committed
241
        return pack(f(self.alpha, "alpha"));
Khalique's avatar
Khalique committed
242
243
244
245
246
247
248
249
250
251
252
253
    }
};

struct elu
{
    std::string name() const { return "elu"; }
    float alpha;
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs, *this}.has(1);
        return inputs.front();
    }
Khalique's avatar
Khalique committed
254
255
256

    template <class Self, class F>
    static auto reflect(Self& self, F f)
Khalique's avatar
Khalique committed
257
    {
Khalique's avatar
Khalique committed
258
        return pack(f(self.alpha, "alpha"));
Khalique's avatar
Khalique committed
259
    }
Khalique's avatar
Khalique committed
260
261
};

262
263
264
struct transpose
{
    std::vector<int64_t> dims;
Paul's avatar
Paul committed
265
266
267
268

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
269
        return pack(f(self.dims, "dims"));
Paul's avatar
Paul committed
270
271
    }

272
273
274
    std::string name() const { return "transpose"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
275
        check_shapes{inputs, *this}.has(1);
276
        auto input         = inputs.at(0);
277
        auto input_lens    = input.lens();
278
279
        auto input_strides = input.strides();
        auto t             = input.type();
Paul's avatar
Paul committed
280
281
        if(dims.size() != input_lens.size())
        {
Paul's avatar
Paul committed
282
            MIGRAPH_THROW("Permutation has wrong number of axes");
283
284
285
        }
        std::vector<int64_t> axes(dims.size());
        std::iota(axes.begin(), axes.end(), 0);
Paul's avatar
Paul committed
286
287
        if(!std::is_permutation(axes.begin(), axes.end(), dims.begin()))
        {
Paul's avatar
Paul committed
288
            MIGRAPH_THROW("Invalid permutation");
289
        }
290
291
        std::vector<size_t> output_lens(input_lens.size());
        std::vector<size_t> output_strides(input_lens.size());
Paul's avatar
Paul committed
292
293
294
        for(int i = 0; i < output_lens.size(); i++)
        {
            output_lens[i]    = input_lens[dims[i]];
295
296
            output_strides[i] = input_strides[dims[i]];
        }
297
        return {t, output_lens, output_strides};
298
    }
Paul's avatar
Paul committed
299
    argument compute(context&, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
300
    {
Paul's avatar
Paul committed
301
        return {std::move(output_shape), std::move(args.front().data)};
Paul's avatar
Paul committed
302
    }
Paul's avatar
Paul committed
303
    int output_alias(const std::vector<shape>&) const { return 0; }
304
305
};

Paul's avatar
Paul committed
306
struct contiguous
307
308
309
310
{
    std::string name() const { return "contiguous"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
311
        check_shapes{inputs, *this}.has(1);
Paul's avatar
Paul committed
312
313
        auto lens = inputs.at(0).lens();
        auto t    = inputs.at(0).type();
314
315
316
317
        return {t, lens};
    }
};

318
319
320
321
struct concat
{
    std::size_t axis = 0;
    std::string name() const { return "concat"; }
322
323
324
325
326
327
328
329
330
331
332
333
334
    std::vector<std::size_t> compute_offsets(const shape& output_shape,
                                             const std::vector<argument> args) const
    {
        std::vector<std::size_t> offsets;
        std::vector<std::size_t> offset(args[0].get_shape().lens().size(), 0);
        offset[axis] = 0;
        for(const auto& arg : args)
        {
            offsets.push_back(output_shape.index(offset));
            offset[axis] += arg.get_shape().lens()[axis];
        }
        return offsets;
    }
335
336
    shape compute_shape(std::vector<shape> inputs) const
    {
Scott Thornton's avatar
Scott Thornton committed
337
        if(inputs.empty())
338
339
340
341
342
        {
            MIGRAPH_THROW("Number of input tensors should exceed 0");
        }

        const auto& first_shape_lens = inputs.front().lens();
Scott Thornton's avatar
Scott Thornton committed
343
344
345
346
347
348
349
350
351
352
        const auto& type             = inputs.front().type();
        for(std::size_t l = 0; l < first_shape_lens.size(); l++)
        {
            if(l != axis)
            {
                if(!std::all_of(inputs.begin(), inputs.end(), [&](auto s) {
                       return s.lens()[l] == first_shape_lens[l];
                   }))
                {
                    MIGRAPH_THROW("Non-axis dimensions should match");
353
354
355
356
                }
            }
        }
        std::size_t new_dim_axis = 0;
Scott Thornton's avatar
Scott Thornton committed
357
        for(const auto& input : inputs)
358
359
360
361
362
363
364
365
366
        {
            const auto& lens = input.lens();
            new_dim_axis += lens[axis];
        }
        std::vector<std::size_t> new_lens;
        std::copy(first_shape_lens.begin(), first_shape_lens.end(), std::back_inserter(new_lens));
        new_lens[axis] = new_dim_axis;
        return {type, new_lens};
    }
Paul's avatar
Paul committed
367
    int output_alias(const std::vector<shape>&) const { return 0; }
368
369
};

370
371
372
373
374
struct slice
{
    std::vector<int64_t> axes;
    std::vector<int64_t> starts;
    std::vector<int64_t> ends;
Paul's avatar
Paul committed
375
376
377
378

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
379
        return pack(f(self.axes, "axes"), f(self.starts, "starts"), f(self.ends, "ends"));
Paul's avatar
Paul committed
380
381
    }

382
    std::string name() const { return "slice"; }
Scott Thornton's avatar
Scott Thornton committed
383
384

    auto fix_index(const std::vector<std::size_t>& lens, std::size_t axis, int64_t index) const
385
    {
Scott Thornton's avatar
Scott Thornton committed
386
        int64_t r = std::min(index, static_cast<int64_t>(lens[axis]));
Scott Thornton's avatar
Scott Thornton committed
387
388
        if(r < 0)
            r += lens[axis];
Scott Thornton's avatar
Scott Thornton committed
389
        return std::size_t(r);
Scott Thornton's avatar
Scott Thornton committed
390
391
392
393
394
395
396
    }

    auto compute_offset(const shape& s) const
    {
        const std::vector<std::size_t>& lens    = s.lens();
        const std::vector<std::size_t>& strides = s.strides();
        auto offset                             = 0;
Scott Thornton's avatar
Scott Thornton committed
397
        if(!axes.empty())
Scott Thornton's avatar
Scott Thornton committed
398
        {
Scott Thornton's avatar
Scott Thornton committed
399
400
401
402
403
            for(std::size_t i = 0; i < axes.size(); i++)
            {
                auto axis = axes[i];
                offset += fix_index(lens, axis, starts[i]) * strides[axis];
            }
404
        }
Scott Thornton's avatar
Scott Thornton committed
405
406
        else
        {
Scott Thornton's avatar
Scott Thornton committed
407
408
409
410
            for(std::size_t axis = 0; axis < lens.size(); axis++)
            {
                offset += fix_index(lens, axis, starts[axis]) * strides[axis];
            }
411
        }
Scott Thornton's avatar
Scott Thornton committed
412
413
414
415
416
        return offset;
    }

    shape compute_shape(std::vector<shape> inputs) const
    {
Scott Thornton's avatar
Scott Thornton committed
417
418
419
420
        auto input_shape        = inputs[0];
        auto t                  = input_shape.type();
        const auto& old_lens    = input_shape.lens();
        const auto& old_strides = input_shape.strides();
Scott Thornton's avatar
Scott Thornton committed
421
422
423
424
425
426
427
428
429
430
        // std::vector<int64_t> t_axes(old_lens.size());
        // if(axes.size() == 0)
        // {
        //     std::iota(t_axes.begin(), t_axes.end(), 0);
        // }
        // else
        // {
        //     std::copy(axes.begin(), axes.end(), t_axes.begin());
        // }
        if(starts.size() != axes.size() || axes.size() != ends.size())
Scott Thornton's avatar
Scott Thornton committed
431
        {
432
433
            MIGRAPH_THROW("inconsistent sizes");
        }
Scott Thornton's avatar
Scott Thornton committed
434
435
        std::vector<std::size_t> new_lens = old_lens;
        for(std::size_t i = 0; i < axes.size(); i++)
Scott Thornton's avatar
Scott Thornton committed
436
        {
Scott Thornton's avatar
Scott Thornton committed
437
438
439
            auto axis = axes[i];
            new_lens[axis] =
                fix_index(old_lens, axis, ends[i]) - fix_index(old_lens, axis, starts[i]);
440
441
442
443
444
        }
        return shape{t, new_lens, old_strides};
    }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
Scott Thornton's avatar
Scott Thornton committed
445
446
447
        auto input  = args[0];
        auto offset = compute_offset(input.get_shape()) * output_shape.type_size();
        return {std::move(output_shape), [=] { return input.data() + offset; }};
448
    }
Paul's avatar
Paul committed
449
    int output_alias(const std::vector<shape>&) const { return 0; }
450
451
452
453
454
};

struct squeeze
{
    std::vector<int64_t> axes;
Paul's avatar
Paul committed
455
456
457
458

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
459
        return pack(f(self.axes, "axes"));
Paul's avatar
Paul committed
460
461
    }

462
463
464
465
    std::string name() const { return "squeeze"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        auto input_shape = inputs[0];
Scott Thornton's avatar
Scott Thornton committed
466
467
        auto type        = input_shape.type();
        auto old_lens    = input_shape.lens();
wsttiger's avatar
wsttiger committed
468
469
        if(std::any_of(
               axes.begin(), axes.end(), [&](auto axis) { return input_shape.lens()[axis] != 1; }))
Scott Thornton's avatar
Scott Thornton committed
470
        {
wsttiger's avatar
wsttiger committed
471
            MIGRAPH_THROW("squeeze axis dimension should be equal to 1");
472
473
        }
        std::vector<std::size_t> new_lens;
Scott Thornton's avatar
Scott Thornton committed
474
        if(axes.empty())
Scott Thornton's avatar
Scott Thornton committed
475
        {
wsttiger's avatar
wsttiger committed
476
477
478
479
            std::copy_if(old_lens.begin(),
                         old_lens.end(),
                         std::back_inserter(new_lens),
                         [](auto len) { return len != 1; });
480
        }
Scott Thornton's avatar
Scott Thornton committed
481
482
483
484
485
486
        else
        {
            for(std::size_t i = 0; i < old_lens.size(); i++)
            {
                if(std::find(axes.begin(), axes.end(), i) == axes.end())
                {
487
488
489
490
491
492
493
494
495
                    new_lens.push_back(old_lens[i]);
                }
            }
        }
        return shape{type, new_lens};
    }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        return {std::move(output_shape), std::move(args.front().data)};
Scott Thornton's avatar
Scott Thornton committed
496
    }
Paul's avatar
Paul committed
497
    int output_alias(const std::vector<shape>&) const { return 0; }
498
499
500
501
502
};

struct unsqueeze
{
    std::vector<int64_t> axes;
Paul's avatar
Paul committed
503
504
505
506

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
507
        return pack(f(self.axes, "axes"));
Paul's avatar
Paul committed
508
509
    }

510
511
512
    std::string name() const { return "unsqueeze"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Scott Thornton's avatar
Scott Thornton committed
513
514
515
        auto input_shape     = inputs[0];
        auto type            = input_shape.type();
        auto old_lens        = input_shape.lens();
516
517
518
        std::size_t new_size = old_lens.size() + axes.size();
        std::vector<std::size_t> new_lens(new_size);
        std::size_t p = 0;
Scott Thornton's avatar
Scott Thornton committed
519
520
521
522
        for(std::size_t i = 0; i < new_size; i++)
        {
            if(std::find(axes.begin(), axes.end(), i) != axes.end())
            {
523
                new_lens[i] = 1;
Scott Thornton's avatar
Scott Thornton committed
524
525
526
            }
            else
            {
527
528
529
530
531
532
533
534
535
                new_lens[i] = old_lens[p++];
            }
        }
        return shape{type, new_lens};
    }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        return {std::move(output_shape), std::move(args.front().data)};
    }
Paul's avatar
Paul committed
536
    int output_alias(const std::vector<shape>&) const { return 0; }
537
538
};

Paul's avatar
Paul committed
539
540
541
struct reshape
{
    std::vector<int64_t> dims;
Paul's avatar
Paul committed
542
543
544
545

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
546
        return pack(f(self.dims, "dims"));
Paul's avatar
Paul committed
547
548
    }

Paul's avatar
Paul committed
549
    std::string name() const { return "reshape"; }
Paul's avatar
Paul committed
550
551
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
552
        check_shapes{inputs, *this}.has(1);
Paul's avatar
Paul committed
553
554
        auto&& idims = inputs.front().lens();
        std::vector<std::size_t> rdims(dims.begin(), dims.end());
555
556
557
        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
558
        for(std::size_t i = 0; i < dims.size(); i++)
Paul's avatar
Paul committed
559
560
561
562
        {
            if(dims[i] == 0)
                rdims[i] = idims[i];
        }
563
564
565
566
567
568
569
570
571
572
573
        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
574
575
576
        if(dims.back() == -1)
        {
            rdims.pop_back();
Paul's avatar
Paul committed
577
            std::copy(idims.begin() + rdims.size(), idims.end(), std::back_inserter(rdims));
Paul's avatar
Paul committed
578
        }
Scott Thornton's avatar
Scott Thornton committed
579
        shape s{inputs.front().type(), rdims};
Paul's avatar
Paul committed
580
        if(s.elements() != inputs.front().elements())
Paul's avatar
Paul committed
581
            MIGRAPH_THROW("Wrong number of elements for reshape");
Scott Thornton's avatar
Scott Thornton committed
582
        return s;
Paul's avatar
Paul committed
583
    }
Paul's avatar
Paul committed
584
    argument compute(context&, shape output_shape, std::vector<argument> args) const
Paul's avatar
Paul committed
585
    {
Paul's avatar
Paul committed
586
        return {std::move(output_shape), std::move(args.front().data)};
Paul's avatar
Paul committed
587
    }
Paul's avatar
Paul committed
588
    int output_alias(const std::vector<shape>&) const { return 0; }
Paul's avatar
Paul committed
589
590
};

Shucai Xiao's avatar
Shucai Xiao committed
591
struct dot
592
{
Paul's avatar
Paul committed
593
    float alpha = 1.0;
Paul's avatar
Paul committed
594
    float beta  = 0.0;
Paul's avatar
Paul committed
595
596
597
598

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
599
        return pack(f(self.alpha, "alpha"), f(self.beta, "beta"));
Paul's avatar
Paul committed
600
601
    }

Shucai Xiao's avatar
Shucai Xiao committed
602
    std::string name() const { return "dot"; }
603
604
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
605
        check_shapes{inputs, *this}.has(2).same_type();
606
607
        const shape& a = inputs.at(0);
        const shape& b = inputs.at(1);
Scott Thornton's avatar
Scott Thornton committed
608
        auto t         = a.type();
609

610
        if(a.lens()[1] != b.lens()[0])
Paul's avatar
Paul committed
611
612
            MIGRAPH_THROW("Inner dimensions do not match: {" + to_string_range(a.lens()) + "} x {" +
                          to_string_range(b.lens()) + "}");
Scott Thornton's avatar
Scott Thornton committed
613
        return {t, {a.lens()[0], b.lens()[1]}};
614
615
616
    }
};

617
struct unary
Scott Thornton's avatar
Scott Thornton committed
618
{
619
620
    shape compute_shape(std::vector<shape> inputs) const
    {
621
622
        check_shapes{inputs}.has(1);
        return inputs.at(0);
623
    }
Scott Thornton's avatar
Scott Thornton committed
624
625
};

626
struct identity
627
{
628
    std::string name() const { return "identity"; }
Scott Thornton's avatar
Scott Thornton committed
629
    shape compute_shape(std::vector<shape> inputs) const { return inputs.at(0); }
630
631
632
633
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        return {std::move(output_shape), std::move(args.at(0).data)};
    }
Paul's avatar
Paul committed
634
    int output_alias(const std::vector<shape>&) const { return 0; }
635
636
637
};

struct abs : unary
Scott Thornton's avatar
Scott Thornton committed
638
{
639
    std::string name() const { return "abs"; }
Scott Thornton's avatar
Scott Thornton committed
640
641
};

642
struct exp : unary
Scott Thornton's avatar
Scott Thornton committed
643
{
644
    std::string name() const { return "exp"; }
Scott Thornton's avatar
Scott Thornton committed
645
646
};

647
struct sin : unary
Scott Thornton's avatar
Scott Thornton committed
648
{
649
    std::string name() const { return "sin"; }
Scott Thornton's avatar
Scott Thornton committed
650
651
};

652
struct cos : unary
Scott Thornton's avatar
Scott Thornton committed
653
{
654
    std::string name() const { return "cos"; }
Scott Thornton's avatar
Scott Thornton committed
655
656
};

657
struct tan : unary
Scott Thornton's avatar
Scott Thornton committed
658
{
659
    std::string name() const { return "tan"; }
Scott Thornton's avatar
Scott Thornton committed
660
661
};

662
struct asin : unary
Scott Thornton's avatar
Scott Thornton committed
663
{
664
    std::string name() const { return "asin"; }
Scott Thornton's avatar
Scott Thornton committed
665
666
};

667
struct acos : unary
Scott Thornton's avatar
Scott Thornton committed
668
{
669
    std::string name() const { return "acos"; }
Scott Thornton's avatar
Scott Thornton committed
670
671
};

672
struct atan : unary
Scott Thornton's avatar
Scott Thornton committed
673
{
674
    std::string name() const { return "atan"; }
Scott Thornton's avatar
Scott Thornton committed
675
676
};

677
678
679
680
681
682
683
684
685
686
struct sinh : unary
{
    std::string name() const { return "sinh"; }
};

struct cosh : unary
{
    std::string name() const { return "cosh"; }
};

687
struct tanh : unary
Scott Thornton's avatar
Scott Thornton committed
688
{
689
    std::string name() const { return "tanh"; }
Scott Thornton's avatar
Scott Thornton committed
690
691
};

692
struct sigmoid : unary
Scott Thornton's avatar
Scott Thornton committed
693
{
694
    std::string name() const { return "sigmoid"; }
Scott Thornton's avatar
Scott Thornton committed
695
696
};

697
struct neg : unary
Scott Thornton's avatar
Scott Thornton committed
698
{
699
    std::string name() const { return "neg"; }
Scott Thornton's avatar
Scott Thornton committed
700
701
};

Khalique's avatar
Khalique committed
702
703
704
705
706
struct relu : unary
{
    std::string name() const { return "relu"; }
};

Paul's avatar
Paul committed
707
708
709
710
711
712
713
714
715
716
struct softmax
{
    std::string name() const { return "softmax"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(1).only_dims(4);
        return inputs.at(0);
    }
};

717
struct flatten
Scott Thornton's avatar
Scott Thornton committed
718
{
Paul's avatar
Paul committed
719
    uint64_t axis = 0;
Paul's avatar
Paul committed
720
721
722
723

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
724
        return pack(f(self.axis, "axis"));
Paul's avatar
Paul committed
725
726
    }

Scott Thornton's avatar
Scott Thornton committed
727
    std::string name() const { return "flatten"; }
Paul's avatar
Paul committed
728
729
730
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(1);
Paul's avatar
Paul committed
731
732
        auto&& lens = inputs.front().lens();

Paul's avatar
Paul committed
733
        if(axis > lens.size())
Paul's avatar
Paul committed
734
        {
Paul's avatar
Paul committed
735
            MIGRAPH_THROW("axis for flatten must be less than tensor rank");
Paul's avatar
Paul committed
736
        }
Paul's avatar
Paul committed
737
738
739
740
        auto x =
            std::accumulate(lens.begin(), lens.begin() + axis, std::size_t{1}, std::multiplies<>{});
        auto y =
            std::accumulate(lens.begin() + axis, lens.end(), std::size_t{1}, std::multiplies<>{});
741
        return {inputs.at(0).type(), {x, y}};
Paul's avatar
Paul committed
742
743
744
    }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
Paul's avatar
Paul committed
745
        return {std::move(output_shape), std::move(args.front().data)};
Paul's avatar
Paul committed
746
    }
Paul's avatar
Paul committed
747
    int output_alias(const std::vector<shape>&) const { return 0; }
Scott Thornton's avatar
Scott Thornton committed
748
};
749
750
751
struct broadcast
{
    uint64_t axis = 0;
Paul's avatar
Paul committed
752
753
754
755

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
756
        return pack(f(self.axis, "axis"));
Paul's avatar
Paul committed
757
758
    }

Scott Thornton's avatar
Scott Thornton committed
759
    shape broadcast_shape;
760
761
762
    std::string name() const { return "broadcast"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
Scott Thornton's avatar
Scott Thornton committed
763
764
        auto t     = inputs.at(0).type();
        auto input = inputs.at(0);
Paul's avatar
Paul committed
765

Scott Thornton's avatar
Scott Thornton committed
766
        std::vector<size_t> bcast_strides(broadcast_shape.lens().size(), 0);
767

Scott Thornton's avatar
Scott Thornton committed
768
769
770
        if(std::all_of(broadcast_shape.lens().cbegin(), broadcast_shape.lens().cend(), [&](auto x) {
               return x == 1;
           }))
771
        {
Scott Thornton's avatar
Scott Thornton committed
772
            if(axis != 0)
Paul's avatar
Paul committed
773
                MIGRAPH_THROW("when broadcasting tensor of size 1, axis should be 0");
Scott Thornton's avatar
Scott Thornton committed
774
            return {t, broadcast_shape.lens(), std::move(bcast_strides)};
775
776
777
        }
        else
        {
Scott Thornton's avatar
Scott Thornton committed
778
            assert(broadcast_shape.lens().size() - axis >= input.lens().size());
Scott Thornton's avatar
Scott Thornton committed
779
780
            if(!std::equal(
                   input.lens().begin(), input.lens().end(), broadcast_shape.lens().begin() + axis))
Paul's avatar
Paul committed
781
                MIGRAPH_THROW("when broadcasting success sizes must match");
Paul's avatar
Paul committed
782
            std::copy(input.strides().begin(), input.strides().end(), bcast_strides.begin() + axis);
Scott Thornton's avatar
Scott Thornton committed
783
            return {t, broadcast_shape.lens(), std::move(bcast_strides)};
784
785
        }
    }
Paul's avatar
Paul committed
786
    argument compute(context&, shape output_shape, std::vector<argument> args) const
Scott Thornton's avatar
Scott Thornton committed
787
    {
Scott Thornton's avatar
Scott Thornton committed
788
        return {std::move(output_shape), std::move(args.at(0).data)};
Scott Thornton's avatar
Scott Thornton committed
789
    }
Paul's avatar
Paul committed
790
    int output_alias(const std::vector<shape>&) const { return 0; }
791
792
};

Scott Thornton's avatar
Scott Thornton committed
793
794
795
struct multibroadcast
{
    std::vector<std::size_t> output_lens;
796
797
798
799
800
801
802

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return pack(f(self.output_lens, "output_lens"));
    }

Scott Thornton's avatar
Scott Thornton committed
803
    std::string name() const { return "multibroadcast"; }
804

Scott Thornton's avatar
Scott Thornton committed
805
806
807
808
809
810
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs, *this}.has(1);
        auto t     = inputs.at(0).type();
        auto input = inputs.at(0);

wsttiger's avatar
wsttiger committed
811
        if(input.lens().empty())
Scott Thornton's avatar
Scott Thornton committed
812
813
            MIGRAPH_THROW("inputs dimensions should be > 0");

Scott Thornton's avatar
Scott Thornton committed
814
        if(input.lens().size() > output_lens.size())
Scott Thornton's avatar
Scott Thornton committed
815
816
817
            MIGRAPH_THROW("inputs dimensions should <= output size");

        std::vector<size_t> bcast_strides(output_lens.size(), 0);
Scott Thornton's avatar
Scott Thornton committed
818
819
        auto offset = output_lens.size() - input.lens().size();
        for(int i = input.lens().size() - 1; i >= 0; i--)
Scott Thornton's avatar
Scott Thornton committed
820
        {
Scott Thornton's avatar
Scott Thornton committed
821
            if(output_lens[i + offset] == input.lens()[i])
Scott Thornton's avatar
Scott Thornton committed
822
            {
Scott Thornton's avatar
Scott Thornton committed
823
                bcast_strides[i + offset] = input.strides()[i];
Scott Thornton's avatar
Scott Thornton committed
824
825
826
827
828
829
830
831
832
833
834
            }
        }
        return {t, output_lens, bcast_strides};
    }
    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        return {std::move(output_shape), std::move(args.at(0).data)};
    }
    int output_alias(const std::vector<shape>&) const { return 0; }
};

Khalique's avatar
Khalique committed
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
struct scalar
{
    shape scalar_bcast;

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

    shape compute_shape(std::vector<shape> inputs) const
    {
        assert(check_shapes{inputs}.has(1).only_dims(1).size() == 1);
        auto t     = inputs.at(0).type();
        auto input = inputs.at(0);
        std::vector<std::size_t> strides(scalar_bcast.lens().size(), 0);
        return {t, scalar_bcast.lens(), strides};
    }

    argument compute(context&, shape output_shape, std::vector<argument> args) const
    {
        return {std::move(output_shape), std::move(args.at(0).data)};
    }
Paul's avatar
Paul committed
854
    int output_alias(const std::vector<shape>&) const { return 0; }
Khalique's avatar
Khalique committed
855
856
};

857
struct binary
Scott Thornton's avatar
Scott Thornton committed
858
{
859
860
    shape compute_shape(std::vector<shape> inputs) const
    {
861
        check_shapes{inputs}.has(2).same_type().same_dims();
Scott Thornton's avatar
Scott Thornton committed
862
        auto t    = inputs.at(0).type();
863
864
        auto lens = inputs.at(0).lens();
        return {t, lens};
865
    }
Scott Thornton's avatar
Scott Thornton committed
866
867
};

868
869
870
871
872
873
struct add : binary
{
    std::string name() const { return "add"; }
};

struct sub : binary
Scott Thornton's avatar
Scott Thornton committed
874
875
876
877
{
    std::string name() const { return "sub"; }
};

878
struct mul : binary
Scott Thornton's avatar
Scott Thornton committed
879
880
881
882
{
    std::string name() const { return "mul"; }
};

883
struct div : binary
Scott Thornton's avatar
Scott Thornton committed
884
885
886
887
{
    std::string name() const { return "div"; }
};

Paul's avatar
Paul committed
888
889
890
891
struct load
{
    shape s;
    std::size_t offset = 0;
Paul's avatar
Paul committed
892
893
894
895

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
896
        return pack(f(self.s, "shape"), f(self.offset, "offset"));
Paul's avatar
Paul committed
897
898
    }

Paul's avatar
Paul committed
899
900
901
902
903
904
905
906
907
908
    std::string name() const { return "load"; }
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        check_shapes{inputs}.has(1);
        return s;
    }
    argument compute(context&, const shape&, const std::vector<argument>& args) const
    {
        return {s, args[0].data() + offset};
    }
Paul's avatar
Paul committed
909
    int output_alias(const std::vector<shape>&) const { return 0; }
Paul's avatar
Paul committed
910
911
};

Paul's avatar
Paul committed
912
struct outline
Scott Thornton's avatar
Scott Thornton committed
913
{
Paul's avatar
Paul committed
914
    shape s;
Paul's avatar
Paul committed
915
916
917
918

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
Paul's avatar
Paul committed
919
        return pack(f(self.s, "shape"));
Paul's avatar
Paul committed
920
921
    }

Paul's avatar
Paul committed
922
    std::string name() const { return "outline"; }
Paul's avatar
Paul committed
923
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
924
    {
Paul's avatar
Paul committed
925
        check_shapes{inputs, *this}.has(0);
Paul's avatar
Paul committed
926
927
        return s;
    }
Paul's avatar
Paul committed
928
929
930
931
    argument compute(context&, const shape&, const std::vector<argument>&) const
    {
        return {s, nullptr};
    }
Scott Thornton's avatar
Scott Thornton committed
932
933
};

934
} // namespace op
935
} // namespace MIGRAPH_INLINE_NS
Paul's avatar
Paul committed
936
} // namespace migraphx
Paul's avatar
Paul committed
937
938

#endif