From a02513958cd4d42e7412c774fabe6b25995f5c00 Mon Sep 17 00:00:00 2001 From: silentCoder-dev Date: Wed, 7 Jan 2026 12:40:16 +0800 Subject: [PATCH 1/5] hoist constr_visitor --- src/transform/common/constr_visitor.h | 179 ++++++++++++++++++++++++++ src/transform/verify_parallel_loop.cc | 165 +----------------------- 2 files changed, 182 insertions(+), 162 deletions(-) create mode 100644 src/transform/common/constr_visitor.h diff --git a/src/transform/common/constr_visitor.h b/src/transform/common/constr_visitor.h new file mode 100644 index 0000000000..00db697eb9 --- /dev/null +++ b/src/transform/common/constr_visitor.h @@ -0,0 +1,179 @@ +#include "tvm/arith/analyzer.h" +#include "tvm/ffi/base_details.h" +#include "tvm/ffi/object.h" +#include "tvm/ir/expr.h" +#include "tvm/tir/op.h" +#include "tvm/tir/stmt.h" +#include "tvm/tir/var.h" +#include +#include +#include +#include +#include + +namespace tvm::tl { + +using namespace tir; + +struct Constr { + + enum Kind { + kConstr, + kBindValue, + kBindRange, + } kind; + bool is_assume = false; + Var var; + PrimExpr value; + Range range; + + Constr(PrimExpr constr, bool is_assume = false) + : kind(kConstr), value(constr), is_assume(is_assume) {}; + Constr(Var var, PrimExpr val) : kind(kBindValue), var(var), value(val) {}; + Constr(Var var, Range range) : kind(kBindRange), var(var), range(range) {}; + + Constr() = default; + Constr(const Constr &other) = default; + Constr(Constr &&other) = default; + Constr &operator=(const Constr &other) = default; + + PrimExpr ToGenericConstr() const { + switch (kind) { + case kConstr: + return value; + case kBindValue: + return var == value; + case kBindRange: + return And(var >= range->min, var < (range->min + range->extent)); + } + LOG(FATAL) << "Unreachable"; + return PrimExpr(); + } + Constr Substitute(ffi::Map subs) const { + return Constr(tir::Substitute(ToGenericConstr(), subs)); + } + void Populate(arith::Analyzer &analyzer) const { + switch (kind) { + case kConstr: + analyzer.EnterConstraint(value); + break; + case kBindValue: + analyzer.Bind(var, value); + break; + case kBindRange: + analyzer.Bind(var, range); + break; + } + } +}; + +struct ConstrSet { + ConstrSet Substitute(ffi::Map subs) const { + ConstrSet new_set; + for (const auto &c : constrs_) { + new_set.constrs_.push_back(c.Substitute(subs)); + } + return new_set; + } + void Populate(arith::Analyzer &analyzer) const { + for (const auto &c : constrs_) { + c.Populate(analyzer); + } + } + bool CanProve(const PrimExpr &expr) const { + arith::Analyzer analyzer; + Populate(analyzer); + return analyzer.CanProve(expr); + } + template void AddConstr(Args... args) { + constrs_.push_back(Constr(args...)); + } + void Extend(const ConstrSet &other) { + for (const auto &c : other.constrs_) { + constrs_.push_back(c); + } + } + std::vector constrs_; +}; + +struct ConstrVisitor : public StmtExprVisitor { +private: + using Base = StmtExprVisitor; + struct Guard { + std::vector &constrs; + ~Guard() { constrs.pop_back(); } + }; + template Guard MakeGuard(const Args... args) { + constr_stack_.push_back(Constr(args...)); + return Guard{constr_stack_}; + } + +public: + void VisitIfThenElseExpr(const PrimExpr cond, const PrimExpr true_value, + const PrimExpr false_value) { + { + auto guard = MakeGuard(cond); + Base::VisitExpr(true_value); + } + { + auto guard = MakeGuard(Not(cond)); + Base::VisitExpr(false_value); + } + } + void VisitStmt_(const LetStmtNode *op) override { + auto guard = MakeGuard(op->var, op->value); + Base::VisitStmt_(op); + } + void VisitStmt_(const AttrStmtNode *op) override { + if (op->attr_key == tir::attr::tilelang_assume) { + auto expr = Downcast(op->node); + auto guard = MakeGuard(expr, true); + Base::VisitStmt_(op); + } else if (op->attr_key == tir::attr::thread_extent || + op->attr_key == tir::attr::virtual_thread) { + IterVar iv = Downcast(op->node); + Range dom = Range::FromMinExtent(make_zero(op->value.dtype()), op->value); + auto guard = MakeGuard(iv->var, dom); + Base::VisitStmt_(op); + } else { + Base::VisitStmt_(op); + } + } + void VisitStmt_(const AssertStmtNode *op) override { + auto guard = MakeGuard(op->condition); + Base::VisitStmt_(op); + } + void VisitStmt_(const IfThenElseNode *op) override { + { + auto guard = MakeGuard(op->condition); + Base::VisitStmt(op->then_case); + } + if (op->else_case) { + auto guard = MakeGuard(Not(op->condition)); + Base::VisitStmt(op->else_case.value()); + } + } + void VisitExpr_(const SelectNode *op) override { + VisitIfThenElseExpr(op->condition, op->true_value, op->false_value); + } + void VisitExpr_(const CallNode *op) override { + static auto op_if_then_else = Op::Get("tir.if_then_else"); + if (op->op.same_as(op_if_then_else)) { + VisitIfThenElseExpr(op->args[0], op->args[1], op->args[2]); + } else { + Base::VisitExpr_(op); + } + } + void VisitStmt_(const ForNode *op) override { + if (op->kind == ForKind::kParallel || op->kind == ForKind::kVectorized) { + auto guard_1 = + MakeGuard(op->loop_var, Range::FromMinExtent(op->min, op->extent)); + auto guard_2 = MakeGuard(op->extent > 0); + Base::VisitStmt_(op); + } else { + Base::VisitStmt_(op); + } + } + std::vector constr_stack_; +}; +} \ No newline at end of file diff --git a/src/transform/verify_parallel_loop.cc b/src/transform/verify_parallel_loop.cc index 8cf2e48876..672ffa07c0 100644 --- a/src/transform/verify_parallel_loop.cc +++ b/src/transform/verify_parallel_loop.cc @@ -1,3 +1,4 @@ +#include "common/constr_visitor.h" #include "layout_reducer.h" #include "tvm/arith/analyzer.h" #include "tvm/ffi/base_details.h" @@ -17,168 +18,8 @@ namespace tvm::tl { using namespace tir; namespace { - -struct Constr { - - enum Kind { - kConstr, - kBindValue, - kBindRange, - } kind; - bool is_assume = false; - Var var; - PrimExpr value; - Range range; - - Constr(PrimExpr constr, bool is_assume = false) - : kind(kConstr), value(constr), is_assume(is_assume) {}; - Constr(Var var, PrimExpr val) : kind(kBindValue), var(var), value(val) {}; - Constr(Var var, Range range) : kind(kBindRange), var(var), range(range) {}; - - Constr() = default; - Constr(const Constr &other) = default; - Constr(Constr &&other) = default; - Constr &operator=(const Constr &other) = default; - - PrimExpr ToGenericConstr() const { - switch (kind) { - case kConstr: - return value; - case kBindValue: - return var == value; - case kBindRange: - return And(var >= range->min, var < (range->min + range->extent)); - } - LOG(FATAL) << "Unreachable"; - return PrimExpr(); - } - Constr Substitute(ffi::Map subs) const { - return Constr(tir::Substitute(ToGenericConstr(), subs)); - } - void Populate(arith::Analyzer &analyzer) const { - switch (kind) { - case kConstr: - analyzer.EnterConstraint(value); - break; - case kBindValue: - analyzer.Bind(var, value); - break; - case kBindRange: - analyzer.Bind(var, range); - break; - } - } -}; - -struct ConstrSet { - ConstrSet Substitute(ffi::Map subs) const { - ConstrSet new_set; - for (const auto &c : constrs_) { - new_set.constrs_.push_back(c.Substitute(subs)); - } - return new_set; - } - void Populate(arith::Analyzer &analyzer) const { - for (const auto &c : constrs_) { - c.Populate(analyzer); - } - } - bool CanProve(const PrimExpr &expr) const { - arith::Analyzer analyzer; - Populate(analyzer); - return analyzer.CanProve(expr); - } - template void AddConstr(Args... args) { - constrs_.push_back(Constr(args...)); - } - void Extend(const ConstrSet &other) { - for (const auto &c : other.constrs_) { - constrs_.push_back(c); - } - } - std::vector constrs_; -}; - -struct ConstrVisitor : public StmtExprVisitor { -private: - using Base = StmtExprVisitor; - struct Guard { - std::vector &constrs; - ~Guard() { constrs.pop_back(); } - }; - template Guard MakeGuard(const Args... args) { - constr_stack_.push_back(Constr(args...)); - return Guard{constr_stack_}; - } - -public: - void VisitIfThenElseExpr(const PrimExpr cond, const PrimExpr true_value, - const PrimExpr false_value) { - { - auto guard = MakeGuard(cond); - Base::VisitExpr(true_value); - } - { - auto guard = MakeGuard(Not(cond)); - Base::VisitExpr(false_value); - } - } - void VisitStmt_(const LetStmtNode *op) override { - auto guard = MakeGuard(op->var, op->value); - Base::VisitStmt_(op); - } - void VisitStmt_(const AttrStmtNode *op) override { - if (op->attr_key == tir::attr::tilelang_assume) { - auto expr = Downcast(op->node); - auto guard = MakeGuard(expr, true); - Base::VisitStmt_(op); - } else if (op->attr_key == tir::attr::thread_extent || - op->attr_key == tir::attr::virtual_thread) { - IterVar iv = Downcast(op->node); - Range dom = Range::FromMinExtent(make_zero(op->value.dtype()), op->value); - auto guard = MakeGuard(iv->var, dom); - Base::VisitStmt_(op); - } else { - Base::VisitStmt_(op); - } - } - void VisitStmt_(const AssertStmtNode *op) override { - auto guard = MakeGuard(op->condition); - Base::VisitStmt_(op); - } - void VisitStmt_(const IfThenElseNode *op) override { - { - auto guard = MakeGuard(op->condition); - Base::VisitStmt(op->then_case); - } - if (op->else_case) { - auto guard = MakeGuard(Not(op->condition)); - Base::VisitStmt(op->else_case.value()); - } - } - void VisitExpr_(const SelectNode *op) override { - VisitIfThenElseExpr(op->condition, op->true_value, op->false_value); - } - void VisitExpr_(const CallNode *op) override { - static auto op_if_then_else = Op::Get("tir.if_then_else"); - if (op->op.same_as(op_if_then_else)) { - VisitIfThenElseExpr(op->args[0], op->args[1], op->args[2]); - } else { - Base::VisitExpr_(op); - } - } - void VisitStmt_(const ForNode *op) override { - if (op->kind == ForKind::kParallel || op->kind == ForKind::kVectorized) { - auto guard_1 = - MakeGuard(op->loop_var, Range::FromMinExtent(op->min, op->extent)); - auto guard_2 = MakeGuard(op->extent > 0); - Base::VisitStmt_(op); - } else { - Base::VisitStmt_(op); - } - } - std::vector constr_stack_; -}; +using tvm::tl::ConstrSet; +using tvm::tl::ConstrVisitor; struct ParallelLoopVerifier : public ConstrVisitor { std::vector parallel_loop_vars_; From 58491660fadabae9209203df54e62909765680e6 Mon Sep 17 00:00:00 2001 From: silentCoder-dev Date: Wed, 7 Jan 2026 12:44:38 +0800 Subject: [PATCH 2/5] hoist constr_visitor --- src/transform/common/constr_visitor.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transform/common/constr_visitor.h b/src/transform/common/constr_visitor.h index 00db697eb9..79d2f72e5b 100644 --- a/src/transform/common/constr_visitor.h +++ b/src/transform/common/constr_visitor.h @@ -176,4 +176,4 @@ struct ConstrVisitor : public StmtExprVisitor { } std::vector constr_stack_; }; -} \ No newline at end of file +} // namespace tvm::tl From f89e7de7e80dc222b209a76965dcec587ce752ca Mon Sep 17 00:00:00 2001 From: silentCoder-dev Date: Wed, 7 Jan 2026 12:51:32 +0800 Subject: [PATCH 3/5] remove "using namespace tir;" --- src/transform/common/constr_visitor.h | 46 ++++++++++++++------------- 1 file changed, 24 insertions(+), 22 deletions(-) diff --git a/src/transform/common/constr_visitor.h b/src/transform/common/constr_visitor.h index 79d2f72e5b..7f49bd8575 100644 --- a/src/transform/common/constr_visitor.h +++ b/src/transform/common/constr_visitor.h @@ -13,8 +13,6 @@ namespace tvm::tl { -using namespace tir; - struct Constr { enum Kind { @@ -23,14 +21,16 @@ struct Constr { kBindRange, } kind; bool is_assume = false; - Var var; + tir::Var var; PrimExpr value; Range range; Constr(PrimExpr constr, bool is_assume = false) : kind(kConstr), value(constr), is_assume(is_assume) {}; - Constr(Var var, PrimExpr val) : kind(kBindValue), var(var), value(val) {}; - Constr(Var var, Range range) : kind(kBindRange), var(var), range(range) {}; + Constr(tir::Var var, PrimExpr val) + : kind(kBindValue), var(var), value(val) {}; + Constr(tir::Var var, Range range) + : kind(kBindRange), var(var), range(range) {}; Constr() = default; Constr(const Constr &other) = default; @@ -44,12 +44,12 @@ struct Constr { case kBindValue: return var == value; case kBindRange: - return And(var >= range->min, var < (range->min + range->extent)); + return tir::And(var >= range->min, var < (range->min + range->extent)); } LOG(FATAL) << "Unreachable"; return PrimExpr(); } - Constr Substitute(ffi::Map subs) const { + Constr Substitute(ffi::Map subs) const { return Constr(tir::Substitute(ToGenericConstr(), subs)); } void Populate(arith::Analyzer &analyzer) const { @@ -68,7 +68,7 @@ struct Constr { }; struct ConstrSet { - ConstrSet Substitute(ffi::Map subs) const { + ConstrSet Substitute(ffi::Map subs) const { ConstrSet new_set; for (const auto &c : constrs_) { new_set.constrs_.push_back(c.Substitute(subs)); @@ -96,9 +96,9 @@ struct ConstrSet { std::vector constrs_; }; -struct ConstrVisitor : public StmtExprVisitor { +struct ConstrVisitor : public tir::StmtExprVisitor { private: - using Base = StmtExprVisitor; + using Base = tir::StmtExprVisitor; struct Guard { std::vector &constrs; ~Guard() { constrs.pop_back(); } @@ -116,47 +116,48 @@ struct ConstrVisitor : public StmtExprVisitor { Base::VisitExpr(true_value); } { - auto guard = MakeGuard(Not(cond)); + auto guard = MakeGuard(tir::Not(cond)); Base::VisitExpr(false_value); } } - void VisitStmt_(const LetStmtNode *op) override { + void VisitStmt_(const tir::LetStmtNode *op) override { auto guard = MakeGuard(op->var, op->value); Base::VisitStmt_(op); } - void VisitStmt_(const AttrStmtNode *op) override { + void VisitStmt_(const tir::AttrStmtNode *op) override { if (op->attr_key == tir::attr::tilelang_assume) { auto expr = Downcast(op->node); auto guard = MakeGuard(expr, true); Base::VisitStmt_(op); } else if (op->attr_key == tir::attr::thread_extent || op->attr_key == tir::attr::virtual_thread) { - IterVar iv = Downcast(op->node); - Range dom = Range::FromMinExtent(make_zero(op->value.dtype()), op->value); + tir::IterVar iv = Downcast(op->node); + Range dom = + Range::FromMinExtent(tir::make_zero(op->value.dtype()), op->value); auto guard = MakeGuard(iv->var, dom); Base::VisitStmt_(op); } else { Base::VisitStmt_(op); } } - void VisitStmt_(const AssertStmtNode *op) override { + void VisitStmt_(const tir::AssertStmtNode *op) override { auto guard = MakeGuard(op->condition); Base::VisitStmt_(op); } - void VisitStmt_(const IfThenElseNode *op) override { + void VisitStmt_(const tir::IfThenElseNode *op) override { { auto guard = MakeGuard(op->condition); Base::VisitStmt(op->then_case); } if (op->else_case) { - auto guard = MakeGuard(Not(op->condition)); + auto guard = MakeGuard(tir::Not(op->condition)); Base::VisitStmt(op->else_case.value()); } } - void VisitExpr_(const SelectNode *op) override { + void VisitExpr_(const tir::SelectNode *op) override { VisitIfThenElseExpr(op->condition, op->true_value, op->false_value); } - void VisitExpr_(const CallNode *op) override { + void VisitExpr_(const tir::CallNode *op) override { static auto op_if_then_else = Op::Get("tir.if_then_else"); if (op->op.same_as(op_if_then_else)) { VisitIfThenElseExpr(op->args[0], op->args[1], op->args[2]); @@ -164,8 +165,9 @@ struct ConstrVisitor : public StmtExprVisitor { Base::VisitExpr_(op); } } - void VisitStmt_(const ForNode *op) override { - if (op->kind == ForKind::kParallel || op->kind == ForKind::kVectorized) { + void VisitStmt_(const tir::ForNode *op) override { + if (op->kind == tir::ForKind::kParallel || + op->kind == tir::ForKind::kVectorized) { auto guard_1 = MakeGuard(op->loop_var, Range::FromMinExtent(op->min, op->extent)); auto guard_2 = MakeGuard(op->extent > 0); From ff0bd7cecf8727e7646e09fbe6fb914573ccfcc8 Mon Sep 17 00:00:00 2001 From: silentCoder-dev Date: Wed, 7 Jan 2026 12:59:09 +0800 Subject: [PATCH 4/5] add TVM_TL_TRANSFORM_COMMON_CONSTR_VISITOR_H_ --- src/transform/common/constr_visitor.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/transform/common/constr_visitor.h b/src/transform/common/constr_visitor.h index 7f49bd8575..b43513dd95 100644 --- a/src/transform/common/constr_visitor.h +++ b/src/transform/common/constr_visitor.h @@ -1,3 +1,6 @@ +#ifndef TVM_TL_TRANSFORM_COMMON_CONSTR_VISITOR_H_ +#define TVM_TL_TRANSFORM_COMMON_CONSTR_VISITOR_H_ + #include "tvm/arith/analyzer.h" #include "tvm/ffi/base_details.h" #include "tvm/ffi/object.h" @@ -179,3 +182,5 @@ struct ConstrVisitor : public tir::StmtExprVisitor { std::vector constr_stack_; }; } // namespace tvm::tl + +#endif // TVM_TL_TRANSFORM_COMMON_CONSTR_VISITOR_H_ From a4f2919b97ce1b12987ba44b24d4d2d298f4f6c5 Mon Sep 17 00:00:00 2001 From: silentCoder-dev Date: Wed, 7 Jan 2026 13:04:00 +0800 Subject: [PATCH 5/5] add default case --- src/transform/common/constr_visitor.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/transform/common/constr_visitor.h b/src/transform/common/constr_visitor.h index b43513dd95..f54540c4f3 100644 --- a/src/transform/common/constr_visitor.h +++ b/src/transform/common/constr_visitor.h @@ -66,6 +66,8 @@ struct Constr { case kBindRange: analyzer.Bind(var, range); break; + default: + LOG(FATAL) << "Unreachable"; } } };