| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # pylint: disable=unused-wildcard-import |
| import numpy as np |
| |
| import tvm |
| from tvm import relay |
| from tvm.relay.build_module import bind_params_by_name |
| from tvm.relay.dataflow_pattern import * |
| from tvm.relay.testing import run_opt_pass |
| |
| # NB: 1 corresponds to the C++ enum that specicfies this |
| # we loose the type safety due to the Python/C++ calling |
| # convention. |
| K_ELEMWISE = 0 |
| K_BROADCAST = 1 |
| |
| |
| ## NODE TESTS |
| def test_expr_pattern(): |
| ep = is_expr(relay.var("x", shape=(4, 1))) |
| assert isinstance(ep, ExprPattern) |
| assert isinstance(ep.expr, relay.Var) |
| |
| |
| def test_var_pattern(): |
| v = is_var("x") |
| assert isinstance(v, VarPattern) |
| assert v.name == "x" |
| |
| |
| def test_constant_pattern(): |
| c = is_constant() |
| assert isinstance(c, ConstantPattern) |
| |
| |
| def test_wildcard_pattern(): |
| wc = wildcard() |
| assert isinstance(wc, WildcardPattern) |
| |
| |
| def test_CallPattern(): |
| wc1 = wildcard() |
| wc2 = wildcard() |
| c = is_op("add")(wc1, wc2) |
| assert isinstance(c, CallPattern) |
| assert isinstance(c.args[0], WildcardPattern) |
| assert isinstance(c.args[1], WildcardPattern) |
| |
| |
| def test_FunctionPattern(): |
| wc1 = wildcard() |
| wc2 = wildcard() |
| c = is_op("add")(wc1, wc2) |
| f = FunctionPattern([wc1, wc2], c) |
| assert isinstance(f, FunctionPattern) |
| assert isinstance(f.params[0], WildcardPattern) |
| assert isinstance(f.params[1], WildcardPattern) |
| assert isinstance(f.body, CallPattern) |
| assert isinstance(f.body.args[0], WildcardPattern) |
| assert isinstance(f.body.args[1], WildcardPattern) |
| |
| |
| def test_TuplePattern(): |
| wc1 = wildcard() |
| wc2 = wildcard() |
| t = is_tuple([wc1, wc2]) |
| assert isinstance(t, TuplePattern) |
| assert isinstance(t.fields[0], WildcardPattern) |
| assert isinstance(t.fields[1], WildcardPattern) |
| |
| |
| def test_TupleGetItemPattern(): |
| wc1 = wildcard() |
| wc2 = wildcard() |
| t = is_tuple([wc1, wc2]) |
| tgi = is_tuple_get_item(t, 1) |
| assert isinstance(tgi, TupleGetItemPattern) |
| assert isinstance(tgi.tuple, TuplePattern) |
| assert isinstance(tgi.tuple.fields[0], WildcardPattern) |
| assert isinstance(tgi.tuple.fields[1], WildcardPattern) |
| |
| |
| def test_AltPattern(): |
| is_add_or_sub = is_op("add") | is_op("subtract") |
| assert isinstance(is_add_or_sub, AltPattern) |
| |
| |
| def test_TypePattern(): |
| ttype = relay.TensorType((10, 10), "float32") |
| ty_pat = has_type(ttype) |
| assert isinstance(ty_pat, TypePattern) |
| assert ty_pat.type == ttype |
| |
| |
| def test_DataTypePattern(): |
| dtype = "float16" |
| pattern = has_dtype(dtype) |
| assert isinstance(pattern, DataTypePattern) |
| assert pattern.dtype == dtype |
| |
| |
| def test_ShapePattern(): |
| shape = [10, 10] |
| pattern = has_shape(shape) |
| assert isinstance(pattern, ShapePattern) |
| assert tvm.ir.structural_equal(pattern.shape, shape) |
| |
| |
| def test_AttrPattern(): |
| op = is_op("add").has_attr({"TOpPattern": K_ELEMWISE}) |
| assert isinstance(op, AttrPattern) |
| assert op.attrs["TOpPattern"] == K_ELEMWISE |
| |
| |
| def test_IfPattern(): |
| x = is_var("x") |
| y = is_var("y") |
| pat = is_if(is_op("less")(x, y), x, y) |
| |
| assert isinstance(pat, IfPattern) |
| assert isinstance(pat.cond, CallPattern) |
| assert isinstance(pat.true_branch, VarPattern) |
| assert isinstance(pat.false_branch, VarPattern) |
| |
| |
| def test_LetPattern(): |
| x = is_var("x") |
| y = is_var("y") |
| let_var = is_var("let") |
| pat = is_let(let_var, is_op("less")(x, y), let_var) |
| |
| assert isinstance(pat, LetPattern) |
| assert isinstance(pat.var, VarPattern) |
| assert isinstance(pat.value, CallPattern) |
| assert isinstance(pat.body, VarPattern) |
| |
| |
| ## MATCHER TESTS |
| |
| |
| def test_match_op(): |
| assert is_op("add").match(relay.op.op.get("add")) |
| |
| |
| def test_no_match_op(): |
| assert not is_op("add").match(relay.op.op.get("subtract")) |
| |
| |
| def test_match_op_or(): |
| is_add_or_sub = is_op("add") | is_op("subtract") |
| assert is_add_or_sub.match(relay.op.op.get("add")) |
| assert is_add_or_sub.match(relay.op.op.get("subtract")) |
| |
| |
| def test_match_call_commutive(): |
| x = relay.var("x") |
| y = relay.var("y") |
| add_pattern = is_op("add")(is_var("x"), is_var("y")) |
| assert add_pattern.match(x + y) |
| assert add_pattern.match(y + x) |
| mul_pattern = is_op("multiply")(is_var("x"), is_var("y")) |
| assert mul_pattern.match(x * y) |
| assert mul_pattern.match(y * x) |
| |
| |
| def test_no_match_call_commutive(): |
| x = relay.var("x") |
| y = relay.var("y") |
| add_pattern = is_op("subtract")(is_var("x"), is_var("y")) |
| assert add_pattern.match(x - y) |
| assert not add_pattern.match(y - x) |
| add_pattern = is_op("divide")(is_var("x"), is_var("y")) |
| assert add_pattern.match(x / y) |
| assert not add_pattern.match(y / x) |
| |
| |
| def test_match_call(): |
| x = relay.var("x") |
| y = relay.var("y") |
| add_pattern = is_op("add")(wildcard(), wildcard()) |
| assert add_pattern.match(x + y) |
| |
| # Match call with any number of inputs |
| call_pattern = wildcard()(None) |
| assert call_pattern.match(relay.op.nn.relu(x)) |
| assert call_pattern.match(relay.op.add(x, y)) |
| |
| |
| def test_no_match_call(): |
| x = relay.var("x") |
| y = relay.var("y") |
| add_pattern = is_op("add")(wildcard(), wildcard()) |
| assert not add_pattern.match(x - y) |
| |
| |
| def test_match_func(): |
| x = relay.var("x") |
| y = relay.var("y") |
| wc1 = wildcard() |
| wc2 = wildcard() |
| func_pattern = FunctionPattern([wc1, wc2], wc1 + wc2) |
| assert func_pattern.match(relay.Function([x, y], x + y)) |
| |
| # Match Function with any number of inputs |
| func_pattern = FunctionPattern(None, wildcard()) |
| assert func_pattern.match(relay.Function([x], x)) |
| assert func_pattern.match(relay.Function([x, y], x + y)) |
| |
| |
| def test_no_match_func(): |
| x = relay.var("x") |
| y = relay.var("y") |
| wc1 = wildcard() |
| wc2 = wildcard() |
| func_pattern = FunctionPattern([wc1, wc2], wc1 + wc2) |
| assert not func_pattern.match(relay.Function([x, y], x - y)) |
| |
| |
| def test_match_if(): |
| x = is_var("x") |
| y = is_var("y") |
| pat = is_if(is_op("less")(x, y), x, y) |
| |
| x = relay.var("x") |
| y = relay.var("y") |
| cond = x < y |
| |
| assert pat.match(relay.expr.If(cond, x, y)) |
| |
| |
| def test_no_match_if(): |
| x = is_var("x") |
| y = is_var("y") |
| pat = is_if(is_op("less")(x, y), x, y) |
| |
| x = relay.var("x") |
| y = relay.var("y") |
| |
| assert not pat.match(relay.expr.If(x > y, x, y)) |
| assert not pat.match(relay.expr.If(x < y, y, x)) |
| |
| |
| def test_match_let(): |
| x = is_var("x") |
| y = is_var("y") |
| let_var = is_var("let") |
| pat = is_let(let_var, is_op("less")(x, y), let_var) |
| |
| x = relay.var("x") |
| y = relay.var("y") |
| lv = relay.var("let") |
| cond = x < y |
| assert pat.match(relay.expr.Let(lv, cond, lv)) |
| |
| |
| def test_no_match_let(): |
| x = is_var("x") |
| y = is_var("y") |
| let_var = is_var("let") |
| pat = is_let(let_var, is_op("less")(x, y), let_var) |
| |
| x = relay.var("x") |
| y = relay.var("y") |
| lv = relay.var("let") |
| |
| assert not pat.match(relay.expr.Let(lv, x > y, lv)) |
| assert not pat.match(relay.expr.Let(lv, x < y, lv * x)) |
| |
| |
| def test_match_option(): |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| pattern = is_op("nn.relu")( |
| is_op("nn.conv2d")(wildcard(), wildcard()).optional( |
| lambda x: is_op("nn.bias_add")(x, wildcard()) |
| ) |
| ) |
| |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.nn.relu(conv2d) |
| assert pattern.match(relu) |
| |
| conv2d = relay.op.nn.conv2d(x, w) |
| bias_add = relay.op.nn.bias_add(conv2d, b) |
| relu = relay.op.nn.relu(bias_add) |
| assert pattern.match(relu) |
| |
| pattern = is_op("nn.conv2d")(wildcard(), wildcard()) |
| pattern = pattern.optional(is_op("nn.relu")).optional(is_op("tanh")) |
| |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.nn.relu(conv2d) |
| tanh = relay.op.tanh(conv2d) |
| tanh2 = relay.op.tanh(relu) |
| relu2 = relay.op.nn.relu(tanh) |
| assert pattern.match(conv2d) |
| assert pattern.match(relu) |
| assert pattern.match(tanh) |
| assert pattern.match(tanh2) |
| assert not pattern.match(relu2) |
| |
| |
| def test_no_match_option(): |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| pattern = is_op("nn.relu")( |
| is_op("nn.conv2d")(wildcard(), wildcard()).optional( |
| lambda x: is_op("nn.bias_add")(x, wildcard()) |
| ) |
| ) |
| |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.tanh(conv2d) |
| assert not pattern.match(relu) |
| |
| conv2d = relay.op.nn.dense(x, w) |
| relu = relay.op.tanh(conv2d) |
| assert not pattern.match(relu) |
| |
| conv2d = relay.op.nn.dense(x, w) |
| bias_add = relay.op.nn.bias_add(conv2d, b) |
| relu = relay.op.nn.relu(bias_add) |
| assert not pattern.match(relu) |
| |
| conv2d = relay.op.nn.conv2d(x, w) |
| bias_add = conv2d + w |
| relu = relay.op.nn.relu(bias_add) |
| assert not pattern.match(relu) |
| |
| |
| def test_match_const(): |
| conv2d = is_op("nn.conv2d")(wildcard(), is_constant()) |
| pattern = is_op("nn.bias_add")(conv2d, wildcard()) |
| |
| x = relay.var("x", shape=(1, 3, 224, 224)) |
| w = relay.var("w", shape=(3, 3, 3, 3)) |
| b = relay.var("b", shape=(3,)) |
| conv2d = relay.op.nn.conv2d(x, w) |
| out = relay.op.nn.bias_add(conv2d, b) |
| func = relay.Function([x, w, b], out) |
| mod = tvm.IRModule.from_expr(func) |
| |
| assert not pattern.match(mod["main"].body) |
| mod["main"] = bind_params_by_name(mod["main"], {"w": tvm.nd.array(np.ones(shape=(3, 3, 3, 3)))}) |
| assert pattern.match(mod["main"].body) |
| |
| |
| def test_match_tuple(): |
| x = relay.var("x") |
| y = relay.var("y") |
| z = relay.op.op.get("add") |
| tuple_pattern = is_tuple((is_var("x"), wildcard(), is_op("add"))) |
| assert tuple_pattern.match(relay.expr.Tuple((x, y, z))) |
| |
| tuple_pattern = is_tuple((is_var("x"), wildcard(), is_op("add"))) |
| tuple_get_item_pattern = is_tuple_get_item(tuple_pattern, 1) |
| assert tuple_get_item_pattern.match(relay.expr.TupleGetItem(relay.expr.Tuple((x, y, z)), 1)) |
| |
| tuple_get_item_pattern = is_tuple_get_item(tuple_pattern) # Match any index |
| assert tuple_get_item_pattern.match(relay.expr.TupleGetItem(relay.expr.Tuple((x, y, z)), 0)) |
| assert tuple_get_item_pattern.match(relay.expr.TupleGetItem(relay.expr.Tuple((x, y, z)), 1)) |
| assert tuple_get_item_pattern.match(relay.expr.TupleGetItem(relay.expr.Tuple((x, y, z)), 2)) |
| |
| # Match tuple with any inputs |
| tuple_pattern = is_tuple(None) |
| concat_pattern = is_op("concatenate")(tuple_pattern) |
| assert concat_pattern.match(relay.op.concatenate(relay.expr.Tuple((x,)), axis=0)) |
| assert concat_pattern.match(relay.op.concatenate(relay.expr.Tuple((x, y)), axis=0)) |
| assert concat_pattern.match(relay.op.concatenate(relay.expr.Tuple((x, y, z)), axis=0)) |
| |
| |
| def test_no_match_tuple(): |
| x = relay.var("x") |
| y = relay.var("y") |
| z = relay.op.op.get("add") |
| tuple_pattern = is_tuple((is_var("x"), wildcard(), is_op("add"), wildcard())) |
| assert not tuple_pattern.match(relay.expr.Tuple((x, y, z))) |
| |
| tuple_pattern = is_tuple((is_var("x"), wildcard(), is_op("add"))) |
| tuple_get_item_pattern = is_tuple_get_item(tuple_pattern, 1) |
| assert not tuple_get_item_pattern.match(relay.expr.TupleGetItem(relay.expr.Tuple((x, y, z)), 2)) |
| |
| |
| def test_match_type(): |
| x = relay.var("x", shape=(10, 10), dtype="float32") |
| ty_pat = has_type(relay.TensorType((10, 10), "float32")) |
| assert ty_pat.match(x) |
| |
| |
| def test_no_match_type(): |
| x = relay.var("x", shape=(10, 10), dtype="int32") |
| ty_pat = has_type(relay.TensorType((10, 10), "float32")) |
| assert not ty_pat.match(x) |
| |
| |
| def test_match_dtype(): |
| x = relay.var("x", shape=(10, 10), dtype="float32") |
| ty_pat = has_dtype("float32") |
| assert ty_pat.match(x) |
| |
| |
| def test_no_match_dtype(): |
| x = relay.var("x", shape=(10, 10), dtype="int32") |
| ty_pat = has_dtype("float32") |
| assert not ty_pat.match(x) |
| |
| |
| def test_match_shape(): |
| x = relay.var("x", shape=(10, 10), dtype="float32") |
| ty_pat = has_shape((10, 10)) |
| assert ty_pat.match(x) |
| |
| |
| def test_no_match_shape(): |
| x = relay.var("x", shape=(10, 10), dtype="int32") |
| ty_pat = has_shape((10, 5)) |
| assert not ty_pat.match(x) |
| |
| |
| def test_match_op_attr(): |
| op = is_op("add").has_attr({"TOpPattern": K_BROADCAST}) |
| op_pat = op(wildcard(), wildcard()) |
| x = relay.var("x") |
| y = relay.var("y") |
| assert op_pat.match(x + y) |
| |
| |
| def test_no_match_op_attr(): |
| op = is_op("nn.dense").has_attr({"TOpPattern": K_ELEMWISE}) |
| op_pat = op(wildcard(), wildcard()) |
| x = relay.var("x") |
| y = relay.var("y") |
| assert not op_pat.match(relay.op.nn.dense(x, y)) |
| op = is_op("add").has_attr({"TOpPattern": K_BROADCAST}) |
| op_pat = op(wildcard(), wildcard()) |
| x = relay.var("x") |
| y = relay.var("y") |
| assert not op_pat.match(x - y) |
| z = relay.var("z") |
| assert not op_pat.match(relay.Let(z, x + y, z)) |
| |
| |
| def test_match_func_attr(): |
| pattern = wildcard().has_attr({"Composite": "add"}) |
| x = relay.var("x") |
| y = relay.var("y") |
| f = relay.Function([x, y], x + y).with_attr("Composite", "add") |
| assert pattern.match(f) |
| |
| |
| def test_no_match_func_attr(): |
| pattern = wildcard().has_attr({"Composite": "add"}) |
| x = relay.var("x") |
| y = relay.var("y") |
| |
| f = relay.Function([x, y], x + y).with_attr("RandomTest", "add") |
| assert not pattern.match(f) |
| f = relay.Function([x, y], x + y).with_attr("Composite", "conv_bias") |
| assert not pattern.match(f) |
| |
| |
| def test_match_call_attr(): |
| # String attr |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()).has_attr({"data_layout": "NCHW"}) |
| x = relay.var("x") |
| y = relay.var("y") |
| assert is_conv2d.match(relay.op.nn.conv2d(x, y)) |
| |
| # Array attr |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()).has_attr({"kernel_size": [3, 3]}) |
| out = relay.op.nn.conv2d(x, y, kernel_size=[3, 3]) |
| assert is_conv2d.match(out) |
| |
| # non-operator call |
| attr_dict = {"call_attr": "attr"} |
| call_has_attr = wildcard()(wildcard()).has_attr(attr_dict) |
| call_attr = tvm.ir.make_node("DictAttrs", **attr_dict) |
| a = relay.Var("a") |
| b = relay.Var("b") |
| assert call_has_attr.match(relay.Call(a, [b], attrs=call_attr)) |
| |
| # empty attrs should match anything |
| empty_attrs = tvm.ir.make_node("DictAttrs", **{}) |
| call_has_empty_attrs = wildcard()(wildcard()).has_attr({}) |
| assert call_has_empty_attrs.match(relay.Call(a, [b], attrs=empty_attrs)) |
| assert call_has_empty_attrs.match(relay.Call(a, [b], attrs=call_attr)) |
| |
| |
| def test_no_match_call_attr(): |
| x = relay.var("x") |
| y = relay.var("y") |
| |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()).has_attr({"data_layout": "NHWC"}) |
| assert not is_conv2d.match(relay.op.nn.conv2d(x, y)) |
| |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()).has_attr({"RandomAttr": "NCHW"}) |
| assert not is_conv2d.match(relay.op.nn.conv2d(x, y)) |
| |
| # Array attr |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()).has_attr({"kernel_size": [3, 3]}) |
| out = relay.op.nn.conv2d(x, y, kernel_size=[2, 1]) |
| assert not is_conv2d.match(out) |
| |
| # non-operator calls |
| call_has_attr = wildcard()(wildcard()).has_attr({"call_attr": "attr"}) |
| wrong_key = tvm.ir.make_node("DictAttrs", **{"wrong": "attr"}) |
| wrong_value = tvm.ir.make_node("DictAttrs", **{"call_attr": "wrong"}) |
| empty_attrs = tvm.ir.make_node("DictAttrs", **{}) |
| |
| a = relay.Var("a") |
| b = relay.Var("b") |
| # attrs left undefined |
| assert not call_has_attr.match(relay.Call(a, [b])) |
| # wrong attrs |
| assert not call_has_attr.match(relay.Call(a, [b], attrs=wrong_key)) |
| assert not call_has_attr.match(relay.Call(a, [b], attrs=wrong_value)) |
| assert not call_has_attr.match(relay.Call(a, [b], attrs=empty_attrs)) |
| |
| |
| def test_match_call_attr_dtype(): |
| is_cast = is_op("cast")(wildcard()).has_attr({"dtype": "float32"}) |
| x = relay.var("x") |
| assert is_cast.match(relay.op.cast(x, "float32")) |
| |
| |
| def test_match_diamond(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| path1 = is_op("nn.relu")(is_conv2d) |
| path2 = is_op("nn.leaky_relu")(is_conv2d) |
| diamond = is_op("add")(path1, path2) |
| |
| # Expr |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert diamond.match(out) |
| |
| |
| def test_no_match_diamond(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| path1 = is_op("nn.relu")(is_conv2d) |
| path2 = is_op("nn.leaky_relu")(is_conv2d) |
| diamond = is_op("add")(path1, path2) |
| |
| # Expr |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| |
| # Check |
| assert not diamond.match(leaky_relu) |
| assert not diamond.match(relu) |
| |
| |
| def test_match_fake_diamond(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| path1 = is_op("nn.relu")(is_conv2d) |
| path2 = is_op("nn.leaky_relu")(is_conv2d) |
| diamond = is_op("add")(path1, path2) |
| |
| # Expr |
| input1 = relay.var("input1") |
| weight1 = relay.var("weight1") |
| conv2d1 = relay.op.nn.conv2d(input1, weight1) |
| inp2 = relay.var("input2") |
| weight2 = relay.var("weight2") |
| conv2d2 = relay.op.nn.conv2d(inp2, weight2) |
| relu = relay.op.nn.relu(conv2d1) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d2, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert not diamond.match(out) |
| |
| |
| def test_at_most_one_parent(): |
| # Pattern |
| P = is_op("nn.conv2d")(wildcard(), wildcard()) # 'parent' |
| I = is_op("nn.relu")(wildcard()) # 'intermediate' ('path' in the code) |
| C = is_op("add")(wildcard(), wildcard()) # 'child' |
| pattern = dominates(P, I, C) |
| |
| # n6(P) |
| # / \ |
| # n7 \ |
| # / \ |
| # n8(P) n10(I) |
| # \ / |
| # n9(I) / |
| # \ / |
| # n11(C) |
| |
| x = relay.var("x") |
| w = relay.var("w") |
| n6 = relay.op.nn.conv2d(x, w) # matches P |
| n7 = relay.op.tanh(n6) # does not match I |
| n8 = relay.op.nn.conv2d(n7, w) # matches P |
| n9 = relay.op.nn.relu(n8) # matches I |
| n10 = relay.op.nn.relu(n6) # matches I |
| n11 = relay.add(n9, n10) # matches C |
| |
| # Does not match: Can't match the parent pattern P at both 8 and 6. |
| # Note that if we did allow P to be used twice the implementation would |
| # need to be changed to not 'jump over' n7. |
| assert not pattern.match(n11) |
| |
| |
| def test_match_dominator(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()) |
| reduction = is_op("add")(wildcard(), wildcard()) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| # Classic Diamond |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert diamond.match(out) |
| |
| # Deeper Branch |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| relu = relay.op.tanh(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert diamond.match(out) |
| |
| # Single Branch |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| tanh = relay.op.tanh(relu) |
| out = relu + tanh |
| |
| # Check |
| assert diamond.match(out) |
| |
| # Fuzzy path/nested Diamond |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()) | is_op( |
| "add" |
| )(wildcard(), wildcard()) |
| reduction = is_op("add")(wildcard(), wildcard()) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relu + relu |
| tanh = relay.op.tanh(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = tanh + leaky_relu |
| |
| assert diamond.match(out) |
| |
| |
| def test_not_match_dominator(): |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()) |
| reduction = is_op("add")(wildcard(), wildcard()) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| # Fake Diamond |
| input1 = relay.var("input1") |
| weight1 = relay.var("weight1") |
| conv2d1 = relay.op.nn.conv2d(input1, weight1) |
| inp2 = relay.var("input2") |
| weight2 = relay.var("weight2") |
| conv2d2 = relay.op.nn.conv2d(inp2, weight2) |
| relu = relay.op.nn.relu(conv2d1) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d2, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert not diamond.match(out) |
| |
| # Add op that doesn't match K_ELEMWISE |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relu + relu |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert not diamond.match(out) |
| |
| # Relu on the input instead of the conv |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(inp) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert not diamond.match(out) |
| |
| # No conv |
| inp = relay.var("input") |
| relu = relay.op.nn.relu(inp) |
| relu = relay.op.nn.relu(relu) |
| tanh = relay.op.tanh(relu) |
| out = relu + tanh |
| |
| # Check |
| assert not diamond.match(out) |
| |
| |
| def test_not_match_dominator2(): |
| # Pattern |
| P = is_op("nn.conv2d")(wildcard(), wildcard()) # 'parent' |
| I = is_op("nn.relu")(wildcard()) # 'intermediate' ('path' in the code) |
| C = is_op("add")(wildcard(), wildcard()) # 'child' |
| pattern = dominates(P, I, C) |
| |
| # n6(P) |
| # / \ |
| # n7 \ |
| # / \ |
| # n8(P) n9(I) |
| # \ / |
| # \ / |
| # \ / |
| # n10(C) |
| |
| x = relay.var("x") |
| w = relay.var("w") |
| n6 = relay.op.nn.conv2d(x, w) # matches P |
| n7 = relay.op.tanh(n6) # does not match I |
| n8 = relay.op.nn.conv2d(n7, w) # matches P |
| n9 = relay.op.nn.relu(n6) # matches I |
| n10 = relay.add(n8, n9) # matches C |
| |
| # Does not match: Can't match the parent pattern P at both 8 and 6. |
| # Note that if we did allow P to be used twice the implementation would |
| # need to be changed to not 'jump over' n7. |
| assert not pattern.match(n10) |
| |
| |
| def test_match_typed_dominator(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()).has_dtype( |
| "float32" |
| ) |
| reduction = is_op("add")(wildcard(), wildcard()).has_shape([1, 3, 10, 10]) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| # Classic Diamond |
| inp = relay.var("input", relay.TensorType((1, 3, 12, 12), "float32")) |
| weight = relay.var("weight", relay.TensorType((3, 3, 3, 3), "float32")) |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Check |
| assert diamond.match(out) |
| |
| |
| def test_no_match_typed_dominator(): |
| # Classic Diamond |
| inp = relay.var("input", relay.TensorType((1, 3, 12, 12), "float32")) |
| weight = relay.var("weight", relay.TensorType((3, 3, 3, 3), "float32")) |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()).has_dtype( |
| "float32" |
| ) |
| reduction = is_op("add")(wildcard(), wildcard()).has_shape([1, 1, 10, 10]) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| # Check |
| assert not diamond.match(out) |
| |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()).has_dtype( |
| "float16" |
| ) |
| reduction = is_op("add")(wildcard(), wildcard()).has_shape([1, 3, 10, 10]) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| # Check |
| assert not diamond.match(out) |
| |
| |
| def test_rewrite(): |
| x = relay.var("x") |
| y = relay.var("y") |
| add_pattern = is_op("add")(wildcard(), wildcard()) |
| sub_pattern = is_op("subtract")(wildcard(), wildcard()) |
| |
| class TestRewrite(DFPatternCallback): |
| def __init__(self): |
| super(TestRewrite, self).__init__() |
| self.pattern = add_pattern |
| |
| def callback(self, pre, post, node_map): |
| return post.args[0] - post.args[1] |
| |
| out = rewrite(TestRewrite(), x + y) |
| assert sub_pattern.match(out) |
| |
| |
| def test_rewrite_func(): |
| x = relay.var("x") |
| w = relay.var("w") |
| y = relay.var("y") |
| add_pattern = is_op("add")(wildcard(), wildcard()) |
| sub_pattern = is_op("subtract")(wildcard(), wildcard()) |
| |
| class TestRewrite(DFPatternCallback): |
| def __init__(self): |
| super(TestRewrite, self).__init__() |
| self.pattern = add_pattern |
| |
| def callback(self, pre, post, node_map): |
| return post.args[0] - post.args[1] |
| |
| inpf = relay.var("input") |
| weightf = relay.var("weight") |
| func = relay.Function( |
| [inpf, weightf], relay.op.nn.relu(relay.op.nn.conv2d(inpf, weightf)), attrs=None |
| ) |
| out = rewrite(TestRewrite(), func(x, w) + y) |
| assert sub_pattern.match(out) |
| |
| |
| def test_rewrite_func_with_attr(): |
| x = relay.var("x") |
| y = relay.var("y") |
| f = relay.Function([x, y], x + y).with_attr("Composite", "add") |
| |
| a = relay.var("a") |
| b = relay.var("b") |
| c = relay.Call(f, [a, b]) |
| c_abs = relay.abs(c) |
| |
| class TestRewrite(DFPatternCallback): |
| def __init__(self): |
| super(TestRewrite, self).__init__() |
| self.pattern = wildcard().has_attr({"Composite": "add"})(wildcard(), wildcard()) |
| |
| def callback(self, pre, post, node_map): |
| return post.args[0] + post.args[1] |
| |
| out = rewrite(TestRewrite(), c_abs) |
| inlined_add_pattern = is_op("abs")(is_op("add")(wildcard(), wildcard())) |
| assert inlined_add_pattern.match(out) |
| |
| |
| def test_nested_rewrite(): |
| class PatternCallback(DFPatternCallback): |
| def __init__(self, pattern): |
| super(PatternCallback, self).__init__() |
| self.pattern = pattern |
| |
| def callback(self, pre, post, node_map): |
| return post |
| |
| def gen(): |
| x = relay.var("x") |
| y = relay.var("y") |
| y_add = relay.add(y, y) |
| n0 = relay.add(x, y_add) |
| n1 = relay.add(x, n0) |
| return relay.add(n1, n0) |
| |
| def pattern(): |
| a = wildcard() |
| b = wildcard() |
| n0 = is_op("add")(a, b) |
| n1 = is_op("add")(n0, a) |
| return is_op("add")(n0, n1) |
| |
| out = gen() |
| pat = pattern() |
| new_out = rewrite(PatternCallback(pat), out) |
| |
| assert tvm.ir.structural_equal(out, new_out) |
| |
| |
| def test_not_fuse_multi_diamond(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| path1 = is_op("nn.relu")(is_conv2d) |
| path2 = is_op("nn.leaky_relu")(is_conv2d) |
| diamond = is_op("add")(path1, path2) |
| |
| # Expr |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| out = out + conv2d |
| # Check |
| assert not diamond.match(out) |
| |
| |
| class BatchnormCallback(DFPatternCallback): |
| def __init__(self): |
| super(BatchnormCallback, self).__init__() |
| self.x = wildcard() |
| self.var = wildcard() |
| self.mean = wildcard() |
| self.beta = wildcard() |
| self.gamma = wildcard() |
| self.eps = is_constant() |
| |
| self.pattern = ( |
| self.gamma * (self.x - self.mean) / is_op("sqrt")(self.var + self.eps) + self.beta |
| ) |
| |
| def callback(self, pre, post, node_map): |
| x = node_map[self.x][0] |
| var = node_map[self.var][0] |
| mean = node_map[self.mean][0] |
| beta = node_map[self.beta][0] |
| gamma = node_map[self.gamma][0] |
| eps = node_map[self.eps][0] |
| return relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=eps.data.numpy().item())[0] |
| |
| |
| def test_fuse_batchnorm(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| |
| BN = gamma * (x - mean) / relay.op.sqrt(var + relay.const(1e-5)) + beta |
| |
| out = rewrite(BatchnormCallback(), BN) |
| assert tvm.ir.structural_equal( |
| out, relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=1e-5)[0] |
| ) |
| |
| |
| def test_no_fuse_batchnorm(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| |
| fake_BN = gamma * (x - mean) / relay.op.sqrt(var + relay.const(1e-5)) - beta |
| |
| out = rewrite(BatchnormCallback(), fake_BN) |
| assert tvm.ir.structural_equal(out, fake_BN) |
| |
| |
| def test_fuse_double_batchnorm(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| |
| BN = gamma * (x - mean) / relay.op.sqrt(var + relay.const(1e-5)) + beta |
| BN2 = gamma * (BN - mean) / relay.op.sqrt(var + relay.const(1e-5)) + beta |
| |
| out = rewrite(BatchnormCallback(), BN2) |
| |
| bn = relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=1e-5)[0] |
| bn2 = relay.op.nn.batch_norm(bn, gamma, beta, mean, var, epsilon=1e-5)[0] |
| |
| assert tvm.ir.structural_equal(out, bn2) |
| |
| |
| def test_partial_fuse_double_batchnorm(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| |
| BN = gamma * (x - mean) / relay.op.sqrt(var + relay.const(1e-5)) - beta |
| BN2 = gamma * (BN - mean) / relay.op.sqrt(var + relay.const(1e-5)) + beta |
| |
| out = rewrite(BatchnormCallback(), BN2) |
| |
| bn2 = relay.op.nn.batch_norm(BN, gamma, beta, mean, var, epsilon=1e-5)[0] |
| |
| assert tvm.ir.structural_equal(out, bn2) |
| |
| |
| def test_fuse_batchnorm_commutation(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| |
| # commute add |
| BN = beta + gamma * (x - mean) / relay.op.sqrt(var + relay.const(1e-5)) |
| out = rewrite(BatchnormCallback(), BN) |
| assert tvm.ir.structural_equal( |
| out, relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=1e-5)[0] |
| ) |
| |
| # associate divide/multiply |
| BN = (gamma * (x - mean)) / relay.op.sqrt(var + relay.const(1e-5)) + beta |
| out = rewrite(BatchnormCallback(), BN) |
| assert tvm.ir.structural_equal( |
| out, relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=1e-5)[0] |
| ) |
| |
| # associate multiply/divide |
| BN = gamma * ((x - mean) / relay.op.sqrt(var + relay.const(1e-5))) + beta |
| out = rewrite(BatchnormCallback(), BN) |
| assert tvm.ir.structural_equal( |
| out, relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=1e-5)[0] |
| ) |
| |
| |
| def test_quadruple_rewrite_dominator(): |
| class DominatorRemovalCallback(DFPatternCallback): |
| def __init__(self): |
| super(DominatorRemovalCallback, self).__init__() |
| self.inp = wildcard() |
| self.weight = wildcard() |
| is_conv2d = is_op("nn.conv2d")(self.inp, self.weight) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))( |
| wildcard() |
| ) | is_op("add")(wildcard(), wildcard()) |
| reduction = is_op("add")(wildcard(), wildcard()) |
| self.pattern = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| def callback(self, pre, post, node_map): |
| inp = node_map[self.inp][0] |
| weight = node_map[self.weight][0] |
| return relay.op.nn.conv2d(inp, weight) |
| |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| # Classic Diamond |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Deeper Branch |
| conv2d = relay.op.nn.conv2d(out, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| relu = relay.op.tanh(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = relu + leaky_relu |
| |
| # Single Branch |
| conv2d = relay.op.nn.conv2d(out, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| tanh = relay.op.tanh(relu) |
| out = relu + tanh |
| |
| # Fuzzy path/nested Diamond |
| conv2d = relay.op.nn.conv2d(out, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relu + relu |
| tanh = relay.op.tanh(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| out = tanh + leaky_relu |
| one = relay.op.nn.conv2d(inp, weight) |
| two = relay.op.nn.conv2d(one, weight) |
| three = relay.op.nn.conv2d(two, weight) |
| four = relay.op.nn.conv2d(three, weight) |
| |
| assert tvm.ir.structural_equal(DominatorRemovalCallback().rewrite(out), four) |
| |
| |
| def algebraic_simplify(expr): |
| zero = is_expr(relay.const(0)) | is_expr(relay.const(0.0)) |
| one = is_expr(relay.const(1)) | is_expr(relay.const(1.0)) |
| |
| class ElwiseNullCallback(DFPatternCallback): |
| def callback(self, pre, post, node_map): |
| return node_map[self.x][0] # pylint: disable=no-member |
| |
| class AddCallback(ElwiseNullCallback): |
| def __init__(self): |
| super(AddCallback, self).__init__() |
| self.x = wildcard() |
| self.pattern = self.x + zero |
| |
| class SubCallback(ElwiseNullCallback): |
| def __init__(self): |
| super(SubCallback, self).__init__() |
| self.x = wildcard() |
| self.pattern = self.x - zero |
| |
| class MulCallback(ElwiseNullCallback): |
| def __init__(self): |
| super(MulCallback, self).__init__() |
| self.x = wildcard() |
| self.pattern = self.x * one |
| |
| class DivCallback(ElwiseNullCallback): |
| def __init__(self): |
| super(DivCallback, self).__init__() |
| self.x = wildcard() |
| self.pattern = self.x / one |
| |
| class MulZeroCallback(ElwiseNullCallback): |
| def __init__(self): |
| super(MulZeroCallback, self).__init__() |
| self.x = zero |
| self.pattern = self.x * wildcard() |
| |
| class ZeroDivCallback(ElwiseNullCallback): |
| def __init__(self): |
| super(ZeroDivCallback, self).__init__() |
| self.x = zero |
| self.pattern = self.x / wildcard() |
| |
| return rewrite( |
| [ |
| AddCallback(), |
| SubCallback(), |
| MulCallback(), |
| DivCallback(), |
| MulZeroCallback(), |
| ZeroDivCallback(), |
| ], |
| expr, |
| ) |
| |
| |
| def test_algebraic_simplify(): |
| x = relay.Var("x") |
| y = relay.Var("y") |
| |
| one = relay.const(1) |
| zero = relay.const(0) |
| onef = relay.const(1.0) |
| zerof = relay.const(0.0) |
| |
| assert algebraic_simplify(x + zero) == x |
| assert algebraic_simplify(x + zerof) == x |
| assert algebraic_simplify(zero + x) == x |
| assert algebraic_simplify(zerof + x) == x |
| |
| assert algebraic_simplify(x - zero) == x |
| assert algebraic_simplify(x - zerof) == x |
| |
| assert algebraic_simplify(x * one) == x |
| assert algebraic_simplify(x * onef) == x |
| assert algebraic_simplify(one * x) == x |
| assert algebraic_simplify(onef * x) == x |
| assert algebraic_simplify(x * zero) == zero |
| assert algebraic_simplify(x * zerof) == zerof |
| |
| assert algebraic_simplify(x / one) == x |
| assert algebraic_simplify(x / onef) == x |
| assert algebraic_simplify(zero / x) == zero |
| assert algebraic_simplify(zerof / x) == zerof |
| |
| assert tvm.ir.structural_equal( |
| algebraic_simplify((x + zero * y) / one + (y * one) - zero / x), x + y |
| ) |
| |
| |
| def test_double_partition(): |
| # Pattern 1 |
| conv2d_p = is_op("nn.conv2d")(wildcard(), wildcard()) |
| bias_add_p = is_op("nn.bias_add")(conv2d_p, wildcard()) |
| relu_p = is_op("nn.relu")(bias_add_p) |
| |
| # Graph |
| x = relay.var("input") |
| w = relay.var("weight") |
| b = relay.var("bias") |
| w2 = relay.var("weight") |
| b2 = relay.var("bias") |
| conv2d = relay.op.nn.conv2d(x, w) |
| bias_add = relay.op.nn.bias_add(conv2d, b) |
| relu = relay.op.nn.relu(bias_add) |
| conv2d2 = relay.op.nn.conv2d(relu, w2) |
| bias_add2 = relay.op.nn.bias_add(conv2d2, b2) |
| |
| partitioned = bias_add2 |
| for pat, label in [(relu_p, "conv_bias_relu"), (bias_add_p, "conv_bias")]: |
| partitioned = pat.partition(partitioned, {"Composite": label}) |
| |
| inpf = relay.var("input") |
| weightf = relay.var("weight") |
| biasf = relay.var("bias") |
| func0 = ( |
| relay.Function( |
| [inpf, weightf, biasf], |
| relay.op.nn.relu(relay.op.nn.bias_add(relay.op.nn.conv2d(inpf, weightf), biasf)), |
| ) |
| .with_attr("Composite", "conv_bias_relu") |
| .with_attr("PartitionedFromPattern", "nn.conv2d_nn.bias_add_nn.relu_") |
| ) |
| inpf = relay.var("input") |
| weightf = relay.var("weight") |
| biasf = relay.var("bias") |
| func1 = ( |
| relay.Function( |
| [inpf, weightf, biasf], relay.op.nn.bias_add(relay.op.nn.conv2d(inpf, weightf), biasf) |
| ) |
| .with_attr("Composite", "conv_bias") |
| .with_attr("PartitionedFromPattern", "nn.conv2d_nn.bias_add_") |
| ) |
| |
| expected = func1(func0(x, w, b), w2, b2) |
| assert tvm.ir.structural_equal(partitioned, expected) |
| |
| |
| def test_partition_dominator(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()) |
| reduction = is_op("add")(wildcard(), wildcard()) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| # Classic Diamond |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| |
| def generate_diamond(inp, weight): |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| return relu + leaky_relu |
| |
| out = generate_diamond(inp * inp, weight * weight) |
| # Check |
| partitioned = diamond.partition(out) |
| |
| i = relay.Var("input") |
| w = relay.Var("weight") |
| f = relay.Function([i, w], generate_diamond(i, w)).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_nn.relu_nn.relu_nn.leaky_relu_add_" |
| ) |
| assert tvm.ir.structural_equal(partitioned, f(inp * inp, weight * weight)) |
| |
| |
| def test_quadruple_partition_dominator(): |
| # Pattern |
| is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard()) | is_op( |
| "add" |
| )(wildcard(), wildcard()) |
| reduction = is_op("add")(wildcard(), wildcard()) |
| diamond = dominates(is_conv2d, is_unary_elemwise, reduction) |
| |
| inp = relay.var("input") |
| weight = relay.var("weight") |
| |
| # Classic Diamond |
| def classic_diamond(inp, weight): |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| return relu + leaky_relu |
| |
| # Deeper Branch |
| def deeper_diamond(inp, weight): |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| relu = relay.op.tanh(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| return relu + leaky_relu |
| |
| # Single Branch |
| def single_branch(inp, weight): |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relay.op.nn.relu(relu) |
| tanh = relay.op.tanh(relu) |
| return relu + tanh |
| |
| # Fuzzy path/nested Diamond |
| def nested_diamond(inp, weight): |
| conv2d = relay.op.nn.conv2d(inp, weight) |
| relu = relay.op.nn.relu(conv2d) |
| relu = relu + relu |
| tanh = relay.op.tanh(relu) |
| leaky_relu = relay.op.nn.leaky_relu(conv2d, alpha=0) |
| return tanh + leaky_relu |
| |
| partitioned = diamond.partition( |
| nested_diamond( |
| single_branch(deeper_diamond(classic_diamond(inp, weight), weight), weight), weight |
| ) |
| ) |
| |
| functions = [] |
| partition_names = [ |
| "nn.conv2d_nn.relu_nn.relu_nn.leaky_relu_add_", |
| "nn.conv2d_nn.relu_nn.relu_tanh_nn.leaky_relu_add_", |
| "nn.conv2d_nn.relu_nn.relu_tanh_add_", |
| "nn.conv2d_nn.relu_add_tanh_nn.leaky_relu_add_", |
| ] |
| for i, f in enumerate([classic_diamond, deeper_diamond, single_branch, nested_diamond]): |
| inpf = relay.var("input") |
| weightf = relay.var("weight") |
| functions.append( |
| relay.Function([inpf, weightf], f(inpf, weightf)).with_attr( |
| "PartitionedFromPattern", partition_names[i] |
| ) |
| ) |
| |
| reference = functions[3]( |
| functions[2](functions[1](functions[0](inp, weight), weight), weight), weight |
| ) |
| assert tvm.ir.structural_equal(partitioned, reference) |
| |
| |
| def get_BN(x, var, mean, beta, gamma, eps): |
| return gamma * (x - mean) / relay.op.sqrt(var + eps) + beta |
| |
| |
| def test_partition_batchnorm(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| eps = relay.const(1e-5) |
| BN = get_BN(x, var, mean, beta, gamma, eps) |
| |
| xf = relay.var("xf") |
| varf = relay.var("varf") |
| meanf = relay.var("meanf") |
| betaf = relay.var("betaf") |
| gammaf = relay.var("gammaf") |
| # Put the arguments in toplogological order for the reference |
| f = relay.Function( |
| [gammaf, xf, meanf, varf, betaf], get_BN(xf, varf, meanf, betaf, gammaf, eps) |
| ).with_attr("PartitionedFromPattern", "subtract_multiply_add_sqrt_divide_add_") |
| |
| partitioned = BatchnormCallback().pattern.partition(BN) |
| reference = f(gamma, x, mean, var, beta) |
| assert tvm.ir.structural_equal(partitioned, reference) |
| |
| |
| def test_partition_double_batchnorm(): |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| eps = relay.const(1e-5) |
| |
| BN = gamma * (x - mean) / relay.op.sqrt(var + eps) + beta |
| BN2 = gamma * (BN - mean) / relay.op.sqrt(var + eps) + beta |
| |
| xf = relay.var("xf") |
| varf = relay.var("varf") |
| meanf = relay.var("meanf") |
| betaf = relay.var("betaf") |
| gammaf = relay.var("gammaf") |
| f1 = relay.Function( |
| [gammaf, xf, meanf, varf, betaf], get_BN(xf, varf, meanf, betaf, gammaf, eps) |
| ).with_attr("PartitionedFromPattern", "subtract_multiply_add_sqrt_divide_add_") |
| # The partitioner doesn't replace duplicates, so we use two copies of the function |
| xf2 = relay.var("xf2") |
| varf2 = relay.var("varf2") |
| meanf2 = relay.var("meanf2") |
| betaf2 = relay.var("betaf2") |
| gammaf2 = relay.var("gammaf2") |
| f2 = relay.Function( |
| [gammaf2, xf2, meanf2, varf2, betaf2], get_BN(xf2, varf2, meanf2, betaf2, gammaf2, eps) |
| ).with_attr("PartitionedFromPattern", "subtract_multiply_add_sqrt_divide_add_") |
| |
| partitioned = BatchnormCallback().pattern.partition(BN2) |
| reference = f2(gamma, f1(gamma, x, mean, var, beta), mean, var, beta) |
| assert tvm.ir.structural_equal(partitioned, reference) |
| |
| |
| def test_overlappting_partitions(): |
| 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) |
| |
| x = relay.var("x") |
| var = relay.var("var") |
| mean = relay.var("mean") |
| beta = relay.var("beta") |
| gamma = relay.var("gamma") |
| BN = relay.op.nn.batch_norm(x, gamma, beta, mean, var, epsilon=1e-5) |
| T1 = BN[0] |
| T2 = BN[0] |
| add = T1 + T2 |
| |
| assert tuple_get_item_node.partition(add) == add |
| |
| |
| def test_partition_overused(): |
| pattern = is_op("nn.relu")(is_op("nn.conv2d")(wildcard(), wildcard())) |
| |
| x = relay.var("input") |
| w = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.nn.relu(conv2d) |
| out = relu + conv2d |
| |
| assert pattern.partition(out) == out |
| |
| |
| def test_partition_fuzzy_tuple(): |
| x = relay.var("x") |
| y = relay.var("y") |
| z = x + y |
| tuple_pattern = is_tuple(None) |
| concat_pattern = is_op("concatenate")(tuple_pattern) |
| |
| xp = relay.var("xp") |
| yp = relay.var("yp") |
| zp = relay.var("zp") |
| |
| def create_func(args, body): |
| return relay.Function(args, body).with_attr("PartitionedFromPattern", "Tuple_concatenate_") |
| |
| def concat(*args): |
| return relay.op.concatenate(relay.expr.Tuple(args), axis=0) |
| |
| one = concat_pattern.partition(concat(x)) |
| assert tvm.ir.structural_equal(one, create_func([xp], concat(xp))(x)) |
| two = concat_pattern.partition(concat(x, y)) |
| assert tvm.ir.structural_equal(two, create_func([xp, yp], concat(xp, yp))(x, y)) |
| three = concat_pattern.partition(concat(x, y, z)) |
| assert tvm.ir.structural_equal(three, create_func([xp, yp, zp], concat(xp, yp, zp))(x, y, z)) |
| |
| |
| def test_partition_fuzzy_function_args(): |
| func_pattern = FunctionPattern(None, wildcard() + wildcard())(None) + wildcard() |
| x = relay.var("x") |
| y = relay.var("y") |
| z = relay.var("z") |
| b = relay.var("b") |
| xp = relay.var("xp") |
| yp = relay.var("yp") |
| zp = relay.var("zp") |
| |
| def create_func(call): |
| N = len(call.op.params) |
| new_params = [relay.var(str(i)) for i in range(N + 1)] |
| label = "add_FunctionCall_add_" |
| if N == 3: |
| label = "add_" + label |
| return relay.Function( |
| new_params, relay.Call(call.op, (new_params[0:-1])) + new_params[-1] |
| ).with_attr("PartitionedFromPattern", label)(*([x, y, z][0:N] + [b])) |
| |
| f1 = relay.Function([xp], xp + xp)(x) |
| one = func_pattern.partition(f1 + b) |
| assert tvm.ir.structural_equal(one, create_func(f1)) |
| f2 = relay.Function([xp, yp], xp + yp)(x, y) |
| two = func_pattern.partition(f2 + b) |
| assert tvm.ir.structural_equal(two, create_func(f2)) |
| f3 = relay.Function([xp, yp, zp], xp + yp + zp)(x, y, z) |
| three = func_pattern.partition(f3 + b) |
| assert tvm.ir.structural_equal(three, create_func(f3)) |
| |
| |
| def test_partition_check(): |
| pattern = is_op("nn.relu")(is_op("nn.conv2d")(is_var("input"), wildcard())) |
| |
| def check(pre): |
| return pre.args[0].attrs.data_layout == "NCHW" |
| |
| x = relay.var("input") |
| w = relay.var("weight") |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.nn.relu(conv2d) |
| |
| xf = relay.var("input") |
| wf = relay.var("weight") |
| conv2df = relay.op.nn.conv2d(xf, wf) |
| reluf = relay.op.nn.relu(conv2df) |
| func = relay.Function([xf, wf], reluf).with_attr("PartitionedFromPattern", "nn.conv2d_nn.relu_") |
| |
| reference = func(x, w) |
| partitioned = pattern.partition(relu, check=check) |
| assert tvm.ir.structural_equal(partitioned, reference) |
| |
| conv2d = relay.op.nn.conv2d(x, w, data_layout="NHWC") |
| relu = relay.op.nn.relu(conv2d) |
| assert relu == pattern.partition(relu, check=check) |
| |
| |
| def test_partition_check_types(): |
| pattern = is_op("nn.relu")(is_op("nn.conv2d")(wildcard(), wildcard())) |
| |
| def check(pre): |
| conv = pre.args[0] |
| return (conv.attrs.data_layout == "NCHW") and bool(conv.checked_type.shape[0] == 1) |
| |
| x = relay.var("input", shape=(1, 10, 10, 10)) |
| w = relay.var("weight", shape=(10, 10, 3, 3)) |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.nn.relu(conv2d) |
| relu = run_opt_pass(relu, relay.transform.InferType()) |
| |
| partitioned = pattern.partition(relu, check=check) |
| assert partitioned.op.attrs["PartitionedFromPattern"] == "nn.conv2d_nn.relu_" |
| |
| conv2d = relay.op.nn.conv2d(x, w, data_layout="NHWC") |
| relu = relay.op.nn.relu(conv2d) |
| relu = run_opt_pass(relu, relay.transform.InferType()) |
| assert relu == pattern.partition(relu, check=check) |
| |
| x = relay.var("input", shape=(2, 10, 10, 10)) |
| w = relay.var("weight", shape=(10, 10, 3, 3)) |
| conv2d = relay.op.nn.conv2d(x, w) |
| relu = relay.op.nn.relu(conv2d) |
| relu = run_opt_pass(relu, relay.transform.InferType()) |
| assert relu == pattern.partition(relu, check=check) |
| |
| |
| def conv_bias_relu(x, w, b): |
| conv2d = relay.op.nn.conv2d(x, w) |
| bias_add = relay.op.nn.bias_add(conv2d, b) |
| relu = relay.op.nn.relu(bias_add) |
| return relu |
| |
| |
| def test_partition_option(): |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| |
| conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| bias = conv2d.optional(lambda x: is_op("nn.bias_add")(x, wildcard())) |
| pattern1 = is_op("nn.relu")(bias) |
| |
| conv2d = is_op("nn.conv2d")(wildcard(), wildcard()) |
| bias = is_op("nn.bias_add")(conv2d, wildcard()) |
| pattern2 = bias.optional(lambda x: is_op("nn.relu")(x)) |
| |
| relu = conv_bias_relu(x, w, b) |
| |
| xf = relay.var("x") |
| wf = relay.var("w") |
| bf = relay.var("b") |
| func = relay.Function([xf, wf, bf], conv_bias_relu(xf, wf, bf)).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_nn.bias_add_nn.relu_" |
| ) |
| |
| assert pattern1.match(relu) |
| assert tvm.ir.structural_equal(func(x, w, b), pattern1.partition(relu)) |
| |
| assert pattern2.match(relu) |
| assert tvm.ir.structural_equal(func(x, w, b), pattern2.partition(relu)) |
| |
| |
| def test_partition_function(): |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| |
| x1 = relay.var("x1") |
| w1 = relay.var("w1") |
| |
| wc_x = wildcard() |
| wc_w = wildcard() |
| wc_b = wildcard() |
| wc_x1 = wildcard() |
| wc_w1 = wildcard() |
| |
| func_pattern = FunctionPattern([wc_x1, wc_w1], is_op("nn.conv2d")(wc_x1, wc_w1)) |
| pattern = func_pattern(wc_x, wc_w) + wc_b |
| |
| func = relay.Function([x1, w1], relay.nn.conv2d(x1, w1)) |
| expr = func(x, w) + b + b |
| |
| x2 = relay.var("x2") |
| w2 = relay.var("w2") |
| b2 = relay.var("b2") |
| func2 = relay.Function([x2, w2, b2], func(x2, w2) + b2).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_FunctionCall_add_" |
| ) |
| expr2 = func2(x, w, b) + b |
| assert tvm.ir.structural_equal(pattern.partition(expr), expr2) |
| |
| |
| def test_partition_optional_function(): |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| |
| x1 = relay.var("x1") |
| w1 = relay.var("w1") |
| |
| wc_x = wildcard() |
| wc_w = wildcard() |
| wc_x1 = wildcard() |
| wc_w1 = wildcard() |
| |
| func_pattern0 = FunctionPattern( |
| [wc_x1, wc_w1], is_op("sigmoid")(is_op("nn.conv2d")(wc_x1, wc_w1)) |
| ) |
| func_pattern1 = FunctionPattern( |
| [wc_x1, wc_w1], is_op("nn.relu")(is_op("nn.conv2d")(wc_x1, wc_w1)) |
| ) |
| pattern = func_pattern0(wc_x, wc_w) | func_pattern1(wc_x, wc_w) |
| |
| func = relay.Function([x1, w1], relay.nn.relu(relay.nn.conv2d(x1, w1))) |
| expr = func(x, w) + b |
| |
| x2 = relay.var("x2") |
| w2 = relay.var("w2") |
| func2 = relay.Function([x2, w2], func(x2, w2)).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_nn.relu_FunctionCall_" |
| ) |
| expr2 = func2(x, w) + b |
| assert tvm.ir.structural_equal(pattern.partition(expr), expr2) |
| |
| |
| def test_rewrite_function_with_fuzzy_body(): |
| """Allow Rewriting a function with a fuzzy body via dominator analysis""" |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| |
| x1 = relay.var("x1") |
| w1 = relay.var("w1") |
| |
| wc_x = wildcard() |
| wc_w = wildcard() |
| wc_b = wildcard() |
| wc_x1 = wildcard() |
| wc_w1 = wildcard() |
| |
| func_pattern = FunctionPattern([wc_x1, wc_w1], wildcard()) |
| pattern = func_pattern(wc_x, wc_w) + wc_b |
| |
| func = relay.Function([x1, w1], relay.nn.conv2d(x1, w1)) |
| expr = func(x, w) + b + b |
| |
| class TestRewrite(DFPatternCallback): |
| def __init__(self): |
| super(TestRewrite, self).__init__() |
| self.pattern = pattern |
| |
| def callback(self, pre, post, node_map): |
| return x + w |
| |
| out = rewrite(TestRewrite(), expr) |
| assert tvm.ir.structural_equal(out, x + w + b) |
| |
| |
| def test_partition_function_with_fuzzy_body(): |
| """ |
| Allow Rewriting a function with a fuzzy body via dominator analysis |
| """ |
| x = relay.var("x") |
| w = relay.var("w") |
| b = relay.var("b") |
| |
| x1 = relay.var("x1") |
| w1 = relay.var("w1") |
| |
| wc_x = wildcard() |
| wc_w = wildcard() |
| wc_b = wildcard() |
| wc_x1 = wildcard() |
| wc_w1 = wildcard() |
| |
| func_pattern = FunctionPattern([wc_x1, wc_w1], wildcard()) |
| pattern = func_pattern(wc_x, wc_w) + wc_b |
| |
| func = relay.Function([x1, w1], relay.nn.conv2d(x1, w1)) |
| expr = func(x, w) + b + b |
| |
| x2 = relay.var("x2") |
| w2 = relay.var("w2") |
| b2 = relay.var("b2") |
| func2 = relay.Function([x2, w2, b2], func(x2, w2) + b2).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_FunctionCall_add_" |
| ) |
| expr2 = func2(x, w, b) + b |
| assert tvm.ir.structural_equal(pattern.partition(expr), expr2) |
| |
| |
| def test_match_match(): |
| add_pattern = is_op("add")(wildcard(), wildcard()) |
| |
| class TestRewrite(DFPatternCallback): |
| def __init__(self): |
| super(TestRewrite, self).__init__() |
| self.pattern = add_pattern |
| |
| def callback(self, pre, post, node_map): |
| return post.args[0] - post.args[1] |
| |
| mod = tvm.IRModule({}) |
| tvm.relay.prelude.Prelude(mod) |
| # Apply rewrite on IR including relay.Match |
| out = rewrite(TestRewrite(), mod["tensor_concatenate_int64"]) |
| assert tvm.ir.structural_equal(mod["tensor_concatenate_int64"], out) |
| |
| |
| def test_partition_constant_embedding(): |
| x = relay.var("x") |
| w = relay.var("w") |
| wc = relay.const(1) |
| b = relay.var("b") |
| |
| xf = relay.var("x") |
| wf = relay.var("w") |
| bf = relay.var("b") |
| embeded_func = relay.Function([xf, bf], conv_bias_relu(xf, wc, bf)).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_nn.bias_add_nn.relu_" |
| ) |
| xf = relay.var("x") |
| wf = relay.var("w") |
| bf = relay.var("b") |
| lifted_func = relay.Function([xf, wf, bf], conv_bias_relu(xf, wf, bf)).with_attr( |
| "PartitionedFromPattern", "nn.conv2d_nn.bias_add_nn.relu_" |
| ) |
| relu = conv_bias_relu(x, w, b) |
| reluc = conv_bias_relu(x, wc, b) |
| |
| # Check lifting of wildcard matches |
| pattern = is_op("nn.relu")( |
| is_op("nn.bias_add")(is_op("nn.conv2d")(wildcard(), wildcard()), wildcard()) |
| ) |
| assert tvm.ir.structural_equal(lifted_func(x, w, b), pattern.partition(relu)) |
| assert tvm.ir.structural_equal(lifted_func(x, wc, b), pattern.partition(reluc)) |
| |
| # Check lifting of input matches |
| pattern = is_op("nn.relu")( |
| is_op("nn.bias_add")(is_op("nn.conv2d")(wildcard(), is_var()), wildcard()) |
| ) |
| assert tvm.ir.structural_equal(lifted_func(x, w, b), pattern.partition(relu)) |
| assert tvm.ir.structural_equal(reluc, pattern.partition(reluc)) # Constants are not Inputs |
| |
| # Check embedding of constant matches |
| pattern = is_op("nn.relu")( |
| is_op("nn.bias_add")(is_op("nn.conv2d")(wildcard(), is_constant()), wildcard()) |
| ) |
| assert tvm.ir.structural_equal(relu, pattern.partition(relu)) |
| assert tvm.ir.structural_equal(embeded_func(x, b), pattern.partition(reluc)) |
| |
| # Check embedding of constant ExprPatterns |
| pattern = is_op("nn.relu")( |
| is_op("nn.bias_add")(is_op("nn.conv2d")(wildcard(), is_expr(wc)), wildcard()) |
| ) |
| assert tvm.ir.structural_equal(relu, pattern.partition(relu)) |
| assert tvm.ir.structural_equal(embeded_func(x, b), pattern.partition(reluc)) |
| |
| # Check lifting/embedding of Alt matches |
| pattern = is_op("nn.relu")( |
| is_op("nn.bias_add")(is_op("nn.conv2d")(wildcard(), is_var() | is_constant()), wildcard()) |
| ) |
| assert tvm.ir.structural_equal(lifted_func(x, w, b), pattern.partition(relu)) |
| assert tvm.ir.structural_equal(embeded_func(x, b), pattern.partition(reluc)) |
| |
| # Check lifting/embedding of Alt matches with the other ordering |
| pattern = is_op("nn.relu")( |
| is_op("nn.bias_add")(is_op("nn.conv2d")(wildcard(), is_constant() | is_var()), wildcard()) |
| ) |
| assert tvm.ir.structural_equal(lifted_func(x, w, b), pattern.partition(relu)) |
| assert tvm.ir.structural_equal(embeded_func(x, b), pattern.partition(reluc)) |
| |
| |
| def test_rewrite_once(): |
| # This class recursively removes the arguments to concat until there is nothing left to concatenate. |
| class ConcatRewriter(DFPatternCallback): |
| def __init__(self, rewrite_once): |
| super().__init__(rewrite_once=rewrite_once) |
| self.pattern = is_op("concatenate")(None) |
| |
| def callback(self, pre, post, node_map): |
| concat_args = post.args[0] |
| # Remove the last argument |
| new_args = [concat_args[i] for i in range(len(concat_args) - 1)] |
| if new_args: |
| return relay.op.concatenate(relay.expr.Tuple(new_args), axis=0) |
| else: |
| return concat_args[0] |
| |
| x = relay.var("x") |
| y = relay.var("y") |
| z = relay.var("z") |
| concat = relay.op.concatenate(relay.expr.Tuple([x, y, z]), axis=0) |
| |
| def test_one_callback(): |
| # Let the rewriter run recursively |
| out = rewrite(ConcatRewriter(False), concat) |
| expected = x |
| assert tvm.ir.structural_equal(out, expected) |
| |
| # Run the rewriter once |
| out = rewrite(ConcatRewriter(True), concat) |
| expected = relay.op.concatenate(relay.expr.Tuple([x, y]), axis=0) |
| assert tvm.ir.structural_equal(out, expected) |
| |
| def test_multi_callbacks(): |
| # This class recursively add a nn.relu operator after nn.softmax |
| class OneMoreReluRewriter(DFPatternCallback): |
| def __init__(self, rewrite_once): |
| super().__init__(rewrite_once=rewrite_once) |
| self.pattern = is_op("nn.softmax")(None) |
| |
| def callback(self, pre, post, node_map): |
| return relay.nn.relu(post) |
| |
| def before(): |
| # Before: |
| # x y z |
| # | | | |
| # concat |
| # | |
| # softmax |
| return relay.nn.softmax(concat) |
| |
| def once_concat(): |
| # ConcatRewrite once, OneMoreReluRewrite once |
| # Expected: |
| # x y |
| # | | |
| # concat |
| # | |
| # softmax |
| # | |
| # relu |
| return relay.nn.relu( |
| relay.nn.softmax(relay.op.concatenate(relay.expr.Tuple([x, y]), axis=0)) |
| ) |
| |
| def recursive_concat(): |
| # ConcatRewrite recursively, OneMoreReluRewrite once |
| # Expected: |
| # x |
| # | |
| # softmax |
| # | |
| # relu |
| return relay.nn.relu(relay.nn.softmax(x)) |
| |
| # Run ConcatRewriter once, OneMoreReluRewriter once |
| out = rewrite( |
| [OneMoreReluRewriter(True), ConcatRewriter(True)], |
| before(), |
| ) |
| assert tvm.ir.structural_equal(out, once_concat()) |
| |
| # Run ConcatRewriter recursively, OneMoreReluRewriter once |
| out = rewrite( |
| [OneMoreReluRewriter(True), ConcatRewriter(False)], |
| before(), |
| ) |
| assert tvm.ir.structural_equal(out, recursive_concat()) |
| |
| test_one_callback() |
| test_multi_callbacks() |
| |
| |
| def test_matched_outside_but_dominated(): |
| """In this example the pattern matches the nn.conv2d/add/multiply flow. Even though the |
| add output is consumed by the sigmoid, the sigmoid itself is dominated by the multiply. |
| So partitioning can proceed, all be it with a duplication of the add.""" |
| in_mod = tvm.relay.parse( |
| """ |
| #[version = "0.0.5"] |
| def @main(%data: Tensor[(16, 16, 32, 32), float16], %weight: Tensor[(32, 16, 3, 3), float16], %bias: Tensor[(32), float32]) -> Tensor[(16, 32, 32, 32), float32] { |
| %0 = layout_transform(%data, src_layout="NCHW", dst_layout="NHWC"); |
| %1 = layout_transform(%weight, src_layout="OIHW", dst_layout="OHWI"); |
| %2 = expand_dims(%bias, axis=1, num_newaxis=2); |
| %3 = expand_dims(%2, axis=0); |
| %4 = nn.conv2d(%0, %1, padding=[1, 1, 1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="OHWI", out_dtype="float32"); |
| %5 = layout_transform(%3, src_layout="NCHW", dst_layout="NHWC"); |
| %6 = add(%4, %5); |
| %7 = sigmoid(%6); |
| %8 = multiply(%6, %7); |
| layout_transform(%8, src_layout="NHWC", dst_layout="NCHW") |
| } |
| """ |
| ) |
| expected_mod = tvm.relay.parse( |
| """ |
| #[version = "0.0.5"] |
| def @main(%data: Tensor[(16, 16, 32, 32), float16], %weight: Tensor[(32, 16, 3, 3), float16], %bias: Tensor[(32), float32]) -> Tensor[(16, 32, 32, 32), float32] { |
| %2 = expand_dims(%bias, axis=1, num_newaxis=2); |
| %3 = expand_dims(%2, axis=0); |
| %4 = layout_transform(%data, src_layout="NCHW", dst_layout="NHWC"); |
| %5 = layout_transform(%weight, src_layout="OIHW", dst_layout="OHWI"); |
| %6 = nn.conv2d(%4, %5, padding=[1, 1, 1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="OHWI", out_dtype="float32"); |
| %7 = layout_transform(%3, src_layout="NCHW", dst_layout="NHWC"); |
| %8 = add(%6, %7); |
| %9 = sigmoid(%8); |
| %10 = fn (%FunctionVar_0_0, %FunctionVar_0_1, %FunctionVar_0_2, %FunctionVar_0_3, PartitionedFromPattern="nn.conv2d_add_multiply_") { |
| %0 = nn.conv2d(%FunctionVar_0_0, %FunctionVar_0_1, padding=[1, 1, 1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="OHWI", out_dtype="float32"); |
| %1 = add(%0, %FunctionVar_0_2); |
| multiply(%1, %FunctionVar_0_3) |
| }; |
| %11 = %10(%4, %5, %7, %9); |
| layout_transform(%11, src_layout="NHWC", dst_layout="NCHW") |
| } |
| """ |
| ) |
| pattern = is_op("multiply")( |
| is_op("add")(is_op("nn.conv2d")(wildcard(), wildcard()), wildcard()), wildcard() |
| ) |
| actual_mod = tvm.IRModule.from_expr(pattern.partition(in_mod["main"])) |
| actual_mod = relay.transform.InferType()(actual_mod) |
| tvm.ir.assert_structural_equal(actual_mod, expected_mod) |
| |
| |
| def test_partition_parallel_branch_with_same_input(): |
| """In this example, conv2d's two consumer(add and multiply) on two different branches are |
| merged into one partition, make sure that the partitioned function has no redundant parameters""" |
| # Pattern |
| path1 = is_op("multiply")(wildcard(), wildcard()) |
| path2 = is_op("add")(wildcard(), wildcard()) |
| pattern = is_op("add")(path1, path2) |
| |
| i = relay.Var("input") |
| w = relay.Var("weight") |
| l = relay.Var("left") |
| r = relay.Var("right") |
| |
| conv2d = relay.op.nn.conv2d(i, w) |
| branch1 = relay.multiply(l, conv2d) |
| branch2 = relay.add(conv2d, r) |
| add = relay.add(branch1, branch2) |
| |
| lf = relay.Var("leftf") |
| mf = relay.Var("midf") |
| rf = relay.Var("rightf") |
| f = relay.Function([lf, mf, rf], (lf * mf) + (mf + rf)).with_attr( |
| "PartitionedFromPattern", "multiply_add_add_" |
| ) |
| |
| partitioned = pattern.partition(add) |
| reference = f(l, conv2d, r) |
| assert tvm.ir.structural_equal(partitioned, reference) |
| |
| |
| def test_rewrite_with_pattern_recursion(): |
| data = relay.var("data", relay.TensorType((2, 8), "float32")) |
| dense_weight = relay.const(np.zeros((4, 8))) |
| feat = relay.nn.dense(data, dense_weight) |
| feat = relay.cast(feat, "float32") |
| feat = relay.cast(feat, "float32") |
| feat = relay.cast(feat, "float32") |
| feat = relay.cast(feat, "float32") |
| feat = relay.cast(feat, "float32") |
| oup = relay.cast(feat, "float32") |
| |
| expected = relay.nn.relu(oup) |
| |
| class TheRewrite(DFPatternCallback): |
| def __init__(self, pattern): |
| super(TheRewrite, self).__init__(rewrite_once=True) |
| self.pattern = pattern |
| |
| def callback(self, pre, post, node_map): |
| return relay.nn.relu(post) |
| |
| def test_reset_call_args(): |
| dense_pattern = is_op("nn.dense")(wildcard(), wildcard()) |
| wildcard_redirect = wildcard() |
| the_pattern = is_op("cast")(wildcard_redirect) |
| the_pattern2 = the_pattern | dense_pattern |
| wildcard_redirect.redirect_to(the_pattern2) |
| |
| actual = rewrite(TheRewrite(the_pattern), oup) |
| tvm.ir.assert_structural_equal(actual, expected) |
| |
| def test_reset_alt_left(): |
| dense_pattern = is_op("nn.dense")(wildcard(), wildcard()) |
| wildcard_redirect = wildcard() |
| or_pattern = wildcard_redirect | dense_pattern |
| the_pattern = is_op("cast")(or_pattern) |
| wildcard_redirect.redirect_to(the_pattern) |
| |
| actual = rewrite(TheRewrite(the_pattern), oup) |
| tvm.ir.assert_structural_equal(actual, expected) |
| |
| def test_reset_alt_right(): |
| dense_pattern = is_op("nn.dense")(wildcard(), wildcard()) |
| wildcard_redirect = wildcard() |
| or_pattern = dense_pattern | wildcard_redirect |
| the_pattern = is_op("cast")(or_pattern) |
| wildcard_redirect.redirect_to(the_pattern) |
| |
| actual = rewrite(TheRewrite(the_pattern), oup) |
| tvm.ir.assert_structural_equal(actual, expected) |
| |
| test_reset_call_args() |
| test_reset_alt_left() |
| test_reset_alt_right() |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
