From e6f0238ff03c5f0b5b7d9f4be774548dd4ad0082 Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 19 May 2022 11:45:29 -0500 Subject: [PATCH 1/4] [TIR] Additional Stmt/Expr simplication rules - Enabled simplification of `A[i] = A[i] + 0` into no-op. This was a bug introduced in https://github.com/apache/tvm/pull/9727, which applied this rewrite only to `A[i] = A[i]`, and not to statements which simplify to `A[i] = A[i]`. Regression test added to prevent reoccurrence of this bug. - Enabled simplification of `x - x` to zero for floating point types. Previously, this simplification was applied only for data types that could be used as buffer indices. --- src/arith/rewrite_simplify.cc | 14 +++--- src/tir/transforms/simplify.cc | 12 +++--- .../unittest/test_arith_rewrite_simplify.py | 8 ++++ .../unittest/test_tir_transform_simplify.py | 43 ++++++++++++++++--- 4 files changed, 60 insertions(+), 17 deletions(-) diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 4d8b6ff769cf..ac107861ece4 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -256,13 +256,16 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { TVM_TRY_REWRITE(broadcast(x, lanes) - broadcast(y, lanes), broadcast(x - y, lanes)); } + // cancelation rules + TVM_TRY_REWRITE_IF(x - x, ZeroWithTypeLike(x), + SideEffect(x.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF((x + y) - y, x, SideEffect(y.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF((x + y) - x, y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF(x - (y + x), 0 - y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF(x - (x + y), 0 - y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); + if (IsIndexType(op->dtype)) { // Index rules - // cancelation rules - TVM_TRY_REWRITE((x + y) - y, x); - TVM_TRY_REWRITE((x + y) - x, y); - TVM_TRY_REWRITE(x - (y + x), 0 - y); - TVM_TRY_REWRITE(x - (x + y), 0 - y); TVM_TRY_REWRITE(min(x, y) - x, min(0, y - x)); TVM_TRY_REWRITE(min(x, y) - y, min(x - y, 0)); @@ -275,7 +278,6 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { TVM_TRY_REWRITE(y - min(x, y), max(y - x, 0)); // mul co-efficient folding - TVM_TRY_REWRITE(x - x, ZeroWithTypeLike(x)); TVM_TRY_REWRITE(x * y - x, x * (y - 1)); TVM_TRY_REWRITE(y * x - x, x * (y - 1)); TVM_TRY_REWRITE(x - y * x, x * (1 - y)); diff --git a/src/tir/transforms/simplify.cc b/src/tir/transforms/simplify.cc index 7d4fac8d7b2d..85f405be447a 100644 --- a/src/tir/transforms/simplify.cc +++ b/src/tir/transforms/simplify.cc @@ -90,12 +90,12 @@ class StmtSimplifier : public IRMutatorWithAnalyzer { // eliminate useless stores Stmt VisitStmt_(const BufferStoreNode* op) final { BufferStore store = Downcast(Parent::VisitStmt_(op)); - if (const BufferLoadNode* load = op->value.as()) { - if (load->buffer->data.same_as(op->buffer->data) && - ArrayDeepEqual(load->indices, op->indices) && - tir::ExprDeepEqual()(load->buffer->elem_offset, op->buffer->elem_offset) && - ArrayDeepEqual(load->buffer->shape, op->buffer->shape) && - ArrayDeepEqual(load->buffer->strides, op->buffer->strides)) { + if (const BufferLoadNode* load = store->value.as()) { + if (load->buffer->data.same_as(store->buffer->data) && + ArrayDeepEqual(load->indices, store->indices) && + tir::ExprDeepEqual()(load->buffer->elem_offset, store->buffer->elem_offset) && + ArrayDeepEqual(load->buffer->shape, store->buffer->shape) && + ArrayDeepEqual(load->buffer->strides, store->buffer->strides)) { return Evaluate(0); } } diff --git a/tests/python/unittest/test_arith_rewrite_simplify.py b/tests/python/unittest/test_arith_rewrite_simplify.py index 855635b3f962..8d26710f40db 100644 --- a/tests/python/unittest/test_arith_rewrite_simplify.py +++ b/tests/python/unittest/test_arith_rewrite_simplify.py @@ -972,5 +972,13 @@ def test_div_zero_simplify(): assert "division by zero" in str(cm.execption) +def test_sub_bufferload(): + ck = RewriteChecker() + buf = tvm.tir.decl_buffer([1], dtype="float32") + load = tvm.tir.BufferLoad(buf, [0]) + expr = load - load + ck.verify(expr, 0.0) + + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/python/unittest/test_tir_transform_simplify.py b/tests/python/unittest/test_tir_transform_simplify.py index 824bef4f32f9..f1401f0e77f1 100644 --- a/tests/python/unittest/test_tir_transform_simplify.py +++ b/tests/python/unittest/test_tir_transform_simplify.py @@ -16,6 +16,7 @@ # under the License. import tvm from tvm import te +from tvm.script import tir as T def test_stmt_simplify(): @@ -133,9 +134,41 @@ def sls(n, d): assert "if" not in str(stmt) +def test_load_store_noop(): + """Store of a value that was just read from the same location is a no-op.""" + + @T.prim_func + def before(A: T.Buffer[(1,), "float32"]): + A[0] = A[0] + + @T.prim_func + def expected(A: T.Buffer[(1,), "float32"]): + T.evaluate(0) + + after = tvm.tir.transform.Simplify()(tvm.IRModule.from_expr(before))["main"] + tvm.ir.assert_structural_equal(after, expected) + + +def test_load_store_noop_after_simplify(): + """As test_load_store_noop, but requiring simplification to identify. + + Previously, a bug caused the self-assignment of a buffer to + checked based on the pre-simplification assignment, not the + post-simplification. This test is to identify any similar + regression. + """ + + @T.prim_func + def before(A: T.Buffer[(1,), "float32"]): + A[0] = A[0] + (5.0 - 5.0) + + @T.prim_func + def expected(A: T.Buffer[(1,), "float32"]): + T.evaluate(0) + + after = tvm.tir.transform.Simplify()(tvm.IRModule.from_expr(before))["main"] + tvm.ir.assert_structural_equal(after, expected) + + if __name__ == "__main__": - test_stmt_simplify() - test_thread_extent_simplify() - test_if_likely() - test_basic_likely_elimination() - test_complex_likely_elimination() + sys.exit(pytest.main(sys.argv)) From f38b09c86b86f650faa33bc0bb085292bd50710a Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Tue, 24 May 2022 09:03:54 -0500 Subject: [PATCH 2/4] Updated to maintain separate int/float simplification paths --- src/arith/rewrite_simplify.cc | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index ac107861ece4..142743ece0dd 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -266,6 +266,11 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { if (IsIndexType(op->dtype)) { // Index rules + // cancelation rules + TVM_TRY_REWRITE((x + y) - y, x); + TVM_TRY_REWRITE((x + y) - x, y); + TVM_TRY_REWRITE(x - (y + x), 0 - y); + TVM_TRY_REWRITE(x - (x + y), 0 - y); TVM_TRY_REWRITE(min(x, y) - x, min(0, y - x)); TVM_TRY_REWRITE(min(x, y) - y, min(x - y, 0)); @@ -278,6 +283,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { TVM_TRY_REWRITE(y - min(x, y), max(y - x, 0)); // mul co-efficient folding + TVM_TRY_REWRITE(x - x, ZeroWithTypeLike(x)); TVM_TRY_REWRITE(x * y - x, x * (y - 1)); TVM_TRY_REWRITE(y * x - x, x * (y - 1)); TVM_TRY_REWRITE(x - y * x, x * (1 - y)); @@ -413,6 +419,15 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { TVM_TRY_RECURSIVE_REWRITE((x + c1) - y, (x - y) + c1); TVM_TRY_RECURSIVE_REWRITE(x - (y - z), (x + z) - y); TVM_TRY_RECURSIVE_REWRITE(x - y * c1, x + y * (0 - c1)); + } else if (op->dtype.is_float()) { + // Cancellation rules. Deliberately off of the integer path, to + // avoid introducing checks on the side effects for the fast path. + TVM_TRY_REWRITE_IF(x - x, ZeroWithTypeLike(x), + SideEffect(x.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF((x + y) - y, x, SideEffect(y.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF((x + y) - x, y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF(x - (y + x), 0 - y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); + TVM_TRY_REWRITE_IF(x - (x + y), 0 - y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); } // condition rules. From bd7e8cd3468080cfacf1ea3205b9aad6e28dd69e Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Wed, 25 May 2022 10:05:50 -0500 Subject: [PATCH 3/4] Updated to use tvm.testing.main --- tests/python/unittest/test_tir_transform_simplify.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/python/unittest/test_tir_transform_simplify.py b/tests/python/unittest/test_tir_transform_simplify.py index f1401f0e77f1..01cc41c7cec7 100644 --- a/tests/python/unittest/test_tir_transform_simplify.py +++ b/tests/python/unittest/test_tir_transform_simplify.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. import tvm +import tvm.testing + from tvm import te from tvm.script import tir as T @@ -171,4 +173,4 @@ def expected(A: T.Buffer[(1,), "float32"]): if __name__ == "__main__": - sys.exit(pytest.main(sys.argv)) + tvm.testing.main() From 069fba9a1cb786c1bd94140d2fa45162a63ae37c Mon Sep 17 00:00:00 2001 From: Eric Lunderberg Date: Thu, 26 May 2022 07:01:20 -0500 Subject: [PATCH 4/4] Remove duplicate rewrite rules --- src/arith/rewrite_simplify.cc | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 142743ece0dd..dab78c77a0a1 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -256,14 +256,6 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) { TVM_TRY_REWRITE(broadcast(x, lanes) - broadcast(y, lanes), broadcast(x - y, lanes)); } - // cancelation rules - TVM_TRY_REWRITE_IF(x - x, ZeroWithTypeLike(x), - SideEffect(x.Eval()) <= CallEffectKind::kReadState); - TVM_TRY_REWRITE_IF((x + y) - y, x, SideEffect(y.Eval()) <= CallEffectKind::kReadState); - TVM_TRY_REWRITE_IF((x + y) - x, y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); - TVM_TRY_REWRITE_IF(x - (y + x), 0 - y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); - TVM_TRY_REWRITE_IF(x - (x + y), 0 - y, SideEffect(x.Eval()) <= CallEffectKind::kReadState); - if (IsIndexType(op->dtype)) { // Index rules // cancelation rules