diff --git a/python/tvm/relax/transform/remove_redundant_reshape.py b/python/tvm/relax/transform/remove_redundant_reshape.py index 2274f8e5dafa..a48923df78a2 100644 --- a/python/tvm/relax/transform/remove_redundant_reshape.py +++ b/python/tvm/relax/transform/remove_redundant_reshape.py @@ -66,13 +66,18 @@ def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRM continue def rewriter(expr, matches): - args = matches[self.pattern] + arg = matches[self.input1] + if self.repeated_reshape in matches: - return relax.op.reshape(matches[self.input1], args.args[1]) + output_shape = matches[self.repeated_reshape].args[1] + return relax.op.reshape(arg, output_shape) + elif self.no_op_reshape in matches: - if args.args[0].struct_info.shape: - if structural_equal(args.args[0].struct_info.shape, args.args[1]): - return args.args[0] + output_shape = matches[self.no_op_reshape].args[1] + if arg.struct_info.shape and structural_equal( + arg.struct_info.shape, output_shape + ): + return arg return expr updated_func = rewrite_call(self.pattern, rewriter, funct) diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc index e85e3c4d5182..d1edb945baf3 100644 --- a/src/relax/ir/dataflow_matcher.cc +++ b/src/relax/ir/dataflow_matcher.cc @@ -42,6 +42,7 @@ #include #include +#include "../transform/utils.h" #include "dataflow_matcher_impl.h" namespace tvm { @@ -937,57 +938,76 @@ class PatternRewriter : ExprMutator { return Downcast(RemoveAllUnused(rewriter.VisitExpr(f))); } - void VisitBinding_(const VarBindingNode* binding) final { - bindings_.Set(binding->var, binding->value); - ExprMutator::VisitBinding_(binding); - if (auto it = memo_.find(binding->value.get()); it != memo_.end()) { - // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten - // expression can be subject to further pattern matchings. - bindings_.Set(binding->var, it->second); + Expr VisitExpr_(const SeqExprNode* seq) override { + if (ctx_) { + return ExprMutator::VisitExpr_(seq); } - } - Expr VisitExpr(const Expr& expr) final { - auto node = ExprMutator::VisitExpr(expr); - if (pattern_) { - if (auto matches_opt = ExtractMatchedExpr(pattern_.value(), node, bindings_)) { - Expr rewritten_expr = rewriter_func_(node, matches_opt.value()); - if (!rewritten_expr.same_as(node)) { - rewritten_expr = builder_->Normalize(rewritten_expr); - - // If the rewriter returns a variable (e.g. when rewriting - // from `R.add(x, R.const(0.0))` to `x`), the variable - // should be dereferenced to avoid trivial `var_2 = var_1` - // bindings. This lookup is done using the builder_ instead - // of the bindings_, as the previous `builder_->Normalize` - // call may have introduced variable bindings. - if (auto opt_var = rewritten_expr.as()) { - if (auto binding = builder_->LookupBinding(opt_var.value())) { - rewritten_expr = binding.value(); - } - } - memo_[expr.get()] = rewritten_expr; - return rewritten_expr; + auto cache = bindings_; + SeqExpr prev = GetRef(seq); + + StructuralEqual struct_equal; + + while (true) { + SeqExpr next = Downcast(builder_->Normalize(ExprMutator::VisitExpr_(prev.get()))); + if (struct_equal(prev, next)) { + return std::move(next); + } + + // Canonicalization may result in two previously-different + // expressions being recognized as identical. Elimination of + // common subexpressions may result in trival var-to-var + // bindings that can be canonicalized. Therefore, iterate the + // simplification steps until converged. + while (true) { + auto start_of_loop = next; + next = Downcast(CanonicalizeBindings(next)); + next = Downcast(EliminateCommonSubexpr(next)); + next = Downcast(RemoveAllUnused(next)); + if (struct_equal(start_of_loop, next)) { + break; } } + + if (struct_equal(prev, next)) { + return std::move(next); + } + + // Reset all knowledge of bindings that were collected from + // this DataflowBlock. The collected bindings are only after + // the point where they were collected, and we are repeating + // the mutation of this DataflowBlock. + bindings_ = cache; + prev = next; } - return node; } - BindingBlock VisitBindingBlock_(const DataflowBlockNode* block_node) final { + BindingBlock VisitBindingBlock_(const DataflowBlockNode* block_node) override { if (ctx_) { return RewriteDataflowBlockFixedPoint(GetRef(block_node)); + } else { + return ExprMutator::VisitBindingBlock_(block_node); } + } - DataflowBlock prev = GetRef(block_node); - while (true) { - DataflowBlock next = Downcast(ExprMutator::VisitBindingBlock_(prev.get())); - if (StructuralEqual()(prev, next)) { - return std::move(next); - } else { - prev = next; + void VisitBinding_(const VarBindingNode* binding) override { + auto expr = VisitExpr(binding->value); + bindings_.Set(binding->var, expr); + ReEmitBinding(binding, expr); + } + + Expr VisitExpr(const Expr& expr) override { + auto node = ExprMutator::VisitExpr(expr); + + if (pattern_) { + if (auto matches_opt = ExtractMatchedExpr(pattern_.value(), node, bindings_)) { + Expr rewritten_expr = rewriter_func_(node, matches_opt.value()); + if (!rewritten_expr.same_as(node)) { + return builder_->Normalize(rewritten_expr); + } } } + return node; } private: @@ -1076,7 +1096,6 @@ class PatternRewriter : ExprMutator { PackedFunc rewriter_func_; std::unordered_set params_; Map bindings_; - std::unordered_map memo_; }; Function RewriteBindings(const PatternContext& ctx, PackedFunc rewriter, Function f) { diff --git a/src/relax/transform/eliminate_common_subexpr.cc b/src/relax/transform/eliminate_common_subexpr.cc index 842470c46357..2addb60697bf 100644 --- a/src/relax/transform/eliminate_common_subexpr.cc +++ b/src/relax/transform/eliminate_common_subexpr.cc @@ -105,7 +105,27 @@ class SubexprCounter : public ExprVisitor { count_map_[e] = count + 1; } } - ExprVisitor::VisitExpr(e); + + // Only visit the interior of objects that we might still keep + // around. Otherwise, double-counting these would lead to extra + // variable bindings. + // + // Before: + // y = f(a+b) + // z = f(a+b) + // + // Expected: + // y = f(a+b) // De-duped from (y==z) + // z = y + // + // Erroneous output: + // c = a+b // Incorrect, a+b only has a single usage. + // y = f(c) // De-duped from + // z = y + // + if (auto it = count_map_.find(e); it == count_map_.end() || it->second < 2) { + ExprVisitor::VisitExpr(e); + } } // do not visit inner functions: we will do CSE within those diff --git a/tests/python/relax/test_dataflow_pattern.py b/tests/python/relax/test_dataflow_pattern.py index 7b68655a2f57..49b7d11a804e 100644 --- a/tests/python/relax/test_dataflow_pattern.py +++ b/tests/python/relax/test_dataflow_pattern.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. +import functools +import math + import pytest import tvm.testing @@ -1372,23 +1375,51 @@ def rewriter(_expr, matches): tvm.ir.assert_structural_equal(after, expected) -def test_rewrite_without_trivial_binding(): - """rewrite_call should avoid producing trivial "y = x" bindings""" +bind_to_dataflow_var = tvm.testing.parameter( + by_dict={"var-to-var": False, "var-to-dataflow-var": True} +) - @R.function(private=True) - def before(x: R.Tensor((1024,))): - with R.dataflow(): + +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(): + a = R.add(x, x) + b = a + 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,)) - R.output(b) - return b + return b - @R.function(private=True) - def expected(x: R.Tensor((1024,))): - with R.dataflow(): + @R.function(private=True) + def expected(x: R.Tensor((1024,))): a = R.add(x, x) - R.output(a) - return a + return a pattern_arg = wildcard() pattern_shape_expr = wildcard() @@ -1490,5 +1521,189 @@ def func( 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 = is_op("relax.strided_slice")(pattern_arg).has_attr( + { + "axes": [0], + "strides": [T.int64(1)], + } + ) + + def rewriter(expr, matches): + arg = matches[pattern_arg] + strided_slice = matches[pattern] + + if arg.struct_info.shape is None: + return expr + + size = arg.struct_info.shape[0] + begin = strided_slice.attrs.begin[0] + end = strided_slice.attrs.end[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) + + if __name__ == "__main__": tvm.testing.main() diff --git a/tests/python/relax/test_optimize_layout_transform.py b/tests/python/relax/test_optimize_layout_transform.py index 56c11984ee99..08c9e3110705 100644 --- a/tests/python/relax/test_optimize_layout_transform.py +++ b/tests/python/relax/test_optimize_layout_transform.py @@ -26,19 +26,15 @@ def _run_pass_compare_output(Before, Expected): - fused_mod = OptimizeLayoutTransform()(Before) - if not relax.analysis.well_formed(fused_mod): - print("IRModule is not well-formed") + After = tvm.ir.transform.Sequential( + [ + OptimizeLayoutTransform(), + DeadCodeElimination(), + FuseTIR(), + ] + )(Before) - fused_mod = DeadCodeElimination()(fused_mod) - if not relax.analysis.well_formed(fused_mod): - print("IRModule is not well-formed") - - fused_mod = FuseTIR()(fused_mod) - if not relax.analysis.well_formed(fused_mod): - print("IRModule is not well-formed") - - tvm.ir.assert_structural_equal(Expected, fused_mod) + tvm.ir.assert_structural_equal(Expected, After) def test_optimize_transform_layout_pass_one_arg(): @@ -129,12 +125,9 @@ def main( (lv, lv1), out_sinfo=R.Tensor((4, 4), dtype="float32"), ) - lv4: R.Tensor((4, 4), dtype="float32") = R.layout_transform( - y, index_map=lambda i: (i // 4, i % 4), pad_value=None - ) lv5 = R.call_tir( Expected.relax_add_replacement, - (lv4, lv2), + (lv1, lv2), out_sinfo=R.Tensor((4, 4), dtype="float32"), ) lv2_1: R.Tensor((16,), dtype="float32") = R.layout_transform( diff --git a/tests/python/relax/test_remove_redundant_reshape.py b/tests/python/relax/test_remove_redundant_reshape.py index 806b563cb8c4..11e8c87cf1aa 100644 --- a/tests/python/relax/test_remove_redundant_reshape.py +++ b/tests/python/relax/test_remove_redundant_reshape.py @@ -41,9 +41,9 @@ def main( with R.dataflow(): lv: R.Tensor((1, 1001), dtype="float16") = R.reshape(x, R.shape([1, 1001])) lv1: R.Tensor((1, 1001), dtype="float16") = R.reshape(lv, R.shape([1, 1001])) - lv2: R.Tensor((1, 1001), dtype="float16") = R.reshape(lv1, R.shape([1, 1001])) - R.output(lv2) - return lv2 + gv: R.Tensor((1, 1001), dtype="float16") = R.reshape(lv1, R.shape([1, 1001])) + R.output(gv) + return gv @I.ir_module class Expected: @@ -52,9 +52,10 @@ def main( x: R.Tensor((1, 1001, 1, 1), dtype="float16") ) -> R.Tensor((1, 1001), dtype="float16"): with R.dataflow(): - lv1: R.Tensor((1, 1001), dtype="float16") = R.reshape(x, R.shape([1, 1001])) - R.output(lv1) - return lv1 + lv: R.Tensor((1, 1001), dtype="float16") = R.reshape(x, R.shape([1, 1001])) + gv: R.Tensor((1, 1001), dtype="float16") = lv + R.output(gv) + return gv _run_pass_compare_output(Before, Expected) @@ -106,10 +107,7 @@ class Expected: def main( x: R.Tensor((1, 1001, 1, 1), dtype="float16") ) -> R.Tensor((1, 1001, 1, 1), dtype="float16"): - with R.dataflow(): - lv: R.Tensor((1, 1001, 1, 1), dtype="float16") = x - R.output(lv) - return lv + return x _run_pass_compare_output(Before, Expected)