fuse_ops.cpp 23.4 KB
Newer Older
kahmed10's avatar
kahmed10 committed
1
2
#include <migraphx/pass_manager.hpp>
#include <migraphx/dead_code_elimination.hpp>
Paul's avatar
Paul committed
3
4
5
#include <migraphx/gpu/fuse_ops.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/gpu/miopen.hpp>
kahmed10's avatar
kahmed10 committed
6
#include <migraphx/gpu/clip.hpp>
Paul's avatar
Paul committed
7
#include <migraphx/gpu/convolution.hpp>
8
#include <migraphx/gpu/oper.hpp>
kahmed10's avatar
kahmed10 committed
9
10
11
#include <migraphx/gpu/add.hpp>
#include <migraphx/gpu/mul.hpp>
#include <migraphx/gpu/device/layernorm.hpp>
kahmed10's avatar
kahmed10 committed
12
#include <migraphx/gpu/device/gelu.hpp>
Paul's avatar
Paul committed
13
#include <migraphx/gpu/device/mul_add.hpp>
14
15
16
17
18
#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
19
#include <migraphx/gpu/device/add.hpp>
Paul's avatar
Paul committed
20
#include <migraphx/instruction.hpp>
21
#include <migraphx/register_op.hpp>
Paul's avatar
Paul committed
22
#include <migraphx/array.hpp>
kahmed10's avatar
kahmed10 committed
23
#include <migraphx/op/clip.hpp>
kahmed10's avatar
kahmed10 committed
24
#include <cmath>
Paul's avatar
Paul committed
25
26

namespace migraphx {
Paul's avatar
Paul committed
27
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
28
29
namespace gpu {

30
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_MIOPEN_FUSION)
kahmed10's avatar
kahmed10 committed
31
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_FAST_GELU)
32

Paul's avatar
Paul committed
33
34
35
36
37
38
39
40
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
41
    template <class T>
Paul's avatar
Paul committed
42
43
44
45
46
47
48
    auto keep_alive(T x)
    {
        auto result = share(std::move(x));
        storage.push_back(result);
        return result;
    }

49
50
    fusion() = default;

Paul's avatar
Paul committed
51
52
53
    fusion(const shape& input)
    // : fp(make_fusion_plan(input))
    {
54
        assert(input.standard());
Paul's avatar
Paul committed
55
        auto t = make_tensor(input);
Paul's avatar
Paul committed
56
        fp     = make_fusion_plan(t);
Paul's avatar
Paul committed
57
58
59
60
61
62
63
64
        keep_alive(std::move(t));
    }

    op_t operator[](std::size_t i) const
    {
        op_t result;
        auto status = miopenFusionPlanGetOp(fp.get(), i, &result);
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
65
            MIGRAPHX_THROW("Failed retrieving operator at " + std::to_string(i));
Paul's avatar
Paul committed
66
67
68
        return result;
    }

Paul's avatar
Paul committed
69
    auto get() const { return fp.get(); }
Paul's avatar
Paul committed
70
71
72
73

    op_t create_bias(const shape& bias)
    {
        op_t result;
Paul's avatar
Paul committed
74
75
        auto b      = shape{bias.type(), {1, bias.lens().at(1), 1, 1}};
        auto t      = keep_alive(make_tensor(b));
Paul's avatar
Paul committed
76
77
        auto status = miopenCreateOpBiasForward(fp.get(), &result, t.get());
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
78
            MIGRAPHX_THROW("Creating operator failed");
Paul's avatar
Paul committed
79
80
81
82
83
84
85
86
        return result;
    }

    op_t create_relu()
    {
        op_t result;
        auto status = miopenCreateOpActivationForward(fp.get(), &result, miopenActivationRELU);
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
87
            MIGRAPHX_THROW("Creating operator failed");
Paul's avatar
Paul committed
88
89
90
91
92
93
        return result;
    }

    op_t create_conv(const op::convolution& op, const shape& weights)
    {
        op_t result;
Paul's avatar
Paul committed
94
95
        auto cd     = keep_alive(make_conv(op));
        auto t      = keep_alive(make_tensor(weights));
Paul's avatar
Paul committed
96
97
        auto status = miopenCreateOpConvForward(fp.get(), &result, cd.get(), t.get());
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
98
            MIGRAPHX_THROW("Creating operator failed");
Paul's avatar
Paul committed
99
100
        return result;
    }
Paul's avatar
Paul committed
101
102
103
104
105
106
107
108

    shape get_workspace(context&)
    {
        // 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
109
110
        // miopenFusionPlanGetWorkSpaceSize(ctx.get_stream().get_miopen(), fp.get(), &ws_size,
        // algo);
Paul's avatar
Paul committed
111
112
113
114
115
        return shape{shape::int8_type, {ws_size}};
    }

    void compile(context& ctx)
    {
Paul's avatar
Paul committed
116
        auto status = miopenCompileFusionPlan(ctx.get_stream().get_miopen(), fp.get());
Paul's avatar
Paul committed
117
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
118
            MIGRAPHX_THROW("Compiling fusion plan failed");
Paul's avatar
Paul committed
119
120
    }

Paul's avatar
Paul committed
121
122
123
124
    argument execute(context& ctx,
                     const fused_operator_args& fargs,
                     const argument& x,
                     const argument& y) const
Paul's avatar
Paul committed
125
    {
Paul's avatar
Paul committed
126
127
        auto x_td   = make_tensor(x.get_shape());
        auto y_td   = make_tensor(y.get_shape());
Paul's avatar
Paul committed
128
        auto status = miopenExecuteFusionPlan(ctx.get_stream().get_miopen(),
Paul's avatar
Paul committed
129
130
131
132
133
134
                                              fp.get(),
                                              x_td.get(),
                                              x.implicit(),
                                              y_td.get(),
                                              y.implicit(),
                                              fargs.get());
Paul's avatar
Paul committed
135
        if(status != miopenStatusSuccess)
Paul's avatar
Paul committed
136
            MIGRAPHX_THROW("Failed to execute fusion plan");
Paul's avatar
Paul committed
137
138
        return y;
    }
Paul's avatar
Paul committed
139
140
};

Paul's avatar
Paul committed
141
MIGRAPHX_PRED_MATCHER(bias_shape, instruction_ref ins)
Paul's avatar
Paul committed
142
143
{
    auto&& s = ins->get_shape();
Paul's avatar
Paul committed
144
145
    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
146
147
}

Paul's avatar
Paul committed
148
MIGRAPHX_PRED_MATCHER(fusable_conv, instruction_ref ins)
Paul's avatar
Paul committed
149
{
150
151
    if(enabled(MIGRAPHX_DISABLE_MIOPEN_FUSION{}))
        return false;
Paul's avatar
Paul committed
152
153
    if(ins->name() != "gpu::convolution")
        return false;
Paul's avatar
Paul committed
154
155
    if(ins->get_shape().type() != shape::float_type)
        return false;
Paul's avatar
Paul committed
156
157
158
    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
159
    if(conv.op.group > 1)
Khalique's avatar
Khalique committed
160
        return false;
Paul's avatar
Paul committed
161
    if(wei.lens()[1] > 512 and conv.algo != miopenConvolutionFwdAlgoWinograd)
Paul's avatar
Paul committed
162
        return false;
163
164
165
166
167
168

    // 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
169
    auto op = conv.op;
170
171
    // 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
172
       wei.lens()[3] != 3 and contains({{1, 1}}, op.stride))
173
        return false;
Paul's avatar
Paul committed
174
    return contains({{0, 0}, {1, 1}, {2, 2}}, op.padding) and
175
           contains({{0, 0}, {1, 1}}, op.stride) and contains({{1, 1}}, op.dilation);
Paul's avatar
Paul committed
176
177
}

178
struct hip_triadd : ternary_device<hip_triadd, &device::add>
Paul's avatar
Paul committed
179
180
{
};
181
MIGRAPHX_REGISTER_OP(hip_triadd)
Paul's avatar
Paul committed
182

183
struct hip_triadd_clip : quinary_device<hip_triadd_clip, &device::add_clip>
kahmed10's avatar
kahmed10 committed
184
185
{
};
186
MIGRAPHX_REGISTER_OP(hip_triadd_clip)
kahmed10's avatar
kahmed10 committed
187

188
struct hip_add_clip : quaternary_device<hip_add_clip, &device::add_clip>
kahmed10's avatar
kahmed10 committed
189
190
{
};
191
MIGRAPHX_REGISTER_OP(hip_add_clip)
kahmed10's avatar
kahmed10 committed
192

193
struct hip_triadd_relu : ternary_device<hip_triadd_relu, &device::add_relu>
Paul's avatar
Paul committed
194
195
{
};
196
MIGRAPHX_REGISTER_OP(hip_triadd_relu)
Paul's avatar
Paul committed
197

198
199
200
struct hip_triadd_sigmoid : ternary_device<hip_triadd_sigmoid, &device::add_sigmoid>
{
};
201
MIGRAPHX_REGISTER_OP(hip_triadd_sigmoid)
202
203
204
205

struct hip_triadd_tanh : ternary_device<hip_triadd_tanh, &device::add_tanh>
{
};
206
MIGRAPHX_REGISTER_OP(hip_triadd_tanh)
207
208
209
210

struct hip_add_relu : binary_device<hip_add_relu, &device::add_relu>
{
};
211
MIGRAPHX_REGISTER_OP(hip_add_relu)
212
213
214
215

struct hip_add_sigmoid : binary_device<hip_add_relu, &device::add_sigmoid>
{
};
216
MIGRAPHX_REGISTER_OP(hip_add_sigmoid)
217
218

struct hip_add_tanh : binary_device<hip_add_tanh, &device::add_tanh>
Paul's avatar
Paul committed
219
220
{
};
221
MIGRAPHX_REGISTER_OP(hip_add_tanh)
Paul's avatar
Paul committed
222

kahmed10's avatar
kahmed10 committed
223
224
225
struct hip_layernorm : unary_device<hip_layernorm, &device::layernorm>
{
};
226
MIGRAPHX_REGISTER_OP(hip_layernorm)
kahmed10's avatar
kahmed10 committed
227

kahmed10's avatar
kahmed10 committed
228
229
230
struct hip_gelu : unary_device<hip_gelu, &device::gelu>
{
};
231
MIGRAPHX_REGISTER_OP(hip_gelu)
kahmed10's avatar
kahmed10 committed
232
233
234
235

struct hip_add_gelu : binary_device<hip_add_gelu, &device::add_gelu>
{
};
236
MIGRAPHX_REGISTER_OP(hip_add_gelu)
kahmed10's avatar
kahmed10 committed
237
238
239
240

struct hip_gelu_new : unary_device<hip_gelu_new, &device::gelu_new>
{
};
241
MIGRAPHX_REGISTER_OP(hip_gelu_new)
kahmed10's avatar
kahmed10 committed
242
243
244
245

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

248
struct hip_mul_add : ternary_device<hip_mul_add, &device::mul_add>
Paul's avatar
Paul committed
249
250
{
};
251
MIGRAPHX_REGISTER_OP(hip_mul_add)
Paul's avatar
Paul committed
252

253
struct hip_mul_add_relu : ternary_device<hip_mul_add_relu, &device::mul_add_relu>
Paul's avatar
Paul committed
254
255
{
};
256
MIGRAPHX_REGISTER_OP(hip_mul_add_relu)
Paul's avatar
Paul committed
257

Paul's avatar
Paul committed
258
259
260
void move_broadcasted_back(std::vector<instruction_ref>& args)
{
    // Ensure the last arguments is the broadcasted one
Paul's avatar
Paul committed
261
    auto last = std::prev(args.end());
Paul's avatar
Paul committed
262
263
    auto it =
        std::find_if(args.begin(), last, [](auto arg) { return arg->get_shape().broadcasted(); });
Paul's avatar
Paul committed
264
265
    if(it != last)
        std::swap(*it, *std::prev(last));
Paul's avatar
Paul committed
266
267
268
269
270
}

void move_standard_front(std::vector<instruction_ref>& args)
{
    // Ensure the first arguments is the standard one
Paul's avatar
Paul committed
271
    auto last = std::prev(args.end());
Paul's avatar
Paul committed
272
273
    auto it =
        std::find_if(args.begin(), last, [](auto arg) { return arg->get_shape().standard(); });
Paul's avatar
Paul committed
274
    if(it != last)
Paul's avatar
Paul committed
275
276
277
        std::swap(*it, args.front());
}

kahmed10's avatar
kahmed10 committed
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
struct find_layernorm
{
    template <class... Ts>
    static auto multibroadcast_op(Ts... xs)
    {
        return match::name("multibroadcast")(match::arg(0)(xs...));
    }

    static auto x_minus_mean()
    {
        return match::name("gpu::sub")(
            match::arg(0)(match::any().bind("x")),
            match::arg(1)(multibroadcast_op(match::name("gpu::reduce_mean"))));
    }

    static auto variance()
    {
        return match::name("gpu::reduce_mean")(match::arg(0)(
            match::name("gpu::pow")(match::arg(0)(x_minus_mean()),
                                    match::arg(1)(multibroadcast_op(match::has_value(2.0f))))));
    }

    static auto layernorm_onnx()
    {
        return match::name("gpu::div")(
            match::arg(0)(x_minus_mean()),

            match::arg(1)(multibroadcast_op(
                match::name("gpu::sqrt")(match::arg(0)(match::name("gpu::add")(match::either_arg(
                    0, 1)(variance(), multibroadcast_op(match::has_value(1e-12f)))))))));
    }

    auto matcher() const { return layernorm_onnx(); }

    void apply(program& p, match::matcher_result r) const
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

        p.replace_instruction(ins, hip_layernorm{}, x_ins, args.back());
    }
};

kahmed10's avatar
kahmed10 committed
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
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
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
struct find_gelu
{

    static auto erf_fn()
    {
        return match::name("gpu::erf")(
            match::used_once(),
            match::arg(0)(match::used_once(),
                          match::name("gpu::mul")(match::either_arg(0, 1)(
                              match::none_of(match::has_value(M_SQRT1_2)).bind("x"),
                              match::has_value(M_SQRT1_2)))));
    }

    auto matcher() const
    {
        return match::name("gpu::mul")(match::either_arg(0, 1)(
            match::name("gpu::mul")(match::any_arg(0, 1)(match::args(match::has_value(0.5f)))),
            match::name("gpu::add")(
                match::used_once(),
                match::either_arg(0, 1)(erf_fn(), match::args(match::has_value(1.0f))))));
    }

    void apply(program& p, match::matcher_result r) const
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

        p.replace_instruction(ins, hip_gelu{}, x_ins, args.back());
    }
};

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

    void apply(program& p, match::matcher_result r) const
    {
        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();
        p.replace_instruction(ins, hip_add_gelu{}, args);
    }
};

struct find_gelu_new
{

    static auto pow_fn()
    {
        return match::name("gpu::pow")(match::used_once(),
                                       match::arg(1)(match::args(match::has_value(3.0f))));
    }

    static auto tanh_fn()
    {
        return match::name("gpu::tanh")(
            match::used_once(),
            match::arg(0)(match::name("gpu::mul")(match::either_arg(0, 1)(
                match::args(match::has_value(sqrt(M_2_PI))),
                match::name("gpu::add")(
                    match::any_arg(0, 1)(match::name("gpu::mul")(match::either_arg(0, 1)(
                        match::args(match::has_value(0.044715f)), pow_fn()))))))));
    }

    auto matcher() const
    {
        return match::name("gpu::mul")(
            match::used_once(),
            match::either_arg(0, 1)(
                match::any().bind("x"),
                match::name("gpu::add")(match::any_arg(0, 1)(match::name("gpu::mul")(
                    match::either_arg(0, 1)(match::args(match::has_value(0.5f)), tanh_fn()))))));
    }

    void apply(program& p, match::matcher_result r) const
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

        if(enabled(MIGRAPHX_DISABLE_FAST_GELU{}))
            p.replace_instruction(ins, hip_gelu_new{}, x_ins, args.back());
        else
            p.replace_instruction(ins, hip_gelu{}, x_ins, args.back());
    }
};

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

    void apply(program& p, match::matcher_result r) const
    {
        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();
        p.replace_instruction(ins, hip_add_gelu_new{}, args);
    }
};

kahmed10's avatar
kahmed10 committed
437
438
439
440
441
442
443
444
445
446
447
448
449
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"),
                                        match::name("hip::triadd"),
                                        match::any_of[match::inputs()](match::standard_shape()))
                              .bind("add")));
    }

    void apply(program& p, match::matcher_result r) const
    {
kahmed10's avatar
kahmed10 committed
450
451
452
453
454
455
456
457
458
459
        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
460
        if(add_ins->name() == "gpu::add")
kahmed10's avatar
kahmed10 committed
461
            p.replace_instruction(ins, hip_add_clip{}, add_args);
kahmed10's avatar
kahmed10 committed
462
        else if(add_ins->name() == "hip::triadd")
kahmed10's avatar
kahmed10 committed
463
            p.replace_instruction(ins, hip_triadd_clip{}, add_args);
kahmed10's avatar
kahmed10 committed
464
465
466
    }
};

467
struct find_add_unary
Paul's avatar
Paul committed
468
{
469
470
471
    std::string op_name;
    operation binary_add_op;
    operation ternary_add_op;
Paul's avatar
Paul committed
472
473
    auto matcher() const
    {
474
        return match::name(op_name)(match::arg(0)(
Paul's avatar
Paul committed
475
            match::used_once(),
Paul's avatar
Paul committed
476
477
478
479
480
            match::any_of(match::name("gpu::add"),
                          match::name("hip::triadd"),
                          match::any_of(match::name("@literal"),
                                        match::any_of[match::inputs()](match::standard_shape())))
                .bind("add")));
Paul's avatar
Paul committed
481
    }
Paul's avatar
Paul committed
482

Paul's avatar
Paul committed
483
484
    void apply(program& p, match::matcher_result r) const
    {
Paul's avatar
Paul committed
485
        auto add_ins = r.instructions["add"];
Paul's avatar
Paul committed
486
487
        auto ins     = r.result;
        auto args    = add_ins->inputs();
Paul's avatar
Paul committed
488
489
490
        move_standard_front(args);
        move_broadcasted_back(args);

Paul's avatar
Paul committed
491
        // Use the allocation from the relu operator
Paul's avatar
Paul committed
492
        args.back() = ins->inputs().back();
Paul's avatar
Paul committed
493
        if(add_ins->name() == "gpu::add")
494
            p.replace_instruction(ins, binary_add_op, args);
Paul's avatar
Paul committed
495
        else if(add_ins->name() == "hip::triadd")
496
            p.replace_instruction(ins, ternary_add_op, args);
Paul's avatar
Paul committed
497
498
499
    }
};

Paul's avatar
Paul committed
500
struct find_triadd
Paul's avatar
Paul committed
501
502
503
{
    auto matcher() const
    {
Paul's avatar
Paul committed
504
        return match::name("gpu::add")(match::either_arg(0, 1)(
Paul's avatar
Paul committed
505
            match::name("gpu::add")(match::used_once()).bind("add"),
Paul's avatar
Paul committed
506
507
508
            match::any(match::any_of(match::name("@literal"),
                                     match::any_of[match::inputs()](match::standard_shape())))
                .bind("input")));
Paul's avatar
Paul committed
509
510
511
512
    }

    void apply(program& p, match::matcher_result r) const
    {
Paul's avatar
Paul committed
513
514
515
516
        auto add_ins   = r.instructions["add"];
        auto input_ins = r.instructions["input"];
        auto ins       = r.result;
        auto args      = add_ins->inputs();
517
518
        assert(add_ins != input_ins);

Paul's avatar
Paul committed
519
520
521
522
        auto is_broadcasted = [](auto arg) { return arg->get_shape().broadcasted(); };
        if(std::count_if(args.begin(), args.end(), is_broadcasted) > 1)
            return;
        args.insert(args.begin(), input_ins);
Paul's avatar
Paul committed
523
524
525
        move_standard_front(args);
        move_broadcasted_back(args);

Paul's avatar
Paul committed
526
527
        args.back() = ins->inputs().back();
        p.replace_instruction(ins, hip_triadd{}, args);
Paul's avatar
Paul committed
528
    }
Paul's avatar
Paul committed
529
530
};

Paul's avatar
Paul committed
531
532
533
534
struct find_mul_add
{
    auto matcher() const
    {
Paul's avatar
Paul committed
535
536
        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
537
538
539
540
    }

    void apply(program& p, match::matcher_result r) const
    {
Paul's avatar
Paul committed
541
542
543
544
        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
545
546
547
548
549
550
551
552
553
554
555
        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();
        p.replace_instruction(ins, hip_mul_add{}, args);
    }
};

Paul's avatar
Paul committed
556
557
558
559
struct find_mul_add_relu
{
    auto matcher() const
    {
Paul's avatar
Paul committed
560
561
        return match::name("gpu::relu")(
            match::arg(0)(match::name("hip::mul_add")(match::used_once()).bind("mul_add")));
Paul's avatar
Paul committed
562
563
564
565
566
    }

    void apply(program& p, match::matcher_result r) const
    {
        auto mul_add_ins = r.instructions["mul_add"];
Paul's avatar
Paul committed
567
568
        auto ins         = r.result;
        auto args        = mul_add_ins->inputs();
Paul's avatar
Paul committed
569
570
571
572
573
574
575

        // Use the allocation from the relu operator
        args.back() = ins->inputs().back();
        p.replace_instruction(ins, hip_mul_add_relu{}, args);
    }
};

Paul's avatar
Paul committed
576
577
578
579
580
581
582
struct miopen_conv_bias
{
    op::convolution op;
    fusion f;
    fusion::op_t conv;
    fusion::op_t bias;

Paul's avatar
Paul committed
583
584
585
586
587
588
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return op::convolution::reflect(self.op, f);
    }

589
590
    miopen_conv_bias() = default;

Paul's avatar
Paul committed
591
    miopen_conv_bias(op::convolution c, const shape& input, const shape& weights, const shape& b)
592
        : op(std::move(c)), f(input)
Paul's avatar
Paul committed
593
    {
Paul's avatar
Paul committed
594
595
        conv = f.create_conv(op, weights);
        bias = f.create_bias(b);
Paul's avatar
Paul committed
596
597
598
599
600
601
602
603
604
    }

    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
        return op.compute_shape({inputs.at(0), inputs.at(1)});
    }
Paul's avatar
Paul committed
605
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
Paul's avatar
Paul committed
606
    {
Paul's avatar
Paul committed
607
        auto fargs  = make_fused_args();
Paul's avatar
Paul committed
608
        float alpha = 1;
Paul's avatar
Paul committed
609
        float beta  = 0;
Paul's avatar
Paul committed
610
611
        miopenSetOpArgsConvForward(fargs.get(), conv, &alpha, &beta, args[1].implicit());
        miopenSetOpArgsBiasForward(fargs.get(), bias, &alpha, &beta, args[3].implicit());
Paul's avatar
Paul committed
612
        return f.execute(ctx, fargs, args[0], args[4]);
Paul's avatar
Paul committed
613
614
    }

Paul's avatar
Paul committed
615
616
    void finalize(context& ctx, const shape&, const std::vector<shape>&) { f.compile(ctx); }
    shape get_workspace(context& ctx) { return f.get_workspace(ctx); }
Paul's avatar
Paul committed
617
618
619
620
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
Paul's avatar
Paul committed
621
};
622
MIGRAPHX_REGISTER_OP(miopen_conv_bias)
Paul's avatar
Paul committed
623

Paul's avatar
Add cbr  
Paul committed
624
625
626
627
628
629
struct miopen_conv_bias_relu
{
    op::convolution op;
    fusion f;
    fusion::op_t conv;
    fusion::op_t bias;
Paul's avatar
Paul committed
630
    fusion::op_t relu;
Paul's avatar
Add cbr  
Paul committed
631

Paul's avatar
Paul committed
632
633
634
635
636
637
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return op::convolution::reflect(self.op, f);
    }

638
639
    miopen_conv_bias_relu() = default;

Paul's avatar
Paul committed
640
641
642
643
    miopen_conv_bias_relu(op::convolution c,
                          const shape& input,
                          const shape& weights,
                          const shape& b)
644
        : op(std::move(c)), f(input)
Paul's avatar
Add cbr  
Paul committed
645
    {
Paul's avatar
Paul committed
646
647
648
        conv = f.create_conv(op, weights);
        bias = f.create_bias(b);
        relu = f.create_relu();
Paul's avatar
Add cbr  
Paul committed
649
650
651
652
653
654
655
656
657
    }

    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
        return op.compute_shape({inputs.at(0), inputs.at(1)});
    }
Paul's avatar
Paul committed
658
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
Paul's avatar
Add cbr  
Paul committed
659
660
    {
        auto fargs  = make_fused_args();
Paul's avatar
Paul committed
661
        float alpha = 1;
Paul's avatar
Paul committed
662
        float beta  = 0;
Paul's avatar
Add cbr  
Paul committed
663
664
        miopenSetOpArgsConvForward(fargs.get(), conv, &alpha, &beta, args[1].implicit());
        miopenSetOpArgsBiasForward(fargs.get(), bias, &alpha, &beta, args[3].implicit());
Paul's avatar
Paul committed
665
666
        miopenSetOpArgsActivForward(fargs.get(), relu, &alpha, &beta, 0, 0, 0);
        return f.execute(ctx, fargs, args[0], args[4]);
Paul's avatar
Add cbr  
Paul committed
667
    }
Paul's avatar
Paul committed
668
669
    void finalize(context& ctx, const shape&, const std::vector<shape>&) { f.compile(ctx); }
    shape get_workspace(context& ctx) { return f.get_workspace(ctx); }
Paul's avatar
Paul committed
670
671
672
673
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
Paul's avatar
Add cbr  
Paul committed
674
};
675
MIGRAPHX_REGISTER_OP(miopen_conv_bias_relu)
Paul's avatar
Add cbr  
Paul committed
676

Paul's avatar
Paul committed
677
template <class... Ms>
Paul's avatar
Add cbr  
Paul committed
678
679
auto conv_bias(Ms... ms)
{
Paul's avatar
Paul committed
680
    return match::name("gpu::add")(
Paul's avatar
Paul committed
681
682
        match::either_arg(0, 1)(bias_shape(match::used_once()).bind("bias"),
                                fusable_conv(match::used_once()).bind("conv")),
Paul's avatar
Paul committed
683
        ms...);
Paul's avatar
Paul committed
684
685
}

Paul's avatar
Paul committed
686
template <class Op>
Paul's avatar
Paul committed
687
688
689
690
691
692
693
694
695
696
697
void apply_conv_bias(context& ctx, program& p, match::matcher_result r)
{
    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);

Paul's avatar
Paul committed
698
    Op cb{conv_op, input_ins->get_shape(), weights_ins->get_shape(), bias_ins->get_shape()};
Paul's avatar
Paul committed
699
    // TODO: Insert ws allocation
Paul's avatar
Paul committed
700
    auto ws = cb.get_workspace(ctx);
Paul's avatar
Paul committed
701
    (void)ws;
Paul's avatar
Paul committed
702
    p.replace_instruction(ins, cb, input_ins, weights_ins, old_ws_ins, bias_ins, alloc_ins);
Paul's avatar
Add cbr  
Paul committed
703
704
}

Paul's avatar
Paul committed
705
struct find_conv_bias
Paul's avatar
Paul committed
706
{
Paul's avatar
Paul committed
707
    context* ctx = nullptr;
Paul's avatar
Paul committed
708
709
    auto matcher() const
    {
kahmed10's avatar
kahmed10 committed
710
711
        return conv_bias(match::none_of(
            match::output(match::name(std::unordered_set<std::string>{"gpu::relu"}))));
Paul's avatar
Paul committed
712
713
714
715
    }

    void apply(program& p, match::matcher_result r) const
    {
Paul's avatar
Paul committed
716
        apply_conv_bias<miopen_conv_bias>(*ctx, p, std::move(r));
Paul's avatar
Paul committed
717
718
719
    }
};

Paul's avatar
Paul committed
720
struct find_conv_bias_relu
Paul's avatar
Add cbr  
Paul committed
721
722
{
    context* ctx = nullptr;
Paul's avatar
Paul committed
723
    auto matcher() const { return match::name("gpu::relu")(match::arg(0)(conv_bias())); }
Paul's avatar
Add cbr  
Paul committed
724
725
726

    void apply(program& p, match::matcher_result r) const
    {
Paul's avatar
Paul committed
727
        apply_conv_bias<miopen_conv_bias_relu>(*ctx, p, std::move(r));
Paul's avatar
Add cbr  
Paul committed
728
729
730
    }
};

Paul's avatar
Paul committed
731
732
void fuse_ops::apply(program& p) const
{
kahmed10's avatar
kahmed10 committed
733
734
    match::find_matches(p, find_gelu{}, find_gelu_new{});
    run_passes(p, {dead_code_elimination{}});
Paul's avatar
Paul committed
735
    match::find_matches(p, find_triadd{});
736
    match::find_matches(p,
kahmed10's avatar
kahmed10 committed
737
                        find_layernorm{},
738
739
740
741
742
743
744
745
746
747
                        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{});
Paul's avatar
Paul committed
748
    // clang-format on
Paul's avatar
Paul committed
749
}
Paul's avatar
Paul committed
750
751

} // namespace gpu
Paul's avatar
Paul committed
752
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
753
} // namespace migraphx