compile_ops.cpp 7.42 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 */
24
25
26
27
28
#include <migraphx/gpu/compile_ops.hpp>
#include <migraphx/gpu/context.hpp>
#include <migraphx/module.hpp>
#include <migraphx/iterator_for.hpp>
#include <migraphx/instruction.hpp>
Paul Fultz II's avatar
Paul Fultz II committed
29
#include <migraphx/par_for.hpp>
30
31
#include <migraphx/register_op.hpp>
#include <migraphx/op/identity.hpp>
32
#include <migraphx/gpu/compiler.hpp>
Paul's avatar
Paul committed
33
#include <migraphx/gpu/time_op.hpp>
34
35
36
37
38

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {

39
40
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_GPU_COMPILE_PARALLEL);

41
42
struct precompile_op
{
43
44
45
    operation op                = op::identity{};
    std::size_t additional_args = 1;
    bool ignore_modules         = false;
46
47
48
49

    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
50
51
52
        return pack(f(self.op, "op"),
                    f(self.additional_args, "additional_args"),
                    f(self.ignore_modules, "ignore_modules"));
53
54
55
56
57
58
    }

    std::string name() const { return "gpu::precompile_op"; }

    shape compute_shape(std::vector<shape> inputs, const std::vector<module_ref>& mods) const
    {
59
60
61
62
        // Pop off additional args
        inputs.resize(inputs.size() - additional_args);
        if(ignore_modules)
            return op.compute_shape(inputs);
63
64
65
66
67
68
69
70
71
72
73
        return op.compute_shape(inputs, mods);
    }

    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
};

MIGRAPHX_REGISTER_OP(precompile_op);

Paul Fultz II's avatar
Paul Fultz II committed
74
75
struct compiled_result
{
76
    compiler_replace replace;
Paul Fultz II's avatar
Paul Fultz II committed
77
78
79
    instruction_ref ins;
};

80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
struct problem_cache
{
    bool has(const std::string& name, const value& problem) const
    {
        return contains(cache, create_key(name, problem));
    }
    void insert(const std::string& name, const value& problem, const value& solution)
    {
        cache[create_key(name, problem)] = solution;
    }
    optional<value> get(const std::string& name, const value& problem) const
    {
        auto it = cache.find(create_key(name, problem));
        if (it == cache.end())
            return nullopt;
        return it->second;
    }
    static value create_key(const std::string& name, const value& problem)
    {
        return {{"name", name}, {"problem", problem}};
    }
    std::unordered_map<value, value> cache;
};

104
105
106
107
108
struct compile_plan
{
    context* ctx;
    operation preop;
    instruction_ref ins;
Paul's avatar
Paul committed
109
    optional<tuning_config> config       = nullopt;
110
    std::vector<compiled_result> results = {};
Paul's avatar
Paul committed
111
112
    void update_config() { config = get_tuning_config(*ctx, ins, preop); }
    template <class Vector>
113
    void add_compiles(Vector& compiles, problem_cache& pc)
114
    {
Paul's avatar
Paul committed
115
        if(config.has_value())
116
        {
117
118
            const auto& problem = config.value().problem;
            if(auto sol = pc.get(preop.name(), problem))
119
            {
120
121
122
123
                auto solution = sol.value();
                // No solution yet until benchmarked so skip for now
                if(solution.empty())
                    return;
124
                compiles.emplace_back([=] {
125
                    results[0] = compiled_result{compile(*ctx, ins, preop, solution), ins};
126
127
                });
            }
128
129
130
131
132
133
134
135
136
137
138
139
            else
            {
                const auto& solutions = config.value().solutions;
                results.resize(solutions.size());
                for(auto i : range(solutions.size()))
                {
                    auto solution = solutions[i];
                    compiles.emplace_back([=] {
                        results[i] = compiled_result{compile(*ctx, ins, preop, solution), ins};
                    });
                }
            }
140
141
142
143
144
145
146
147
148
        }
        else
        {
            results.resize(1);
            compiles.emplace_back([=] {
                results[0] = compiled_result{compile(*ctx, ins, preop, value{}), ins};
            });
        }
    }
149
    const compiled_result& benchmark(problem_cache& pc) const
150
    {
Paul's avatar
Paul committed
151
152
        if(results.empty())
            MIGRAPHX_THROW("No configs to tune");
153
        if(results.size() == 1)
Paul's avatar
Paul committed
154
            return results.front();
Paul's avatar
Paul committed
155
156
        std::cout << "Benchmarking " << preop.name() << ": " << results.size() << " configs"
                  << std::endl;
Paul's avatar
Paul committed
157
        std::vector<double> times;
Paul's avatar
Paul committed
158
        for(const auto& cr : results)
159
        {
Paul's avatar
Paul committed
160
161
            times.push_back(
                time_op(*ctx, cr.replace.code_object, to_shapes(cr.ins->inputs()), 20).first);
162
        }
Paul's avatar
Paul committed
163
        auto i = std::distance(times.begin(), std::min_element(times.begin(), times.end()));
164
        pc.insert(preop.name(), config.value().problem, config.value().solutions[i]);
Paul's avatar
Paul committed
165
166
        return results[i];
    }
167
    void replace(module& m, problem_cache& pc) const
Paul's avatar
Paul committed
168
    {
169
        const auto& cr = benchmark(pc);
Paul's avatar
Paul committed
170
        cr.replace.replace(m, cr.ins);
171
172
173
    }
};

174
175
176
177
178
179
180
181
template <class F>
void par_compile(std::size_t n, F f)
{
    if(n == 0)
        return;
    par_for(n, n / value_of(MIGRAPHX_GPU_COMPILE_PARALLEL{}, n), f);
}

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
struct compile_manager
{
    problem_cache pc;
    std::vector<compile_plan> cps;

    template<class... Ts>
    void add_plan(Ts&&... xs)
    {
        cps.push_back({std::forward<Ts>(xs)...});
    }

    void update_configs()
    {
        par_compile(cps.size(), [&](auto i) { cps[i].update_config(); });
    }

    void compile(module& m)
    {
        std::vector<std::function<void()>> compiles;
        for(auto& cp : cps)
        {
            cp.add_compiles(compiles, pc);
        }
        par_compile(compiles.size(), [&](auto i) { compiles[i](); });

        // Replace and/or benchmark
        for(const auto& cp : cps)
        {
            if(cp.results.empty())
                continue;
            cp.replace(m, pc);
        }

        // Remove compile_plan already executed
        cps.erase(std::remove_if(cps.begin(), cps.end(), [](const auto& cp) {
            return not cp.results.empty();
        }), cps.end());
    }
};

222
223
void compile_ops::apply(module& m) const
{
224
225
    compile_manager cm;
    problem_cache pc;
226
    std::vector<compile_plan> cps;
Paul's avatar
Paul committed
227
    // Find all precompile opes
228
229
230
231
232
    for(auto ins : iterator_for(m))
    {
        if(ins->name() != "gpu::precompile_op")
            continue;
        operation preop = any_cast<precompile_op>(ins->get_operator()).op;
233
        cm.add_plan(ctx, preop, ins);
234
    }
235
236
237
238
    cm.update_configs();
    cm.compile(m);
    // Compile already tuned configs
    cm.compile(m);
239
240
241
242
243
244
}

} // namespace gpu

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx