[GitHub] [tvm] FrozenGene commented on a change in pull request #7752: [ARITH] detect iter affine map with predicate

[GitBox <gi...@apache.org>]

FrozenGene commented on a change in pull request #7752:
URL: https://github.com/apache/tvm/pull/7752#discussion_r602653607



##########
File path: src/arith/iter_affine_map.cc
##########
@@ -459,27 +665,107 @@ class IterMapRewriter : public ExprMutator {
   }
 };
 
+/*! \brief An internal struct to represent range extent on iterators(iter < upper_bound). */
+struct IterConstraint {
+  // The expr of the iter
+  PrimExpr iter;
+  // The expr of the upper_bound
+  PrimExpr upper_bound;
+  // The size of the iter, which is the number of nodes
+  size_t expr_size = 0;
+
+  IterConstraint(PrimExpr iter, PrimExpr upper_bound, size_t size)
+      : iter(std::move(iter)), upper_bound(std::move(upper_bound)), expr_size(size) {}
+};
+
+/*!
+ * \brief Split the predicate into `(a < b) && (c < d) && ...`
+ * \param pred The predicate to be split.
+ * \return A list of pairs, each element of which are lhs and rhs of the '<' sign,
+ *         empty if the split failed.
+ */
+std::vector<IterConstraint> MatchUpperBoundConstraints(PrimExpr pred) {
+  std::vector<IterConstraint> result;
+  arith::PVar<PrimExpr> lhs, rhs, rest;
+  for (;;) {
+    if ((rest && (lhs < rhs)).Match(pred)) {
+      result.emplace_back(lhs.Eval(), rhs.Eval(), 0);
+      pred = rest.Eval();
+    } else if ((lhs < rhs).Match(pred)) {
+      result.emplace_back(lhs.Eval(), rhs.Eval(), 0);
+      break;
+    } else {
+      return std::vector<IterConstraint>();
+    }
+  }
+  return result;
+}
+
+/*! \brief Count the size of the PrimExpr. */
+class PrimExprSizeCounter : public ExprVisitor {

Review comment:
       Yes. In my previous pr #7469 I have done something. We could borrow it. My pr has the left problem of max ops counter, however, it is not an urgent problem IMO, which just a more safer guard for complex expr analysis.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org