program.cpp 6.2 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
#include <migraph/program.hpp>
#include <migraph/stringutils.hpp>
#include <migraph/instruction.hpp>
Paul's avatar
Paul committed
4
#include <iostream>
Paul's avatar
Paul committed
5
6
#include <algorithm>

Paul's avatar
Paul committed
7
namespace migraph {
Paul's avatar
Paul committed
8

Paul's avatar
Paul committed
9
10
11
12
struct program_impl
{
    // A list is used to keep references to an instruction stable
    std::list<instruction> instructions;
Paul's avatar
Paul committed
13
    context ctx;
Paul's avatar
Paul committed
14
15
};

Paul's avatar
Paul committed
16
const operation& get_operation(instruction_ref ins) { return ins->op; }
Paul's avatar
Paul committed
17

Paul's avatar
Paul committed
18
program::program() : impl(std::make_unique<program_impl>()) {}
Paul's avatar
Paul committed
19

Paul's avatar
Paul committed
20
program::program(program&&) noexcept = default;
Paul's avatar
Paul committed
21
22
program& program::operator=(program&&) noexcept = default;
program::~program() noexcept                    = default;
Paul's avatar
Paul committed
23

Paul's avatar
Paul committed
24
instruction_ref program::add_instruction(operation op, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
25
26
27
{
    return insert_instruction(impl->instructions.end(), std::move(op), std::move(args));
}
Paul's avatar
Paul committed
28
29
instruction_ref
program::insert_instruction(instruction_ref ins, operation op, std::vector<instruction_ref> args)
Paul's avatar
Paul committed
30
{
Paul's avatar
Paul committed
31
32
33
    assert(std::all_of(
               args.begin(), args.end(), [&](instruction_ref x) { return has_instruction(x); }) &&
           "Argument is not an exisiting instruction");
Paul's avatar
Paul committed
34
35
    // TODO: Use move
    shape r     = compute_shape(op, args);
Paul's avatar
Paul committed
36
    auto result = impl->instructions.insert(ins, {op, r, args});
Paul's avatar
Paul committed
37
    backreference(result);
Paul's avatar
Paul committed
38
    assert(result->arguments == args);
Paul's avatar
Paul committed
39
    return result;
Paul's avatar
Paul committed
40
41
}

Paul's avatar
Paul committed
42
43
44
45
46
47
48
instruction_ref
program::replace_instruction(instruction_ref ins, operation op, std::vector<instruction_ref> args)
{
    assert(std::all_of(
               args.begin(), args.end(), [&](instruction_ref x) { return has_instruction(x); }) &&
           "Argument is not an exisiting instruction");

Paul's avatar
Paul committed
49
    shape r = compute_shape(op, args);
Paul's avatar
Paul committed
50
51
52
53
54
    ins->replace(op, r, args);
    backreference(ins);
    return ins;
}

Paul's avatar
Paul committed
55
56
57
58
59
60
61
62
63
instruction_ref
program::remove_instruction(instruction_ref ins)
{
    assert(has_instruction(ins));
    assert(ins->output.empty());
    ins->clear_arguments();
    return impl->instructions.erase(ins);
}

Paul's avatar
Paul committed
64
instruction_ref program::add_literal(literal l)
Paul's avatar
Paul committed
65
{
Paul's avatar
Paul committed
66
67
68
69
70
71
72
73
    impl->instructions.emplace_front(std::move(l));
    return impl->instructions.begin();
}

instruction_ref program::add_outline(shape s)
{
    impl->instructions.push_front({builtin::outline{s}, s, {}});
    return impl->instructions.begin();
Paul's avatar
Paul committed
74
75
}

Paul's avatar
Paul committed
76
instruction_ref program::add_parameter(std::string name, shape s)
Paul's avatar
Paul committed
77
{
Paul's avatar
Paul committed
78
79
80
81
82
83
84
    impl->instructions.push_front({builtin::param{std::move(name)}, s, {}});
    return impl->instructions.begin();
}

shape program::get_parameter_shape(std::string name)
{
    auto ins = std::find_if(
Paul's avatar
Paul committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
        impl->instructions.begin(), impl->instructions.end(), [&](const instruction& x) {
            if(x.op.name() == "@param")
            {
                return any_cast<builtin::param>(x.op).parameter == name;
            }
            else
            {
                return false;
            }
        });
    if(ins != this->end())
        return ins->result;
    else
        return {};
Paul's avatar
Paul committed
99
100
}

Paul's avatar
Paul committed
101
bool program::has_instruction(instruction_ref ins) const
Paul's avatar
Paul committed
102
{
Paul's avatar
Paul committed
103
104
105
106
    return std::find_if(
               impl->instructions.begin(), impl->instructions.end(), [&](const instruction& x) {
                   return std::addressof(*ins) == std::addressof(x);
               }) != impl->instructions.end();
Paul's avatar
Paul committed
107
108
}

Paul's avatar
Paul committed
109
110
instruction_ref program::begin() { return impl->instructions.begin(); }
instruction_ref program::end() { return impl->instructions.end(); }
111

Paul's avatar
Paul committed
112
113
instruction_ref program::validate() const
{
Paul's avatar
Paul committed
114
115
116
    return std::find_if(impl->instructions.begin(),
                        impl->instructions.end(),
                        [](const instruction& i) { return i.valid(); });
Paul's avatar
Paul committed
117
118
}

Paul's avatar
Paul committed
119
120
121
void program::compile(const target& t)
{
    assert(this->validate() != impl->instructions.end());
Paul's avatar
Paul committed
122
    this->impl->ctx = t.get_context();
Paul's avatar
Paul committed
123
    for(auto&& p : t.get_passes(this->impl->ctx))
Paul's avatar
Paul committed
124
125
126
127
128
129
130
    {
        p.apply(*this);
#ifndef NDEBUG
        if(this->validate() == impl->instructions.end())
            MIGRAPH_THROW(p.name() + " pass produces invalid program");
#endif
    }
Paul's avatar
Paul committed
131
    if(this->validate() == impl->instructions.end())
Paul's avatar
Paul committed
132
        MIGRAPH_THROW("Invalid program from compilation");
Paul's avatar
Paul committed
133
134
}

135
argument program::eval(std::unordered_map<std::string, argument> params) const
Paul's avatar
Paul committed
136
{
Paul's avatar
Paul committed
137
    assert(this->validate() != impl->instructions.end());
Paul's avatar
Paul committed
138
139
    std::unordered_map<const instruction*, argument> results;
    argument result;
Paul's avatar
Paul committed
140
    for(auto& ins : impl->instructions)
Paul's avatar
Paul committed
141
    {
Paul's avatar
Paul committed
142
        if(ins.op.name() == "@literal")
Paul's avatar
Paul committed
143
144
145
        {
            result = ins.lit.get_argument();
        }
Paul's avatar
Paul committed
146
        else if(ins.op.name() == "@param")
Paul's avatar
Paul committed
147
        {
Paul's avatar
Paul committed
148
            result = params.at(any_cast<builtin::param>(ins.op).parameter);
Paul's avatar
Paul committed
149
        }
Paul's avatar
Paul committed
150
151
152
153
        else if(ins.op.name() == "@outline")
        {
            result = argument{ins.result, nullptr};
        }
Paul's avatar
Paul committed
154
155
156
        else
        {
            std::vector<argument> values(ins.arguments.size());
Paul's avatar
Paul committed
157
158
159
            std::transform(ins.arguments.begin(),
                           ins.arguments.end(),
                           values.begin(),
Paul's avatar
Paul committed
160
                           [&](instruction_ref i) { return results.at(std::addressof(*i)); });
Paul's avatar
Paul committed
161
            result = ins.op.compute(this->impl->ctx, ins.result, values);
Paul's avatar
Paul committed
162
163
164
        }
        results.emplace(std::addressof(ins), result);
    }
165
    return result;
Paul's avatar
Paul committed
166
167
}

Paul's avatar
Paul committed
168
std::ostream& operator<<(std::ostream& os, const program& p)
Paul's avatar
Paul committed
169
170
171
172
{
    std::unordered_map<const instruction*, std::string> names;
    int count = 0;

Paul's avatar
Paul committed
173
    for(auto& ins : p.impl->instructions)
Paul's avatar
Paul committed
174
175
    {
        std::string var_name = "@" + std::to_string(count);
Paul's avatar
Paul committed
176
        if(ins.op.name() == "@param")
Paul's avatar
Paul committed
177
        {
Paul's avatar
Paul committed
178
            var_name = any_cast<builtin::param>(ins.op).parameter;
Paul's avatar
Paul committed
179
180
        }

Paul's avatar
Paul committed
181
        os << var_name << " = ";
Paul's avatar
Paul committed
182

Paul's avatar
Paul committed
183
        os << ins.op;
Paul's avatar
Paul committed
184
185
186

        if(ins.op.name() == "@literal")
        {
Paul's avatar
Paul committed
187
            if(ins.lit.get_shape().elements() > 10)
Paul's avatar
Paul committed
188
                os << "{ ... }";
Paul's avatar
Paul committed
189
            else
Paul's avatar
Paul committed
190
                os << "{" << ins.lit << "}";
Paul's avatar
Paul committed
191
192
193
194
195
        }

        if(!ins.arguments.empty())
        {
            char delim = '(';
Paul's avatar
Paul committed
196
            for(auto&& arg : ins.arguments)
Paul's avatar
Paul committed
197
            {
Paul's avatar
Paul committed
198
199
                assert(p.has_instruction(arg) && "Instruction not found");
                os << delim << names.at(std::addressof(*arg));
Paul's avatar
Paul committed
200
201
                delim = ',';
            }
Paul's avatar
Paul committed
202
            os << ")";
Paul's avatar
Paul committed
203
204
        }

Paul's avatar
Paul committed
205
        os << " -> " << ins.result;
Paul's avatar
Paul committed
206

Paul's avatar
Paul committed
207
        os << std::endl;
Paul's avatar
Paul committed
208
209

        names.emplace(std::addressof(ins), var_name);
Paul's avatar
Paul committed
210
        count++;
Paul's avatar
Paul committed
211
    }
Paul's avatar
Paul committed
212
    return os;
Paul's avatar
Paul committed
213
214
}

Paul's avatar
Paul committed
215
} // namespace migraph