[Feature] Reimplement Threadsync with ConstrVisitor

src/transform/common/constr_visitor.h

@@ -8,6 +8,7 @@
 #include "tvm/tir/op.h"
 #include "tvm/tir/stmt.h"
 #include "tvm/tir/var.h"
+#include <ostream>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/tir/analysis.h>
 #include <tvm/tir/builtin.h>
@@ -40,6 +41,31 @@ struct Constr {
   Constr(Constr &&other) = default;
   Constr &operator=(const Constr &other) = default;
 
+  void format(std::ostream &os) const {
+    os << "Constr(kind=";
+    switch (kind) {
+    case kConstr:
+      os << "kConstr";
+      os << ", is_assume=" << (is_assume ? "true" : "false");
+      os << ", value=" << value;
+      break;
+    case kBindValue:
+      os << "kBindValue";
+      os << ", var=" << var->name_hint;
+      os << ", value=" << value;
+      break;
+    case kBindRange:
+      os << "kBindRange";
+      os << ", var=" << var->name_hint;
+      os << ", range=Range(min=" << range->min;
+      os << ", extent=" << range->extent << ")";
+      break;
+    default:
+      os << "Unknown";
+    }
+    os << ")";
+  }
+
   PrimExpr ToGenericConstr() const {
     switch (kind) {
     case kConstr:
@@ -98,22 +124,38 @@ struct ConstrSet {
       constrs_.push_back(c);
     }
   }
+
+  void format(std::ostream &os) const {
+    os << "ConstrSet(size=" << constrs_.size() << ") {\n";
+    for (size_t i = 0; i < constrs_.size(); ++i) {
+      os << "  [" << i << "] ";
+      constrs_[i].format(os);
+      os << "\n";
+    }
+    os << "}";
+  }
+
   std::vector<Constr> constrs_;
 };
 
 struct ConstrVisitor : public tir::StmtExprVisitor {
 private:
   using Base = tir::StmtExprVisitor;
+
   struct Guard {
     std::vector<Constr> &constrs;
     ~Guard() { constrs.pop_back(); }
   };
+
+protected:
   template <typename... Args> Guard MakeGuard(const Args... args) {
     constr_stack_.push_back(Constr(args...));
     return Guard{constr_stack_};
   }
 
 public:
+  using StmtExprVisitor::VisitExpr_;
+  using StmtExprVisitor::VisitStmt_;
   void VisitIfThenElseExpr(const PrimExpr cond, const PrimExpr true_value,
                            const PrimExpr false_value) {
     {
@@ -178,9 +220,18 @@ struct ConstrVisitor : public tir::StmtExprVisitor {
       auto guard_2 = MakeGuard(op->extent > 0);
       Base::VisitStmt_(op);
     } else {
+      auto guard_1 =
+          MakeGuard(op->loop_var, Range::FromMinExtent(op->min, op->extent));
+      auto guard_2 = MakeGuard(op->extent > 0);
       Base::VisitStmt_(op);
     }
   }
+  void VisitStmt_(const tir::WhileNode *op) override {
+    {
+      auto guard = MakeGuard(op->condition);
+      Base::VisitStmt(op->body);
+    }
+  }
   std::vector<Constr> constr_stack_;
 };
 } // namespace tvm::tl

src/transform/eliminate_storage_sync_for_mbarrier.cc

@@ -2,7 +2,6 @@
  * \file eliminate_storage_sync_for_mbarrier.cc
  */
 #include "../op/builtin.h"
-#include "./storage_access.h"
 #include "arith/ir_mutator_with_analyzer.h"
 #include "arith/ir_visitor_with_analyzer.h"
 #include <tvm/ffi/function.h>

src/transform/inject_ptx_async_copy.cc

@@ -29,7 +29,6 @@
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
-#include "storage_access.h"
 #include "tir/ir/buffer_common.h"
 #include "tvm/tir/stmt.h"