diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 4d8b6ff769cf..dab78c77a0a1 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -411,6 +411,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. 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..01cc41c7cec7 100644 --- a/tests/python/unittest/test_tir_transform_simplify.py +++ b/tests/python/unittest/test_tir_transform_simplify.py @@ -15,7 +15,10 @@ # 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 def test_stmt_simplify(): @@ -133,9 +136,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() + tvm.testing.main()