compile_ops.cpp 9.49 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
        }
    }
    const compiled_result& benchmark(problem_cache& pc) const
    {
171
        const auto trace_level = value_of(MIGRAPHX_TRACE_BENCHMARKING{});
172
173
174
        if(results.empty())
            MIGRAPHX_THROW("No configs to tune");
        if(results.size() == 1)
175
176
177
178
179
        {
            if(not results.front().has_value())
                MIGRAPHX_THROW("No configs to tune");
            return *results.front();
        }
180
181
        if(not config)
            MIGRAPHX_THROW("Multiple kernels without config");
182
183
184
185
        if(trace_level > 0)
            std::cout << "Benchmarking " << preop.name() << ": " << results.size() << " configs"
                      << std::endl;
        if(trace_level > 1)
Paul Fultz II's avatar
Paul Fultz II committed
186
            std::cout << "Problem: " << config->problem << std::endl;
187
188
        std::vector<double> times;
        times.reserve(results.size());
Paul Fultz II's avatar
Paul Fultz II committed
189
190
191
192
193
        std::transform(results.begin(),
                       results.end(),
                       config->solutions.begin(),
                       std::back_inserter(times),
                       [&](const auto& cr, const auto& solution) {
194
                           if(trace_level > 1)
Paul Fultz II's avatar
Paul Fultz II committed
195
196
197
                               std::cout << "Benchmarking solution: " << solution << std::endl;
                           if(not cr.has_value())
                           {
198
                               if(trace_level > 1)
Paul Fultz II's avatar
Paul Fultz II committed
199
200
201
202
203
                                   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);
204
                           if(trace_level > 1)
Paul Fultz II's avatar
Paul Fultz II committed
205
206
207
                               std::cout << t << "ms" << std::endl;
                           return t;
                       });
208
        auto i = std::distance(times.begin(), std::min_element(times.begin(), times.end()));
209
210
        if(trace_level > 0)
            std::cout << "Fastest solution: " << config->solutions.at(i) << std::endl;
211
        pc.insert(preop.name(), config->problem, config->solutions.at(i));
212
213
214
        if(not results[i].has_value())
            MIGRAPHX_THROW("No valid tuned compilation.");
        return *results[i];
215
216
217
218
219
220
221
222
    }
    void replace(module& m, problem_cache& pc) const
    {
        const auto& cr = benchmark(pc);
        cr.replace.replace(m, cr.ins);
    }
};

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

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

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

    void update_configs()
    {
248
        par_compile(cps.size(), [&](auto i) { cps[i].update_config(exhaustive); });
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
278
279
280
    }

    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
281
282
283
284
285
    for(auto ins : iterator_for(m))
    {
        if(ins->name() != "gpu::precompile_op")
            continue;
        operation preop = any_cast<precompile_op>(ins->get_operator()).op;
286
        cm.add_plan(ctx, preop, ins);
287
    }
288
289
290
291
292
    cm.update_configs();
    cm.compile(m);
    // Compile already tuned configs
    cm.compile(m);
    assert(cm.cps.empty());
293
294
295
296
297
298
}

} // namespace gpu

} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx