lowering.cpp 15.1 KB
Newer Older
1
2
3
/*
 * The MIT License (MIT)
 *
Ted Themistokleous's avatar
Ted Themistokleous committed
4
 * Copyright (c) 2015-2023 Advanced Micro Devices, Inc. All rights reserved.
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
 *
 * 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.
 */
Shucai Xiao's avatar
Shucai Xiao committed
24
#include <iterator>
25
26
27
28
29
#include <utility>
#include <functional>
#include <algorithm>
#include <map>

Paul's avatar
Paul committed
30
31
#include <migraphx/manage_ptr.hpp>
#include <migraphx/instruction.hpp>
32
#include <migraphx/make_op.hpp>
33
34
#include <migraphx/instruction_ref.hpp>
#include <migraphx/stringutils.hpp>
35
36
37
#include <migraphx/pass_manager.hpp>
#include <migraphx/iterator_for.hpp>
#include <migraphx/program.hpp>
38
39

#include <migraphx/op/dot.hpp>
Shucai Xiao's avatar
Shucai Xiao committed
40
#include <migraphx/op/if_op.hpp>
41
42
#include <migraphx/op/reshape.hpp>
#include <migraphx/op/quant_dot.hpp>
Ted Themistokleous's avatar
Ted Themistokleous committed
43
#include <migraphx/op/reshape_lazy.hpp>
44

Paul's avatar
Paul committed
45
#include <migraphx/gpu/context.hpp>
46
#include <migraphx/gpu/lowering.hpp>
47
#include <migraphx/gpu/device_name.hpp>
Paul's avatar
Paul committed
48
#include <migraphx/gpu/gemm.hpp>
49
50
#include <migraphx/gpu/miopen.hpp>
#include <migraphx/gpu/rocblas.hpp>
51
#include <migraphx/gpu/compiler.hpp>
Paul's avatar
Paul committed
52

Paul's avatar
Paul committed
53
namespace migraphx {
Paul's avatar
Paul committed
54
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
55
namespace gpu {
Paul's avatar
Paul committed
56
57
58

struct miopen_apply
{
59
60
61
    module* mod              = nullptr;
    module_pass_manager* mpm = nullptr;
    const lowering* pass     = nullptr;
Shucai Xiao's avatar
Shucai Xiao committed
62
    std::unordered_map<std::string, std::function<instruction_ref(instruction_ref)>> apply_map{};
Shucai Xiao's avatar
Shucai Xiao committed
63
    instruction_ref last{};
64
65
    bool offload_copy = false;
    bool compute_fp32 = false;
Paul's avatar
Paul committed
66

67
    context& get_context() const
68
69
70
71
72
73
    {
        assert(pass != nullptr);
        assert(pass->ctx != nullptr);
        return *pass->ctx;
    }

Paul's avatar
Paul committed
74
75
76
77
78
79
80
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

81
82
    void init()
    {
Shucai Xiao's avatar
Shucai Xiao committed
83
        assert(mod != nullptr);
84
        assert(pass != nullptr);
85

86
87
        compute_fp32 = get_compute_fp32_flag();
        offload_copy = (mod == mpm->get_root_module()) ? pass->offload_copy : false;
Paul's avatar
Paul committed
88

89
90
91
92
        add_generic_op("contiguous");
        add_extend_op("argmax");
        add_extend_op("argmin");
        add_extend_op("logsoftmax");
Shucai Xiao's avatar
Shucai Xiao committed
93
        add_extend_op("lrn");
turneram's avatar
turneram committed
94
        add_extend_op("multinomial");
Shucai Xiao's avatar
Shucai Xiao committed
95
        add_extend_op("nonzero");
96
        add_extend_op("pooling");
97
        add_extend_op("prefix_scan_sum");
Cagri Eryilmaz's avatar
Cagri Eryilmaz committed
98
        add_extend_op("reverse");
99
100
101
        add_extend_op("rnn_var_sl_last_output");
        add_extend_op("rnn_var_sl_shift_output");
        add_extend_op("rnn_var_sl_shift_sequence");
102
        add_extend_op("scatter_none");
Shucai Xiao's avatar
Shucai Xiao committed
103
        add_extend_op("topk");
104

105
        add_convolution_op("convolution");
106
        add_convolution_op("convolution_backwards");
107
        add_convolution_op("quant_convolution");
Shucai Xiao's avatar
Shucai Xiao committed
108
109
        add_gemm_op<op::dot>("dot");
        add_gemm_op<op::quant_dot>("quant_dot");
Shucai Xiao's avatar
Shucai Xiao committed
110
        add_if_op();
Shucai Xiao's avatar
Shucai Xiao committed
111
        add_loop_op();
Shucai Xiao's avatar
Shucai Xiao committed
112
        add_neg_op();
113
        add_nms_op();
Charlie Lin's avatar
Charlie Lin committed
114
        add_select_module_op();
Ted Themistokleous's avatar
Ted Themistokleous committed
115
        add_reshape_lazy_op();
116
117
    }

118
    void copy_params() const
119
    {
Shucai Xiao's avatar
Shucai Xiao committed
120
        if(not offload_copy)
121
            return;
122

Shucai Xiao's avatar
Shucai Xiao committed
123
        for(auto ins : iterator_for(*mod))
124
125
126
        {
            if(ins->name() != "@param")
                continue;
127

Shucai Xiao's avatar
Shucai Xiao committed
128
129
130
131
            // parameter no outputs, no need to insert copy to gpu
            if(ins->outputs().empty())
                continue;

132
133
            auto pos = std::next(ins);
            auto a   = insert_allocation(pos, ins->get_shape());
134
            auto c   = mod->insert_instruction(pos, make_op("hip::copy_to_gpu"), ins, a);
Shucai Xiao's avatar
Shucai Xiao committed
135
            mod->replace_instruction(ins, c);
136
        }
137
138

        // return instruction
Shucai Xiao's avatar
Shucai Xiao committed
139
        auto ret = std::prev(mod->end());
140
141
        if(ret->name() == "@return")
        {
142
            const auto& inputs = ret->inputs();
143
144
145

            // each input of ret need to be copied from gpu to host, and replace
            // output with copy output
146
            for(const auto& in : inputs)
147
            {
148
                auto p_output = mod->insert_instruction(ret, make_op("hip::copy_from_gpu"), in);
149
150
151
152
153
154
                instruction::replace_argument(ret, in, p_output);
            }
        }
        // else branch to handle legacy program without the return instruction
        else
        {
155
            mod->add_instruction(make_op("hip::copy_from_gpu"), ret);
156
        }
157
158
    }

Paul's avatar
Paul committed
159
160
    void apply()
    {
161
        init();
Shucai Xiao's avatar
Shucai Xiao committed
162
        for(auto it = mod->begin(); it != mod->end(); it++)
Paul's avatar
Paul committed
163
        {
164
165
            auto s     = it->get_shape();
            auto attrs = it->get_operator().attributes();
166
            if(apply_map.count(it->name()) > 0)
167
            {
168
                check_shape(s, apply_map.at(it->name())(it));
Paul's avatar
Paul committed
169
            }
170
171
172
173
            else if(has_compiler_for(it->name()))
            {
                check_shape(s, insert_precompile_op(it));
            }
174
175
176
177
            else if(attrs.contains("target"))
            {
                check_shape(s, insert_custom_op(it, attrs));
            }
Paul's avatar
Paul committed
178
        }
179
        copy_params();
Paul's avatar
Paul committed
180
181
    }

182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
    instruction_ref insert_custom_op(instruction_ref ins, const value& attrs) const
    {
        const auto& custom_op = ins->get_operator();
        if(attrs.at("target") == "cpu")
        {
            auto s = ins->get_shape();
            std::vector<instruction_ref> cpu_inputs;
            auto inputs = ins->inputs();
            auto output = inputs.back();
            std::transform(
                inputs.begin(), inputs.end(), std::back_inserter(cpu_inputs), [&](auto in) {
                    return mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), in);
                });
            cpu_inputs.front() =
                mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_inputs);
            auto cpu_out = mod->insert_instruction(ins, custom_op, cpu_inputs);
            auto gpu_out =
                mod->insert_instruction(ins, make_op("hip::copy_to_gpu"), cpu_out, output);
            return mod->replace_instruction(ins, gpu_out);
        }
        return ins;
    }

205
    instruction_ref insert_precompile_op(instruction_ref ins) const
206
207
208
209
210
211
212
213
214
215
216
217
    {
        auto output                       = insert_allocation(ins, ins->get_shape());
        std::vector<instruction_ref> refs = ins->inputs();
        refs.push_back(output);

        return mod->replace_instruction(
            ins,
            make_op("gpu::precompile_op", {{"op", to_value(ins->get_operator())}}),
            refs,
            ins->module_inputs());
    }

218
    instruction_ref insert_allocation(instruction_ref ins, const shape& s) const
Paul's avatar
Paul committed
219
    {
220
        return mod->insert_instruction(ins, make_op("allocate", {{"shape", to_value(s)}}));
Paul's avatar
Paul committed
221
222
    }

223
224
    template <typename Op>
    void add_gemm_op(const std::string& name)
225
226
    {
        apply_map.emplace(name, [=](instruction_ref ins) {
227
            std::vector<instruction_ref> refs = ins->inputs();
228
229
230
            assert(refs.size() == 2);
            auto output = insert_allocation(ins, ins->get_shape());
            refs.push_back(output);
231
            return mod->replace_instruction(ins, rocblas_gemm<Op>{Op{}, 1, 0, compute_fp32}, refs);
232
233
234
        });
    }

235
    void add_convolution_op(const std::string& name)
236
    {
237
        apply_map.emplace(name, [=](instruction_ref ins) {
238
239
            operation conv = make_op("gpu::" + name, {{"op", ins->get_operator().to_value()}});
            auto output    = insert_allocation(ins, ins->get_shape());
240

241
242
243
244
245
            return mod->replace_instruction(ins,
                                            make_op("gpu::miopen_op", {{"op", to_value(conv)}}),
                                            ins->inputs().at(0),
                                            ins->inputs().at(1),
                                            output);
Shucai Xiao's avatar
Shucai Xiao committed
246
247
248
        });
    }

249
250
251
    // add_generic_op just constructs the operator with no fields whereas add_extend_op copies over
    // the fields Since it doesn't have fields its default constructed

252
253
254
    void add_generic_op(const std::string& name) { add_generic_op(name, "gpu::" + name); }

    void add_generic_op(const std::string& op_name, const std::string& gpu_name)
Paul's avatar
Paul committed
255
    {
256
        apply_map.emplace(op_name, [=](instruction_ref ins) {
257
258
259
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
260

Shucai Xiao's avatar
Shucai Xiao committed
261
            return mod->replace_instruction(ins, make_op(gpu_name), refs);
262
        });
Paul's avatar
Paul committed
263
    }
Paul's avatar
Paul committed
264

265
266
267
    void add_extend_op(const std::string& name) { add_extend_op(name, "gpu::" + name); }

    void add_extend_op(const std::string& op_name, const std::string& gpu_name)
Khalique's avatar
Khalique committed
268
    {
269
270
        apply_map.emplace(op_name, [=](instruction_ref ins) {
            auto&& op                         = ins->get_operator();
271
272
273
            auto output                       = insert_allocation(ins, ins->get_shape());
            std::vector<instruction_ref> refs = ins->inputs();
            refs.push_back(output);
Paul's avatar
Paul committed
274

Shucai Xiao's avatar
Shucai Xiao committed
275
            return mod->replace_instruction(ins, make_op(gpu_name, op.to_value()), refs);
276
        });
Khalique's avatar
Khalique committed
277
278
    }

Shucai Xiao's avatar
Shucai Xiao committed
279
280
281
282
283
284
    // use 0 - input to represent neg
    void add_neg_op()
    {
        apply_map.emplace("neg", [=](instruction_ref ins) {
            auto s = ins->get_shape();
            std::vector<float> zeros(s.elements(), 0.0f);
Shucai Xiao's avatar
Shucai Xiao committed
285
            auto l0     = mod->add_literal(literal(s, zeros));
Shucai Xiao's avatar
Shucai Xiao committed
286
            auto output = insert_allocation(ins, s);
Shucai Xiao's avatar
Shucai Xiao committed
287
            return mod->replace_instruction(
288
                ins, make_op("gpu::sub"), l0, ins->inputs().front(), output);
Shucai Xiao's avatar
Shucai Xiao committed
289
290
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
291

Shucai Xiao's avatar
Shucai Xiao committed
292
    // add input and output argument for the if operator
Shucai Xiao's avatar
Shucai Xiao committed
293
294
295
296
    void add_if_op()
    {
        apply_map.emplace("if", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
297
298
299
            auto cpu_cond =
                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.front());
            auto sync_cond = mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_cond);
Shucai Xiao's avatar
Shucai Xiao committed
300
301
            inputs.front() = sync_cond;

302
            return mod->replace_instruction(ins, ins->get_operator(), inputs, ins->module_inputs());
Shucai Xiao's avatar
Shucai Xiao committed
303
304
        });
    }
Shucai Xiao's avatar
Shucai Xiao committed
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320

    // replace the loop operator with gpu_loop operator
    void add_loop_op()
    {
        apply_map.emplace("loop", [=](instruction_ref ins) {
            std::vector<instruction_ref> inputs = ins->inputs();
            // copy max_iter from gpu to cpu
            auto cpu_max_iter =
                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.at(0));
            auto cpu_cond =
                mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), inputs.at(1));
            auto synced_max_iter =
                mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_max_iter, cpu_cond);
            inputs.at(0)     = synced_max_iter;
            inputs.at(1)     = cpu_cond;
            auto copy_inputs = inputs;
321
322
323
324
            std::transform(copy_inputs.begin(),
                           copy_inputs.end(),
                           std::back_inserter(inputs),
                           [&](auto in) { return insert_allocation(ins, in->get_shape()); });
Shucai Xiao's avatar
Shucai Xiao committed
325
326
327
328
329

            auto mod_args = ins->module_inputs();
            auto output   = insert_allocation(ins, ins->get_shape());

            const auto* sub_mod = mod_args.front();
330
331
            auto cond_out       = insert_allocation(ins, sub_mod->get_output_shapes().front());

Shucai Xiao's avatar
Shucai Xiao committed
332
333
334
335
336
337
338
339
            // add cond and mod outputs to the argument list
            inputs.push_back(cond_out);
            inputs.push_back(output);

            return mod->replace_instruction(
                ins, make_op("gpu::loop", ins->get_operator().to_value()), inputs, mod_args);
        });
    }
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359

    void add_nms_op()
    {
        apply_map.emplace("nonmaxsuppression", [=](instruction_ref ins) {
            auto s      = ins->get_shape();
            auto output = insert_allocation(ins, s);
            std::vector<instruction_ref> cpu_inputs;
            auto inputs = ins->inputs();
            std::transform(
                inputs.begin(), inputs.end(), std::back_inserter(cpu_inputs), [&](auto in) {
                    return mod->insert_instruction(ins, make_op("hip::copy_from_gpu"), in);
                });
            cpu_inputs.front() =
                mod->insert_instruction(ins, make_op("hip::sync_stream"), cpu_inputs);
            auto cpu_out = mod->insert_instruction(ins, ins->get_operator(), cpu_inputs);
            auto gpu_out =
                mod->insert_instruction(ins, make_op("hip::copy_to_gpu"), cpu_out, output);
            return mod->replace_instruction(ins, gpu_out);
        });
    }
Charlie Lin's avatar
Charlie Lin committed
360
361

    /**
Charlie Lin's avatar
Charlie Lin committed
362
     * Adds dynamic allocation for submodule output parameter.
Charlie Lin's avatar
Charlie Lin committed
363
364
365
366
     */
    void add_select_module_op()
    {
        apply_map.emplace("select_module", [=](instruction_ref ins) {
Charlie Lin's avatar
Charlie Lin committed
367
368
            auto s                              = ins->get_shape();
            auto output                         = insert_allocation(ins, s);
Charlie Lin's avatar
Charlie Lin committed
369
            std::vector<instruction_ref> inputs = ins->inputs();
Charlie Lin's avatar
Charlie Lin committed
370
371
            inputs.push_back(output);
            return mod->replace_instruction(ins, ins->get_operator(), inputs, ins->module_inputs());
Charlie Lin's avatar
Charlie Lin committed
372
373
        });
    }
Ted Themistokleous's avatar
Ted Themistokleous committed
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399

    /**
     *  Adds reshape lazy to reshape ops that can be aliased instead of copied.
     *  `gpu::contiguous` are added before and after the reshape; these contiguous
     *  instructions can be removed by the eliminate_contiguous pass.
     */
    void add_reshape_lazy_op()
    {
        apply_map.emplace("reshape", [=](instruction_ref ins) {
            std::vector<instruction_ref> before_contiguous_args = ins->inputs();
            auto before_alloc = insert_allocation(ins, std::prev(ins)->get_shape());
            before_contiguous_args.push_back(before_alloc);
            auto before_contig =
                mod->insert_instruction(ins, make_op("gpu::contiguous"), {before_contiguous_args});

            auto new_lazy_reshape = mod->insert_instruction(
                ins,
                make_op("reshape_lazy", {{"dims", {ins->get_operator().to_value().at("dims")}}}),
                before_contig);

            std::vector<instruction_ref> after_contiguous_args = {new_lazy_reshape};
            auto after_alloc = insert_allocation(new_lazy_reshape, new_lazy_reshape->get_shape());
            after_contiguous_args.push_back(after_alloc);
            return mod->replace_instruction(ins, make_op("gpu::contiguous"), after_contiguous_args);
        });
    }
Paul's avatar
Paul committed
400
401
};

402
403
404
405
void lowering::apply(module_pass_manager& mpm) const
{
    miopen_apply{&mpm.get_module(), &mpm, this}.apply();
}
Shucai Xiao's avatar
Shucai Xiao committed
406

Paul's avatar
Paul committed
407
} // namespace gpu
Paul's avatar
Paul committed
408
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
409
} // namespace migraphx