operation.hpp 10.5 KB
Newer Older
Paul's avatar
Paul committed
1
2
#ifndef MIGRAPH_GUARD_MIGRAPHLIB_OPERAND_HPP
#define MIGRAPH_GUARD_MIGRAPHLIB_OPERAND_HPP
Paul's avatar
Paul committed
3

Paul's avatar
Paul committed
4
#include <cassert>
Paul's avatar
Paul committed
5
#include <string>
Paul's avatar
Paul committed
6
#include <functional>
Paul's avatar
Paul committed
7
8
9
#include <memory>
#include <type_traits>
#include <utility>
Paul's avatar
Paul committed
10
#include <migraph/shape.hpp>
Paul's avatar
Paul committed
11
#include <migraph/rank.hpp>
Paul's avatar
Paul committed
12
13
#include <migraph/argument.hpp>
#include <migraph/context.hpp>
Paul's avatar
Paul committed
14
#include <migraph/auto_any_cast.hpp>
Paul's avatar
Paul committed
15

Paul's avatar
Paul committed
16
namespace migraph {
Paul's avatar
Paul committed
17

Paul's avatar
Paul committed
18
19
#ifdef DOXYGEN

Paul's avatar
Paul committed
20
/// The operation interface represents an action an instruction will perform. All
Paul's avatar
Paul committed
21
22
23
24
25
26
27
28
/// operation classes must be CopyConstructible.
struct operation
{
    /// A unique name identifying the operation
    std::string name() const;
    /// This is used to compute the resulting shape from an operation. If an
    /// operation cannot be run with input shapes, then it should throw an
    /// exception.
Paul's avatar
Paul committed
29
    shape compute_shape(const std::vector<shape>& input) const;
Paul's avatar
Paul committed
30
31
32
33
34
35
36
37
38
39
40
    /**
     * @brief This performs the operation's computation
     *
     * @param ctx This is the context created by the `target` during compilation. Implementations
     * can use the target's `context` class rather than the `context` interface class.
     * @param output This is the output shape. It is equivalent to running `compute_shape` with each
     * `shape` of the `argument`.
     * @param input This is the `argument` result from the previous instuction's computation.
     * @return Return an `argument` of the result computation. The `shape` of `argument` should be
     * the same the `output` shape.
     */
Paul's avatar
Paul committed
41
    argument compute(context& ctx, const shape& output, const std::vector<argument>& input) const;
Paul's avatar
Paul committed
42
43
44
45
46
47
48
    /// An optional stream operator to print the operation. When this is not
    /// implemented, it will just print the operation's name.
    friend std::ostream& operator<<(std::ostream& os, const operation& op);
};

#else

Paul's avatar
Paul committed
49
50
51
52
53
54
55
56
57
58
namespace operation_stream {

template <class T>
auto operator<<(std::ostream& os, const T& x) -> decltype(os << x.name())
{
    return os << x.name();
}

} // namespace operation_stream

Paul's avatar
Paul committed
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
template <class T>
auto compute_op(rank<1>,
                const T& x,
                context& ctx,
                const shape& output_shape,
                const std::vector<argument>& input)
    -> decltype(x.compute(auto_any_cast(ctx), output_shape, input))
{
    return x.compute(auto_any_cast(ctx), output_shape, input);
}

template <class T>
argument compute_op(rank<0>, const T& x, context&, const shape&, const std::vector<argument>&)
{
    MIGRAPH_THROW("Not computable: " + x.name());
}

Paul's avatar
Paul committed
76
template <class T>
Paul's avatar
Paul committed
77
78
argument
compute_op(const T& x, context& ctx, const shape& output_shape, const std::vector<argument>& input)
Paul's avatar
Paul committed
79
{
Paul's avatar
Paul committed
80
    return compute_op(rank<1>{}, x, ctx, output_shape, input);
Paul's avatar
Paul committed
81
82
}

Paul's avatar
Paul committed
83
/*
Paul's avatar
Paul committed
84
85
 * Type-erased interface for:
 *
Paul's avatar
Paul committed
86
 * struct operation
Paul's avatar
Paul committed
87
 * {
Paul's avatar
Paul committed
88
 *      std::string name() const;
Paul's avatar
Paul committed
89
90
 *      shape compute_shape(const std::vector<shape>& input) const;
 *      argument compute(context& ctx,const shape& output,const std::vector<argument>& input) const;
Paul's avatar
Paul committed
91
 *     friend std::ostream & operator<<(std::ostream & os,const operation & op) ;
Paul's avatar
Paul committed
92
93
94
 * };
 *
 */
Paul's avatar
Paul committed
95

Paul's avatar
Paul committed
96
struct operation
Paul's avatar
Paul committed
97
{
Paul's avatar
Paul committed
98
    // Constructors
Paul's avatar
Paul committed
99
    operation() = default;
Paul's avatar
Paul committed
100

Paul's avatar
Paul committed
101
    template <typename PrivateDetailTypeErasedT>
Paul's avatar
Paul committed
102
    operation(PrivateDetailTypeErasedT value)
Paul's avatar
Paul committed
103
104
105
106
        : private_detail_te_handle_mem_var(
              std::make_shared<private_detail_te_handle_type<
                  typename std::remove_reference<PrivateDetailTypeErasedT>::type>>(
                  std::forward<PrivateDetailTypeErasedT>(value)))
Paul's avatar
Paul committed
107
108
109
110
    {
    }

    // Assignment
Paul's avatar
Paul committed
111
    template <typename PrivateDetailTypeErasedT>
Paul's avatar
Paul committed
112
    operation& operator=(PrivateDetailTypeErasedT value)
Paul's avatar
Paul committed
113
    {
Paul's avatar
Paul committed
114
115
116
117
118
        if(private_detail_te_handle_mem_var.unique())
            *private_detail_te_handle_mem_var = std::forward<PrivateDetailTypeErasedT>(value);
        else if(!private_detail_te_handle_mem_var)
            private_detail_te_handle_mem_var = std::make_shared<PrivateDetailTypeErasedT>(
                std::forward<PrivateDetailTypeErasedT>(value));
Paul's avatar
Paul committed
119
120
121
        return *this;
    }

Paul's avatar
Paul committed
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
    // Cast
    template <typename PrivateDetailTypeErasedT>
    PrivateDetailTypeErasedT* any_cast()
    {
        return private_detail_te_get_handle().type() == typeid(PrivateDetailTypeErasedT)
                   ? std::addressof(static_cast<private_detail_te_handle_type<
                                        typename std::remove_cv<PrivateDetailTypeErasedT>::type>&>(
                                        private_detail_te_get_handle())
                                        .private_detail_te_value)
                   : nullptr;
    }

    template <typename PrivateDetailTypeErasedT>
    const typename std::remove_cv<PrivateDetailTypeErasedT>::type* any_cast() const
    {
        return private_detail_te_get_handle().type() == typeid(PrivateDetailTypeErasedT)
                   ? std::addressof(static_cast<const private_detail_te_handle_type<
                                        typename std::remove_cv<PrivateDetailTypeErasedT>::type>&>(
                                        private_detail_te_get_handle())
                                        .private_detail_te_value)
                   : nullptr;
    }

Paul's avatar
Paul committed
145
146
147
148
149
150
151
152
    const std::type_info& type_id() const
    {
        if(private_detail_te_handle_empty())
            return typeid(std::nullptr_t);
        else
            return private_detail_te_get_handle().type();
    }

Paul's avatar
Paul committed
153
154
    std::string name() const
    {
Paul's avatar
Paul committed
155
156
        assert((*this).private_detail_te_handle_mem_var);
        return (*this).private_detail_te_get_handle().name();
Paul's avatar
Paul committed
157
158
    }

Paul's avatar
Paul committed
159
    shape compute_shape(const std::vector<shape>& input) const
Paul's avatar
Paul committed
160
    {
Paul's avatar
Paul committed
161
        assert((*this).private_detail_te_handle_mem_var);
Paul's avatar
Paul committed
162
        return (*this).private_detail_te_get_handle().compute_shape(input);
Paul's avatar
Paul committed
163
164
    }

Paul's avatar
Paul committed
165
    argument compute(context& ctx, const shape& output, const std::vector<argument>& input) const
Paul's avatar
Paul committed
166
    {
Paul's avatar
Paul committed
167
        assert((*this).private_detail_te_handle_mem_var);
Paul's avatar
Paul committed
168
        return (*this).private_detail_te_get_handle().compute(ctx, output, input);
Paul's avatar
Paul committed
169
170
171
172
173
174
    }

    friend std::ostream& operator<<(std::ostream& os, const operation& op)
    {
        assert(op.private_detail_te_handle_mem_var);
        return op.private_detail_te_get_handle().operator_shift_left(os);
Paul's avatar
Paul committed
175
176
177
    }

    private:
Paul's avatar
Paul committed
178
    struct private_detail_te_handle_base_type
Paul's avatar
Paul committed
179
    {
Paul's avatar
Paul committed
180
181
        virtual ~private_detail_te_handle_base_type() {}
        virtual std::shared_ptr<private_detail_te_handle_base_type> clone() const = 0;
Paul's avatar
Paul committed
182
        virtual const std::type_info& type() const                                = 0;
Paul's avatar
Paul committed
183

Paul's avatar
Paul committed
184
185
186
187
188
        virtual std::string name() const                                   = 0;
        virtual shape compute_shape(const std::vector<shape>& input) const = 0;
        virtual argument
        compute(context& ctx, const shape& output, const std::vector<argument>& input) const = 0;
        virtual std::ostream& operator_shift_left(std::ostream& os) const                    = 0;
Paul's avatar
Paul committed
189
190
    };

Paul's avatar
Paul committed
191
192
    template <typename PrivateDetailTypeErasedT>
    struct private_detail_te_handle_type : private_detail_te_handle_base_type
Paul's avatar
Paul committed
193
    {
Paul's avatar
Paul committed
194
195
196
197
198
199
        template <typename PrivateDetailTypeErasedU = PrivateDetailTypeErasedT>
        private_detail_te_handle_type(
            PrivateDetailTypeErasedT value,
            typename std::enable_if<std::is_reference<PrivateDetailTypeErasedU>::value>::type* =
                nullptr)
            : private_detail_te_value(value)
Paul's avatar
Paul committed
200
201
202
        {
        }

Paul's avatar
Paul committed
203
204
205
206
207
208
        template <typename PrivateDetailTypeErasedU = PrivateDetailTypeErasedT>
        private_detail_te_handle_type(
            PrivateDetailTypeErasedT value,
            typename std::enable_if<!std::is_reference<PrivateDetailTypeErasedU>::value,
                                    int>::type* = nullptr) noexcept
            : private_detail_te_value(std::move(value))
Paul's avatar
Paul committed
209
210
211
        {
        }

Paul's avatar
Paul committed
212
        std::shared_ptr<private_detail_te_handle_base_type> clone() const override
Paul's avatar
Paul committed
213
        {
Paul's avatar
Paul committed
214
            return std::make_shared<private_detail_te_handle_type>(private_detail_te_value);
Paul's avatar
Paul committed
215
216
        }

Paul's avatar
Paul committed
217
        const std::type_info& type() const override { return typeid(private_detail_te_value); }
Paul's avatar
Paul committed
218

Paul's avatar
Paul committed
219
        std::string name() const override { return private_detail_te_value.name(); }
Paul's avatar
Paul committed
220

Paul's avatar
Paul committed
221
        shape compute_shape(const std::vector<shape>& input) const override
Paul's avatar
Paul committed
222
        {
Paul's avatar
Paul committed
223

Paul's avatar
Paul committed
224
            return private_detail_te_value.compute_shape(input);
Paul's avatar
Paul committed
225
226
        }

Paul's avatar
Paul committed
227
228
229
        argument compute(context& ctx,
                         const shape& output,
                         const std::vector<argument>& input) const override
Paul's avatar
Paul committed
230
        {
Paul's avatar
Paul committed
231

Paul's avatar
Paul committed
232
            return compute_op(private_detail_te_value, ctx, output, input);
Paul's avatar
Paul committed
233
234
        }

Paul's avatar
Paul committed
235
236
        std::ostream& operator_shift_left(std::ostream& os) const override
        {
Paul's avatar
Paul committed
237
            using migraph::operation_stream::operator<<;
Paul's avatar
Paul committed
238
239
240
            return os << private_detail_te_value;
        }

Paul's avatar
Paul committed
241
        PrivateDetailTypeErasedT private_detail_te_value;
Paul's avatar
Paul committed
242
243
    };

Paul's avatar
Paul committed
244
245
246
    template <typename PrivateDetailTypeErasedT>
    struct private_detail_te_handle_type<std::reference_wrapper<PrivateDetailTypeErasedT>>
        : private_detail_te_handle_type<PrivateDetailTypeErasedT&>
Paul's avatar
Paul committed
247
    {
Paul's avatar
Paul committed
248
249
        private_detail_te_handle_type(std::reference_wrapper<PrivateDetailTypeErasedT> ref)
            : private_detail_te_handle_type<PrivateDetailTypeErasedT&>(ref.get())
Paul's avatar
Paul committed
250
251
252
253
        {
        }
    };

Paul's avatar
Paul committed
254
255
256
257
258
    bool private_detail_te_handle_empty() const
    {
        return private_detail_te_handle_mem_var == nullptr;
    }

Paul's avatar
Paul committed
259
260
    const private_detail_te_handle_base_type& private_detail_te_get_handle() const
    {
Paul's avatar
Paul committed
261
        assert(private_detail_te_handle_mem_var != nullptr);
Paul's avatar
Paul committed
262
263
        return *private_detail_te_handle_mem_var;
    }
Paul's avatar
Paul committed
264

Paul's avatar
Paul committed
265
    private_detail_te_handle_base_type& private_detail_te_get_handle()
Paul's avatar
Paul committed
266
    {
Paul's avatar
Paul committed
267
        assert(private_detail_te_handle_mem_var != nullptr);
Paul's avatar
Paul committed
268
269
270
        if(!private_detail_te_handle_mem_var.unique())
            private_detail_te_handle_mem_var = private_detail_te_handle_mem_var->clone();
        return *private_detail_te_handle_mem_var;
Paul's avatar
Paul committed
271
272
    }

Paul's avatar
Paul committed
273
    std::shared_ptr<private_detail_te_handle_base_type> private_detail_te_handle_mem_var;
Paul's avatar
Paul committed
274
275
};

Paul's avatar
Paul committed
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
template <typename ValueType>
inline const ValueType* any_cast(const operation* x)
{
    return x->any_cast<ValueType>();
}

template <typename ValueType>
inline ValueType* any_cast(operation* x)
{
    return x->any_cast<ValueType>();
}

template <typename ValueType>
inline ValueType& any_cast(operation& x)
{
Paul's avatar
Paul committed
291
    auto* y = x.any_cast<typename std::remove_reference<ValueType>::type>();
Paul's avatar
Paul committed
292
293
294
295
296
297
298
299
    if(y == nullptr)
        throw std::bad_cast();
    return *y;
}

template <typename ValueType>
inline const ValueType& any_cast(const operation& x)
{
Paul's avatar
Paul committed
300
    const auto* y = x.any_cast<typename std::remove_reference<ValueType>::type>();
Paul's avatar
Paul committed
301
302
303
304
305
    if(y == nullptr)
        throw std::bad_cast();
    return *y;
}

Paul's avatar
Paul committed
306
307
#endif

Paul's avatar
Paul committed
308
} // namespace migraph
Paul's avatar
Paul committed
309
310

#endif