[GitHub] [tvm] junrushao1994 commented on a change in pull request #8114: [TensorIR][M2a] Verification of cached flags

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

junrushao1994 commented on a change in pull request #8114:
URL: https://github.com/apache/tvm/pull/8114#discussion_r638112423



##########
File path: src/tir/schedule/analysis/analysis.cc
##########
@@ -21,6 +21,79 @@
 namespace tvm {
 namespace tir {
 
+/******** Binding ********/
+
+bool IsAffineBinding(const BlockRealize& realize, const Map<Var, Range>& loop_var_ranges,
+                     arith::Analyzer* analyzer) {
+  if (loop_var_ranges.empty()) {
+    return true;
+  }
+  Array<arith::IterSumExpr> results = arith::DetectIterMap(
+      /*indices=*/realize->iter_values,
+      /*input_iters=*/loop_var_ranges,
+      /*predicate=*/realize->predicate,
+      /*require_bijective=*/false,
+      /*analyzer=*/analyzer);
+  if (results.empty()) {
+    return false;
+  }
+  for (const arith::IterSumExpr& sum_expr : results) {
+    const Array<arith::IterSplitExpr>& args = sum_expr->args;
+    if (!args.empty() && !is_one(args[0]->scale)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+Map<Var, Range> LoopDomainOfSRefTreePath(const StmtSRef& low_inclusive,
+                                         const Optional<StmtSRef>& high_exclusive,
+                                         const runtime::StorageScope& extra_relax_scope) {
+  Map<Var, Range> result;
+  const StmtSRefNode* p = low_inclusive.get();
+  const StmtSRefNode* limit = static_cast<const StmtSRefNode*>(high_exclusive.get());
+  for (; p != limit; p = p->parent) {
+    const ForNode* loop = p->StmtAs<ForNode>();
+    if (loop == nullptr) {
+      break;
+    }
+    result.Set(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
+  }
+  if (extra_relax_scope.rank != runtime::StorageRank::kGlobal) {
+    for (; p; p = p->parent) {
+      if (const ForNode* loop = p->StmtAs<ForNode>()) {
+        if (loop->kind == ForKind::kThreadBinding) {

Review comment:
       Perhaps I should just implement the logic rather than leaving a TODO for future




-- 
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