[Cherry-Pick][ARITH] Enhance CanonicalSimplify to Simplify ProdDiv (apache#174)

This PR enhances canonical simplify to simplify product division pattern
where both side involves symbolic variables.

Test cases are added.

Co-authored-by: Tianqi Chen <[email protected]>

Author: MasterJH5574

src/arith/canonical_simplify.cc

@@ -574,6 +574,23 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
       p->stream << ')';
     });
 
+/*!
+ * \brief Unpack reduction by calling each leaf via fleaf
+ *
+ * \param value The expression value.
+ * \tparam TNode the reduction node to match.
+ * \tparam FLeaf The callback function at leaf.
+ */
+template <typename TNode, typename FLeaf>
+void UnpackReduction(const PrimExpr& value, FLeaf fleaf) {
+  if (const TNode* node = value.as<TNode>()) {
+    UnpackReduction<TNode, FLeaf>(node->a, fleaf);
+    UnpackReduction<TNode, FLeaf>(node->b, fleaf);
+  } else {
+    fleaf(value);
+  }
+}
+
 // Sub-class RewriteSimplifier::Impl to take benefit of
 // rewriter for condition simplification etc.
 class CanonicalSimplifier::Impl : public RewriteSimplifier::Impl {
@@ -633,6 +650,27 @@ class CanonicalSimplifier::Impl : public RewriteSimplifier::Impl {
    */
   void SeparateDivisibleParts(const SumExprNode* psum, int64_t coeff, SumExpr* out_divisible,
                               SumExpr* out_non_divisible);
+  /*!
+   * \brief Pattern match and check whether lhs is fully divisible by
+   *        rhs using prod pattern simiplification expressions.
+   *
+   * The following two relations holds for floordiv/mod and truncdiv/mod
+   * Note that the relation do not hold for euclidean divide and mod.
+   *
+   * This is because the floordiv/mod and truncdiv/mod result can be
+   * uniquely determined by the value of the realdiv result and the
+   * relation holds for realdiv.
+   *
+   * - div((a0 * a1 * c), (b0 * b1 * c)) = div((a0 * a1), (b0 * b1))
+   * - mod((a0 * a1 * c), (b0 * b1 * c)) = mod((a0 * a1), (b0 * b1)) * c
+   *
+   * \param lhs The left operand to be updated.
+   * \param rhs The right operand to be updated.
+   * \param common_scale The common scale between lhs and rhs.
+   * \returns The simplified result if it is successful.
+   * \note This simplification mainly target when rhs is symbolic.
+   */
+  bool ProdDivSimplify(PrimExpr* lhs, PrimExpr* rhs, PrimExpr* common_scale);
   /*!
    * \brief Normalize expr to normal expr.
    * \param expr The input expression.
@@ -862,6 +900,66 @@ SplitExpr CanonicalSimplifier::Impl::SplitDivConst(SplitExpr lhs, int64_t cval,
   return lhs;
 }
 
+bool CanonicalSimplifier::Impl::ProdDivSimplify(PrimExpr* plhs, PrimExpr* prhs,
+                                                PrimExpr* common_scale) {
+  // the constant rhs case is covered by other simplifier so
+  // we just skip to save the time
+  if (prhs->as<IntImmNode>()) return false;
+  // collect lhs products and try to eliminate by matching them to prod in rhs
+  Array<Optional<PrimExpr>> lhs_prods;
+  PrimExpr new_rhs = make_const(prhs->dtype(), 1);
+  PrimExpr new_common_scale = make_const(prhs->dtype(), 1);
+  int64_t lhs_cscale = 1, rhs_cscale = 1;
+  int num_elimination = 0;
+
+  // collect lhs product and constant scale.
+  auto fcollect_lhs = [&](PrimExpr value) {
+    if (auto* intimm = value.as<tir::IntImmNode>()) {
+      lhs_cscale *= intimm->value;
+    } else {
+      lhs_prods.push_back(value);
+    }
+  };
+  UnpackReduction<tir::MulNode>(*plhs, fcollect_lhs);
+
+  // collect rhs product and try to eliminate when possible
+  PEqualChecker<PrimExpr> deep_equal;
+  auto fcollect_rhs = [&](PrimExpr value) {
+    if (auto* intimm = value.as<tir::IntImmNode>()) {
+      rhs_cscale *= intimm->value;
+    } else {
+      // try eliminate from lhs
+      for (size_t i = 0; i < lhs_prods.size(); ++i) {
+        if (lhs_prods[i].defined() && deep_equal(value, lhs_prods[i].value())) {
+          lhs_prods.Set(i, NullOpt);
+          ++num_elimination;
+          new_common_scale = new_common_scale * value;
+          return;
+        }
+      }
+      // if elimination is not possible then construct the expression.
+      new_rhs = new_rhs * value;
+    }
+  };
+  UnpackReduction<tir::MulNode>(*prhs, fcollect_rhs);
+  // find gcd of const scales.
+  int64_t cscale_gcd = ZeroAwareGCD(lhs_cscale, rhs_cscale);
+  lhs_cscale /= cscale_gcd;
+  rhs_cscale /= cscale_gcd;
+  // if no elimination is possible
+  if (num_elimination == 0 && cscale_gcd == 1) return false;
+
+  // construct prod via canonical form
+  PrimExpr new_lhs = make_const(plhs->dtype(), 1);
+  for (Optional<PrimExpr> val : lhs_prods) {
+    if (val.defined()) new_lhs = new_lhs * val.value();
+  }
+  *plhs = new_lhs * make_const(plhs->dtype(), lhs_cscale);
+  *prhs = new_rhs * make_const(prhs->dtype(), rhs_cscale);
+  *common_scale = new_common_scale * make_const(prhs->dtype(), cscale_gcd);
+  return true;
+}
+
 PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const DivNode* op) {
   if (!IsIndexType(op->dtype)) {
     return Rewriter::VisitExpr_(op);
@@ -913,6 +1011,12 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const DivNode* op) {
   // normal path
   a = Normalize(a);
   b = Normalize(b);
+  PrimExpr scale;
+  // note this is the case where b is not constant
+  if (ProdDivSimplify(&a, &b, &scale)) {
+    // use operator ver so it can constant fold if b == 1
+    return truncdiv(a, b);
+  }
   if (op->a.same_as(a) && op->b.same_as(b)) {
     return GetRef<PrimExpr>(op);
   } else {
@@ -967,6 +1071,11 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const FloorDivNode* op) {
   // normal path
   a = Normalize(a);
   b = Normalize(b);
+  PrimExpr scale;
+  if (ProdDivSimplify(&a, &b, &scale)) {
+    // use operator ver so it can const fold.
+    return floordiv(a, b);
+  }
   if (op->a.same_as(a) && op->b.same_as(b)) {
     return GetRef<PrimExpr>(op);
   } else {
@@ -1088,6 +1197,13 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const ModNode* op) {
   // normal path
   a = Normalize(a);
   b = Normalize(b);
+
+  PrimExpr scale;
+  if (ProdDivSimplify(&a, &b, &scale)) {
+    // use operator version here so it can const fold b == 1
+    return truncmod(a, b) * scale;
+  }
+
   if (op->a.same_as(a) && op->b.same_as(b)) {
     return GetRef<PrimExpr>(op);
   } else {
@@ -1146,6 +1262,13 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const FloorModNode* op) {
   // normal path
   a = Normalize(a);
   b = Normalize(b);
+
+  PrimExpr scale;
+  if (ProdDivSimplify(&a, &b, &scale)) {
+    // use operator version here so it can const fold b == 1
+    return floormod(a, b) * scale;
+  }
+
   if (op->a.same_as(a) && op->b.same_as(b)) {
     return GetRef<PrimExpr>(op);
   } else {

tests/python/relax/test_op_manipulate.py

@@ -134,7 +134,7 @@ def test_reshape_infer_struct_info_shape_symbolic():
     _check_inference(
         bb,
         relax.op.reshape(x, (d, c, b, -1)),
-        relax.TensorStructInfo((d, c, b, tir.floordiv(a * b * c * d, d * c * b)), "float32"),
+        relax.TensorStructInfo((d, c, b, a), "float32"),
     )
     _check_inference(
         bb,
@@ -144,12 +144,12 @@ def test_reshape_infer_struct_info_shape_symbolic():
     _check_inference(
         bb,
         relax.op.reshape(x, (2, -1, a)),
-        relax.TensorStructInfo((2, tir.floordiv(a * b * c * d, a * 2), a), "float32"),
+        relax.TensorStructInfo((2, tir.floordiv(b * c * d, 2), a), "float32"),
     )
     _check_inference(
         bb,
         relax.op.reshape(x, (c, -1, d, b)),
-        relax.TensorStructInfo((c, tir.floordiv(a * b * c * d, c * d * b), d, b), "float32"),
+        relax.TensorStructInfo((c, a, d, b), "float32"),
     )
     _check_inference(
         bb,
@@ -159,9 +159,7 @@ def test_reshape_infer_struct_info_shape_symbolic():
     _check_inference(
         bb,
         relax.op.reshape(x, (c, a * b * d, -1)),
-        relax.TensorStructInfo(
-            (c, a * b * d, tir.floordiv(a * b * c * d, c * (a * b * d))), "float32"
-        ),
+        relax.TensorStructInfo((c, a * b * d, 1), "float32"),
     )
     # Remove Var from StructInfo when we can
     _check_inference(bb, relax.op.reshape(x, s0), relax.TensorStructInfo((c, a, d, b), "float32"))

tests/python/unittest/test_arith_canonical_simplify.py

@@ -386,5 +386,34 @@ def test_simplify_normalize_min_value_expr():
     ck.verify(0 == x + te.min_value("int32"), False)
 
 
+def test_proddiv_simplify():
+    ck = CanonicalChecker()
+    flm = tvm.te.floormod
+    fld = tvm.te.floordiv
+    tdiv = tvm.te.truncdiv
+    tmod = tvm.te.truncmod
+
+    x, y, z = te.var("x"), te.var("y"), te.var("y")
+
+    ck.verify(flm(x * 32 * x, x), 0)
+    ck.verify(flm(z * x * 32 * x * y, x * z), 0)
+    ck.verify(flm(z * x * 32 * x * y, x * z * y * 8 * x), 0)
+    ck.verify(flm(z * x * 32 * (x * y), 6 * x * z), flm(x * y * 16, 3) * (x * z * 2))
+    ck.verify(flm(x * 32 * x, x * z), flm(x * 32, z) * x)
+
+    ck.verify(tmod(x * 32 * x, x), 0)
+    ck.verify(tmod(z * x * 32 * x * y, x * z), 0)
+    ck.verify(tmod(z * x * 32 * (x * y), 6 * x * z), tmod(x * y * 16, 3) * (x * z * 2))
+    ck.verify(tmod(x * 32 * x, x * z), tmod(x * 32, z) * x)
+
+    ck.verify(fld(x * 2 * x * z, 4 * x * x * x), fld(z, x * 2))
+    ck.verify(fld(x * (2 * y) * 3, 3 * y), x * 2)
+    ck.verify(fld(x * (2 * y) * 3, 3 * y * z), fld(x * 2, z))
+
+    ck.verify(tdiv(x * 2 * x * z, 4 * x * x * x), tdiv(z, x * 2))
+    ck.verify(tdiv(x * (2 * y) * 3, 3 * y), x * 2)
+    ck.verify(tdiv(x * (2 * y) * 3, 3 * y * z), tdiv(x * 2, z))
+
+
 if __name__ == "__main__":
     tvm.testing.main()