From 5859fdff2fc9bc738928bd312f1279dce01a124f Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 12 Sep 2023 21:26:54 +0000 Subject: [PATCH 1/3] [Unity][Transform] Canonicalize and use CSE between pattern matches The `PatternRewriter` is intended to iterate until no matching patterns remain. Prior to this commit, this only involved repeating the pattern match rewrite rules. However, intermediate results produced by pattern replacement could cause the iterative pattern matching to terminate early. * If two rewrite rules each introduce the same intermediate, there will exist two copies of that intermediate, which can prevent `only_used_by` patterns from matching. Applying `EliminateCommonSubexpr` allows the pattern matching to continue. * Applying a rewrite rule may result in dangling intermediates that are no longer used. These dangling intermediates may prevent the next application of a rewrite rule that uses the `only_used_by` constraint. Applying `RemoveAllUnused` allows the pattern matching to continue. * A rewrite rule that returns a `relax::Var` or `relax::TupleGetItem` as the replacement introduces trivial var-to-var rebinding, which are not tracked by `PatternRewriter`. Applying `CanonicalizeBindings` allows the pattern matching to continue. While this could be fixed externally by repeatedly applying `rewrite_call`, this would require re-inspecting the entire function, and not just the dataflow block in which the replacement occurred. --- src/relax/ir/dataflow_matcher.cc | 97 ++++---- tests/python/relax/test_dataflow_pattern.py | 239 +++++++++++++++++++- 2 files changed, 285 insertions(+), 51 deletions(-) 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/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() From 6cd00688c7152f475281ffa97fbe0ed645e856dc Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Wed, 11 Oct 2023 19:04:47 -0500 Subject: [PATCH 2/3] Fix tests for removing redundant reshapes --- .../transform/remove_redundant_reshape.py | 15 ++++++++++----- .../relax/test_remove_redundant_reshape.py | 18 ++++++++---------- 2 files changed, 18 insertions(+), 15 deletions(-) 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/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) From 001298d1cc0fa8a3688143b40ecddf0627b1d5eb Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Wed, 11 Oct 2023 21:15:14 -0500 Subject: [PATCH 3/3] Fixed failing unit tests, along with edge case in CSE --- .../transform/eliminate_common_subexpr.cc | 22 +++++++++++++++- .../relax/test_optimize_layout_transform.py | 25 +++++++------------ 2 files changed, 30 insertions(+), 17 deletions(-) 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_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(