operators.hpp 7.43 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;
    }

Paul's avatar
Paul committed
59
    template <class F>
Paul's avatar
Paul committed
60
61
62
63
64
65
    bool same(F f) const
    {
        assert(shapes != nullptr);
        if(shapes->empty())
            return true;
        auto&& key = f(shapes->front());
Paul's avatar
Paul committed
66
        return this->all_of([&](const shape& s) { return f(s) == key; });
Paul's avatar
Paul committed
67
68
    }

Paul's avatar
Paul committed
69
    template <class Predicate>
Paul's avatar
Paul committed
70
71
72
73
74
75
76
    bool all_of(Predicate p) const
    {
        assert(shapes != nullptr);
        return std::all_of(shapes->begin(), shapes->end(), p);
    }
};

Paul's avatar
Paul committed
77
78
struct not_computable
{
Paul's avatar
Paul committed
79
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
80
81
};

Paul's avatar
Paul committed
82
struct convolution
Paul's avatar
Paul committed
83
{
Paul's avatar
Paul committed
84
85
86
    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
87
    std::string name() const { return "convolution"; }
Paul's avatar
Paul committed
88
89
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
90
91
        check_shapes{inputs}.has(2).same_type().same_dims().only_dims(4);

Paul's avatar
Paul committed
92
        const shape& input   = inputs.at(0);
Paul's avatar
Paul committed
93
        const shape& weights = inputs.at(1);
Paul's avatar
Paul committed
94
        auto t               = input.type();
Paul's avatar
Paul committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
        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
112
    }
Paul's avatar
Paul committed
113

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

Paul's avatar
Paul committed
116
    friend std::ostream& operator<<(std::ostream& os, const convolution& op)
Paul's avatar
Paul committed
117
    {
Paul's avatar
Paul committed
118
119
120
121
122
        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
123
124
        return os;
    }
Paul's avatar
Paul committed
125
126
};

Paul's avatar
Paul committed
127
struct pooling
Paul's avatar
Paul committed
128
129
{
    std::string mode;
Paul's avatar
Paul committed
130
131
132
    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
133
    std::string name() const { return "pooling"; }
Paul's avatar
Paul committed
134
135
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
136
        check_shapes{inputs}.has(1).only_dims(4);
Paul's avatar
Paul committed
137

Paul's avatar
Paul committed
138
        const shape& input = inputs.at(0);
Paul's avatar
Paul committed
139
        auto t             = input.type();
Paul's avatar
Paul committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        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
155
    }
Paul's avatar
Paul committed
156

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

Paul's avatar
Paul committed
159
    friend std::ostream& operator<<(std::ostream& os, const pooling& op)
Paul's avatar
Paul committed
160
    {
Paul's avatar
Paul committed
161
162
163
164
165
        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
166
167
        return os;
    }
Paul's avatar
Paul committed
168
169
};

Paul's avatar
Paul committed
170
struct activation
Paul's avatar
Paul committed
171
172
{
    std::string mode;
Paul's avatar
Paul committed
173
    std::string name() const { return "activation"; }
Paul's avatar
Paul committed
174
175
    shape compute_shape(std::vector<shape> inputs) const
    {
Paul's avatar
Paul committed
176
        check_shapes{inputs}.has(1);
Paul's avatar
Paul committed
177
178
        return inputs.front();
    }
Paul's avatar
Paul committed
179

Paul's avatar
Paul committed
180
    argument compute(shape, std::vector<argument>) const { RTG_THROW("not computable"); }
Paul's avatar
Paul committed
181
    friend std::ostream& operator<<(std::ostream& os, const activation& op)
Paul's avatar
Paul committed
182
    {
Paul's avatar
Paul committed
183
        os << op.name() << ":" << op.mode;
Paul's avatar
Paul committed
184
185
        return os;
    }
Paul's avatar
Paul committed
186
187
};

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

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

Paul's avatar
Paul committed
213
    friend std::ostream& operator<<(std::ostream& os, const reshape& op)
Paul's avatar
Paul committed
214
    {
Paul's avatar
Paul committed
215
216
217
        os << op.name() << "[";
        os << "dims={" << stream_range(op.dims) << "}, ";
        os << "]";
Paul's avatar
Paul committed
218
219
        return os;
    }
Paul's avatar
Paul committed
220
221
};

222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
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
    {
        check_shapes{inputs}.has(2).same_type().same_dims().only_dims(2);
        const shape& A = inputs.at(0); 
        const shape& B = inputs.at(1); 
        
        auto t         = A.type();
        if (A.lens()[1] != B.lens()[0])
            RTG_THROW("Inner dimensions do not match");
        return {t, {A.lens()[0], B.lens()[1]}};
    }
  
    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 << "]"; 
    }
};

Paul's avatar
Paul committed
249
250
251
} // namespace rtg

#endif