operation.hpp 2.54 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
4
5
6
7
8

#include <string>
#include <functional>
#include <memory>
#include <type_traits>
#include <utility>
Paul's avatar
Paul committed
9
10
11
#include <migraph/shape.hpp>
#include <migraph/argument.hpp>
#include <migraph/context.hpp>
Paul's avatar
Paul committed
12
#include <migraph/auto_any_cast.hpp>
Paul's avatar
Paul committed
13

Paul's avatar
Paul committed
14
namespace migraph {
Paul's avatar
Paul committed
15

Paul's avatar
Paul committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
#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;
    /// 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.
    shape compute_shape(std::vector<shape> input) const;
    /**
     * @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.
     */
    argument compute(context& ctx, shape output, std::vector<argument> input) const;
    /// 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
44
45
namespace operation_stream {

Paul's avatar
Paul committed
46
47
template <class T>
auto operator<<(std::ostream& os, const T& x) -> decltype(os << x.name())
Paul's avatar
Paul committed
48
{
Paul's avatar
Paul committed
49
    return os << x.name();
Paul's avatar
Paul committed
50
51
}

Paul's avatar
Paul committed
52
} // namespace operation_stream
Paul's avatar
Paul committed
53

Paul's avatar
Paul committed
54
template <class T>
Paul's avatar
Paul committed
55
56
57
58
59
argument compute_op(const T& x, context& ctx, shape output_shape, std::vector<argument> input)
{
    return x.compute(auto_any_cast(ctx), output_shape, input);
}

60
<%
Paul's avatar
Paul committed
61
62
63
interface('operation',
    virtual('name', returns='std::string', const=True),
    virtual('compute_shape', returns='shape', input='std::vector<shape>', const=True),
Paul's avatar
Paul committed
64
    virtual('compute', returns='argument', ctx='context&', output='shape', input='std::vector<argument>', const=True, default='compute_op'),
Paul's avatar
Paul committed
65
    friend('operator<<', returns='std::ostream &', os='std::ostream &', op='const operation &', using='migraph::operation_stream::operator<<')
Paul's avatar
Paul committed
66
)
67
%>
Paul's avatar
Paul committed
68

Paul's avatar
Paul committed
69
70
#endif

Paul's avatar
Paul committed
71
} // namespace migraph
Paul's avatar
Paul committed
72
73

#endif