[TIR] Allow compute_at create block predicate for non-trivial bounds and support floordiv pattern (#9527)

* allow generate block predicate in compute_at schedule

* revert #9880 and add more testcases

Author: wrongtest-intellif

include/tvm/arith/int_set.h

@@ -93,6 +93,11 @@ class IntSet : public ObjectRef {
   bool CanProveNonPositive() const;
   /*! \return Whether the set is proved to be larger than or equal to 0 */
   bool CanProveNonNegative() const;
+  /*! \return Whether the set has upper bound. */
+  bool HasUpperBound() const;
+  /*! \return Whether the set has lower bound. */
+  bool HasLowerBound() const;
+
   /*!
    * \brief The single point value, call only if IsSinglePoint is true
    * \return The point value.

src/arith/int_set.cc

@@ -574,6 +574,20 @@ bool IntSet::CanProveNonNegative() const {
   return false;
 }
 
+bool IntSet::HasLowerBound() const {
+  if (const IntervalSetNode* s_int = (*this).as<IntervalSetNode>()) {
+    return s_int->HasLowerBound();
+  }
+  return false;
+}
+
+bool IntSet::HasUpperBound() const {
+  if (const IntervalSetNode* s_int = (*this).as<IntervalSetNode>()) {
+    return s_int->HasUpperBound();
+  }
+  return false;
+}
+
 SignType IntSet::GetSignType() const {
   if (CanProvePositive()) {
     return kPositive;

src/arith/iter_affine_map.cc

@@ -503,6 +503,7 @@ class IterMapRewriter : public ExprMutator {
       if (predicate_induced_min.defined()) predicate_induced_min = predicate_induced_min - base;
       if (predicate_induced_max.defined()) predicate_induced_max = predicate_induced_max - base;
     }
+    if (expr->args.size() < 1) return expr;
     Optional<IterSumExpr> opt = TryFuseIters(expr);
     ICHECK(!opt.defined() || opt.value()->args.size() == 1);
     // scale should be 1

src/tir/analysis/block_access_region_detector.cc

@@ -145,22 +145,10 @@ void BlockReadWriteDetector::VisitExpr_(const LoadNode* op) {
   ExprVisitor::VisitExpr_(op);
 }
 
-arith::IntSet BlockReadWriteDetector::RelaxAccessIndex(const PrimExpr& index) {
-  arith::IntSet relaxed = arith::EvalSet(index, dom_map_);
-  if (!hint_map_.empty()) {
-    // take non-relaxed var bound hints into considerations
-    // eg, if i * 4 + j with i >= 10 and j in [0, 4), only j in domain scope
-    // then the index region can be relaxed to [i*4, i*4+4) ^ [40, inf)
-    arith::IntSet hint_bound = arith::EvalSet(relaxed, hint_map_);
-    relaxed = arith::Intersect({relaxed, hint_bound});
-  }
-  return relaxed;
-}
-
 void BlockReadWriteDetector::VisitExpr_(const BufferLoadNode* op) {
   std::vector<arith::IntSet> relaxed_region;
   for (const PrimExpr& index : op->indices) {
-    relaxed_region.push_back(RelaxAccessIndex(index));
+    relaxed_region.push_back(arith::EvalSet(index, dom_map_));
   }
   Update(&read_buffers_, &read_regions_, op->buffer, relaxed_region);
   ExprVisitor::VisitExpr_(op);
@@ -213,7 +201,7 @@ void BlockReadWriteDetector::VisitStmt_(const StoreNode* op) {
 void BlockReadWriteDetector::VisitStmt_(const BufferStoreNode* op) {
   std::vector<arith::IntSet> relaxed_region;
   for (const PrimExpr& index : op->indices) {
-    relaxed_region.push_back(RelaxAccessIndex(index));
+    relaxed_region.push_back(arith::EvalSet(index, dom_map_));
   }
   Update(&writes_buffers_, &write_regions_, op->buffer, relaxed_region);
   StmtVisitor::VisitStmt_(op);

src/tir/schedule/primitive/compute_at.cc

@@ -163,6 +163,49 @@ int FindInsertionPoint(
   return split.first_consumer_position;
 }
 
+/*!
+ * \brief Represent the iteration domain to fully cover the required region of Intersect(dom, bound)
+ * The bound region may not get directly intersected with dom region, instead we try to generate
+ * extra predicates for non-trivial bound. The domain info class can also union with each other.
+ */
+struct BlockVarDomainInfo {
+  arith::IntSet dom{arith::IntSet::Nothing()};  // dom is ensured to be bounded
+  arith::IntSet bound{arith::IntSet::Nothing()};
+
+  /*! \brief Relaxed union operation */
+  void Union(const BlockVarDomainInfo& other) {
+    // just relax (d0 ^ b0) v (d1 ^ b1) to (d0 v d1) ^ (b0 v b1)
+    dom = arith::Union({dom, other.dom});
+    bound = arith::Union({bound, other.bound});
+  }
+
+  /*! \brief Simplify domain info */
+  void Simplify(arith::Analyzer* analyzer) {
+    auto to_simplified = [analyzer](const arith::IntSet& set) {
+      PrimExpr min = set.HasLowerBound() ? analyzer->Simplify(set.min()) : set.min();
+      PrimExpr max = set.HasUpperBound() ? analyzer->Simplify(set.max()) : set.max();
+      return arith::IntSet::Interval(min, max);
+    };
+    // if no dom specified, try use bound as dom
+    if (dom.IsNothing()) {
+      if (bound.HasLowerBound() && bound.HasUpperBound()) {
+        bound = to_simplified(bound);
+        std::swap(dom, bound);
+      }
+      return;
+    }
+    // simplify intsets
+    dom = to_simplified(dom);
+    bound = to_simplified(bound);
+    // if can proof the dom is within bound, remove bound
+    auto intersect = to_simplified(arith::Intersect({dom, bound}));
+    if (analyzer->CanProveEqual(dom.min(), intersect.min()) &&
+        analyzer->CanProveEqual(dom.max(), intersect.max())) {
+      bound = arith::IntSet::Nothing();
+    }
+  }
+};
+
 /*!
  * \brief A helper to reconstruct the block scope where the given block is moved under the given
  * loop, and the given block's induced loop nest is regenerated to satisfy the required region.
@@ -179,29 +222,44 @@ class ScopeReconstructor : private StmtMutator {
    * \param insert_position The position among the subtrees where the block and its induced loop
    * nest is inserted
    * \param iter_doms The domain of each block var
+   * \param analyzer The arithmetic analyzer
    * \param preserve_unit_loops Whether to generate unit loops where the loop extent is 1
    */
-  void MakeNewLoop(int insert_position, std::vector<Range> iter_doms, bool preserve_unit_loops) {
+  void MakeNewLoop(int insert_position, std::vector<BlockVarDomainInfo> iter_doms,
+                   arith::Analyzer* analyzer, bool preserve_unit_loops) {
     int n_iters = iter_doms.size();
     Array<Var> loop_vars;
     Array<PrimExpr> loop_extents;
     Array<PrimExpr> iter_values;
     loop_vars.reserve(n_iters);
     loop_extents.reserve(n_iters);
     iter_values.reserve(n_iters);
+    PrimExpr predicate = const_true();
     for (int i = 0; i < n_iters; ++i) {
-      const Range& iter_dom = iter_doms[i];
+      Range iter_dom = iter_doms[i].dom.CoverRange(block_->iter_vars[i]->dom);
       if (preserve_unit_loops || !is_one(iter_dom->extent)) {
         Var var("ax" + std::to_string(loop_vars.size()), DataType::Int(32));
         loop_vars.push_back(var);
         loop_extents.push_back(iter_dom->extent);
         iter_values.push_back(iter_dom->min + var);
+        analyzer->Bind(var, Range::FromMinExtent(0, iter_dom->extent));
       } else {
         iter_values.push_back(iter_dom->min);
       }
+      const arith::IntSet& pred_bound = iter_doms[i].bound;
+      if (!pred_bound.IsNothing()) {
+        if (pred_bound.HasLowerBound()) {
+          PrimExpr lower_bound = iter_values[i] >= pred_bound.min();
+          predicate = predicate && lower_bound;
+        }
+        if (pred_bound.HasUpperBound()) {
+          PrimExpr upper_bound = iter_values[i] < pred_bound.max() + 1;
+          predicate = predicate && upper_bound;
+        }
+      }
     }
     this->new_block_realize_ =
-        BlockRealize(std::move(iter_values), const_true(), std::move(block_));
+        BlockRealize(std::move(iter_values), analyzer->Simplify(predicate), std::move(block_));
     Stmt new_subtree = this->new_block_realize_;
     for (int i = static_cast<int>(loop_vars.size()) - 1; i >= 0; --i) {
       const Var& loop_var = loop_vars[i];
@@ -310,39 +368,74 @@ void RelaxBufferRegions(const Map<Var, PrimExpr>& binding,
  * domain
  * \param provided The provided integer set to cover the required domain
  * \param required The required domain to be covered
- * \param iter_doms The result iteration domains to be updated
  * \param analyzer The arithmetic analyzer
  */
-void UpdateBlockVarDomain(const arith::IntSet& provided, const arith::IntSet& required,
-                          std::unordered_map<const VarNode*, std::vector<arith::IntSet>>* iter_doms,
-                          arith::Analyzer* analyzer) {
+std::pair<Var, arith::IntSet> SolveBlockVarDomain(const arith::IntSet& provided,
+                                                  const arith::IntSet& required,
+                                                  arith::Analyzer* analyzer) {
   PrimExpr provided_min = analyzer->Simplify(provided.min());
-  PrimExpr provided_extent = analyzer->Simplify(provided.max() - provided_min + 1);
+  PrimExpr provided_max = analyzer->Simplify(provided.max());
   PrimExpr required_min = analyzer->Simplify(required.min());
-  PrimExpr required_extent = analyzer->Simplify(required.max() - required_min + 1);
-  PrimExpr dom_min{nullptr}, dom_extent{nullptr};
+  PrimExpr required_max = analyzer->Simplify(required.max());
+  PrimExpr dom_min{nullptr}, dom_max{nullptr};
   Var dom_var{ObjectPtr<VarNode>{nullptr}};
   arith::PVar<Var> p_v;
   arith::PVar<PrimExpr> p_e;
   if ((p_v * p_e).Match(provided_min) || (p_e * p_v).Match(provided_min)) {
     PrimExpr e = p_e.Eval();
     dom_var = p_v.Eval();
     dom_min = floordiv(required_min, e);
-    dom_extent = analyzer->Simplify((required_extent + e - 1) / e);
-  } else if (analyzer->CanProveEqual(provided_extent, 1) && p_v.Match(provided_min)) {
-    dom_var = p_v.Eval();
-    dom_min = required_min;
-    dom_extent = required_extent;
-  } else {
-    ICHECK(false) << "ValueError: BufferRegion pattern match failed";
+    dom_max = floordiv(required_max, e);
+  } else if (analyzer->CanProveEqual(provided_min, provided_max)) {
+    if (p_v.Match(provided_min)) {
+      dom_var = p_v.Eval();
+      dom_min = required_min;
+      dom_max = required_max;
+    } else {
+      arith::PVar<PrimExpr> p_f;
+      if ((floordiv(p_v, p_f)).Match(provided_min)) {
+        // a <= (x // factor) <= b, fac > 0 ==> (a * fac) <= x <= (b * fac + fac - 1)
+        PrimExpr fac = p_f.Eval();
+        if (analyzer->CanProveGreaterEqual(fac, 1)) {
+          dom_var = p_v.Eval();
+          dom_min = required_min * fac;
+          dom_max = analyzer->Simplify(required_max * fac + fac - 1);
+        }
+      } else if ((floormod(p_v, p_f).Match(provided_min))) {
+        // generally domain of (x % fac) enforce no constraints to domain of x
+        dom_var = p_v.Eval();
+        return std::make_pair(dom_var, arith::IntSet::Nothing());
+      }
+    }
   }
-  auto it = iter_doms->find(dom_var.get());
+  ICHECK(dom_var.defined()) << "ValueError: BufferRegion pattern match failed: " << provided_min;
+  return std::make_pair(dom_var, arith::IntSet::Interval(dom_min, dom_max));
+}
+
+/*!
+ * \brief Calculate and update the iteration domain info to fully cover the required domain
+ * \param provided The provided integer set to cover the required domain
+ * \param required The required domain to be covered
+ * \param required_bound The additional region bound of the required domain to be covered
+ * \param iter_doms The result iteration domains to be updated
+ * \param analyzer The arithmetic analyzer
+ */
+void UpdateBlockVarDomain(const arith::IntSet& provided, const arith::IntSet& required,
+                          const arith::IntSet& required_bound,
+                          std::unordered_map<const VarNode*, BlockVarDomainInfo>* iter_doms,
+                          arith::Analyzer* analyzer) {
+  auto var_with_dom = SolveBlockVarDomain(provided, required, analyzer);
+  auto var_with_bound = SolveBlockVarDomain(provided, required_bound, analyzer);
+  const Var& var = var_with_dom.first;
+  const auto& var_dom = var_with_dom.second;
+  const auto& var_bound = var_with_bound.second;
+  ICHECK(var.same_as(var_with_bound.first));
+  auto it = iter_doms->find(var.get());
   if (it != iter_doms->end()) {
-    std::vector<arith::IntSet>& doms = it->second;
-    doms.push_back(arith::IntSet::FromMinExtent(dom_min, dom_extent));
+    it->second.Union({var_dom, var_bound});
   } else {
-    ICHECK(analyzer->CanProveEqual(provided_min, required_min));
-    ICHECK(analyzer->CanProveEqual(provided_extent, required_extent));
+    ICHECK(analyzer->CanProveEqual(provided.min(), required.min()));
+    ICHECK(analyzer->CanProveEqual(provided.max(), required.max()));
   }
 }
 
@@ -352,19 +445,19 @@ void UpdateBlockVarDomain(const arith::IntSet& provided, const arith::IntSet& re
  * \param provided_regions The region provided by one iteration instance of the block vars
  * \param required_regions The region required to be covered
  * \param analyzer The arithmetic analyzer
- * \return A list of iteration domain corresponding to the given list of block vars
+ * \return A list of iteration domain info corresponding to the given list of block vars
  */
-std::vector<Range> CalculateBlockVarDomain(
+std::vector<BlockVarDomainInfo> CalculateBlockVarDomain(
     const Array<IterVar>& iter_vars,
     std::unordered_map<const BufferNode*, std::vector<NDIntSet>> provided_regions,
     std::unordered_map<const BufferNode*, std::vector<NDIntSet>> required_regions,
     arith::Analyzer* analyzer) {
   int n_iters = iter_vars.size();
   // Step 1. Construct the mapping from block var to their iteration domain (initialized to empty)
-  std::unordered_map<const VarNode*, std::vector<arith::IntSet>> iter_doms;
+  std::unordered_map<const VarNode*, BlockVarDomainInfo> iter_doms;
   iter_doms.reserve(n_iters);
   for (const IterVar& iter_var : iter_vars) {
-    iter_doms[iter_var->var.get()] = {};
+    iter_doms[iter_var->var.get()] = BlockVarDomainInfo();
   }
   // Step 2. For each buffer, update the domain according to the provided and required regions
   for (const auto& kv : provided_regions) {
@@ -384,23 +477,23 @@ std::vector<Range> CalculateBlockVarDomain(
     for (int i = 0; i < ndim; ++i) {
       arith::IntSet provided = provided_region[i];
       arith::IntSet required = required_region[i];
-      required = arith::Intersect(
-          {std::move(required), arith::IntSet::FromMinExtent(Integer(0), buffer->shape[i])});
-      UpdateBlockVarDomain(provided, required, &iter_doms, analyzer);
+      arith::IntSet required_bound = arith::IntSet::FromMinExtent(Integer(0), buffer->shape[i]);
+      UpdateBlockVarDomain(provided, required, required_bound, &iter_doms, analyzer);
     }
   }
   // Union the iter var domains, put them in the same order of block vars, and return
-  std::vector<Range> result;
+  std::vector<BlockVarDomainInfo> result;
   result.reserve(n_iters);
   for (const IterVar& iter_var : iter_vars) {
-    const std::vector<arith::IntSet>& doms = iter_doms.at(iter_var->var.get());
-    arith::IntSet dom = arith::IntSet::FromRange(iter_var->dom);
-    if (!doms.empty()) {
-      dom = arith::Intersect({std::move(dom), arith::Union(doms)});
+    BlockVarDomainInfo& info = iter_doms.at(iter_var->var.get());
+    if (info.bound.IsNothing()) {
+      info.bound = arith::IntSet::FromRange(iter_var->dom);
+    } else {
+      info.bound = arith::Intersect({info.bound, arith::IntSet::FromRange(iter_var->dom)});
     }
-    PrimExpr min = analyzer->Simplify(dom.min());
-    PrimExpr extent = analyzer->Simplify(dom.max() - min + 1);
-    result.push_back(Range::FromMinExtent(min, extent));
+    info.Simplify(analyzer);
+    ICHECK(!info.dom.IsNothing());
+    result.push_back(info);
   }
   return result;
 }
@@ -498,14 +591,14 @@ void ComputeAtOrReverseComputeAtImpl(ScheduleState self, const StmtSRef& block_s
       /*consumer_srefs=*/std::move(consumer_srefs),
       /*provided_regions=*/&provided_regions, /*required_regions=*/&required_regions);
   // Step 5. Calculate the iteration domain for each block var
-  std::vector<Range> iter_doms =
+  std::vector<BlockVarDomainInfo> iter_doms =
       CalculateBlockVarDomain(/*iter_vars=*/block->iter_vars,
                               /*provided_regions=*/std::move(provided_regions),
                               /*required_regions=*/std::move(required_regions),
                               /*analyzer=*/analyzer);
   // Step 6. Create the new scope according to the iteration domain
   reconstructor.MakeNewLoop(/*insert_position=*/insert_position, /*iter_doms=*/std::move(iter_doms),
-                            /*preserve_unit_loops=*/preserve_unit_loops);
+                            /*analyzer=*/analyzer, /*preserve_unit_loops=*/preserve_unit_loops);
   Block new_scope_root = Downcast<Block>(reconstructor(scope_root));
 
   // Step 7. Do the actual replacement

tests/python/unittest/test_arith_iter_affine_map.py

@@ -263,6 +263,14 @@ def test_predicate():
     )
     assert len(res) == 0
 
+    # irrelavant predicate
+    res = tvm.arith.detect_iter_map(
+        [i + j],
+        var_dom([(i, 1)]),
+        j <= 24,
+    )
+    assert_iter_sum_pattern(res[0], 1, j)
+
     # constraint on nested fused iters
     res = tvm.arith.detect_iter_map(
         [i * 8 + j * 2 + k],

tests/python/unittest/test_tir_analysis_get_block_access_region.py

@@ -138,29 +138,15 @@ def access_of_padding_pattern() -> None:
     for i, j in T.grid(32, 32):
         with T.block("padding"):
             vi, vj = T.axis.remap("SS", [i, j])
-            T.reads(
-                [
-                    X[
-                        T.max(vi - 2, 0) : T.min(vi - 2, 27) + 1,
-                        T.max(vj - 2, 0) : T.min(vj - 2, 27) + 1,
-                    ]
-                ]
-            )
+            T.reads([X[vi - 2, vj - 2]])
             T.writes([X_pad[vi, vj]])
             X_pad[vi, vj] = T.if_then_else(
                 2 <= vi and vi < 30 and 2 <= vj and vj < 30, X[vi - 2, vj - 2], 0.0, dtype="float32"
             )
         with T.block("padding_reverse"):
             vi, vj = T.axis.remap("SS", [i, j])
-            T.reads([X_pad[T.max(vi, 2) : T.min(vi, 29) + 1, T.max(vj, 2) : T.min(vj, 29) + 1]])
-            T.writes(
-                [
-                    Y[
-                        T.max(vi - 2, 0) : T.min(vi - 2, 27) + 1,
-                        T.max(vj - 2, 0) : T.min(vj - 2, 27) + 1,
-                    ]
-                ]
-            )
+            T.reads([X_pad[vi, vj]])
+            T.writes([Y[vi - 2, vj - 2]])
             if 2 <= vi and vi < 30 and 2 <= vj and vj < 30:
                 Y[vi - 2, vj - 2] = X_pad[vi, vj]