fuse_ops.cpp 32.7 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.
 */
kahmed10's avatar
kahmed10 committed
24
25
#include <migraphx/pass_manager.hpp>
#include <migraphx/dead_code_elimination.hpp>
Paul's avatar
Paul committed
26
27
28
#include <migraphx/gpu/fuse_ops.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/gpu/miopen.hpp>
kahmed10's avatar
kahmed10 committed
29
#include <migraphx/gpu/clip.hpp>
Paul's avatar
Paul committed
30
#include <migraphx/gpu/convolution.hpp>
31
#include <migraphx/gpu/device_name.hpp>
32
#include <migraphx/gpu/oper.hpp>
kahmed10's avatar
kahmed10 committed
33
34
#include <migraphx/gpu/add.hpp>
#include <migraphx/gpu/mul.hpp>
35
#include <migraphx/gpu/gemm.hpp>
kahmed10's avatar
kahmed10 committed
36
#include <migraphx/gpu/device/layernorm.hpp>
kahmed10's avatar
kahmed10 committed
37
#include <migraphx/gpu/device/gelu.hpp>
Paul's avatar
Paul committed
38
#include <migraphx/gpu/device/mul_add.hpp>
39
40
41
42
43
#include <migraphx/gpu/device/add_clip.hpp>
#include <migraphx/gpu/device/add_relu.hpp>
#include <migraphx/gpu/device/add_sigmoid.hpp>
#include <migraphx/gpu/device/add_tanh.hpp>
#include <migraphx/gpu/device/mul_add_relu.hpp>
Paul's avatar
Paul committed
44
#include <migraphx/gpu/device/add.hpp>
45
46
47
#include <migraphx/match/layernorm.hpp>
#include <migraphx/match/gelu_erf.hpp>
#include <migraphx/match/gelu_tanh.hpp>
Paul's avatar
Paul committed
48
#include <migraphx/instruction.hpp>
49
#include <migraphx/register_op.hpp>
Paul's avatar
Paul committed
50
#include <migraphx/array.hpp>
kahmed10's avatar
kahmed10 committed
51
#include <migraphx/op/clip.hpp>
kahmed10's avatar
kahmed10 committed
52
#include <cmath>
53
#include <set>
Paul's avatar
Paul committed
54
55

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

59
60
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_MIOPEN_FUSION)

Paul's avatar
Paul committed
61
62
63
64
65
66
67
68
struct fusion
{
    using op_t = miopenFusionOpDescriptor_t;
    shared<fusion_plan_descriptor> fp;

    // Used as a temporary hack to keep descriptor references alive
    std::vector<std::shared_ptr<void>> storage;

Paul's avatar
Paul committed
69
    template <class T>
Paul's avatar
Paul committed
70
71
72
73
74
75
76
    auto keep_alive(T x)
    {
        auto result = share(std::move(x));
        storage.push_back(result);
        return result;
    }

77
78
    fusion() = default;

Paul's avatar
Paul committed
79
80
    fusion(const shape& input)
    {
81
        assert(input.standard());
Paul's avatar
Paul committed
82
        auto t = make_tensor(input);
Paul's avatar
Paul committed
83
        fp     = make_fusion_plan(t);
84
        assert(fp);
Paul's avatar
Paul committed
85
86
87
        keep_alive(std::move(t));
    }

88
89
    bool empty() const { return fp == nullptr; }

Paul's avatar
Paul committed
90
91
    op_t operator[](std::size_t i) const
    {
92
        assert(fp);
Paul's avatar
Paul committed
93
94
95
        op_t result;
        auto status = miopenFusionPlanGetOp(fp.get(), i, &result);
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
96
            MIGRAPHX_THROW("Failed retrieving operator at " + std::to_string(i));
Paul's avatar
Paul committed
97
98
99
        return result;
    }

100
101
102
103
104
    auto get() const
    {
        assert(fp);
        return fp.get();
    }
Paul's avatar
Paul committed
105
106
107

    op_t create_bias(const shape& bias)
    {
108
        assert(fp);
Paul's avatar
Paul committed
109
        op_t result;
Paul's avatar
Paul committed
110
111
        auto b      = shape{bias.type(), {1, bias.lens().at(1), 1, 1}};
        auto t      = keep_alive(make_tensor(b));
Paul's avatar
Paul committed
112
113
        auto status = miopenCreateOpBiasForward(fp.get(), &result, t.get());
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
114
            MIGRAPHX_THROW("Creating operator failed");
Paul's avatar
Paul committed
115
116
117
118
119
        return result;
    }

    op_t create_relu()
    {
120
        assert(fp);
Paul's avatar
Paul committed
121
122
123
        op_t result;
        auto status = miopenCreateOpActivationForward(fp.get(), &result, miopenActivationRELU);
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
124
            MIGRAPHX_THROW("Creating operator failed");
Paul's avatar
Paul committed
125
126
127
128
129
        return result;
    }

    op_t create_conv(const op::convolution& op, const shape& weights)
    {
130
        assert(fp);
Paul's avatar
Paul committed
131
        op_t result;
Paul's avatar
Paul committed
132
133
        auto cd     = keep_alive(make_conv(op));
        auto t      = keep_alive(make_tensor(weights));
Paul's avatar
Paul committed
134
135
        auto status = miopenCreateOpConvForward(fp.get(), &result, cd.get(), t.get());
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
136
            MIGRAPHX_THROW("Creating operator failed");
Paul's avatar
Paul committed
137
138
        return result;
    }
Paul's avatar
Paul committed
139
140
141

    shape get_workspace(context&)
    {
142
        // assert(fp);
Paul's avatar
Paul committed
143
144
145
146
147
        // TODO: Use zero workspace for now
        std::size_t ws_size = 0;
        // int algo_count = 1;
        // miopenConvFwdAlgorithm_t algo;
        // miopenFusionPlanConvolutionGetAlgo(fp.get(), 1, &algo_count, &algo);
Paul's avatar
Paul committed
148
149
        // miopenFusionPlanGetWorkSpaceSize(ctx.get_stream().get_miopen(), fp.get(), &ws_size,
        // algo);
Paul's avatar
Paul committed
150
151
152
        return shape{shape::int8_type, {ws_size}};
    }

153
    bool compile(context& ctx)
Paul's avatar
Paul committed
154
    {
155
        assert(fp);
156
157
        return miopenCompileFusionPlan(ctx.get_stream().get_miopen(), fp.get()) ==
               miopenStatusSuccess;
Paul's avatar
Paul committed
158
159
    }

Paul's avatar
Paul committed
160
161
162
163
    argument execute(context& ctx,
                     const fused_operator_args& fargs,
                     const argument& x,
                     const argument& y) const
Paul's avatar
Paul committed
164
    {
165
        assert(fp);
Paul's avatar
Paul committed
166
167
        auto x_td   = make_tensor(x.get_shape());
        auto y_td   = make_tensor(y.get_shape());
Paul's avatar
Paul committed
168
        auto status = miopenExecuteFusionPlan(ctx.get_stream().get_miopen(),
Paul's avatar
Paul committed
169
170
171
172
173
174
                                              fp.get(),
                                              x_td.get(),
                                              x.implicit(),
                                              y_td.get(),
                                              y.implicit(),
                                              fargs.get());
Paul's avatar
Paul committed
175
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
176
            MIGRAPHX_THROW("Failed to execute fusion plan");
Paul's avatar
Paul committed
177
178
        return y;
    }
Paul's avatar
Paul committed
179
180
};

181
182
183
184
185
186
const std::unordered_set<std::string>& get_supported_archs()
{
    static std::unordered_set<std::string> supported_archs{"gfx900", "gfx906", "gfx908", "gfx1030"};
    return supported_archs;
}

Paul's avatar
Paul committed
187
MIGRAPHX_PRED_MATCHER(bias_shape, instruction_ref ins)
Paul's avatar
Paul committed
188
189
{
    auto&& s = ins->get_shape();
Paul's avatar
Paul committed
190
191
    return s.broadcasted() and s.strides().size() == 4 and s.strides()[0] == 0 and
           s.strides()[1] != 0 and s.strides()[2] == 0 and s.strides()[3] == 0;
Paul's avatar
Paul committed
192
193
}

Paul's avatar
Paul committed
194
MIGRAPHX_PRED_MATCHER(fusable_conv, instruction_ref ins)
Paul's avatar
Paul committed
195
{
196
    const auto device_name = trim(split_string(get_device_name(), ':').front());
197
198
    if(not contains(get_supported_archs(), device_name))
        return false;
199
200
    if(enabled(MIGRAPHX_DISABLE_MIOPEN_FUSION{}))
        return false;
Paul's avatar
Paul committed
201
202
    if(ins->name() != "gpu::convolution")
        return false;
Paul's avatar
Paul committed
203
204
    if(ins->get_shape().type() != shape::float_type)
        return false;
Paul's avatar
Paul committed
205
206
207
    auto wei = ins->inputs().at(1)->get_shape();
    assert(wei.lens().size() == 4);
    auto conv = any_cast<miopen_convolution>(ins->get_operator());
Khalique's avatar
Khalique committed
208
    if(conv.op.group > 1)
Khalique's avatar
Khalique committed
209
        return false;
Paul's avatar
Paul committed
210
    if(wei.lens()[1] > 512 and conv.algo != miopenConvolutionFwdAlgoWinograd)
Paul's avatar
Paul committed
211
        return false;
212
213
214
215
216
217

    // Do not fuse non-symmetric input
    auto input_lens = ins->inputs().at(0)->get_shape().lens();
    if(input_lens[2] != input_lens[3] or wei.lens()[2] != wei.lens()[3])
        return false;

Paul's avatar
Paul committed
218
    auto op = conv.op;
219
220
    // Dont fuse winograd for non-3x3s since there is no fused windograd for those configs
    if(conv.algo == miopenConvolutionFwdAlgoWinograd and wei.lens()[2] != 3 and
221
       wei.lens()[3] != 3 and contains({{1, 1}}, op.stride))
222
        return false;
kahmed10's avatar
kahmed10 committed
223
    return contains({{0, 0, 0, 0}, {1, 1, 1, 1}, {2, 2, 2, 2}}, op.padding) and
224
           contains({{0, 0}, {1, 1}}, op.stride) and contains({{1, 1}}, op.dilation);
Paul's avatar
Paul committed
225
226
}

227
struct hip_triadd : ternary_device<hip_triadd, &device::add>
Paul's avatar
Paul committed
228
229
{
};
230
MIGRAPHX_REGISTER_OP(hip_triadd)
Paul's avatar
Paul committed
231

232
struct hip_triadd_clip : quinary_device<hip_triadd_clip, &device::add_clip>
kahmed10's avatar
kahmed10 committed
233
234
{
};
235
MIGRAPHX_REGISTER_OP(hip_triadd_clip)
kahmed10's avatar
kahmed10 committed
236

237
struct hip_add_clip : quaternary_device<hip_add_clip, &device::add_clip>
kahmed10's avatar
kahmed10 committed
238
239
{
};
240
MIGRAPHX_REGISTER_OP(hip_add_clip)
kahmed10's avatar
kahmed10 committed
241

242
struct hip_triadd_relu : ternary_device<hip_triadd_relu, &device::add_relu>
Paul's avatar
Paul committed
243
244
{
};
245
MIGRAPHX_REGISTER_OP(hip_triadd_relu)
Paul's avatar
Paul committed
246

247
248
249
struct hip_triadd_sigmoid : ternary_device<hip_triadd_sigmoid, &device::add_sigmoid>
{
};
250
MIGRAPHX_REGISTER_OP(hip_triadd_sigmoid)
251
252
253
254

struct hip_triadd_tanh : ternary_device<hip_triadd_tanh, &device::add_tanh>
{
};
255
MIGRAPHX_REGISTER_OP(hip_triadd_tanh)
256
257
258
259

struct hip_add_relu : binary_device<hip_add_relu, &device::add_relu>
{
};
260
MIGRAPHX_REGISTER_OP(hip_add_relu)
261
262
263
264

struct hip_add_sigmoid : binary_device<hip_add_relu, &device::add_sigmoid>
{
};
265
MIGRAPHX_REGISTER_OP(hip_add_sigmoid)
266
267

struct hip_add_tanh : binary_device<hip_add_tanh, &device::add_tanh>
Paul's avatar
Paul committed
268
269
{
};
270
MIGRAPHX_REGISTER_OP(hip_add_tanh)
Paul's avatar
Paul committed
271

kahmed10's avatar
kahmed10 committed
272
273
struct hip_layernorm : unary_device<hip_layernorm, &device::layernorm>
{
274
275
    // Empty finalize to skip dimension reduction
    void finalize(context&, const shape&, const std::vector<shape>&) {}
kahmed10's avatar
kahmed10 committed
276
};
277
MIGRAPHX_REGISTER_OP(hip_layernorm)
kahmed10's avatar
kahmed10 committed
278

Paul Fultz II's avatar
Paul Fultz II committed
279
280
281
282
283
284
285
struct hip_triadd_layernorm : ternary_device<hip_triadd_layernorm, &device::triadd_layernorm>
{
    // Empty finalize to skip dimension reduction
    void finalize(context&, const shape&, const std::vector<shape>&) {}
};
MIGRAPHX_REGISTER_OP(hip_triadd_layernorm)

kahmed10's avatar
kahmed10 committed
286
287
288
struct hip_gelu : unary_device<hip_gelu, &device::gelu>
{
};
289
MIGRAPHX_REGISTER_OP(hip_gelu)
kahmed10's avatar
kahmed10 committed
290
291
292
293

struct hip_add_gelu : binary_device<hip_add_gelu, &device::add_gelu>
{
};
294
MIGRAPHX_REGISTER_OP(hip_add_gelu)
kahmed10's avatar
kahmed10 committed
295
296
297
298

struct hip_gelu_new : unary_device<hip_gelu_new, &device::gelu_new>
{
};
299
MIGRAPHX_REGISTER_OP(hip_gelu_new)
kahmed10's avatar
kahmed10 committed
300
301
302
303

struct hip_add_gelu_new : binary_device<hip_add_gelu_new, &device::add_gelu_new>
{
};
304
MIGRAPHX_REGISTER_OP(hip_add_gelu_new)
kahmed10's avatar
kahmed10 committed
305

306
struct hip_mul_add : ternary_device<hip_mul_add, &device::mul_add>
Paul's avatar
Paul committed
307
308
{
};
309
MIGRAPHX_REGISTER_OP(hip_mul_add)
Paul's avatar
Paul committed
310

311
struct hip_mul_add_relu : ternary_device<hip_mul_add_relu, &device::mul_add_relu>
Paul's avatar
Paul committed
312
313
{
};
314
MIGRAPHX_REGISTER_OP(hip_mul_add_relu)
Paul's avatar
Paul committed
315

Paul's avatar
Paul committed
316
317
318
void move_broadcasted_back(std::vector<instruction_ref>& args)
{
    // Ensure the last arguments is the broadcasted one
Paul's avatar
Paul committed
319
    auto last = std::prev(args.end());
Paul's avatar
Paul committed
320
321
    auto it =
        std::find_if(args.begin(), last, [](auto arg) { return arg->get_shape().broadcasted(); });
Paul's avatar
Paul committed
322
323
    if(it != last)
        std::swap(*it, *std::prev(last));
Paul's avatar
Paul committed
324
325
326
327
328
}

void move_standard_front(std::vector<instruction_ref>& args)
{
    // Ensure the first arguments is the standard one
Paul's avatar
Paul committed
329
    auto last = std::prev(args.end());
Paul's avatar
Paul committed
330
331
    auto it =
        std::find_if(args.begin(), last, [](auto arg) { return arg->get_shape().standard(); });
Paul's avatar
Paul committed
332
    if(it != last)
Paul's avatar
Paul committed
333
334
335
        std::swap(*it, args.front());
}

336
337
auto gpu_name(const std::string& s) { return match::name("gpu::" + s); }

kahmed10's avatar
kahmed10 committed
338
339
struct find_layernorm
{
340
    auto matcher() const { return match::layernorm(&gpu_name); }
kahmed10's avatar
kahmed10 committed
341

342
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
343
344
345
346
347
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

348
349
350
351
352
353
354
355
356
        // We dont fuse for non-standard layouts
        if(not x_ins->get_shape().standard())
            return;

        auto relements = x_ins->get_shape().lens().back();

        if(relements > 1024 or (relements % 4 != 0 and relements > 256))
            return;

357
        m.replace_instruction(ins, hip_layernorm{}, x_ins, args.back());
kahmed10's avatar
kahmed10 committed
358
359
360
    }
};

Paul Fultz II's avatar
Paul Fultz II committed
361
362
363
364
365
366
367
368
struct find_triadd_layernorm
{
    auto matcher() const
    {
        return match::name("gpu::layernorm")(match::arg(0)(match::name("gpu::triadd")(
            match::used_once(), match::all_of[match::inputs()](match::standard_shape()))));
    }

369
    void apply(module& m, const match::matcher_result& r) const
Paul Fultz II's avatar
Paul Fultz II committed
370
371
372
    {
        auto ins    = r.result;
        auto triadd = ins->inputs().front();
373
        m.replace_instruction(ins, hip_triadd_layernorm{}, triadd->inputs());
Paul Fultz II's avatar
Paul Fultz II committed
374
375
376
    }
};

kahmed10's avatar
kahmed10 committed
377
378
struct find_gelu
{
379
    auto matcher() const { return match::gelu_erf(&gpu_name); }
kahmed10's avatar
kahmed10 committed
380

381
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
382
383
384
385
386
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

387
        m.replace_instruction(ins, hip_gelu{}, x_ins, args.back());
kahmed10's avatar
kahmed10 committed
388
389
390
391
392
393
394
395
396
397
    }
};

struct find_add_gelu
{
    auto matcher() const
    {
        return match::name("gpu::gelu")(match::arg(0)(match::name("gpu::add").bind("add")));
    }

398
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
399
400
401
402
403
404
405
406
    {
        auto add_ins = r.instructions["add"];
        auto ins     = r.result;
        auto args    = add_ins->inputs();
        move_standard_front(args);
        move_broadcasted_back(args);

        args.back() = ins->inputs().back();
407
        m.replace_instruction(ins, hip_add_gelu{}, args);
kahmed10's avatar
kahmed10 committed
408
409
410
411
412
    }
};

struct find_gelu_new
{
kahmed10's avatar
kahmed10 committed
413
    bool fast_math = true;
kahmed10's avatar
kahmed10 committed
414

415
    auto matcher() const { return match::gelu_tanh(&gpu_name); }
kahmed10's avatar
kahmed10 committed
416

417
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
418
419
420
421
422
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

Paul Fultz II's avatar
Paul Fultz II committed
423
        if(fast_math)
424
            m.replace_instruction(ins, hip_gelu{}, x_ins, args.back());
Paul Fultz II's avatar
Paul Fultz II committed
425
        else
426
            m.replace_instruction(ins, hip_gelu_new{}, x_ins, args.back());
kahmed10's avatar
kahmed10 committed
427
428
429
430
431
432
433
434
435
436
    }
};

struct find_add_gelu_new
{
    auto matcher() const
    {
        return match::name("gpu::gelu_new")(match::arg(0)(match::name("gpu::add").bind("add")));
    }

437
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
438
439
440
441
442
443
444
445
    {
        auto add_ins = r.instructions["add"];
        auto ins     = r.result;
        auto args    = add_ins->inputs();
        move_standard_front(args);
        move_broadcasted_back(args);

        args.back() = ins->inputs().back();
446
        m.replace_instruction(ins, hip_add_gelu_new{}, args);
kahmed10's avatar
kahmed10 committed
447
448
449
    }
};

kahmed10's avatar
kahmed10 committed
450
451
452
453
454
455
struct find_add_clip
{
    auto matcher() const
    {
        return match::name(std::unordered_set<std::string>{"gpu::clip", "gpu::clipped_relu"})(
            match::arg(0)(match::any_of(match::name("gpu::add"),
kahmed10's avatar
kahmed10 committed
456
                                        match::name("gpu::triadd"),
kahmed10's avatar
kahmed10 committed
457
458
459
460
                                        match::any_of[match::inputs()](match::standard_shape()))
                              .bind("add")));
    }

461
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
462
    {
kahmed10's avatar
kahmed10 committed
463
464
465
466
467
468
469
470
471
472
        auto add_ins  = r.instructions["add"];
        auto ins      = r.result;
        auto ins_args = ins->inputs();
        auto add_args = add_ins->inputs();
        move_standard_front(add_args);
        move_broadcasted_back(add_args);

        // Use the allocation from the clip operator
        add_args.pop_back();
        add_args.insert(add_args.end(), std::next(ins_args.begin()), ins_args.end());
kahmed10's avatar
kahmed10 committed
473
        if(add_ins->name() == "gpu::add")
474
            m.replace_instruction(ins, hip_add_clip{}, add_args);
kahmed10's avatar
kahmed10 committed
475
        else if(add_ins->name() == "gpu::triadd")
476
            m.replace_instruction(ins, hip_triadd_clip{}, add_args);
kahmed10's avatar
kahmed10 committed
477
478
479
    }
};

480
struct find_add_unary
Paul's avatar
Paul committed
481
{
482
483
484
    std::string op_name;
    operation binary_add_op;
    operation ternary_add_op;
Paul's avatar
Paul committed
485
486
    auto matcher() const
    {
487
        return match::name(op_name)(match::arg(0)(
Paul's avatar
Paul committed
488
            match::used_once(),
Paul's avatar
Paul committed
489
            match::any_of(match::name("gpu::add"),
kahmed10's avatar
kahmed10 committed
490
                          match::name("gpu::triadd"),
Paul's avatar
Paul committed
491
492
493
                          match::any_of(match::name("@literal"),
                                        match::any_of[match::inputs()](match::standard_shape())))
                .bind("add")));
Paul's avatar
Paul committed
494
    }
Paul's avatar
Paul committed
495

496
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
497
    {
Paul's avatar
Paul committed
498
        auto add_ins = r.instructions["add"];
Paul's avatar
Paul committed
499
500
        auto ins     = r.result;
        auto args    = add_ins->inputs();
Paul's avatar
Paul committed
501
502
503
        move_standard_front(args);
        move_broadcasted_back(args);

Paul's avatar
Paul committed
504
        // Use the allocation from the relu operator
Paul's avatar
Paul committed
505
        args.back() = ins->inputs().back();
Paul's avatar
Paul committed
506
        if(add_ins->name() == "gpu::add")
507
            m.replace_instruction(ins, binary_add_op, args);
kahmed10's avatar
kahmed10 committed
508
        else if(add_ins->name() == "gpu::triadd")
509
            m.replace_instruction(ins, ternary_add_op, args);
Paul's avatar
Paul committed
510
511
512
    }
};

Paul's avatar
Paul committed
513
struct find_triadd
Paul's avatar
Paul committed
514
515
516
{
    auto matcher() const
    {
Paul's avatar
Paul committed
517
        return match::name("gpu::add")(match::either_arg(0, 1)(
Paul's avatar
Paul committed
518
            match::name("gpu::add")(match::used_once()).bind("add"),
Paul's avatar
Paul committed
519
520
521
            match::any(match::any_of(match::name("@literal"),
                                     match::any_of[match::inputs()](match::standard_shape())))
                .bind("input")));
Paul's avatar
Paul committed
522
523
    }

524
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
525
    {
Paul's avatar
Paul committed
526
527
528
529
        auto add_ins   = r.instructions["add"];
        auto input_ins = r.instructions["input"];
        auto ins       = r.result;
        auto args      = add_ins->inputs();
530

Paul's avatar
Paul committed
531
        auto is_broadcasted = [](auto arg) { return arg->get_shape().broadcasted(); };
532
        if(std::count_if(args.begin(), args.end(), is_broadcasted) > 2)
Paul's avatar
Paul committed
533
534
            return;
        args.insert(args.begin(), input_ins);
Paul's avatar
Paul committed
535
536
537
        move_standard_front(args);
        move_broadcasted_back(args);

Paul's avatar
Paul committed
538
        args.back() = ins->inputs().back();
539
        m.replace_instruction(ins, hip_triadd{}, args);
Paul's avatar
Paul committed
540
    }
Paul's avatar
Paul committed
541
542
};

Paul's avatar
Paul committed
543
544
545
546
struct find_mul_add
{
    auto matcher() const
    {
Paul's avatar
Paul committed
547
548
        return match::name("gpu::add")(match::either_arg(0, 1)(
            match::name("gpu::mul")(match::used_once()).bind("mul"), match::any().bind("b")));
Paul's avatar
Paul committed
549
550
    }

551
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
552
    {
Paul's avatar
Paul committed
553
554
555
556
        auto mul_ins = r.instructions["mul"];
        auto b_ins   = r.instructions["b"];
        auto ins     = r.result;
        auto args    = mul_ins->inputs();
Paul's avatar
Paul committed
557
558
559
560
561
562
563
        assert(mul_ins != b_ins);

        move_standard_front(args);
        move_broadcasted_back(args);
        args.insert(std::prev(args.end()), b_ins);

        args.back() = ins->inputs().back();
564
        m.replace_instruction(ins, hip_mul_add{}, args);
Paul's avatar
Paul committed
565
566
567
    }
};

Paul's avatar
Paul committed
568
569
570
571
struct find_mul_add_relu
{
    auto matcher() const
    {
Paul's avatar
Paul committed
572
        return match::name("gpu::relu")(
kahmed10's avatar
kahmed10 committed
573
            match::arg(0)(match::name("gpu::mul_add")(match::used_once()).bind("mul_add")));
Paul's avatar
Paul committed
574
575
    }

576
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
577
578
    {
        auto mul_add_ins = r.instructions["mul_add"];
Paul's avatar
Paul committed
579
580
        auto ins         = r.result;
        auto args        = mul_add_ins->inputs();
Paul's avatar
Paul committed
581
582
583

        // Use the allocation from the relu operator
        args.back() = ins->inputs().back();
584
        m.replace_instruction(ins, hip_mul_add_relu{}, args);
Paul's avatar
Paul committed
585
586
587
    }
};

588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
struct miopen_fusion
{
    struct fuse_op_data
    {
        operation op;
        float alpha = 1;
        float beta  = 0;
    };
    struct fuse_op : fuse_op_data, reflect_equality<fuse_op>, reflect_stream<fuse_op>
    {
        template <class Self, class F>
        static auto reflect(Self& self, F f)
        {
            return pack(f(self.op, "op"), f(self.alpha, "alpha"), f(self.beta, "beta"));
        }
    };
    std::vector<fuse_op> ops = {};
    fusion f                 = {};
    std::function<void(context&, const fusion&, const std::vector<argument>&)> execute;
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return pack(f(self.ops, "ops"));
    }

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

618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
    value compile(context& ctx, const shape&, std::vector<shape> inputs)
    {
        // Compensate for allocation
        inputs.pop_back();
        std::size_t i = 0;
        f             = fusion(inputs[i]);
        i++;
        std::vector<std::function<void(const fused_operator_args&, const std::vector<argument>&)>>
            invokers;
        for(auto&& fop : ops)
        {
            if(i > inputs.size())
            {
                f = {};
                return {};
            }
            if(fop.op.name() == "convolution")
            {
                auto* mop = f.create_conv(any_cast<op::convolution>(fop.op), inputs[i]);
                invokers.push_back(
                    [=](const fused_operator_args& fargs, const std::vector<argument>& args) {
                        miopenSetOpArgsConvForward(
                            fargs.get(), mop, &fop.alpha, &fop.beta, args[i].implicit());
                    });
                i++;
            }
            else if(fop.op.name() == "add")
            {
                auto* mop = f.create_bias(inputs[i]);
                invokers.push_back(
                    [=](const fused_operator_args& fargs, const std::vector<argument>& args) {
                        miopenSetOpArgsBiasForward(
                            fargs.get(), mop, &fop.alpha, &fop.beta, args[i].implicit());
                    });
                i++;
            }
            else if(fop.op.name() == "relu")
            {
                auto* mop = f.create_relu();
                invokers.push_back([=](const fused_operator_args& fargs,
                                       const std::vector<argument>&) {
                    miopenSetOpArgsActivForward(fargs.get(), mop, &fop.alpha, &fop.beta, 0, 0, 0);
                });
            }
            else
            {
                f = {};
                return {};
            }
        }
        if(not f.compile(ctx))
        {
            f = {};
            return {};
        }
        execute = [invokers](context& c, const fusion& ff, const std::vector<argument>& args) {
            auto fargs = make_fused_args();
            for(auto&& invoker : invokers)
                invoker(fargs, args);
            ff.execute(c, fargs, args.front(), args.back());
        };
        return {{"workspace", f.get_workspace(ctx).bytes()}};
    }
    void finalize(context& ctx, const shape& output_shape, const std::vector<shape>& inputs)
    {
        if(not f.empty())
            return;
        auto v = compile(ctx, output_shape, inputs);
        if(not v.is_object())
            MIGRAPHX_THROW("Failed to compile fusion plan");
    }
    std::string name() const { return "gpu::miopen_fusion"; }
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        if(ops.empty())
            return {};
        // TODO: Check number of arguments
        return ops.front().op.compute_shape({inputs[0], inputs[1]});
    }
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
    {
        execute(ctx, f, args);
        return args.back();
    }
};

Paul's avatar
Paul committed
704
705
706
struct miopen_conv_bias
{
    op::convolution op;
707
    fusion fp         = {};
708
709
    fusion::op_t conv = {};
    fusion::op_t bias = {};
Paul's avatar
Paul committed
710

Paul's avatar
Paul committed
711
712
713
714
715
716
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return op::convolution::reflect(self.op, f);
    }

Paul's avatar
Paul committed
717
718
719
720
721
    std::string name() const { return "gpu::conv_bias"; }
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        check_shapes{inputs, *this}.has(5);
        // TODO: Check slices
kahmed10's avatar
kahmed10 committed
722
        return op.normalize_compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
723
    }
Paul's avatar
Paul committed
724
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
Paul's avatar
Paul committed
725
    {
Paul's avatar
Paul committed
726
        auto fargs  = make_fused_args();
Paul's avatar
Paul committed
727
        float alpha = 1;
Paul's avatar
Paul committed
728
        float beta  = 0;
Paul's avatar
Paul committed
729
730
        miopenSetOpArgsConvForward(fargs.get(), conv, &alpha, &beta, args[1].implicit());
        miopenSetOpArgsBiasForward(fargs.get(), bias, &alpha, &beta, args[3].implicit());
731
        return fp.execute(ctx, fargs, args[0], args[4]);
Paul's avatar
Paul committed
732
733
    }

734
735
    void finalize(context& ctx, const shape&, const std::vector<shape>& inputs)
    {
736
737
738
739
        fp   = fusion(inputs[0]);
        conv = fp.create_conv(op, inputs[1]);
        bias = fp.create_bias(inputs[3]);
        if(not fp.compile(ctx))
740
            MIGRAPHX_THROW("Failed to compile fusion plan");
741
742
    }

743
    shape get_workspace(context& ctx) { return fp.get_workspace(ctx); }
Paul's avatar
Paul committed
744
745
746
747
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
Paul's avatar
Paul committed
748
};
749
MIGRAPHX_REGISTER_OP(miopen_conv_bias)
Paul's avatar
Paul committed
750

Paul's avatar
Add cbr  
Paul committed
751
752
753
struct miopen_conv_bias_relu
{
    op::convolution op;
754
    fusion fp         = {};
755
756
757
    fusion::op_t conv = {};
    fusion::op_t bias = {};
    fusion::op_t relu = {};
Paul's avatar
Add cbr  
Paul committed
758

Paul's avatar
Paul committed
759
760
761
762
763
764
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return op::convolution::reflect(self.op, f);
    }

Paul's avatar
Add cbr  
Paul committed
765
766
767
768
769
    std::string name() const { return "gpu::conv_bias_relu"; }
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        check_shapes{inputs, *this}.has(5);
        // TODO: Check slices
kahmed10's avatar
kahmed10 committed
770
        return op.normalize_compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Add cbr  
Paul committed
771
    }
Paul's avatar
Paul committed
772
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
Paul's avatar
Add cbr  
Paul committed
773
774
    {
        auto fargs  = make_fused_args();
Paul's avatar
Paul committed
775
        float alpha = 1;
Paul's avatar
Paul committed
776
        float beta  = 0;
Paul's avatar
Add cbr  
Paul committed
777
778
        miopenSetOpArgsConvForward(fargs.get(), conv, &alpha, &beta, args[1].implicit());
        miopenSetOpArgsBiasForward(fargs.get(), bias, &alpha, &beta, args[3].implicit());
Paul's avatar
Paul committed
779
        miopenSetOpArgsActivForward(fargs.get(), relu, &alpha, &beta, 0, 0, 0);
780
        return fp.execute(ctx, fargs, args[0], args[4]);
Paul's avatar
Add cbr  
Paul committed
781
    }
782
783
    void finalize(context& ctx, const shape&, const std::vector<shape>& inputs)
    {
784
785
786
787
788
        fp   = fusion(inputs[0]);
        conv = fp.create_conv(op, inputs[1]);
        bias = fp.create_bias(inputs[3]);
        relu = fp.create_relu();
        fp.compile(ctx);
789
790
    }

791
    shape get_workspace(context& ctx) { return fp.get_workspace(ctx); }
Paul's avatar
Paul committed
792
793
794
795
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
Paul's avatar
Add cbr  
Paul committed
796
};
797
MIGRAPHX_REGISTER_OP(miopen_conv_bias_relu)
Paul's avatar
Add cbr  
Paul committed
798

Paul's avatar
Paul committed
799
template <class... Ms>
Paul's avatar
Add cbr  
Paul committed
800
801
auto conv_bias(Ms... ms)
{
Paul's avatar
Paul committed
802
    return match::name("gpu::add")(
Paul's avatar
Paul committed
803
804
        match::either_arg(0, 1)(bias_shape(match::used_once()).bind("bias"),
                                fusable_conv(match::used_once()).bind("conv")),
Paul's avatar
Paul committed
805
        ms...);
Paul's avatar
Paul committed
806
807
}

Paul's avatar
Paul committed
808
template <class Op>
809
void apply_conv_bias(context& ctx, module& m, const match::matcher_result& r)
Paul's avatar
Paul committed
810
811
812
813
814
815
816
817
818
819
{
    auto conv_ins    = r.instructions["conv"];
    auto bias_ins    = r.instructions["bias"];
    auto ins         = r.result;
    auto input_ins   = conv_ins->inputs().at(0);
    auto weights_ins = conv_ins->inputs().at(1);
    auto conv_op     = any_cast<miopen_convolution>(conv_ins->get_operator()).op;
    auto alloc_ins   = ins->inputs().back();
    auto old_ws_ins  = conv_ins->inputs().at(2);

820
    Op cb{conv_op};
Paul's avatar
Paul committed
821
    // TODO: Insert ws allocation
Paul's avatar
Paul committed
822
    auto ws = cb.get_workspace(ctx);
Paul's avatar
Paul committed
823
    (void)ws;
824
    m.replace_instruction(ins, cb, input_ins, weights_ins, old_ws_ins, bias_ins, alloc_ins);
Paul's avatar
Add cbr  
Paul committed
825
826
}

827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
inline auto precompile_name(std::string s) // NOLINT
{
    return match::make_basic_pred_matcher([=](instruction_ref ins) {
        if(ins->name() != "gpu::precompile_op")
            return false;
        auto op = from_value<operation>(ins->get_operator().to_value().at("op"));
        return (op.name() == s);
    });
}

template <class... Ms>
auto conv_bias_pointwise(Ms... ms)
{
    return precompile_name("pointwise")(
        match::either_arg(0, 1)(bias_shape(match::used_once()).bind("bias"),
                                fusable_conv(match::used_once()).bind("conv")),
        ms...);
}

Paul's avatar
Paul committed
846
struct find_conv_bias
Paul's avatar
Paul committed
847
{
Paul's avatar
Paul committed
848
    context* ctx = nullptr;
Paul's avatar
Paul committed
849
850
    auto matcher() const
    {
kahmed10's avatar
kahmed10 committed
851
852
        return conv_bias(match::none_of(
            match::output(match::name(std::unordered_set<std::string>{"gpu::relu"}))));
Paul's avatar
Paul committed
853
854
    }

855
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
856
    {
857
        apply_conv_bias<miopen_conv_bias>(*ctx, m, r);
Paul's avatar
Paul committed
858
859
860
    }
};

Paul's avatar
Paul committed
861
struct find_conv_bias_relu
Paul's avatar
Add cbr  
Paul committed
862
863
{
    context* ctx = nullptr;
Paul's avatar
Paul committed
864
    auto matcher() const { return match::name("gpu::relu")(match::arg(0)(conv_bias())); }
Paul's avatar
Add cbr  
Paul committed
865

866
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Add cbr  
Paul committed
867
    {
868
        apply_conv_bias<miopen_conv_bias_relu>(*ctx, m, r);
Paul's avatar
Add cbr  
Paul committed
869
870
    }
};
871

872
873
874
875
876
877
878
879
880
881
882
struct find_conv_pointwise
{
    context* ctx = nullptr;
    auto matcher() const
    {
        return precompile_name("pointwise")(
            match::nargs(3),
            match::either_arg(0, 1)(bias_shape(match::used_once()).bind("bias"),
                                    fusable_conv(match::used_once()).bind("conv")));
    }

883
    void apply(module& m, const match::matcher_result& r) const
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
    {
        auto conv_ins    = r.instructions["conv"];
        auto bias_ins    = r.instructions["bias"];
        auto ins         = r.result;
        auto input_ins   = conv_ins->inputs().at(0);
        auto weights_ins = conv_ins->inputs().at(1);
        auto conv_op     = any_cast<miopen_convolution>(conv_ins->get_operator()).op;
        auto alloc_ins   = ins->inputs().back();

        module_ref pm = ins->module_inputs().front();

        miopen_fusion op{};
        op.ops.push_back({{conv_op}});
        for(auto&& i : *pm)
        {
            if(i.name()[0] == '@')
                continue;
            op.ops.push_back({{i.get_operator()}});
        }
        std::vector<instruction_ref> inputs = {input_ins, weights_ins, bias_ins, alloc_ins};
        auto v                              = op.compile(*ctx, ins->get_shape(), to_shapes(inputs));
        if(not v.is_object())
            return;
        m.replace_instruction(ins, op, inputs);
    }
};

911
912
913
914
915
916
917
918
919
920
struct find_gemm_add
{
    auto matcher() const
    {
        return match::name("gpu::add")(
            match::all_of[match::inputs()](match::standard_shape()),
            match::either_arg(0, 1)(match::used_once().bind("c"),
                                    match::name("gpu::gemm")(match::nargs(3)).bind("gemm")));
    }

921
    void apply(module& m, const match::matcher_result& r) const
922
923
924
925
926
927
928
929
    {
        auto ins      = r.result;
        auto gemm_ins = r.instructions["gemm"];
        auto c_ins    = r.instructions["c"];

        auto gemm = any_cast<rocblas_gemm<op::dot>>(gemm_ins->get_operator());

        // Already fused gemm
930
        if(not float_equal(gemm.beta, 0))
931
932
933
934
935
936
937
938
            return;

        auto inputs = gemm_ins->inputs();
        inputs.pop_back();

        auto copy_ins = c_ins;

        // Insert copy
939
        if(ins == m.end() or c_ins->outputs().size() > 1 or c_ins->inputs().empty())
940
        {
941
            copy_ins = m.insert_instruction(ins, hip_copy{}, c_ins, ins->inputs().back());
942
943
944
945
        }
        inputs.push_back(copy_ins);
        inputs.push_back(copy_ins);

946
        gemm.beta = 1;
947
        m.replace_instruction(ins, gemm, inputs);
948
949
950
    }
};

951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
auto pointwise_name(const std::string& s)
{
    return precompile_name("pointwise")(match::make_basic_pred_matcher([=](auto ins) {
        module_ref pm = ins->module_inputs().front();
        auto n = std::count_if(pm->begin(), pm->end(), [&](auto& i) { return i.name() == s; });
        if(n != 1)
            return false;
        return std::all_of(pm->begin(), pm->end(), [&](auto& i) {
            return starts_with(i.name(), "@") or i.name() == s;
        });
    }));
}

struct find_gemm_pointwise
{
    auto matcher() const
    {
        return pointwise_name("add")(
            match::nargs(3),
            match::all_of[match::inputs()](match::standard_shape()),
            match::either_arg(0, 1)(match::used_once().bind("c"),
                                    match::name("gpu::gemm")(match::nargs(3)).bind("gemm")));
    }

    void apply(module& m, const match::matcher_result& r) const
    {
        auto ins      = r.result;
        auto gemm_ins = r.instructions["gemm"];
        auto c_ins    = r.instructions["c"];

        auto gemm = any_cast<rocblas_gemm<op::dot>>(gemm_ins->get_operator());

        // Already fused gemm
        if(not float_equal(gemm.beta, 0))
            return;

        auto inputs = gemm_ins->inputs();
        inputs.pop_back();

        inputs.push_back(c_ins);
991
        inputs.push_back(ins->inputs().back());
992
993
994
995
996
997

        gemm.beta = 1;
        m.replace_instruction(ins, gemm, inputs);
    }
};

998
999
1000
1001
1002
1003
1004
struct find_commutative_broadcast
{
    auto matcher() const
    {
        return match::name("gpu::add", "gpu::mul")(match::arg(1)(match::broadcast_shape()));
    }

1005
    void apply(module& m, const match::matcher_result& r) const
1006
1007
1008
1009
1010
    {
        auto ins  = r.result;
        auto args = ins->inputs();
        move_broadcasted_back(args);

1011
        m.replace_instruction(ins, ins->get_operator(), args);
1012
1013
1014
    }
};

1015
void fuse_ops::apply(module& m) const
Paul's avatar
Paul committed
1016
{
1017
1018
1019
1020
    match::find_matches(m, find_gelu{}, find_gelu_new{fast_math});
    run_passes(m, {dead_code_elimination{}});
    match::find_matches(m, find_triadd{});
    match::find_matches(m,
kahmed10's avatar
kahmed10 committed
1021
                        find_layernorm{},
1022
                        find_conv_pointwise{ctx},
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
                        find_conv_bias_relu{ctx},
                        find_conv_bias{ctx},
                        find_add_gelu{},
                        find_add_gelu_new{},
                        find_mul_add{},
                        find_mul_add_relu{},
                        find_add_unary{"gpu::relu", hip_add_relu{}, hip_triadd_relu{}},
                        find_add_unary{"gpu::sigmoid", hip_add_sigmoid{}, hip_triadd_sigmoid{}},
                        find_add_unary{"gpu::tanh", hip_add_tanh{}, hip_triadd_tanh{}},
                        find_add_clip{});
1033
    run_passes(m, {dead_code_elimination{}});
1034
1035
1036
1037
1038
    match::find_matches(m,
                        find_triadd_layernorm{},
                        find_gemm_add{},
                        find_gemm_pointwise{},
                        find_commutative_broadcast{});
Paul's avatar
Paul committed
1039
}
Paul's avatar
Paul committed
1040
1041

} // namespace gpu
Paul's avatar
Paul committed
1042
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1043
} // namespace migraphx