fuse_ops.cpp 38.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>
Paul's avatar
Paul committed
51
#include <migraphx/permutation.hpp>
52
#include <migraphx/make_op.hpp>
kahmed10's avatar
kahmed10 committed
53
#include <migraphx/op/clip.hpp>
Paul's avatar
Paul committed
54
#include <migraphx/op/contiguous.hpp>
kahmed10's avatar
kahmed10 committed
55
#include <cmath>
56
#include <set>
Paul's avatar
Paul committed
57
58

namespace migraphx {
Paul's avatar
Paul committed
59
inline namespace MIGRAPHX_INLINE_NS {
Paul's avatar
Paul committed
60
61
namespace gpu {

62
63
MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_DISABLE_MIOPEN_FUSION)

Paul's avatar
Paul committed
64
65
66
67
68
69
70
71
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
72
    template <class T>
Paul's avatar
Paul committed
73
74
75
76
77
78
79
    auto keep_alive(T x)
    {
        auto result = share(std::move(x));
        storage.push_back(result);
        return result;
    }

80
81
    fusion() = default;

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

91
92
    bool empty() const { return fp == nullptr; }

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

103
104
105
106
107
    auto get() const
    {
        assert(fp);
        return fp.get();
    }
Paul's avatar
Paul committed
108
109
110

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

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

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

    shape get_workspace(context&)
    {
145
        // assert(fp);
Paul's avatar
Paul committed
146
147
148
149
150
        // 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
151
152
        // miopenFusionPlanGetWorkSpaceSize(ctx.get_stream().get_miopen(), fp.get(), &ws_size,
        // algo);
Paul's avatar
Paul committed
153
154
155
        return shape{shape::int8_type, {ws_size}};
    }

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

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

184
185
186
187
188
189
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
190
MIGRAPHX_PRED_MATCHER(bias_shape, instruction_ref ins)
Paul's avatar
Paul committed
191
192
{
    auto&& s = ins->get_shape();
Paul's avatar
Paul committed
193
194
    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
195
196
}

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

    // 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
221
    auto op = conv.op;
222
223
    // 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
224
       wei.lens()[3] != 3 and contains({{1, 1}}, op.stride))
225
        return false;
kahmed10's avatar
kahmed10 committed
226
    return contains({{0, 0, 0, 0}, {1, 1, 1, 1}, {2, 2, 2, 2}}, op.padding) and
227
           contains({{0, 0}, {1, 1}}, op.stride) and contains({{1, 1}}, op.dilation);
Paul's avatar
Paul committed
228
229
}

230
struct hip_triadd : ternary_device<hip_triadd, &device::add>
Paul's avatar
Paul committed
231
232
{
};
233
MIGRAPHX_REGISTER_OP(hip_triadd)
Paul's avatar
Paul committed
234

235
struct hip_triadd_clip : quinary_device<hip_triadd_clip, &device::add_clip>
kahmed10's avatar
kahmed10 committed
236
237
{
};
238
MIGRAPHX_REGISTER_OP(hip_triadd_clip)
kahmed10's avatar
kahmed10 committed
239

240
struct hip_add_clip : quaternary_device<hip_add_clip, &device::add_clip>
kahmed10's avatar
kahmed10 committed
241
242
{
};
243
MIGRAPHX_REGISTER_OP(hip_add_clip)
kahmed10's avatar
kahmed10 committed
244

245
struct hip_triadd_relu : ternary_device<hip_triadd_relu, &device::add_relu>
Paul's avatar
Paul committed
246
247
{
};
248
MIGRAPHX_REGISTER_OP(hip_triadd_relu)
Paul's avatar
Paul committed
249

250
251
252
struct hip_triadd_sigmoid : ternary_device<hip_triadd_sigmoid, &device::add_sigmoid>
{
};
253
MIGRAPHX_REGISTER_OP(hip_triadd_sigmoid)
254
255
256
257

struct hip_triadd_tanh : ternary_device<hip_triadd_tanh, &device::add_tanh>
{
};
258
MIGRAPHX_REGISTER_OP(hip_triadd_tanh)
259
260
261
262

struct hip_add_relu : binary_device<hip_add_relu, &device::add_relu>
{
};
263
MIGRAPHX_REGISTER_OP(hip_add_relu)
264
265
266
267

struct hip_add_sigmoid : binary_device<hip_add_relu, &device::add_sigmoid>
{
};
268
MIGRAPHX_REGISTER_OP(hip_add_sigmoid)
269
270

struct hip_add_tanh : binary_device<hip_add_tanh, &device::add_tanh>
Paul's avatar
Paul committed
271
272
{
};
273
MIGRAPHX_REGISTER_OP(hip_add_tanh)
Paul's avatar
Paul committed
274

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

Paul Fultz II's avatar
Paul Fultz II committed
282
283
struct hip_triadd_layernorm : ternary_device<hip_triadd_layernorm, &device::triadd_layernorm>
{
284
285
286
287
288
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        check_shapes{inputs, *this}.has(4).standard();
        return inputs[0];
    }
Paul Fultz II's avatar
Paul Fultz II committed
289
290
291
292
293
    // 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
294
295
296
struct hip_gelu : unary_device<hip_gelu, &device::gelu>
{
};
297
MIGRAPHX_REGISTER_OP(hip_gelu)
kahmed10's avatar
kahmed10 committed
298
299
300
301

struct hip_add_gelu : binary_device<hip_add_gelu, &device::add_gelu>
{
};
302
MIGRAPHX_REGISTER_OP(hip_add_gelu)
kahmed10's avatar
kahmed10 committed
303
304
305
306

struct hip_gelu_new : unary_device<hip_gelu_new, &device::gelu_new>
{
};
307
MIGRAPHX_REGISTER_OP(hip_gelu_new)
kahmed10's avatar
kahmed10 committed
308
309
310
311

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

314
struct hip_mul_add : ternary_device<hip_mul_add, &device::mul_add>
Paul's avatar
Paul committed
315
316
{
};
317
MIGRAPHX_REGISTER_OP(hip_mul_add)
Paul's avatar
Paul committed
318

319
struct hip_mul_add_relu : ternary_device<hip_mul_add_relu, &device::mul_add_relu>
Paul's avatar
Paul committed
320
321
{
};
322
MIGRAPHX_REGISTER_OP(hip_mul_add_relu)
Paul's avatar
Paul committed
323

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

void move_standard_front(std::vector<instruction_ref>& args)
{
    // Ensure the first arguments is the standard one
Paul's avatar
Paul committed
337
    auto last = std::prev(args.end());
Paul's avatar
Paul committed
338
339
    auto it =
        std::find_if(args.begin(), last, [](auto arg) { return arg->get_shape().standard(); });
Paul's avatar
Paul committed
340
    if(it != last)
Paul's avatar
Paul committed
341
342
343
        std::swap(*it, args.front());
}

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

kahmed10's avatar
kahmed10 committed
346
347
struct find_layernorm
{
348
    auto matcher() const { return match::layernorm(&gpu_name); }
kahmed10's avatar
kahmed10 committed
349

350
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
351
352
353
354
355
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

356
357
358
359
360
361
362
363
364
        // 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;

365
        m.replace_instruction(ins, hip_layernorm{}, x_ins, args.back());
kahmed10's avatar
kahmed10 committed
366
367
368
    }
};

Paul Fultz II's avatar
Paul Fultz II committed
369
370
371
372
373
374
375
376
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()))));
    }

377
    void apply(module& m, const match::matcher_result& r) const
Paul Fultz II's avatar
Paul Fultz II committed
378
379
380
    {
        auto ins    = r.result;
        auto triadd = ins->inputs().front();
381
        m.replace_instruction(ins, hip_triadd_layernorm{}, triadd->inputs());
Paul Fultz II's avatar
Paul Fultz II committed
382
383
384
    }
};

kahmed10's avatar
kahmed10 committed
385
386
struct find_gelu
{
387
    auto matcher() const { return match::gelu_erf(&gpu_name); }
kahmed10's avatar
kahmed10 committed
388

389
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
390
391
392
393
394
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

395
        m.replace_instruction(ins, hip_gelu{}, x_ins, args.back());
kahmed10's avatar
kahmed10 committed
396
397
398
399
400
401
402
403
404
405
    }
};

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

406
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
407
408
409
410
411
412
413
414
    {
        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();
415
        m.replace_instruction(ins, hip_add_gelu{}, args);
kahmed10's avatar
kahmed10 committed
416
417
418
419
420
    }
};

struct find_gelu_new
{
kahmed10's avatar
kahmed10 committed
421
    bool fast_math = true;
kahmed10's avatar
kahmed10 committed
422

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

425
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
426
427
428
429
430
    {
        auto ins   = r.result;
        auto x_ins = r.instructions["x"];
        auto args  = ins->inputs();

Paul Fultz II's avatar
Paul Fultz II committed
431
        if(fast_math)
432
            m.replace_instruction(ins, hip_gelu{}, x_ins, args.back());
Paul Fultz II's avatar
Paul Fultz II committed
433
        else
434
            m.replace_instruction(ins, hip_gelu_new{}, x_ins, args.back());
kahmed10's avatar
kahmed10 committed
435
436
437
438
439
440
441
442
443
444
    }
};

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

445
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
446
447
448
449
450
451
452
453
    {
        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();
454
        m.replace_instruction(ins, hip_add_gelu_new{}, args);
kahmed10's avatar
kahmed10 committed
455
456
457
    }
};

kahmed10's avatar
kahmed10 committed
458
459
460
461
462
463
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
464
                                        match::name("gpu::triadd"),
kahmed10's avatar
kahmed10 committed
465
466
467
468
                                        match::any_of[match::inputs()](match::standard_shape()))
                              .bind("add")));
    }

469
    void apply(module& m, const match::matcher_result& r) const
kahmed10's avatar
kahmed10 committed
470
    {
kahmed10's avatar
kahmed10 committed
471
472
473
474
475
476
477
478
479
480
        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
481
        if(add_ins->name() == "gpu::add")
482
            m.replace_instruction(ins, hip_add_clip{}, add_args);
kahmed10's avatar
kahmed10 committed
483
        else if(add_ins->name() == "gpu::triadd")
484
            m.replace_instruction(ins, hip_triadd_clip{}, add_args);
kahmed10's avatar
kahmed10 committed
485
486
487
    }
};

488
struct find_add_unary
Paul's avatar
Paul committed
489
{
490
491
492
    std::string op_name;
    operation binary_add_op;
    operation ternary_add_op;
Paul's avatar
Paul committed
493
494
    auto matcher() const
    {
495
        return match::name(op_name)(match::arg(0)(
Paul's avatar
Paul committed
496
            match::used_once(),
Paul's avatar
Paul committed
497
            match::any_of(match::name("gpu::add"),
kahmed10's avatar
kahmed10 committed
498
                          match::name("gpu::triadd"),
Paul's avatar
Paul committed
499
500
501
                          match::any_of(match::name("@literal"),
                                        match::any_of[match::inputs()](match::standard_shape())))
                .bind("add")));
Paul's avatar
Paul committed
502
    }
Paul's avatar
Paul committed
503

504
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
505
    {
Paul's avatar
Paul committed
506
        auto add_ins = r.instructions["add"];
Paul's avatar
Paul committed
507
508
        auto ins     = r.result;
        auto args    = add_ins->inputs();
Paul's avatar
Paul committed
509
510
511
        move_standard_front(args);
        move_broadcasted_back(args);

Paul's avatar
Paul committed
512
        // Use the allocation from the relu operator
Paul's avatar
Paul committed
513
        args.back() = ins->inputs().back();
Paul's avatar
Paul committed
514
        if(add_ins->name() == "gpu::add")
515
            m.replace_instruction(ins, binary_add_op, args);
kahmed10's avatar
kahmed10 committed
516
        else if(add_ins->name() == "gpu::triadd")
517
            m.replace_instruction(ins, ternary_add_op, args);
Paul's avatar
Paul committed
518
519
520
    }
};

Paul's avatar
Paul committed
521
struct find_triadd
Paul's avatar
Paul committed
522
523
524
{
    auto matcher() const
    {
Paul's avatar
Paul committed
525
        return match::name("gpu::add")(match::either_arg(0, 1)(
Paul's avatar
Paul committed
526
            match::name("gpu::add")(match::used_once()).bind("add"),
Paul's avatar
Paul committed
527
528
529
            match::any(match::any_of(match::name("@literal"),
                                     match::any_of[match::inputs()](match::standard_shape())))
                .bind("input")));
Paul's avatar
Paul committed
530
531
    }

532
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
533
    {
Paul's avatar
Paul committed
534
535
536
537
        auto add_ins   = r.instructions["add"];
        auto input_ins = r.instructions["input"];
        auto ins       = r.result;
        auto args      = add_ins->inputs();
538

Paul's avatar
Paul committed
539
        auto is_broadcasted = [](auto arg) { return arg->get_shape().broadcasted(); };
540
        if(std::count_if(args.begin(), args.end(), is_broadcasted) > 2)
Paul's avatar
Paul committed
541
542
            return;
        args.insert(args.begin(), input_ins);
Paul's avatar
Paul committed
543
544
545
        move_standard_front(args);
        move_broadcasted_back(args);

Paul's avatar
Paul committed
546
        args.back() = ins->inputs().back();
547
        m.replace_instruction(ins, hip_triadd{}, args);
Paul's avatar
Paul committed
548
    }
Paul's avatar
Paul committed
549
550
};

Paul's avatar
Paul committed
551
552
553
554
struct find_mul_add
{
    auto matcher() const
    {
Paul's avatar
Paul committed
555
556
        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
557
558
    }

559
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
560
    {
Paul's avatar
Paul committed
561
562
563
564
        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
565
566
567
568
569
570
571
        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();
572
        m.replace_instruction(ins, hip_mul_add{}, args);
Paul's avatar
Paul committed
573
574
575
    }
};

Paul's avatar
Paul committed
576
577
578
579
struct find_mul_add_relu
{
    auto matcher() const
    {
Paul's avatar
Paul committed
580
        return match::name("gpu::relu")(
kahmed10's avatar
kahmed10 committed
581
            match::arg(0)(match::name("gpu::mul_add")(match::used_once()).bind("mul_add")));
Paul's avatar
Paul committed
582
583
    }

584
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
585
586
    {
        auto mul_add_ins = r.instructions["mul_add"];
Paul's avatar
Paul committed
587
588
        auto ins         = r.result;
        auto args        = mul_add_ins->inputs();
Paul's avatar
Paul committed
589
590
591

        // Use the allocation from the relu operator
        args.back() = ins->inputs().back();
592
        m.replace_instruction(ins, hip_mul_add_relu{}, args);
Paul's avatar
Paul committed
593
594
595
    }
};

596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
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"));
    }

621
622
623
624
625
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }

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
704
705
706
707
708
709
710
    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();
    }
};
711
MIGRAPHX_REGISTER_OP(miopen_fusion)
712

Paul's avatar
Paul committed
713
714
715
struct miopen_conv_bias
{
    op::convolution op;
716
    fusion fp         = {};
717
718
    fusion::op_t conv = {};
    fusion::op_t bias = {};
Paul's avatar
Paul committed
719

Paul's avatar
Paul committed
720
721
722
723
724
725
    template <class Self, class F>
    static auto reflect(Self& self, F f)
    {
        return op::convolution::reflect(self.op, f);
    }

Paul's avatar
Paul committed
726
727
728
729
730
    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
731
        return op.normalize_compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
732
    }
Paul's avatar
Paul committed
733
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
Paul's avatar
Paul committed
734
    {
Paul's avatar
Paul committed
735
        auto fargs  = make_fused_args();
Paul's avatar
Paul committed
736
        float alpha = 1;
Paul's avatar
Paul committed
737
        float beta  = 0;
Paul's avatar
Paul committed
738
739
        miopenSetOpArgsConvForward(fargs.get(), conv, &alpha, &beta, args[1].implicit());
        miopenSetOpArgsBiasForward(fargs.get(), bias, &alpha, &beta, args[3].implicit());
740
        return fp.execute(ctx, fargs, args[0], args[4]);
Paul's avatar
Paul committed
741
742
    }

743
744
    void finalize(context& ctx, const shape&, const std::vector<shape>& inputs)
    {
745
746
747
748
        fp   = fusion(inputs[0]);
        conv = fp.create_conv(op, inputs[1]);
        bias = fp.create_bias(inputs[3]);
        if(not fp.compile(ctx))
749
            MIGRAPHX_THROW("Failed to compile fusion plan");
750
751
    }

752
    shape get_workspace(context& ctx) { return fp.get_workspace(ctx); }
Paul's avatar
Paul committed
753
754
755
756
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
Paul's avatar
Paul committed
757
};
758
MIGRAPHX_REGISTER_OP(miopen_conv_bias)
Paul's avatar
Paul committed
759

Paul's avatar
Add cbr  
Paul committed
760
761
762
struct miopen_conv_bias_relu
{
    op::convolution op;
763
    fusion fp         = {};
764
765
766
    fusion::op_t conv = {};
    fusion::op_t bias = {};
    fusion::op_t relu = {};
Paul's avatar
Add cbr  
Paul committed
767

Paul's avatar
Paul committed
768
769
770
771
772
773
    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
774
775
776
777
778
    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
779
        return op.normalize_compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Add cbr  
Paul committed
780
    }
Paul's avatar
Paul committed
781
    argument compute(context& ctx, const shape&, const std::vector<argument>& args) const
Paul's avatar
Add cbr  
Paul committed
782
783
    {
        auto fargs  = make_fused_args();
Paul's avatar
Paul committed
784
        float alpha = 1;
Paul's avatar
Paul committed
785
        float beta  = 0;
Paul's avatar
Add cbr  
Paul committed
786
787
        miopenSetOpArgsConvForward(fargs.get(), conv, &alpha, &beta, args[1].implicit());
        miopenSetOpArgsBiasForward(fargs.get(), bias, &alpha, &beta, args[3].implicit());
Paul's avatar
Paul committed
788
        miopenSetOpArgsActivForward(fargs.get(), relu, &alpha, &beta, 0, 0, 0);
789
        return fp.execute(ctx, fargs, args[0], args[4]);
Paul's avatar
Add cbr  
Paul committed
790
    }
791
792
    void finalize(context& ctx, const shape&, const std::vector<shape>& inputs)
    {
793
794
795
796
797
        fp   = fusion(inputs[0]);
        conv = fp.create_conv(op, inputs[1]);
        bias = fp.create_bias(inputs[3]);
        relu = fp.create_relu();
        fp.compile(ctx);
798
799
    }

800
    shape get_workspace(context& ctx) { return fp.get_workspace(ctx); }
Paul's avatar
Paul committed
801
802
803
804
    std::ptrdiff_t output_alias(const std::vector<shape>& shapes) const
    {
        return shapes.size() - 1;
    }
Paul's avatar
Add cbr  
Paul committed
805
};
806
MIGRAPHX_REGISTER_OP(miopen_conv_bias_relu)
Paul's avatar
Add cbr  
Paul committed
807

Paul's avatar
Paul committed
808
template <class... Ms>
Paul's avatar
Add cbr  
Paul committed
809
810
auto conv_bias(Ms... ms)
{
Paul's avatar
Paul committed
811
    return match::name("gpu::add")(
Paul's avatar
Paul committed
812
813
        match::either_arg(0, 1)(bias_shape(match::used_once()).bind("bias"),
                                fusable_conv(match::used_once()).bind("conv")),
Paul's avatar
Paul committed
814
        ms...);
Paul's avatar
Paul committed
815
816
}

Paul's avatar
Paul committed
817
template <class Op>
818
void apply_conv_bias(context& ctx, module& m, const match::matcher_result& r)
Paul's avatar
Paul committed
819
820
821
822
823
824
825
826
827
828
{
    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);

829
    Op cb{conv_op};
Paul's avatar
Paul committed
830
    // TODO: Insert ws allocation
Paul's avatar
Paul committed
831
    auto ws = cb.get_workspace(ctx);
Paul's avatar
Paul committed
832
    (void)ws;
833
    m.replace_instruction(ins, cb, input_ins, weights_ins, old_ws_ins, bias_ins, alloc_ins);
Paul's avatar
Add cbr  
Paul committed
834
835
}

836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
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
855
struct find_conv_bias
Paul's avatar
Paul committed
856
{
Paul's avatar
Paul committed
857
    context* ctx = nullptr;
Paul's avatar
Paul committed
858
859
    auto matcher() const
    {
kahmed10's avatar
kahmed10 committed
860
861
        return conv_bias(match::none_of(
            match::output(match::name(std::unordered_set<std::string>{"gpu::relu"}))));
Paul's avatar
Paul committed
862
863
    }

864
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Paul committed
865
    {
866
        apply_conv_bias<miopen_conv_bias>(*ctx, m, r);
Paul's avatar
Paul committed
867
868
869
    }
};

Paul's avatar
Paul committed
870
struct find_conv_bias_relu
Paul's avatar
Add cbr  
Paul committed
871
872
{
    context* ctx = nullptr;
Paul's avatar
Paul committed
873
    auto matcher() const { return match::name("gpu::relu")(match::arg(0)(conv_bias())); }
Paul's avatar
Add cbr  
Paul committed
874

875
    void apply(module& m, const match::matcher_result& r) const
Paul's avatar
Add cbr  
Paul committed
876
    {
877
        apply_conv_bias<miopen_conv_bias_relu>(*ctx, m, r);
Paul's avatar
Add cbr  
Paul committed
878
879
    }
};
880

881
882
883
884
885
886
887
888
889
890
891
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")));
    }

892
    void apply(module& m, const match::matcher_result& r) const
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
    {
        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);
    }
};

920
921
922
923
924
925
926
927
928
929
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")));
    }

930
    void apply(module& m, const match::matcher_result& r) const
931
932
933
934
935
936
937
938
    {
        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
939
        if(not float_equal(gemm.beta, 0))
940
941
942
943
944
945
946
947
            return;

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

        auto copy_ins = c_ins;

        // Insert copy
948
        if(ins == m.end() or c_ins->outputs().size() > 1 or c_ins->inputs().empty())
949
        {
950
            copy_ins = m.insert_instruction(ins, hip_copy{}, c_ins, ins->inputs().back());
951
952
953
954
        }
        inputs.push_back(copy_ins);
        inputs.push_back(copy_ins);

955
        gemm.beta = 1;
956
        m.replace_instruction(ins, gemm, inputs);
957
958
959
    }
};

960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
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
    {
Paul's avatar
Paul committed
977
        return precompile_name("pointwise")(
978
            match::nargs(3),
Paul's avatar
Format  
Paul committed
979
            match::either_arg(0, 1)(
Paul's avatar
Paul committed
980
                match::any_of(match::standard_shape(), match::is_constant()).bind("c"),
Paul's avatar
Format  
Paul committed
981
                match::name("gpu::gemm")(match::nargs(3), match::used_once()).bind("gemm")));
982
983
    }

Paul's avatar
Paul committed
984
985
986
987
    // TODO: Move to matcher.hpp
    static auto match_param(const std::string& name)
    {
        return match::make_basic_pred_matcher([=](auto ins) {
Paul's avatar
Format  
Paul committed
988
            if(ins->name() != "@param")
Paul's avatar
Paul committed
989
990
991
992
993
994
                return false;
            auto p = any_cast<builtin::param>(ins->get_operator());
            return p.parameter == name;
        });
    }

Paul's avatar
Format  
Paul committed
995
    template <class M>
Paul's avatar
Paul committed
996
997
    static auto match_mul_const(M m, const std::string& var)
    {
Paul's avatar
Format  
Paul committed
998
999
        return match::name("mul")(match::either_arg(0, 1)(match::name("@literal").bind(var), m))
            .bind(var + "_mul");
Paul's avatar
Paul committed
1000
1001
1002
1003
    }

    static auto match_add(const std::string& input, const std::string& output)
    {
Paul's avatar
Format  
Paul committed
1004
1005
1006
1007
1008
1009
        auto param     = match::name("@param");
        auto add       = match::name("add")(match::args(param, param));
        auto inner_mul = match::any_of(match_mul_const(match_param(input), "alpha"),
                                       match_mul_const(match_param(output), "beta"));
        auto mul_add   = match::name("add")(match::either_arg(0, 1)(inner_mul, param));
        auto add_mul   = match_mul_const(add, "gamma");
Paul's avatar
Paul committed
1010
1011
1012
        return match::name("@return")(match::args(match::any_of(add, mul_add, add_mul)));
    }

Paul's avatar
Format  
Paul committed
1013
    static float get_float(instruction_ref ins) { return ins->get_literal().at<float>(); }
Paul's avatar
Paul committed
1014

Paul's avatar
Format  
Paul committed
1015
    template <class Gemm>
Paul's avatar
Paul committed
1016
1017
1018
1019
1020
1021
1022
    static bool update_gemm(Gemm& gemm, module_ref pm, unsigned input)
    {
        auto names = pm->get_parameter_names();
        if(names.size() != 2)
            return false;
        std::sort(names.begin(), names.end());
        unsigned output = input == 0 ? 1 : 0;
Paul's avatar
Format  
Paul committed
1023
1024
1025
        auto mr         = match::match_instruction(
            *pm, std::prev(pm->end()), match_add(names[input], names[output]));
        if(mr.result == pm->end())
Paul's avatar
Paul committed
1026
            return false;
Paul's avatar
Format  
Paul committed
1027
        if(contains(mr.instructions, "alpha_mul"))
Paul's avatar
Paul committed
1028
            gemm.alpha *= get_float(mr.instructions["alpha"]);
Paul's avatar
Format  
Paul committed
1029
        else if(contains(mr.instructions, "beta_mul"))
Paul's avatar
Paul committed
1030
            gemm.beta *= get_float(mr.instructions["beta"]);
Paul's avatar
Format  
Paul committed
1031
        else if(contains(mr.instructions, "gamma_mul"))
Paul's avatar
Paul committed
1032
1033
1034
1035
1036
1037
1038
        {
            gemm.alpha *= get_float(mr.instructions["gamma"]);
            gemm.beta *= get_float(mr.instructions["gamma"]);
        }
        return true;
    }

1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
    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;
Paul's avatar
Paul committed
1050
1051
        gemm.beta = 1;

Paul's avatar
Format  
Paul committed
1052
1053
        if(not update_gemm(
               gemm, ins->module_inputs().front(), ins->inputs().front() == gemm_ins ? 0 : 1))
Paul's avatar
Paul committed
1054
            return;
1055

Paul's avatar
Paul committed
1056
        // const-fold input if not standard shape since rocblas can't handle it
Paul's avatar
Format  
Paul committed
1057
        if(not c_ins->get_shape().standard())
Paul's avatar
Paul committed
1058
        {
Paul's avatar
Format  
Paul committed
1059
            auto c = op::contiguous{};
Paul's avatar
Paul committed
1060
            auto l = c.compute(c.compute_shape({c_ins->get_shape()}), {c_ins->eval()});
Paul's avatar
Format  
Paul committed
1061
            c_ins  = m.add_literal(l.get_shape(), l.data());
Paul's avatar
Paul committed
1062
1063
        }

1064
1065
1066
1067
        auto inputs = gemm_ins->inputs();
        inputs.pop_back();

        inputs.push_back(c_ins);
1068
        inputs.push_back(ins->inputs().back());
1069
1070
1071
1072
1073

        m.replace_instruction(ins, gemm, inputs);
    }
};

Paul's avatar
Paul committed
1074
1075
1076
1077
struct find_contiguous_tranpose_gemm
{
    auto matcher() const
    {
Paul's avatar
Format  
Paul committed
1078
1079
1080
1081
        return match::name("gpu::contiguous")(match::arg(0)(
            match::name("transpose")(
                match::arg(0)(match::name("gpu::gemm")(match::used_once()).bind("gemm")))
                .bind("transpose")));
Paul's avatar
Paul committed
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
    }

    template <class Vector>
    static bool is_swapped(const Vector& perm, std::size_t i, std::size_t j)
    {
        if(i >= perm.size() or j >= perm.size())
            return false;
        auto perm2 = perm;
        std::iota(perm2.begin(), perm2.end(), 0);
        std::swap(perm2[i], perm2[j]);
        return perm2 == perm;
    }

    void apply(module& m, const match::matcher_result& r) const
    {
Paul's avatar
Format  
Paul committed
1097
1098
1099
        auto ins       = r.result;
        auto gemm      = r.instructions["gemm"];
        auto alloc     = gemm->inputs().back();
Paul's avatar
Paul committed
1100
        auto transpose = r.instructions["transpose"];
Paul's avatar
Format  
Paul committed
1101
1102
        auto perm      = transpose->get_operator().to_value()["permutation"].to_vector<int64_t>();
        auto iperm     = invert_permutation(perm);
Paul's avatar
Paul committed
1103

Paul's avatar
Format  
Paul committed
1104
        if(perm.size() < 3)
Paul's avatar
Paul committed
1105
1106
            return;

Paul's avatar
Format  
Paul committed
1107
        if(not is_swapped(perm, perm.size() - 3, perm.size() - 2))
Paul's avatar
Paul committed
1108
1109
1110
            return;

        auto lens = gemm->get_shape().lens();
Paul's avatar
Format  
Paul committed
1111
1112
        if(lens.size() > 3 and
           not std::all_of(lens.begin(), lens.end() - 3, [](auto i) { return i == 1; }))
Paul's avatar
Paul committed
1113
1114
            return;

Paul's avatar
Format  
Paul committed
1115
        auto gemmv           = gemm->get_operator().to_value();
Paul's avatar
Paul committed
1116
1117
1118
1119
        gemmv["trans_batch"] = 1;

        auto s = shape{alloc->get_shape().type(), reorder_dims(alloc->get_shape().lens(), iperm)};
        auto new_alloc = m.insert_instruction(gemm, make_op("allocate", {{"shape", to_value(s)}}));
Paul's avatar
Format  
Paul committed
1120
1121
        auto alloc_transpose =
            m.insert_instruction(gemm, make_op("transpose", {{"permutation", perm}}), new_alloc);
Paul's avatar
Paul committed
1122

Paul's avatar
Format  
Paul committed
1123
1124
1125
        auto inputs        = gemm->inputs();
        inputs.back()      = alloc_transpose;
        auto new_gemm      = m.insert_instruction(gemm, make_op("gpu::gemm", gemmv), inputs);
Paul's avatar
Paul committed
1126
1127
1128
1129
1130
1131
        auto gemm_transpoe = m.insert_instruction(gemm, transpose->get_operator(), new_gemm);

        m.replace_instruction(ins, gemm_transpoe);
    }
};

1132
1133
1134
1135
1136
1137
1138
struct find_commutative_broadcast
{
    auto matcher() const
    {
        return match::name("gpu::add", "gpu::mul")(match::arg(1)(match::broadcast_shape()));
    }

1139
    void apply(module& m, const match::matcher_result& r) const
1140
1141
1142
1143
1144
    {
        auto ins  = r.result;
        auto args = ins->inputs();
        move_broadcasted_back(args);

1145
        m.replace_instruction(ins, ins->get_operator(), args);
1146
1147
1148
    }
};

1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
struct find_contiguous
{
    auto matcher() const { return match::name("gpu::contiguous"); }

    void apply(module& m, const match::matcher_result& r) const
    {
        auto ins = r.result;

        m.replace_instruction(
            ins,
            make_op("gpu::precompile_op", {{"op", to_value(make_op("contiguous"))}}),
            ins->inputs());
    }
};

struct find_contiguous_pointwise
{
    auto matcher() const
    {
        return match::name("gpu::contiguous")(match::arg(0)(precompile_name("pointwise")));
    }

    void apply(module& m, const match::matcher_result& r) const
    {
        auto ins    = r.result;
        auto pw     = ins->inputs().front();
        auto alloc  = ins->inputs().back();
        auto args   = pw->inputs();
        args.back() = alloc;

        m.replace_instruction(ins, pw->get_operator(), args, pw->module_inputs());
    }
};

1183
void fuse_ops::apply(module& m) const
Paul's avatar
Paul committed
1184
{
1185
    match::find_matches(m, find_contiguous_pointwise{}, find_gelu{}, find_gelu_new{fast_math});
1186
1187
1188
    run_passes(m, {dead_code_elimination{}});
    match::find_matches(m, find_triadd{});
    match::find_matches(m,
kahmed10's avatar
kahmed10 committed
1189
                        find_layernorm{},
1190
                        find_conv_pointwise{ctx},
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
                        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{});
1201
    run_passes(m, {dead_code_elimination{}});
1202
1203
1204
1205
    match::find_matches(m,
                        find_triadd_layernorm{},
                        find_gemm_add{},
                        find_gemm_pointwise{},
Paul's avatar
Paul committed
1206
                        find_contiguous_tranpose_gemm{},
1207
                        find_commutative_broadcast{});
1208
    match::find_matches(m, find_contiguous{});
Paul's avatar
Paul committed
1209
}
Paul's avatar
Paul committed
1210
1211

} // namespace gpu
Paul's avatar
Paul committed
1212
} // namespace MIGRAPHX_INLINE_NS
Paul's avatar
Paul committed
1213
} // namespace migraphx