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 e11e355ed3 [REFACTOR][TIRX] Migrate StmtFunctor dispatch and organize
analyzer helpers (#20370)
e11e355ed3 is described below
commit e11e355ed3dd015b80b3ee618a580dab70a57fc4
Author: Tianqi Chen <[email protected]>
AuthorDate: Thu Sep 17 07:14:18 2026 -0400
[REFACTOR][TIRX] Migrate StmtFunctor dispatch and organize analyzer helpers
(#20370)
Give TIRx statement functors the shared `Dispatch` API and extensible,
finalized dispatch tables. Update path visitors, code generators, FLOP
estimation, and tensorization comparators while preserving their
explicit traversal and backend algorithms.
Move analyzer-aware visitor and mutator helpers into `src/tirx/ir`
alongside the other IR utilities.
---
include/tvm/tirx/stmt_functor.h | 204 +++++++++++++--------
src/backend/cuda/codegen/codegen_cuda.cc | 26 +--
src/backend/cuda/codegen/codegen_cuda.h | 12 +-
src/backend/cuda/codegen/llvm/codegen_nvptx.cc | 2 +-
.../hexagon/codegen/llvm/codegen_hexagon.cc | 2 +-
src/backend/metal/codegen/codegen_metal.cc | 6 +-
src/backend/metal/codegen/codegen_metal.h | 4 +-
src/backend/opencl/codegen/codegen_opencl.cc | 4 +-
src/backend/opencl/codegen/codegen_opencl.h | 2 +-
src/backend/rocm/codegen/llvm/codegen_amdgpu.cc | 2 +-
src/backend/trn/codegen/codegen_trn.cc | 16 +-
src/backend/trn/codegen/codegen_trn.h | 15 +-
src/backend/trn/transform/lower_trainium_layout.cc | 2 +-
src/backend/vulkan/codegen/codegen_spirv.cc | 38 ++--
src/backend/vulkan/codegen/codegen_spirv.h | 25 +--
src/backend/webgpu/codegen/codegen_webgpu.cc | 16 +-
src/backend/webgpu/codegen/codegen_webgpu.h | 16 +-
src/relax/transform/split_call_tir_by_pattern.cc | 30 +--
src/s_tir/analysis/domain_touched.cc | 2 +-
src/s_tir/analysis/estimate_flops.cc | 50 ++---
src/s_tir/analysis/identify_memcpy.cc | 2 +-
src/s_tir/analysis/is_pure_function.cc | 8 +-
src/s_tir/analysis/oob_checker.cc | 2 +-
src/s_tir/backend/adreno/inject_texture_alloc.cc | 2 +-
src/s_tir/backend/adreno/texture_flatten.cc | 2 +-
src/s_tir/schedule/analysis/analysis.cc | 2 +-
src/s_tir/schedule/ir_comparator.cc | 36 ++--
src/s_tir/schedule/ir_comparator.h | 18 +-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 2 +-
.../schedule/primitive/layout_transformation.cc | 2 +-
src/s_tir/schedule/transform.h | 2 +-
src/s_tir/transform/hoist_expression.cc | 2 +-
src/s_tir/transform/inject_permuted_layout.cc | 2 +-
src/s_tir/transform/inject_virtual_thread.cc | 2 +-
src/s_tir/transform/lower_async_dma.cc | 2 +-
src/s_tir/transform/renormalize_split_pattern.cc | 2 +-
.../transform/using_assume_to_reduce_branches.cc | 2 +-
src/target/llvm/codegen_aarch64.cc | 10 +-
src/target/llvm/codegen_cpu.cc | 24 +--
src/target/llvm/codegen_cpu.h | 7 +-
src/target/llvm/codegen_llvm.cc | 40 ++--
src/target/llvm/codegen_llvm.h | 27 +--
src/target/source/codegen_c.cc | 28 +--
src/target/source/codegen_c.h | 33 ++--
src/target/source/codegen_c_host.cc | 2 +-
src/target/source/codegen_c_host.h | 2 +-
src/tirx/analysis/verify_tirx_well_formed.cc | 29 ++-
src/tirx/{ => ir}/ir_mutator_with_analyzer.cc | 4 +-
src/tirx/{ => ir}/ir_mutator_with_analyzer.h | 10 +-
src/tirx/{ => ir}/ir_visitor_with_analyzer.cc | 2 +-
src/tirx/{ => ir}/ir_visitor_with_analyzer.h | 8 +-
src/tirx/ir/tir_visitor_with_path.cc | 36 ++--
src/tirx/ir/tir_visitor_with_path.h | 42 ++---
src/tirx/transform/flatten_buffer.cc | 2 +-
src/tirx/transform/lower_intrin.cc | 2 +-
src/tirx/transform/lower_tirx_cleanup.cc | 2 +-
src/tirx/transform/remove_no_op.cc | 2 +-
src/tirx/transform/stmt_simplify.cc | 2 +-
src/tirx/transform/tile_primitive_dispatch.cc | 4 +-
tests/cpp/ir_functor_test.cc | 6 +-
tests/cpp/stmt_functor_test.cc | 131 +++++++++++++
61 files changed, 606 insertions(+), 413 deletions(-)
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 61c17cf574..c1147f5f80 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -40,105 +40,149 @@
namespace tvm {
namespace tirx {
/*!
- * \brief Same as ExprFunctor except it is applied on statements
+ * \brief Type-dispatched statement functor with a caller-selected signature.
+ *
+ * Override Dispatch_ for a node type or DispatchDefault_ for default behavior.
+ * This functor does not traverse children automatically. Dispatch may use a
+ * registered ancestor. Derived extensions can initialize a fresh inherited
+ * table with InitVTable and register additional hooks with SetDispatch.
* \tparam FType The function signature.
* \sa ExprFunctor
*/
template <typename FType>
class StmtFunctor;
-#define STMT_FUNCTOR_DEFAULT \
- { \
- return VisitStmtDefault_(op, std::forward<Args>(args)...); \
- }
-
-#define IR_STMT_FUNCTOR_DISPATCH(OP)
\
- vtable.template SetDispatch<OP>([](const ffi::ObjectRef& n, TSelf* self,
Args... args) { \
- return self->VisitStmt_(static_cast<const OP*>(n.get()),
std::forward<Args>(args)...); \
- });
-
template <typename R, typename... Args>
-class StmtFunctor<R(const Stmt& n, Args... args)> {
+class StmtFunctor<R(const Stmt&, Args...)> {
private:
- using TSelf = StmtFunctor<R(const Stmt& n, Args... args)>;
- using FType = ObjectFunctor<R(const ffi::ObjectRef& n, TSelf* self, Args...
args)>;
+ using TSelf = StmtFunctor<R(const Stmt&, Args...)>;
public:
- /*! \brief the result type of this functor */
+ /*! \brief The result type of this functor. */
using result_type = R;
- /*! \brief virtual destructor */
- virtual ~StmtFunctor() {}
- /*!
- * \brief Same as call.
- * \param n The stmt node.
- * \param args Additional arguments.
- * \return The result of the call
- */
- R operator()(const Stmt& n, Args... args) { return VisitStmt(n,
std::forward<Args>(args)...); }
- /*!
- * \brief The functor call.
- * \param n The stmt node.
- * \param args Additional arguments.
- * \return The result of the call
- */
- virtual R VisitStmt(const Stmt& n, Args... args) {
- static FType vtable = InitVTable();
- return vtable(n, this, std::forward<Args>(args)...);
+ /*! \brief Construct a functor with the TIRx statement hooks. */
+ StmtFunctor() : StmtFunctor(GlobalVTable()) {}
+ /*! \brief Destroy through the statement functor base. */
+ virtual ~StmtFunctor() = default;
+ /*! \brief Dispatch a statement, forwarding additional arguments to its
hook. */
+ TVM_FFI_INLINE R operator()(const Stmt& node, Args... args) {
+ return Dispatch(node, std::forward<Args>(args)...);
+ }
+ /*! \brief Dispatch to a node hook, including registered ancestor hooks. */
+ TVM_FFI_INLINE virtual R Dispatch(const Stmt& node, Args... args) {
+ TVM_FFI_ICHECK(node.defined()) << "Cannot dispatch a null statement";
+ return (*vtable_)(node, this, std::forward<Args>(args)...);
+ }
+
+ virtual R Dispatch_(const BindNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const AttrStmtNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const IfThenElseNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const ForNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const WhileNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const ReturnNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const BreakNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const ContinueNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
}
- // Functions that can be overriden by subclass
- virtual R VisitStmt_(const BindNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const AttrStmtNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const IfThenElseNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const ForNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const WhileNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const ReturnNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const BreakNode* op, Args... args) STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const ContinueNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const AllocBufferNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const DeclBufferNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const BufferStoreNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const AssertStmtNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const SeqStmtNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const EvaluateNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const SBlockNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const SBlockRealizeNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const ScopeIdDefStmtNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmt_(const tirx::TilePrimitiveCallNode* op, Args... args)
STMT_FUNCTOR_DEFAULT;
- virtual R VisitStmtDefault_(const ffi::Object* op, Args...) {
- TVM_FFI_THROW(InternalError) << "Do not have a default for " <<
op->GetTypeKey();
+ virtual R Dispatch_(const AllocBufferNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const DeclBufferNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const BufferStoreNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const AssertStmtNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const SeqStmtNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const EvaluateNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const SBlockNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const SBlockRealizeNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const ScopeIdDefStmtNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const tirx::TilePrimitiveCallNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ /*! \brief Default behavior for statement hooks not overridden by a
subclass. */
+ virtual R DispatchDefault_(const ffi::Object* node, Args...) {
+ TVM_FFI_THROW(InternalError) << "Do not have a default for " <<
node->GetTypeKey();
TVM_FFI_UNREACHABLE();
}
+ protected:
+ /*! \brief Dispatch table shared by this signature and its subclasses. */
+ using VTable = ObjectFunctor<R(const ffi::ObjectRef&, TSelf*, Args...)>;
+ /*! \brief Construct with a finalized table that outlives the functor. */
+ explicit StmtFunctor(const VTable* vtable) : vtable_(vtable) {}
+ /*! \brief Register statement hooks in a fresh mutable table. */
+ static void InitVTable(VTable* vtable) {
+ SetDispatch<TSelf, BindNode>(vtable);
+ SetDispatch<TSelf, AttrStmtNode>(vtable);
+ SetDispatch<TSelf, IfThenElseNode>(vtable);
+ SetDispatch<TSelf, ForNode>(vtable);
+ SetDispatch<TSelf, WhileNode>(vtable);
+ SetDispatch<TSelf, ReturnNode>(vtable);
+ SetDispatch<TSelf, BreakNode>(vtable);
+ SetDispatch<TSelf, ContinueNode>(vtable);
+ SetDispatch<TSelf, AllocBufferNode>(vtable);
+ SetDispatch<TSelf, DeclBufferNode>(vtable);
+ SetDispatch<TSelf, BufferStoreNode>(vtable);
+ SetDispatch<TSelf, AssertStmtNode>(vtable);
+ SetDispatch<TSelf, SeqStmtNode>(vtable);
+ SetDispatch<TSelf, EvaluateNode>(vtable);
+ SetDispatch<TSelf, SBlockNode>(vtable);
+ SetDispatch<TSelf, SBlockRealizeNode>(vtable);
+ SetDispatch<TSelf, ScopeIdDefStmtNode>(vtable);
+ SetDispatch<TSelf, tirx::TilePrimitiveCallNode>(vtable);
+ }
+ /*! \brief Register an additional node hook implemented by Self. */
+ template <typename Self, typename Node>
+ static void SetDispatch(VTable* vtable) {
+ vtable->template SetDispatch<Node>(
+ [](const ffi::ObjectRef& node, TSelf* self, Args... args) -> R {
+ return static_cast<Self*>(self)->Dispatch_(static_cast<const
Node*>(node.get()),
+
std::forward<Args>(args)...);
+ });
+ }
+
private:
- // initialize the vtable.
- static FType InitVTable() {
- FType vtable;
- IR_STMT_FUNCTOR_DISPATCH(BindNode);
- IR_STMT_FUNCTOR_DISPATCH(AttrStmtNode);
- IR_STMT_FUNCTOR_DISPATCH(IfThenElseNode);
- IR_STMT_FUNCTOR_DISPATCH(ForNode);
- IR_STMT_FUNCTOR_DISPATCH(WhileNode);
- IR_STMT_FUNCTOR_DISPATCH(ReturnNode);
- IR_STMT_FUNCTOR_DISPATCH(BreakNode);
- IR_STMT_FUNCTOR_DISPATCH(ContinueNode);
- IR_STMT_FUNCTOR_DISPATCH(AllocBufferNode);
- IR_STMT_FUNCTOR_DISPATCH(DeclBufferNode);
- IR_STMT_FUNCTOR_DISPATCH(AssertStmtNode);
- IR_STMT_FUNCTOR_DISPATCH(SeqStmtNode);
- IR_STMT_FUNCTOR_DISPATCH(EvaluateNode);
- IR_STMT_FUNCTOR_DISPATCH(BufferStoreNode);
- IR_STMT_FUNCTOR_DISPATCH(SBlockNode);
- IR_STMT_FUNCTOR_DISPATCH(SBlockRealizeNode);
- IR_STMT_FUNCTOR_DISPATCH(ScopeIdDefStmtNode);
- IR_STMT_FUNCTOR_DISPATCH(tirx::TilePrimitiveCallNode);
- vtable.Finalize();
- return vtable;
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
}
+ const VTable* vtable_;
};
-#undef IR_STMT_FUNCTOR_DISPATCH
-#undef STMT_FUNCTOR_DEFAULT
-
/*!
* \brief Native visitor for TIRx statements and their expression operands.
*
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index d60eea8d29..c35e4d1a6d 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -305,10 +305,10 @@ void CodeGenCUDA::PrintExtraAttrs(const PrimFunc& f,
std::ostream& os) {
}
}
-void CodeGenCUDA::VisitStmt_(const ReturnNode* op) {
+void CodeGenCUDA::Dispatch_(const ReturnNode* op) {
if (!in_kernel_launch_) {
// __device__ subroutines return real values.
- CodeGenC::VisitStmt_(op);
+ CodeGenC::Dispatch_(op);
return;
}
const auto* value = op->value.as<IntImmNode>();
@@ -336,7 +336,7 @@ std::string CodeGenCUDA::Finish() {
return CodeGenC::Finish();
}
-void CodeGenCUDA::VisitStmt_(const tirx::ForNode* op) {
+void CodeGenCUDA::Dispatch_(const tirx::ForNode* op) {
// Materialize the loop bounds before emitting an unroll pragma. PrintExpr
// may introduce temporaries (for example, for a Select expression). CUDA
// requires #pragma unroll to immediately precede the loop it controls; if
@@ -380,7 +380,7 @@ void CodeGenCUDA::VisitStmt_(const tirx::ForNode* op) {
stream << "}\n";
}
-void CodeGenCUDA::VisitStmt_(const WhileNode* op) {
+void CodeGenCUDA::Dispatch_(const WhileNode* op) {
PrintIndent();
// Match CodeGenC: dynamic-trip-count loops must not be unrolled.
stream << "#pragma unroll 1\n";
@@ -1612,7 +1612,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
}
}
-void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenCUDA::Dispatch_(const AttrStmtNode* op) {
if (op->attr_key == s_tir::attr::fragment_shape) {
const VarNode* buffer = op->node.as<VarNode>();
const prim::StringImmNode* shape_str = op->value.as<prim::StringImmNode>();
@@ -1625,7 +1625,7 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) {
const IntImmNode* queue_id = op->value.as<IntImmNode>();
TVM_FFI_ICHECK(queue_id && queue_id->value == 0)
<< "For CUDA, the index of an async queue must be 0.";
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
static const Op& ptx_cp_async_commit_group_op =
Op::Get("tirx.ptx.cp_async_commit_group");
// ptx Call layout: [operands...] [slot tokens] [pred marker ""].
auto commit_group =
@@ -1655,12 +1655,12 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) {
this->stream << ";\n";
auto inner = op->body.as<AttrStmtNode>();
TVM_FFI_ICHECK(inner);
- this->VisitStmt(inner->body);
+ this->Dispatch(inner->body);
return;
} else if (op->attr_key == "disable_unroll") {
PrintIndent();
stream << "#pragma unroll 1\n";
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
return;
} else if (op->attr_key == "pragma_unroll") {
PrintIndent();
@@ -1669,14 +1669,14 @@ void CodeGenCUDA::VisitStmt_(const AttrStmtNode* op) {
stream << " " << count->value;
}
stream << "\n";
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
return;
} else if (op->attr_key == tirx::attr::thread_extent) {
}
- CodeGenC::VisitStmt_(op);
+ CodeGenC::Dispatch_(op);
}
-void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenCUDA::Dispatch_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get(), op->buffer.name() + "_ptr");
@@ -1748,7 +1748,7 @@ void CodeGenCUDA::VisitStmt_(const AllocBufferNode* op) {
}
}
-void CodeGenCUDA::VisitStmt_(const EvaluateNode* op) {
+void CodeGenCUDA::Dispatch_(const EvaluateNode* op) {
if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
const CallNode* call = op->value.as<CallNode>();
if (call && call->op.same_as(tirx::builtin::tvm_global_barrier_kinit())) {
@@ -1761,7 +1761,7 @@ void CodeGenCUDA::VisitStmt_(const EvaluateNode* op) {
PrintIndent();
stream << "}\n";
} else {
- CodeGenC::VisitStmt_(op);
+ CodeGenC::Dispatch_(op);
}
}
diff --git a/src/backend/cuda/codegen/codegen_cuda.h
b/src/backend/cuda/codegen/codegen_cuda.h
index fa69f4b20f..c7703f9e0b 100644
--- a/src/backend/cuda/codegen/codegen_cuda.h
+++ b/src/backend/cuda/codegen/codegen_cuda.h
@@ -52,8 +52,8 @@ class CodeGenCUDA final : public CodeGenC {
void PrintFunctionSignature(const ffi::String& function_name, const
PrimFunc& func,
std::ostream& os) final;
void PrintExtraAttrs(const PrimFunc& f, std::ostream& os) final; //
NOLINT(*)
- void VisitStmt_(const ForNode* op) final;
- void VisitStmt_(const WhileNode* op) final;
+ void Dispatch_(const ForNode* op) final;
+ void Dispatch_(const WhileNode* op) final;
void PrintStorageSync(const CallNode* op) final;
void PrintStorageScope(const std::string& scope, std::ostream& os) final;
// NOLINT(*)
using CodeGenC::PrintType;
@@ -77,10 +77,10 @@ class CodeGenCUDA final : public CodeGenC {
void Dispatch_(const FloatImmNode* op, std::ostream& os) final;
void Dispatch_(const CallNode* op, std::ostream& os) final;
void Dispatch_(const prim::CastNode* op, std::ostream& os) final;
- void VisitStmt_(const EvaluateNode* op) final;
- void VisitStmt_(const ReturnNode* op) final;
- void VisitStmt_(const AllocBufferNode* op) final;
- void VisitStmt_(const AttrStmtNode* op) final;
+ void Dispatch_(const EvaluateNode* op) final;
+ void Dispatch_(const ReturnNode* op) final;
+ void Dispatch_(const AllocBufferNode* op) final;
+ void Dispatch_(const AttrStmtNode* op) final;
// Target
Target target;
diff --git a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
index e248fed902..6d2bdd0b63 100644
--- a/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
+++ b/src/backend/cuda/codegen/llvm/codegen_nvptx.cc
@@ -78,7 +78,7 @@ class CodeGenNVPTX : public CodeGenLLVM {
llvm::ValueAsMetadata::get(ConstInt32(1))}));
}
- void VisitStmt_(const AllocBufferNode* op) final {
+ void Dispatch_(const AllocBufferNode* op) final {
llvm::Value* buf = nullptr;
StorageInfo& info = alloc_storage_info_[op->buffer.get()];
// maximum necessary alignment in the NV devices
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index 61173f52b8..97083d33f0 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -79,7 +79,7 @@ class CodeGenHexagon final : public CodeGenCPU {
bool target_c_runtime) override;
void InitTarget() final;
- using CodeGenCPU::VisitStmt_;
+ using CodeGenCPU::Dispatch_;
llvm::Value* Dispatch_(const TensorLoadNode* op) override;
llvm::Value* CreateIntrinsic(const CallNode* op) override;
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index e25d33db26..a40b188c37 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -330,10 +330,10 @@ void CodeGenMetal::PrintStorageScope(const std::string&
scope, std::ostream& os)
}
}
-void CodeGenMetal::VisitStmt_(const BindNode* op) {
+void CodeGenMetal::Dispatch_(const BindNode* op) {
const auto* pointer_type = op->var->ty.as<PointerTypeNode>();
if (pointer_type == nullptr || pointer_type->storage_scope.empty()) {
- return CodeGenC::VisitStmt_(op);
+ return CodeGenC::Dispatch_(op);
}
const std::string& storage_scope = pointer_type->storage_scope;
@@ -355,7 +355,7 @@ void CodeGenMetal::VisitStmt_(const BindNode* op) {
stream << "*)" << value << ";\n";
}
-void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenMetal::Dispatch_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get());
diff --git a/src/backend/metal/codegen/codegen_metal.h
b/src/backend/metal/codegen/codegen_metal.h
index 23991df89d..7fad4eb398 100644
--- a/src/backend/metal/codegen/codegen_metal.h
+++ b/src/backend/metal/codegen/codegen_metal.h
@@ -52,8 +52,8 @@ class CodeGenMetal final : public CodeGenC {
void PrintVecElemStore(const std::string& vec, const PrimType& t, int i,
const std::string& value) final;
// overload visitor
- void VisitStmt_(const BindNode* op) final; //
NOLINT(*)
- void VisitStmt_(const AllocBufferNode* op) final; //
NOLINT(*)
+ void Dispatch_(const BindNode* op) final; //
NOLINT(*)
+ void Dispatch_(const AllocBufferNode* op) final; //
NOLINT(*)
void Dispatch_(const prim::SelectNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const prim::BroadcastNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const CallNode* op, std::ostream& os) final; //
NOLINT(*)
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc
b/src/backend/opencl/codegen/codegen_opencl.cc
index 692b7cd0a0..8aecd841a5 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -445,7 +445,7 @@ std::string CodeGenOpenCL::CastTo(std::string value, const
PrimType& target) {
}
}
-void CodeGenOpenCL::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenOpenCL::Dispatch_(const AllocBufferNode* op) {
// Compute constant_size from buffer shape
size_t constant_size = 1;
for (const auto& dim : op->buffer->shape) {
@@ -454,7 +454,7 @@ void CodeGenOpenCL::VisitStmt_(const AllocBufferNode* op) {
constant_size *= dim_imm->value;
}
allocation_size_.insert({op->buffer.get(), constant_size *
op->buffer->dtype.lanes()});
- CodeGenC::VisitStmt_(op);
+ CodeGenC::Dispatch_(op);
}
void CodeGenOpenCL::Dispatch_(const CallNode* op, std::ostream& os) {
diff --git a/src/backend/opencl/codegen/codegen_opencl.h
b/src/backend/opencl/codegen/codegen_opencl.h
index b3b20ea8b4..6de95420f7 100644
--- a/src/backend/opencl/codegen/codegen_opencl.h
+++ b/src/backend/opencl/codegen/codegen_opencl.h
@@ -64,7 +64,7 @@ class CodeGenOpenCL final : public CodeGenC {
void SetTextureScope(const std::unordered_map<const VarNode*,
std::string>&); // NOLINT(*)
// overload visitor
- void VisitStmt_(const AllocBufferNode* op) final; //
NOLINT(*)
+ void Dispatch_(const AllocBufferNode* op) final; //
NOLINT(*)
void Dispatch_(const prim::BroadcastNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const prim::RampNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const CallNode* op, std::ostream& os) final; //
NOLINT(*)
diff --git a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
index 3272835f33..1e91582e56 100644
--- a/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
+++ b/src/backend/rocm/codegen/llvm/codegen_amdgpu.cc
@@ -96,7 +96,7 @@ class CodeGenAMDGPU : public CodeGenLLVM {
function_->addFnAttr("amdgpu-flat-work-group-size", attr.str());
}
- void VisitStmt_(const AllocBufferNode* op) final {
+ void Dispatch_(const AllocBufferNode* op) final {
llvm::Value* buf = nullptr;
StorageInfo& info = alloc_storage_info_[op->buffer.get()];
auto storage_scope =
runtime::StorageScope::Create(GetPtrStorageScope(op->buffer.var()));
diff --git a/src/backend/trn/codegen/codegen_trn.cc
b/src/backend/trn/codegen/codegen_trn.cc
index bf36859bf6..ee8dbe5f6c 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -214,7 +214,7 @@ std::string CodeGenTrainium::GetStorageScopeStr(const
std::string& scope) { //
}
}
-void CodeGenTrainium::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenTrainium::Dispatch_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get(), op->buffer.name() + "_ptr");
@@ -262,7 +262,7 @@ void CodeGenTrainium::VisitStmt_(const AllocBufferNode* op)
{
}
}
-void CodeGenTrainium::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenTrainium::Dispatch_(const AttrStmtNode* op) {
if (op->attr_key == tirx::attr::tensorized_nki_instruction) {
ctx_.tensorizing = true;
ctx_.mask = PrimExpr(nullptr);
@@ -275,7 +275,7 @@ void CodeGenTrainium::VisitStmt_(const AttrStmtNode* op) {
}
}
-void CodeGenTrainium::VisitStmt_(const ForNode* op) {
+void CodeGenTrainium::Dispatch_(const ForNode* op) {
bool is_outermost_loop = is_outermost_loop_;
is_outermost_loop_ = false;
std::string extent = PrintExpr(op->extent);
@@ -335,11 +335,11 @@ std::string CodeGenTrainium::PrintIndices(const
Array<PrimExpr>& indices) {
return os.str();
}
-void CodeGenTrainium::VisitStmt_(const BufferStoreNode* op) {
+void CodeGenTrainium::Dispatch_(const BufferStoreNode* op) {
LOG(FATAL) << "Trainium codegen does not support buffer store";
}
-void CodeGenTrainium::VisitStmt_(const EvaluateNode* op) {
+void CodeGenTrainium::Dispatch_(const EvaluateNode* op) {
if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
std::string vid = this->PrintExpr(op->value);
if (vid != "") {
@@ -614,7 +614,7 @@ void CodeGenTrainium::Dispatch_(const prim::FloorModNode*
op, std::ostream& os)
os << PrintExpr(op->a) << " % " << PrintExpr(op->b);
}
-void CodeGenTrainium::VisitStmt_(const DeclBufferNode* op) {
+void CodeGenTrainium::Dispatch_(const DeclBufferNode* op) {
if (op->buffer.scope() == "trn.psum" || op->buffer.scope() == "trn.sbuf") {
return;
}
@@ -684,12 +684,12 @@ ffi::Module BuildTrainium(IRModule mod, Target target) {
return codegen::DeviceSourceModuleCreate(source_maker.str(), fmt,
ExtractFuncInfo(mod), "nki");
}
-void CodeGenTrainium::VisitStmt_(const IfThenElseNode* op) {
+void CodeGenTrainium::Dispatch_(const IfThenElseNode* op) {
if (ctx_.tensorizing) {
TVM_FFI_ICHECK(!op->else_case.has_value()) << "Else not allowed in
tensorized instruction";
TVM_FFI_ICHECK(!ctx_.mask.defined()) << "Only one if stmt allowed in
tensorized instruction";
ctx_.mask = op->condition;
- VisitStmt(op->then_case);
+ Dispatch(op->then_case);
return;
}
std::string cond = PrintExpr(op->condition);
diff --git a/src/backend/trn/codegen/codegen_trn.h
b/src/backend/trn/codegen/codegen_trn.h
index 429c5b9fb5..7e216f96c3 100644
--- a/src/backend/trn/codegen/codegen_trn.h
+++ b/src/backend/trn/codegen/codegen_trn.h
@@ -51,7 +51,6 @@ class CodeGenTrainium final : public CodeGenC {
explicit CodeGenTrainium(Target target);
using CodeGenC::Dispatch_;
using CodeGenC::PrintType;
- using CodeGenC::VisitStmt_;
// override print thread tag.
void PrintArgUnionDecl();
void AddFunction(const GlobalVar& gvar, const PrimFunc& func) final;
@@ -59,11 +58,11 @@ class CodeGenTrainium final : public CodeGenC {
std::string GetStorageScopeStr(const std::string& scope); //
NOLINT(*)
void Dispatch_(const VarNode* op, std::ostream& os) final; //
NOLINT(*)
void PrintType(const PrimType& t, std::ostream& os) final; //
NOLINT(*)
- void VisitStmt_(const AllocBufferNode* op) final; //
NOLINT(*)
- void VisitStmt_(const AttrStmtNode* op) final; //
NOLINT(*)
- void VisitStmt_(const ForNode* op) final; //
NOLINT(*)
- void VisitStmt_(const BufferStoreNode* op) final; //
NOLINT(*)=
- void VisitStmt_(const EvaluateNode* op) final; //
NOLINT(*)
+ void Dispatch_(const AllocBufferNode* op) final; //
NOLINT(*)
+ void Dispatch_(const AttrStmtNode* op) final; //
NOLINT(*)
+ void Dispatch_(const ForNode* op) final; //
NOLINT(*)
+ void Dispatch_(const BufferStoreNode* op) final; //
NOLINT(*)=
+ void Dispatch_(const EvaluateNode* op) final; //
NOLINT(*)
std::string PrintIndices(const ffi::Array<PrimExpr>& indices); //
NOLINT(*)
void Dispatch_(const TensorLoadNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const CallNode* op, std::ostream& os) final; //
NOLINT(*)
@@ -71,8 +70,8 @@ class CodeGenTrainium final : public CodeGenC {
void Dispatch_(const prim::CastNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const prim::FloorDivNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const prim::FloorModNode* op, std::ostream& os) final; //
NOLINT(*)
- void VisitStmt_(const DeclBufferNode* op) final; //
NOLINT(*)
- void VisitStmt_(const IfThenElseNode* op) final; //
NOLINT(*)
+ void Dispatch_(const DeclBufferNode* op) final; //
NOLINT(*)
+ void Dispatch_(const IfThenElseNode* op) final; //
NOLINT(*)
void Dispatch_(const prim::AndNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const prim::OrNode* op, std::ostream& os) final; //
NOLINT(*)
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc
b/src/backend/trn/transform/lower_trainium_layout.cc
index 30ed8a5bf4..6d44daf95f 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -37,7 +37,7 @@
#include <utility>
#include <vector>
-#include "../../../tirx/ir_mutator_with_analyzer.h"
+#include "../../../tirx/ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace tirx {
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 45bfe016b2..62892c96f3 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -137,7 +137,7 @@ runtime::SPIRVShader CodeGenSPIRV::BuildFunction(const
PrimFunc& f, const std::s
}
}
}
- this->VisitStmt(f->body);
+ this->Dispatch(f->body);
builder_->SetLocalSize(func_ptr, workgroup_size_);
builder_->MakeInst(spv::OpReturn);
builder_->MakeInst(spv::OpFunctionEnd);
@@ -684,7 +684,7 @@ spirv::Value CodeGenSPIRV::Dispatch_(const
prim::ShuffleNode* op) {
return element;
}
-void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
+void CodeGenSPIRV::Dispatch_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "SPIR-V codegen expects flat
memory buffers";
Var buffer_var = op->buffer.var();
PrimExpr prim_index = op->indices[0];
@@ -733,7 +733,7 @@ void CodeGenSPIRV::VisitStmt_(const BufferStoreNode* op) {
}
}
-void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
+void CodeGenSPIRV::Dispatch_(const ForNode* op) {
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
spirv::Value init_value = MakeValue(op->min);
PrimExpr end = is_zero(op->min) ? op->extent : analyzer_->Simplify(op->min +
op->extent);
@@ -775,7 +775,7 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
// loop body
builder_->StartLabel(body_label);
var_map_[op->loop_var.get()] = spirv::Value(loop_var);
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
builder_->MakeInst(spv::OpBranch, continue_label);
// loop continue
@@ -788,7 +788,7 @@ void CodeGenSPIRV::VisitStmt_(const ForNode* op) {
builder_->StartLabel(merge_label);
}
-void CodeGenSPIRV::VisitStmt_(const WhileNode* op) {
+void CodeGenSPIRV::Dispatch_(const WhileNode* op) {
spirv::Label head_label = builder_->NewLabel();
spirv::Label condition_label = builder_->NewLabel();
spirv::Label body_label = builder_->NewLabel();
@@ -812,7 +812,7 @@ void CodeGenSPIRV::VisitStmt_(const WhileNode* op) {
// loop body
builder_->StartLabel(body_label);
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
builder_->MakeInst(spv::OpBranch, continue_label);
// loop continue
@@ -823,7 +823,7 @@ void CodeGenSPIRV::VisitStmt_(const WhileNode* op) {
builder_->StartLabel(merge_label);
}
-void CodeGenSPIRV::VisitStmt_(const IfThenElseNode* op) {
+void CodeGenSPIRV::Dispatch_(const IfThenElseNode* op) {
spirv::Value cond = MakeValue(op->condition);
spirv::Label then_label = builder_->NewLabel();
spirv::Label merge_label = builder_->NewLabel();
@@ -833,11 +833,11 @@ void CodeGenSPIRV::VisitStmt_(const IfThenElseNode* op) {
builder_->MakeInst(spv::OpBranchConditional, cond, then_label, else_label);
// then block
builder_->StartLabel(then_label);
- this->VisitStmt(op->then_case);
+ this->Dispatch(op->then_case);
builder_->MakeInst(spv::OpBranch, merge_label);
// else block
builder_->StartLabel(else_label);
- this->VisitStmt(op->else_case.value());
+ this->Dispatch(op->else_case.value());
builder_->MakeInst(spv::OpBranch, merge_label);
} else {
builder_->MakeInst(spv::OpSelectionMerge, merge_label,
spv::SelectionControlMaskNone);
@@ -845,14 +845,14 @@ void CodeGenSPIRV::VisitStmt_(const IfThenElseNode* op) {
weight_likely_branch_, 1);
// then block
builder_->StartLabel(then_label);
- this->VisitStmt(op->then_case);
+ this->Dispatch(op->then_case);
builder_->MakeInst(spv::OpBranch, merge_label);
}
// start merge label;
builder_->StartLabel(merge_label);
}
-void CodeGenSPIRV::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenSPIRV::Dispatch_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(!op->buffer->dtype.IsVoid());
const IntImmNode* dim_imm = op->buffer->shape[0].as<IntImmNode>();
TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack allocation
in GPU";
@@ -915,7 +915,7 @@ void CodeGenSPIRV::VisitStmt_(const AllocBufferNode* op) {
}
}
-void CodeGenSPIRV::VisitStmt_(const DeclBufferNode* op) {
+void CodeGenSPIRV::Dispatch_(const DeclBufferNode* op) {
const VarNode* buffer_var = op->buffer.get();
TVM_FFI_ICHECK(!var_map_.count(buffer_var))
<< "Buffer variable " << op->buffer.name() << " is already defined";
@@ -947,7 +947,7 @@ void CodeGenSPIRV::VisitStmt_(const DeclBufferNode* op) {
storage_info_[buffer_var] = std::move(info);
}
-void CodeGenSPIRV::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenSPIRV::Dispatch_(const AttrStmtNode* op) {
if (op->attr_key == tirx::attr::thread_extent) {
auto iv_opt = op->node.as<IterVar>();
TVM_FFI_ICHECK(iv_opt);
@@ -964,14 +964,14 @@ void CodeGenSPIRV::VisitStmt_(const AttrStmtNode* op) {
const prim::StringImmNode* shape_str = op->value.as<prim::StringImmNode>();
fragment_info_[buffer] = {shape_str->value};
}
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
}
-void CodeGenSPIRV::VisitStmt_(const AssertStmtNode* op) {
+void CodeGenSPIRV::Dispatch_(const AssertStmtNode* op) {
// AssertStmt is a leaf — no body to visit.
}
-void CodeGenSPIRV::VisitStmt_(const BindNode* op) {
+void CodeGenSPIRV::Dispatch_(const BindNode* op) {
TVM_FFI_ICHECK(!var_map_.count(op->var.get()));
if (auto prim_type = op->var->ty.as<PrimType>()) {
TVM_FFI_ICHECK(!prim_type.value().IsVoid());
@@ -984,13 +984,13 @@ void CodeGenSPIRV::VisitStmt_(const BindNode* op) {
}
}
-void CodeGenSPIRV::VisitStmt_(const SeqStmtNode* op) {
+void CodeGenSPIRV::Dispatch_(const SeqStmtNode* op) {
for (Stmt stmt : op->seq) {
- this->VisitStmt(stmt);
+ this->Dispatch(stmt);
}
}
-void CodeGenSPIRV::VisitStmt_(const EvaluateNode* op) { MakeValue(op->value); }
+void CodeGenSPIRV::Dispatch_(const EvaluateNode* op) { MakeValue(op->value); }
spirv::SType CodeGenSPIRV::GetFragmentSType(const VarNode* buffer, const
PrimType& dtype) {
TVM_FFI_ICHECK(fragment_info_.count(buffer));
diff --git a/src/backend/vulkan/codegen/codegen_spirv.h
b/src/backend/vulkan/codegen/codegen_spirv.h
index b88e5f828f..e5ee1321b2 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.h
+++ b/src/backend/vulkan/codegen/codegen_spirv.h
@@ -53,6 +53,9 @@ using namespace tirx;
class CodeGenSPIRV : public tirx::ExprFunctor<spirv::Value(const Expr&)>,
public StmtFunctor<void(const Stmt&)> {
public:
+ using tirx::ExprFunctor<spirv::Value(const Expr&)>::Dispatch;
+ using StmtFunctor::Dispatch;
+
/*!
* \brief Initialize the codegen based on a specific target.
*
@@ -105,17 +108,17 @@ class CodeGenSPIRV : public
tirx::ExprFunctor<spirv::Value(const Expr&)>,
spirv::Value Dispatch_(const TensorLoadNode* op) override;
spirv::Value Dispatch_(const prim::ShuffleNode* op) override;
// stmt
- void VisitStmt_(const BufferStoreNode* op) override;
- void VisitStmt_(const ForNode* op) override;
- void VisitStmt_(const WhileNode* op) override;
- void VisitStmt_(const IfThenElseNode* op) override;
- void VisitStmt_(const DeclBufferNode* op) override;
- void VisitStmt_(const AllocBufferNode* op) override;
- void VisitStmt_(const AttrStmtNode* op) override;
- void VisitStmt_(const AssertStmtNode* op) override;
- void VisitStmt_(const BindNode* op) override;
- void VisitStmt_(const SeqStmtNode* op) override;
- void VisitStmt_(const EvaluateNode* op) override;
+ void Dispatch_(const BufferStoreNode* op) override;
+ void Dispatch_(const ForNode* op) override;
+ void Dispatch_(const WhileNode* op) override;
+ void Dispatch_(const IfThenElseNode* op) override;
+ void Dispatch_(const DeclBufferNode* op) override;
+ void Dispatch_(const AllocBufferNode* op) override;
+ void Dispatch_(const AttrStmtNode* op) override;
+ void Dispatch_(const AssertStmtNode* op) override;
+ void Dispatch_(const BindNode* op) override;
+ void Dispatch_(const SeqStmtNode* op) override;
+ void Dispatch_(const EvaluateNode* op) override;
protected:
/*! \brief Storage information for a buffer */
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 9d345c27fd..1a4a12267e 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -645,7 +645,7 @@ void CodeGenWebGPU::Dispatch_(const TensorLoadNode* op,
std::ostream& os) { //
}
}
-void CodeGenWebGPU::VisitStmt_(const BindNode* op) {
+void CodeGenWebGPU::Dispatch_(const BindNode* op) {
// use ssa form.
if (print_ssa_form_) {
std::string value = PrintExpr(op->value);
@@ -660,7 +660,7 @@ void CodeGenWebGPU::VisitStmt_(const BindNode* op) {
}
}
-void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op) {
+void CodeGenWebGPU::Dispatch_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not
supported.";
PrimType value_ty = op->value.ty();
@@ -720,7 +720,7 @@ void CodeGenWebGPU::VisitStmt_(const BufferStoreNode* op) {
}
}
-void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenWebGPU::Dispatch_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get());
size_t constant_size = 1;
@@ -782,7 +782,7 @@ void CodeGenWebGPU::VisitStmt_(const AllocBufferNode* op) {
}
}
-void CodeGenWebGPU::VisitStmt_(const ForNode* op) {
+void CodeGenWebGPU::Dispatch_(const ForNode* op) {
std::string begin_str = PrintExpr(op->min);
PrimExpr end = is_zero(op->min) ? op->extent :
arith::Analyzer()->Simplify(op->min + op->extent);
std::string end_str = PrintExpr(end);
@@ -805,11 +805,11 @@ void CodeGenWebGPU::VisitStmt_(const ForNode* op) {
stream << "}\n";
}
-void CodeGenWebGPU::VisitStmt_(const AssertStmtNode* op) {
+void CodeGenWebGPU::Dispatch_(const AssertStmtNode* op) {
// skip assert — AssertStmt is a leaf, nothing to emit.
}
-void CodeGenWebGPU::VisitStmt_(const WhileNode* op) {
+void CodeGenWebGPU::Dispatch_(const WhileNode* op) {
PrintIndent();
stream << "while (true) {\n";
int while_scope = BeginScope();
@@ -822,12 +822,12 @@ void CodeGenWebGPU::VisitStmt_(const WhileNode* op) {
stream << "}\n";
}
-void CodeGenWebGPU::VisitStmt_(const BreakNode* op) {
+void CodeGenWebGPU::Dispatch_(const BreakNode* op) {
PrintIndent();
stream << "break;\n";
}
-void CodeGenWebGPU::VisitStmt_(const ContinueNode* op) {
+void CodeGenWebGPU::Dispatch_(const ContinueNode* op) {
PrintIndent();
stream << "continue;\n";
}
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.h
b/src/backend/webgpu/codegen/codegen_webgpu.h
index 9f906b98c4..e7e4bc1d06 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.h
+++ b/src/backend/webgpu/codegen/codegen_webgpu.h
@@ -76,14 +76,14 @@ class CodeGenWebGPU final : public CodeGenC {
void Dispatch_(const IntImmNode* op, std::ostream& os) final; //
NOLINT(*)
// stmt printing
- void VisitStmt_(const BindNode* op) final;
- void VisitStmt_(const BufferStoreNode* op) final;
- void VisitStmt_(const ForNode* op) final;
- void VisitStmt_(const AllocBufferNode* op) final;
- void VisitStmt_(const AssertStmtNode* op) final;
- void VisitStmt_(const WhileNode* op) final;
- void VisitStmt_(const BreakNode* op) final;
- void VisitStmt_(const ContinueNode* op) final;
+ void Dispatch_(const BindNode* op) final;
+ void Dispatch_(const BufferStoreNode* op) final;
+ void Dispatch_(const ForNode* op) final;
+ void Dispatch_(const AllocBufferNode* op) final;
+ void Dispatch_(const AssertStmtNode* op) final;
+ void Dispatch_(const WhileNode* op) final;
+ void Dispatch_(const BreakNode* op) final;
+ void Dispatch_(const ContinueNode* op) final;
private:
/*!
diff --git a/src/relax/transform/split_call_tir_by_pattern.cc
b/src/relax/transform/split_call_tir_by_pattern.cc
index 5f24da7cec..c2c621fb03 100644
--- a/src/relax/transform/split_call_tir_by_pattern.cc
+++ b/src/relax/transform/split_call_tir_by_pattern.cc
@@ -49,12 +49,12 @@ using namespace tvm::prim;
using relax::FCodegen;
using relax::MatchResult;
using relax::TIRPattern;
-using s_tir::ExprComparator;
using s_tir::TensorizeComparator;
/*! \brief helper to match a for stmt to a pattern*/
class ForMatcher : public TensorizeComparator {
public:
+ using TensorizeComparator::Dispatch;
using SymbolMap = std::unordered_map<Var, PrimExpr>;
explicit ForMatcher(const tirx::PrimFunc& pattern, const ffi::Array<Var>&
pattern_vars)
: TensorizeComparator(IRModule({{GlobalVar(""), pattern}}), false),
pattern_(pattern) {
@@ -67,7 +67,7 @@ class ForMatcher : public TensorizeComparator {
bool Match(const For& top) {
const ForNode* pattern_top =
pattern_->body.as<SBlockRealizeNode>()->block->body.as<ForNode>();
TVM_FFI_ICHECK(pattern_top) << "Invalid pattern function";
- if (!VisitStmt(top, ffi::GetRef<Stmt>(pattern_top))) {
+ if (!Dispatch(top, ffi::GetRef<Stmt>(pattern_top))) {
return false;
}
// Get evaluated symbols, buffers from the pattern.
@@ -85,7 +85,7 @@ class ForMatcher : public TensorizeComparator {
std::vector<BufferVar> evaluated_buffers;
private:
- using ExprComparator::Dispatch_;
+ using TensorizeComparator::Dispatch_;
ffi::Optional<PrimExpr> QueryEvaluatedSymbols(const Var& var) {
for (const SymbolMap& symbol_map : evaluated_symbols) {
@@ -249,7 +249,7 @@ class ForMatcher : public TensorizeComparator {
return true;
}
- bool VisitStmt_(const tirx::ForNode* op, const Stmt& other) final {
+ bool Dispatch_(const tirx::ForNode* op, const Stmt& other) final {
const auto* rhs = other.as<ForNode>();
loop_stack_lhs_.push_back(ffi::GetRef<For>(op));
loop_stack_rhs_.push_back(ffi::GetRef<For>(rhs));
@@ -269,10 +269,10 @@ class ForMatcher : public TensorizeComparator {
if (!op->annotations.empty() || !rhs->annotations.empty()) return false;
// Match the extents of loops
if (!Dispatch(op->extent, rhs->extent)) return false;
- return VisitStmt(op->body, rhs->body);
+ return Dispatch(op->body, rhs->body);
}
- bool VisitStmt_(const tirx::SBlockNode* op, const Stmt& other) final {
+ bool Dispatch_(const tirx::SBlockNode* op, const Stmt& other) final {
const auto* rhs = other.as<SBlockNode>();
// Check block equality.
// All iter vars and buffer regions including the order should match.
@@ -296,12 +296,12 @@ class ForMatcher : public TensorizeComparator {
if (op->init.has_value() && !rhs->init.has_value()) return false;
if (!op->init.has_value() && rhs->init.has_value()) return false;
if (op->init.has_value() && rhs->init.has_value()) {
- if (!VisitStmt(op->init.value(), rhs->init.value())) return false;
+ if (!Dispatch(op->init.value(), rhs->init.value())) return false;
}
- return VisitStmt(op->body, rhs->body);
+ return Dispatch(op->body, rhs->body);
}
- bool VisitStmt_(const SBlockRealizeNode* op, const Stmt& other) final {
+ bool Dispatch_(const SBlockRealizeNode* op, const Stmt& other) final {
const auto* rhs = other.as<SBlockRealizeNode>();
// Only allow trivial bindings
for (size_t i = 0; i < op->iter_values.size(); ++i) {
@@ -312,10 +312,10 @@ class ForMatcher : public TensorizeComparator {
}
// Disallow predicates now
if (!is_one(op->predicate) || !is_one(rhs->predicate)) return false;
- return VisitStmt(op->block, rhs->block);
+ return Dispatch(op->block, rhs->block);
}
- bool VisitStmt_(const BufferStoreNode* op, const Stmt& other) {
+ bool Dispatch_(const BufferStoreNode* op, const Stmt& other) {
const auto* rhs = other.as<BufferStoreNode>();
return CompareBufferAccess(op, rhs) && Dispatch(op->value, rhs->value);
}
@@ -357,7 +357,9 @@ class ForMatcher : public TensorizeComparator {
template <typename T>
bool CompareBufferAccess(const T* lhs, const T* rhs) {
if (!CompareBuffer(lhs->buffer, rhs->buffer)) return false;
- return CompareArray(lhs->indices, rhs->indices, &ForMatcher::Dispatch);
+ return CompareArray(
+ lhs->indices, rhs->indices,
+ static_cast<bool (ForMatcher::*)(const Expr&, const
PrimExpr&)>(&ForMatcher::Dispatch));
}
bool CompareBufferAccess(const TensorLoadNode* lhs, const TensorLoadNode*
rhs) {
@@ -366,7 +368,9 @@ class ForMatcher : public TensorizeComparator {
rhs->source.as_or_throw<BufferVar>())) {
return false;
}
- return CompareArray(lhs->indices, rhs->indices, &ForMatcher::Dispatch);
+ return CompareArray(
+ lhs->indices, rhs->indices,
+ static_cast<bool (ForMatcher::*)(const Expr&, const
PrimExpr&)>(&ForMatcher::Dispatch));
}
template <typename T, typename Self, typename F>
diff --git a/src/s_tir/analysis/domain_touched.cc
b/src/s_tir/analysis/domain_touched.cc
index 11c45d6b28..2e67be88ea 100644
--- a/src/s_tir/analysis/domain_touched.cc
+++ b/src/s_tir/analysis/domain_touched.cc
@@ -34,7 +34,7 @@
#include <unordered_map>
#include <unordered_set>
-#include "../../tirx/ir_visitor_with_analyzer.h"
+#include "../../tirx/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/analysis/estimate_flops.cc
b/src/s_tir/analysis/estimate_flops.cc
index 9e410064fa..0da2814819 100644
--- a/src/s_tir/analysis/estimate_flops.cc
+++ b/src/s_tir/analysis/estimate_flops.cc
@@ -90,7 +90,7 @@ class FlopEstimator : private tirx::ExprFunctor<TResult(const
Expr& n)>,
public:
using tirx::ExprFunctor<TResult(const Expr&)>::Dispatch;
- TResult VisitStmt(const Stmt& stmt) override { return
StmtFunctor::VisitStmt(stmt); }
+ TResult Dispatch(const Stmt& stmt) override { return
StmtFunctor::Dispatch(stmt); }
#define TVM_TIR_ESTIMATE_FLOP_VISIT_BINARY(Node) \
TResult Dispatch_(const Node* op) final { \
@@ -138,49 +138,49 @@ class FlopEstimator : private
tirx::ExprFunctor<TResult(const Expr& n)>,
}
TResult Dispatch_(const TensorLoadNode* op) override { return TResult(); }
- TResult VisitStmt_(const AttrStmtNode* op) override {
- TResult result = VisitStmt(op->body);
+ TResult Dispatch_(const AttrStmtNode* op) override {
+ TResult result = Dispatch(op->body);
result += Dispatch(op->value);
return result;
}
- TResult VisitStmt_(const BufferStoreNode* store) override { return
Dispatch(store->value); }
- TResult VisitStmt_(const SBlockRealizeNode* block) override {
- return VisitStmt(block->block->body);
+ TResult Dispatch_(const BufferStoreNode* store) override { return
Dispatch(store->value); }
+ TResult Dispatch_(const SBlockRealizeNode* block) override {
+ return Dispatch(block->block->body);
}
- TResult VisitStmt_(const SBlockNode* block) override {
+ TResult Dispatch_(const SBlockNode* block) override {
TResult result;
if (block->init.has_value()) {
- result += VisitStmt(block->init.value());
+ result += Dispatch(block->init.value());
}
- result += VisitStmt(block->body);
+ result += Dispatch(block->body);
return result;
}
- TResult VisitStmt_(const ForNode* loop) override {
+ TResult Dispatch_(const ForNode* loop) override {
ana->Bind(loop->loop_var, Range::FromMinExtent(loop->min, loop->extent));
const auto int_imm = GetLoopExtent(loop, ana);
- TResult result = VisitStmt(loop->body);
+ TResult result = Dispatch(loop->body);
result *= int_imm;
return result;
}
- TResult VisitStmt_(const IfThenElseNode* branch) override {
+ TResult Dispatch_(const IfThenElseNode* branch) override {
TResult cond = Dispatch(branch->condition);
if (branch->else_case) {
- cond +=
VisitStmt(branch->then_case).MaxWith(VisitStmt(branch->else_case.value()));
+ cond +=
Dispatch(branch->then_case).MaxWith(Dispatch(branch->else_case.value()));
} else {
- cond += VisitStmt(branch->then_case);
+ cond += Dispatch(branch->then_case);
}
return cond;
}
- TResult VisitStmt_(const WhileNode* op) override {
+ TResult Dispatch_(const WhileNode* op) override {
// TODO(jikechao): Improve while loop FLOP estimation with loop bound
analysis
TResult result = Dispatch(op->condition);
- result += VisitStmt(op->body);
+ result += Dispatch(op->body);
return result;
}
- TResult VisitStmt_(const BindNode* let) override {
+ TResult Dispatch_(const BindNode* let) override {
if (auto value = let->value.as<PrimExpr>()) return Dispatch(value.value());
return TResult();
}
@@ -191,7 +191,7 @@ class FlopEstimator : private
tirx::ExprFunctor<TResult(const Expr& n)>,
return cond;
}
- TResult VisitStmt_(const AssertStmtNode* op) override {
+ TResult Dispatch_(const AssertStmtNode* op) override {
TResult result = Dispatch(op->condition);
return result;
}
@@ -201,14 +201,14 @@ class FlopEstimator : private
tirx::ExprFunctor<TResult(const Expr& n)>,
TResult Dispatch_(const FloatImmNode* op) override { return TResult(); }
TResult Dispatch_(const prim::StringImmNode* op) override { return
TResult(); }
TResult Dispatch_(const prim::CastNode* op) override { return
Dispatch(op->value); }
- TResult VisitStmt_(const AllocBufferNode* op) override { return TResult(); }
- TResult VisitStmt_(const DeclBufferNode* op) override { return TResult(); }
- TResult VisitStmt_(const EvaluateNode* op) override { return TResult(); }
+ TResult Dispatch_(const AllocBufferNode* op) override { return TResult(); }
+ TResult Dispatch_(const DeclBufferNode* op) override { return TResult(); }
+ TResult Dispatch_(const EvaluateNode* op) override { return TResult(); }
- TResult VisitStmt_(const SeqStmtNode* seq) override {
+ TResult Dispatch_(const SeqStmtNode* seq) override {
TResult result;
for (const Stmt& stmt : seq->seq) {
- result += VisitStmt(stmt);
+ result += Dispatch(stmt);
}
return result;
}
@@ -232,7 +232,7 @@ double PostprocessResults(const TResult& result) {
double EstimateTIRFlops(const Stmt& stmt) {
FlopEstimator counter;
- return PostprocessResults(counter.VisitStmt(stmt));
+ return PostprocessResults(counter.Dispatch(stmt));
}
double EstimateTIRFlops(const IRModule& mod) {
@@ -243,7 +243,7 @@ double EstimateTIRFlops(const IRModule& mod) {
if (auto cached = f->attrs.GetAttr<int64_t>("estimated_flops")) {
cached_result += cached.value();
} else {
- result += counter.VisitStmt(f->body); //
+ result += counter.Dispatch(f->body); //
}
});
return PostprocessResults(result) + cached_result;
diff --git a/src/s_tir/analysis/identify_memcpy.cc
b/src/s_tir/analysis/identify_memcpy.cc
index f965745fa6..62be9ec6c4 100644
--- a/src/s_tir/analysis/identify_memcpy.cc
+++ b/src/s_tir/analysis/identify_memcpy.cc
@@ -38,7 +38,7 @@
#include <string>
#include <variant>
-#include "../../tirx/ir_visitor_with_analyzer.h"
+#include "../../tirx/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/analysis/is_pure_function.cc
b/src/s_tir/analysis/is_pure_function.cc
index d0c4785cc4..3c7acb10e6 100644
--- a/src/s_tir/analysis/is_pure_function.cc
+++ b/src/s_tir/analysis/is_pure_function.cc
@@ -45,13 +45,13 @@ class PurityChecker : TIRVisitorWithPath {
private:
explicit PurityChecker(bool assert_on_error) :
assert_on_error_(assert_on_error) {}
- void VisitStmt_(const AllocBufferNode* op, ffi::reflection::AccessPath path)
override {
+ void Dispatch_(const AllocBufferNode* op, ffi::reflection::AccessPath path)
override {
internal_allocations_.insert(op->buffer.var());
- TIRVisitorWithPath::VisitStmt_(op, path);
+ TIRVisitorWithPath::Dispatch_(op, path);
}
- void VisitStmt_(const BufferStoreNode* op, ffi::reflection::AccessPath path)
override {
- TIRVisitorWithPath::VisitStmt_(op, path);
+ void Dispatch_(const BufferStoreNode* op, ffi::reflection::AccessPath path)
override {
+ TIRVisitorWithPath::Dispatch_(op, path);
if (!internal_allocations_.count(op->buffer.var())) {
is_pure_ = false;
diff --git a/src/s_tir/analysis/oob_checker.cc
b/src/s_tir/analysis/oob_checker.cc
index bb40c9951e..55bf73e344 100644
--- a/src/s_tir/analysis/oob_checker.cc
+++ b/src/s_tir/analysis/oob_checker.cc
@@ -24,7 +24,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/transform.h>
-#include "../../tirx/ir_visitor_with_analyzer.h"
+#include "../../tirx/ir/ir_visitor_with_analyzer.h"
#include "../schedule/error.h"
namespace tvm {
diff --git a/src/s_tir/backend/adreno/inject_texture_alloc.cc
b/src/s_tir/backend/adreno/inject_texture_alloc.cc
index 13ee5520fc..0602e47a77 100644
--- a/src/s_tir/backend/adreno/inject_texture_alloc.cc
+++ b/src/s_tir/backend/adreno/inject_texture_alloc.cc
@@ -27,7 +27,7 @@
#include <tvm/tirx/stmt_functor.h>
#include "../../../backend/opencl/runtime/texture.h"
-#include "../../../tirx/ir_mutator_with_analyzer.h"
+#include "../../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../../tirx/transform/ir_utils.h"
namespace tvm {
diff --git a/src/s_tir/backend/adreno/texture_flatten.cc
b/src/s_tir/backend/adreno/texture_flatten.cc
index 3e76a21fe3..616c307a47 100644
--- a/src/s_tir/backend/adreno/texture_flatten.cc
+++ b/src/s_tir/backend/adreno/texture_flatten.cc
@@ -35,7 +35,7 @@
#include "../../../backend/opencl/runtime/texture.h"
#include "../../../runtime/thread_storage_scope.h"
-#include "../../../tirx/ir_visitor_with_analyzer.h"
+#include "../../../tirx/ir/ir_visitor_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/schedule/analysis/analysis.cc
b/src/s_tir/schedule/analysis/analysis.cc
index 6d695274b5..acc17917e0 100644
--- a/src/s_tir/schedule/analysis/analysis.cc
+++ b/src/s_tir/schedule/analysis/analysis.cc
@@ -2159,7 +2159,7 @@ bool CheckAutoTensorizeApplicable(const ScheduleState&
state, const tirx::StmtSR
arith::Analyzer analyzer;
auto desc_info = ExtractTensorIntrinDescInfo(analyzer.get(), desc_func);
- return extractor->VisitStmt(block->block, desc_info.desc_block->block);
+ return extractor->Dispatch(block->block, desc_info.desc_block->block);
}
bool CheckAutoTensorizeApplicable(const s_tir::Schedule& sch, const
s_tir::SBlockRV& block_rv,
diff --git a/src/s_tir/schedule/ir_comparator.cc
b/src/s_tir/schedule/ir_comparator.cc
index 81da621e37..1b8b1851cf 100644
--- a/src/s_tir/schedule/ir_comparator.cc
+++ b/src/s_tir/schedule/ir_comparator.cc
@@ -85,9 +85,9 @@ class TensorIntrinMismatchError : public
ScheduleErrorContextObj {
};
/* Override the dispatcher to make sure RHS is always valid */
-bool TensorizeComparator::VisitStmt(const Stmt& n, const Stmt& other) {
+bool TensorizeComparator::Dispatch(const Stmt& n, const Stmt& other) {
bool equal = n.same_as(other) ||
- ((n->type_index() == other->type_index()) &&
StmtComparator::VisitStmt(n, other));
+ ((n->type_index() == other->type_index()) &&
StmtComparator::Dispatch(n, other));
if (!equal && assert_mode_ && (n->IsInstance<ForNode>() ||
n->IsInstance<SBlockNode>())) {
throw MakeScheduleError<TensorIntrinMismatchError>(lhs_mod_, n, other,
std::move(error_messages_));
@@ -171,7 +171,7 @@ bool TensorizeComparator::Dispatch_(const CallNode* op,
const PrimExpr& other) {
return true;
}
-bool TensorizeComparator::VisitStmt_(const ForNode* op, const Stmt& other) {
+bool TensorizeComparator::Dispatch_(const ForNode* op, const Stmt& other) {
const auto* rhs = other.as<ForNode>();
if (!DefEqual(op->loop_var, rhs->loop_var)) {
if (assert_mode_) {
@@ -230,23 +230,27 @@ bool TensorizeComparator::VisitStmt_(const ForNode* op,
const Stmt& other) {
}
return false;
}
- return VisitStmt(op->body, rhs->body);
+ return Dispatch(op->body, rhs->body);
}
-bool TensorizeComparator::VisitStmt_(const SeqStmtNode* op, const Stmt& other)
{
+bool TensorizeComparator::Dispatch_(const SeqStmtNode* op, const Stmt& other) {
const auto* rhs = other.as<SeqStmtNode>();
- return CompareArray(op->seq, rhs->seq, &TensorizeComparator::VisitStmt);
+ return CompareArray(op->seq, rhs->seq,
+ static_cast<bool (TensorizeComparator::*)(const Stmt&,
const Stmt&)>(
+ &TensorizeComparator::Dispatch));
}
-bool TensorizeComparator::VisitStmt_(const BufferStoreNode* op, const Stmt&
other) {
+bool TensorizeComparator::Dispatch_(const BufferStoreNode* op, const Stmt&
other) {
const auto* rhs = other.as<BufferStoreNode>();
return CompareBufferAccess(op, rhs) && Dispatch(op->value, rhs->value);
}
-bool TensorizeComparator::VisitStmt_(const SBlockRealizeNode* op, const Stmt&
other) {
+bool TensorizeComparator::Dispatch_(const SBlockRealizeNode* op, const Stmt&
other) {
const auto* rhs = other.as<SBlockRealizeNode>();
if (!is_scope_block) {
- if (!CompareArray(op->iter_values, rhs->iter_values,
&TensorizeComparator::Dispatch)) {
+ if (!CompareArray(op->iter_values, rhs->iter_values,
+ static_cast<bool (TensorizeComparator::*)(const Expr&,
const PrimExpr&)>(
+ &TensorizeComparator::Dispatch))) {
if (assert_mode_) {
std::ostringstream os;
os << "BlockRealizeNode iter_values do not match: op->iter_values=" <<
op->iter_values
@@ -256,10 +260,10 @@ bool TensorizeComparator::VisitStmt_(const
SBlockRealizeNode* op, const Stmt& ot
return false;
}
}
- return Dispatch(op->predicate, rhs->predicate) && VisitStmt(op->block,
rhs->block);
+ return Dispatch(op->predicate, rhs->predicate) && Dispatch(op->block,
rhs->block);
}
-bool TensorizeComparator::VisitStmt_(const SBlockNode* op, const Stmt& other) {
+bool TensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt& other) {
const auto* rhs = other.as<SBlockNode>();
for (const IterVar& iter : op->iter_vars) {
lhs_analyzer_->Bind(iter->var, iter->dom);
@@ -306,7 +310,7 @@ bool TensorizeComparator::VisitStmt_(const SBlockNode* op,
const Stmt& other) {
return false;
}
is_scope_block = false;
- return VisitStmt(op->body, rhs->body);
+ return Dispatch(op->body, rhs->body);
}
// Exprs
@@ -709,11 +713,11 @@ bool AutoTensorizeComparator::DispatchDefault_(const
ffi::Object* op, const Prim
return false;
}
-bool AutoTensorizeComparator::VisitStmtDefault_(const ffi::Object* op, const
Stmt& other) {
+bool AutoTensorizeComparator::DispatchDefault_(const ffi::Object* op, const
Stmt& other) {
return false;
}
-bool AutoTensorizeComparator::VisitStmt_(const SBlockNode* op, const Stmt&
other) {
+bool AutoTensorizeComparator::Dispatch_(const SBlockNode* op, const Stmt&
other) {
const auto* rhs = other.as<SBlockNode>();
// Check block equality.
// All iter vars and buffer regions including the order should match.
@@ -753,7 +757,7 @@ bool AutoTensorizeComparator::VisitStmt_(const SBlockNode*
op, const Stmt& other
}
}
is_scope_block = false;
- return VisitStmt(op->body, rhs->body);
+ return Dispatch(op->body, rhs->body);
}
bool AutoTensorizeComparator::CompareBuffer(const BufferVar& lhs, const
BufferVar& rhs) {
@@ -783,7 +787,7 @@ bool AutoTensorizeComparator::CompareBuffer(const
BufferVar& lhs, const BufferVa
return equal;
}
-bool AutoTensorizeComparator::VisitStmt_(const BufferStoreNode* op, const
Stmt& other) {
+bool AutoTensorizeComparator::Dispatch_(const BufferStoreNode* op, const Stmt&
other) {
const auto* rhs = other.as<BufferStoreNode>();
return CompareBufferAccess(op, rhs) && Dispatch(op->value, rhs->value);
}
diff --git a/src/s_tir/schedule/ir_comparator.h
b/src/s_tir/schedule/ir_comparator.h
index 28c8332341..75303df5b8 100644
--- a/src/s_tir/schedule/ir_comparator.h
+++ b/src/s_tir/schedule/ir_comparator.h
@@ -47,14 +47,14 @@ class TensorizeComparator : public ExprComparator, public
StmtComparator {
: lhs_mod_(std::move(lhs_mod)), assert_mode_(assert_mode) {}
bool Dispatch(const Expr& n, const PrimExpr& other) override;
- bool VisitStmt(const Stmt& n, const Stmt& other) override;
+ bool Dispatch(const Stmt& n, const Stmt& other) override;
bool Dispatch_(const CallNode* op, const PrimExpr& other) override;
- bool VisitStmt_(const ForNode* op, const Stmt& other) override;
- bool VisitStmt_(const SeqStmtNode* op, const Stmt& other) override;
- bool VisitStmt_(const BufferStoreNode* op, const Stmt& other) override;
- bool VisitStmt_(const SBlockRealizeNode* op, const Stmt& other) override;
- bool VisitStmt_(const SBlockNode* op, const Stmt& other) override;
+ bool Dispatch_(const ForNode* op, const Stmt& other) override;
+ bool Dispatch_(const SeqStmtNode* op, const Stmt& other) override;
+ bool Dispatch_(const BufferStoreNode* op, const Stmt& other) override;
+ bool Dispatch_(const SBlockRealizeNode* op, const Stmt& other) override;
+ bool Dispatch_(const SBlockNode* op, const Stmt& other) override;
bool Dispatch_(const AddNode* op, const PrimExpr& other) override;
bool Dispatch_(const SubNode* op, const PrimExpr& other) override;
@@ -139,10 +139,10 @@ class AutoTensorizeComparator : public
TensorizeComparator {
private:
bool DispatchDefault_(const ffi::Object* op, const PrimExpr& other) override;
- bool VisitStmtDefault_(const ffi::Object* op, const Stmt& other) override;
+ bool DispatchDefault_(const ffi::Object* op, const Stmt& other) override;
- bool VisitStmt_(const SBlockNode* op, const Stmt& other) override;
- bool VisitStmt_(const BufferStoreNode* op, const Stmt& other) override;
+ bool Dispatch_(const SBlockNode* op, const Stmt& other) override;
+ bool Dispatch_(const BufferStoreNode* op, const Stmt& other) override;
bool Dispatch_(const TensorLoadNode* op, const PrimExpr& other) override;
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc
b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index 830909374d..572f8c4804 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -835,7 +835,7 @@ void Tensorize(ScheduleState self, const StmtSRef& sref,
const TensorIntrin& int
->Rewrite(intrin_impl);
// Step 2: Structural pattern matching
TensorizeComparator comparator(self->mod, /*assert_mode=*/true);
- comparator.VisitStmt(block_realize, intrin_desc->body);
+ comparator.Dispatch(block_realize, intrin_desc->body);
// Step 3: Prepare necessary mapping
// 1) BufferVar mapping from intrin impl buffers to intrin desc buffers.
// 2) BufferVar mapping from intrin impl buffers to buffers in the current
AST.
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index 72475f2465..207b226ef8 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -27,7 +27,7 @@
#include <optional>
#include <variant>
-#include "../../../tirx/ir_mutator_with_analyzer.h"
+#include "../../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../utils.h"
namespace tvm {
diff --git a/src/s_tir/schedule/transform.h b/src/s_tir/schedule/transform.h
index 72c76daaa9..8e0cd2098a 100644
--- a/src/s_tir/schedule/transform.h
+++ b/src/s_tir/schedule/transform.h
@@ -28,7 +28,7 @@
#include <utility>
#include "../../tirx/ir/functor_common.h"
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/hoist_expression.cc
b/src/s_tir/transform/hoist_expression.cc
index 305bda9be8..9af41d25c5 100644
--- a/src/s_tir/transform/hoist_expression.cc
+++ b/src/s_tir/transform/hoist_expression.cc
@@ -37,7 +37,7 @@
#include "../../arith/interval_set.h"
#include "../../runtime/thread_storage_scope.h"
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
diff --git a/src/s_tir/transform/inject_permuted_layout.cc
b/src/s_tir/transform/inject_permuted_layout.cc
index 55e71985f8..11bcdbcebe 100644
--- a/src/s_tir/transform/inject_permuted_layout.cc
+++ b/src/s_tir/transform/inject_permuted_layout.cc
@@ -31,7 +31,7 @@
#include "../../runtime/thread_storage_scope.h"
#include "../../support/utils.h"
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
diff --git a/src/s_tir/transform/inject_virtual_thread.cc
b/src/s_tir/transform/inject_virtual_thread.cc
index 6397723504..5c4dd96d02 100644
--- a/src/s_tir/transform/inject_virtual_thread.cc
+++ b/src/s_tir/transform/inject_virtual_thread.cc
@@ -33,7 +33,7 @@
#include <unordered_set>
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
diff --git a/src/s_tir/transform/lower_async_dma.cc
b/src/s_tir/transform/lower_async_dma.cc
index 12fdd9e13a..2e176067bc 100644
--- a/src/s_tir/transform/lower_async_dma.cc
+++ b/src/s_tir/transform/lower_async_dma.cc
@@ -36,7 +36,7 @@
#include <optional>
#include <set>
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"
namespace tvm {
diff --git a/src/s_tir/transform/renormalize_split_pattern.cc
b/src/s_tir/transform/renormalize_split_pattern.cc
index a81cc0c2d8..eadb6e4540 100644
--- a/src/s_tir/transform/renormalize_split_pattern.cc
+++ b/src/s_tir/transform/renormalize_split_pattern.cc
@@ -30,7 +30,7 @@
#include <tvm/tirx/stmt_functor.h>
#include "../../arith/pattern_match.h"
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace s_tir {
diff --git a/src/s_tir/transform/using_assume_to_reduce_branches.cc
b/src/s_tir/transform/using_assume_to_reduce_branches.cc
index 4592203dd1..8d8e9e5d54 100644
--- a/src/s_tir/transform/using_assume_to_reduce_branches.cc
+++ b/src/s_tir/transform/using_assume_to_reduce_branches.cc
@@ -49,7 +49,7 @@
#include <unordered_map>
#include "../../arith/constraint_extract.h"
-#include "../../tirx/ir_mutator_with_analyzer.h"
+#include "../../tirx/ir/ir_mutator_with_analyzer.h"
#include "tvm/ir/expr.h"
namespace tvm {
namespace s_tir {
diff --git a/src/target/llvm/codegen_aarch64.cc
b/src/target/llvm/codegen_aarch64.cc
index 1ac7450ea2..6523b6be91 100644
--- a/src/target/llvm/codegen_aarch64.cc
+++ b/src/target/llvm/codegen_aarch64.cc
@@ -40,7 +40,7 @@ class CodeGenAArch64 final : public CodeGenCPU {
CodeGenAArch64() = default;
virtual ~CodeGenAArch64() = default;
- void VisitStmt_(const AttrStmtNode* op);
+ void Dispatch_(const AttrStmtNode* op);
void AddFunction(const GlobalVar& gvar, const PrimFunc& f);
void SetTargetAttributes(llvm::Function* func);
@@ -84,16 +84,16 @@ void CodeGenAArch64::SetTargetAttributes(llvm::Function*
func) {
* \brief Visit and handle AArch64 specific pragmas. To be AArch64 specific,
* the expectation is that they are prepended with "pragma_aarch64".
*/
-void CodeGenAArch64::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenAArch64::Dispatch_(const AttrStmtNode* op) {
std::string attr_key = op->attr_key;
if (!tirx::attr::IsPragmaKey(attr_key)) {
- CodeGenCPU::VisitStmt_(op);
+ CodeGenCPU::Dispatch_(op);
return;
}
bool is_aarch64_specific_pragma = attr_key.substr(7, 7) == "aarch64";
if (!is_aarch64_specific_pragma) {
- CodeGenCPU::VisitStmt_(op);
+ CodeGenCPU::Dispatch_(op);
return;
}
@@ -117,7 +117,7 @@ void CodeGenAArch64::VisitStmt_(const AttrStmtNode* op) {
} else {
LOG(WARNING) << "Unknown pragma " << op->attr_key;
}
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
}
TVM_FFI_STATIC_INIT_BLOCK() {
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 840fdf12e3..ed3a4c9c29 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -604,7 +604,7 @@ void CodeGenCPU::CreateComputeScope(const AttrStmtNode* op)
{
CreateDebugFunction(MakeStringRef(value->value), debug_param_types,
PrimType::Int(32));
auto* compute_entry = llvm::BasicBlock::Create(*ctx, "entry", function_);
builder_->SetInsertPoint(compute_entry);
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
builder_->CreateRet(ConstInt32(0));
builder_->SetInsertPoint(compute_call_end);
@@ -690,7 +690,7 @@ void CodeGenCPU::CreateParallelLaunch(const Stmt& body, int
num_task, std::strin
std::swap(parallel_env_, par_env);
std::swap(analyzer_, new_analyzer);
std::swap(var_map_, new_vmap);
- this->VisitStmt(body);
+ this->Dispatch(body);
builder_->CreateRet(ConstInt32(0));
// swap the var map back, now we are back on track.
std::swap(var_map_, new_vmap);
@@ -1113,7 +1113,7 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode*
op) {
}
}
-void CodeGenCPU::VisitStmt_(const AssertStmtNode* op) {
+void CodeGenCPU::Dispatch_(const AssertStmtNode* op) {
EmitDebugLocation(op);
llvm::Value* cond = MakeValue(op->condition);
llvm::LLVMContext* ctx = llvm_target_->GetContext();
@@ -1146,10 +1146,10 @@ void CodeGenCPU::VisitStmt_(const AssertStmtNode* op) {
builder_->CreateRet(ConstInt32(-1));
// otherwise set it to be new end.
builder_->SetInsertPoint(end_block);
- CodeGenLLVM::VisitStmt_(op);
+ CodeGenLLVM::Dispatch_(op);
}
-void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenCPU::Dispatch_(const AttrStmtNode* op) {
EmitDebugLocation(op);
if (op->attr_key == tirx::attr::compute_scope) {
this->CreateComputeScope(op);
@@ -1158,7 +1158,7 @@ void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
TVM_FFI_ICHECK(parallel_env_.penv != nullptr)
<< "Pragma parallel_stride_pattern only valid in parallel launch";
parallel_env_.stride_pattern = true;
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
} else if (op->attr_key == "pragma_parallel_launch_point") {
CreateParallelLaunch(op->body, 0, "pragma_parallel");
} else if (op->attr_key == "pragma_parallel_barrier_when_finish") {
@@ -1167,7 +1167,7 @@ void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop)
<< "Cannot not place within parallel loop as the workload may
differ, "
<< " place it between parallel and parallel_launch_point";
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
auto bar_callee =
llvm::FunctionCallee(ftype_tvm_parallel_barrier_,
RuntimeTVMParallelBarrier());
builder_->CreateCall(bar_callee, {MakeValue(parallel_env_.task_id),
parallel_env_.penv});
@@ -1175,20 +1175,20 @@ void CodeGenCPU::VisitStmt_(const AttrStmtNode* op) {
const prim::StringImmNode* value = op->value.as<prim::StringImmNode>();
TVM_FFI_ICHECK(value != nullptr);
this->HandleImport(value->value);
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
} else {
LOG(WARNING) << "Unknown pragma " << op->attr_key;
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
}
} else {
- CodeGenLLVM::VisitStmt_(op);
+ CodeGenLLVM::Dispatch_(op);
}
}
-void CodeGenCPU::VisitStmt_(const ForNode* op) {
+void CodeGenCPU::Dispatch_(const ForNode* op) {
EmitDebugLocation(op);
if (op->kind == ForKind::kSerial || op->kind == ForKind::kUnrolled) {
- CodeGenLLVM::VisitStmt_(op);
+ CodeGenLLVM::Dispatch_(op);
} else if (op->kind == ForKind::kParallel) {
TVM_FFI_ICHECK(is_zero(op->min))
<< "Parallel launch require canonical loop with zero start index";
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index 88799b7779..4d253ba38b 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -61,6 +61,7 @@ class LLVMTarget;
// CPU host code generation
class CodeGenCPU : public CodeGenLLVM {
public:
+ using CodeGenLLVM::Dispatch_;
CodeGenCPU();
virtual ~CodeGenCPU();
@@ -70,9 +71,9 @@ class CodeGenCPU : public CodeGenLLVM {
void AddFunction(const GlobalVar& gvar, const PrimFunc& f) override;
void AddMainFunction(const std::string& entry_func_name) override;
std::unique_ptr<llvm::Module> Finish() override;
- void VisitStmt_(const AssertStmtNode* op) override;
- void VisitStmt_(const AttrStmtNode* op) override;
- void VisitStmt_(const ForNode* op) override;
+ void Dispatch_(const AssertStmtNode* op) override;
+ void Dispatch_(const AttrStmtNode* op) override;
+ void Dispatch_(const ForNode* op) override;
llvm::Value* CreateIntrinsic(const CallNode* op) override;
llvm::Value* CreateCallExtern(Type ret_type, ffi::String global_symbol,
const ffi::Array<Expr>& args, bool
skip_first_arg) override;
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 92892ec7bf..2301ff43e9 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -340,7 +340,7 @@ void CodeGenLLVM::AddFunctionInternal(const GlobalVar&
gvar, const PrimFunc& f)
llvm::LLVMContext* ctx = llvm_target_->GetContext();
llvm::BasicBlock* entry = llvm::BasicBlock::Create(*ctx, "entry", function_);
builder_->SetInsertPoint(entry);
- this->VisitStmt(f->body);
+ this->Dispatch(f->body);
// Add alignment attribute if needed.
for (size_t i = 0; i < f->params.size(); ++i) {
@@ -901,7 +901,7 @@ void CodeGenLLVM::CreateSerialFor(llvm::Value* begin,
llvm::Value* end, llvm::Va
EmitDebugLocation(body->span);
PushLoopFrame(for_next, for_end);
- this->VisitStmt(body);
+ this->Dispatch(body);
PopLoopFrame();
var_map_.erase(loop_var.get());
@@ -2071,7 +2071,7 @@ llvm::Value* CodeGenLLVM::Dispatch_(const
prim::BroadcastNode* op) {
return builder_->CreateShuffleVector(value, undef, mask);
}
-void CodeGenLLVM::VisitStmt_(const BufferStoreNode* op) {
+void CodeGenLLVM::Dispatch_(const BufferStoreNode* op) {
EmitDebugLocation(op);
PrimType value_dtype = PrimType(op->value.ty()->dtype);
Var buffer_var = op->buffer.var();
@@ -2111,7 +2111,7 @@ void CodeGenLLVM::VisitStmt_(const BufferStoreNode* op) {
BufferAccessHelper(op->buffer, op->indices, std::nullopt, value_dtype,
make_store);
}
-void CodeGenLLVM::VisitStmt_(const ForNode* op) {
+void CodeGenLLVM::Dispatch_(const ForNode* op) {
EmitDebugLocation(op);
analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
if (op->kind == ForKind::kUnrolled) {
@@ -2127,7 +2127,7 @@ void CodeGenLLVM::VisitStmt_(const ForNode* op) {
CreateSerialFor(begin_value, end_value, MakeValue(step), op->loop_var,
op->body);
}
-void CodeGenLLVM::VisitStmt_(const WhileNode* op) {
+void CodeGenLLVM::Dispatch_(const WhileNode* op) {
EmitDebugLocation(op);
llvm::LLVMContext* ctx = llvm_target_->GetContext();
auto* while_cond = llvm::BasicBlock::Create(*ctx, "while_cond", function_);
@@ -2138,13 +2138,13 @@ void CodeGenLLVM::VisitStmt_(const WhileNode* op) {
builder_->CreateCondBr(MakeValue(op->condition), while_body, while_merge);
builder_->SetInsertPoint(while_body);
PushLoopFrame(while_cond, while_merge);
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
PopLoopFrame();
builder_->CreateBr(while_cond);
builder_->SetInsertPoint(while_merge);
}
-void CodeGenLLVM::VisitStmt_(const ReturnNode* op) {
+void CodeGenLLVM::Dispatch_(const ReturnNode* op) {
EmitDebugLocation(op);
auto const* val = op->value.as<IntImmNode>();
TVM_FFI_ICHECK(val) << "Return should be transformed to return zero "
@@ -2159,7 +2159,7 @@ void CodeGenLLVM::VisitStmt_(const ReturnNode* op) {
builder_->SetInsertPoint(ret_dummy);
}
-void CodeGenLLVM::VisitStmt_(const IfThenElseNode* op) {
+void CodeGenLLVM::Dispatch_(const IfThenElseNode* op) {
EmitDebugLocation(op);
llvm::Value* cond = MakeValue(op->condition);
llvm::LLVMContext* ctx = llvm_target_->GetContext();
@@ -2169,21 +2169,21 @@ void CodeGenLLVM::VisitStmt_(const IfThenElseNode* op) {
auto* else_block = llvm::BasicBlock::Create(*ctx, "if_else", function_);
builder_->CreateCondBr(cond, then_block, else_block);
builder_->SetInsertPoint(then_block);
- this->VisitStmt(op->then_case);
+ this->Dispatch(op->then_case);
builder_->CreateBr(end_block);
builder_->SetInsertPoint(else_block);
- this->VisitStmt(op->else_case.value());
+ this->Dispatch(op->else_case.value());
builder_->CreateBr(end_block);
} else {
builder_->CreateCondBr(cond, then_block, end_block,
md_very_likely_branch_);
builder_->SetInsertPoint(then_block);
- this->VisitStmt(op->then_case);
+ this->Dispatch(op->then_case);
builder_->CreateBr(end_block);
}
builder_->SetInsertPoint(end_block);
}
-void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenLLVM::Dispatch_(const AllocBufferNode* op) {
EmitDebugLocation(op);
TVM_FFI_ICHECK_EQ(op->buffer->shape.size(), 1)
<< "LLVM codegen only supports flat 1-d buffer allocation, but
allocation of "
@@ -2230,7 +2230,7 @@ void CodeGenLLVM::VisitStmt_(const AllocBufferNode* op) {
}
}
-void CodeGenLLVM::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenLLVM::Dispatch_(const AttrStmtNode* op) {
EmitDebugLocation(op);
if (op->attr_key == tirx::attr::thread_extent) {
IterVar iv = op->node.as_or_throw<IterVar>();
@@ -2249,16 +2249,16 @@ void CodeGenLLVM::VisitStmt_(const AttrStmtNode* op) {
alloc_storage_info_[v].alignment);
}
}
- this->VisitStmt(op->body);
+ this->Dispatch(op->body);
}
-void CodeGenLLVM::VisitStmt_(const AssertStmtNode* op) {
+void CodeGenLLVM::Dispatch_(const AssertStmtNode* op) {
EmitDebugLocation(op);
// AssertStmt is a leaf — no body to visit.
// Constraint scoping is handled by ScopeStack in analysis passes.
}
-void CodeGenLLVM::VisitStmt_(const BindNode* op) {
+void CodeGenLLVM::Dispatch_(const BindNode* op) {
EmitDebugLocation(op);
const VarNode* v = op->var.get();
TVM_FFI_ICHECK(!var_map_.count(v));
@@ -2297,14 +2297,14 @@ void CodeGenLLVM::VisitStmt_(const BindNode* op) {
AddDebugInformation(value, op->var);
}
-void CodeGenLLVM::VisitStmt_(const SeqStmtNode* op) {
+void CodeGenLLVM::Dispatch_(const SeqStmtNode* op) {
EmitDebugLocation(op);
for (Stmt stmt : op->seq) {
- this->VisitStmt(stmt);
+ this->Dispatch(stmt);
}
}
-void CodeGenLLVM::VisitStmt_(const DeclBufferNode* op) {
+void CodeGenLLVM::Dispatch_(const DeclBufferNode* op) {
EmitDebugLocation(op);
const VarNode* buffer = op->buffer.get();
TVM_FFI_ICHECK(!var_map_.count(buffer));
@@ -2338,7 +2338,7 @@ void CodeGenLLVM::VisitStmt_(const DeclBufferNode* op) {
}
}
-void CodeGenLLVM::VisitStmt_(const EvaluateNode* op) {
+void CodeGenLLVM::Dispatch_(const EvaluateNode* op) {
EmitDebugLocation(op);
MakeValue(op->value);
}
diff --git a/src/target/llvm/codegen_llvm.h b/src/target/llvm/codegen_llvm.h
index 6839811799..c3e4989a5a 100644
--- a/src/target/llvm/codegen_llvm.h
+++ b/src/target/llvm/codegen_llvm.h
@@ -94,6 +94,9 @@ using namespace tirx;
class CodeGenLLVM : public tirx::ExprFunctor<llvm::Value*(const Expr&)>,
public StmtFunctor<void(const Stmt&)> {
public:
+ using tirx::ExprFunctor<llvm::Value*(const Expr&)>::Dispatch;
+ using StmtFunctor::Dispatch;
+
CodeGenLLVM(); // Do not make it default here.
virtual ~CodeGenLLVM(); // Do not make it default here.
@@ -227,18 +230,18 @@ class CodeGenLLVM : public
tirx::ExprFunctor<llvm::Value*(const Expr&)>,
llvm::Value* Dispatch_(const prim::ShuffleNode* op) override;
llvm::Value* Dispatch_(const prim::BroadcastNode* op) override;
// stmt
- void VisitStmt_(const BufferStoreNode* op) override;
- void VisitStmt_(const ForNode* op) override;
- void VisitStmt_(const WhileNode* op) override;
- void VisitStmt_(const ReturnNode* op) override;
- void VisitStmt_(const IfThenElseNode* op) override;
- void VisitStmt_(const AllocBufferNode* op) override;
- void VisitStmt_(const AttrStmtNode* op) override;
- void VisitStmt_(const AssertStmtNode* op) override;
- void VisitStmt_(const BindNode* op) override;
- void VisitStmt_(const SeqStmtNode* op) override;
- void VisitStmt_(const EvaluateNode* op) override;
- void VisitStmt_(const DeclBufferNode* op) override;
+ void Dispatch_(const BufferStoreNode* op) override;
+ void Dispatch_(const ForNode* op) override;
+ void Dispatch_(const WhileNode* op) override;
+ void Dispatch_(const ReturnNode* op) override;
+ void Dispatch_(const IfThenElseNode* op) override;
+ void Dispatch_(const AllocBufferNode* op) override;
+ void Dispatch_(const AttrStmtNode* op) override;
+ void Dispatch_(const AssertStmtNode* op) override;
+ void Dispatch_(const BindNode* op) override;
+ void Dispatch_(const SeqStmtNode* op) override;
+ void Dispatch_(const EvaluateNode* op) override;
+ void Dispatch_(const DeclBufferNode* op) override;
// Get constant string
llvm::Constant* GetConstString(const std::string& str);
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 4cca7b6255..8ac80b4f74 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -926,7 +926,7 @@ void CodeGenC::PrintVecBinaryOp(const std::string& op,
const PrimType& t, PrimEx
}
}
-void CodeGenC::VisitStmt_(const DeclBufferNode* op) {
+void CodeGenC::Dispatch_(const DeclBufferNode* op) {
const VarNode* source = op->data.as<VarNode>();
if (const auto* call = op->data.as<CallNode>();
call && call->op.same_as(tirx::builtin::buffer_data()) &&
call->args.size() == 1) {
@@ -1039,7 +1039,7 @@ void CodeGenC::Dispatch_(const TensorLoadNode* op,
std::ostream& os) { // NOLIN
}
}
-void CodeGenC::VisitStmt_(const BufferStoreNode* op) {
+void CodeGenC::Dispatch_(const BufferStoreNode* op) {
TVM_FFI_ICHECK_EQ(op->indices.size(), 1) << "Store to non-flat memory not
supported.";
PrimType value_ty = op->value.ty();
@@ -1228,7 +1228,7 @@ void CodeGenC::Dispatch_(const prim::SelectNode* op,
std::ostream& os) { // NOL
os << ")";
}
-void CodeGenC::VisitStmt_(const BindNode* op) {
+void CodeGenC::Dispatch_(const BindNode* op) {
RegisterHandleTypeFromPointer(op->var, &op->value);
std::string value = PrintExpr(op->value);
if (print_ssa_form_) {
@@ -1249,7 +1249,7 @@ void CodeGenC::VisitStmt_(const BindNode* op) {
}
}
-void CodeGenC::VisitStmt_(const AllocBufferNode* op) {
+void CodeGenC::Dispatch_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get(), op->buffer.name() + "_ptr");
@@ -1276,7 +1276,7 @@ void CodeGenC::VisitStmt_(const AllocBufferNode* op) {
}
}
-void CodeGenC::VisitStmt_(const AttrStmtNode* op) {
+void CodeGenC::Dispatch_(const AttrStmtNode* op) {
if (op->attr_key == tirx::attr::thread_extent) {
IterVar iv = op->node.as_or_throw<IterVar>();
if (iv->thread_tag.length() != 0) {
@@ -1340,7 +1340,7 @@ void CodeGenC::PrintEscapedCString(const std::string&
str, std::ostream& os) {
os << "\"";
}
-void CodeGenC::VisitStmt_(const AssertStmtNode* op) {
+void CodeGenC::Dispatch_(const AssertStmtNode* op) {
std::string cond = PrintExpr(op->condition);
PrintIndent();
int num_parts = static_cast<int>(op->message_parts.size());
@@ -1368,7 +1368,7 @@ void CodeGenC::VisitStmt_(const AssertStmtNode* op) {
}
}
-void CodeGenC::VisitStmt_(const ForNode* op) {
+void CodeGenC::Dispatch_(const ForNode* op) {
std::string begin_str = PrintExpr(op->min);
PrimExpr end = is_zero(op->min) ? op->extent :
arith::Analyzer()->Simplify(op->min + op->extent);
std::string end_str = PrintExpr(end);
@@ -1391,7 +1391,7 @@ void CodeGenC::VisitStmt_(const ForNode* op) {
stream << "}\n";
}
-void CodeGenC::VisitStmt_(const WhileNode* op) {
+void CodeGenC::Dispatch_(const WhileNode* op) {
PrintIndent();
stream << "#pragma unroll 1\n";
PrintIndent();
@@ -1406,24 +1406,24 @@ void CodeGenC::VisitStmt_(const WhileNode* op) {
stream << "}\n";
}
-void CodeGenC::VisitStmt_(const ReturnNode* op) {
+void CodeGenC::Dispatch_(const ReturnNode* op) {
PrintIndent();
stream << "return ";
PrintExpr(op->value, stream);
stream << ";\n";
}
-void CodeGenC::VisitStmt_(const BreakNode* op) {
+void CodeGenC::Dispatch_(const BreakNode* op) {
PrintIndent();
stream << "break;\n";
}
-void CodeGenC::VisitStmt_(const ContinueNode* op) {
+void CodeGenC::Dispatch_(const ContinueNode* op) {
PrintIndent();
stream << "continue;\n";
}
-void CodeGenC::VisitStmt_(const IfThenElseNode* op) {
+void CodeGenC::Dispatch_(const IfThenElseNode* op) {
std::string cond = PrintExpr(op->condition);
PrintIndent();
if (cond[0] == '(' && cond[cond.length() - 1] == ')') {
@@ -1446,13 +1446,13 @@ void CodeGenC::VisitStmt_(const IfThenElseNode* op) {
stream << "}\n";
}
-void CodeGenC::VisitStmt_(const SeqStmtNode* op) {
+void CodeGenC::Dispatch_(const SeqStmtNode* op) {
for (Stmt stmt : op->seq) {
PrintStmt(stmt);
}
}
-void CodeGenC::VisitStmt_(const EvaluateNode* op) {
+void CodeGenC::Dispatch_(const EvaluateNode* op) {
if (auto value = op->value.as<PrimExpr>(); value &&
is_const_int(value.value())) return;
const CallNode* call = op->value.as<CallNode>();
if (call) {
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index b5ac52aa67..b2e0a99b56 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -61,6 +61,9 @@ class CodeGenC : public tirx::ExprFunctor<void(const Expr&,
std::ostream&)>,
public StmtFunctor<void(const Stmt&)>,
public CodeGenSourceBase {
public:
+ using tirx::ExprFunctor<void(const Expr&, std::ostream&)>::Dispatch;
+ using StmtFunctor::Dispatch;
+
/*!
* \brief Initialize the code generator.
* \param output_ssa Whether output SSA.
@@ -102,7 +105,7 @@ class CodeGenC : public tirx::ExprFunctor<void(const Expr&,
std::ostream&)>,
* \brief Print the Stmt n to CodeGenC->stream
* \param n The statement to be printed.
*/
- void PrintStmt(const Stmt& n) { VisitStmt(n); }
+ void PrintStmt(const Stmt& n) { Dispatch(n); }
/*!
* \brief Print the expression n(or its ssa id if in ssa mode) into os
* \param n The expression to be printed.
@@ -194,20 +197,20 @@ class CodeGenC : public tirx::ExprFunctor<void(const
Expr&, std::ostream&)>,
void Dispatch_(const FloatImmNode* op, std::ostream& os) override;
// NOLINT(*)
void Dispatch_(const prim::StringImmNode* op, std::ostream& os) override;
// NOLINT(*)
// statment
- void VisitStmt_(const BindNode* op) override;
- void VisitStmt_(const BufferStoreNode* op) override;
- void VisitStmt_(const ForNode* op) override;
- void VisitStmt_(const WhileNode* op) override;
- void VisitStmt_(const ReturnNode* op) override;
- void VisitStmt_(const BreakNode* op) override;
- void VisitStmt_(const ContinueNode* op) override;
- void VisitStmt_(const IfThenElseNode* op) override;
- void VisitStmt_(const AllocBufferNode* op) override;
- void VisitStmt_(const AttrStmtNode* op) override;
- void VisitStmt_(const AssertStmtNode* op) override;
- void VisitStmt_(const EvaluateNode* op) override;
- void VisitStmt_(const SeqStmtNode* op) override;
- void VisitStmt_(const DeclBufferNode* op) override;
+ void Dispatch_(const BindNode* op) override;
+ void Dispatch_(const BufferStoreNode* op) override;
+ void Dispatch_(const ForNode* op) override;
+ void Dispatch_(const WhileNode* op) override;
+ void Dispatch_(const ReturnNode* op) override;
+ void Dispatch_(const BreakNode* op) override;
+ void Dispatch_(const ContinueNode* op) override;
+ void Dispatch_(const IfThenElseNode* op) override;
+ void Dispatch_(const AllocBufferNode* op) override;
+ void Dispatch_(const AttrStmtNode* op) override;
+ void Dispatch_(const AssertStmtNode* op) override;
+ void Dispatch_(const EvaluateNode* op) override;
+ void Dispatch_(const SeqStmtNode* op) override;
+ void Dispatch_(const DeclBufferNode* op) override;
/*!
* \brief Print expr representing the thread tag
diff --git a/src/target/source/codegen_c_host.cc
b/src/target/source/codegen_c_host.cc
index 599f45d890..8e6b4d1fe5 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -320,7 +320,7 @@ void CodeGenCHost::Dispatch_(const CallNode* op,
std::ostream& os) { // NOLINT(
}
}
-void CodeGenCHost::VisitStmt_(const AssertStmtNode* op) { // NOLINT(*)
+void CodeGenCHost::Dispatch_(const AssertStmtNode* op) { // NOLINT(*)
if (emit_asserts_) {
std::string cond = PrintExpr(op->condition);
PrintIndent();
diff --git a/src/target/source/codegen_c_host.h
b/src/target/source/codegen_c_host.h
index 1176a86827..34ffa3c809 100644
--- a/src/target/source/codegen_c_host.h
+++ b/src/target/source/codegen_c_host.h
@@ -69,7 +69,7 @@ class CodeGenCHost : public CodeGenC {
void Dispatch_(const prim::MinNode* op, std::ostream& os) final; //
NOLINT(*)
void Dispatch_(const prim::MaxNode* op, std::ostream& os) final; //
NOLINT(*)
- void VisitStmt_(const AssertStmtNode* op) final; // NOLINT(*)
+ void Dispatch_(const AssertStmtNode* op) final; // NOLINT(*)
void GenerateForwardFunctionDeclarations(ffi::String global_symbol,
const ffi::Array<Type>& arg_types,
diff --git a/src/tirx/analysis/verify_tirx_well_formed.cc
b/src/tirx/analysis/verify_tirx_well_formed.cc
index 889aba164f..1fab6449cd 100644
--- a/src/tirx/analysis/verify_tirx_well_formed.cc
+++ b/src/tirx/analysis/verify_tirx_well_formed.cc
@@ -51,16 +51,15 @@ class ExecScopeVerifier : public
Verifier<ExecScopeVerifier> {
private:
using Verifier::Visit;
- void VisitStmt_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
+ void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at "
<< path;
}
- void VisitStmt_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
+ void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True
mode at " << path;
}
- void VisitStmt_(const tirx::TilePrimitiveCallNode* op,
- ffi::reflection::AccessPath path) override {
+ void Dispatch_(const tirx::TilePrimitiveCallNode* op,
ffi::reflection::AccessPath path) override {
static const auto& category_map =
Op::GetAttrMap<tirx::TIRxOpCategory>("TIRxOpCategory");
Verify(category_map.get(op->op, ffi::String("")) == "tile_primitive")
<< "TIRxError: TilePrimitiveCall at " << path << " has non-tile op "
<< op->op;
@@ -74,12 +73,12 @@ class ScopeIdVerifier : public Verifier<ScopeIdVerifier> {
private:
using Verifier::Visit;
- void VisitStmt_(const AttrStmtNode* op, ffi::reflection::AccessPath path)
override {
+ void Dispatch_(const AttrStmtNode* op, ffi::reflection::AccessPath path)
override {
if (op->attr_key == tvm::tirx::attr::kDeviceEntry) {
// Device-region marker: defs gathered from the body are verified when
// the AttrStmt exits, with launch-param sanity enforced as ``is_root``.
size_t baseline = scope_id_def_.size();
- Verifier::VisitStmt_(op, path);
+ Verifier::Dispatch_(op, path);
size_t total = scope_id_def_.size();
if (total > baseline) {
RunScopeIdVerify(path, baseline, /*is_root=*/true);
@@ -89,7 +88,7 @@ class ScopeIdVerifier : public Verifier<ScopeIdVerifier> {
}
return;
}
- Verifier::VisitStmt_(op, path);
+ Verifier::Dispatch_(op, path);
}
void RunScopeIdVerify(ffi::reflection::AccessPath path, size_t baseline,
bool is_root) {
@@ -118,9 +117,9 @@ class ScopeIdVerifier : public Verifier<ScopeIdVerifier> {
}
}
- void VisitStmt_(const ScopeIdDefStmtNode* op, ffi::reflection::AccessPath
path) override {
+ void Dispatch_(const ScopeIdDefStmtNode* op, ffi::reflection::AccessPath
path) override {
scope_id_def_.push_back(op->def);
- Verifier::VisitStmt_(op, path);
+ Verifier::Dispatch_(op, path);
}
Array<ScopeIdDef> scope_id_def_;
@@ -134,11 +133,11 @@ class LayoutVerifier : public Verifier<LayoutVerifier> {
private:
using Verifier::Visit;
- void VisitStmt_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
+ void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at "
<< path;
}
- void VisitStmt_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
+ void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True
mode at " << path;
}
};
@@ -150,11 +149,11 @@ class AsyncStructsVerifier : public
Verifier<AsyncStructsVerifier> {
private:
using Verifier::Visit;
- void VisitStmt_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
+ void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at "
<< path;
}
- void VisitStmt_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
+ void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True
mode at " << path;
}
};
@@ -166,11 +165,11 @@ class DeviceFuncVerifier : public
Verifier<DeviceFuncVerifier> {
private:
using Verifier::Visit;
- void VisitStmt_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
+ void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path)
override {
Verify(false) << "TIRxError: SBlock is not allowed in tirx=True mode at "
<< path;
}
- void VisitStmt_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
+ void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override {
Verify(false) << "TIRxError: SBlockRealize is not allowed in tirx=True
mode at " << path;
}
};
diff --git a/src/tirx/ir_mutator_with_analyzer.cc
b/src/tirx/ir/ir_mutator_with_analyzer.cc
similarity index 99%
rename from src/tirx/ir_mutator_with_analyzer.cc
rename to src/tirx/ir/ir_mutator_with_analyzer.cc
index 279ab2c29c..e26fdf5282 100644
--- a/src/tirx/ir_mutator_with_analyzer.cc
+++ b/src/tirx/ir/ir_mutator_with_analyzer.cc
@@ -18,7 +18,7 @@
*/
/*!
- * \file tirx/ir_mutator_with_analyzer.cc
+ * \file tirx/ir/ir_mutator_with_analyzer.cc
*/
#include "ir_mutator_with_analyzer.h"
@@ -30,7 +30,7 @@
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/op.h>
-#include "../arith/constraint_helpers.h"
+#include "../../arith/constraint_helpers.h"
namespace tvm {
namespace tirx {
diff --git a/src/tirx/ir_mutator_with_analyzer.h
b/src/tirx/ir/ir_mutator_with_analyzer.h
similarity index 95%
rename from src/tirx/ir_mutator_with_analyzer.h
rename to src/tirx/ir/ir_mutator_with_analyzer.h
index 5463bdb6ec..20a6dd2587 100644
--- a/src/tirx/ir_mutator_with_analyzer.h
+++ b/src/tirx/ir/ir_mutator_with_analyzer.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tirx/ir_mutator_with_analyzer.h
+ * \file tirx/ir/ir_mutator_with_analyzer.h
* \brief IR mutator base-class with an analyzer context.
*/
-#ifndef TVM_TIRX_IR_MUTATOR_WITH_ANALYZER_H_
-#define TVM_TIRX_IR_MUTATOR_WITH_ANALYZER_H_
+#ifndef TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_
+#define TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
@@ -45,7 +45,7 @@ namespace tirx {
* It will populates scope-related info such as bounds of loop-variables and
constraints
* for the analyzer, so that the child class can do accurate context-dependent
analysis.
*
- * \sa src/tirx/ir_mutator_with_analyzer.cc
+ * \sa src/tirx/ir/ir_mutator_with_analyzer.cc
*/
class IRMutatorWithAnalyzer : public StmtExprMutator {
public:
@@ -123,4 +123,4 @@ class IRMutatorWithAnalyzer : public StmtExprMutator {
};
} // namespace tirx
} // namespace tvm
-#endif // TVM_TIRX_IR_MUTATOR_WITH_ANALYZER_H_
+#endif // TVM_TIRX_IR_IR_MUTATOR_WITH_ANALYZER_H_
diff --git a/src/tirx/ir_visitor_with_analyzer.cc
b/src/tirx/ir/ir_visitor_with_analyzer.cc
similarity index 99%
rename from src/tirx/ir_visitor_with_analyzer.cc
rename to src/tirx/ir/ir_visitor_with_analyzer.cc
index 639768f6c0..9bdc895722 100644
--- a/src/tirx/ir_visitor_with_analyzer.cc
+++ b/src/tirx/ir/ir_visitor_with_analyzer.cc
@@ -18,7 +18,7 @@
*/
/*!
- * \file tirx/ir_visitor_with_analyzer.cc
+ * \file tirx/ir/ir_visitor_with_analyzer.cc
*/
#include "ir_visitor_with_analyzer.h"
diff --git a/src/tirx/ir_visitor_with_analyzer.h
b/src/tirx/ir/ir_visitor_with_analyzer.h
similarity index 92%
rename from src/tirx/ir_visitor_with_analyzer.h
rename to src/tirx/ir/ir_visitor_with_analyzer.h
index d1e36b86cf..33a895acf6 100644
--- a/src/tirx/ir_visitor_with_analyzer.h
+++ b/src/tirx/ir/ir_visitor_with_analyzer.h
@@ -18,12 +18,12 @@
*/
/*!
- * \file tirx/ir_visitor_with_analyzer.h
+ * \file tirx/ir/ir_visitor_with_analyzer.h
* \brief IR visitor class with an analyzer context.
*/
-#ifndef TVM_TIRX_IR_VISITOR_WITH_ANALYZER_H_
-#define TVM_TIRX_IR_VISITOR_WITH_ANALYZER_H_
+#ifndef TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_
+#define TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_
#include <tvm/arith/analyzer.h>
#include <tvm/ir/prim/expr.h>
@@ -70,4 +70,4 @@ class IRVisitorWithAnalyzer : public StmtExprVisitor {
} // namespace tirx
} // namespace tvm
-#endif // TVM_TIRX_IR_VISITOR_WITH_ANALYZER_H_
+#endif // TVM_TIRX_IR_IR_VISITOR_WITH_ANALYZER_H_
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index 53fe95f12c..0f0f684a50 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -179,14 +179,14 @@ void TIRVisitorWithPath::Visit(const Range& range,
AccessPath path) {
Visit(range->extent, path->Attr("extent"));
}
-void TIRVisitorWithPath::VisitStmt_(const BindNode* op, AccessPath path) {
+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::VisitStmt_(const AttrStmtNode* op, AccessPath path) {
+void TIRVisitorWithPath::Dispatch_(const AttrStmtNode* op, AccessPath path) {
Visit(op->value, path->Attr("value"));
std::vector<std::variant<DefContext<IterVar>, DefContext<Var>,
DefContext<BufferVar>>> context;
@@ -207,68 +207,68 @@ void TIRVisitorWithPath::VisitStmt_(const AttrStmtNode*
op, AccessPath path) {
}
}
-void TIRVisitorWithPath::VisitStmt_(const ForNode* op, AccessPath path) {
+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::VisitStmt_(const WhileNode* op, AccessPath path) {
+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::VisitStmt_(const ReturnNode* op, AccessPath path) {
+void TIRVisitorWithPath::Dispatch_(const ReturnNode* op, AccessPath path) {
Visit(op->value, path->Attr("value"));
}
-void TIRVisitorWithPath::VisitStmt_(const BreakNode* op, AccessPath path) {}
+void TIRVisitorWithPath::Dispatch_(const BreakNode* op, AccessPath path) {}
-void TIRVisitorWithPath::VisitStmt_(const ContinueNode* op, AccessPath path) {}
+void TIRVisitorWithPath::Dispatch_(const ContinueNode* op, AccessPath path) {}
-void TIRVisitorWithPath::VisitStmt_(const AllocBufferNode* op, AccessPath
path) {
+void TIRVisitorWithPath::Dispatch_(const AllocBufferNode* op, AccessPath path)
{
// Push definitions into the current scope so they are visible to subsequent
siblings.
auto buf_path = path->Attr("buffer");
bind_scope_.Current().push_back(WithDef(op->buffer, buf_path));
}
-void TIRVisitorWithPath::VisitStmt_(const DeclBufferNode* op, AccessPath path)
{
+void TIRVisitorWithPath::Dispatch_(const DeclBufferNode* op, AccessPath path) {
Visit(op->data, path->Attr("data"));
// Push buffer definition into the current scope so it is visible to
subsequent siblings.
bind_scope_.Current().push_back(WithDef(op->buffer, path->Attr("buffer")));
}
-void TIRVisitorWithPath::VisitStmt_(const BufferStoreNode* op, AccessPath
path) {
+void TIRVisitorWithPath::Dispatch_(const BufferStoreNode* op, AccessPath path)
{
Visit(op->value, path->Attr("value"));
VisitBufferUse(op->buffer, path->Attr("buffer"));
Visit(op->indices, path->Attr("indices"));
}
-void TIRVisitorWithPath::VisitStmt_(const IfThenElseNode* op, AccessPath path)
{
+void TIRVisitorWithPath::Dispatch_(const IfThenElseNode* 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::VisitStmt_(const AssertStmtNode* op, AccessPath path)
{
+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::VisitStmt_(const SeqStmtNode* op, AccessPath path) {
+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::VisitStmt_(const EvaluateNode* op, AccessPath path) {
+void TIRVisitorWithPath::Dispatch_(const EvaluateNode* op, AccessPath path) {
Visit(op->value, path->Attr("value"));
}
-void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op, AccessPath path) {
+void TIRVisitorWithPath::Dispatch_(const SBlockNode* op, AccessPath path) {
std::vector<std::variant<DefContext<Var>, DefContext<IterVar>,
DefContext<BufferVar>>> context;
{
@@ -313,13 +313,13 @@ void TIRVisitorWithPath::VisitStmt_(const SBlockNode* op,
AccessPath path) {
while (context.size()) context.pop_back();
}
-void TIRVisitorWithPath::VisitStmt_(const SBlockRealizeNode* op, AccessPath
path) {
+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"));
}
-void TIRVisitorWithPath::VisitStmt_(const tirx::TilePrimitiveCallNode* op,
AccessPath path) {
+void TIRVisitorWithPath::Dispatch_(const tirx::TilePrimitiveCallNode* op,
AccessPath path) {
for (size_t i = 0; i < op->args.size(); i++) {
if (op->args[i] == nullptr) {
continue;
@@ -336,7 +336,7 @@ void TIRVisitorWithPath::VisitStmt_(const
tirx::TilePrimitiveCallNode* op, Acces
}
}
-void TIRVisitorWithPath::VisitStmt_(const ScopeIdDefStmtNode* op, AccessPath
path) {
+void TIRVisitorWithPath::Dispatch_(const ScopeIdDefStmtNode* op, AccessPath
path) {
// Flat stmt -- no body. Visit extents and preferred_extents (if present),
// then push the bound Var(s) into the current scope so subsequent siblings
// see them as defined.
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index ef54a7c35e..f4547becd6 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -75,9 +75,9 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
TVM_FFI_THROW(TypeError) << "Unsupported non-primitive TIR expression "
<< obj.GetTypeKey();
}
}
- // Delegate to ExprFunctor::VisitStmt for Stmt, and any subclasses
+ // Delegate to StmtFunctor::Dispatch for Stmt, and any subclasses
virtual inline void Visit(const Stmt& obj, ffi::reflection::AccessPath path)
{
- VisitStmt(obj, path);
+ Dispatch(obj, path);
}
// Visit a buffer at a use site (BufferLoad, BufferStore, reads/writes).
@@ -132,25 +132,25 @@ class TIRVisitorWithPath : protected
ExprFunctor<void(const Expr&, ffi::reflecti
}
}
- using StmtFunctor::VisitStmt;
- void VisitStmt_(const BindNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const AttrStmtNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const IfThenElseNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const ForNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const WhileNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const ReturnNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const BreakNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const ContinueNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const AllocBufferNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const DeclBufferNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const BufferStoreNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const AssertStmtNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const SeqStmtNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const EvaluateNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const SBlockNode* op, ffi::reflection::AccessPath path)
override;
- void VisitStmt_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override;
- void VisitStmt_(const tirx::TilePrimitiveCallNode* op,
ffi::reflection::AccessPath path) override;
- void VisitStmt_(const ScopeIdDefStmtNode* op, ffi::reflection::AccessPath
path) override;
+ using StmtFunctor::Dispatch;
+ void Dispatch_(const BindNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const AttrStmtNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const IfThenElseNode* 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 AllocBufferNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const DeclBufferNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const BufferStoreNode* 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;
+ void Dispatch_(const SBlockNode* op, ffi::reflection::AccessPath path)
override;
+ void Dispatch_(const SBlockRealizeNode* op, ffi::reflection::AccessPath
path) override;
+ void Dispatch_(const tirx::TilePrimitiveCallNode* op,
ffi::reflection::AccessPath path) override;
+ void Dispatch_(const ScopeIdDefStmtNode* op, ffi::reflection::AccessPath
path) override;
using ExprFunctor::Dispatch;
void Dispatch_(const VarNode* op, ffi::reflection::AccessPath path) override;
diff --git a/src/tirx/transform/flatten_buffer.cc
b/src/tirx/transform/flatten_buffer.cc
index 7baad9cc3a..0bf7e3cb43 100644
--- a/src/tirx/transform/flatten_buffer.cc
+++ b/src/tirx/transform/flatten_buffer.cc
@@ -32,7 +32,7 @@
#include <unordered_set>
-#include "../ir_mutator_with_analyzer.h"
+#include "../ir/ir_mutator_with_analyzer.h"
#include "ir_utils.h"
namespace tvm {
diff --git a/src/tirx/transform/lower_intrin.cc
b/src/tirx/transform/lower_intrin.cc
index 8baf89355e..53af3e51c2 100644
--- a/src/tirx/transform/lower_intrin.cc
+++ b/src/tirx/transform/lower_intrin.cc
@@ -38,7 +38,7 @@
#include <unordered_set>
#include "../../arith/pattern_match.h"
-#include "../ir_mutator_with_analyzer.h"
+#include "../ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace tirx {
diff --git a/src/tirx/transform/lower_tirx_cleanup.cc
b/src/tirx/transform/lower_tirx_cleanup.cc
index 929a9e079e..9a05731c93 100644
--- a/src/tirx/transform/lower_tirx_cleanup.cc
+++ b/src/tirx/transform/lower_tirx_cleanup.cc
@@ -38,7 +38,7 @@
#include <utility>
#include <vector>
-#include "../ir_mutator_with_analyzer.h"
+#include "../ir/ir_mutator_with_analyzer.h"
#include "ir_utils.h"
namespace tvm {
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 0a58927e1b..a688acd074 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -37,7 +37,7 @@
#include "../../arith/const_fold.h"
#include "../analysis/var_use_def_analysis.h"
-#include "../ir_mutator_with_analyzer.h"
+#include "../ir/ir_mutator_with_analyzer.h"
#include "ir_utils.h"
namespace tvm {
diff --git a/src/tirx/transform/stmt_simplify.cc
b/src/tirx/transform/stmt_simplify.cc
index 3d9dea2b68..f6904e9b56 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -36,7 +36,7 @@
#include <tvm/tirx/op.h>
#include <tvm/tirx/transform.h>
-#include "../ir_mutator_with_analyzer.h"
+#include "../ir/ir_mutator_with_analyzer.h"
namespace tvm {
namespace arith {
diff --git a/src/tirx/transform/tile_primitive_dispatch.cc
b/src/tirx/transform/tile_primitive_dispatch.cc
index c0535bd17c..73634b18fc 100644
--- a/src/tirx/transform/tile_primitive_dispatch.cc
+++ b/src/tirx/transform/tile_primitive_dispatch.cc
@@ -246,7 +246,7 @@ class NoOpCallVerifier : public Verifier<NoOpCallVerifier> {
private:
using Verifier::Visit;
- void VisitStmt_(const tirx::TilePrimitiveCallNode* obj,
ffi::reflection::AccessPath path) final {
+ void Dispatch_(const tirx::TilePrimitiveCallNode* obj,
ffi::reflection::AccessPath path) final {
Verify(false) << "TIRxError: TilePrimitiveCall at " << path
<< " is not allowed in TIRx before lowering";
}
@@ -322,7 +322,7 @@ class TilePrimitiveDispatcher : public StmtExprMutator {
launch_params_.clear();
// Pre-dispatch: only populate ``launch_params_`` + synthesize
- // ``warp_id_in_cta``. The dispatch impls (run via ``VisitStmt`` below)
+ // ``warp_id_in_cta``. The dispatch impls (run via ``Dispatch`` below)
// read ``launch_params_`` through ``sctx``, so this much must happen
// first. The per-def Bind resolution is deferred to AFTER dispatch so
// it can pick up any ``ScopeIdDef`` declared inside dispatched impls.
diff --git a/tests/cpp/ir_functor_test.cc b/tests/cpp/ir_functor_test.cc
index abbdf56035..5456119c8a 100644
--- a/tests/cpp/ir_functor_test.cc
+++ b/tests/cpp/ir_functor_test.cc
@@ -181,6 +181,8 @@ TEST(IRF, ExprVisit) {
class MyVisitor : public tirx::ExprFunctor<void(const Expr&)>,
public tirx::StmtFunctor<void(const Stmt&)> {
public:
+ using tirx::ExprFunctor<void(const Expr&)>::Dispatch;
+ using tirx::StmtFunctor<void(const Stmt&)>::Dispatch;
int count = 0;
// implementation
void Dispatch_(const VarNode* op) final { ++count; }
@@ -189,10 +191,10 @@ TEST(IRF, ExprVisit) {
Dispatch(op->a);
Dispatch(op->b);
}
- void VisitStmt_(const EvaluateNode* op) final { Dispatch(op->value); }
+ void Dispatch_(const EvaluateNode* op) final { Dispatch(op->value); }
};
MyVisitor v;
- v.VisitStmt(Evaluate(z));
+ v.Dispatch(Evaluate(z));
TVM_FFI_ICHECK_EQ(v.count, 1);
}
diff --git a/tests/cpp/stmt_functor_test.cc b/tests/cpp/stmt_functor_test.cc
new file mode 100644
index 0000000000..d799819003
--- /dev/null
+++ b/tests/cpp/stmt_functor_test.cc
@@ -0,0 +1,131 @@
+/*
+ * 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 <gtest/gtest.h>
+#include <tvm/tirx/stmt_functor.h>
+
+#include <memory>
+
+namespace tvm {
+namespace tirx {
+namespace {
+
+class ExtensionStmtNode : public StmtNode {
+ public:
+ static constexpr uint32_t _type_child_slots = 1;
+ TVM_FFI_DECLARE_OBJECT_INFO("test.ExtensionStmt", ExtensionStmtNode,
StmtNode);
+};
+
+class ChildExtensionStmtNode : public ExtensionStmtNode {
+ public:
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.ChildExtensionStmt",
ChildExtensionStmtNode,
+ ExtensionStmtNode);
+};
+
+using ForwardFunctor = StmtFunctor<int(const Stmt&, std::unique_ptr<int>,
int&)>;
+
+class ExtendedFunctor : public ForwardFunctor {
+ public:
+ using ForwardFunctor::Dispatch_;
+ ExtendedFunctor() : ForwardFunctor(Table()) {}
+
+ int Dispatch_(const EvaluateNode*, std::unique_ptr<int> value, int& count)
override {
+ ++count;
+ return *value;
+ }
+ virtual int Dispatch_(const ExtensionStmtNode*, std::unique_ptr<int> value,
int& count) {
+ ++count;
+ return *value + 10;
+ }
+ int DispatchDefault_(const ffi::Object*, std::unique_ptr<int> value, int&
count) override {
+ ++count;
+ return -*value;
+ }
+
+ private:
+ static const VTable* Table() {
+ static const VTable table = [] {
+ VTable table;
+ ForwardFunctor::InitVTable(&table);
+ SetDispatch<ExtendedFunctor, ExtensionStmtNode>(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+};
+
+TEST(StmtFunctor, ExtendedTableForwardingAndAncestorDispatch) {
+ ExtendedFunctor functor;
+ int calls = 0;
+ auto value = std::make_unique<int>(3);
+ EXPECT_EQ(functor(Evaluate(IntImm::Int32(1)), std::move(value), calls), 3);
+ EXPECT_EQ(value, nullptr);
+ EXPECT_EQ(calls, 1);
+
+ Stmt extension(ffi::make_object<ExtensionStmtNode>());
+ EXPECT_EQ(functor.Dispatch(extension, std::make_unique<int>(4), calls), 14);
+ Stmt child(ffi::make_object<ChildExtensionStmtNode>());
+ EXPECT_EQ(functor(child, std::make_unique<int>(5), calls), 15);
+ EXPECT_EQ(calls, 3);
+
+ // An inherited default hook receives the same move-only and reference
arguments.
+ EXPECT_EQ(functor(Break(Span()), std::make_unique<int>(6), calls), -6);
+ EXPECT_EQ(calls, 4);
+ // Extending one finalized table does not alter the base signature's default
table.
+ ForwardFunctor base;
+ EXPECT_THROW(base(extension, std::make_unique<int>(1), calls), ffi::Error);
+ EXPECT_EQ(calls, 4);
+}
+
+TEST(StmtFunctor, EntryPointLifetimeAndNoRecursion) {
+ class Functor : public StmtFunctor<int(const Stmt&)> {
+ public:
+ explicit Functor(bool* destroyed) : destroyed_(destroyed) {}
+ ~Functor() override { *destroyed_ = true; }
+ int Dispatch(const Stmt& node) override { return 10 +
StmtFunctor::Dispatch(node); }
+ int Dispatch_(const SeqStmtNode*) override { return 2; }
+
+ private:
+ bool* destroyed_;
+ };
+ bool destroyed = false;
+ {
+ std::unique_ptr<StmtFunctor<int(const Stmt&)>> functor =
std::make_unique<Functor>(&destroyed);
+ // Evaluate would throw if the dispatcher recursed into the sequence.
+ EXPECT_EQ((*functor)(SeqStmt({Evaluate(IntImm::Int32(1)),
Evaluate(IntImm::Int32(2))})), 12);
+ EXPECT_THROW((*functor)(Evaluate(IntImm::Int32(1))), ffi::Error);
+ EXPECT_THROW((*functor)(Stmt(nullptr)), ffi::Error);
+ }
+ EXPECT_TRUE(destroyed);
+}
+
+TEST(StmtFunctor, ReferenceResult) {
+ class Functor : public StmtFunctor<int&(const Stmt&, int&)> {
+ public:
+ int& Dispatch_(const EvaluateNode*, int& value) override { return value; }
+ };
+ Functor functor;
+ int value = 3;
+ functor(Evaluate(IntImm::Int32(1)), value) = 7;
+ EXPECT_EQ(value, 7);
+}
+
+} // namespace
+} // namespace tirx
+} // namespace tvm
