from tvm.script import relax as R
from tvm.script import tir as T
from tvm import relax as rx
from tvm import relay, tir
from tvm.relax.analysis import get_var2val
import tvm.testing
from tvm.relax.dpl import *
same_shape_func_type = tvm.testing.parameter(
"same_static_shape",
"same_dynamic_shape",
"different_static_shape",
"different_dynamic_shape",
)
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.struct_info.shape is None:
return expr
if len(axes) != 1:
return expr
axis = axes[0].value
begin = begin[0].value
end = end[0].value
stride = strides[0].value
if stride != 1:
return expr
size = arg.struct_info.shape[0]
if (
isinstance(size, tir.IntImm)
and isinstance(begin, tir.IntImm)
and isinstance(end, tir.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.ir.structural_equal(args[0].struct_info, args[1].struct_info):
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)
