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

Paul's avatar
Paul committed
4
#include <cassert>
Paul's avatar
Paul committed
5
6
7
8
9
#include <string>
#include <functional>
#include <memory>
#include <type_traits>
#include <utility>
Paul's avatar
Paul committed
10
11
#include <migraphx/reflect.hpp>
#include <migraphx/streamutils.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
12
#include <migraphx/normalize_attributes.hpp>
Paul's avatar
Paul committed
13
#include <migraphx/argument.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
14
#include <migraphx/module_ref.hpp>
15
#include <migraphx/serialize.hpp>
Paul's avatar
Paul committed
16
#include <migraphx/auto_any_cast.hpp>
Paul's avatar
Paul committed
17
#include <migraphx/config.hpp>
Paul's avatar
Paul committed
18

Paul's avatar
Paul committed
19
namespace migraphx {
Paul's avatar
Paul committed
20
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
21

Paul's avatar
Paul committed
22
23
struct context;

Paul's avatar
Paul committed
24
25
26
27
28
29
30
31
#ifdef DOXYGEN

/// The operation interface represents an action an instruction will perform. All
/// operation classes must be CopyConstructible.
struct operation
{
    /// A unique name identifying the operation
    std::string name() const;
Paul's avatar
Paul committed
32
33
    /// An optional method that can be used to finalize the operator before running
    void finalize(context& ctx);
Paul's avatar
Paul committed
34
35
36
    /// 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
37
    shape compute_shape(const std::vector<shape>& input) const;
Paul's avatar
Paul committed
38
    /**
Paul's avatar
Paul committed
39
40
     * @brief This performs the operation's computation.
     *
Paul's avatar
Paul committed
41
42
     * This method can be optional when the operation is only used as a placeholder to be lowered
     * later on.
Paul's avatar
Paul committed
43
44
45
46
47
     *
     * @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`.
Paul's avatar
Paul committed
48
     * @param input This is the `argument` result from the previous instruction's computation.
Paul's avatar
Paul committed
49
50
     * @return Return an `argument` of the result computation. The `shape` of `argument` should be
     * the same the `output` shape.
Paul's avatar
Paul committed
51
     */
Paul's avatar
Paul committed
52
    argument compute(context& ctx, const shape& output, const std::vector<argument>& input) const;
Paul's avatar
Paul committed
53
54
    /// An optional method to return which argument the output will alias. If
    /// there is no aliased output then -1 can be returned.
Paul's avatar
Paul committed
55
    std::ptrdiff_t output_alias(const std::vector<shape>& input) const;
Paul's avatar
Paul committed
56
57
    /// An optional stream operator to print the operation. When this is not
    /// implemented, it will just print the operation's name.
Paul's avatar
Paul committed
58
    friend std::ostream& operator<<(std::ostream& os, const operation& op);
Paul's avatar
Paul committed
59
60
};

Paul's avatar
Paul committed
61
62
/// Returns true if operation does not require a context to run compute
bool is_context_free(const operation& x);
Shucai Xiao's avatar
Shucai Xiao committed
63
64
/// Returns true if operation needs normalization before running compute
bool need_normalization(const operation& x);
Paul's avatar
Paul committed
65
66
/// Returns true if the operation has a finalize method
bool has_finalize(const operation& x);
Paul's avatar
Paul committed
67

Paul's avatar
Paul committed
68
69
#else

70
71
72
namespace detail {

namespace operation_operators {
Paul's avatar
Paul committed
73

Paul's avatar
Paul committed
74
75
template <class T>
auto operator<<(std::ostream& os, const T& x) -> decltype(os << x.name())
Paul's avatar
Paul committed
76
{
Paul's avatar
Paul committed
77
78
    os << x.name();
    char delim = '[';
Paul's avatar
Paul committed
79
    reflect_each(x, [&](auto&& y, auto name) {
Paul's avatar
Paul committed
80
        os << delim;
Paul's avatar
Paul committed
81
82
        os << name << "=";
        stream_write_value(os, y);
Paul's avatar
Paul committed
83
84
        delim = ',';
    });
Paul's avatar
Paul committed
85
86
    if(delim == ',')
        os << "]";
Paul's avatar
Paul committed
87
    return os;
Paul's avatar
Paul committed
88
89
}

Paul's avatar
Paul committed
90
91
92
template <class T, class U>
auto operator==(const T& x, const U& y) -> decltype(x.name() == y.name())
{
Paul's avatar
Paul committed
93
94
    static_assert(is_reflectable<T>{} or sizeof(T) <= 1,
                  "Missing equality operator or reflect method.");
Paul's avatar
Paul committed
95
96
97
98
99
100
    if(x.name() != y.name())
        return false;
    const auto& yy = any_cast<T>(y);
    return reflect_tie(x) == reflect_tie(yy);
}

101
} // namespace operation_operators
Paul's avatar
Paul committed
102

Shucai Xiao's avatar
Shucai Xiao committed
103
template <class T>
Shucai Xiao's avatar
Shucai Xiao committed
104
105
auto normalize_compute_shape_op(rank<1>, const T& x, const std::vector<shape>& inputs)
    -> decltype(x.normalize_compute_shape(inputs))
Shucai Xiao's avatar
Shucai Xiao committed
106
107
108
109
110
111
{
    dependent_type<operation, T> y = x;
    normalize_attributes(y, inputs[0].lens());
    return any_cast<T>(y).normalize_compute_shape(inputs);
}

Shucai Xiao's avatar
Shucai Xiao committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
template <class T>
shape normalize_compute_shape_op(rank<0>, const T& x, const std::vector<shape>&)
{
    std::string name = x.name();
    MIGRAPHX_THROW("Shape not computable: " + name);
}

template <class T>
shape normalize_compute_shape_op(const T& x, const std::vector<shape>& inputs)
{
    return normalize_compute_shape_op(rank<1>{}, x, inputs);
}

template <class T>
auto compute_shape_op(rank<1>,
                      const T& x,
                      const std::vector<shape>& inputs,
                      const std::vector<module_ref>& mod_args)
    -> decltype(x.compute_shape(inputs, mod_args))
{
    return x.compute_shape(inputs, mod_args);
}

template <class T>
shape
    compute_shape_op(rank<0>, const T& x, const std::vector<shape>&, const std::vector<module_ref>&)
{
    std::string name = x.name();
    MIGRAPHX_THROW("Shape not computable: " + name);
}

template <class T>
shape compute_shape_op(const T& x,
                       const std::vector<shape>& inputs,
                       const std::vector<module_ref>& mod_args)
{
    return compute_shape_op(rank<1>{}, x, inputs, mod_args);
}

template <class T>
auto normalize_compute_shape_op(rank<1>,
                                const T& x,
                                const std::vector<shape>& inputs,
                                std::vector<module_ref>& mod_args)
    -> decltype(x.normalize_compute_shape(inputs, mod_args))
{
    return x.normalize_compute_shape(inputs, mod_args);
}

template <class T>
shape normalize_compute_shape_op(rank<0>,
                                 const T& x,
                                 const std::vector<shape>&,
                                 const std::vector<module_ref>&)
{
    std::string name = x.name();
    MIGRAPHX_THROW("Shape not computable: " + name);
}

template <class T>
shape normalize_compute_shape_op(const T& x,
                                 const std::vector<shape>& inputs,
                                 std::vector<module_ref>& mod_args)
{
    return normalize_compute_shape_op(rank<1>{}, x, inputs, mod_args);
}

Paul's avatar
Paul committed
179
template <class T>
Paul's avatar
Paul committed
180
auto compute_op(rank<2>,
Paul's avatar
Paul committed
181
182
183
184
185
                const T& x,
                context& ctx,
                const shape& output_shape,
                const std::vector<argument>& input)
    -> decltype(x.compute(auto_any_cast(ctx), output_shape, input))
Paul's avatar
Paul committed
186
187
188
189
{
    return x.compute(auto_any_cast(ctx), output_shape, input);
}

Paul's avatar
Paul committed
190
template <class T>
Paul's avatar
Paul committed
191
192
auto compute_op(
    rank<1>, const T& x, context&, const shape& output_shape, const std::vector<argument>& input)
Paul's avatar
Paul committed
193
194
195
196
197
    -> decltype(x.compute(output_shape, input))
{
    return x.compute(output_shape, input);
}

Paul's avatar
Paul committed
198
template <class T>
Paul's avatar
Paul committed
199
argument compute_op(rank<0>, const T& x, context&, const shape&, const std::vector<argument>&)
Paul's avatar
Paul committed
200
{
Paul's avatar
Paul committed
201
    std::string name = x.name();
Paul's avatar
Paul committed
202
    MIGRAPHX_THROW("Not computable: " + name);
Paul's avatar
Paul committed
203
204
}

Paul's avatar
Paul committed
205
template <class T>
Paul's avatar
Paul committed
206
207
argument
compute_op(const T& x, context& ctx, const shape& output_shape, const std::vector<argument>& input)
Paul's avatar
Paul committed
208
{
Paul's avatar
Paul committed
209
210
211
212
    return compute_op(rank<2>{}, x, ctx, output_shape, input);
}

template <class T>
Paul's avatar
Paul committed
213
auto compute_op(rank<2>, const T& x, const shape& output_shape, const std::vector<argument>& input)
Paul's avatar
Paul committed
214
215
216
217
218
219
    -> decltype(x.compute(output_shape, input))
{
    return x.compute(output_shape, input);
}

template <class T>
Paul's avatar
Paul committed
220
auto compute_op(rank<1>, const T& x, const shape& output_shape, const std::vector<argument>& input)
Paul's avatar
Paul committed
221
222
223
224
225
226
227
228
229
230
231
232
233
234
    -> decltype(x.compute(auto_any_cast(std::declval<context&>()), output_shape, input))
{
    std::string name = x.name();
    MIGRAPHX_THROW("Not computable without a context: " + name);
}

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

template <class T>
Paul's avatar
Paul committed
235
argument compute_op(const T& x, const shape& output_shape, const std::vector<argument>& input)
Paul's avatar
Paul committed
236
237
238
239
240
{
    return compute_op(rank<2>{}, x, output_shape, input);
}

template <class T>
Paul's avatar
Paul committed
241
242
243
244
245
auto is_context_free_op(rank<1>,
                        const T& x,
                        const shape& output_shape,
                        const std::vector<argument>& input)
    -> decltype(x.compute(output_shape, input), std::true_type{});
Paul's avatar
Paul committed
246
247

template <class T>
Paul's avatar
Paul committed
248
249
auto is_context_free_op(rank<0>, const T&, const shape&, const std::vector<argument>&)
    -> std::false_type;
Paul's avatar
Paul committed
250
251

template <class T>
Paul's avatar
Paul committed
252
253
auto is_context_free_op(const T& x) -> decltype(is_context_free_op(
    rank<1>{}, x, std::declval<const shape&>(), std::declval<std::vector<argument>>()))
Paul's avatar
Paul committed
254
255
{
    return {};
Paul's avatar
Paul committed
256
257
}

Shucai Xiao's avatar
Shucai Xiao committed
258
259
260
261
262
263
264
265
266
267
268
269
270
271
template <class T>
auto need_normalization_op(rank<1>, const T& x, const std::vector<shape>& inputs)
    -> decltype(x.normalize_compute_shape(inputs), std::true_type{});

template <class T>
auto need_normalization_op(rank<0>, const T&, const std::vector<shape>&) -> std::false_type;

template <class T>
auto need_normalization_op(const T& x)
    -> decltype(need_normalization_op(rank<1>{}, x, std::declval<std::vector<shape>>()))
{
    return {};
}

Paul's avatar
Paul committed
272
template <class T>
273
std::ptrdiff_t output_alias_op(const T&, const std::vector<shape>&)
Paul's avatar
Paul committed
274
275
276
277
{
    return -1;
}

Paul's avatar
Paul committed
278
template <class T>
Paul's avatar
Paul committed
279
280
auto finalize_op(
    rank<1>, T& x, context& ctx, const shape& output_shape, const std::vector<shape>& input)
Paul's avatar
Paul committed
281
282
283
284
285
286
287
    -> decltype(x.finalize(auto_any_cast(ctx), output_shape, input), void())
{
    x.finalize(auto_any_cast(ctx), output_shape, input);
}

template <class T>
void finalize_op(rank<0>, T&, context&, const shape&, const std::vector<shape>&)
Paul's avatar
Paul committed
288
289
{
}
Paul's avatar
Paul committed
290
291
292
293
294
295
296
297

template <class T>
void finalize_op(T& x, context& ctx, const shape& output_shape, const std::vector<shape>& input)
{
    finalize_op(rank<1>{}, x, ctx, output_shape, input);
}

template <class T>
Paul's avatar
Paul committed
298
299
auto has_finalize_op(
    rank<1>, T& x, context& ctx, const shape& output_shape, const std::vector<shape>& input)
Paul's avatar
Paul committed
300
301
302
303
304
305
306
    -> decltype(x.finalize(auto_any_cast(ctx), output_shape, input), std::true_type{});

template <class T>
auto has_finalize_op(rank<0>, T&, context&, const shape&, const std::vector<shape>&)
    -> std::false_type;

template <class T>
Paul's avatar
Paul committed
307
308
309
310
311
auto has_finalize_op(const T&) -> decltype(has_finalize_op(rank<1>{},
                                                           std::declval<T&>(),
                                                           std::declval<context&>(),
                                                           std::declval<const shape&>(),
                                                           std::declval<std::vector<shape>>()))
Paul's avatar
Paul committed
312
313
314
315
{
    return {};
}

316
317
318
319
320
321
template <class T>
value attributes_op(const T&)
{
    return value::object{};
}

322
323
324
325
326
327
328
329
330
template <class T>
value to_value_op(const T& x)
{
    return migraphx::to_value(x);
}

template <class T>
void from_value_op(T& x, const value& v)
{
331
332
    if(not(v.is_object() or (v.empty() and v.is_array())))
        MIGRAPHX_THROW("Value is not an object");
333
334
335
    return migraphx::from_value(v, x);
}

336
337
} // namespace detail

338
<%
Paul's avatar
Paul committed
339
340
341
 interface(
     'operation',
     virtual('name', returns = 'std::string', const = True),
Paul's avatar
Paul committed
342
343
     virtual(
         'is_context_free', returns = 'bool', const = True, default = 'detail::is_context_free_op'),
Shucai Xiao's avatar
Shucai Xiao committed
344
345
346
347
     virtual('need_normalization',
             returns = 'bool',
             const   = True,
             default = 'detail::need_normalization_op'),
348
     virtual('has_finalize', returns = 'bool', const = True, default = 'detail::has_finalize_op'),
Paul's avatar
Paul committed
349
     virtual('output_alias',
Paul's avatar
Paul committed
350
             returns = 'std::ptrdiff_t',
Paul's avatar
Paul committed
351
352
             input   = 'const std::vector<shape>&',
             const   = True,
353
             default = 'detail::output_alias_op'),
Paul's avatar
Paul committed
354
355
356
357
     virtual('finalize',
             ctx     = 'context&',
             output  = 'const shape&',
             input   = 'const std::vector<shape>&',
358
             default = 'detail::finalize_op'),
Shucai Xiao's avatar
Shucai Xiao committed
359
360
361
362
363
     virtual('compute_shape',
             returns = 'shape',
             input   = 'const std::vector<shape>&',
             const   = True,
             default = 'detail::normalize_compute_shape_op'),
Shucai Xiao's avatar
Shucai Xiao committed
364
365
366
367
368
369
     virtual('compute_shape',
             returns  = 'shape',
             inputs   = 'const std::vector<shape>&',
             mod_args = 'const std::vector<module_ref>&',
             const    = True,
             default  = 'detail::compute_shape_op'),
Paul's avatar
Paul committed
370
371
372
373
374
375
     virtual('compute',
             returns = 'argument',
             ctx     = 'context&',
             output  = 'const shape&',
             input   = 'const std::vector<argument>&',
             const   = True,
376
             default = 'detail::compute_op'),
Paul's avatar
Paul committed
377
378
379
380
381
     virtual('compute',
             returns = 'argument',
             output  = 'const shape&',
             input   = 'const std::vector<argument>&',
             const   = True,
382
             default = 'detail::compute_op'),
383
384
     virtual('to_value', returns = 'value', const = True, default = 'detail::to_value_op'),
     virtual('from_value', v = 'const value&', default = 'detail::from_value_op'),
385
     virtual('attributes', returns = 'value', const = True, default = 'detail::attributes_op'),
Paul's avatar
Paul committed
386
387
388
389
     friend('operator<<',
            returns = 'std::ostream &',
            os      = 'std::ostream &',
            op      = 'const operation &',
390
            using   = 'migraphx::detail::operation_operators::operator<<'),
Paul's avatar
Paul committed
391
392
393
394
     friend('operator==',
            returns = 'bool',
            x       = 'const operation &',
            y       = 'const operation &',
395
            using   = 'migraphx::detail::operation_operators::operator==')) %>
Paul's avatar
Paul committed
396
397

    inline bool operator!=(const operation& x, const operation& y)
Paul's avatar
Paul committed
398
399
400
401
{
    return !(x == y);
}

402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
inline shape compute_shape(const operation& op, const std::vector<shape>& inputs)
{
    return op.compute_shape(inputs);
}

template <class T>
inline auto compute_shape(const T& op, const std::vector<shape>& inputs)
    -> decltype(op.compute_shape(inputs))
{
    return op.compute_shape(inputs);
}

template <class T>
inline auto compute_shape(const T& op, const std::vector<shape>& inputs)
    -> decltype(op.normalize_compute_shape(inputs))
{
    return detail::normalize_compute_shape_op(op, inputs);
}

Shucai Xiao's avatar
Shucai Xiao committed
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
inline shape compute_shape(const operation& op,
                           const std::vector<shape>& inputs,
                           const std::vector<module_ref>& mod_args)
{
    return op.compute_shape(inputs, mod_args);
}

template <class T>
inline auto compute_shape(const T& op,
                          const std::vector<shape>& inputs,
                          const std::vector<module_ref>& mod_args)
    -> decltype(op.compute_shape(inputs, mod_args))
{
    return op.compute_shape(inputs, mod_args);
}

template <class T>
inline auto compute_shape(const T& op,
                          const std::vector<shape>& inputs,
                          const std::vector<module_ref>& mod_args)
    -> decltype(op.normalize_compute_shape(inputs, mod_args))
{
    return detail::normalize_compute_shape_op(op, inputs, mod_args);
}

Paul's avatar
Paul committed
446
inline bool is_context_free(const operation& op) { return op.is_context_free(); }
Paul's avatar
Paul committed
447

Paul's avatar
Paul committed
448
template <class T>
Paul's avatar
Paul committed
449
450
bool is_context_free(const T& x)
{
451
    return detail::is_context_free_op(x);
Paul's avatar
Paul committed
452
453
}

Shucai Xiao's avatar
Shucai Xiao committed
454
455
456
457
458
459
460
461
inline bool need_normalization(const operation& op) { return op.need_normalization(); }

template <class T>
bool need_normalization(const T& x)
{
    return detail::need_normalization_op(x);
}

Paul's avatar
Paul committed
462
463
464
465
466
inline bool has_finalize(const operation& op) { return op.has_finalize(); }

template <class T>
bool has_finalize(const T& x)
{
467
    return detail::has_finalize_op(x);
Paul's avatar
Paul committed
468
469
}

470
471
472
void migraphx_to_value(value& v, const operation& op);
void migraphx_from_value(const value& v, operation& op);

Paul's avatar
Paul committed
473
474
#endif

Paul's avatar
Paul committed
475
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
476
} // namespace migraphx
Paul's avatar
Paul committed
477
478

#endif