compile_ops.cpp 9.44 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>
33
#include <migraphx/gpu/time_op.hpp>
34
35
36
37
38

namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
namespace gpu {

39
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_GPU_COMPILE_PARALLEL);
Paul Fultz II's avatar
Paul Fultz II committed
40
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TRACE_BENCHMARKING);
41

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

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

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

    shape compute_shape(std::vector<shape> inputs, const std::vector<module_ref>& mods) const
    {
60
61
62
63
        // Pop off additional args
        inputs.resize(inputs.size() - additional_args);
        if(ignore_modules)
            return op.compute_shape(inputs);
64
65
66
67
68
69
70
71
72
73
74
        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
75
76
struct compiled_result
{
77
    compiler_replace replace;
Paul Fultz II's avatar
Paul Fultz II committed
78
79
80
    instruction_ref ins;
};

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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)
    {
        assert(not solution.is_null());
        cache[create_key(name, problem)] = solution;
    }
    void mark(const std::string& name, const value& problem)
    {
        cache.insert(std::make_pair(create_key(name, problem), value{}));
    }
    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;
};

struct compile_plan
{
    context* ctx;
    operation preop;
    instruction_ref ins;
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
    optional<tuning_config> config                 = nullopt;
    std::vector<optional<compiled_result>> results = {};
    void update_config(bool exhaustive)
    {
        config = get_tuning_config(*ctx, ins, preop, exhaustive);
    }
    template <class Vector>
    void insert_compiles(Vector& compiles, const value& solution, std::size_t i)
    {
        compiles.emplace_back([=] {
            try
            {
                results[i] = compiled_result{compile(*ctx, ins, preop, solution), ins};
            }
            catch(...)
            {
                results[i] = nullopt;
            }
        });
    }

136
137
138
139
140
141
142
143
144
145
146
147
148
    template <class Vector>
    void add_compiles(Vector& compiles, problem_cache& pc)
    {
        if(config.has_value())
        {
            const auto& problem = config->problem;
            if(auto sol = pc.get(preop.name(), problem))
            {
                auto solution = sol.value();
                // No solution yet until benchmarked so skip for now
                if(solution.is_null())
                    return;
                results.resize(1);
149
                insert_compiles(compiles, solution, 0);
150
151
152
153
154
155
156
157
158
            }
            else
            {
                pc.mark(preop.name(), problem);
                const auto& solutions = config->solutions;
                results.resize(solutions.size());
                for(auto i : range(solutions.size()))
                {
                    auto solution = solutions[i];
159
                    insert_compiles(compiles, solution, i);
160
161
162
163
164
165
                }
            }
        }
        else
        {
            results.resize(1);
166
            insert_compiles(compiles, value{}, 0);
167
168
169
170
171
172
173
        }
    }
    const compiled_result& benchmark(problem_cache& pc) const
    {
        if(results.empty())
            MIGRAPHX_THROW("No configs to tune");
        if(results.size() == 1)
174
175
176
177
178
        {
            if(not results.front().has_value())
                MIGRAPHX_THROW("No configs to tune");
            return *results.front();
        }
179
180
181
182
        if(not config)
            MIGRAPHX_THROW("Multiple kernels without config");
        std::cout << "Benchmarking " << preop.name() << ": " << results.size() << " configs"
                  << std::endl;
Paul Fultz II's avatar
Paul Fultz II committed
183
184
        if(enabled(MIGRAPHX_TRACE_BENCHMARKING{}))
            std::cout << "Problem: " << config->problem << std::endl;
185
186
        std::vector<double> times;
        times.reserve(results.size());
Paul Fultz II's avatar
Paul Fultz II committed
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
        std::transform(results.begin(),
                       results.end(),
                       config->solutions.begin(),
                       std::back_inserter(times),
                       [&](const auto& cr, const auto& solution) {
                           if(enabled(MIGRAPHX_TRACE_BENCHMARKING{}))
                               std::cout << "Benchmarking solution: " << solution << std::endl;
                           if(not cr.has_value())
                           {
                               if(enabled(MIGRAPHX_TRACE_BENCHMARKING{}))
                                   std::cout << "No binary" << std::endl;
                               return std::numeric_limits<double>::max();
                           }
                           auto t = time_op(
                               *ctx, cr->replace.code_object, to_shapes(cr->ins->inputs()), 20);
                           if(enabled(MIGRAPHX_TRACE_BENCHMARKING{}))
                               std::cout << t << "ms" << std::endl;
                           return t;
                       });
206
        auto i = std::distance(times.begin(), std::min_element(times.begin(), times.end()));
207
        std::cout << "Fastest solution: " << config->solutions.at(i) << std::endl;
208
        pc.insert(preop.name(), config->problem, config->solutions.at(i));
209
210
211
        if(not results[i].has_value())
            MIGRAPHX_THROW("No valid tuned compilation.");
        return *results[i];
212
213
214
215
216
217
218
219
    }
    void replace(module& m, problem_cache& pc) const
    {
        const auto& cr = benchmark(pc);
        cr.replace.replace(m, cr.ins);
    }
};

220
221
222
223
224
template <class F>
void par_compile(std::size_t n, F f)
{
    if(n == 0)
        return;
225
226
227
228
    auto d = value_of(MIGRAPHX_GPU_COMPILE_PARALLEL{});
    if(d == 0)
        d = n;
    par_for(n, n / d, f);
229
230
}

231
struct compile_manager
232
{
233
234
235
    problem_cache pc;
    std::vector<compile_plan> cps;
    bool exhaustive = false;
Paul Fultz II's avatar
Paul Fultz II committed
236

237
238
239
240
241
242
243
244
    template <class... Ts>
    void add_plan(Ts&&... xs)
    {
        cps.push_back({std::forward<Ts>(xs)...});
    }

    void update_configs()
    {
245
        par_compile(cps.size(), [&](auto i) { cps[i].update_config(exhaustive); });
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
    }

    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());
    }
};

void compile_ops::apply(module& m) const
{
    compile_manager cm;
    cm.exhaustive = exhaustive_tune;
    // Find all precompile opes
278
279
280
281
282
    for(auto ins : iterator_for(m))
    {
        if(ins->name() != "gpu::precompile_op")
            continue;
        operation preop = any_cast<precompile_op>(ins->get_operator()).op;
283
        cm.add_plan(ctx, preop, ins);
284
    }
285
286
287
288
289
    cm.update_configs();
    cm.compile(m);
    // Compile already tuned configs
    cm.compile(m);
    assert(cm.cps.empty());
290
291
292
293
294
295
}

} // namespace gpu

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx