lowering.cpp 17 KB
Newer Older
1
#include <rocblas.h>
Paul's avatar
Paul committed
2
#include <migraph/gpu/lowering.hpp>
Paul's avatar
Paul committed
3
4
5
#include <migraph/manage_ptr.hpp>
#include <migraph/instruction.hpp>
#include <migraph/operators.hpp>
Paul's avatar
Paul committed
6
#include <migraph/generate.hpp>
Paul's avatar
Paul committed
7
#include <migraph/shape_for_each.hpp>
Paul's avatar
Paul committed
8
9
#include <migraph/gpu/miopen.hpp>
#include <migraph/gpu/hip.hpp>
Paul's avatar
Paul committed
10
#include <migraph/dfor.hpp>
11
#include <migraph/gpu/device/contiguous.hpp>
Paul's avatar
Paul committed
12
#include <migraph/gpu/device/add.hpp>
Paul's avatar
Paul committed
13
#include <migraph/iterator_for.hpp>
Paul's avatar
Paul committed
14
15
#include <migraph/gpu/rocblas.hpp>
#include <migraph/gpu/context.hpp>
Paul's avatar
Paul committed
16
#include <utility>
Paul's avatar
Paul committed
17
18

namespace migraph {
Paul's avatar
Paul committed
19
namespace gpu {
Paul's avatar
Paul committed
20

21
22
23
24
25
26
struct miopen_batch_norm_inference
{
    batch_norm_inference op;

    std::string name() const { return "gpu::batch_norm_inference"; }

Paul's avatar
Paul committed
27
    shape compute_shape(const std::vector<shape>& inputs) const
28
29
30
31
32
33
    {
        check_shapes{inputs, *this}.has(6);
        return op.compute_shape(
            {inputs.at(0), inputs.at(1), inputs.at(2), inputs.at(3), inputs.at(4)});
    }

Paul's avatar
Paul committed
34
35
    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
36
    {
wsttiger's avatar
wsttiger committed
37
38
        auto x_desc  = make_tensor(args[0].get_shape());
        auto y_desc  = make_tensor(output_shape);
39
        auto bn_desc = make_tensor(args[3].get_shape());
40
41
42
43
44
45
46
47
48
49
50

        float alpha = 1.0, beta = 0.0f;

        miopenBatchNormalizationForwardInference(ctx.handle.get(),
                                                 miopenBatchNormMode_t(op.bn_mode),
                                                 &alpha,
                                                 &beta,
                                                 x_desc.get(),
                                                 args[0].implicit(),
                                                 y_desc.get(),
                                                 args[5].implicit(),
51
                                                 bn_desc.get(),
52
53
                                                 args[1].implicit(),
                                                 args[2].implicit(),
Paul's avatar
Paul committed
54
55
                                                 args[3].implicit(),
                                                 args[4].implicit(),
56
                                                 op.epsilon);
57
58
59
60
61

        return args[5];
    }
};

Paul's avatar
Paul committed
62
63
64
struct miopen_convolution
{
    convolution op;
Paul's avatar
Paul committed
65
    shared<convolution_descriptor> cd;
Paul's avatar
Paul committed
66
    miopenConvFwdAlgorithm_t algo{};
Paul's avatar
Paul committed
67

Paul's avatar
Paul committed
68
    std::string name() const { return "gpu::convolution"; }
Paul's avatar
Paul committed
69
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
70
    {
Paul's avatar
Paul committed
71
        check_shapes{inputs, *this}.has(4).standard();
Paul's avatar
Paul committed
72
        return op.compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
73
    }
Paul's avatar
Paul committed
74
75
    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
76
    {
Paul's avatar
Paul committed
77
78
        auto x_desc = make_tensor(args[0].get_shape());
        auto w_desc = make_tensor(args[1].get_shape());
Paul's avatar
Paul committed
79
80
        auto y_desc = make_tensor(output_shape);

Paul's avatar
Paul committed
81
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
82
        miopenConvolutionForward(ctx.handle.get(),
Paul's avatar
Paul committed
83
                                 &alpha,
Paul's avatar
Paul committed
84
                                 x_desc.get(),
Paul's avatar
Paul committed
85
                                 args[0].implicit(),
Paul's avatar
Paul committed
86
                                 w_desc.get(),
Paul's avatar
Paul committed
87
                                 args[1].implicit(),
Paul's avatar
Paul committed
88
                                 cd.get(),
Paul's avatar
Paul committed
89
                                 algo,
Paul's avatar
Paul committed
90
                                 &beta,
Paul's avatar
Paul committed
91
                                 y_desc.get(),
Paul's avatar
Paul committed
92
                                 args[3].implicit(),
Paul's avatar
Paul committed
93
                                 args[2].implicit(),
Paul's avatar
Paul committed
94
95
                                 args[2].get_shape().bytes());
        return args[3];
Paul's avatar
Paul committed
96
    }
Paul's avatar
Paul committed
97

Paul's avatar
Paul committed
98
    shape compile(context& ctx, const shape& output_shape, std::vector<instruction_ref> inputs)
Paul's avatar
Paul committed
99
    {
Paul's avatar
Paul committed
100
        shape workspace_shape{};
Paul's avatar
Paul committed
101
102
103
104
        auto x_desc = make_tensor(inputs[0]->get_shape());
        auto w_desc = make_tensor(inputs[1]->get_shape());
        auto y_desc = make_tensor(output_shape);

Paul's avatar
Paul committed
105
        std::size_t workspace_size = 0;
Paul's avatar
Paul committed
106
        miopenConvolutionForwardGetWorkSpaceSize(
Paul's avatar
Paul committed
107
            ctx.handle.get(), w_desc.get(), x_desc.get(), cd.get(), y_desc.get(), &workspace_size);
Paul's avatar
Paul committed
108
109
        workspace_shape = shape{shape::int8_type, {workspace_size}};

Paul's avatar
Paul committed
110
111
112
        auto x         = to_gpu(generate_argument(inputs[0]->get_shape()));
        auto w         = to_gpu(generate_argument(inputs[1]->get_shape()));
        auto y         = to_gpu(generate_argument(output_shape));
Paul's avatar
Paul committed
113
        auto workspace = allocate_gpu(workspace_shape);
Paul's avatar
Paul committed
114

Paul's avatar
Paul committed
115
        int algo_count = 1;
Paul's avatar
Paul committed
116
117
118
119
120
121
122
123
124
125
126
127
        miopenConvAlgoPerf_t perf;
        miopenFindConvolutionForwardAlgorithm(ctx.handle.get(),
                                              x_desc.get(),
                                              x.implicit(),
                                              w_desc.get(),
                                              w.implicit(),
                                              cd.get(),
                                              y_desc.get(),
                                              y.implicit(),
                                              1,
                                              &algo_count,
                                              &perf,
Paul's avatar
Paul committed
128
129
                                              workspace.implicit(),
                                              workspace_size,
Paul's avatar
Paul committed
130
131
                                              false);
        algo = perf.fwd_algo;
132
        return shape{shape::int8_type, {perf.memory}};
Paul's avatar
Paul committed
133
    }
Paul's avatar
Paul committed
134
135
};

Paul's avatar
Paul committed
136
137
138
139
140
struct miopen_pooling
{
    pooling op;
    shared<pooling_descriptor> pd;

Paul's avatar
Paul committed
141
    std::string name() const { return "gpu::pooling"; }
Paul's avatar
Paul committed
142
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
143
    {
Paul's avatar
Paul committed
144
        check_shapes{inputs, *this}.has(2).standard();
Paul's avatar
Paul committed
145
        return op.compute_shape({inputs.at(0)});
Paul's avatar
Paul committed
146
    }
Paul's avatar
Paul committed
147
148
    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
149
    {
Paul's avatar
Paul committed
150
        auto x_desc = make_tensor(args[0].get_shape());
Paul's avatar
Paul committed
151
152
153
154
        auto y_desc = make_tensor(output_shape);

        float alpha = 1, beta = 0;

Paul's avatar
Paul committed
155
        miopenPoolingForward(ctx.handle.get(),
Paul's avatar
Paul committed
156
157
158
                             pd.get(),
                             &alpha,
                             x_desc.get(),
Paul's avatar
Paul committed
159
                             args[0].implicit(),
Paul's avatar
Paul committed
160
161
                             &beta,
                             y_desc.get(),
Paul's avatar
Paul committed
162
                             args[1].implicit(),
Paul's avatar
Paul committed
163
164
165
                             false,
                             nullptr,
                             0);
Paul's avatar
Paul committed
166

Paul's avatar
Paul committed
167
        return args[1];
Paul's avatar
Paul committed
168
169
170
    }
};

Paul's avatar
Paul committed
171
struct hip_add
Paul's avatar
Paul committed
172
{
Paul's avatar
Paul committed
173
    std::string name() const { return "gpu::add"; }
Paul's avatar
Paul committed
174
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
175
    {
Paul's avatar
Paul committed
176
        // check_shapes{inputs, *this}.has(3).standard();
Paul's avatar
Paul committed
177
        check_shapes{inputs, *this}.has(3);
Paul's avatar
Paul committed
178
        return inputs.at(0);
Paul's avatar
Paul committed
179
180
    }

Paul's avatar
Paul committed
181
    argument compute(context&, const shape&, const std::vector<argument>& args) const
Paul's avatar
Paul committed
182
    {
Paul's avatar
Paul committed
183
        device::add(args[2], args[0], args[1]);
Paul's avatar
Paul committed
184
        return args[2];
Paul's avatar
Paul committed
185
186
187
188
189
    }
};

struct miopen_add
{
Paul's avatar
Paul committed
190
    std::string name() const { return "gpu::add"; }
Paul's avatar
Paul committed
191
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
192
    {
Paul's avatar
Paul committed
193
        check_shapes{inputs, *this}.has(3).not_broadcasted();
Paul's avatar
Paul committed
194
        return inputs.at(0);
Paul's avatar
Paul committed
195
196
    }

Paul's avatar
Paul committed
197
198
    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
199
    {
Paul's avatar
Paul committed
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
        float alpha = 1, beta = 0;
        auto a_desc = make_tensor(args[0].get_shape());
        auto b_desc = make_tensor(args[1].get_shape());
        auto c_desc = make_tensor(output_shape);
        miopenOpTensor(ctx.handle.get(),
                       miopenTensorOpAdd,
                       &alpha,
                       a_desc.get(),
                       args[0].implicit(),
                       &alpha,
                       b_desc.get(),
                       args[1].implicit(),
                       &beta,
                       c_desc.get(),
                       args[2].implicit());
        return args[2];
Paul's avatar
Paul committed
216
217
218
    }
};

Paul's avatar
Paul committed
219
220
221
struct miopen_gemm
{
    gemm op;
222
    std::string name() const { return "gpu::gemm"; }
Paul's avatar
Paul committed
223
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
224
    {
Paul's avatar
Paul committed
225
226
        check_shapes{inputs, *this}.has(3);
        return op.compute_shape({inputs.at(0), inputs.at(1)});
Paul's avatar
Paul committed
227
    }
Paul's avatar
Paul committed
228
229
    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
230
    {
231
232
        float alpha     = 1.0f;
        float beta      = 0.0f;
Paul's avatar
Paul committed
233
234
        bool transa     = args[0].get_shape().transposed();
        bool transb     = args[1].get_shape().transposed();
235
236
237
        rocblas_int lda = args[0].get_shape().strides()[transa ? 1 : 0];
        rocblas_int ldb = args[1].get_shape().strides()[transb ? 1 : 0];
        rocblas_int ldc = args[2].get_shape().strides()[0];
238
239
240
        rocblas_int m   = output_shape.lens()[0];
        rocblas_int n   = output_shape.lens()[1];
        rocblas_int k   = args[0].get_shape().lens()[1];
241
        rocblas_sgemm(ctx.rbhandle.get(),
242
243
                      transb ? rocblas_operation_transpose : rocblas_operation_none,
                      transa ? rocblas_operation_transpose : rocblas_operation_none,
244
245
246
247
248
249
250
251
252
253
254
255
                      n,
                      m,
                      k,
                      &alpha,
                      args[1].implicit(),
                      ldb,
                      args[0].implicit(),
                      lda,
                      &beta,
                      args[2].implicit(),
                      ldc);
        return args[2];
Paul's avatar
Paul committed
256
257
258
    }
};

259
260
261
struct miopen_contiguous
{
    contiguous op;
Paul's avatar
Paul committed
262
    std::string name() const { return "gpu::contiguous"; }
Paul's avatar
Paul committed
263
    shape compute_shape(const std::vector<shape>& inputs) const
264
265
266
267
    {
        check_shapes{inputs, *this}.has(2);
        return op.compute_shape({inputs.at(0)});
    }
Paul's avatar
Paul committed
268
    argument compute(context&, shape output_shape, const std::vector<argument>& args) const
269
    {
Paul's avatar
Paul committed
270
271
        assert(output_shape == args[1].get_shape());
        assert(output_shape.standard());
Paul's avatar
Paul committed
272
        (void)output_shape;
273
        device::contiguous(args.at(1), args.at(0));
274
        return args.at(1);
275
276
277
    }
};

Paul's avatar
Paul committed
278
279
280
struct miopen_relu
{
    shared<activation_descriptor> ad;
Paul's avatar
Paul committed
281
    std::string name() const { return "gpu::relu"; }
Paul's avatar
Paul committed
282
    shape compute_shape(const std::vector<shape>& inputs) const
Paul's avatar
Paul committed
283
    {
Paul's avatar
Paul committed
284
        check_shapes{inputs, *this}.has(2).not_broadcasted();
Paul's avatar
Paul committed
285
        return inputs.at(1);
Paul's avatar
Paul committed
286
287
    }

Paul's avatar
Paul committed
288
289
    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
Paul's avatar
Paul committed
290
291
    {
        float alpha = 1, beta = 0;
Paul's avatar
Paul committed
292
        auto x_desc = make_tensor(args[0].get_shape());
Paul's avatar
Paul committed
293
        auto y_desc = make_tensor(output_shape);
Paul's avatar
Paul committed
294
        miopenActivationForward(ctx.handle.get(),
Paul's avatar
Paul committed
295
296
297
                                ad.get(),
                                &alpha,
                                x_desc.get(),
Paul's avatar
Paul committed
298
                                args[0].implicit(),
Paul's avatar
Paul committed
299
300
                                &beta,
                                y_desc.get(),
Paul's avatar
Paul committed
301
                                args[1].implicit());
Paul's avatar
Paul committed
302

Paul's avatar
Paul committed
303
        return args[1];
Paul's avatar
Paul committed
304
305
306
    }
};

Paul's avatar
Paul committed
307
308
309
310
311
312
313
struct miopen_softmax
{
    softmax op;
    std::string name() const { return "gpu::softmax"; }
    shape compute_shape(const std::vector<shape>& inputs) const
    {
        check_shapes{inputs, *this}.has(2).standard();
Paul's avatar
Paul committed
314
        return op.compute_shape({inputs.at(0)});
Paul's avatar
Paul committed
315
316
317
318
319
320
321
322
323
    }

    argument
    compute(context& ctx, const shape& output_shape, const std::vector<argument>& args) const
    {
        float alpha = 1, beta = 0;
        auto x_desc = make_tensor(args[0].get_shape());
        auto y_desc = make_tensor(output_shape);
        miopenSoftmaxForward(ctx.handle.get(),
Paul's avatar
Paul committed
324
325
326
327
328
329
                             &alpha,
                             x_desc.get(),
                             args[0].implicit(),
                             &beta,
                             y_desc.get(),
                             args[1].implicit());
Paul's avatar
Paul committed
330
331
332
333
334

        return args[1];
    }
};

Paul's avatar
Paul committed
335
336
struct miopen_apply
{
Paul's avatar
Paul committed
337
    program* prog = nullptr;
Paul's avatar
Paul committed
338
    context ctx{};
Paul's avatar
Paul committed
339

Paul's avatar
Paul committed
340
341
342
343
344
345
346
    void check_shape(shape x, instruction_ref i)
    {
        assert(x == i->get_shape());
        (void)x;
        (void)i;
    }

Paul's avatar
Paul committed
347
348
    void apply()
    {
Paul's avatar
Paul committed
349
350
        for(auto it = prog->begin(); it != prog->end(); it++)
        {
Paul's avatar
Paul committed
351
            auto s = it->get_shape();
Paul's avatar
Paul committed
352
            if(it->name() == "convolution")
Paul's avatar
Paul committed
353
            {
Paul's avatar
Paul committed
354
                check_shape(s, apply_convolution(it));
Paul's avatar
Paul committed
355
            }
Paul's avatar
Paul committed
356
            else if(it->name() == "activation")
Paul's avatar
Paul committed
357
            {
Paul's avatar
Paul committed
358
                check_shape(s, apply_activation(it));
Paul's avatar
Paul committed
359
            }
Paul's avatar
Paul committed
360
            else if(it->name() == "pooling")
Paul's avatar
Paul committed
361
            {
Paul's avatar
Paul committed
362
                check_shape(s, apply_pooling(it));
Paul's avatar
Paul committed
363
            }
Paul's avatar
Paul committed
364
            else if(it->name() == "add")
Paul's avatar
Paul committed
365
            {
Paul's avatar
Paul committed
366
                check_shape(s, apply_add(it));
Paul's avatar
Paul committed
367
            }
Paul's avatar
Paul committed
368
            else if(it->name() == "gemm")
Paul's avatar
Paul committed
369
            {
Paul's avatar
Paul committed
370
                check_shape(s, apply_gemm(it));
Paul's avatar
Paul committed
371
            }
Paul's avatar
Paul committed
372
            else if(it->name() == "contiguous")
373
            {
Paul's avatar
Paul committed
374
                check_shape(s, apply_contiguous(it));
375
            }
Paul's avatar
Paul committed
376
            else if(it->name() == "batch_norm_inference")
377
            {
Paul's avatar
Paul committed
378
                check_shape(s, apply_batch_norm_inference(it));
379
            }
Paul's avatar
Paul committed
380
381
382
383
            else if(it->name() == "softmax")
            {
                check_shape(s, apply_softmax(it));
            }
Paul's avatar
Paul committed
384
385
386
        }
    }

Paul's avatar
Paul committed
387
    instruction_ref insert_allocation(instruction_ref ins, const shape& s, std::string tag = "")
Paul's avatar
Paul committed
388
    {
Paul's avatar
Paul committed
389
        if(ins == --prog->end() and tag.empty())
Paul's avatar
Paul committed
390
391
392
393
394
        {
            return prog->add_parameter("output", s);
        }
        else
        {
Paul's avatar
Paul committed
395
            auto is     = prog->add_outline(s);
Paul's avatar
Paul committed
396
            auto result = prog->insert_instruction(ins, hip_allocate{std::move(tag)}, is);
Paul's avatar
Paul committed
397
398
399
400
            return result;
        }
    }

Paul's avatar
Paul committed
401
    instruction_ref apply_convolution(instruction_ref ins)
Paul's avatar
Paul committed
402
    {
403
        auto&& op = any_cast<convolution>(ins->get_operator());
Paul's avatar
Paul committed
404

Paul's avatar
Paul committed
405
        auto conv = miopen_convolution{op, make_conv(op)};
Paul's avatar
Paul committed
406
        auto ws   = conv.compile(ctx, ins->get_shape(), ins->inputs());
Paul's avatar
Paul committed
407

408
        auto workspace = insert_allocation(ins, ws, "workspace");
Paul's avatar
Paul committed
409
        auto output    = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
410

Paul's avatar
Paul committed
411
        return prog->replace_instruction(
Paul's avatar
Paul committed
412
            ins, conv, ins->inputs().at(0), ins->inputs().at(1), workspace, output);
Paul's avatar
Paul committed
413
414
    }

Paul's avatar
Paul committed
415
    instruction_ref apply_pooling(instruction_ref ins)
Paul's avatar
Paul committed
416
    {
417
        auto&& op   = any_cast<pooling>(ins->get_operator());
Paul's avatar
Paul committed
418
        auto pd     = make_pooling(op);
Paul's avatar
Paul committed
419
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
420

Paul's avatar
Paul committed
421
        return prog->replace_instruction(
Paul's avatar
Paul committed
422
            ins, miopen_pooling{op, std::move(pd)}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
423
424
    }

Paul's avatar
Paul committed
425
    instruction_ref apply_activation(instruction_ref ins)
Paul's avatar
Paul committed
426
    {
427
        auto&& op = any_cast<activation>(ins->get_operator());
Paul's avatar
Paul committed
428
429
        auto ad   = make_relu();
        if(op.mode == "relu")
Paul's avatar
Paul committed
430
        {
Paul's avatar
Paul committed
431
            auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
432
            return prog->replace_instruction(
Paul's avatar
Paul committed
433
                ins, miopen_relu{std::move(ad)}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
434
        }
Paul's avatar
Paul committed
435
        return ins;
Paul's avatar
Paul committed
436
    }
Paul's avatar
Paul committed
437

Paul's avatar
Paul committed
438
439
    instruction_ref apply_softmax(instruction_ref ins)
    {
Paul's avatar
Paul committed
440
        auto&& op   = any_cast<softmax>(ins->get_operator());
Paul's avatar
Paul committed
441
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
442
        return prog->replace_instruction(ins, miopen_softmax{op}, ins->inputs().at(0), output);
Paul's avatar
Paul committed
443
444
    }

Paul's avatar
Paul committed
445
    instruction_ref apply_add(instruction_ref ins)
Paul's avatar
Paul committed
446
    {
Paul's avatar
Paul committed
447
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
448
        return prog->replace_instruction(
Paul's avatar
Paul committed
449
            ins, hip_add{}, ins->inputs().at(0), ins->inputs().at(1), output);
Paul's avatar
Paul committed
450
    }
Paul's avatar
Paul committed
451

Paul's avatar
Paul committed
452
    instruction_ref apply_gemm(instruction_ref ins)
Paul's avatar
Paul committed
453
    {
454
        auto&& op   = any_cast<gemm>(ins->get_operator());
Paul's avatar
Paul committed
455
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
456
        return prog->replace_instruction(
Paul's avatar
Paul committed
457
            ins, miopen_gemm{op}, ins->inputs().at(0), ins->inputs().at(1), output);
Paul's avatar
Paul committed
458
    }
459

Paul's avatar
Paul committed
460
    instruction_ref apply_contiguous(instruction_ref ins)
461
    {
462
        auto&& op   = any_cast<contiguous>(ins->get_operator());
Paul's avatar
Paul committed
463
        auto output = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
464
        return prog->replace_instruction(ins, miopen_contiguous{op}, ins->inputs().at(0), output);
465
    }
466

Paul's avatar
Paul committed
467
    instruction_ref apply_batch_norm_inference(instruction_ref ins)
468
    {
469
        auto&& op       = any_cast<batch_norm_inference>(ins->get_operator());
Paul's avatar
Paul committed
470
        auto output     = insert_allocation(ins, ins->get_shape());
Paul's avatar
Paul committed
471
        shape old_shape = ins->inputs().at(1)->get_shape();
wsttiger's avatar
wsttiger committed
472
        std::vector<int64_t> new_shape{1, static_cast<int64_t>(old_shape.elements()), 1, 1};
Paul's avatar
Paul committed
473
474
        auto reshape_op = reshape{new_shape};
        std::vector<instruction_ref> reshapes;
Paul's avatar
Paul committed
475
476
        std::transform(ins->inputs().begin() + 1,
                       ins->inputs().end(),
Paul's avatar
Paul committed
477
478
                       std::back_inserter(reshapes),
                       [&](auto i) { return prog->insert_instruction(ins, reshape_op, i); });
Paul's avatar
Paul committed
479
        return prog->replace_instruction(ins,
Paul's avatar
Paul committed
480
                                         miopen_batch_norm_inference{op},
Paul's avatar
Paul committed
481
                                         ins->inputs().at(0),
Paul's avatar
Paul committed
482
483
484
485
486
                                         reshapes[0],
                                         reshapes[1],
                                         reshapes[2],
                                         reshapes[3],
                                         output);
487
    }
Paul's avatar
Paul committed
488
489
};

Paul's avatar
Paul committed
490
void lowering::apply(program& p) const { miopen_apply{&p, ctx}.apply(); }
Paul's avatar
Paul committed
491

Paul's avatar
Paul committed
492
} // namespace gpu
Paul's avatar
Paul committed
493

Paul's avatar
Paul committed
494
} // namespace migraph