IsAffineBinding(...)

tlc-pack · May 10, 2021 · f9f0389 · f9f0389
1 parent 8f25349
commit f9f0389
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/src/tir/schedule/analysis.h b/src/tir/schedule/analysis.h
@@ -22,6 +22,7 @@
 #include <tvm/tir/op.h>
 #include <tvm/tir/schedule/state.h>
 #include <tvm/tir/stmt_functor.h>
+#include <tvm/arith/iter_affine_map.h>
 
 namespace tvm {
 namespace tir {
@@ -60,7 +61,8 @@ bool IsReductionBlock(const ScheduleState& self, const StmtSRef& block_sref,
 
 /******** Binding ********/
 // Todo
-bool ValidateBlockBinding(const BlockRealize& realize, const Map<Var, Range>& loop_var_ranges);
+bool IsAffineBinding(const BlockRealize& realize, const Map<Var, Range>& loop_var_ranges,
+                     arith::Analyzer* analyzer);
 
 /******** Block-loop relation ********/
 /*!

diff --git a/src/tir/schedule/analysis/analysis.cc b/src/tir/schedule/analysis/analysis.cc
@@ -110,21 +110,17 @@ bool IsReductionBlock(const ScheduleState& self, const StmtSRef& block_sref,
 
 /******** Binding ********/
 
-bool ValidateBlockBinding(const BlockRealize& realize, const Map<Var, Range>& loop_var_ranges) {
+bool IsAffineBinding(const BlockRealize& realize, const Map<Var, Range>& loop_var_ranges,
+                     arith::Analyzer* analyzer) {
   if (loop_var_ranges.empty()) {
     return true;
   }
-  arith::Analyzer analyzer;
-  Array<PrimExpr> leaf_iters;
-  leaf_iters.reserve(static_cast<int64_t>(realize->block->iter_vars.size()));
-  for (const IterVar& iter_var : realize->block->iter_vars) {
-    leaf_iters.push_back(iter_var);
-  }
   Array<arith::IterSumExpr> results = arith::DetectIterMap(
-      /*leaf_iters=*/leaf_iters,
-      /*bindings=*/realize->iter_values, // Todo
-      /*root_iters=*/loop_var_ranges,
-      /*input_pred=*/realize->predicate, /*analyzer=*/&analyzer);
+      /*indices=*/realize->iter_values,
+      /*input_iters=*/loop_var_ranges,
+      /*predicate=*/realize->predicate,
+      /*require_bijective=*/false,
+      /*analyzer=*/analyzer);
   if (results.empty()) {
     return false;
   }
@@ -400,7 +396,9 @@ void UpdateAffineFlag(ScheduleState self, const StmtSRef& block_sref) {
     }
   }
   ICHECK(self->block_info.count(block_sref));
-  self->block_info[block_sref].affine_binding = ValidateBlockBinding(realize, loop_var_ranges);
+  arith::Analyzer analyzer;
+  self->block_info[block_sref].affine_binding =
+      IsAffineBinding(realize, loop_var_ranges, &analyzer);
 }
 
 /******** Pattern Matcher ********/