operation.hpp 3.18 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
6
7
8
9
#include <string>
#include <functional>
#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
20
21
22
23
24
25
26
27
28
#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.
Paul's avatar
Paul committed
29
    shape compute_shape(const std::vector<shape>& input) const;
Paul's avatar
Paul committed
30
31
    /**
     * @brief This performs the operation's computation
Paul's avatar
Paul committed
32
33
34
35
36
     *
     * @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
37
     * @param input This is the `argument` result from the previous instuction's computation.
Paul's avatar
Paul committed
38
39
     * @return Return an `argument` of the result computation. The `shape` of `argument` should be
     * the same the `output` shape.
Paul's avatar
Paul committed
40
     */
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
    /// 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
44
    friend std::ostream& operator<<(std::ostream& os, const operation& op);
Paul's avatar
Paul committed
45
46
47
48
};

#else

Paul's avatar
Paul committed
49
50
namespace operation_stream {

Paul's avatar
Paul committed
51
52
template <class T>
auto operator<<(std::ostream& os, const T& x) -> decltype(os << x.name())
Paul's avatar
Paul committed
53
{
Paul's avatar
Paul committed
54
    return os << x.name();
Paul's avatar
Paul committed
55
56
}

Paul's avatar
Paul committed
57
} // namespace operation_stream
Paul's avatar
Paul committed
58

Paul's avatar
Paul committed
59
template <class T>
Paul's avatar
Paul committed
60
61
62
63
64
65
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))
Paul's avatar
Paul committed
66
67
68
69
70
{
    return x.compute(auto_any_cast(ctx), output_shape, input);
}

template <class T>
Paul's avatar
Paul committed
71
argument compute_op(rank<0>, const T& x, context&, const shape&, const std::vector<argument>&)
Paul's avatar
Paul committed
72
{
Paul's avatar
Paul committed
73
74
    std::string name = x.name();
    MIGRAPH_THROW("Not computable: " + name);
Paul's avatar
Paul committed
75
76
}

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

84
<%
Paul's avatar
Paul committed
85
86
interface('operation',
    virtual('name', returns='std::string', const=True),
Paul's avatar
Paul committed
87
88
    virtual('compute_shape', returns='shape', input='const std::vector<shape>&', const=True),
    virtual('compute', returns='argument', ctx='context&', output='const shape&', input='const std::vector<argument>&', const=True, default='compute_op'),
Paul's avatar
Paul committed
89
    friend('operator<<', returns='std::ostream &', os='std::ostream &', op='const operation &', using='migraph::operation_stream::operator<<')
Paul's avatar
Paul committed
90
)
91
%>
Paul's avatar
Paul committed
92

Paul's avatar
Paul committed
93
94
#endif

Paul's avatar
Paul committed
95
} // namespace migraph
Paul's avatar
Paul committed
96
97

#endif