From 795c9098b82440d49beb3537c4ec1b6adda8ddbc Mon Sep 17 00:00:00 2001 From: Ghosts381937 Date: Wed, 19 Mar 2025 14:35:12 +0800 Subject: [PATCH] [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));