From 795c9098b82440d49beb3537c4ec1b6adda8ddbc Mon Sep 17 00:00:00 2001 From: Ghosts381937 Date: Wed, 19 Mar 2025 14:35:12 +0800 Subject: [PATCH 1/2] [Fix] Enhance floormod simplification rules for better expression matching --- src/arith/rewrite_simplify.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 3682054e8e4b..6805a427c19d 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -1230,7 +1230,8 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) { TVM_TRY_REWRITE_IF(floormod(x + y * c1, c2), floormod(x + y * floormod(c1, c2), c2), c2.Eval()->value > 0); - TVM_TRY_REWRITE_IF(floormod(x * c1, x * c2), x * floormod(c1, c2), c2.Eval()->value != 0); + TVM_TRY_REWRITE_IF(matches_one_of(floormod(x * c1, x * c2), floormod(c1 * x, c2 * x)), + floormod(c1, c2), c2.Eval()->value != 0); TVM_TRY_REWRITE(matches_one_of(floormod(x * y, y), floormod(y * x, y)), ZeroWithTypeLike(y)); From a2119aba1773d1e341cf0e33211e29eba45e7503 Mon Sep 17 00:00:00 2001 From: Ghosts381937 Date: Sat, 10 May 2025 22:44:32 +0800 Subject: [PATCH 2/2] [Enhance] Add rewrite rule for coefficient folding in addition --- src/arith/rewrite_simplify.cc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/arith/rewrite_simplify.cc b/src/arith/rewrite_simplify.cc index 6805a427c19d..6fae80fa166f 100644 --- a/src/arith/rewrite_simplify.cc +++ b/src/arith/rewrite_simplify.cc @@ -446,6 +446,8 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const AddNode* op) { // mul co-efficient folding TVM_TRY_REWRITE(x + x, x * 2); + TVM_TRY_REWRITE(x * c1 + c1, (x + 1) * c1); + TVM_TRY_REWRITE(matches_one_of(x * y + x, y * x + x, x + y * x, x + x * y), x * (y + 1)); TVM_TRY_REWRITE(matches_one_of(x * y + x * z, y * x + x * z, x * y + z * x, y * x + z * x),