| 594 | |
| 595 | |
| 596 | def test_inception_like(): |
| 597 | def before(): |
| 598 | bb = relax.BlockBuilder() |
| 599 | |
| 600 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 601 | w0 = relax.Var("w0", R.Tensor((16, 16, 3, 3), "float32")) |
| 602 | w1 = relax.Var("w1", R.Tensor((16, 16, 3, 3), "float32")) |
| 603 | w2 = relax.Var("w2", R.Tensor((16, 32, 3, 3), "float32")) |
| 604 | w3 = relax.Var("w3", R.Tensor((16, 32, 3, 3), "float32")) |
| 605 | with bb.function("main", [x, w0, w1, w2, w3]): |
| 606 | with bb.dataflow(): |
| 607 | lv0 = bb.emit_te(topi.nn.conv2d, x, w0, strides=1, padding=1, dilation=1) |
| 608 | lv1 = bb.emit_te(topi.nn.relu, lv0) |
| 609 | lv2 = bb.emit_te(topi.nn.conv2d, x, w1, strides=1, padding=1, dilation=1) |
| 610 | lv3 = bb.emit_te(topi.nn.relu, lv2) |
| 611 | lv4 = bb.emit_te(topi.concatenate, (lv1, lv3), axis=1) |
| 612 | lv5 = bb.emit_te(topi.nn.conv2d, lv4, w2, strides=1, padding=1, dilation=1) |
| 613 | lv6 = bb.emit_te(topi.nn.relu, lv5) |
| 614 | lv7 = bb.emit_te(topi.nn.conv2d, lv4, w3, strides=1, padding=1, dilation=1) |
| 615 | lv8 = bb.emit_te(topi.nn.relu, lv7) |
| 616 | gv = bb.emit_output(bb.call_te(topi.concatenate, (lv6, lv8), axis=1)) |
| 617 | bb.emit_func_output(gv) |
| 618 | |
| 619 | return bb.get() |
| 620 | |
| 621 | def expected(): |
| 622 | bb = relax.BlockBuilder() |
| 623 | |
| 624 | # Grouped function 1 |
| 625 | x = relax.Var("x", R.Tensor((1, 16, 64, 64), "float32")) |
| 626 | w = relax.Var("w", R.Tensor((16, 16, 3, 3), "float32")) |
| 627 | with bb.function("fused_conv2d_relu", [x, w], attrs={"Primitive": True}, private=True): |
| 628 | with bb.dataflow(): |
| 629 | lv0 = bb.emit_te( |
| 630 | topi.nn.conv2d, |
| 631 | x, |
| 632 | w, |
| 633 | strides=1, |
| 634 | padding=1, |
| 635 | dilation=1, |
| 636 | primfunc_name_hint="conv2d", |
| 637 | ) |
| 638 | gv = bb.emit_output(bb.call_te(topi.nn.relu, lv0)) |
| 639 | bb.emit_func_output(gv) |
| 640 | |
| 641 | # Grouped function 2 |
| 642 | x = relax.Var("x", R.Tensor((1, 32, 64, 64), "float32")) |
| 643 | w = relax.Var("w", R.Tensor((16, 32, 3, 3), "float32")) |
| 644 | with bb.function("fused_conv2d1_relu", [x, w], attrs={"Primitive": True}, private=True): |
| 645 | with bb.dataflow(): |
| 646 | lv0 = bb.emit_te( |
| 647 | topi.nn.conv2d, |
| 648 | x, |
| 649 | w, |
| 650 | strides=1, |
| 651 | padding=1, |
| 652 | dilation=1, |
| 653 | primfunc_name_hint="conv2d1", |