-
Notifications
You must be signed in to change notification settings - Fork 3
/
pass_args.py
118 lines (101 loc) · 2.54 KB
/
pass_args.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
# -*- coding:utf-8 -*-
# Author: liyanpeng
# Email: [email protected]
# Datetime: 2022/10/10 10:44
# Filename: pass_args.py
from tvm import relay
from tvm.relay.dataflow_pattern import TupleGetItemPattern, is_op, wildcard
def make_add_sub_mul_pattern():
r"""Create a pattern to match the following graph.
add sub
\ /
\ /
mul
"""
x = wildcard()
y = wildcard()
return (x + y) * (x - y)
def make_add_relu_pattern():
r"""Create a pattern to match the following graph.
add
|
relu
"""
add_node = wildcard() + wildcard()
r = is_op("nn.relu")(add_node)
return r
def make_conv_bias_relu_pattern():
r"""Create a pattern to match the following graph.
conv2d
|
bias_add
|
relu
"""
x = wildcard()
y = wildcard()
z = wildcard()
conv_node = is_op("nn.conv2d")(x, y)
bias_node = is_op("nn.bias_add")(conv_node, z)
r = is_op("nn.relu")(bias_node)
return r
def make_pattern_with_optional():
r"""Create a pattern to match the following graph. Note that relu is optinal.
conv2d
|
bias_add
|
(relu)
"""
x = wildcard()
y = wildcard()
z = wildcard()
conv_node = is_op("nn.conv2d")(x, y)
bias_node = is_op("nn.bias_add")(conv_node, z)
r = bias_node.optional(lambda x: is_op("nn.relu")(x))
return r
def make_add_add_add_pattern():
r"""Create a pattern to match the following graph.
Useful for testing re-using a call node.
x y
/ \ /
| add
\ | \
add |
| /
add
"""
x = wildcard()
y = wildcard()
add_node = is_op("add")(x, y)
add_node_1 = is_op("add")(x, add_node)
r = is_op("add")(add_node_1, add_node)
return r
def make_bn_relu_pattern():
r"""Create a pattern to match the following graph.
batch_norm
|
TupleGetItem(0)
|
relu
"""
x = wildcard()
gamma = wildcard()
beta = wildcard()
moving_mean = wildcard()
moving_var = wildcard()
bn_node = is_op("nn.batch_norm")(x, gamma, beta, moving_mean, moving_var)
tuple_get_item_node = TupleGetItemPattern(bn_node, 0)
r = is_op("nn.relu")(tuple_get_item_node)
return r
def fskip(expr):
if isinstance(expr, relay.expr.Call) and expr.op.name == "add":
return True
return False
pattern_table = [
("conv2d_bias_relu", make_conv_bias_relu_pattern()),
("add_relu", make_add_relu_pattern()),
]
desired_layouts = {
"nn.conv2d": ["NCHW", "default"]
}