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

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

Paul's avatar
Paul committed
10
11
namespace rtg {

Paul's avatar
Paul committed
12
13
14
15
struct check_shapes
{
    const std::vector<shape>* shapes;

Paul's avatar
Paul committed
16
    check_shapes(const std::vector<shape>& s) : shapes(&s) {}
Paul's avatar
Paul committed
17
18
19
20
21

    const check_shapes& has(std::size_t n) const
    {
        assert(shapes != nullptr);
        if(shapes->size() != n)
Paul's avatar
Paul committed
22
23
            RTG_THROW("Wrong number of arguments: expected " + std::to_string(n) + " but given " +
                      std::to_string(shapes->size()));
Paul's avatar
Paul committed
24
25
26
27
28
29
        return *this;
    }

    const check_shapes& only_dims(std::size_t n) const
    {
        assert(shapes != nullptr);
Paul's avatar
Paul committed
30
31
        if(!shapes->empty())
        {
Paul's avatar
Paul committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
            if(shapes->front().lens().size() != n)
                RTG_THROW("Only " + std::to_string(n) + "d supported");
        }
        return *this;
    }

    const check_shapes& same_shape() const
    {
        if(!this->same([](const shape& s) { return s; }))
            RTG_THROW("Shapes do not match");
        return *this;
    }

    const check_shapes& same_type() const
    {
        if(!this->same([](const shape& s) { return s.type(); }))
            RTG_THROW("Types do not match");
        return *this;
    }

    const check_shapes& same_dims() const
    {
        if(!this->same([](const shape& s) { return s.lens(); }))
            RTG_THROW("Dimensions do not match");
        return *this;
    }

59
60
61
62
63
64
65
    const check_shapes& same_ndims() const
    {
        if(!this->same([](const shape& s) { return s.lens().size(); }))
            RTG_THROW("Dimensions do not match");
        return *this;
    }

Paul's avatar
Paul committed
66
    template <class F>
Paul's avatar
Paul committed
67
68
69
70
71
72
    bool same(F f) const
    {
        assert(shapes != nullptr);
        if(shapes->empty())
            return true;
        auto&& key = f(shapes->front());
Paul's avatar
Paul committed
73
        return this->all_of([&](const shape& s) { return f(s) == key; });
Paul's avatar
Paul committed
74
75
    }

Paul's avatar
Paul committed
76
    template <class Predicate>
Paul's avatar
Paul committed
77
78
79
80
81
82
83
    bool all_of(Predicate p) const
    {
        assert(shapes != nullptr);
        return std::all_of(shapes->begin(), shapes->end(), p);
    }
};

Paul's avatar
Paul committed
84
85
struct not_computable
{
Paul's avatar
Paul committed
86
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
87
88
};

Paul's avatar
Paul committed
89
struct convolution
Paul's avatar
Paul committed
90
{
Paul's avatar
Paul committed
91
92
93
    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
94
    std::string name() const { return "convolution"; }
Paul's avatar
Paul committed
95
96
    shape compute_shape(std::vector<shape> inputs) const
    {
97
        check_shapes{inputs}.has(2).same_type().same_ndims().only_dims(4);
Paul's avatar
Paul committed
98

Paul's avatar
Paul committed
99
        const shape& input   = inputs.at(0);
Paul's avatar
Paul committed
100
        const shape& weights = inputs.at(1);
Paul's avatar
Paul committed
101
        auto t               = input.type();
Paul's avatar
Paul committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        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
119
    }
Paul's avatar
Paul committed
120

Paul's avatar
Paul committed
121
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
122

Paul's avatar
Paul committed
123
    friend std::ostream& operator<<(std::ostream& os, const convolution& op)
Paul's avatar
Paul committed
124
    {
Paul's avatar
Paul committed
125
126
127
128
129
        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
130
131
        return os;
    }
Paul's avatar
Paul committed
132
133
};

Paul's avatar
Paul committed
134
struct pooling
Paul's avatar
Paul committed
135
136
{
    std::string mode;
Paul's avatar
Paul committed
137
138
139
    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
140
    std::string name() const { return "pooling"; }
Paul's avatar
Paul committed
141
142
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
143
        check_shapes{inputs}.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
148
149
150
151
152
153
154
155
156
157
158
159
160
161
        return {t,
                {
                    input.lens()[0],
                    input.lens()[1],
                    std::size_t(std::max<std::ptrdiff_t>(
                        1,
                        std::ceil((input.lens()[3] + 2 * padding[0] - lengths[0]) /
                                  static_cast<float>(stride[0])) +
                            1)),
                    std::size_t(std::max<std::ptrdiff_t>(
                        1,
                        std::ceil((input.lens()[4] + 2 * padding[1] - lengths[1]) /
                                  static_cast<float>(stride[1])) +
                            1)),
                }};
Paul's avatar
Paul committed
162
    }
Paul's avatar
Paul committed
163

Paul's avatar
Paul committed
164
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
165

Paul's avatar
Paul committed
166
    friend std::ostream& operator<<(std::ostream& os, const pooling& op)
Paul's avatar
Paul committed
167
    {
Paul's avatar
Paul committed
168
169
170
171
172
        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
173
174
        return os;
    }
Paul's avatar
Paul committed
175
176
};

Paul's avatar
Paul committed
177
struct activation
Paul's avatar
Paul committed
178
179
{
    std::string mode;
Paul's avatar
Paul committed
180
    std::string name() const { return "activation"; }
Paul's avatar
Paul committed
181
182
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
183
        check_shapes{inputs}.has(1);
Paul's avatar
Paul committed
184
185
        return inputs.front();
    }
Paul's avatar
Paul committed
186

Paul's avatar
Paul committed
187
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
188
    friend std::ostream& operator<<(std::ostream& os, const activation& op)
Paul's avatar
Paul committed
189
    {
Paul's avatar
Paul committed
190
        os << op.name() << ":" << op.mode;
Paul's avatar
Paul committed
191
192
        return os;
    }
Paul's avatar
Paul committed
193
194
};

Paul's avatar
Paul committed
195
196
197
struct reshape
{
    std::vector<int64_t> dims;
Paul's avatar
Paul committed
198
    std::string name() const { return "reshape"; }
Paul's avatar
Paul committed
199
200
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
201
        if(inputs.empty())
Paul's avatar
Paul committed
202
            RTG_THROW("Wrong number of arguments");
Paul's avatar
Paul committed
203
204
        auto&& idims = inputs.front().lens();
        std::vector<std::size_t> rdims(dims.begin(), dims.end());
Paul's avatar
Paul committed
205
        for(std::size_t i = 0; i < dims.size(); i++)
Paul's avatar
Paul committed
206
207
208
209
210
211
212
        {
            if(dims[i] == 0)
                rdims[i] = idims[i];
        }
        if(dims.back() == -1)
        {
            rdims.pop_back();
Paul's avatar
Paul committed
213
            std::copy(idims.begin() + rdims.size(), idims.end(), std::back_inserter(rdims));
Paul's avatar
Paul committed
214
215
216
217
        }
        return {inputs.front().type(), rdims};
    }

Paul's avatar
Paul committed
218
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
219

Paul's avatar
Paul committed
220
    friend std::ostream& operator<<(std::ostream& os, const reshape& op)
Paul's avatar
Paul committed
221
    {
Paul's avatar
Paul committed
222
223
224
        os << op.name() << "[";
        os << "dims={" << stream_range(op.dims) << "}, ";
        os << "]";
Paul's avatar
Paul committed
225
226
        return os;
    }
Paul's avatar
Paul committed
227
228
};

229
230
231
232
233
234
235
236
struct gemm
{
    std::string name() const { return "gemm";}
    std::size_t lda = 1; 
    std::size_t ldb = 1; 
    std::size_t ldc = 1; 
    shape compute_shape(std::vector<shape> inputs) const
    {
237
        check_shapes{inputs}.has(2).same_type();
Scott Thornton's avatar
Scott Thornton committed
238
239
240
        const shape& a = inputs.at(0); 
        const shape& b = inputs.at(1); 
        auto t         = a.type();
241

Scott Thornton's avatar
Scott Thornton committed
242
        if (a.lens()[1] != b.lens()[0])
243
            RTG_THROW("Inner dimensions do not match");
Scott Thornton's avatar
Scott Thornton committed
244
        return {t, {a.lens()[0], b.lens()[1]}};
245
246
247
248
249
250
251
252
    }
  
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
  
    friend std::ostream& operator<<(std::ostream& os, const gemm& op) 
    {
        os << op.name() << "[";
        os << "]"; 
Scott Thornton's avatar
Scott Thornton committed
253
        return os;
254
255
256
    }
};

257
struct identity
Scott Thornton's avatar
Scott Thornton committed
258
{
Scott Thornton's avatar
Scott Thornton committed
259
    std::string name() const {return "identity"; }
260
261
262
263
264
265
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
266
267
};

268
struct abs
Scott Thornton's avatar
Scott Thornton committed
269
270
{
    std::string name() const {return "abs"; }
271
272
273
274
275
276
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
277
278
};

279
struct exp 
Scott Thornton's avatar
Scott Thornton committed
280
{
281
282
283
284
285
286
287
    std::string name() const { return "exp"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
288
289
};

290
struct sin
Scott Thornton's avatar
Scott Thornton committed
291
292
{
    std::string name() const {return "sin"; }
293
294
295
296
297
298
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
299
300
};

301
struct cos
Scott Thornton's avatar
Scott Thornton committed
302
303
{
    std::string name() const {return "cos"; }
304
305
306
307
308
309
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
310
311
};

312
struct tan
Scott Thornton's avatar
Scott Thornton committed
313
314
{
    std::string name() const {return "tan"; }
315
316
317
318
319
320
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
321
322
};

323
struct asin
Scott Thornton's avatar
Scott Thornton committed
324
325
{
    std::string name() const {return "asin"; }
326
327
328
329
330
331
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
332
333
};

334
struct acos
Scott Thornton's avatar
Scott Thornton committed
335
336
{
    std::string name() const {return "acos"; }
337
338
339
340
341
342
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
343
344
};

345
struct atan
Scott Thornton's avatar
Scott Thornton committed
346
347
{
    std::string name() const {return "atan"; }
348
349
350
351
352
353
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
354
355
};

356
struct softmax
Scott Thornton's avatar
Scott Thornton committed
357
358
{
    std::string name() const {return "softmax"; }
359
360
    shape compute_shape(std::vector<shape> inputs) const
    {
361
      check_shapes{inputs}.has(1).only_dims(4);
362
363
364
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
365
366
};

367
struct tanh
Scott Thornton's avatar
Scott Thornton committed
368
369
{
    std::string name() const {return "tanh"; }
370
371
372
373
374
375
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
376
377
};

378
struct sigmoid
Scott Thornton's avatar
Scott Thornton committed
379
380
{
    std::string name() const {return "sigmoid"; }
381
382
383
384
385
386
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
387
388
};

389
struct neg
Scott Thornton's avatar
Scott Thornton committed
390
{
Scott Thornton's avatar
Scott Thornton committed
391
    std::string name() const {return "neg"; }
Scott Thornton's avatar
Scott Thornton committed
392
393
394
395
396
    shape compute_shape(std::vector<shape> inputs) const
    {
      check_shapes{inputs}.has(1);
      return inputs.at(0);
    }
397
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Scott Thornton's avatar
Scott Thornton committed
398
399
400
401
402
403
404
};

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

405
struct add
Scott Thornton's avatar
Scott Thornton committed
406
407
{
    std::string name() const { return "add"; }
408
409
410
411
412
413
    shape compute_shape(std::vector<shape> inputs) const
    {
      // TODO(wsttiger@gmail.com) Check this for numpy-style broadcasting operations
      check_shapes{inputs}.has(2).same_type().same_dims();
      return inputs.at(0);
    }
Scott Thornton's avatar
Scott Thornton committed
414
415
};

416
struct sub
Scott Thornton's avatar
Scott Thornton committed
417
418
{
    std::string name() const { return "sub"; }
419
420
421
422
423
424
    shape compute_shape(std::vector<shape> inputs) const
    {
      // TODO(wsttiger@gmail.com) Check this for numpy-style broadcasting operations
      check_shapes{inputs}.has(2).same_type().same_dims();
      return inputs.at(0);
    }
Scott Thornton's avatar
Scott Thornton committed
425
426
};

427
struct mul
Scott Thornton's avatar
Scott Thornton committed
428
429
{
    std::string name() const { return "mul"; }
430
431
432
433
434
435
    shape compute_shape(std::vector<shape> inputs) const
    {
      // TODO(wsttiger@gmail.com) Check this for numpy-style broadcasting operations
      check_shapes{inputs}.has(2).same_type().same_dims();
      return inputs.at(0);
    }
Scott Thornton's avatar
Scott Thornton committed
436
437
};

438
struct div
Scott Thornton's avatar
Scott Thornton committed
439
440
441
442
443
444
445
446
447
448
{
    std::string name() const { return "div"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
      // TODO(wsttiger@gmail.com) Check this for numpy-style broadcasting operations
      check_shapes{inputs}.has(2).same_type().same_dims();
      return inputs.at(0);
    }
};

Paul's avatar
Paul committed
449
struct outline
Scott Thornton's avatar
Scott Thornton committed
450
{
Paul's avatar
Paul committed
451
452
453
454
455
456
457
    shape s;
    std::string name() const { return "outline"; }
    shape compute_shape(std::vector<shape> inputs) const
    {
        check_shapes{inputs}.has(0);
        return s;
    }
Paul's avatar
Paul committed
458
    argument compute(shape, std::vector<argument>) const { return {s, nullptr}; }
Scott Thornton's avatar
Scott Thornton committed
459
460
};

Paul's avatar
Paul committed
461
462
463
} // namespace rtg

#endif