# 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. # ruff: noqa: F403, F405, F841 import functools import math import pytest import tvm_ffi import tvm.testing from tvm import relax as rx from tvm import tirx from tvm.relax.analysis import get_var2val from tvm.relax.dpl import * from tvm.script import relax as R from tvm.script import tirx as T @tvm.script.ir_module class Module: @T.prim_func(s_tir=True) def tir_matmul(x: T.handle, y: T.handle, z: T.handle) -> None: T.func_attr({"global_symbol": "tir_matmul"}) k = T.int32() A = T.match_buffer(x, (32, 32)) B = T.match_buffer(y, (32, 32)) C = T.match_buffer(z, (32, 32)) for i0, j0, k0 in T.grid(32, 32, 32): with T.sblock(): i, j, k = T.axis.remap("SSR", [i0, j0, k0]) with T.init(): C[i, j] = 0.0 C[i, j] += A[i, k] * B[j, k] @T.prim_func(s_tir=True) def tir_relu(x: T.handle, y: T.handle): T.func_attr({"global_symbol": "tir_relu"}) A = T.match_buffer(x, (32, 32)) B = T.match_buffer(y, (32, 32)) for i, j in T.grid(32, 32): with T.sblock(): vi, vj = T.axis.remap("SS", [i, j]) B[vi, vj] = T.max(A[vi, vj], 0.0) @T.prim_func(s_tir=True) def tir_zeros(x: T.handle, n: T.int64): T.func_attr({"global_symbol": "tir_zeros"}) A = T.match_buffer(x, [n]) for i in range(n): with T.sblock(): vi = T.axis.remap("S", [i]) A[vi] = 1.0 @R.function def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tuple: cls = Module with R.dataflow(): lv0 = R.call_tir(cls.tir_matmul, (x, w), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_tir(cls.tir_relu, (lv0), R.Tensor((32, 32), dtype="float32")) lv2 = R.call_tir( cls.tir_zeros, [], R.Tensor((32,), dtype="float32"), tir_vars=R.ShapeExpr([32]) ) gv = (lv1, lv2) R.output(gv) return gv main_fn = Module["main"] bindings = main_fn.body.blocks[0].bindings ## Node-wise Matching def test_expr_pattern(): ep = is_expr(rx.Var("x")) assert isinstance(ep, ExprPattern) assert isinstance(ep.expr, rx.Var) def test_var_pattern(): v = is_var("x") assert isinstance(v, VarPattern) assert v.name == "x" assert v.match(rx.Var("x")) assert is_var().match(rx.Var("x")) assert is_var().match(rx.DataflowVar("x")) # DataflowVar is also a Var assert not v.match(rx.GlobalVar("x")) def test_dataflow_var_pattern(): v = is_dfv("x") assert isinstance(v, DataflowVarPattern) assert v.name == "x" assert v.match(rx.DataflowVar("x")) assert not v.match(rx.GlobalVar("x")) assert is_dfv().match(bindings[0].var) def test_global_var_pattern(): assert is_gv("x").match(rx.GlobalVar("x")) # TODO: disabled as regex is not supported due to # symbol conflict with PyTorch # assert is_gv("x.*").match(rx.GlobalVar("x_2")) assert is_gv().match(rx.GlobalVar("x")) assert not is_gv("x").match(rx.GlobalVar("y")) assert not is_gv("x").match(rx.Var("x")) def test_constant_pattern(): c = is_const() assert isinstance(c, ConstantPattern) assert c.match(rx.const([[0.1, 1.1, 2.1], [3.1, 4.1, 5.1]])) def test_wildcard_pattern(): wc = wildcard() assert isinstance(wc, WildcardPattern) assert wc.match(rx.Var("x")) def test_call_pattern(): wc1 = wildcard() wc2 = wildcard() c = is_op("relax.add")(wc1, wc2) assert isinstance(c, CallPattern) assert isinstance(c.args[0], WildcardPattern) assert isinstance(c.args[1], WildcardPattern) assert c.match(rx.op.add(rx.Var("x"), rx.Var("y"))) def test_function_pattern(): wc1 = wildcard() wc2 = wildcard() f = FunctionPattern([wc1, wc2], is_op("relax.add")(wc1, wc2)) 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) x = rx.Var("x", R.Tensor("float32")) y = rx.Var("y", R.Tensor("float32")) assert f.match(rx.Function([x, y], rx.op.add(x, y), ret_ty=R.Tensor("float32"))) assert not f.match(rx.Function([x, y], rx.op.multiply(x, y), ret_ty=R.Tensor("float32"))) def test_tuple_pattern(): wc1 = wildcard() wc2 = is_dfv() t = is_tuple([wc1, wc2]) assert isinstance(t, TuplePattern) assert isinstance(t.fields[0], WildcardPattern) assert isinstance(t.fields[1], DataflowVarPattern) assert t.match(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")])) assert not t.match(rx.Tuple([rx.DataflowVar("x"), rx.GlobalVar("y")])) assert not t.match(rx.Tuple([])) assert t[0].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0)) assert t[1].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 1)) # Negative index is also allowed assert t[-1].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 1)) # None means any index. assert t[None].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0)) assert t[None].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 1)) with pytest.raises(IndexError): t[2] # index cannot be greater than or equal to the tuple size. def test_unordered_tuple_pattern(): t = is_tuple([is_const(), is_dfv()], unordered=True) assert isinstance(t, UnorderedTuplePattern) assert isinstance(t.fields[0], ConstantPattern) assert isinstance(t.fields[1], DataflowVarPattern) assert t.match(rx.Tuple([rx.const([]), rx.DataflowVar("x")])) assert t.match(rx.Tuple([rx.DataflowVar("x"), rx.const([])])) assert not t.match(rx.Tuple([rx.DataflowVar("x"), rx.DataflowVar("y")])) assert not t.match(rx.Tuple([])) def test_tuple_get_item_pattern(): assert is_tuple_get_item(is_tuple([is_gv("x"), is_dfv("y")]), 0).match( rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0) ) assert is_tuple_get_item(is_tuple([is_gv("x"), is_dfv("y")]), 0).match( rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0) ) def test_or_pattern(): dfv_or_gv = is_dfv("x") | is_gv("x") assert isinstance(dfv_or_gv, OrPattern) assert dfv_or_gv.match(rx.DataflowVar("x")) assert dfv_or_gv.match(rx.GlobalVar("x")) assert not dfv_or_gv.match(rx.Var("x")) assert not dfv_or_gv.match(rx.DataflowVar("y")) assert not dfv_or_gv.match(rx.GlobalVar("y")) def test_and_pattern(): # float[2, 3, 3] f32_233 = wildcard().has_shape((2, 3, 3)) & has_dtype("float32") assert isinstance(f32_233, AndPattern) assert f32_233.match(rx.Var("x", R.Tensor((2, 3, 3), "float32"))) assert not f32_233.match(rx.Var("x", R.Tensor((3, 3, 3), "float32"))) assert not f32_233.match(rx.Var("x", R.Tensor("float32", ndim=3))) def test_not_pattern(): no_shape233 = ~wildcard().has_shape((2, 3, 3)) assert isinstance(no_shape233, NotPattern) assert no_shape233.match(rx.Var("x", R.Tensor((3, 3, 3), "float32"))) assert not no_shape233.match(rx.Var("x", R.Tensor((2, 3, 3), "float32"))) def test_dtype_pattern(): dtype = "float16" pattern = has_dtype(dtype) assert isinstance(pattern, DataTypePattern) assert pattern.dtype == dtype assert has_dtype("float32").match(bindings[0].var) def test_shape_pattern(): shape = [32, 32] pattern = wildcard().has_shape(shape) assert isinstance(pattern, ShapePattern) tvm_ffi.structural_equal(pattern.shape, shape) assert pattern.match(bindings[0].var) assert wildcard().has_shape([32, 32]).match(bindings[0].var) n, m = tirx.Var("n", dtype="int64"), tirx.Var("m", dtype="int64") symsh_var = rx.Var("x", R.Tensor([n, m, n + m], "float32")) assert wildcard().has_shape([n, m, n + m]).match(symsh_var) assert wildcard().has_shape([n, m, m + n]).match(symsh_var) # + is commutative. assert not wildcard().has_shape([1, 2, 3]).match(symsh_var) assert not wildcard().has_shape([m, n, n + m]).match(symsh_var) def test_prim_arr_pattern(): """ The difference between is_shape and has_shape is that: 1) is_shape directly matches a shape (e.g., as an argument); 2) has_shape matches a tensor and puts assumptions on the tensor's shape. """ pattern = is_shape([32, 32]) assert pattern[0] == 32 assert pattern[1] == 32 assert isinstance(pattern, PrimArrPattern) assert pattern.match(rx.get_shape_of(bindings[0].var)) n, m = tirx.Var("n", dtype="int64"), tirx.Var("m", dtype="int64") symbolic_shape = rx.ShapeExpr([n, m, n + m]) assert is_shape([n, m, n + m]).match(symbolic_shape) assert not is_shape([n, m, n * m]).match(symbolic_shape) def test_extern_fn_pattern(): pattern = ExternFuncPattern("test.blockbuilder.nop") assert pattern.match(rx.ExternFunc("test.blockbuilder.nop")) def test_op_attr(): x = rx.Var("x", R.Tensor("float32")) y = rx.Var("y", R.Tensor("float32")) conv2d = rx.op.nn.conv2d(x, y, strides=(3, 3)) xp = is_var("x") yp = is_var("y") assert is_op("relax.nn.conv2d")(xp, yp).has_attr({"strides": [3, 3]}).match(conv2d) assert not is_op("relax.nn.conv2d")(xp, yp).has_attr({"strides": [4, 3]}).match(conv2d) def test_match_call_attr(): x = rx.Var("x", R.Tensor("float32")) y = rx.Var("y", R.Tensor("float32")) fn = rx.Function([x, y], rx.op.add(x, y), ret_ty=R.Tensor("float32")) annotated_fn = fn.with_attr({"Codegen": "test-codegen", "global_symbol": "test-symbol"}) xp = is_var("x") yp = is_var("y") root_pattern = FunctionPattern([xp, yp], is_op("relax.add")(xp, yp)) assert root_pattern.has_attr({"Codegen": "test-codegen", "global_symbol": "test-symbol"}).match( annotated_fn ) assert root_pattern.has_attr({"Codegen": "test-codegen"}).match(annotated_fn) assert not root_pattern.has_attr({"ping": "pong"}).match(annotated_fn) assert root_pattern.has_attr({}).match(annotated_fn) def test_is_call_tir(): lv1_val = bindings[1].value lv2_val = bindings[2].value var2val = get_var2val(Module["main"]) assert is_call_tir("tir_relu").match(lv1_val) assert is_call_tir("tir_relu", [is_call_tir("tir_matmul")]).match(lv1_val, var2val=var2val) assert not is_call_tir("tir_relu", [is_call_tir("tir_relu")]).match(lv1_val, var2val=var2val) assert is_call_tir("tir_zeros", wildcard(), wildcard()).match(lv2_val, var2val=var2val) @R.function(pure=False) def simple_call_packed( x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32") ) -> R.Tensor: gv0 = R.call_packed("test.vm.mul", x, w, ty_args=(R.Tensor(ndim=2, dtype="float32"))) return gv0 def test_varg_default_wildcard(): expr = simple_call_packed.body.blocks[0].bindings[0].value yes_pattern_explicit = ExternFuncPattern("test.vm.mul")(wildcard(), wildcard()) yes_pattern_implicit = ExternFuncPattern("test.vm.mul")(varg_default_wildcard=True) no_pattern = ExternFuncPattern("test.vm.mul")(wildcard()) assert yes_pattern_explicit.match(expr) assert yes_pattern_implicit.match(expr) assert not no_pattern.match(expr) def test_simple_call_packed(): expr = simple_call_packed.body.blocks[0].bindings[0].value assert is_call_packed("test.vm.mul").match(expr) assert is_call_packed("test.vm.mul", [is_var("x"), is_var("w")]).match(expr) ## Graph-wise Matching def test_simple_used_by(): with PatternContext() as ctx: n0 = is_var("x") # x is a free var (fn arg) n1 = wildcard() n0 ^ n1 dfb = main_fn.body.blocks[0] matched = ctx.match_dfb(dfb) assert matched assert matched[n0] == main_fn.params[0] assert matched[n1] == dfb.bindings[0].var def test_simple_call_tir_edge(): with PatternContext() as ctx: n0 = is_call_tir("tir_matmul") n1 = is_call_tir("tir_relu") n0.used_by(n1) dfb = main_fn.body.blocks[0] matched = ctx.match_dfb(dfb) assert matched assert matched[n0] == dfb.bindings[0].var assert matched[n1] == dfb.bindings[1].var def test_simple_oub(): with PatternContext() as ctx: n0 = is_call_tir("tir_matmul") n1 = is_call_tir("tir_relu") n0 >> n1 dfb = main_fn.body.blocks[0] matched = ctx.match_dfb(dfb) assert matched assert matched[n0] == dfb.bindings[0].var assert matched[n1] == dfb.bindings[1].var def test_counter_syntax_match(): with PatternContext() as ctx: n0 = is_call_dps_packed("extern_matmul") n1 = is_call_dps_packed("extern_impossible") n0 >> n1 dfb = main_fn.body.blocks[0] assert not ctx.match_dfb(dfb) with PatternContext() as ctx: n0 = is_call_dps_packed("extern_matmul") n1 = is_call_dps_packed("extern_impossible") n0 ^ n1 dfb = main_fn.body.blocks[0] assert not ctx.match_dfb(dfb) @tvm.script.ir_module class Diamond: @R.function def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor: with R.dataflow(): # matmul # / \ # relu sigmoid # \ / # add lv0 = R.call_dps_packed("extern_matmul", (x, w), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("extern_relu", (lv0,), R.Tensor((32, 32), dtype="float32")) lv2 = R.call_dps_packed("extern_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32")) lv3 = R.call_dps_packed("extern_add", (lv1, lv2), R.Tensor((32, 32), dtype="float32")) R.output(lv3) return lv3 def test_diamond(): with PatternContext() as ctx: n0 = is_call_dps_packed("extern_matmul") n1 = is_call_dps_packed("extern_relu") n2 = is_call_dps_packed("extern_sigmoid") n3 = is_call_dps_packed("extern_add") n0 ^ n1 n0 ^ n2 n1 >> n3 n2 >> n3 dfb = Diamond["main"].body.blocks[0] assert ctx.match_dfb(dfb) # simplify it with fork_to with PatternContext() as ctx: n1 = is_call_dps_packed("extern_relu") n2 = is_call_dps_packed("extern_sigmoid") n3 = is_call_dps_packed("extern_add") is_call_dps_packed("extern_matmul").fork_to(n1, n2) n1 >> n3 n2 >> n3 dfb = Diamond["main"].body.blocks[0] assert ctx.match_dfb(dfb) def test_diamond_counter_oub(): with PatternContext() as ctx: n0 = is_call_dps_packed("extern_matmul") n1 = is_call_dps_packed("extern_relu") n2 = is_call_dps_packed("extern_sigmoid") n3 = is_call_dps_packed("extern_add") n0 >> n1 n0 >> n2 n1 >> n3 n2 >> n3 dfb = Diamond["main"].body.blocks[0] assert not ctx.match_dfb(dfb) @tvm.script.ir_module class SmallDiamond: @R.function def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor: with R.dataflow(): # relu # / \ # \ / # add lv0 = R.call_dps_packed("my_relu", (x,), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("my_add", (lv0, lv0), R.Tensor((32, 32), dtype="float32")) R.output(lv1) return lv1 @tvm.script.ir_module class SmallParallel: @R.function def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor: with R.dataflow(): # relu relu # \ / # add lv0 = R.call_dps_packed("my_relu", (x,), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("my_relu", (x,), R.Tensor((32, 32), dtype="float32")) lv2 = R.call_dps_packed("my_add", (lv0, lv1), R.Tensor((32, 32), dtype="float32")) R.output(lv2) return lv2 def test_distinguish_diamond_and_parallel(): # pattern lang cannot distinguish the two cases above. diamond = SmallDiamond["main"].body.blocks[0] parallel = SmallParallel["main"].body.blocks[0] with PatternContext() as ctx: # describe a diamond pattern fork = is_call_dps_packed("my_relu") join = is_call_dps_packed("my_add") fork.only_used_by(join, index=0) fork.only_used_by(join, index=1) assert ctx.match_dfb(diamond) assert not ctx.match_dfb(parallel) with PatternContext() as ctx: # describe a parallel pattern join = is_call_dps_packed("my_add") # Due to one-one matching: # is_call_dps_packed("my_relu") creates the 1st relu is_call_dps_packed("my_relu") >> join # is_call_dps_packed("my_relu") # creates the another different relu (obj address is different) is_call_dps_packed("my_relu") >> join assert ctx.match_dfb(parallel) assert not ctx.match_dfb(diamond) @tvm.script.ir_module class CBRx2: @R.function def main( x: R.Tensor((32, 32), "float32"), w0: R.Tensor((1, 1), "float32"), bias0: R.Tensor((32, 32), "float32"), w1: R.Tensor((1, 1), "float32"), bias1: R.Tensor((32, 32), "float32"), ) -> R.Tensor: # R.TensorRT's CBR Optimization Pattern # input # / \ # cbr0 cbr1 # \ / # concat with R.dataflow(): lv0 = R.call_dps_packed("conv1x1", (x, w0), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("bias_add", (lv0, bias0), R.Tensor((32, 32), dtype="float32")) lv2 = R.call_dps_packed("my_relu", (lv1), R.Tensor((32, 32), dtype="float32")) lv3 = R.call_dps_packed("conv1x1", (x, w1), R.Tensor((32, 32), dtype="float32")) lv4 = R.call_dps_packed("bias_add", (lv3, bias1), R.Tensor((32, 32), dtype="float32")) lv5 = R.call_dps_packed("my_relu", (lv4), R.Tensor((32, 32), dtype="float32")) lv6 = R.call_dps_packed("concat", (lv2, lv5), R.Tensor((32, 64), dtype="float32")) R.output(lv6) return lv6 def test_nested_context(): dfb = CBRx2["main"].body.blocks[0] with PatternContext() as ctx0: ( is_call_dps_packed("conv1x1") >> is_call_dps_packed("bias_add") >> is_call_dps_packed("my_relu") ) with PatternContext() as ctx1: is_call_dps_packed("conv1x1") >> is_call_dps_packed("my_relu") # pattern to miss with PatternContext() as ctx2: is_call_dps_packed("bias_add") >> is_call_dps_packed("my_relu") assert ctx2.match_dfb(dfb) assert PatternContext.current() == ctx2 assert not ctx1.match_dfb(dfb) assert PatternContext.current() == ctx1 assert ctx0.match_dfb(dfb) assert PatternContext.current() == ctx0 def test_two_cbr(): with PatternContext() as ctx: cbr0 = ( is_call_dps_packed("conv1x1") >> is_call_dps_packed("bias_add") >> is_call_dps_packed("my_relu") ) cbr1 = cbr0.dup() assert cbr0.patterns[0] != cbr1.patterns[0] assert cbr0.patterns[1] != cbr1.patterns[1] assert cbr0.patterns[2] != cbr1.patterns[2] is_var("x").fork_to(cbr0, cbr1) dfb = CBRx2["main"].body.blocks[0] assert ctx.match_dfb(dfb) with PatternContext() as ctx: # Deny the pattern cbr0 = ( is_call_dps_packed("conv1x1") >> is_call_dps_packed("bias_add") >> is_call_dps_packed("my_relu") ) cbr1 = cbr0.dup() # input has no fork at y. is_var("y").fork_to(cbr0, cbr1) dfb = CBRx2["main"].body.blocks[0] assert not ctx.match_dfb(dfb) def test_two_matmul(): # Same as Figure 2(a) in TASO paper. @tvm.script.ir_module class MatMul2: @R.function def main( a: R.Tensor((32, 16), "float32"), b: R.Tensor((16, 48), "float32"), c: R.Tensor((48, 32), "float32"), ) -> R.Tensor: with R.dataflow(): lv0 = R.call_dps_packed("matmul", (a, b), R.Tensor((32, 48), dtype="float32")) lv1 = R.call_dps_packed("matmul", (lv0, c), R.Tensor((32, 32), dtype="float32")) R.output(lv1) return lv1 with PatternContext() as ctx: is_call_dps_packed("matmul") >> is_call_dps_packed("matmul") dfb = MatMul2["main"].body.blocks[0] assert ctx.match_dfb(dfb) with PatternContext() as ctx: is_call_dps_packed("matmul").has_shape([32, 48]) >> is_call_dps_packed("matmul").has_shape( [32, 32] ) dfb = MatMul2["main"].body.blocks[0] assert ctx.match_dfb(dfb) with PatternContext() as ctx: is_call_dps_packed("matmul") >> is_call_dps_packed("matmul") >> is_call_dps_packed("matmul") dfb = MatMul2["main"].body.blocks[0] # Three MatMul cannot match assert not ctx.match_dfb(dfb) def test_concat_mm_split(): # Same as Figure 2(b) in TASO paper. @tvm.script.ir_module class CMS: @R.function def main( a: R.Tensor((32, 32), "float32"), b: R.Tensor((16, 32), "float32"), c: R.Tensor((16, 32), "float32"), ) -> R.Tensor: with R.dataflow(): lv0 = R.call_dps_packed("my_concat", (b, c), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("my_matmul", (a, lv0), R.Tensor((32, 32), dtype="float32")) lv2 = R.call_dps_packed( "my_split", (lv1,), [R.Tensor((16, 32), dtype="float32"), R.Tensor((16, 32), dtype="float32")], ) lv3 = R.TupleGetItem(lv2, 0) lv4 = R.TupleGetItem(lv2, 1) lv5 = R.add(lv3, lv4) R.output(lv5) return lv5 with PatternContext() as ctx: ( is_call_dps_packed("my_concat") >> is_call_dps_packed("my_matmul") >> is_call_dps_packed("my_split") ) dfb = CMS["main"].body.blocks[0] assert ctx.match_dfb(dfb) with PatternContext() as ctx: split = is_call_dps_packed("my_split") lv3 = TupleGetItemPattern(split, 0).has_shape([16, 32]) lv4 = TupleGetItemPattern(split, 1).has_shape([16, 32]) split.fork_to(lv3, lv4) add = is_op("relax.add")(lv3, lv4) # TODO(@ganler): simplify this through implicit graph pattern. lv3 >> add lv4 >> add dfb = CMS["main"].body.blocks[0] assert ctx.match_dfb(dfb) def test_self_attention(): # The example comes from. # https://developer.nvidia.com/blog/nlu-with-tensorrt-bert/ @tvm.script.ir_module class SelfAttention: @R.function def main( x: R.Tensor(("b", "s", "n", "h"), "float32"), wq: R.Tensor(("h", "h"), "float32"), wk: R.Tensor(("h", "h"), "float32"), wv: R.Tensor(("h", "h"), "float32"), ) -> R.Tensor: b, s, n, h = T.int64(), T.int64(), T.int64(), T.int64() with R.dataflow(): fcq = R.call_dps_packed("my_fc", (x, wq), R.Tensor((b, s, n, h), dtype="float32")) tpq = R.call_dps_packed( "my_transpose", (fcq,), R.Tensor((b, s, h, n), dtype="float32") ) fck = R.call_dps_packed("my_fc", (x, wk), R.Tensor((b, s, n, h), dtype="float32")) tpk = R.call_dps_packed( "my_transpose", (fck,), R.Tensor((b, s, h, n), dtype="float32") ) mul = R.multiply(tpq, tpk) scale = R.multiply(mul, R.const(1.1, "float32")) softmax = R.call_dps_packed( "softmax", (scale,), R.Tensor((b, s, n, h), dtype="float32") ) fcv = R.call_dps_packed("my_fc", (x, wv), R.Tensor((b, s, n, h), dtype="float32")) tpv = R.call_dps_packed( "my_transpose", (fcv,), R.Tensor((b, s, h, n), dtype="float32") ) out = R.multiply(softmax, tpv) R.output(out) return out with PatternContext() as ctx: fc_trans_q = is_call_dps_packed("my_fc") >> is_call_dps_packed("my_transpose") fc_trans_k = fc_trans_q.dup() fc_trans_v = fc_trans_q.dup() is_var("x").fork_to(fc_trans_q, fc_trans_k, fc_trans_v) dfb = SelfAttention["main"].body.blocks[0] assert ctx.match_dfb(dfb) def test_nested_diamond(): @tvm.script.ir_module class DiamondInDiamond: @R.function def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor: with R.dataflow(): # matmul0 matmul1 # / \ / \ # sigmoid2 add4 sigmoid3 # \ / \ / # add5 add6 # \ / # add7 lv0 = R.call_dps_packed( "extern_matmul", (x, w), R.Tensor((32, 32), dtype="float32") ) lv1 = R.call_dps_packed( "extern_matmul", (x, w), R.Tensor((32, 32), dtype="float32") ) lv2 = R.call_dps_packed( "extern_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32") ) lv3 = R.call_dps_packed( "extern_sigmoid", (lv1), R.Tensor((32, 32), dtype="float32") ) lv4 = R.call_dps_packed( "extern_add", (lv0, lv1), R.Tensor((32, 32), dtype="float32") ) lv5 = R.call_dps_packed( "extern_add", (lv2, lv4), R.Tensor((32, 32), dtype="float32") ) lv6 = R.call_dps_packed( "extern_add", (lv3, lv4), R.Tensor((32, 32), dtype="float32") ) lv7 = R.call_dps_packed( "extern_add", (lv5, lv6), R.Tensor((32, 32), dtype="float32") ) R.output(lv7) return lv7 # match matmul0 diamond with PatternContext() as ctx: sigmoid2 = is_call_dps_packed("extern_sigmoid") add4 = is_call_dps_packed("extern_add") is_call_dps_packed("extern_matmul").fork_to(sigmoid2, add4) add5 = is_call_dps_packed("extern_add") sigmoid2 >> add5 add4 ^ add5 assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0]) # counter case: mis-match matmul0 diamond with PatternContext() as ctx: sigmoid2 = is_call_dps_packed("extern_sigmoid") add4 = is_call_dps_packed("extern_add") is_call_dps_packed("extern_matmul").fork_to(sigmoid2, add4) add5 = is_call_dps_packed("extern_add") sigmoid2 >> add5 add4 >> add5 # not only-used-by relation assert not ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0]) # match matmul1 diamond with PatternContext() as ctx: sigmoid3 = is_call_dps_packed("extern_sigmoid") add4 = is_call_dps_packed("extern_add") is_call_dps_packed("extern_matmul").fork_to(sigmoid3, add4) add6 = is_call_dps_packed("extern_add") sigmoid3 >> add6 add4 ^ add6 assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0]) # match add-4-5-6-7 with PatternContext() as ctx: add5, add6, add7 = ( is_call_dps_packed("extern_add"), is_call_dps_packed("extern_add"), is_call_dps_packed("extern_add"), ) is_call_dps_packed("extern_add").fork_to(add5, add6) # add4 add5 >> add7 add6 >> add7 assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0]) def test_incremental_solving(): @R.function def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor: with R.dataflow(): # relu -> sigmoid -> neg lv0 = R.call_dps_packed("extern_relu", (x), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("extern_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32")) lv2 = R.call_dps_packed("extern_neg", (lv1), R.Tensor((32, 32), dtype="float32")) R.output(lv2) return lv2 relu = is_call_dps_packed("extern_relu") sigmoid = is_call_dps_packed("extern_sigmoid") neg = is_call_dps_packed("extern_neg") with PatternContext() as ctx0: relu >> sigmoid with PatternContext(incremental=True) as ctx1: # because we are doing incremental solving # relu >> sigmoid is still a constraint in this context. # that said the total constraint is: # relu >> sigmoid >> neg sigmoid >> neg assert ctx1.match_dfb(simple_chain.body.blocks[0]) # match relue -> sigmoid assert ctx0.match_dfb(simple_chain.body.blocks[0]) def test_incremental_solving_counter(): @R.function def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor: with R.dataflow(): # sigmoid -> neg lv0 = R.call_dps_packed("extern_sigmoid", (x), R.Tensor((32, 32), dtype="float32")) lv1 = R.call_dps_packed("extern_neg", (lv0), R.Tensor((32, 32), dtype="float32")) R.output(lv1) return lv1 relu = is_call_dps_packed("extern_relu") sigmoid = is_call_dps_packed("extern_sigmoid") neg = is_call_dps_packed("extern_neg") with PatternContext() as ctx0: relu >> sigmoid # cannot match with PatternContext(incremental=False) as ctx1: # total constraint: sigmoid >> neg sigmoid >> neg assert ctx1.match_dfb(simple_chain.body.blocks[0]) with PatternContext(incremental=True) as ctx1: # total constraint: relu >> sigmoid >> neg sigmoid >> neg assert not ctx1.match_dfb(simple_chain.body.blocks[0]) def test_rewrite_simple(): @R.function def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"): with R.dataflow(): x2 = R.add(x, x) x4 = R.add(x2, x2) R.output(x4) return x4 @R.function def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"): with R.dataflow(): lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32")) x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32")) R.output(x4) return x4 @R.function def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"): with R.dataflow(): x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32")) R.output(x4) return x4 x = wildcard() pattern = is_op("relax.add")(x, x) def rewriter(_, matchings): return R.multiply(matchings[x], R.const(2, "float32")) rewritten = rewrite_call(pattern, rewriter, main) tvm.ir.assert_structural_equal(rewritten, expected1.with_attr("global_symbol", "main")) add1 = is_op("relax.add")(x, x) pattern = is_op("relax.add")(add1, add1) def rewriter(_, matchings): return R.multiply(matchings[x], R.const(4, "float32")) rewritten = rewrite_call(pattern, rewriter, main) tvm.ir.assert_structural_equal(rewritten, expected2.with_attr("global_symbol", "main")) # No rewriting, return the original call node as is def rewriter(orig, _): return orig rewritten = rewrite_call(pattern, rewriter, main) tvm.ir.assert_structural_equal(rewritten, main) def test_rewrite_attention(): @R.function def main( Q: R.Tensor((2, 4096, 8, 40), "float32"), K: R.Tensor((2, 4096, 8, 40), "float32"), V: R.Tensor((2, 4096, 8, 40), "float32"), ) -> R.Tensor((2, 4096, 8, 40), "float32"): with R.dataflow(): lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3]) lv59 = R.reshape(lv58, R.shape([16, 4096, 40])) lv61 = R.permute_dims(K, axes=[0, 2, 1, 3]) lv62 = R.reshape(lv61, R.shape([16, 4096, 40])) lv64 = R.permute_dims(V, axes=[0, 2, 1, 3]) lv65 = R.reshape(lv64, R.shape([16, 4096, 40])) lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1]) lv3_1 = R.matmul(lv59, lv62_transposed) lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32")) lv69 = R.nn.softmax(lv68, axis=-1) lv_3 = R.matmul(lv69, lv65) lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40])) lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3]) R.output(lv72) return lv72 @R.function def expected( Q: R.Tensor((2, 4096, 8, 40), dtype="float32"), K: R.Tensor((2, 4096, 8, 40), dtype="float32"), V: R.Tensor((2, 4096, 8, 40), dtype="float32"), ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"): with R.dataflow(): lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K) R.output(lv72) return lv72 def BSNH_to_BSH(tensor): return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard()) def BSH_to_BSNH(tensor): return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard())) Q = wildcard() K = wildcard() V = wildcard() Q_3D = BSNH_to_BSH(Q) V_3D = BSNH_to_BSH(V) K_3D = BSNH_to_BSH(K) matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D)) multiply = is_op("relax.multiply")(matmul1, is_const()) softmax = is_op("relax.nn.softmax")(multiply) matmul2 = is_op("relax.matmul")(softmax, K_3D) pattern = BSH_to_BSNH(matmul2) def rewriter(_, matchings): return R.nn.attention(matchings[Q], matchings[K], matchings[V]) rewritten = rewrite_call(pattern, rewriter, main) tvm.ir.assert_structural_equal(rewritten, expected.with_attr("global_symbol", "main")) def test_attention_qkv(): @tvm.script.ir_module class QKV_proj: @R.function def main( x: R.Tensor((2, 1024, 640), "float32"), w0: R.Tensor((640, 640), "float32"), w1: R.Tensor((640, 640), "float32"), w2: R.Tensor((640, 640), "float32"), ) -> R.Tensor: with R.dataflow(): lv0 = R.matmul(x, w0) lv1 = R.matmul(x, w1) lv2 = R.matmul(x, w2) out = (lv0, lv1, lv2) R.output(out) return out with PatternContext() as ctx: inp_pat = wildcard() Q_weight_pat = wildcard() K_weight_pat = wildcard() V_weight_pat = wildcard() matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat) matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat) matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat) dfb = QKV_proj["main"].body.blocks[0] out = ctx.match_dfb(dfb) assert out[Q_weight_pat].name_hint == "w0" assert out[K_weight_pat].name_hint == "w1" assert out[V_weight_pat].name_hint == "w2" def test_attention_fake_qkv(): @tvm.script.ir_module class QKV_proj: @R.function def main( x1: R.Tensor((2, 1024, 640), "float32"), x2: R.Tensor((2, 1024, 640), "float32"), w0: R.Tensor((640, 640), "float32"), w1: R.Tensor((640, 640), "float32"), w2: R.Tensor((640, 640), "float32"), ) -> R.Tensor: with R.dataflow(): lv0 = R.matmul(x1, w0) lv1 = R.matmul(x2, w1) lv2 = R.matmul(x2, w2) out = (lv0, lv1, lv2) R.output(out) return out with PatternContext() as ctx: inp_pat = wildcard() Q_weight_pat = wildcard() K_weight_pat = wildcard() V_weight_pat = wildcard() matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat) matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat) matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat) dfb = QKV_proj["main"].body.blocks[0] assert ctx.match_dfb(dfb) is None def get_qkv_proj_rewriter(): with PatternContext() as ctx: inp_pat = wildcard() Q_weight_pat = wildcard() K_weight_pat = wildcard() V_weight_pat = wildcard() matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat) matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat) matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat) def qkv_proj_rewriter(matchings, _): inp = matchings[inp_pat] Q_weight = matchings[Q_weight_pat] K_weight = matchings[K_weight_pat] V_weight = matchings[V_weight_pat] width = Q_weight.ty.shape[1] concat = R.concat([Q_weight, K_weight, V_weight], axis=1) matmul = R.matmul(inp, concat) Q = R.strided_slice(matmul, axes=[2], begin=[0], end=[width]) K = R.strided_slice(matmul, axes=[2], begin=[width], end=[width * 2]) V = R.strided_slice(matmul, axes=[2], begin=[width * 2], end=[width * 3]) return {matchings[matmul1]: Q, matchings[matmul2]: K, matchings[matmul3]: V} return ctx, qkv_proj_rewriter def test_combine_matmul_twice(): @R.function(private=True) def qkv_x2( x1: R.Tensor((2, 1024, 640), "float32"), x2: R.Tensor((2, 1024, 640), "float32"), w0: R.Tensor((640, 640), "float32"), w1: R.Tensor((640, 640), "float32"), w2: R.Tensor((640, 640), "float32"), w3: R.Tensor((640, 640), "float32"), w4: R.Tensor((640, 640), "float32"), w5: R.Tensor((640, 640), "float32"), ): with R.dataflow(): lv0 = R.matmul(x1, w0) lv1 = R.matmul(x1, w1) lv2 = R.matmul(x1, w2) lv3 = R.matmul(x2, w3) lv4 = R.matmul(x2, w4) lv5 = R.matmul(x2, w5) out = (lv0, lv1, lv2, lv3, lv4, lv5) R.output(out) return out @R.function(private=True) def expected( x1: R.Tensor((2, 1024, 640), "float32"), x2: R.Tensor((2, 1024, 640), "float32"), w0: R.Tensor((640, 640), "float32"), w1: R.Tensor((640, 640), "float32"), w2: R.Tensor((640, 640), "float32"), w3: R.Tensor((640, 640), "float32"), w4: R.Tensor((640, 640), "float32"), w5: R.Tensor((640, 640), "float32"), ): with R.dataflow(): lv = R.concat((w0, w1, w2), axis=1) lv1 = R.matmul(x1, lv) lv0 = R.strided_slice(lv1, axes=[2], begin=[0], end=[640]) lv1_1 = R.strided_slice(lv1, axes=[2], begin=[640], end=[1280]) lv2 = R.strided_slice(lv1, axes=[2], begin=[1280], end=[1920]) lv2_1 = R.concat((w3, w4, w5), axis=1) lv3 = R.matmul(x2, lv2_1) lv3_1 = R.strided_slice(lv3, axes=[2], begin=[0], end=[640]) lv4 = R.strided_slice(lv3, axes=[2], begin=[640], end=[1280]) lv5 = R.strided_slice(lv3, axes=[2], begin=[1280], end=[1920]) out = lv0, lv1_1, lv2, lv3_1, lv4, lv5 R.output(out) return out ctx, rewriter = get_qkv_proj_rewriter() rewritten = rewrite_bindings(ctx, rewriter, qkv_x2) tvm.ir.assert_structural_equal(rewritten, expected) def test_dataflow_may_start_with_match_cast(): """Inputs to rewrite_bindings may contain R.match_cast This is a regression test. In previous implementations, applying `rewrite_bindings` when `R.match_cast` is the first binding of a `R.dataflow` block would cause a segfault. """ @R.function(private=True) def before( x_untyped: R.Tensor, w0_untyped: R.Tensor, w1_untyped: R.Tensor, w2_untyped: R.Tensor, ): with R.dataflow(): x = R.match_cast(x_untyped, R.Tensor((2, 1024, 640), "float32")) w0 = R.match_cast(w0_untyped, R.Tensor((640, 640), "float32")) w1 = R.match_cast(w1_untyped, R.Tensor((640, 640), "float32")) w2 = R.match_cast(w2_untyped, R.Tensor((640, 640), "float32")) out_0 = R.matmul(x, w0) out_1 = R.matmul(x, w1) out_2 = R.matmul(x, w2) out = (out_0, out_1, out_2) R.output(out) return out @R.function(private=True) def expected( x_untyped: R.Tensor, w0_untyped: R.Tensor, w1_untyped: R.Tensor, w2_untyped: R.Tensor, ): with R.dataflow(): x = R.match_cast(x_untyped, R.Tensor((2, 1024, 640), "float32")) w0 = R.match_cast(w0_untyped, R.Tensor((640, 640), "float32")) w1 = R.match_cast(w1_untyped, R.Tensor((640, 640), "float32")) w2 = R.match_cast(w2_untyped, R.Tensor((640, 640), "float32")) w_concat = R.concat((w0, w1, w2), axis=1) out_concat = R.matmul(x, w_concat) out_0 = R.strided_slice(out_concat, axes=[2], begin=[0], end=[640]) out_1 = R.strided_slice(out_concat, axes=[2], begin=[640], end=[1280]) out_2 = R.strided_slice(out_concat, axes=[2], begin=[1280], end=[1920]) out = (out_0, out_1, out_2) R.output(out) return out ctx, rewriter = get_qkv_proj_rewriter() rewritten = rewrite_bindings(ctx, rewriter, before) tvm.ir.assert_structural_equal(rewritten, expected) def test_combine_matmul_emit_order(): @R.function(private=True) def main( x1: R.Tensor((2, 1024, 640), "float32"), w0: R.Tensor((640, 640), "float32"), w1: R.Tensor((640, 640), "float32"), w2: R.Tensor((640, 640), "float32"), ): with R.dataflow(): w0_t = R.permute_dims(w0, axes=None) lv0 = R.matmul(x1, w0_t) w1_t = R.permute_dims(w1, axes=None) w1_t_t = R.permute_dims(w1_t, axes=None) lv1 = R.matmul(x1, w1_t_t) w2_t = R.permute_dims(w2, axes=None) lv2 = R.matmul(x1, w2_t) out = (lv0, lv1, lv2) R.output(out) return out @R.function(private=True) def expected( x1: R.Tensor((2, 1024, 640), dtype="float32"), w0: R.Tensor((640, 640), dtype="float32"), w1: R.Tensor((640, 640), dtype="float32"), w2: R.Tensor((640, 640), dtype="float32"), ): with R.dataflow(): w0_t = R.permute_dims(w0, axes=None) w1_t = R.permute_dims(w1, axes=None) w1_t_t = R.permute_dims(w1_t, axes=None) w2_t = R.permute_dims(w2, axes=None) lv = R.concat((w0_t, w1_t_t, w2_t), axis=1) lv1 = R.matmul(x1, lv) lv0 = R.strided_slice(lv1, axes=[2], begin=[0], end=[640]) lv1_1 = R.strided_slice(lv1, axes=[2], begin=[640], end=[1280]) lv2 = R.strided_slice(lv1, axes=[2], begin=[1280], end=[1920]) out = lv0, lv1_1, lv2 R.output(out) return out ctx, rewriter = get_qkv_proj_rewriter() rewritten = rewrite_bindings(ctx, rewriter, main) tvm.ir.assert_structural_equal(rewritten, expected) # make sure it builds mod = tvm.IRModule() mod["main"] = rewritten tvm.compile(mod, target="llvm") def test_combine_transposed_matmul_twice(): @R.function(private=True) def main( x1: R.Tensor((2, 1024, 640), "float32"), x2: R.Tensor((2, 1024, 640), "float32"), w0: R.Tensor((640, 640), "float32"), w1: R.Tensor((640, 640), "float32"), w2: R.Tensor((640, 640), "float32"), w3: R.Tensor((640, 640), "float32"), ): with R.dataflow(): w0_t = R.permute_dims(w0, axes=None) lv0 = R.matmul(x1, w0_t) w1_t = R.permute_dims(w1, axes=None) lv1 = R.matmul(x1, w1_t) w2_t = R.permute_dims(w2, axes=None) lv2 = R.matmul(x2, w2_t) w3_t = R.permute_dims(w3, axes=None) lv3 = R.matmul(x2, w3_t) out = (lv0, lv1, lv2, lv3) R.output(out) return out @R.function(private=True) def expected( x1: R.Tensor((2, 1024, 640), dtype="float32"), x2: R.Tensor((2, 1024, 640), dtype="float32"), w0: R.Tensor((640, 640), dtype="float32"), w1: R.Tensor((640, 640), dtype="float32"), w2: R.Tensor((640, 640), dtype="float32"), w3: R.Tensor((640, 640), dtype="float32"), ): with R.dataflow(): lv: R.Tensor((1280, 640), dtype="float32") = R.concat((w0, w1), axis=0) lv1: R.Tensor((640, 1280), dtype="float32") = R.permute_dims(lv, axes=None) lv2: R.Tensor((2, 1024, 1280), dtype="float32") = R.matmul(x1, lv1) lv3: R.Tuple( R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32"), ) = R.split(lv2, indices_or_sections=[640], axis=-1) lv0: R.Tensor((2, 1024, 640), dtype="float32") = lv3[0] lv1_1: R.Tensor((2, 1024, 640), dtype="float32") = lv3[1] lv_1: R.Tensor((1280, 640), dtype="float32") = R.concat((w2, w3), axis=0) lv1_2: R.Tensor((640, 1280), dtype="float32") = R.permute_dims(lv_1, axes=None) lv2_1: R.Tensor((2, 1024, 1280), dtype="float32") = R.matmul(x2, lv1_2) lv3_1: R.Tuple( R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32"), ) = R.split(lv2_1, indices_or_sections=[640], axis=-1) lv2_1_1: R.Tensor((2, 1024, 640), dtype="float32") = lv3_1[0] lv3_1_1: R.Tensor((2, 1024, 640), dtype="float32") = lv3_1[1] out: R.Tuple( R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32"), R.Tensor((2, 1024, 640), dtype="float32"), ) = (lv0, lv1_1, lv2_1_1, lv3_1_1) R.output(out) return out with PatternContext() as ctx: inp_pat = wildcard() w1_pat = wildcard() w2_pat = wildcard() matmul1 = is_op("relax.matmul")(inp_pat, is_op("relax.permute_dims")(w1_pat)) matmul2 = is_op("relax.matmul")(inp_pat, is_op("relax.permute_dims")(w2_pat)) def rewriter(matchings, _): inp = matchings[inp_pat] w1 = matchings[w1_pat] w2 = matchings[w2_pat] concat = R.concat([w1, w2], axis=0) matmul = R.matmul(inp, R.permute_dims(concat)) sections = [w1.ty.shape[0]] chunks = R.split(matmul, sections, -1) return { matchings[matmul1]: chunks[0], matchings[matmul2]: chunks[1], } rewritten = rewrite_bindings(ctx, rewriter, main) tvm.ir.assert_structural_equal(rewritten, expected) # make sure it builds mod = tvm.IRModule() mod["main"] = rewritten print(mod) tvm.compile(mod, target="llvm") def test_commutative_pattern_match(): @R.function(private=True) def before( x: R.Tensor((1024,)), ): with R.dataflow(): y = R.add(x, x) out = R.add(R.const(1.0), y) R.output(out) return out @R.function(private=True) def expected( x: R.Tensor((1024,)), ): with R.dataflow(): y = R.add(x, x) out = R.add(y, R.const(2.0)) R.output(out) return out pattern_add = is_op("relax.add") pattern_mul = is_op("relax.multiply") pattern_op = pattern_add | pattern_mul pattern_arg = wildcard() pattern_const = is_const() pattern = pattern_op(pattern_arg, pattern_const) def rewriter(expr, matches): op = matches[pattern_op] arg = matches[pattern_arg] const = matches[pattern_const].data.numpy() if const.shape == tuple() and const[()] == 1.0: return rx.Call(op, [arg, rx.const(2.0)]) else: return expr after = rewrite_call(pattern, rewriter, before) tvm.ir.assert_structural_equal(after, expected) def test_repeated_pattern_match(): """rewrite_call should iterate until convergence""" @R.function(private=True) def before( x: R.Tensor((1024,)), y: R.Tensor((1024,)), z: R.Tensor((1024,)), ): with R.dataflow(): a = R.add(x, y) b = R.add(a, z) out = R.multiply(b, R.const(5.0)) R.output(out) return out @R.function(private=True) def expected( x: R.Tensor((1024,)), y: R.Tensor((1024,)), z: R.Tensor((1024,)), ): with R.dataflow(): x = R.multiply(x, R.const(5.0)) y = R.multiply(y, R.const(5.0)) a = R.add(x, y) z = R.multiply(z, R.const(5.0)) b = R.add(a, z) R.output(b) return b pattern_add_lhs = wildcard() pattern_add_rhs = wildcard() pattern_add = is_op("relax.add")(pattern_add_lhs, pattern_add_rhs) mul_const = is_const() pattern_mul = is_op("relax.multiply")(pattern_add, mul_const) pattern = pattern_mul def rewriter(_expr, matches): const = matches[mul_const] return (matches[pattern_add_lhs] * const) + (matches[pattern_add_rhs] * const) after = rewrite_call(pattern, rewriter, before) tvm.ir.assert_structural_equal(after, expected) bind_to_dataflow_var = tvm.testing.parameter( by_dict={"var-to-var": False, "var-to-dataflow-var": True} ) def test_rewrite_without_trivial_binding(bind_to_dataflow_var): """rewrite_call should avoid producing trivial "y = x" bindings This may not be possible in all cases, and follows the same rules as CanonicalizeBindings. For example, a `relax.Var` is bound to a `relax.DataflowVar` may not be removed, to ensure that the `relax.DataflowVar` is only used within a `DataflowBlock`. """ if bind_to_dataflow_var: @R.function(private=True) def before(x: R.Tensor((1024,))): with R.dataflow(): a = R.add(x, x) b = R.reshape(a, (1024,)) R.output(b) return b @R.function(private=True) def expected(x: R.Tensor((1024,))): with R.dataflow(): b = R.add(x, x) R.output(b) return b else: @R.function(private=True) def before(x: R.Tensor((1024,))): a = R.add(x, x) b = R.reshape(a, (1024,)) return b @R.function(private=True) def expected(x: R.Tensor((1024,))): a = R.add(x, x) return a pattern_arg = wildcard() pattern_shape_expr = wildcard() pattern = is_op("relax.reshape")(pattern_arg, pattern_shape_expr) def rewriter(expr, matches): arg = matches[pattern_arg] shape_expr = matches[pattern_shape_expr] if tvm_ffi.structural_equal(arg.ty.shape, shape_expr): return arg else: return expr after = rewrite_call(pattern, rewriter, before) tvm.ir.assert_structural_equal(after, expected) same_shape_func_type = tvm.testing.parameter( "same_static_shape", "same_dynamic_shape", "different_static_shape", "different_dynamic_shape", ) def test_same_shape_pattern(same_shape_func_type): if same_shape_func_type == "same_static_shape": @R.function(private=True) def func( a: R.Tensor((1024, 128), "float32"), b: R.Tensor((1024, 128), "float32"), ) -> R.Tensor: with R.dataflow(): c = R.multiply(a, R.const(2.0)) d = R.add(b, c) out = d R.output(out) return out elif same_shape_func_type == "same_dynamic_shape": @R.function(private=True) def func( a: R.Tensor(("n", 128), "float32"), b: R.Tensor(("n", 128), "float32"), ) -> R.Tensor: with R.dataflow(): c = R.multiply(a, R.const(2.0)) d = R.add(b, c) out = d R.output(out) return out elif same_shape_func_type == "different_static_shape": @R.function(private=True) def func( a: R.Tensor((1024, 128), "float32"), b: R.Tensor((1, 128), "float32"), ) -> R.Tensor: with R.dataflow(): c = R.multiply(a, R.const(2.0)) d = R.add(b, c) out = d R.output(out) return out elif same_shape_func_type == "different_dynamic_shape": @R.function(private=True) def func( a: R.Tensor(("n", 128), "float32"), b: R.Tensor(("m", 128), "float32"), ) -> R.Tensor: with R.dataflow(): c = R.multiply(a, R.const(2.0)) d = R.add(b, c) out = d R.output(out) return out else: raise ValueError(f"Unknown value of same_shape_func_type={same_shape_func_type}") with PatternContext() as ctx: pat_lhs = wildcard() pat_rhs = wildcard() pat_sum = is_op("relax.add")(pat_lhs, pat_rhs) pat_lhs.same_shape_as(pat_rhs) block = func.body.blocks[0] match = ctx.match_dfb(block) if "same" in same_shape_func_type: assert match else: assert match is None def test_iterative_rewrite_without_trivial_binding(): """Avoid introducing common sub-expressions Pattern replacement may produce the same intermediate, which should appear only once in the final result. """ @R.function(private=True) def before(x: R.Tensor((1024,))): with R.dataflow(): a = R.strided_slice(x, [0], [0], [512], [1]) b = R.strided_slice(x, [0], [512], [1024], [1]) c = R.add(a, b) R.output(c) return c @R.function(private=True) def expected(x: R.Tensor((1024,))): with R.dataflow(): x_split = R.split(x, 2) a = x_split[0] b = x_split[1] c = R.add(a, b) R.output(c) return c pattern_arg = wildcard() pattern_axes = wildcard() pattern_begin = wildcard() pattern_end = wildcard() pattern_strides = wildcard() pattern = is_op("relax.strided_slice")( pattern_arg, pattern_axes, pattern_begin, pattern_end, pattern_strides ) def rewriter(expr, matches): arg = matches[pattern_arg] axes = matches[pattern_axes] begin = matches[pattern_begin] end = matches[pattern_end] strides = matches[pattern_strides] strided_slice = matches[pattern] if arg.ty.shape is None: return expr if len(axes) != 1: return expr axis = axes[0] begin = begin[0] end = end[0] stride = strides[0] if not isinstance(axis, tirx.IntImm) or axis.value != 0: return expr if not isinstance(stride, tirx.IntImm) or stride.value != 1: return expr size = arg.ty.shape[0] if ( isinstance(size, tirx.IntImm) and isinstance(begin, tirx.IntImm) and isinstance(end, tirx.IntImm) ): size = size.value begin = begin.value end = end.value else: return expr gcd = functools.reduce(math.gcd, [begin, end, size]) if (end - begin) // gcd == 1: return rx.op.split(arg, size // gcd)[begin // gcd] return expr after = rewrite_call(pattern, rewriter, before) tvm.ir.assert_structural_equal(after, expected) def test_iterative_rewrite_with_removed_intermediates(): """Pattern replacement may require canonicalization A pattern may replace a tuple returned by a function with a tuple whose contents are known by Relax. In that case, canonicalization is required to unwrap the TupleGetItem instances into the known contents. This test case shows the intermediate results produced in the process of pattern-matching. """ @R.function(private=True) def before(a: R.Tensor((1024,)), b: R.Tensor((1024,))): with R.dataflow(): c = R.concat([a, b]) d = R.split(c, 2) e = d[0] f = d[1] g = R.add(a, e) h = R.add(f, g) R.output(h) return h # First pattern rewrite. The concat/rewrite can be unwrapped, so # `d` is rewritten from `R.split(c, 2)` into `(a, b)`. # # @R.function(private=True) # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))): # with R.dataflow(): # c = R.concat([a, b]) # d = (a,b) # e = d[0] # f = d[1] # g = R.add(a, e) # h = R.add(f, g) # R.output(h) # Canonicalization step. Because `d` is known to be `(a,b)`, # canonicalization can rewrite `d[0]` into `a` and `d[1]` into # `b`. # # @R.function(private=True) # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))): # with R.dataflow(): # c = R.concat([a, b]) # d = (a,b) # e = a # f = b # g = R.add(a, a) # h = R.add(b, g) # R.output(h) # Dead-code-elimination step. This technically isn't required # until the pattern matching has converged, but performing it now # prevents testing for matches on dead code. # # @R.function(private=True) # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))): # with R.dataflow(): # g = R.add(a, a) # h = R.add(b, g) # R.output(h) # Second pattern-matching step. Now, the `R.add(a,a)` can match # the other option in our pattern, and be rewritten as # `R.multiply(a,R.const(2))`. # # @R.function(private=True) # def intermediate(a: R.Tensor((1024,)), b: R.Tensor((1024,))): # with R.dataflow(): # g = R.multiply(a, R.const(2)) # h = R.add(b, g) # R.output(h) # Canonicalization and dead-code-elimination are applied again, # but have no effect this time. @R.function(private=True) def expected(a: R.Tensor((1024,)), b: R.Tensor((1024,))): with R.dataflow(): g = R.multiply(a, R.const(2)) h = R.add(b, g) R.output(h) return h pat_args = wildcard() op_concat = is_op("relax.concat") pat_concat = op_concat(pat_args).has_attr({"axis": 0}) op_split = is_op("relax.split") pat_split = op_split(pat_concat).has_attr({"axis": 0, "indices_or_sections": T.int64(2)}) pat_unwrap_concat_split = pat_split pat_arg = wildcard() op_add = is_op("relax.add") pat_add_self = op_add(pat_arg, pat_arg) pattern = pat_unwrap_concat_split | pat_add_self def rewriter(expr, matches): if pat_unwrap_concat_split in matches: args = matches[pat_args] if len(args) == 2 and tvm_ffi.structural_equal(args[0].ty, args[1].ty): return args elif pat_add_self in matches: arg = matches[pat_arg] return arg * rx.const(2) return expr after = rewrite_call(pattern, rewriter, before) tvm.ir.assert_structural_equal(expected, after) def test_wildcard_with_ty_updates_when_matching(): """A DFPattern may be restricted to a specific Type""" pat_lhs = wildcard().has_ty(R.Tensor([2, 3])) pat_rhs = wildcard().has_ty(R.Tensor([2, 3])) pat = is_op("relax.add")(pat_lhs, pat_rhs) def rewriter(expr, matches): lhs = matches[pat_lhs] rhs = matches[pat_rhs] return rx.op.multiply(lhs, rhs) @R.function(private=True) def before(): with R.dataflow(): A = R.zeros([2, 3], "int32") B = R.ones([2, 3], "int32") C = R.add(A, B) R.output(C) return C @R.function(private=True) def expected(): with R.dataflow(): A = R.zeros([2, 3], "int32") B = R.ones([2, 3], "int32") C = R.multiply(A, B) R.output(C) return C after = rewrite_call(pat, rewriter, before) tvm.ir.assert_structural_equal(expected, after) def test_wildcard_with_ty_is_no_op_when_not_matching(): """TypePattern requires the Type provided Here, the pattern would match, expect that the function has `R.Tensor([16,32])`, and the pattern requires `R.Tensor([2,3])`. """ pat_lhs = wildcard().has_ty(R.Tensor([2, 3])) pat_rhs = wildcard().has_ty(R.Tensor([2, 3])) pat = is_op("relax.add")(pat_lhs, pat_rhs) def rewriter(expr, matches): lhs = matches[pat_lhs] rhs = matches[pat_rhs] return rx.op.multiply(lhs, rhs) @R.function(private=True) def before(): with R.dataflow(): # This R.add has the same shape as the pattern, and will # be updated. A = R.zeros([16, 32], "int32") B = R.ones([16, 32], "int32") C = R.add(A, B) R.output(C) return C expected = before after = rewrite_call(pat, rewriter, before) tvm.ir.assert_structural_equal(expected, after) def test_wildcard_ty_for_unknown_dtype(): """TensorType with unknown dtype allows any dtype""" pat_lhs = wildcard().has_ty(R.Tensor([2, 3])) pat_rhs = wildcard().has_ty(R.Tensor([2, 3])) pat = is_op("relax.add")(pat_lhs, pat_rhs) def rewriter(expr, matches): lhs = matches[pat_lhs] rhs = matches[pat_rhs] return rx.op.multiply(lhs, rhs) @R.function(private=True) def before(): with R.dataflow(): A = R.zeros([2, 3], "int32") B = R.ones([2, 3], "int32") C = R.add(A, B) D = R.zeros([2, 3], "float32") E = R.ones([2, 3], "float32") F = R.add(D, E) output = (C, F) R.output(output) return output @R.function(private=True) def expected(): with R.dataflow(): A = R.zeros([2, 3], "int32") B = R.ones([2, 3], "int32") C = R.multiply(A, B) D = R.zeros([2, 3], "float32") E = R.ones([2, 3], "float32") F = R.multiply(D, E) output = (C, F) R.output(output) return output after = rewrite_call(pat, rewriter, before) tvm.ir.assert_structural_equal(expected, after) def test_wildcard_ty_with_symbolic_vars(): """TypePattern may define symbolic vars This test finds an elementwise `R.add`, while ignoring a broadcasted `R.add`. """ m = tirx.Var("m", "int64") n = tirx.Var("n", "int64") pat_lhs = wildcard().has_ty(R.Tensor([m, n])) pat_rhs = wildcard().has_ty(R.Tensor([m, n])) pat = is_op("relax.add")(pat_lhs, pat_rhs) def rewriter(expr, matches): lhs = matches[pat_lhs] rhs = matches[pat_rhs] return rx.op.multiply(lhs, rhs) @R.function(private=True) def before(): with R.dataflow(): A = R.zeros([64, 128], "int32") B = R.ones([64, 128], "int32") C = R.add(A, B) D = R.zeros([64, 128], "float32") E = R.ones([1, 128], "float32") F = R.add(D, E) output = (C, F) R.output(output) return output @R.function(private=True) def expected(): with R.dataflow(): A = R.zeros([64, 128], "int32") B = R.ones([64, 128], "int32") C = R.multiply(A, B) D = R.zeros([64, 128], "float32") E = R.ones([1, 128], "float32") F = R.add(D, E) output = (C, F) R.output(output) return output after = rewrite_call(pat, rewriter, before) tvm.ir.assert_structural_equal(expected, after) def test_backtrack_if_rewriter_returns_no_op(): """Rewriter participates in the pattern matching Sometimes, the pattern-matching syntax is insufficient to check if a replacement may be performed. In this case, the `rewriter` function may perform additional validation. If this validation fails, the `rewriter` function can return the original expression, and no replacement is performed. In addition, when the `rewriter` returns the original expression, the pattern match should backtrack to determine if another branch of the match may have produced a replacement. This functionality allows pattern replacements to be composed. """ pat_match_no_rewrite = is_op("relax.add")(wildcard(), wildcard()) pat_arg = wildcard() pat_zeros = is_op("relax.zeros")(wildcard()) pat_add = is_op("relax.add")(pat_arg, pat_zeros) # OR conditions are checked in the order that they occur. Because # `pat_match_no_rewrite` is a superset of `pat_add`, it will # always match first. pat = pat_match_no_rewrite | pat_add def rewriter(expr, matches): if pat_match_no_rewrite in matches: # This branch simulates a rewrite whose precondition has # failed. If the pattern-matching treats this as a # successful match with no replacemen required, then no # rewrite would be performed. On the other hand, if the # pattern-matching treats this as an unsuccessful match, # then it can backtrack and attempt `pat_add` instead. return expr elif pat_add in matches: return matches[pat_arg] else: raise RuntimeError("Pattern matched, but neither branch matched") @R.function(private=True) def before(): with R.dataflow(): A = R.ones([64, 128], "int32") B = R.zeros([64, 128], "int32") C = R.add(A, B) R.output(C) return C @R.function(private=True) def expected(): with R.dataflow(): C = R.ones([64, 128], "int32") R.output(C) return C after = rewrite_call(pat, rewriter, before) tvm.ir.assert_structural_equal(expected, after) def test_backtrack_for_no_op_rewriter_does_not_match_on_var(): """The matches should always contain the bound value This is a regression test. In versions from https://github.com/apache/tvm/pull/16732 to https://github.com/apache/tvm/pull/16828, the `rewrite_call` function could erroneously call the rewriter with `expr` and `matches[pat]` set to a variable (`C`) instead of the value to which it is bound (`R.add(A,B)`). """ pat_a = is_op("relax.add")(wildcard(), wildcard()) pat_b = is_op("relax.add")(wildcard(), wildcard()) pat = pat_a | pat_b def rewriter(expr, matches): assert isinstance(matches[pat], rx.Call) return expr @R.function(private=True) def before(): with R.dataflow(): A = R.ones([64, 128], "int32") B = R.zeros([64, 128], "int32") C = R.add(A, B) R.output(C) return C expected = before after = rewrite_call(pat, rewriter, before) tvm.ir.assert_structural_equal(expected, after) if __name__ == "__main__": tvm.testing.main()