This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f392ee50d7 [REFACTOR] Use local verifier scopes and exception visit
context (#20684)
f392ee50d7 is described below
commit f392ee50d7adf8ab82e98cea9ff029f6e758ff6b
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 19:08:33 2026 -0400
[REFACTOR] Use local verifier scopes and exception visit context (#20684)
Migrate TIRx and S-TIR well-formedness, purity analysis, and remaining
tile-op checks to ordinary visitors. Keep definition scopes and earlier
definition locations within well-formedness checking, and resolve
current diagnostic paths from exception visit context at verifier
entrypoints.
Preserve each analysis's traversal and binding policy while removing the
private path visitor implementations.
---
src/s_tir/analysis/is_pure_function.cc | 159 ++++++-
src/s_tir/analysis/verify_well_formed.cc | 85 +++-
src/s_tir/ir/tir_visitor_with_path.cc | 84 ----
src/s_tir/ir/tir_visitor_with_path.h | 54 ---
src/tirx/analysis/verify_well_formed.cc | 5 +-
src/tirx/analysis/verify_well_formed.h | 711 +++++++++++++++++++++----------
src/tirx/ir/tir_visitor_with_path.cc | 355 ---------------
src/tirx/ir/tir_visitor_with_path.h | 403 ------------------
src/tirx/tile/lower.cc | 114 ++++-
9 files changed, 792 insertions(+), 1178 deletions(-)
diff --git a/src/s_tir/analysis/is_pure_function.cc
b/src/s_tir/analysis/is_pure_function.cc
index f5d6f7d842..bea01754b8 100644
--- a/src/s_tir/analysis/is_pure_function.cc
+++ b/src/s_tir/analysis/is_pure_function.cc
@@ -31,41 +31,59 @@
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op/memory.h>
-#include "../ir/tir_visitor_with_path.h"
-
namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
namespace {
-class PurityChecker : TIRVisitorWithPath {
+class PurityChecker : public StmtExprVisitor {
+ friend class tvm::ObjectVisitor;
+ using Result = ffi::Optional<VisitInterrupt>;
+
public:
static bool Check(const Function& func, bool assert_on_error) {
- PurityChecker visitor(assert_on_error);
- visitor(func);
- return visitor.is_pure_;
+ auto visitor = ffi::make_object<PurityChecker>(assert_on_error);
+ visitor->Visit(func);
+ return visitor->is_pure_;
}
+ explicit PurityChecker(bool assert_on_error)
+ : StmtExprVisitor(GlobalVTable()), assert_on_error_(assert_on_error) {}
+
private:
- explicit PurityChecker(bool assert_on_error) :
assert_on_error_(assert_on_error) {}
+ using StmtExprVisitor::Visit_;
+ Result Visit(ffi::AnyView value) override {
+ bool is_root = !entered_;
+ entered_ = true;
+ if (auto function = value.as<Function>(); function && !is_root) {
+ TVM_FFI_VISIT_THROW(TypeError, function.value())
+ << "Unsupported non-primitive TIR expression " <<
function.value()->GetTypeKey();
+ }
+ if (auto stmt = value.as<Stmt>(); stmt &&
!GlobalVTable()->CanDispatch(value)) {
+ TVM_FFI_VISIT_THROW(InternalError, stmt.value())
+ << "Do not have a default for " << stmt.value()->GetTypeKey();
+ }
+ if (auto expr = value.as<Expr>(); expr &&
!GlobalVTable()->CanDispatch(value)) {
+ if (value.as<OpaqueExprNode>() || value.as<VarNode>()) return
std::nullopt;
+ TVM_FFI_VISIT_THROW(TypeError, expr.value())
+ << "Unsupported non-primitive TIR expression " <<
expr.value()->GetTypeKey();
+ }
+ return StmtExprVisitor::Visit(value);
+ }
- void Dispatch_(const BindNode* op, ffi::reflection::AccessPath path) final {
+ Result Visit_(const BindNode* op) final {
if (const auto* call = op->value.as<CallNode>();
call && call->op.same_as(tirx::alloc_tensor_op())) {
- return DispatchAllocTensor(op, call, path);
+ internal_allocations_.insert(op->var);
+ allocation_calls_.insert(call);
}
- return TIRVisitorWithPath::Dispatch_(op, path);
+ this->Visit(op->value);
+ return this->Visit(op->var->ty);
}
- void DispatchAllocTensor(const BindNode* op, const CallNode* call,
- ffi::reflection::AccessPath path) {
- internal_allocations_.insert(op->var);
- allocation_calls_.insert(call);
- TIRVisitorWithPath::Dispatch_(op, path);
- }
-
- void Dispatch_(const TensorStoreNode* op, ffi::reflection::AccessPath path)
override {
- TIRVisitorWithPath::Dispatch_(op, path);
+ Result Visit_(const TensorStoreNode* op) override {
+ this->Visit(op->value);
+ this->Visit(op->indices);
if (!internal_allocations_.count(op->dest.as_or_throw<TensorVar>().var()))
{
is_pure_ = false;
@@ -76,11 +94,13 @@ class PurityChecker : TIRVisitorWithPath {
<< op->indices << " of value " << op->value;
}
}
+ return std::nullopt;
}
- void Dispatch_(const CallNode* call, ffi::reflection::AccessPath path)
override {
- TIRVisitorWithPath::Dispatch_(call, path);
- if (allocation_calls_.count(call)) return;
+ Result Visit_(const CallNode* call) override {
+ if (call->op.as<GlobalVar>() || call->op.as<OpaqueExprNode>())
this->Visit(call->op);
+ this->Visit(call->args);
+ if (allocation_calls_.count(call)) return std::nullopt;
static auto op_call_effect =
Op::GetAttrMap<TCallEffectKind>(tvm::op_attr::kCallEffectKind);
CallEffectKind effect = [&]() {
@@ -100,8 +120,103 @@ class PurityChecker : TIRVisitorWithPath {
<< ", which has side effect " << effect;
}
}
+ return std::nullopt;
+ }
+
+ // Match the analysis' traversal policy without introducing binder state.
+ Result Visit_(const FunctionNode* op) {
+ for (const auto& param : op->params) {
+ if (!param->ty.as<TensorTypeNode>()) this->Visit(param->ty);
+ }
+ for (const auto& param : op->params) {
+ if (param->ty.as<TensorTypeNode>()) this->Visit(param->ty);
+ }
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const LambdaExprNode* op) {
+ this->Visit(op->ty);
+ for (const auto& var : op->vars) this->Visit(var->ty);
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const RegionStmtNode* op) override {
+ this->Visit(op->args);
+ this->Visit(op->attrs);
+ for (const auto& param : op->body_params) this->Visit(param->ty);
+ this->Visit(op->body);
+ for (const auto& var : op->result_vars) this->Visit(var->ty);
+ return std::nullopt;
+ }
+
+ Result Visit_(const ForNode* op) override {
+ this->Visit(op->min);
+ this->Visit(op->extent);
+ this->Visit(op->loop_var->ty);
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const prim::LetNode* op) override {
+ this->Visit(op->value);
+ this->Visit(op->var->ty);
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const TensorLoadNode* op) override { return
this->Visit(op->indices); }
+
+ Result Visit_(const TensorRegionNode* op) override {
+ if (!op->source.as<TensorVar>()) this->Visit(op->source);
+ return this->Visit(op->region);
+ }
+
+ Result Visit_(const GenericConstNode* op) override { return std::nullopt; }
+
+ Result Visit_(const AssertStmtNode* op) override {
+ this->Visit(op->condition);
+ this->Visit(op->error_kind);
+ return this->Visit(op->message_parts);
+ }
+
+ Result Visit_(const prim::RampNode* op) override {
+ this->Visit(op->base);
+ this->Visit(op->stride);
+ return this->Visit(op->lanes);
+ }
+
+ Result Visit_(const prim::BroadcastNode* op) override {
+ this->Visit(op->value);
+ return this->Visit(op->lanes);
+ }
+
+ Result Visit_(const SBlockNode* op) override {
+ for (const auto& axis : op->iter_vars) {
+ this->Visit(axis->dom);
+ this->Visit(axis->var.as_or_throw<Var>()->ty);
+ }
+ for (const auto& tensor : op->alloc_tensors) this->Visit(tensor.var()->ty);
+ for (const auto& match : op->match_tensors) {
+ this->Visit(match->source);
+ this->Visit(match->tensor.var()->ty);
+ }
+ this->Visit(op->reads);
+ this->Visit(op->writes);
+ this->Visit(op->init);
+ return this->Visit(op->body);
+ }
+
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ StmtExprVisitor::InitVTable(&table);
+ SetDispatch<PurityChecker, FunctionNode>(&table);
+ SetDispatch<PurityChecker, LambdaExprNode>(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
}
+ bool entered_{false};
bool assert_on_error_{false};
bool is_pure_{true};
std::unordered_set<Var> internal_allocations_;
diff --git a/src/s_tir/analysis/verify_well_formed.cc
b/src/s_tir/analysis/verify_well_formed.cc
index d37e0c8134..c1f705fbb9 100644
--- a/src/s_tir/analysis/verify_well_formed.cc
+++ b/src/s_tir/analysis/verify_well_formed.cc
@@ -23,8 +23,6 @@
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt_functor.h>
-#include "../ir/tir_visitor_with_path.h"
-
namespace tvm {
namespace s_tir {
using tirx::Function;
@@ -136,42 +134,87 @@ class BlockVarAccessVerifier : public StmtExprVisitor {
bool has_error_{false};
};
-/*! \brief Validate the scheduling annotation before thread-binding lowering.
*/
-class ThreadBindingVerifier : public tirx::Verifier<ThreadBindingVerifier,
TIRVisitorWithPath> {
- using Verifier = tirx::Verifier<ThreadBindingVerifier, TIRVisitorWithPath>;
+// S-TIR definitions and scopes are part of this verifier, not its visitor
base.
+class WellFormednessVerifier : public
tirx::WellFormednessVerifier<StmtExprVisitor> {
+ using Parent = tirx::WellFormednessVerifier<StmtExprVisitor>;
public:
- using Verifier::Verifier;
- using Verifier::Verify;
+ WellFormednessVerifier(Check check, bool assert_mode)
+ : Parent(GlobalVTable(), check, assert_mode) {}
+
+ protected:
+ using Parent::Visit_;
+ Result Visit_(const ForNode* op) override {
+ if (check_ == Check::kThreadBinding) {
+ if (auto tag = op->annotations.Get(attr::kThreadBinding)) {
+ Verify(op->kind == ForKind::kParallel)
+ << "ValueError: thread_binding requires a parallel loop at
<current location>";
+ Verify(tag->as<ffi::String>().has_value())
+ << "TypeError: thread_binding annotation must be a string at
<current location>";
+ Verify(op->HasTrivialStep())
+ << "ValueError: Thread binding loops require a unit step at
<current location>";
+ }
+ }
+ return Parent::Visit_(op);
+ }
+
+ Result Visit_(const SBlockNode* op) override {
+ return WithScope([&]() -> Result {
+ for (size_t i = 0; i < op->iter_vars.size(); ++i) {
+ const auto& axis = op->iter_vars[i];
+ this->Visit(axis->dom);
+ Define(axis->var.as_or_throw<Var>(),
+
AccessPath::Root()->Attr("iter_vars")->ArrayItem(i)->Attr("var"));
+ }
+ for (size_t i = 0; i < op->alloc_tensors.size(); ++i) {
+ Define(op->alloc_tensors[i].var(),
AccessPath::Root()->Attr("alloc_tensors")->ArrayItem(i));
+ }
+ for (size_t i = 0; i < op->match_tensors.size(); ++i) {
+ const auto& match = op->match_tensors[i];
+ this->Visit(match->source);
+ auto site =
AccessPath::Root()->Attr("match_tensors")->ArrayItem(i)->Attr("tensor");
+ MatchTensorDefinitions(match->tensor, site->Attr("ty"));
+ Define(match->tensor.var(), site);
+ }
+ this->Visit(op->reads);
+ this->Visit(op->writes);
+ WithScope([&]() { return this->Visit(op->init); });
+ return WithScope([&]() { return this->Visit(op->body); });
+ });
+ }
private:
- void Dispatch_(const ForNode* op, AccessPath path) override {
- if (auto tag = op->annotations.Get(attr::kThreadBinding)) {
- Verify(op->kind == ForKind::kParallel)
- << "ValueError: thread_binding requires a parallel loop at " << path;
- Verify(tag->as<ffi::String>().has_value())
- << "TypeError: thread_binding annotation must be a string at " <<
path;
- Verify(op->HasTrivialStep())
- << "ValueError: Thread binding loops require a unit step at " <<
path;
- }
- Verifier::Dispatch_(op, path);
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ Parent::InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
}
};
+bool VerifyThreadBindings(const ffi::ObjectRef& node, bool assert_mode) {
+ return
ffi::make_object<WellFormednessVerifier>(WellFormednessVerifier::Check::kThreadBinding,
+ assert_mode)
+ ->Run(node);
+}
+
bool VerifyWellFormed(const tirx::Function& func, bool assert_mode) {
- return ThreadBindingVerifier::Verify(func, assert_mode) &&
+ return VerifyThreadBindings(func, assert_mode) &&
BlockVarAccessVerifier::Verify(func, assert_mode) &&
- tirx::VerifyWellFormedCommon<TIRVisitorWithPath>(func, assert_mode);
+ tirx::VerifyWellFormedCommon<WellFormednessVerifier>(func,
assert_mode);
}
bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
- if (!ThreadBindingVerifier::Verify(mod, assert_mode)) return false;
+ if (!VerifyThreadBindings(mod, assert_mode)) return false;
for (const auto& [gvar, base_func] : mod->functions) {
if (auto func = base_func.as<tirx::Function>()) {
if (!BlockVarAccessVerifier::Verify(func.value(), assert_mode)) return
false;
}
}
- return tirx::VerifyWellFormedCommon<TIRVisitorWithPath>(mod, assert_mode);
+ return tirx::VerifyWellFormedCommon<WellFormednessVerifier>(mod,
assert_mode);
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/s_tir/ir/tir_visitor_with_path.cc
b/src/s_tir/ir/tir_visitor_with_path.cc
deleted file mode 100644
index 89005fd097..0000000000
--- a/src/s_tir/ir/tir_visitor_with_path.cc
+++ /dev/null
@@ -1,84 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#include "tir_visitor_with_path.h"
-
-#include <tvm/s_tir/stmt.h>
-
-namespace tvm {
-namespace s_tir {
-using namespace tirx;
-using AccessPath = ffi::reflection::AccessPath;
-
-void TIRVisitorWithPath::Dispatch_(const SBlockNode* op, AccessPath path) {
- std::vector<DefContext<Var>> context;
-
- {
- auto iter_path = path->Attr("iter_vars");
- for (size_t i = 0; i < op->iter_vars.size(); i++) {
- const IterVar& axis = op->iter_vars[i];
- auto axis_path = iter_path->ArrayItem(i);
- if (axis->dom.has_value()) Visit(axis->dom.value(),
axis_path->Attr("dom"));
- context.push_back(WithDef(axis->var.as_or_throw<Var>(),
axis_path->Attr("var")));
- }
- }
-
- // Define alloc_tensors before visiting reads/writes, since reads/writes
- // may reference tensors from alloc_tensors (e.g. after transform_layout).
- {
- auto alloc_path = path->Attr("alloc_tensors");
- for (size_t i = 0; i < op->alloc_tensors.size(); i++) {
- auto tensor_path = alloc_path->ArrayItem(i);
- auto tensor = op->alloc_tensors[i];
- context.push_back(WithDef(tensor.var(), tensor_path));
- }
- }
-
- {
- auto match_path = path->Attr("match_tensors");
- for (size_t i = 0; i < op->match_tensors.size(); i++) {
- Visit(op->match_tensors[i]->source,
match_path->ArrayItem(i)->Attr("source"));
- auto tensor = op->match_tensors[i]->tensor;
- auto tensor_path = match_path->ArrayItem(i)->Attr("tensor");
-
- for (auto& def : WithMatchTensorDefs(tensor, tensor_path->Attr("ty"))) {
- context.push_back(std::move(def));
- }
- context.push_back(WithDef(tensor.var(), tensor_path));
- }
- }
-
- // Regions may use allocation and match-tensor definitions in this block.
- Visit(op->reads, path->Attr("reads"));
- Visit(op->writes, path->Attr("writes"));
-
- bind_scope_.WithNewScope([&]() { Visit(op->init, path->Attr("init")); });
- bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); });
-
- while (context.size()) context.pop_back();
-}
-
-void TIRVisitorWithPath::Dispatch_(const SBlockRealizeNode* op, AccessPath
path) {
- Visit(op->iter_values, path->Attr("iter_values"));
- Visit(op->predicate, path->Attr("predicate"));
- Visit(op->block, path->Attr("block"));
-}
-
-} // namespace s_tir
-} // namespace tvm
diff --git a/src/s_tir/ir/tir_visitor_with_path.h
b/src/s_tir/ir/tir_visitor_with_path.h
deleted file mode 100644
index 6556d2eec7..0000000000
--- a/src/s_tir/ir/tir_visitor_with_path.h
+++ /dev/null
@@ -1,54 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-#ifndef TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_
-#define TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_
-#include <tvm/s_tir/stmt.h>
-
-#include "../../tirx/ir/tir_visitor_with_path.h"
-namespace tvm {
-namespace s_tir {
-class TIRVisitorWithPath : public tirx::TIRVisitorWithPath {
- public:
- using Parent = tirx::TIRVisitorWithPath;
- TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(TIRVisitorWithPath, Parent)
- protected:
- using AccessPath = ffi::reflection::AccessPath;
- using Parent::Dispatch_;
- virtual void Dispatch_(const SBlockNode* op, AccessPath path);
- virtual void Dispatch_(const SBlockRealizeNode* op, AccessPath path);
- static void InitVTable(VTable* vtable) {
- Parent::InitVTable(vtable);
- // StmtVisitor is a protected base of the path visitor, so create the
- // downcast thunks here instead of in StmtVisitor::SetDispatch.
- vtable->SetDispatch<SBlockNode>(
- [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) {
- static_cast<TIRVisitorWithPath*>(self)->Dispatch_(
- static_cast<const SBlockNode*>(node.get()), path);
- });
- vtable->SetDispatch<SBlockRealizeNode>(
- [](const ffi::ObjectRef& node, StmtVisitor* self, AccessPath path) {
- static_cast<TIRVisitorWithPath*>(self)->Dispatch_(
- static_cast<const SBlockRealizeNode*>(node.get()), path);
- });
- }
-};
-} // namespace s_tir
-} // namespace tvm
-#endif // TVM_S_TIR_IR_TIR_VISITOR_WITH_PATH_H_
diff --git a/src/tirx/analysis/verify_well_formed.cc
b/src/tirx/analysis/verify_well_formed.cc
index 2ce62ba761..ef3d193325 100644
--- a/src/tirx/analysis/verify_well_formed.cc
+++ b/src/tirx/analysis/verify_well_formed.cc
@@ -35,18 +35,17 @@
#include <tuple>
#include <variant>
-#include "../ir/tir_visitor_with_path.h"
#include "tvm/ir/module.h"
namespace tvm {
namespace tirx {
bool VerifyWellFormed(const Function& func, bool assert_mode) {
- return VerifyWellFormedCommon<TIRVisitorWithPath>(func, assert_mode);
+ return VerifyWellFormedCommon<WellFormednessVerifier<StmtExprVisitor>>(func,
assert_mode);
}
bool VerifyWellFormed(const IRModule& mod, bool assert_mode) {
- return VerifyWellFormedCommon<TIRVisitorWithPath>(mod, assert_mode);
+ return VerifyWellFormedCommon<WellFormednessVerifier<StmtExprVisitor>>(mod,
assert_mode);
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/tirx/analysis/verify_well_formed.h
b/src/tirx/analysis/verify_well_formed.h
index aaae0bdabd..763d74c545 100644
--- a/src/tirx/analysis/verify_well_formed.h
+++ b/src/tirx/analysis/verify_well_formed.h
@@ -20,299 +20,556 @@
#ifndef TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_
#define TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_
-#include <tvm/ir/prim/op.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/ffi/extra/visit_error_context.h>
+#include <tvm/ir/function.h>
+#include <tvm/ir/module.h>
+#include <tvm/ir/scope_stack.h>
+#include <tvm/runtime/logging.h>
#include <tvm/tirx/analysis.h>
-#include <tvm/tirx/op/region.h>
-#include <tvm/tirx/op_attr_types.h>
-#include <tvm/tirx/tile/exec_scope.h>
+#include <tvm/tirx/stmt_functor.h>
+#include <algorithm>
+#include <optional>
+#include <sstream>
+#include <unordered_map>
#include <utility>
+#include <vector>
-#include "../ir/tir_visitor_with_path.h"
namespace tvm {
namespace tirx {
-using AccessPath = ffi::reflection::AccessPath;
-template <typename PathVisitor>
-class UndefinedVarVerifier : public
Verifier<UndefinedVarVerifier<PathVisitor>, PathVisitor> {
- using Verifier = tirx::Verifier<UndefinedVarVerifier<PathVisitor>,
PathVisitor>;
+// This state and traversal belong to well-formedness checking. In particular,
+// ordinary visitors do not impose a definition/use or lexical-scope policy.
+template <typename Visitor>
+class WellFormednessVerifier : public Visitor {
+ friend class tvm::ObjectVisitor;
public:
- // Until templated-this arrives in C++23, the CRTP can't inject a
- // constructor into the child class. Therefore, must explicitly add
- // the constructor.
- using Verifier::Verifier;
- using Verifier::Verify;
-
- private:
- using Verifier::Visit;
- void EnterDef(const Var& var, AccessPath path) override {
- {
- auto it = currently_defined_.find(var);
- auto verify = Verify(it == currently_defined_.end());
- verify << "ValueError: "
- << "TIR is ill-formed, "
- << "due to multiple nested definitions of variable " << var->name
<< ".";
- if (it != currently_defined_.end()) {
- verify << " It was first defined at " << it->second << ", and was
re-defined at " << path;
+ enum class Check { kVariables, kTensors, kTensorLoadType, kLoopControl,
kThreadBinding };
+ using AccessPath = ffi::reflection::AccessPath;
+ using Result = ffi::Optional<VisitInterrupt>;
+
+ WellFormednessVerifier(Check check, bool assert_mode)
+ : WellFormednessVerifier(GlobalVTable(), check, assert_mode) {}
+
+ bool Run(const ffi::ObjectRef& root) {
+ root_ = root;
+ try {
+ this->Visit(root);
+ } catch (const ffi::Error& error) {
+ // Resolve at the verifier boundary for both direct and pass-driven
calls.
+ // Once reported, consume the visit context to avoid a second enrichment.
+ auto context = ffi::VisitErrorContext::TryGetFromError(error);
+ if (!context) throw;
+ auto normalized = ffi::make_object<ffi::VisitErrorContextObj>();
+ normalized->prev_error_context = context.value()->prev_error_context;
+ for (const auto& object : context.value()->reverse_visit_pattern) {
+ if (!IsContainer(object))
normalized->reverse_visit_pattern.push_back(object);
}
+ auto location = ffi::VisitErrorContext(std::move(normalized));
+ auto paths = ffi::VisitErrorContext::FindAccessPaths(root, location);
+ if (paths.empty()) throw;
+ std::ostringstream path;
+ path << paths[0];
+ std::string message = error.message();
+ const std::string marker = "<current location>";
+ size_t pos = 0;
+ bool replaced = false;
+ while ((pos = message.find(marker, pos)) != std::string::npos) {
+ message.replace(pos, marker.size(), path.str());
+ pos += path.str().size();
+ replaced = true;
+ }
+ if (!replaced) message += " at " + path.str();
+ throw ffi::Error(error.kind(), message, error.backtrace(),
error.cause_chain(),
+ location->prev_error_context.has_value()
+ ?
std::optional<ffi::ObjectRef>(location->prev_error_context.value())
+ : std::nullopt);
}
+ return !has_error_;
+ }
- {
- auto it = previously_defined_.find(var);
- auto verify = Verify(it == previously_defined_.end());
- verify << "ValueError: "
- << "TIR is ill-formed, "
- << "due to multiple definitions of variable " << var->name << ".";
- if (it != previously_defined_.end()) {
- verify << " It was first defined at " << it->second << ", and was
later re-defined at "
- << path;
+ protected:
+ using VTable = typename Visitor::VTable;
+ using Visitor::Visit_;
+ explicit WellFormednessVerifier(const VTable* table, Check check, bool
assert_mode)
+ : Visitor(table), check_(check), table_(table),
assert_mode_(assert_mode) {}
+
+ static void InitVTable(VTable* table) {
+ Visitor::InitVTable(table);
+ Visitor::template SetDispatch<WellFormednessVerifier, IRModuleNode>(table);
+ Visitor::template SetDispatch<WellFormednessVerifier, FunctionNode>(table);
+ Visitor::template SetDispatch<WellFormednessVerifier,
LambdaExprNode>(table);
+ }
+
+ Result Visit(ffi::AnyView value) override {
+ auto object = value.as<ffi::ObjectRef>();
+ if (!object) return Visitor::Visit(value);
+ visit_stack_.push_back(object.value());
+ struct Pop {
+ std::vector<ffi::ObjectRef>* stack;
+ ~Pop() { stack->pop_back(); }
+ } pop{&visit_stack_};
+ // Function entrypoints and module members are roots of verification;
+ // a Function embedded as an expression remains unsupported.
+ if (value.as<Function>() && visit_stack_.size() != 1 &&
+ !(visit_stack_.size() == 2 && root_.as<IRModule>())) {
+ TVM_FFI_VISIT_THROW(TypeError, object.value())
+ << "Unsupported non-primitive TIR expression " <<
object.value()->GetTypeKey();
+ }
+ if (value.as<Stmt>() && !table_->CanDispatch(value)) {
+ TVM_FFI_VISIT_BEGIN();
+ Verify(false) << "TIR verifier does not support statement " <<
object.value()->GetTypeKey()
+ << " at <current location>";
+ return std::nullopt;
+ TVM_FFI_VISIT_END(object.value());
+ }
+ if (value.as<Expr>() && !table_->CanDispatch(value)) {
+ if (value.as<OpaqueExprNode>()) return std::nullopt;
+ if (auto var = value.as<VarNode>()) {
+ TVM_FFI_VISIT_BEGIN();
+ return this->Visit_(var);
+ TVM_FFI_VISIT_END(object.value());
}
+ TVM_FFI_VISIT_THROW(TypeError, object.value())
+ << "Unsupported non-primitive TIR expression " <<
object.value()->GetTypeKey();
}
-
- currently_defined_.insert({var, path});
+ return Visitor::Visit(value);
}
- void ExitDef(const Var& var, AccessPath path) override {
- auto active_def = currently_defined_.find(var);
+ class VerifyStream {
+ public:
+ explicit VerifyStream(bool fail) {
+ if (fail) message_.emplace();
+ }
+ VerifyStream(const VerifyStream&) = delete;
+ VerifyStream(VerifyStream&& other) { std::swap(message_, other.message_); }
+ template <typename T>
+ VerifyStream& operator<<(T&& value) {
+ if (message_) *message_ << std::forward<T>(value);
+ return *this;
+ }
+ ~VerifyStream() noexcept(false) {
+ if (message_) LOG(FATAL) << message_->str();
+ }
+
+ private:
+ std::optional<std::ostringstream> message_;
+ };
- if (active_def != currently_defined_.end())
currently_defined_.erase(active_def);
- previously_defined_.insert({var, path});
+ VerifyStream Verify(bool condition) {
+ has_error_ |= !condition;
+ return VerifyStream(!condition && assert_mode_);
}
- void Dispatch_(const VarNode* op, AccessPath path) override {
- auto var = ffi::GetRef<Var>(op);
+ // Only earlier definitions need a saved occurrence. Keep their enclosing
+ // objects and the defining field, rather than searching for a Var that may
+ // also occur at arbitrarily many use sites.
+ struct DefinitionSite {
+ ffi::VisitErrorContext context;
+ AccessPath field;
+ };
+
+ // FindAccessPaths descends through containers without matching their
identity.
+ // Native visiting can record them; keep only matchable object breadcrumbs.
+ static bool IsContainer(const ffi::ObjectRef& object) {
+ switch (object->type_index()) {
+ case ffi::TypeIndex::kTVMFFIArray:
+ case ffi::TypeIndex::kTVMFFIList:
+ case ffi::TypeIndex::kTVMFFIMap:
+ case ffi::TypeIndex::kTVMFFIDict:
+ return true;
+ default:
+ return false;
+ }
+ }
- auto active_def = currently_defined_.find(var);
- auto verify = Verify(active_def != currently_defined_.end());
- verify << "ValueError: "
- << "Invalid use of undefined variable " << var->name << " at " <<
path << ".";
-
- // Check if there was a previous definition, and append the
- // location to the error message if there was. This is to aid in
- // debugging, by distinguishing between a variable that is
- // currently out-of-scope, and a variable that never had a
- // definition in the first place.
- if (auto prev_def = previously_defined_.find(var); prev_def !=
previously_defined_.end()) {
- verify << ". While this variable was previously defined at " <<
prev_def->second
- << ", this definition is no longer in-scope.";
+ DefinitionSite DefinitionAt(AccessPath field) const {
+ auto context = ffi::make_object<ffi::VisitErrorContextObj>();
+ for (auto it = visit_stack_.rbegin(); it != visit_stack_.rend(); ++it) {
+ if (!IsContainer(*it)) context->reverse_visit_pattern.push_back(*it);
}
+ return {ffi::VisitErrorContext(std::move(context)), std::move(field)};
}
- // Variables that are defined in the currently-visited scope.
- std::unordered_map<Var, AccessPath> currently_defined_;
+ AccessPath DefinitionPath(const DefinitionSite& site) const {
+ auto paths = ffi::VisitErrorContext::FindAccessPaths(root_, site.context);
+ auto path = paths.empty() ? AccessPath::Root() : paths[0];
+ for (const auto& step : site.field->ToSteps()) path = path->Extend(step);
+ return path;
+ }
- // Variables that were previously defined, and are now out of scope.
- std::unordered_map<Var, AccessPath> previously_defined_;
-};
+ template <typename F>
+ Result WithScope(F&& body) {
+ return scopes_.WithNewScope([&]() -> Result {
+ Result result = body();
+ // Definition exit never diagnoses an error. On an exception the entire
+ // verifier is discarded, so only normal completion needs history
updates.
+ for (auto it = scopes_.Current().rbegin(); it !=
scopes_.Current().rend(); ++it) {
+ current_.erase(it->first);
+ previous_.insert(*it);
+ }
+ return result;
+ });
+ }
-/*! \brief Verify that tensors with a declaration are not used outside their
declared scope.
- *
- * When a tensor is declared via one of the following sites:
- * - TensorType-annotated Function parameters
- * - DeclTensor statement
- * - Dialect-specific definitions exposed by PathVisitor
- *
- * it must not appear in a TensorLoad, TensorStore, or TensorRegion outside
that declaration's
- * scope.
- *
- * All tensors that appear in TensorLoad or TensorStore must have a prior
declaration.
- */
-template <typename PathVisitor>
-class UndefinedTensorVerifier : public
Verifier<UndefinedTensorVerifier<PathVisitor>, PathVisitor> {
- using Verifier = tirx::Verifier<UndefinedTensorVerifier<PathVisitor>,
PathVisitor>;
+ void Define(const Var& var, AccessPath field, bool visit_type = true) {
+ if (visit_type) this->Visit(var->ty);
+ auto site = DefinitionAt(std::move(field));
+ if (check_ == Check::kVariables) {
+ {
+ auto it = current_.find(var);
+ auto verify = Verify(it == current_.end());
+ verify << "ValueError: TIR is ill-formed, due to multiple nested
definitions of variable "
+ << var->name << ".";
+ if (it != current_.end()) {
+ verify << " It was first defined at " << DefinitionPath(it->second)
+ << ", and was re-defined at <current location>";
+ }
+ }
+ {
+ auto it = previous_.find(var);
+ auto verify = Verify(it == previous_.end());
+ verify << "ValueError: TIR is ill-formed, due to multiple definitions
of variable "
+ << var->name << ".";
+ if (it != previous_.end()) {
+ verify << " It was first defined at " << DefinitionPath(it->second)
+ << ", and was later re-defined at <current location>";
+ }
+ }
+ }
+ current_.insert({var, site});
+ scopes_.Current().emplace_back(var, std::move(site));
+ }
- public:
- using Verifier::Verifier;
- using Verifier::Verify;
+ void MatchTensorDefinitions(const TensorVar& tensor, AccessPath field) {
+ auto define_if_needed = [&](const Var& var, AccessPath site) {
+ if (!current_.count(var)) Define(var, std::move(site));
+ };
+ for (size_t i = 0; i < tensor->shape.size(); ++i) {
+ auto site = field->Attr("shape")->ArrayItem(i);
+ auto walk = [&](const Var& var) -> ffi::Expected<ffi::WalkResult> {
+ define_if_needed(var, site);
+ return ffi::WalkResult::Advance();
+ };
+ ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(tensor->shape[i], walk);
+ }
+ for (size_t i = 0; i < tensor->strides.size(); ++i) {
+ if (auto var = tensor->strides[i].template as<Var>()) {
+ define_if_needed(var.value(), field->Attr("strides")->ArrayItem(i));
+ }
+ }
+ if (auto var = tensor->elem_offset.template as<Var>()) {
+ define_if_needed(var.value(), field->Attr("elem_offset"));
+ }
+ }
- private:
- using Verifier::Visit;
+ Result Visit_(const IRModuleNode* op) {
+ std::vector<GlobalVar> globals;
+ for (const auto& [global, function] : op->functions)
globals.push_back(global);
+ std::sort(globals.begin(), globals.end(), [&](const GlobalVar& a, const
GlobalVar& b) {
+ bool a_public =
+ op->functions[a]->template
GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
+ bool b_public =
+ op->functions[b]->template
GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value();
+ return a_public != b_public ? a_public > b_public : a->name_hint <
b->name_hint;
+ });
+ for (const auto& global : globals) {
+ if (auto function = op->functions[global].template as<Function>())
+ this->Visit(function.value());
+ }
+ return std::nullopt;
+ }
- void Visit(const Function& function, AccessPath path) override {
- Verifier::Visit(function, path);
- // Clear per-function state (tensors should not cross function boundaries).
- currently_defined_.clear();
- previously_defined_.clear();
+ Result Visit_(const FunctionNode* op) {
+ int enclosing_depth = loop_depth_;
+ loop_depth_ = 0;
+ WithScope([&]() -> Result {
+ auto params = AccessPath::Root()->Attr("params");
+ for (size_t i = 0; i < op->params.size(); ++i) {
+ if (!op->params[i]->ty.template as<TensorTypeNode>())
+ Define(op->params[i], params->ArrayItem(i));
+ }
+ for (size_t i = 0; i < op->params.size(); ++i) {
+ if (auto tensor = op->params[i].template as<TensorVar>()) {
+ MatchTensorDefinitions(tensor.value(),
params->ArrayItem(i)->Attr("ty"));
+ }
+ }
+ for (size_t i = 0; i < op->params.size(); ++i) {
+ if (op->params[i].template as<TensorVar>()) Define(op->params[i],
params->ArrayItem(i));
+ }
+ return WithScope([&]() { return this->Visit(op->body); });
+ });
+ loop_depth_ = enclosing_depth;
+ if (check_ == Check::kTensors) {
+ current_.clear();
+ previous_.clear();
+ }
+ return std::nullopt;
}
- void EnterDef(const Var& var, AccessPath path) override {
- if (auto tensor = var.as<TensorVar>()) {
- currently_defined_.insert({tensor.value(), path});
+ Result Visit_(const VarNode* op) override {
+ auto var = ffi::GetRef<Var>(op);
+ if (check_ == Check::kVariables) {
+ auto verify = Verify(current_.count(var));
+ verify << "ValueError: Invalid use of undefined variable " << var->name
+ << " at <current location>.";
+ if (auto it = previous_.find(var); it != previous_.end()) {
+ verify << ". While this variable was previously defined at " <<
DefinitionPath(it->second)
+ << ", this definition is no longer in-scope.";
+ }
+ }
+ return std::nullopt;
+ }
+
+ void TensorUse(const TensorVar& tensor) {
+ if (check_ != Check::kTensors) return;
+ if (current_.count(tensor.var())) return;
+ auto it = previous_.find(tensor.var());
+ auto verify = Verify(false);
+ verify << "TIR is ill-formed: tensor " << tensor.name() << " is used at
<current location>";
+ if (it != previous_.end()) {
+ verify << " but its declaration is no longer in-scope. It was declared
at "
+ << DefinitionPath(it->second) << ".";
+ } else {
+ verify << " without a prior DeclTensor or other declaration.";
}
}
- void ExitDef(const Var& var, AccessPath path) override {
- if (!var->ty.as<TensorTypeNode>()) return;
- auto tensor = var.as_or_throw<TensorVar>();
- auto active_def = currently_defined_.find(tensor);
- if (active_def != currently_defined_.end()) {
- currently_defined_.erase(active_def);
+ Result Visit_(const BindNode* op) override {
+ this->Visit(op->value);
+ Define(op->var, AccessPath::Root()->Attr("var"));
+ return std::nullopt;
+ }
+
+ Result Visit_(const RegionStmtNode* op) override {
+ this->Visit(op->args);
+ this->Visit(op->attrs);
+ for (const auto& param : op->body_params) this->Visit(param->ty);
+ WithScope([&]() -> Result {
+ for (size_t i = 0; i < op->body_params.size(); ++i) {
+ Define(op->body_params[i],
AccessPath::Root()->Attr("body_params")->ArrayItem(i), false);
+ }
+ return this->Visit(op->body);
+ });
+ for (size_t i = 0; i < op->result_vars.size(); ++i) {
+ Define(op->result_vars[i],
AccessPath::Root()->Attr("result_vars")->ArrayItem(i));
}
- previously_defined_.insert({tensor, path});
- }
-
- void VisitTensorUse(const TensorVar& tensor, AccessPath path) override {
- bool is_declared = currently_defined_.count(tensor);
- bool was_declared = previously_defined_.count(tensor);
-
- if (was_declared && !is_declared) {
- // TensorVar was previously declared but is now out of scope — always an
error.
- auto prev_def = previously_defined_.find(tensor);
- Verify(false) << "TIR is ill-formed: tensor " << tensor.name() << " is
used at " << path
- << " but its declaration is no longer in-scope. "
- << "It was declared at " << prev_def->second << ".";
- } else if (!is_declared && !was_declared) {
- // TensorVar was never declared — error.
- Verify(false) << "TIR is ill-formed: tensor " << tensor.name() << " is
used at " << path
- << " without a prior DeclTensor or other declaration.";
+ return std::nullopt;
+ }
+
+ Result Visit_(const ForNode* op) override {
+ this->Visit(op->min);
+ this->Visit(op->extent);
+ if (check_ == Check::kLoopControl) {
+ this->Visit(op->step);
+ ++loop_depth_;
+ this->Visit(op->body);
+ --loop_depth_;
+ return std::nullopt;
}
- // TensorVar fields are visited at definition site (EnterDef), not here.
- Verifier::VisitTensorUse(tensor, path);
+ return WithScope([&]() -> Result {
+ Define(op->loop_var, AccessPath::Root()->Attr("loop_var"));
+ return WithScope([&]() { return this->Visit(op->body); });
+ });
}
- // Tensors defined in the currently-visited scope.
- std::unordered_map<TensorVar, AccessPath, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
- currently_defined_;
- // Tensors that were previously defined and are now out of scope.
- std::unordered_map<TensorVar, AccessPath, ffi::ObjectPtrHash,
ffi::ObjectPtrEqual>
- previously_defined_;
-};
+ Result Visit_(const WhileNode* op) override {
+ this->Visit(op->condition);
+ ++loop_depth_;
+ WithScope([&]() { return this->Visit(op->body); });
+ --loop_depth_;
+ return std::nullopt;
+ }
-/*! \brief Verify the asserted type of each tirx tensor load. */
-template <typename PathVisitor>
-class TensorLoadTypeVerifier : public
Verifier<TensorLoadTypeVerifier<PathVisitor>, PathVisitor> {
- using Verifier = tirx::Verifier<TensorLoadTypeVerifier<PathVisitor>,
PathVisitor>;
+ Result Visit_(const IfNode* op) override {
+ this->Visit(op->condition);
+ WithScope([&]() { return this->Visit(op->then_case); });
+ return WithScope([&]() { return this->Visit(op->else_case); });
+ }
- public:
- using Verifier::Verifier;
- using Verifier::Verify;
+ Result Visit_(const prim::LetNode* op) override {
+ this->Visit(op->value);
+ return WithScope([&]() -> Result {
+ Define(op->var, AccessPath::Root()->Attr("var"));
+ return this->Visit(op->body);
+ });
+ }
- private:
- using Verifier::Visit;
- void Dispatch_(const TensorLoadNode* op, AccessPath path) override {
- auto tensor = op->source.as<TensorVar>();
- auto valid_source = Verify(tensor.has_value());
- valid_source << "TypeError: TIR TensorLoad source at " <<
path->Attr("source")
- << " must be a TensorVar.";
- if (!tensor.has_value()) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
+ Result Visit_(const LambdaExprNode* op) {
+ this->Visit(op->ty);
+ return WithScope([&]() -> Result {
+ for (size_t i = 0; i < op->vars.size(); ++i) {
+ Define(op->vars[i], AccessPath::Root()->Attr("vars")->ArrayItem(i));
+ }
+ return this->Visit(op->body);
+ });
+ }
- bool valid_indices = tensor.value()->shape.size() == op->indices.size();
- auto valid_rank = Verify(valid_indices);
- valid_rank << "ValueError: TIR TensorLoad at " << path << " indexes "
- << tensor.value()->shape.size() << "-dimensional tensor " <<
tensor.value().name()
- << " with " << op->indices.size() << " indices.";
- if (!valid_indices) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
+ Result Visit_(const CallNode* op) override {
+ if (op->op.template as<GlobalVar>() || op->op.template
as<OpaqueExprNode>())
+ this->Visit(op->op);
+ return this->Visit(op->args);
+ }
- for (size_t i = 0; i + 1 < op->indices.size(); ++i) {
- bool is_scalar = op->indices[i].ty().IsScalar();
- auto valid_index = Verify(is_scalar);
- valid_index << "TypeError: TIR TensorLoad index " << i << " at "
- << path->Attr("indices")->ArrayItem(i)
- << " must be scalar because only the final index may be
vector-valued.";
- valid_indices = valid_indices && is_scalar;
- }
- if (!valid_indices) {
- Visit(op->indices, path->Attr("indices"));
- return;
- }
+ Result Visit_(const TensorStoreNode* op) override {
+ this->Visit(op->value);
+ TensorUse(op->dest.template as_or_throw<TensorVar>());
+ return this->Visit(op->indices);
+ }
- ffi::Optional<PrimType> index_ty = op->indices.empty()
- ?
ffi::Optional<PrimType>(tensor.value()->dtype)
- :
op->indices.back().ty().as<PrimType>();
- AccessPath final_index_path = op->indices.empty()
- ? path->Attr("indices")
- :
path->Attr("indices")->ArrayItem(op->indices.size() - 1);
- auto valid_index_type = Verify(index_ty.has_value());
- valid_index_type << "TypeError: TIR TensorLoad final index at " <<
final_index_path
- << " must have a primitive type.";
- if (!index_ty.has_value()) {
- Visit(op->indices, path->Attr("indices"));
- return;
+ Result Visit_(const TensorRegionNode* op) override {
+ if (auto tensor = op->source.template as<TensorVar>()) {
+ TensorUse(tensor.value());
+ } else {
+ this->Visit(op->source);
}
+ return this->Visit(op->region);
+ }
+
+ Result Visit_(const GenericConstNode* op) override { return std::nullopt; }
+
+ Result Visit_(const AssertStmtNode* op) override {
+ this->Visit(op->condition);
+ this->Visit(op->error_kind);
+ return this->Visit(op->message_parts);
+ }
+
+ Result Visit_(const prim::RampNode* op) override {
+ this->Visit(op->base);
+ this->Visit(op->stride);
+ return this->Visit(op->lanes);
+ }
+
+ Result Visit_(const prim::BroadcastNode* op) override {
+ this->Visit(op->value);
+ return this->Visit(op->lanes);
+ }
- bool scalable_compatible = op->indices.empty() ||
!(tensor.value()->dtype.IsScalableVector() &&
-
index_ty.value().IsScalableVector());
- auto valid_scalability = Verify(scalable_compatible);
- valid_scalability << "TypeError: TIR TensorLoad at " << path
- << " cannot combine a scalable tensor dtype with a
scalable index.";
- if (!scalable_compatible) {
- Visit(op->indices, path->Attr("indices"));
- return;
+ Result Visit_(const BreakNode* op) override {
+ if (check_ == Check::kLoopControl) {
+ Verify(loop_depth_ > 0) << "ValueError: break at <current location>
requires an enclosing "
+ "loop in the same function.";
}
+ return std::nullopt;
+ }
- TensorLoad expected = MakeTensorLoad(tensor.value(), op->indices, op->loc);
- ffi::Optional<PrimType> asserted_ty = op->ty.as<PrimType>();
- ffi::Optional<PrimType> expected_ty = expected->ty.as<PrimType>();
- auto valid_type = Verify(asserted_ty.has_value() &&
expected_ty.has_value() &&
- asserted_ty.value() == expected_ty.value());
- valid_type << "TypeError: TIR TensorLoad at " << path << " asserts result
type " << op->ty
- << ", but its source and indices imply " << expected->ty << ".";
- PathVisitor::Dispatch_(op, path);
+ Result Visit_(const ContinueNode* op) override {
+ if (check_ == Check::kLoopControl) {
+ Verify(loop_depth_ > 0) << "ValueError: continue at <current location>
requires an enclosing "
+ "loop in the same function.";
+ }
+ return std::nullopt;
}
-};
-/*! \brief Verify that loop control belongs to a loop body in the same
function. */
-template <typename PathVisitor>
-class LoopControlVerifier : public Verifier<LoopControlVerifier<PathVisitor>,
PathVisitor> {
- using Verifier = tirx::Verifier<LoopControlVerifier<PathVisitor>,
PathVisitor>;
+ Result Visit_(const TensorLoadNode* op) override;
- public:
- using Verifier::Verifier;
- using Verifier::Verify;
+ Check check_;
private:
- using Verifier::Visit;
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
- void Visit(const Function& function, AccessPath path) override {
- int enclosing_depth = loop_depth_;
- loop_depth_ = 0;
- Verifier::Visit(function, path);
- loop_depth_ = enclosing_depth;
+ const VTable* table_;
+ bool assert_mode_;
+ bool has_error_{false};
+ int loop_depth_{0};
+ ffi::ObjectRef root_;
+ std::vector<ffi::ObjectRef> visit_stack_;
+ ScopeStack<std::vector<std::pair<Var, DefinitionSite>>> scopes_;
+ std::unordered_map<Var, DefinitionSite> current_;
+ std::unordered_map<Var, DefinitionSite> previous_;
+};
+
+template <typename Visitor>
+auto WellFormednessVerifier<Visitor>::Visit_(const TensorLoadNode* op) ->
Result {
+ if (check_ != Check::kTensorLoadType) {
+ TensorUse(op->source.template as_or_throw<TensorVar>());
+ return this->Visit(op->indices);
+ }
+ auto tensor = op->source.as<TensorVar>();
+ auto valid_source = Verify(tensor.has_value());
+ valid_source << "TypeError: TIR TensorLoad source at " << "<current
location>.source"
+ << " must be a TensorVar.";
+ if (!tensor.has_value()) {
+ this->Visit(op->indices);
+ return std::nullopt;
}
- void Dispatch_(const ForNode* op, AccessPath path) override {
- Visit(op->min, path->Attr("min"));
- Visit(op->extent, path->Attr("extent"));
- Visit(op->step, path->Attr("step"));
- ++loop_depth_;
- Visit(op->body, path->Attr("body"));
- --loop_depth_;
+ bool valid_indices = tensor.value()->shape.size() == op->indices.size();
+ auto valid_rank = Verify(valid_indices);
+ valid_rank << "ValueError: TIR TensorLoad at " << "<current location>" << "
indexes "
+ << tensor.value()->shape.size() << "-dimensional tensor " <<
tensor.value().name()
+ << " with " << op->indices.size() << " indices.";
+ if (!valid_indices) {
+ this->Visit(op->indices);
+ return std::nullopt;
}
- void Dispatch_(const WhileNode* op, AccessPath path) override {
- Visit(op->condition, path->Attr("condition"));
- ++loop_depth_;
- Visit(op->body, path->Attr("body"));
- --loop_depth_;
+ for (size_t i = 0; i + 1 < op->indices.size(); ++i) {
+ bool is_scalar = op->indices[i].ty().IsScalar();
+ auto valid_index = Verify(is_scalar);
+ valid_index << "TypeError: TIR TensorLoad index " << i << " at "
+ << "<current location>.indices[" << i << "]"
+ << " must be scalar because only the final index may be
vector-valued.";
+ valid_indices = valid_indices && is_scalar;
+ }
+ if (!valid_indices) {
+ this->Visit(op->indices);
+ return std::nullopt;
}
- void Dispatch_(const BreakNode* op, AccessPath path) override {
- Verify(loop_depth_ > 0) << "ValueError: break at " << path
- << " requires an enclosing loop in the same
function.";
+ ffi::Optional<PrimType> index_ty = op->indices.empty()
+ ?
ffi::Optional<PrimType>(tensor.value()->dtype)
+ :
op->indices.back().ty().as<PrimType>();
+ auto valid_index_type = Verify(index_ty.has_value());
+ valid_index_type << "TypeError: TIR TensorLoad final index at <current
location>.indices";
+ if (!op->indices.empty()) valid_index_type << "[" << op->indices.size() - 1
<< "]";
+ valid_index_type << " must have a primitive type.";
+ if (!index_ty.has_value()) {
+ this->Visit(op->indices);
+ return std::nullopt;
}
- void Dispatch_(const ContinueNode* op, AccessPath path) override {
- Verify(loop_depth_ > 0) << "ValueError: continue at " << path
- << " requires an enclosing loop in the same
function.";
+ bool scalable_compatible = op->indices.empty() ||
!(tensor.value()->dtype.IsScalableVector() &&
+
index_ty.value().IsScalableVector());
+ auto valid_scalability = Verify(scalable_compatible);
+ valid_scalability << "TypeError: TIR TensorLoad at " << "<current location>"
+ << " cannot combine a scalable tensor dtype with a
scalable index.";
+ if (!scalable_compatible) {
+ this->Visit(op->indices);
+ return std::nullopt;
}
- int loop_depth_{0};
-};
+ TensorLoad expected = MakeTensorLoad(tensor.value(), op->indices, op->loc);
+ ffi::Optional<PrimType> asserted_ty = op->ty.as<PrimType>();
+ ffi::Optional<PrimType> expected_ty = expected->ty.as<PrimType>();
+ auto valid_type = Verify(asserted_ty.has_value() && expected_ty.has_value()
&&
+ asserted_ty.value() == expected_ty.value());
+ valid_type << "TypeError: TIR TensorLoad at " << "<current location>" << "
asserts result type "
+ << op->ty << ", but its source and indices imply " <<
expected->ty << ".";
+ TensorUse(tensor.value());
+ return this->Visit(op->indices);
+}
-template <typename PathVisitor, typename NodeRef>
+template <typename Checker, typename NodeRef>
bool VerifyWellFormedCommon(const NodeRef& node, bool assert_mode) {
- return UndefinedVarVerifier<PathVisitor>::Verify(node, assert_mode) &&
- UndefinedTensorVerifier<PathVisitor>::Verify(node, assert_mode) &&
- TensorLoadTypeVerifier<PathVisitor>::Verify(node, assert_mode) &&
- LoopControlVerifier<PathVisitor>::Verify(node, assert_mode);
+ using Check = typename Checker::Check;
+ for (Check check :
+ {Check::kVariables, Check::kTensors, Check::kTensorLoadType,
Check::kLoopControl}) {
+ if (!ffi::make_object<Checker>(check, assert_mode)->Run(node)) return
false;
+ }
+ return true;
}
+
} // namespace tirx
} // namespace tvm
#endif // TVM_TIRX_ANALYSIS_VERIFY_WELL_FORMED_H_
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
deleted file mode 100644
index 3bf6f490c6..0000000000
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ /dev/null
@@ -1,355 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tirx/ir/tir_visitor_with_path.cc
- * \brief Provide a TIR visitor that tracks the current location
- */
-#include "tir_visitor_with_path.h"
-
-#include <tvm/ffi/reflection/access_path.h>
-#include <tvm/ffi/reflection/accessor.h>
-#include <tvm/ir/function.h>
-#include <tvm/tirx/stmt.h>
-
-#include <algorithm>
-#include <optional>
-#include <unordered_set>
-#include <variant>
-#include <vector>
-
-namespace tvm {
-namespace tirx {
-using AccessPath = ffi::reflection::AccessPath;
-
-void TIRVisitorWithPath::Visit(const IRModule& mod, AccessPath path) {
- // To ensure deterministic order of visits, sort the GlobalVar first
- // by visibility (public then private), then alphabetically by name.
- std::vector<GlobalVar> gvars;
- std::unordered_set<GlobalVar> externally_exposed;
- for (const auto& [gvar, func] : mod->functions) {
- gvars.push_back(gvar);
- if (func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol).has_value()) {
- externally_exposed.insert(gvar);
- }
- }
-
- std::sort(gvars.begin(), gvars.end(),
- [&externally_exposed](const GlobalVar& a, const GlobalVar& b) {
- bool a_exposed = externally_exposed.count(a);
- bool b_exposed = externally_exposed.count(b);
- if (a_exposed != b_exposed) {
- return a_exposed > b_exposed;
- } else {
- return a->name_hint < b->name_hint;
- }
- });
-
- std::vector<DefContext<GlobalVar>> context;
-
- for (const auto& gvar : gvars) {
- context.push_back(WithDef(gvar,
path->Attr("global_var_map_")->MapItem(gvar->name_hint)));
- }
-
- for (const auto& gvar : gvars) {
- auto base_func = mod->functions[gvar];
- if (auto function = base_func.as<Function>()) {
- Visit(function.value(), path->Attr("functions")->MapItem(gvar));
- }
- }
-
- while (context.size()) context.pop_back();
-}
-
-void TIRVisitorWithPath::Visit(const Function& func, AccessPath path) {
- // TensorType metadata may introduce symbolic dimensions. Define those
- // symbols before entering the tensor parameter itself.
- std::vector<DefContext<Var>> context;
-
- auto ppath = path->Attr("params");
- for (size_t i = 0; i < func->params.size(); i++) {
- const Var& param = func->params[i];
- if (!param->ty.as<TensorTypeNode>()) {
- context.push_back(WithDef(param, ppath->ArrayItem(i)));
- }
- }
-
- for (size_t i = 0; i < func->params.size(); i++) {
- if (auto opt = func->params[i].as<TensorVar>()) {
- auto buf = opt.value();
- auto buf_path = ppath->ArrayItem(i)->Attr("ty");
-
- for (auto& def : WithMatchTensorDefs(buf, buf_path)) {
- context.push_back(std::move(def));
- }
- }
- }
-
- // Only after all the implicit definitions have been visited can we
- // visit the tensor definition itself.
- for (size_t i = 0; i < func->params.size(); i++) {
- if (auto opt = func->params[i].as<TensorVar>()) {
- context.push_back(WithDef(opt.value().var(), ppath->ArrayItem(i)));
- }
- }
-
- bind_scope_.WithNewScope([&]() { Visit(func->body, path->Attr("body")); });
-
- while (context.size()) context.pop_back();
-}
-
-void TIRVisitorWithPath::Visit(const Type& type, AccessPath path) {
- Visit(ffi::AnyView(type), path);
-}
-
-void TIRVisitorWithPath::Visit(ffi::AnyView obj, AccessPath path) {
- if (obj == nullptr) return;
- if (auto expr = obj.as<Expr>()) {
- Visit(expr.value(), path);
- } else if (auto array = obj.as<ffi::Array<ffi::Any>>()) {
- for (size_t i = 0; i < array.value().size(); ++i) {
- Visit(ffi::AnyView(array.value()[i]), path->ArrayItem(i));
- }
- } else if (auto map = obj.as<ffi::Map<ffi::Any, ffi::Any>>()) {
- for (const auto& [key, value] : map.value()) {
- Visit(ffi::AnyView(value), path->MapItem(key));
- }
- } else if (auto object = obj.as<ffi::ObjectRef>()) {
- ffi::reflection::ForEachFieldInfo(
- TVMFFIGetTypeInfo(object.value()->type_index()), [&](const
TVMFFIFieldInfo* field) {
-
Visit(ffi::AnyView(ffi::reflection::FieldGetter(field)(object.value())),
- path->Attr(ffi::String(field->name)));
- });
- }
-}
-
-// Uses retain their definition's identity without revisiting its type
metadata.
-void TIRVisitorWithPath::VisitTensorUse(const TensorVar& tensor, AccessPath
path) {}
-
-void TIRVisitorWithPath::Visit(const TensorRegion& region, AccessPath path) {
- if (auto tensor = region->source.as<TensorVar>()) {
- VisitTensorUse(tensor.value(), path->Attr("source"));
- } else {
- Visit(region->source, path->Attr("source"));
- }
- Visit(region->region, path->Attr("region"));
-}
-
-void TIRVisitorWithPath::Visit(const Range& range, AccessPath path) {
- Visit(range->min, path->Attr("min"));
- Visit(range->extent, path->Attr("extent"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const BindNode* op, AccessPath path) {
- Visit(op->value, path->Attr("value"));
- // Push the Bind's var definition into the current scope.
- // The def lives until the enclosing scope (body-carrying stmt) exits.
- bind_scope_.Current().push_back(WithDef(op->var, path->Attr("var")));
-}
-
-void TIRVisitorWithPath::Dispatch_(const RegionStmtNode* op, AccessPath path) {
- Visit(op->args, path->Attr("args"));
- Visit(ffi::AnyView(op->attrs), path->Attr("attrs"));
- // All parameter types belong to the enclosing scope, including types with
- // symbolic dimensions. No body parameter is visible to another parameter's
type.
- auto params_path = path->Attr("body_params");
- for (size_t i = 0; i < op->body_params.size(); ++i) {
- Visit(op->body_params[i]->ty, params_path->ArrayItem(i)->Attr("ty"));
- }
- bind_scope_.WithNewScope([&]() {
- for (size_t i = 0; i < op->body_params.size(); ++i) {
- bind_scope_.Current().push_back(
- WithDef(op->body_params[i], params_path->ArrayItem(i),
/*visit_type=*/false));
- }
- Visit(op->body, path->Attr("body"));
- });
- auto results_path = path->Attr("result_vars");
- for (size_t i = 0; i < op->result_vars.size(); ++i) {
- bind_scope_.Current().push_back(WithDef(op->result_vars[i],
results_path->ArrayItem(i)));
- }
-}
-
-void TIRVisitorWithPath::Dispatch_(const ForNode* op, AccessPath path) {
- Visit(op->min, path->Attr("min"));
- Visit(op->extent, path->Attr("extent"));
- auto context = WithDef(op->loop_var, path->Attr("loop_var"));
- bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); });
-}
-
-void TIRVisitorWithPath::Dispatch_(const WhileNode* op, AccessPath path) {
- Visit(op->condition, path->Attr("condition"));
- bind_scope_.WithNewScope([&]() { Visit(op->body, path->Attr("body")); });
-}
-
-void TIRVisitorWithPath::Dispatch_(const ReturnNode* op, AccessPath path) {
- Visit(op->value, path->Attr("value"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const BreakNode* op, AccessPath path) {}
-
-void TIRVisitorWithPath::Dispatch_(const ContinueNode* op, AccessPath path) {}
-
-void TIRVisitorWithPath::Dispatch_(const TensorStoreNode* op, AccessPath path)
{
- Visit(op->value, path->Attr("value"));
- VisitTensorUse(op->dest.as_or_throw<TensorVar>(), path->Attr("dest"));
- Visit(op->indices, path->Attr("indices"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const IfNode* op, AccessPath path) {
- Visit(op->condition, path->Attr("condition"));
- bind_scope_.WithNewScope([&]() { Visit(op->then_case,
path->Attr("then_case")); });
- bind_scope_.WithNewScope([&]() { Visit(op->else_case,
path->Attr("else_case")); });
-}
-
-void TIRVisitorWithPath::Dispatch_(const AssertStmtNode* op, AccessPath path) {
- Visit(op->condition, path->Attr("condition"));
- Visit(op->error_kind, path->Attr("error_kind"));
- Visit(op->message_parts, path->Attr("message_parts"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const SeqStmtNode* op, AccessPath path) {
- auto seq_path = path->Attr("seq");
- for (size_t i = 0; i < op->seq.size(); i++) {
- Visit(op->seq[i], seq_path->ArrayItem(i));
- }
-}
-
-void TIRVisitorWithPath::Dispatch_(const EvaluateNode* op, AccessPath path) {
- Visit(op->value, path->Attr("value"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const VarNode* op, AccessPath path) {}
-
-void TIRVisitorWithPath::Dispatch_(const TensorLoadNode* op, AccessPath path) {
- VisitTensorUse(op->source.as_or_throw<tvm::tirx::TensorVar>(),
path->Attr("source"));
- Visit(op->indices, path->Attr("indices"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const TensorRegionNode* op, AccessPath
path) {
- Visit(ffi::GetRef<TensorRegion>(op), path);
-}
-
-void TIRVisitorWithPath::Dispatch_(const OpaqueExprNode* op, AccessPath path)
{}
-
-void TIRVisitorWithPath::VisitLambda(const LambdaExprNode* op, AccessPath
path) {
- Visit(op->ty, path->Attr("ty"));
- std::vector<DefContext<Var>> context;
- for (size_t i = 0; i < op->vars.size(); ++i) {
- context.push_back(WithDef(op->vars[i], path->Attr("vars")->ArrayItem(i)));
- }
- Visit(op->body, path->Attr("body"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const TupleNode* op, AccessPath path) {
- Visit(op->fields, path->Attr("fields"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const TupleGetItemNode* op, AccessPath
path) {
- Visit(op->tuple, path->Attr("tuple"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::LetNode* op, AccessPath path) {
- Visit(op->value, path->Attr("value"));
- auto context = WithDef(op->var, path->Attr("var"));
- Visit(op->body, path->Attr("body"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const CallNode* op, AccessPath path) {
- if (auto gvar = op->op.as<GlobalVar>()) {
- Visit(gvar.value(), path->Attr("op"));
- } else if (op->op.as<OpaqueExprNode>()) {
- Visit(op->op, path->Attr("op"));
- }
- Visit(op->args, path->Attr("args"));
-}
-
-#define DEFINE_BINOP_VISIT_(OP) \
- void TIRVisitorWithPath::Dispatch_(const OP* op, AccessPath path) { \
- Visit(op->a, path->Attr("a")); \
- Visit(op->b, path->Attr("b")); \
- }
-
-DEFINE_BINOP_VISIT_(prim::AddNode);
-DEFINE_BINOP_VISIT_(prim::SubNode);
-DEFINE_BINOP_VISIT_(prim::MulNode);
-DEFINE_BINOP_VISIT_(prim::DivNode);
-DEFINE_BINOP_VISIT_(prim::ModNode);
-DEFINE_BINOP_VISIT_(prim::FloorDivNode);
-DEFINE_BINOP_VISIT_(prim::FloorModNode);
-DEFINE_BINOP_VISIT_(prim::MinNode);
-DEFINE_BINOP_VISIT_(prim::MaxNode);
-DEFINE_BINOP_VISIT_(prim::EQNode);
-DEFINE_BINOP_VISIT_(prim::NENode);
-DEFINE_BINOP_VISIT_(prim::LTNode);
-DEFINE_BINOP_VISIT_(prim::LENode);
-DEFINE_BINOP_VISIT_(prim::GTNode);
-DEFINE_BINOP_VISIT_(prim::GENode);
-DEFINE_BINOP_VISIT_(prim::AndNode);
-DEFINE_BINOP_VISIT_(prim::OrNode);
-
-DEFINE_BINOP_VISIT_(prim::LShiftNode);
-DEFINE_BINOP_VISIT_(prim::RShiftNode);
-DEFINE_BINOP_VISIT_(prim::BitwiseAndNode);
-DEFINE_BINOP_VISIT_(prim::BitwiseOrNode);
-DEFINE_BINOP_VISIT_(prim::BitwiseXorNode);
-
-#undef DEFINE_BINOP_VISIT_
-
-void TIRVisitorWithPath::Dispatch_(const IntImmNode* op, AccessPath path) {}
-void TIRVisitorWithPath::Dispatch_(const FloatImmNode* op, AccessPath path) {}
-void TIRVisitorWithPath::Dispatch_(const StringImmNode* op, AccessPath path) {}
-void TIRVisitorWithPath::Dispatch_(const DataTypeImmNode* op, AccessPath path)
{}
-void TIRVisitorWithPath::Dispatch_(const GenericConstNode* op, AccessPath
path) {}
-
-void TIRVisitorWithPath::Dispatch_(const prim::CastNode* op, AccessPath path) {
- Visit(op->value, path->Attr("value"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::BitwiseNotNode* op, AccessPath
path) {
- Visit(op->a, path->Attr("a"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::NotNode* op, AccessPath path) {
- Visit(op->a, path->Attr("a"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::SelectNode* op, AccessPath
path) {
- Visit(op->condition, path->Attr("condition"));
- Visit(op->true_value, path->Attr("true_value"));
- Visit(op->false_value, path->Attr("false_value"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::RampNode* op, AccessPath path) {
- Visit(op->base, path->Attr("base"));
- Visit(op->stride, path->Attr("stride"));
- Visit(op->lanes, path->Attr("lanes"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::ShuffleNode* op, AccessPath
path) {
- Visit(op->indices, path->Attr("indices"));
- Visit(op->vectors, path->Attr("vectors"));
-}
-
-void TIRVisitorWithPath::Dispatch_(const prim::BroadcastNode* op, AccessPath
path) {
- Visit(op->value, path->Attr("value"));
- Visit(op->lanes, path->Attr("lanes"));
-}
-
-} // namespace tirx
-} // namespace tvm
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
deleted file mode 100644
index c86ae7fa72..0000000000
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ /dev/null
@@ -1,403 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file tirx/ir/tir_visitor_with_path.h
- * \brief Provide a TIR visitor that tracks the current location
- */
-#ifndef TVM_TIRX_IR_TIR_VISITOR_WITH_PATH_H_
-#define TVM_TIRX_IR_TIR_VISITOR_WITH_PATH_H_
-
-#include <tvm/ffi/extra/structural_visit.h>
-#include <tvm/ir/expr_functor.h>
-#include <tvm/ir/module.h>
-#include <tvm/ir/scope_stack.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/tirx/stmt.h>
-#include <tvm/tirx/stmt_functor.h>
-
-#include <exception>
-#include <optional>
-#include <unordered_set>
-#include <utility>
-#include <variant>
-#include <vector>
-
-namespace tvm {
-namespace tirx {
-
-/*! \brief Visit TIR while tracking the ffi::reflection::AccessPath */
-class TIRVisitorWithPath : protected ExprFunctor<void(const Expr&,
ffi::reflection::AccessPath)>,
- protected StmtFunctor<void(const Stmt&,
ffi::reflection::AccessPath)> {
- public:
- TIRVisitorWithPath() = default;
- template <typename TObjectRef>
- void operator()(TObjectRef&& obj) {
- Visit(std::forward<TObjectRef>(obj), ffi::reflection::AccessPath::Root());
- }
-
- protected:
- using StmtVisitor = StmtFunctor<void(const Stmt&,
ffi::reflection::AccessPath)>;
- using VTable = StmtVisitor::VTable;
- explicit TIRVisitorWithPath(const VTable* vtable) : StmtVisitor(vtable) {}
- static void InitVTable(VTable* vtable) { StmtVisitor::InitVTable(vtable); }
- // Delegate to ExprFunctor::Dispatch for PrimExpr, and any subclasses
- virtual inline void Visit(const PrimExpr& obj, ffi::reflection::AccessPath
path) {
- Dispatch(obj, path);
- }
- // Core Call stores arguments as Expr, including pointer-typed Vars.
- virtual inline void Visit(const Expr& obj, ffi::reflection::AccessPath path)
{
- if (auto prim = obj.as<PrimExpr>()) {
- Visit(prim.value(), path);
- } else if (auto* constant = obj.as<GenericConstNode>()) {
- Dispatch_(constant, path);
- } else if (auto* str = obj.as<StringImmNode>()) {
- Dispatch_(str, path);
- } else if (auto* str = obj.as<DataTypeImmNode>()) {
- Dispatch_(str, path);
- } else if (obj.as<OpNode>()) {
- // Registered operations are atomic expression inputs.
- } else if (auto global_var = obj.as<GlobalVar>()) {
- Visit(global_var.value(), path);
- } else if (auto* var = obj.as<VarNode>()) {
- Dispatch_(var, path);
- } else if (auto* call = obj.as<CallNode>()) {
- Dispatch_(call, path);
- } else if (auto* tuple = obj.as<TupleNode>()) {
- Dispatch_(tuple, path);
- } else if (auto* tuple_get_item = obj.as<TupleGetItemNode>()) {
- Dispatch_(tuple_get_item, path);
- } else if (auto* tensor_region = obj.as<TensorRegionNode>()) {
- Dispatch_(tensor_region, path);
- } else if (auto* lambda = obj.as<LambdaExprNode>()) {
- VisitLambda(lambda, path);
- } else if (obj.as<OpaqueExprNode>()) {
- Dispatch(obj, path);
- } else {
- TVM_FFI_THROW(TypeError) << "Unsupported non-primitive TIR expression "
<< obj.GetTypeKey();
- }
- }
- // Delegate to StmtFunctor::Dispatch for Stmt, and any subclasses
- virtual inline void Visit(const Stmt& obj, ffi::reflection::AccessPath path)
{
- Dispatch(obj, path);
- }
-
- // Visit a tensor at a use site (TensorLoad, TensorStore, reads/writes).
- // By default, does not re-visit tensor fields (shape, strides, elem_offset),
- // as those are visited at the definition site via EnterDef.
- virtual void VisitTensorUse(const TensorVar& obj,
ffi::reflection::AccessPath path);
-
- // Visit type metadata through its reflected fields, preserving source
access paths.
- virtual void Visit(const Type& obj, ffi::reflection::AccessPath path);
- virtual void Visit(ffi::AnyView obj, ffi::reflection::AccessPath path);
- void VisitLambda(const LambdaExprNode* op, ffi::reflection::AccessPath path);
-
- // Visitors for TIR constructs that are neither PrimExpr nor Stmt
- virtual void Visit(const IRModule& obj, ffi::reflection::AccessPath path);
- virtual void Visit(const Function& obj, ffi::reflection::AccessPath path);
- virtual void Visit(const GlobalVar& obj, ffi::reflection::AccessPath path) {}
- virtual void Visit(const Range& obj, ffi::reflection::AccessPath path);
- virtual void Visit(const TensorRegion& obj, ffi::reflection::AccessPath
path);
-
- // Called when entering/exiting the scope of a GlobalVar definition.
- virtual void EnterDef(const GlobalVar& var, ffi::reflection::AccessPath
path) {}
- virtual void ExitDef(const GlobalVar& var, ffi::reflection::AccessPath path)
{}
-
- // Called when entering/exiting the scope of a tvm::Var definition.
- virtual void EnterDef(const Var& var, ffi::reflection::AccessPath path) {}
- virtual void ExitDef(const Var& var, ffi::reflection::AccessPath path) {}
-
- // Utility to visit an array of nodes
- template <typename T>
- inline void Visit(const ffi::Array<T>& arr, ffi::reflection::AccessPath
path) {
- for (size_t i = 0; i < arr.size(); i++) {
- Visit(arr[i], path->ArrayItem(i));
- }
- }
-
- // Utility to visit an optional node nodes
- template <typename T>
- inline void Visit(const ffi::Optional<T>& opt, ffi::reflection::AccessPath
path) {
- if (opt) {
- Visit(opt.value(), path);
- }
- }
-
- using StmtFunctor::Dispatch;
- void Dispatch_(const BindNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const RegionStmtNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const IfNode* op, ffi::reflection::AccessPath path) override;
- void Dispatch_(const ForNode* op, ffi::reflection::AccessPath path) override;
- void Dispatch_(const WhileNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const ReturnNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const BreakNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const ContinueNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const TensorStoreNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const AssertStmtNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const SeqStmtNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const EvaluateNode* op, ffi::reflection::AccessPath path)
override;
-
- using ExprFunctor::Dispatch;
- void Dispatch_(const VarNode* op, ffi::reflection::AccessPath path) override;
- void Dispatch_(const TensorLoadNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const TensorRegionNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const OpaqueExprNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const TupleNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const TupleGetItemNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const CallNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const IntImmNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const FloatImmNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const StringImmNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const DataTypeImmNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const GenericConstNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::LetNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::AddNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::SubNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::MulNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::DivNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::ModNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::FloorDivNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::FloorModNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::MinNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::MaxNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::EQNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::NENode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::LTNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::LENode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::GTNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::GENode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::AndNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::OrNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::CastNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::LShiftNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::RShiftNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::BitwiseAndNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::BitwiseOrNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::BitwiseXorNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::BitwiseNotNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::NotNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::SelectNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::RampNode* op, ffi::reflection::AccessPath path)
override;
- void Dispatch_(const prim::BroadcastNode* op, ffi::reflection::AccessPath
path) override;
- void Dispatch_(const prim::ShuffleNode* op, ffi::reflection::AccessPath
path) override;
-
- // Utility to call EnterDef/ExitDef. Used in the implementation of
- // WithDef.
- template <typename T>
- class DefContext {
- public:
- DefContext(DefContext&& other) : obj_(other.obj_) {
swap(std::move(other)); }
- DefContext& operator=(DefContext&& other) {
- swap(std::move(other));
- return *this;
- }
-
- DefContext(const DefContext&) = delete;
- DefContext& operator=(const DefContext&) = delete;
- ~DefContext() noexcept(false) {
- // Checks performed when a definition goes out of scope may
- // raise an exception. If the stack is already being unwound
- // due to another exception being thrown, this would cause a
- // segfault and terminate the program. By checking that no
- // additional exceptions have been thrown between the
- // construction of the DefContext and the destruction, we avoid
- // this case and allow the first error to propagate upward.
- if (self_ && std::uncaught_exceptions() == uncaught_exceptions_) {
- self_->in_scope_definitions_.erase(obj_);
- self_->ExitDef(obj_, path_);
- }
- }
-
- private:
- friend class TIRVisitorWithPath;
-
- DefContext(TIRVisitorWithPath* self, T obj, ffi::reflection::AccessPath
path, bool visit_type)
- : self_(self), obj_(obj), path_(path),
uncaught_exceptions_(std::uncaught_exceptions()) {
- if (auto var = obj_.template as<Var>(); var && visit_type) {
- self_->Visit(var.value()->ty, path_->Attr("ty"));
- }
- self_->in_scope_definitions_.insert(obj_);
- self_->EnterDef(obj_, path_);
- }
-
- void swap(DefContext&& other) {
- std::swap(this->self_, other.self_);
- std::swap(this->obj_, other.obj_);
- std::swap(this->path_, other.path_);
- std::swap(this->uncaught_exceptions_, other.uncaught_exceptions_);
- }
-
- TIRVisitorWithPath* self_{nullptr};
- T obj_;
- ffi::reflection::AccessPath path_{ffi::reflection::AccessPath::Root()};
- int uncaught_exceptions_{-1};
- };
-
- // Utility to track the scope of a node's definition.
- template <typename T>
- DefContext<T> WithDef(T obj, ffi::reflection::AccessPath path, bool
visit_type = true) {
- return DefContext(this, obj, path, visit_type);
- }
-
- /* \brief Utility to track the scope of a node's definition. */
- template <typename T>
- std::optional<DefContext<T>> WithDefIfUndefined(T obj,
ffi::reflection::AccessPath path) {
- if (in_scope_definitions_.count(obj)) {
- return std::nullopt;
- } else {
- return WithDef(obj, path);
- }
- }
-
- std::vector<DefContext<Var>> WithMatchTensorDefs(TensorVar buf,
- ffi::reflection::AccessPath
path) {
- std::vector<DefContext<Var>> context;
-
- auto try_visit_implicit_var_def = [this, &context](const Expr& expr,
-
ffi::reflection::AccessPath path) {
- if (auto opt = expr.as<Var>()) {
- auto var = opt.value();
- if (auto var_def = WithDefIfUndefined(var, path)) {
- context.push_back(std::move(var_def).value());
- }
- }
- };
-
- // A Tensor shape is a match scope. The first shape expression that
- // contains an undefined Var defines it, even when the expression is
- // compound (for example, `n + 1`). Later expressions then see the same
- // Var in `in_scope_definitions_` and reuse it.
- auto shape_path = path->Attr("shape");
- for (size_t i = 0; i < buf->shape.size(); i++) {
- auto dim_path = shape_path->ArrayItem(i);
- auto walk_fn = [this, &context, &dim_path](const Var& var) ->
ffi::Expected<ffi::WalkResult> {
- if (auto var_def = WithDefIfUndefined(var, dim_path)) {
- context.push_back(std::move(var_def).value());
- }
- return ffi::WalkResult::Advance();
- };
- ffi::StructuralWalk<ffi::WalkOrder::kPostOrder>(buf->shape[i], walk_fn);
- }
-
- auto strides_path = path->Attr("strides");
- for (size_t i = 0; i < buf->strides.size(); i++) {
- try_visit_implicit_var_def(buf->strides[i], strides_path->ArrayItem(i));
- }
- try_visit_implicit_var_def(buf->elem_offset, path->Attr("elem_offset"));
-
- return context;
- }
-
- std::unordered_set<ffi::ObjectRef, ffi::ObjectPtrHash, ffi::ObjectPtrEqual>
in_scope_definitions_;
-
- /*! \brief Scope stack for Bind variable definitions.
- *
- * Body-carrying statements (For, If, etc.) push a new scope.
- * BindNode pushes its WithDef into the current scope. When the
- * scope exits, all Bind defs are cleaned up automatically.
- */
- ScopeStack<std::vector<DefContext<Var>>> bind_scope_;
-};
-
-namespace {
-
-template <typename DerivedVerifier, typename PathVisitor = TIRVisitorWithPath>
-class Verifier : protected PathVisitor {
- public:
- template <typename TirNodeRef>
- static bool Verify(const TirNodeRef& node, bool assert_on_error) {
- DerivedVerifier verifier(assert_on_error);
- verifier(node);
- return !verifier.has_error_;
- }
-
- protected:
- explicit Verifier(bool assert_on_error) : assert_on_error_(assert_on_error)
{}
- void DispatchDefault_(const ffi::Object* op, ffi::reflection::AccessPath
path) override {
- Verify(false) << "TIR verifier does not support statement " <<
op->GetTypeKey() << " at "
- << path;
- }
-
- /* \brief Helper class to handle the bool-or-assert handles
- *
- * Each verifier can either return a boolean, or assert on failure.
- * To avoid needing to duplicate this logic at every step, the
- * Verify() method can be used. Similar to `LOG(FATAL)` or
- * `LOG(DEBUG)`, it returns an object that can accept streamed
- * context information.
- *
- * If the error should be raised, then the context is collected
- * identically to `LOG(FATAL)`. If a boolean is returned, or if the
- * condition passes, then the streamed context is discarded.
- *
- * Usage:
- *
- * Verify(value == expected_value)
- * << "ValueError: " << value
- * << " was not the expected value of " << expected_value;
- */
- class VerifyStream {
- public:
- explicit VerifyStream(bool log_fatal) {
- if (log_fatal) {
- log_.emplace();
- }
- }
-
- VerifyStream(const VerifyStream&) = delete;
- VerifyStream& operator=(const VerifyStream&) = delete;
- VerifyStream(VerifyStream&& other) { std::swap(log_, other.log_); }
- VerifyStream& operator=(VerifyStream&& other) {
- std::swap(log_, other.log_);
- return *this;
- }
-
- template <typename T>
- VerifyStream& operator<<(T&& t) {
- if (log_.has_value()) {
- log_.value() << std::forward<T>(t);
- }
- return *this;
- }
-
- ~VerifyStream() noexcept(false) {
- if (log_.has_value()) {
- LOG(FATAL) << log_->str();
- }
- }
-
- std::optional<std::ostringstream> log_{std::nullopt};
- };
-
- // TODO(Lunderberg): Add the filename/linenum with
- // std::source_location when C++20 is available.
- VerifyStream Verify(bool condition) {
- has_error_ = has_error_ || !condition;
- return VerifyStream(!condition && assert_on_error_);
- }
-
- bool assert_on_error_;
- bool has_error_{false};
-};
-
-} // namespace
-
-} // namespace tirx
-} // namespace tvm
-#endif // TVM_TIR_IR_TIR_VISITOR_WITH_PATH_H_
diff --git a/src/tirx/tile/lower.cc b/src/tirx/tile/lower.cc
index aaabf822d5..7a3b436851 100644
--- a/src/tirx/tile/lower.cc
+++ b/src/tirx/tile/lower.cc
@@ -49,29 +49,125 @@
#include <vector>
#include "../../backend/cuda/tile/exec_analysis.h"
-#include "../ir/tir_visitor_with_path.h"
namespace tvm {
namespace tirx {
namespace {
-class NoOpCallVerifier : public Verifier<NoOpCallVerifier> {
+class NoOpCallVerifier : public StmtExprVisitor {
+ friend class tvm::ObjectVisitor;
+ using Result = ffi::Optional<VisitInterrupt>;
+
public:
- using Verifier::Verifier;
+ NoOpCallVerifier() : StmtExprVisitor(GlobalVTable()) {}
+ static bool Verify(const Stmt& stmt) {
+ auto verifier = ffi::make_object<NoOpCallVerifier>();
+ verifier->Visit(stmt);
+ return verifier->valid_;
+ }
private:
- using Verifier::Visit;
+ using StmtExprVisitor::Visit_;
+ Result Visit(ffi::AnyView value) override {
+ if (value.as<Stmt>() && !GlobalVTable()->CanDispatch(value)) {
+ valid_ = false;
+ return std::nullopt;
+ }
+ if (auto expr = value.as<Expr>(); expr &&
!GlobalVTable()->CanDispatch(value)) {
+ if (value.as<OpaqueExprNode>() || value.as<VarNode>()) return
std::nullopt;
+ TVM_FFI_VISIT_THROW(TypeError, expr.value())
+ << "Unsupported non-primitive TIR expression " <<
expr.value()->GetTypeKey();
+ }
+ return StmtExprVisitor::Visit(value);
+ }
- void Dispatch_(const CallNode* call, ffi::reflection::AccessPath path) final
{
+ Result Visit_(const CallNode* call) final {
if (auto op = call->op.as<Op>(); op &&
Op::HasAttrMap(tile::op_attr::kTileOpDispatch)) {
static const auto& dispatchers =
Op::GetAttrMap<tile::FTileOpDispatch>(tile::op_attr::kTileOpDispatch);
- Verify(!dispatchers.count(op.value()))
- << "Unlowered tensor instruction " << op.value()->name << " at " <<
path;
+ valid_ &= !dispatchers.count(op.value());
}
- Verifier::Dispatch_(call, path);
+ if (call->op.as<GlobalVar>() || call->op.as<OpaqueExprNode>())
this->Visit(call->op);
+ return this->Visit(call->args);
+ }
+
+ Result Visit_(const BindNode* op) override {
+ this->Visit(op->value);
+ return this->Visit(op->var->ty);
+ }
+
+ Result Visit_(const RegionStmtNode* op) override {
+ this->Visit(op->args);
+ this->Visit(op->attrs);
+ for (const auto& param : op->body_params) this->Visit(param->ty);
+ this->Visit(op->body);
+ for (const auto& var : op->result_vars) this->Visit(var->ty);
+ return std::nullopt;
+ }
+
+ Result Visit_(const ForNode* op) override {
+ this->Visit(op->min);
+ this->Visit(op->extent);
+ this->Visit(op->loop_var->ty);
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const prim::LetNode* op) override {
+ this->Visit(op->value);
+ this->Visit(op->var->ty);
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const LambdaExprNode* op) {
+ this->Visit(op->ty);
+ for (const auto& var : op->vars) this->Visit(var->ty);
+ return this->Visit(op->body);
+ }
+
+ Result Visit_(const TensorLoadNode* op) override { return
this->Visit(op->indices); }
+
+ Result Visit_(const TensorStoreNode* op) override {
+ this->Visit(op->value);
+ return this->Visit(op->indices);
+ }
+
+ Result Visit_(const TensorRegionNode* op) override {
+ if (!op->source.as<TensorVar>()) this->Visit(op->source);
+ return this->Visit(op->region);
+ }
+
+ Result Visit_(const prim::RampNode* op) override {
+ this->Visit(op->base);
+ this->Visit(op->stride);
+ return this->Visit(op->lanes);
+ }
+
+ Result Visit_(const prim::BroadcastNode* op) override {
+ this->Visit(op->value);
+ return this->Visit(op->lanes);
+ }
+
+ Result Visit_(const GenericConstNode* op) override { return std::nullopt; }
+
+ Result Visit_(const AssertStmtNode* op) override {
+ this->Visit(op->condition);
+ this->Visit(op->error_kind);
+ return this->Visit(op->message_parts);
}
+
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ StmtExprVisitor::InitVTable(&table);
+ SetDispatch<NoOpCallVerifier, LambdaExprNode>(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+
+ bool valid_{true};
};
class TileDispatcher : public StmtExprMutator {
@@ -543,7 +639,7 @@ Pass TileDispatch() {
Target target = ResolveTarget(f);
auto* n = f.CopyOnWrite();
n->body = TileDispatcher::LowerOpCalls(n->body.value(), target);
- if (!NoOpCallVerifier::Verify(n->body.value(), false)) {
+ if (!NoOpCallVerifier::Verify(n->body.value())) {
LOG(FATAL) << "Failed to lower the TIRx program: " << f;
}
return f;
