tqchen commented on a change in pull request #7752:
URL: https://github.com/apache/tvm/pull/7752#discussion_r602342033
##########
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:
move to src/analysis/expr_complexity.cc analysis.h and expose as a
function `size_t ExprComplexity(const PrimExpr& expr)`; document as number of
expressions in the child
--
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:
[email protected]
