This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 7fca2e1716 [REFACTOR][IR] New Functor/Visitor/Mutator Infra at Base
Layer (#20327)
7fca2e1716 is described below
commit 7fca2e17160bae19d0311c327042d5859361a92e
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Sep 12 21:15:00 2026 -0400
[REFACTOR][IR] New Functor/Visitor/Mutator Infra at Base Layer (#20327)
This PR introduces a new `expr_functor.h` with `ExprMutator` and
`ExprVisitor` aligned with the `StructuralMutate`/`StructuralVisit`
system.
Both provide virtual hooks for individual expression kinds. Default
hooks traverse children, with structural fallback for types without
exact registrations. `ExprMutator` supports copy-on-write rewriting and
variable remapping; `ExprVisitor` supports early interruption. The
header also provides `ExprFunctor<R(...)>` for custom return types and
arguments, with caller-controlled recursion.
## Guidelines
- Use `StructuralMap` for common rewriting and `StructuralWalk` or
`StructuralVisit` for common traversal.
- Use `ExprMutator` or `ExprVisitor` for per-node virtual overrides with
default child traversal.
- Use `ExprFunctor` for typed dispatch with custom results, arguments,
or recursion.
## Notes
### Traversal beyond expressions
`ExprMutator` and `ExprVisitor` are structural mutator and visitor
objects with expression hooks. Unlike the older expression visitors and
mutators, they can follow registered structural fields through other IR
nodes without a dedicated hook, reaching expressions nested inside them.
Expression types without exact hook registrations also use this
structural fallback.
```cpp
// Visit expressions inside a function through its registered structural
fields.
auto visitor = ffi::make_object<ExprVisitor>();
visitor->VisitExpected(function).value();
```
### Adding hooks in subclasses
Override existing hooks directly. To add a hook for a custom expression
type, initialize an inherited dispatch table and register the new hook.
Here, `MyExprNode` is a custom `ExprNode` subtype:
```cpp
class MyExprMutator : public ExprMutator {
public:
TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(MyExprMutator, ExprMutator)
using ExprMutator::Mutate_;
virtual Expected<UnchangedOr<ffi::Any>> Mutate_(
const MyExprNode* node, bool allow_inplace);
protected:
static void InitVTable(VTable* table) {
ExprMutator::InitVTable(table);
SetDispatch<MyExprMutator, MyExprNode>(table);
}
};
```
Define the new `Mutate_` hook with the custom rewrite logic. Further
subclasses can override it; `ExprVisitor` supports the same registration
pattern with `Visit_` hooks.
---
3rdparty/tvm-ffi | 2 +-
include/tvm/ir/expr_functor.h | 472 ++++++++++++++++++++++++++
include/tvm/ir/object_functor.h | 398 +++++++++++++++++++++-
src/ir/expr_functor.cc | 722 ++++++++++++++++++++++++++++++++++++++++
tests/cpp/expr_functor_test.cc | 95 ++++++
5 files changed, 1677 insertions(+), 12 deletions(-)
diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi
index c5e636bbfb..be35ec1297 160000
--- a/3rdparty/tvm-ffi
+++ b/3rdparty/tvm-ffi
@@ -1 +1 @@
-Subproject commit c5e636bbfb097cca610acee841e26d751d47a63b
+Subproject commit be35ec1297d6ddf388b7feb4f2e89d2cf030dda8
diff --git a/include/tvm/ir/expr_functor.h b/include/tvm/ir/expr_functor.h
new file mode 100644
index 0000000000..b7115ed6c5
--- /dev/null
+++ b/include/tvm/ir/expr_functor.h
@@ -0,0 +1,472 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+/*!
+ * \file tvm/ir/expr_functor.h
+ * \brief Native visiting and mutation of core IR expressions.
+ */
+#ifndef TVM_IR_EXPR_FUNCTOR_H_
+#define TVM_IR_EXPR_FUNCTOR_H_
+#include <tvm/ir/expr.h>
+#include <tvm/ir/object_functor.h>
+#include <tvm/ir/op.h>
+#include <tvm/ir/prim/expr.h>
+#include <tvm/ir/prim/vector_expr.h>
+
+namespace tvm {
+
+/*!
+ * \brief Type-dispatched expression functor with a caller-selected signature.
+ * \tparam FType A function signature of the form R(const Expr&, Args...).
+ *
+ * Override Dispatch_ for a node type, or DispatchDefault_ for the default
+ * behavior. Unlike structural visitors and mutators, this functor does not
+ * traverse children automatically. Dispatch may use a registered ancestor.
+ * A derived class can add node types with a fresh inherited table.
+ * Use tvm::ExprFunctor explicitly when dialect functors are also in scope.
+ */
+template <typename FType>
+class ExprFunctor;
+
+template <typename R, typename... Args>
+class ExprFunctor<R(const Expr&, Args...)> {
+ private:
+ using TSelf = ExprFunctor<R(const Expr&, Args...)>;
+
+ public:
+ /*! \brief The result type of this functor. */
+ using result_type = R;
+ /*! \brief Construct a functor with the core expression hooks. */
+ ExprFunctor() : ExprFunctor(GlobalVTable()) {}
+ /*!
+ * \brief Dispatch an expression through Dispatch.
+ * \param node The borrowed expression.
+ * \param args Additional arguments forwarded to the selected hook.
+ * \return The hook result.
+ */
+ TVM_FFI_INLINE R operator()(const Expr& node, Args... args) {
+ return Dispatch(node, std::forward<Args>(args)...);
+ }
+ /*!
+ * \brief Dispatch to a node hook, including registered ancestor hooks.
+ * \param node The non-null borrowed expression.
+ * \param args Additional arguments forwarded to the selected hook.
+ * \return The hook result.
+ */
+ TVM_FFI_INLINE virtual R Dispatch(const Expr& node, Args... args) {
+ TVM_FFI_ICHECK(node.defined()) << "Cannot dispatch a null expression";
+ return (*vtable_)(node, this, std::forward<Args>(args)...);
+ }
+
+ virtual R Dispatch_(const OpaqueExprNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const TupleNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const TupleGetItemNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const TensorLoadNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const VarNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const GlobalVarNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const CallNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const IntImmNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const FloatImmNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const OpNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::StringImmNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::CastNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::AddNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::SubNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::MulNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::DivNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::ModNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::FloorDivNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::FloorModNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::MinNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::MaxNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::EQNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::NENode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::LTNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::LENode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::GTNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::GENode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::AndNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::OrNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::NotNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::SelectNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::LetNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::RampNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::BroadcastNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+ virtual R Dispatch_(const prim::ShuffleNode* node, Args... args) {
+ return DispatchDefault_(node, std::forward<Args>(args)...);
+ }
+
+ /*!
+ * \brief Default node behavior, overridden by subclasses that handle
arbitrary core nodes.
+ * \param node The borrowed node.
+ * \param args Additional arguments supplied to the functor.
+ * \return The hook result.
+ */
+ virtual R DispatchDefault_(const ffi::Object* node, Args... 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 a functor with an extended finalized table.
+ * \param vtable The table, which must outlive the functor.
+ */
+ explicit ExprFunctor(const VTable* vtable) : vtable_(vtable) {}
+ /*!
+ * \brief Register core expression hooks in a fresh mutable table.
+ * \param vtable The table to initialize before adding derived registrations.
+ */
+ static void InitVTable(VTable* vtable) {
+ SetDispatch<TSelf, OpaqueExprNode>(vtable);
+ SetDispatch<TSelf, TupleNode>(vtable);
+ SetDispatch<TSelf, TupleGetItemNode>(vtable);
+ SetDispatch<TSelf, TensorLoadNode>(vtable);
+ SetDispatch<TSelf, VarNode>(vtable);
+ SetDispatch<TSelf, GlobalVarNode>(vtable);
+ SetDispatch<TSelf, CallNode>(vtable);
+ SetDispatch<TSelf, IntImmNode>(vtable);
+ SetDispatch<TSelf, FloatImmNode>(vtable);
+ SetDispatch<TSelf, OpNode>(vtable);
+ SetDispatch<TSelf, prim::StringImmNode>(vtable);
+ SetDispatch<TSelf, prim::CastNode>(vtable);
+ SetDispatch<TSelf, prim::AddNode>(vtable);
+ SetDispatch<TSelf, prim::SubNode>(vtable);
+ SetDispatch<TSelf, prim::MulNode>(vtable);
+ SetDispatch<TSelf, prim::DivNode>(vtable);
+ SetDispatch<TSelf, prim::ModNode>(vtable);
+ SetDispatch<TSelf, prim::FloorDivNode>(vtable);
+ SetDispatch<TSelf, prim::FloorModNode>(vtable);
+ SetDispatch<TSelf, prim::MinNode>(vtable);
+ SetDispatch<TSelf, prim::MaxNode>(vtable);
+ SetDispatch<TSelf, prim::EQNode>(vtable);
+ SetDispatch<TSelf, prim::NENode>(vtable);
+ SetDispatch<TSelf, prim::LTNode>(vtable);
+ SetDispatch<TSelf, prim::LENode>(vtable);
+ SetDispatch<TSelf, prim::GTNode>(vtable);
+ SetDispatch<TSelf, prim::GENode>(vtable);
+ SetDispatch<TSelf, prim::AndNode>(vtable);
+ SetDispatch<TSelf, prim::OrNode>(vtable);
+ SetDispatch<TSelf, prim::NotNode>(vtable);
+ SetDispatch<TSelf, prim::SelectNode>(vtable);
+ SetDispatch<TSelf, prim::LetNode>(vtable);
+ SetDispatch<TSelf, prim::RampNode>(vtable);
+ SetDispatch<TSelf, prim::BroadcastNode>(vtable);
+ SetDispatch<TSelf, prim::ShuffleNode>(vtable);
+ }
+ /*!
+ * \brief Register a hook for an additional expression type.
+ * \tparam Self The subclass implementing the hook.
+ * \tparam Node The node type handled by the hook.
+ * \param vtable The mutable table to receive the registration.
+ */
+ 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:
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+ const VTable* vtable_;
+};
+
+/*!
+ * \brief Structural Expr visitor with native overrides for core expressions.
+ *
+ * Prefer StructuralVisit or StructuralWalk for common cases; use ExprVisitor
+ * for extensive per-kind customization or optimization. Default hooks follow
+ * the registered structural child order and skip primitive metadata.
+ * Native hooks match exact node types; derived node types use structural
fallback
+ * unless separately registered in a fresh inherited table with SetDispatch.
+ * Existing hooks only require overriding. For an extra MyExprNode derived from
+ * ExprNode, initialize an inherited table and register the new virtual hook:
+ * \code
+ * class MyExprVisitor : public ExprVisitor {
+ * public:
+ * TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(MyExprVisitor, ExprVisitor)
+ * using ExprVisitor::Visit_;
+ * virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const MyExprNode*
node);
+ *
+ * protected:
+ * static void InitVTable(VTable* vtable) {
+ * ExprVisitor::InitVTable(vtable);
+ * SetDispatch<MyExprVisitor, MyExprNode>(vtable);
+ * }
+ * };
+ * \endcode
+ * Use tvm::ExprVisitor explicitly when dialect visitors are also in scope.
+ */
+class TVM_DLL ExprVisitor : public ObjectVisitor {
+ public:
+ /*! \brief Construct a visitor with the core expression hooks. */
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(ExprVisitor, ObjectVisitor)
+
+ using ObjectVisitor::VisitExpected;
+
+ // Override existing hooks directly. Extra types need a fresh inherited
table.
+ // Hooks borrow the node and return None, an owning interrupt, or an Error.
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const OpaqueExprNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const TupleNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
TupleGetItemNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const TensorLoadNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const VarNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const GlobalVarNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const CallNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const IntImmNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const FloatImmNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const OpNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
prim::StringImmNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::CastNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::AddNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::SubNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::MulNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::DivNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::ModNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
prim::FloorDivNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
prim::FloorModNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::MinNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::MaxNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::EQNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::NENode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::LTNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::LENode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::GTNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::GENode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::AndNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::OrNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::NotNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
prim::SelectNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::LetNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const prim::RampNode*
node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
prim::BroadcastNode* node);
+ virtual Expected<ffi::Optional<VisitInterrupt>> Visit_(const
prim::ShuffleNode* node);
+
+ protected:
+ /*!
+ * \brief Construct a visitor with an extended native dispatch table.
+ * \param vtable The finalized table, which must outlive this visitor.
+ */
+ explicit ExprVisitor(const VTable* vtable) : ObjectVisitor(vtable) {}
+ /*!
+ * \brief Register core expression hooks in a fresh mutable table.
+ * \param vtable The table to initialize before adding derived registrations.
+ */
+ static void InitVTable(VTable* vtable);
+};
+
+/*!
+ * \brief Structural Expr mutator with native overrides for core expressions
+ * and structural fallback for other objects.
+ *
+ * Prefer StructuralMap for common cases; use ExprMutator for extensive
per-kind
+ * customization or optimization.
+ * The default Var hook leaves PrimType variables unchanged before remap
lookup.
+ *
+ * Native hooks match exact node types; derived node types need their own
registrations.
+ * Existing hooks only require overriding. For an extra MyExprNode derived from
+ * ExprNode, initialize an inherited table and register the new virtual hook:
+ * \code
+ * class MyExprMutator : public ExprMutator {
+ * public:
+ * TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(MyExprMutator, ExprMutator)
+ * using ExprMutator::Mutate_;
+ * virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const MyExprNode* node,
bool allow_inplace);
+ *
+ * protected:
+ * static void InitVTable(VTable* vtable) {
+ * ExprMutator::InitVTable(vtable);
+ * SetDispatch<MyExprMutator, MyExprNode>(vtable);
+ * }
+ * };
+ * \endcode
+ * Use tvm::ExprMutator explicitly when dialect mutators are also in scope.
+ */
+class TVM_DLL ExprMutator : public ObjectMutator {
+ public:
+ /*! \brief Construct a mutator with the core expression hooks. */
+ TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(ExprMutator, ObjectMutator)
+
+ using ObjectMutator::MaybeInplaceMutateIfUniqueExpected;
+ using ObjectMutator::MutateExpected;
+
+ /*!
+ * \brief Mutate a borrowed expression without allowing in-place changes.
+ * \param value The borrowed expression to mutate.
+ * \return An Expr replacement, Unchanged, or an Error if mutation fails.
+ * \note This entry trusts the Expr replacement contract of native hooks and
extensions.
+ */
+ TVM_FFI_INLINE Expected<UnchangedOr<Expr>> MutateExpected(const Expr& value)
noexcept {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<Expr>>(
+
ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(ObjectMutator::MutateExpected(value)));
+ }
+ /*!
+ * \brief Forward inherited permission and check the expression's uniqueness.
+ * \param value The borrowed expression to mutate.
+ * \param allow_inplace Whether the path to this expression is already
uniquely owned.
+ * \return An Expr replacement, Unchanged, or an Error if mutation fails.
+ * \note This entry trusts the Expr replacement contract of native hooks and
extensions.
+ */
+ TVM_FFI_INLINE Expected<UnchangedOr<Expr>>
MaybeInplaceMutateIfUniqueExpected(
+ const Expr& value, bool allow_inplace = true) noexcept {
+ return ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<UnchangedOr<Expr>>(
+ ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ ObjectMutator::MaybeInplaceMutateIfUniqueExpected(value,
allow_inplace)));
+ }
+
+ // A downstream class overrides any existing hook without rebuilding the
table.
+ // Extra node types use a fresh inherited table and SetDispatch<Self,
ExtraNode>.
+ // Hooks borrow the node and return an owning Expr replacement in Any,
Unchanged, or Error.
+ // Narrower field types are checked separately. Forward allow_inplace on
every child edge.
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const OpaqueExprNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const TupleNode* node, bool
allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const TupleGetItemNode*
node, bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const TensorLoadNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const VarNode* node, bool
allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const GlobalVarNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const CallNode* node, bool
allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const IntImmNode* node, bool
allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const FloatImmNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const OpNode* node, bool
allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::StringImmNode*
node,
+ bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::CastNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::AddNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::SubNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::MulNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::DivNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::ModNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::FloorDivNode*
node,
+ bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::FloorModNode*
node,
+ bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::MinNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::MaxNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::EQNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::NENode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::LTNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::LENode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::GTNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::GENode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::AndNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::OrNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::NotNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::SelectNode*
node, bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::LetNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::RampNode* node,
bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::BroadcastNode*
node,
+ bool allow_inplace);
+ virtual Expected<UnchangedOr<ffi::Any>> Mutate_(const prim::ShuffleNode*
node,
+ bool allow_inplace);
+
+ protected:
+ /*!
+ * \brief Construct a mutator with an extended native dispatch table.
+ * \param vtable The finalized table, which must outlive this mutator.
+ */
+ explicit ExprMutator(const VTable* vtable) : ObjectMutator(vtable) {}
+ /*!
+ * \brief Register core expression hooks in a fresh mutable table.
+ * \param vtable The table to initialize before adding derived registrations.
+ */
+ static void InitVTable(VTable* vtable);
+};
+
+} // namespace tvm
+#endif // TVM_IR_EXPR_FUNCTOR_H_
diff --git a/include/tvm/ir/object_functor.h b/include/tvm/ir/object_functor.h
index 082c2c1c13..037a1d0673 100644
--- a/include/tvm/ir/object_functor.h
+++ b/include/tvm/ir/object_functor.h
@@ -18,12 +18,15 @@
*/
/*!
* \file tvm/ir/object_functor.h
- * \brief Defines the Functor data structures.
+ * \brief Native object dispatch, visiting and mutation with structural
fallback.
*/
#ifndef TVM_IR_OBJECT_FUNCTOR_H_
#define TVM_IR_OBJECT_FUNCTOR_H_
#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/runtime/base.h>
#include <cstring>
#include <type_traits>
@@ -75,13 +78,18 @@ namespace tvm {
* }
* \endcode
*
- * \tparam FType Function signature, with const ffi::ObjectRef& as its first
argument.
+ * \tparam FType Function signature, with const ffi::ObjectRef& or const
ffi::Object* as its first
+ * argument.
*/
template <typename FType>
class ObjectFunctor;
-template <typename R, typename... Args>
-class ObjectFunctor<R(const ffi::ObjectRef& n, Args...)> {
+template <typename R, typename NodeArg, typename... Args>
+class ObjectFunctor<R(NodeArg, Args...)> {
+ static_assert(std::is_same_v<NodeArg, const ffi::ObjectRef&> ||
+ std::is_same_v<NodeArg, const ffi::Object*>,
+ "ObjectFunctor requires an ObjectRef reference or borrowed
Object pointer");
+
public:
/*! \brief the result type of this functor */
using result_type = R;
@@ -90,7 +98,7 @@ class ObjectFunctor<R(const ffi::ObjectRef& n, Args...)> {
* \param n The object to be dispatched.
* \return Whether a dispatch function is registered for n's type, excluding
ancestors.
*/
- bool CanDispatch(const ffi::ObjectRef& n) const {
+ TVM_FFI_INLINE bool CanDispatch(NodeArg n) const {
uint32_t type_index = n->type_index();
if (type_index < begin_type_index_) return false;
type_index -= begin_type_index_;
@@ -102,7 +110,7 @@ class ObjectFunctor<R(const ffi::ObjectRef& n, Args...)> {
* \param args The additional arguments
* \return The result.
*/
- R operator()(const ffi::ObjectRef& n, Args... args) const {
+ TVM_FFI_INLINE R operator()(NodeArg n, Args... args) const {
uint32_t type_index = n->type_index();
if (type_index >= begin_type_index_) {
uint32_t index = type_index - begin_type_index_;
@@ -121,9 +129,7 @@ class ObjectFunctor<R(const ffi::ObjectRef& n, Args...)> {
}
}
}
- TVM_FFI_THROW(InternalError) << "ObjectFunctor calls un-registered
function on type "
- << n->GetTypeKey();
- throw;
+ ThrowUnregistered(n);
}
/*!
* \brief set the dispatcher for type TNode
@@ -132,7 +138,7 @@ class ObjectFunctor<R(const ffi::ObjectRef& n, Args...)> {
* \return reference to self.
*/
template <typename TNode>
- ObjectFunctor& SetDispatch(R (*f)(const ffi::ObjectRef& n, Args...)) {
+ ObjectFunctor& SetDispatch(R (*f)(NodeArg n, Args...)) {
uint32_t tindex = TNode::RuntimeTypeIndex();
if (func_.size() <= tindex) {
func_.resize(tindex + 1, nullptr);
@@ -178,13 +184,383 @@ class ObjectFunctor<R(const ffi::ObjectRef& n, Args...)>
{
}
private:
+ [[noreturn]] TVM_FFI_COLD_CODE static void ThrowUnregistered(NodeArg n) {
+ TVM_FFI_THROW(InternalError) << "ObjectFunctor calls un-registered
function on type "
+ << n->GetTypeKey();
+ throw;
+ }
+
/*! \brief internal function pointer type */
- using FPointer = R (*)(const ffi::ObjectRef& n, Args...);
+ using FPointer = R (*)(NodeArg n, Args...);
/*! \brief internal function table */
std::vector<FPointer> func_;
/*! \brief start range of func index */
uint32_t begin_type_index_{0};
};
+using ffi::Expected;
+using ffi::UnchangedOr;
+using ffi::VisitInterrupt;
+
+/*!
+ * \brief Native visiting with exact dispatch and structural fallback.
+ *
+ * Allocate visitors with ffi::make_object<Derived>(). Inputs are borrowed;
+ * interrupts and errors own their values. A hook controls descent into its
node.
+ * One visitor must not be shared by overlapping traversals.
+ */
+class TVM_DLL ObjectVisitor : public ffi::StructuralVisitorObj {
+ public:
+ /*! \brief Construct a visitor using structural fallback for every value. */
+ ObjectVisitor() : ObjectVisitor(GlobalVTable()) {}
+ /*! \brief Release the managed visitor state. */
+ ~ObjectVisitor() = default;
+ ObjectVisitor(const ObjectVisitor& other) = delete;
+ ObjectVisitor& operator=(const ObjectVisitor& other) = delete;
+
+ /*!
+ * \brief Visit a borrowed object and propagate an interrupt or error.
+ * \param value The borrowed object to visit.
+ * \return None on completion, an owning VisitInterrupt, or an Error.
+ */
+ TVM_FFI_INLINE Expected<ffi::Optional<VisitInterrupt>> VisitExpected(
+ const ffi::ObjectRef& value) noexcept {
+ return VisitExpected(ffi::AnyView(value));
+ }
+ /*!
+ * \brief Visit a borrowed object or inline value.
+ * \param value The borrowed value to visit.
+ * \return None on completion, an owning VisitInterrupt, or an Error.
+ */
+ TVM_FFI_INLINE Expected<ffi::Optional<VisitInterrupt>> VisitExpected(
+ ffi::AnyView value) noexcept {
+ if (const auto* object = value.as<ffi::Object>()) return Dispatch(object);
+ return StructuralVisitDefault(value);
+ }
+
+ protected:
+ /*! \brief Exact native dispatch table with owning interrupt and error
results. */
+ using VTable =
+ ObjectFunctor<Expected<ffi::Optional<VisitInterrupt>>(const
ffi::Object*, ObjectVisitor*)>;
+
+ /*!
+ * \brief Construct a visitor with a finalized native dispatch table.
+ * \param vtable The immutable table, which must outlive this visitor.
+ */
+ explicit ObjectVisitor(const VTable* vtable)
+ : ffi::StructuralVisitorObj(StructuralVTable()), native_vtable_(vtable)
{}
+ /*!
+ * \brief Initialize the base registrations in a fresh mutable table.
+ * \param vtable The table to initialize before adding derived registrations.
+ */
+ static void InitVTable(VTable* vtable) {}
+
+ /*!
+ * \brief Register an exact node type in a fresh table initialized by its
parent.
+ * Derived node types require separate registrations.
+ * \tparam Self The visitor implementing the hook.
+ * \tparam Node The exact node type handled by the hook.
+ * \param vtable The mutable table to receive the registration.
+ */
+ template <typename Self, typename Node>
+ static void SetDispatch(VTable* vtable) {
+ vtable->template SetDispatch<Node>(DispatchNode<Self, Node>);
+ }
+
+ private:
+ // Native table callback; core instantiations may be shared by the library.
+ template <typename Self, typename Node>
+ static Expected<ffi::Optional<VisitInterrupt>> DispatchNode(const
ffi::Object* node,
+ ObjectVisitor*
self) {
+ return static_cast<Self*>(self)->Visit_(static_cast<const Node*>(node));
+ }
+
+ // The AnyView entry establishes that value is a non-null object.
+ TVM_FFI_INLINE Expected<ffi::Optional<VisitInterrupt>> Dispatch(
+ const ffi::Object* value) noexcept {
+ if (native_vtable_->CanDispatch(value)) {
+ try {
+ return DispatchNative(value);
+ } catch (ffi::Error& error) {
+ return AttachVisitErrorContext(error, value);
+ }
+ }
+ return StructuralVisitDefault(value);
+ }
+ // Keep one named return value in this scope so native results can be
constructed in place.
+ TVM_FFI_INLINE Expected<ffi::Optional<VisitInterrupt>> DispatchNative(const
ffi::Object* value) {
+ Expected<ffi::Optional<VisitInterrupt>> result = (*native_vtable_)(value,
this);
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ UpdateVisitErrorContext(result, value);
+ }
+ return result;
+ }
+ Expected<ffi::Optional<VisitInterrupt>> StructuralVisitDefault(ffi::AnyView
value) noexcept {
+ TVMFFIAny result = ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ ffi::StructuralVisitorObj::DefaultVisitExpected(value));
+ if (TVM_FFI_PREDICT_FALSE(result.type_index ==
ffi::TypeIndex::kTVMFFIError)) {
+ result = ffi::details::AttachStructuralVisitErrorContextRaw(result,
value);
+ }
+ return
ffi::details::ExpectedUnsafe::MoveFromTVMFFIAny<ffi::Optional<VisitInterrupt>>(result);
+ }
+ TVM_FFI_COLD_CODE static Expected<ffi::Optional<VisitInterrupt>>
AttachVisitErrorContext(
+ ffi::Error& error, const ffi::Object* value) {
+ if (value) ffi::details::UpdateVisitErrorContext(error,
ffi::GetRef<ffi::ObjectRef>(value));
+ return ffi::Unexpected(std::move(error));
+ }
+ TVM_FFI_COLD_CODE static void UpdateVisitErrorContext(
+ const Expected<ffi::Optional<VisitInterrupt>>& result, const
ffi::Object* value) {
+ ffi::Error error = result.error();
+ if (value) ffi::details::UpdateVisitErrorContext(error,
ffi::GetRef<ffi::ObjectRef>(value));
+ }
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+ static const ffi::StructuralVisitorVTable* StructuralVTable() {
+ static const ffi::StructuralVisitorVTable table{StructuralVTableVisitImpl};
+ return &table;
+ }
+ static TVMFFIAny StructuralVTableVisitImpl(ffi::StructuralVisitorObj* self,
+ ffi::AnyView value) noexcept {
+ return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ static_cast<ObjectVisitor*>(self)->VisitExpected(value));
+ }
+
+ const VTable* const native_vtable_;
+};
+
+/*!
+ * \brief Native mutation with exact dispatch and structural fallback.
+ *
+ * Allocate mutators with ffi::make_object<Derived>(). Inputs are borrowed;
+ * replacements and errors own their values. Overrides forward allow_inplace
+ * to child calls, following the structural mutation ownership contract.
+ */
+class TVM_DLL ObjectMutator : public ffi::StructuralMapEngineBase {
+ public:
+ /*! \brief Construct a mutator using structural fallback for every object. */
+ ObjectMutator() : ObjectMutator(GlobalVTable()) {}
+ /*! \brief Release the managed mutator state. */
+ ~ObjectMutator() = default;
+ ObjectMutator(const ObjectMutator& other) = delete;
+ ObjectMutator& operator=(const ObjectMutator& other) = delete;
+
+ /*!
+ * \brief Mutate a borrowed value without allowing in-place changes.
+ * \param value The borrowed object to mutate.
+ * \return A replacement, Unchanged, or an Error if mutation fails.
+ */
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>> MutateExpected(
+ const ffi::ObjectRef& value) noexcept {
+ return Dispatch(value.get(), false);
+ }
+ /*!
+ * \brief Mutate a borrowed value without allowing in-place changes.
+ * \param value The borrowed object or inline value to mutate.
+ * \return A replacement, Unchanged, or an Error if mutation fails.
+ */
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>> MutateExpected(ffi::AnyView
value) noexcept {
+ if (const auto* object = value.as<ffi::Object>()) return Dispatch(object,
false);
+ return StructuralMutateDefault(value, false);
+ }
+
+ /*!
+ * \brief Forward inherited permission and check the value's uniqueness.
+ * \param value The borrowed object to mutate.
+ * \param allow_inplace Whether the path to this value is already uniquely
owned.
+ * \return A replacement, Unchanged, or an Error if mutation fails.
+ */
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>>
MaybeInplaceMutateIfUniqueExpected(
+ const ffi::ObjectRef& value, bool allow_inplace = true) noexcept {
+ return Dispatch(value.get(), allow_inplace && value.defined() &&
value->unique());
+ }
+ /*!
+ * \brief Forward inherited permission and check the value's uniqueness.
+ * \param value The borrowed object or inline value to mutate.
+ * \param allow_inplace Whether the path to this value is already uniquely
owned.
+ * \return A replacement, Unchanged, or an Error if mutation fails.
+ */
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>>
MaybeInplaceMutateIfUniqueExpected(
+ ffi::AnyView value, bool allow_inplace = true) noexcept {
+ const auto* object = value.as<ffi::Object>();
+ if (object) return Dispatch(object, allow_inplace && object->unique());
+ return StructuralMutateDefault(value, false);
+ }
+
+ /*!
+ * \brief Look up a replacement in the variable-remap environment.
+ * \param var The borrowed variable identity to look up.
+ * \return The owning replacement, None on a miss, or an Error.
+ */
+ TVM_FFI_INLINE Expected<ffi::Any> VarRemapGetExpected(ffi::AnyView var)
noexcept {
+ return VarRemapGetImpl(var);
+ }
+ /*!
+ * \brief Record a replacement in the variable-remap environment.
+ * \param var The borrowed variable identity to bind.
+ * \param mapped_value The borrowed replacement value.
+ * \return Successful completion or an Error.
+ */
+ TVM_FFI_INLINE Expected<void> VarRemapSetExpected(ffi::AnyView var,
+ ffi::AnyView mapped_value)
noexcept {
+ return VarRemapSetImpl(var, mapped_value);
+ }
+
+ protected:
+ /*! \brief Exact native dispatch table with owning replacement and error
results. */
+ using VTable =
+ ObjectFunctor<Expected<UnchangedOr<ffi::Any>>(const ffi::Object*,
ObjectMutator*, bool)>;
+
+ /*!
+ * \brief Construct a mutator with a finalized native dispatch table.
+ * \param vtable The immutable table, which must outlive this mutator.
+ */
+ explicit ObjectMutator(const VTable* vtable)
+ : ffi::StructuralMapEngineBase(StructuralVTable()),
native_vtable_(vtable) {}
+ /*!
+ * \brief Initialize the base registrations in a fresh mutable table.
+ * \param vtable The table to initialize before adding derived registrations.
+ */
+ static void InitVTable(VTable* vtable) {}
+
+ /*!
+ * \brief Register an exact node type in a fresh table initialized by its
parent.
+ * Derived node types require separate registrations.
+ * \tparam Self The mutator implementing the hook.
+ * \tparam Node The exact node type handled by the hook.
+ * \param vtable The mutable table to receive the registration.
+ */
+ template <typename Self, typename Node>
+ static void SetDispatch(VTable* vtable) {
+ vtable->template SetDispatch<Node>(DispatchNode<Self, Node>);
+ }
+
+ private:
+ // Native table callback; core instantiations may be shared by the library.
+ template <typename Self, typename Node>
+ static Expected<UnchangedOr<ffi::Any>> DispatchNode(const ffi::Object* node,
ObjectMutator* self,
+ bool allow_inplace) {
+ return static_cast<Self*>(self)->Mutate_(static_cast<const Node*>(node),
allow_inplace);
+ }
+
+ using ffi::StructuralMutatorObj::DefaultMaybeInplaceMutateExpected;
+ using ffi::StructuralMutatorObj::DefaultMutateExpected;
+ using ffi::StructuralMutatorObj::MaybeInplaceMutate;
+ using ffi::StructuralMutatorObj::Mutate;
+
+ // Structural ABI entry: the caller guarantees ownership of the entire path.
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>> MaybeInplaceMutateExpected(
+ const ffi::ObjectRef& value) noexcept {
+ return Dispatch(value.get(), true);
+ }
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>> MaybeInplaceMutateExpected(
+ ffi::AnyView value) noexcept {
+ if (const auto* object = value.as<ffi::Object>()) return Dispatch(object,
true);
+ return StructuralMutateDefault(value, true);
+ }
+
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>> Dispatch(const ffi::Object*
value,
+ bool allow_inplace)
noexcept {
+ if (value == nullptr) return ffi::Unchanged();
+ if (native_vtable_->CanDispatch(value)) {
+ try {
+ return DispatchNative(value, allow_inplace);
+ } catch (ffi::Error& error) {
+ return AttachVisitErrorContext(error, value);
+ }
+ }
+ return StructuralMutateDefault(value, allow_inplace);
+ }
+ // Keep one named return value in this scope so native results can be
constructed in place.
+ TVM_FFI_INLINE Expected<UnchangedOr<ffi::Any>> DispatchNative(const
ffi::Object* value,
+ bool
allow_inplace) {
+ Expected<UnchangedOr<ffi::Any>> result = (*native_vtable_)(value, this,
allow_inplace);
+ if (TVM_FFI_PREDICT_FALSE(result.is_err())) {
+ UpdateVisitErrorContext(result, value);
+ }
+ return result;
+ }
+ Expected<UnchangedOr<ffi::Any>> StructuralMutateDefault(ffi::AnyView value,
+ bool allow_inplace)
noexcept {
+ if (allow_inplace) {
+ return
ffi::StructuralMutatorObj::DefaultMaybeInplaceMutateExpected(value);
+ } else {
+ return ffi::StructuralMutatorObj::DefaultMutateExpected(value);
+ }
+ }
+ TVM_FFI_COLD_CODE static Expected<UnchangedOr<ffi::Any>>
AttachVisitErrorContext(
+ ffi::Error& error, const ffi::Object* value) {
+ if (value) ffi::details::UpdateVisitErrorContext(error,
ffi::GetRef<ffi::ObjectRef>(value));
+ return ffi::Unexpected(std::move(error));
+ }
+ TVM_FFI_COLD_CODE static void UpdateVisitErrorContext(
+ const Expected<UnchangedOr<ffi::Any>>& result, const ffi::Object* value)
{
+ ffi::Error error = result.error();
+ if (value) ffi::details::UpdateVisitErrorContext(error,
ffi::GetRef<ffi::ObjectRef>(value));
+ }
+ static const VTable* GlobalVTable() {
+ static const VTable table = [] {
+ VTable table;
+ InitVTable(&table);
+ table.Finalize();
+ return table;
+ }();
+ return &table;
+ }
+ static const ffi::StructuralMutatorVTable* StructuralVTable() {
+ static const ffi::StructuralMutatorVTable table{
+ StructuralVTableMutateImpl, StructuralVTableMaybeInplaceMutateImpl,
+ StructuralVTableVarRemapGetImpl, StructuralVTableVarRemapSetImpl};
+ return &table;
+ }
+ static TVMFFIAny StructuralVTableMutateImpl(ffi::StructuralMutatorObj* self,
+ ffi::AnyView value) noexcept {
+ return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ static_cast<ObjectMutator*>(self)->MutateExpected(value));
+ }
+ static TVMFFIAny
StructuralVTableMaybeInplaceMutateImpl(ffi::StructuralMutatorObj* self,
+ ffi::AnyView value)
noexcept {
+ return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ static_cast<ObjectMutator*>(self)->MaybeInplaceMutateExpected(value));
+ }
+ static TVMFFIAny StructuralVTableVarRemapGetImpl(ffi::StructuralMutatorObj*
self,
+ ffi::AnyView key) noexcept {
+ return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ static_cast<ObjectMutator*>(self)->VarRemapGetImpl(key));
+ }
+ static TVMFFIAny StructuralVTableVarRemapSetImpl(ffi::StructuralMutatorObj*
self,
+ ffi::AnyView key,
ffi::AnyView value) noexcept {
+ return ffi::details::ExpectedUnsafe::MoveToTVMFFIAny(
+ static_cast<ObjectMutator*>(self)->VarRemapSetImpl(key, value));
+ }
+
+ const VTable* const native_vtable_;
+};
+
+/*!
+ * \brief Define a default constructor backed by one initialized native
dispatch table.
+ * \param Class The class whose default constructor is defined.
+ * \param Parent The parent accepting a const VTable pointer in its
constructor.
+ * \note Class must provide InitVTable(VTable*) including inherited
registrations.
+ * The table is initialized and finalized once; explicit table-taking
constructors are separate.
+ */
+#define TVM_DEFINE_OBJECT_FUNCTOR_DEFAULT_CONSTRUCTOR(Class, Parent) \
+ Class() \
+ : Parent([] { \
+ static const VTable table = [] { \
+ VTable table; \
+ Class::InitVTable(&table); \
+ table.Finalize(); \
+ return table; \
+ }(); \
+ return &table; \
+ }()) {}
+
} // namespace tvm
#endif // TVM_IR_OBJECT_FUNCTOR_H_
diff --git a/src/ir/expr_functor.cc b/src/ir/expr_functor.cc
new file mode 100644
index 0000000000..bbdd025639
--- /dev/null
+++ b/src/ir/expr_functor.cc
@@ -0,0 +1,722 @@
+/*
+ * 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 <tvm/ir/expr_functor.h>
+
+namespace tvm {
+
+void ExprVisitor::InitVTable(VTable* vtable) {
+ ObjectVisitor::InitVTable(vtable);
+ SetDispatch<ExprVisitor, OpaqueExprNode>(vtable);
+ SetDispatch<ExprVisitor, TupleNode>(vtable);
+ SetDispatch<ExprVisitor, TupleGetItemNode>(vtable);
+ SetDispatch<ExprVisitor, TensorLoadNode>(vtable);
+ SetDispatch<ExprVisitor, VarNode>(vtable);
+ SetDispatch<ExprVisitor, GlobalVarNode>(vtable);
+ SetDispatch<ExprVisitor, CallNode>(vtable);
+ SetDispatch<ExprVisitor, IntImmNode>(vtable);
+ SetDispatch<ExprVisitor, FloatImmNode>(vtable);
+ SetDispatch<ExprVisitor, OpNode>(vtable);
+ SetDispatch<ExprVisitor, prim::StringImmNode>(vtable);
+ SetDispatch<ExprVisitor, prim::CastNode>(vtable);
+ SetDispatch<ExprVisitor, prim::AddNode>(vtable);
+ SetDispatch<ExprVisitor, prim::SubNode>(vtable);
+ SetDispatch<ExprVisitor, prim::MulNode>(vtable);
+ SetDispatch<ExprVisitor, prim::DivNode>(vtable);
+ SetDispatch<ExprVisitor, prim::ModNode>(vtable);
+ SetDispatch<ExprVisitor, prim::FloorDivNode>(vtable);
+ SetDispatch<ExprVisitor, prim::FloorModNode>(vtable);
+ SetDispatch<ExprVisitor, prim::MinNode>(vtable);
+ SetDispatch<ExprVisitor, prim::MaxNode>(vtable);
+ SetDispatch<ExprVisitor, prim::EQNode>(vtable);
+ SetDispatch<ExprVisitor, prim::NENode>(vtable);
+ SetDispatch<ExprVisitor, prim::LTNode>(vtable);
+ SetDispatch<ExprVisitor, prim::LENode>(vtable);
+ SetDispatch<ExprVisitor, prim::GTNode>(vtable);
+ SetDispatch<ExprVisitor, prim::GENode>(vtable);
+ SetDispatch<ExprVisitor, prim::AndNode>(vtable);
+ SetDispatch<ExprVisitor, prim::OrNode>(vtable);
+ SetDispatch<ExprVisitor, prim::NotNode>(vtable);
+ SetDispatch<ExprVisitor, prim::SelectNode>(vtable);
+ SetDispatch<ExprVisitor, prim::LetNode>(vtable);
+ SetDispatch<ExprVisitor, prim::RampNode>(vtable);
+ SetDispatch<ExprVisitor, prim::BroadcastNode>(vtable);
+ SetDispatch<ExprVisitor, prim::ShuffleNode>(vtable);
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
OpaqueExprNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const TupleNode*
node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->fields));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
TupleGetItemNode* node) {
+ // skips: index
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->tuple));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
TensorLoadNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->source));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->indices));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const VarNode*
node) {
+ // Primitive types have no children; dynamic type fields are visited.
+ if (!node->ty.as<PrimTypeNode>()) {
+ // Clamp Simple for dynamic type fields; Pattern continues through them.
+ if (this->def_region_kind() == kTVMFFIDefRegionKindSimple) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind(
+ kTVMFFIDefRegionKindNone, [&]() { return
this->VisitExpected(node->ty); }));
+ } else {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty));
+ }
+ }
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
GlobalVarNode* node) {
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const CallNode*
node) {
+ // Skip constant attrs and primitive result types.
+ if (!node->ty.as<PrimTypeNode>()) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty));
+ }
+ // Interned operators have no children; function-valued operators are
visited.
+ if (!node->op.as<OpNode>()) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->op));
+ }
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->args));
+ // Empty type arguments are skipped, including the container callback.
+ if (!node->ty_args.empty()) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->ty_args));
+ }
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const IntImmNode*
node) {
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
FloatImmNode* node) {
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const OpNode*
node) {
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::StringImmNode* node) {
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::CastNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->value));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::AddNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::SubNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::MulNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::DivNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::ModNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::FloorDivNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::FloorModNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::MinNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::MaxNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::EQNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::NENode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::LTNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::LENode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::GTNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::GENode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::AndNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::OrNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->b));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::NotNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->a));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::SelectNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->condition));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->true_value));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->false_value));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::LetNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->WithDefRegionKind(
+ kTVMFFIDefRegionKindSimple, [&]() { return
this->VisitExpected(node->var); }));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->value));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->body));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::RampNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->base));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->stride));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->lanes));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::BroadcastNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->value));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->lanes));
+ return std::nullopt;
+}
+
+Expected<ffi::Optional<VisitInterrupt>> ExprVisitor::Visit_(const
prim::ShuffleNode* node) {
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->vectors));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(this->VisitExpected(node->indices));
+ return std::nullopt;
+}
+
+void ExprMutator::InitVTable(VTable* vtable) {
+ ObjectMutator::InitVTable(vtable);
+ SetDispatch<ExprMutator, OpaqueExprNode>(vtable);
+ SetDispatch<ExprMutator, TupleNode>(vtable);
+ SetDispatch<ExprMutator, TupleGetItemNode>(vtable);
+ SetDispatch<ExprMutator, TensorLoadNode>(vtable);
+ SetDispatch<ExprMutator, VarNode>(vtable);
+ SetDispatch<ExprMutator, GlobalVarNode>(vtable);
+ SetDispatch<ExprMutator, CallNode>(vtable);
+ SetDispatch<ExprMutator, IntImmNode>(vtable);
+ SetDispatch<ExprMutator, FloatImmNode>(vtable);
+ SetDispatch<ExprMutator, OpNode>(vtable);
+ SetDispatch<ExprMutator, prim::StringImmNode>(vtable);
+ SetDispatch<ExprMutator, prim::CastNode>(vtable);
+ SetDispatch<ExprMutator, prim::AddNode>(vtable);
+ SetDispatch<ExprMutator, prim::SubNode>(vtable);
+ SetDispatch<ExprMutator, prim::MulNode>(vtable);
+ SetDispatch<ExprMutator, prim::DivNode>(vtable);
+ SetDispatch<ExprMutator, prim::ModNode>(vtable);
+ SetDispatch<ExprMutator, prim::FloorDivNode>(vtable);
+ SetDispatch<ExprMutator, prim::FloorModNode>(vtable);
+ SetDispatch<ExprMutator, prim::MinNode>(vtable);
+ SetDispatch<ExprMutator, prim::MaxNode>(vtable);
+ SetDispatch<ExprMutator, prim::EQNode>(vtable);
+ SetDispatch<ExprMutator, prim::NENode>(vtable);
+ SetDispatch<ExprMutator, prim::LTNode>(vtable);
+ SetDispatch<ExprMutator, prim::LENode>(vtable);
+ SetDispatch<ExprMutator, prim::GTNode>(vtable);
+ SetDispatch<ExprMutator, prim::GENode>(vtable);
+ SetDispatch<ExprMutator, prim::AndNode>(vtable);
+ SetDispatch<ExprMutator, prim::OrNode>(vtable);
+ SetDispatch<ExprMutator, prim::NotNode>(vtable);
+ SetDispatch<ExprMutator, prim::SelectNode>(vtable);
+ SetDispatch<ExprMutator, prim::LetNode>(vtable);
+ SetDispatch<ExprMutator, prim::RampNode>(vtable);
+ SetDispatch<ExprMutator, prim::BroadcastNode>(vtable);
+ SetDispatch<ExprMutator, prim::ShuffleNode>(vtable);
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const OpaqueExprNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Type>, ty,
this->MaybeInplaceMutateIfUniqueExpected(node->ty, allow_inplace));
+ if (ty.UnchangedOrSameAs(node->ty)) return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<OpaqueExprNode*>(node);
+ if (!ty.IsUnchanged()) writable->ty = std::move(ty).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<OpaqueExprNode>(*node);
+ copy->ty = std::move(ty).ValueOrUnchanged(std::move(copy->ty));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const TupleNode* node,
bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Type>, ty,
this->MaybeInplaceMutateIfUniqueExpected(node->ty, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ffi::Array<Expr>>, fields,
+ this->MaybeInplaceMutateIfUniqueExpected(node->fields, allow_inplace));
+ if (ty.UnchangedOrSameAs(node->ty) && fields.UnchangedOrSameAs(node->fields))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<TupleNode*>(node);
+ if (!ty.IsUnchanged()) writable->ty = std::move(ty).ValueUnchecked();
+ if (!fields.IsUnchanged()) writable->fields =
std::move(fields).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<TupleNode>(*node);
+ copy->ty = std::move(ty).ValueOrUnchanged(std::move(copy->ty));
+ copy->fields = std::move(fields).ValueOrUnchanged(std::move(copy->fields));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const TupleGetItemNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Type>, ty,
this->MaybeInplaceMutateIfUniqueExpected(node->ty, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Expr>, tuple,
+
MaybeInplaceMutateIfUniqueExpected(node->tuple, allow_inplace));
+ if (ty.UnchangedOrSameAs(node->ty) && tuple.UnchangedOrSameAs(node->tuple))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<TupleGetItemNode*>(node);
+ if (!ty.IsUnchanged()) writable->ty = std::move(ty).ValueUnchecked();
+ if (!tuple.IsUnchanged()) writable->tuple =
std::move(tuple).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<TupleGetItemNode>(*node);
+ copy->ty = std::move(ty).ValueOrUnchanged(std::move(copy->ty));
+ copy->tuple = std::move(tuple).ValueOrUnchanged(std::move(copy->tuple));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const TensorLoadNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Type>, ty,
this->MaybeInplaceMutateIfUniqueExpected(node->ty, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Expr>, source,
MaybeInplaceMutateIfUniqueExpected(node->source, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ffi::Array<PrimExpr>>, indices,
+ this->MaybeInplaceMutateIfUniqueExpected(node->indices, allow_inplace));
+ if (ty.UnchangedOrSameAs(node->ty) && source.UnchangedOrSameAs(node->source)
&&
+ indices.UnchangedOrSameAs(node->indices))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<TensorLoadNode*>(node);
+ if (!ty.IsUnchanged()) writable->ty = std::move(ty).ValueUnchecked();
+ if (!source.IsUnchanged()) writable->source =
std::move(source).ValueUnchecked();
+ if (!indices.IsUnchanged()) writable->indices =
std::move(indices).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<TensorLoadNode>(*node);
+ copy->ty = std::move(ty).ValueOrUnchanged(std::move(copy->ty));
+ copy->source = std::move(source).ValueOrUnchanged(std::move(copy->source));
+ copy->indices =
std::move(indices).ValueOrUnchanged(std::move(copy->indices));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const GlobalVarNode* node,
+ bool allow_inplace) {
+ // Registry atoms and constant leaves do not descend into metadata or types.
+ return ffi::Unchanged();
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const CallNode* node,
bool allow_inplace) {
+ UnchangedOr<Type> ty = ffi::Unchanged();
+ if (!node->ty.as<PrimTypeNode>()) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Type>, mapped_ty,
+ this->MaybeInplaceMutateIfUniqueExpected(node->ty, allow_inplace));
+ ty = std::move(mapped_ty);
+ }
+ UnchangedOr<Expr> op = ffi::Unchanged();
+ if (!node->op.as<OpNode>()) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Expr>, mapped_op,
+
MaybeInplaceMutateIfUniqueExpected(node->op, allow_inplace));
+ op = std::move(mapped_op);
+ }
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ffi::Array<Expr>>, args,
+ this->MaybeInplaceMutateIfUniqueExpected(node->args, allow_inplace));
+ UnchangedOr<ffi::Array<Type>> ty_args = ffi::Unchanged();
+ if (!node->ty_args.empty()) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ffi::Array<Type>>, mapped_ty_args,
+ this->MaybeInplaceMutateIfUniqueExpected(node->ty_args,
allow_inplace));
+ ty_args = std::move(mapped_ty_args);
+ }
+ if (ty.UnchangedOrSameAs(node->ty) && op.UnchangedOrSameAs(node->op) &&
+ args.UnchangedOrSameAs(node->args) &&
ty_args.UnchangedOrSameAs(node->ty_args))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<CallNode*>(node);
+ if (!ty.IsUnchanged()) writable->ty = std::move(ty).ValueUnchecked();
+ if (!op.IsUnchanged()) writable->op = std::move(op).ValueUnchecked();
+ if (!args.IsUnchanged()) writable->args = std::move(args).ValueUnchecked();
+ if (!ty_args.IsUnchanged()) writable->ty_args =
std::move(ty_args).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<CallNode>(*node);
+ copy->ty = std::move(ty).ValueOrUnchanged(std::move(copy->ty));
+ copy->op = std::move(op).ValueOrUnchanged(std::move(copy->op));
+ copy->args = std::move(args).ValueOrUnchanged(std::move(copy->args));
+ copy->ty_args =
std::move(ty_args).ValueOrUnchanged(std::move(copy->ty_args));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const IntImmNode* node,
bool allow_inplace) {
+ // Registry atoms and constant leaves do not descend into metadata or types.
+ return ffi::Unchanged();
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const FloatImmNode* node,
bool allow_inplace) {
+ // Registry atoms and constant leaves do not descend into metadata or types.
+ return ffi::Unchanged();
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const OpNode* node, bool
allow_inplace) {
+ // Registry atoms and constant leaves do not descend into metadata or types.
+ return ffi::Unchanged();
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const
prim::StringImmNode* node,
+ bool allow_inplace) {
+ // Registry atoms and constant leaves do not descend into metadata or types.
+ return ffi::Unchanged();
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::CastNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, value,
+
MaybeInplaceMutateIfUniqueExpected(node->value, allow_inplace));
+ if (value.UnchangedOrSameAs(node->value)) return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::CastNode*>(node);
+ if (!value.IsUnchanged()) writable->value =
std::move(value).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::CastNode>(*node);
+ copy->value = std::move(value).ValueOrUnchanged(std::move(copy->value));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+#define TVM_IR_BINARY_MUTATE_IMPL(Name)
\
+ Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::Name##Node*
node, \
+ bool allow_inplace) {
\
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, a,
\
+
MaybeInplaceMutateIfUniqueExpected(node->a, allow_inplace)); \
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, b,
\
+
MaybeInplaceMutateIfUniqueExpected(node->b, allow_inplace)); \
+ if (a.UnchangedOrSameAs(node->a) && b.UnchangedOrSameAs(node->b)) return
ffi::Unchanged(); \
+ if (allow_inplace) {
\
+ auto* writable = const_cast<prim::Name##Node*>(node);
\
+ if (!a.IsUnchanged()) writable->a = std::move(a).ValueUnchecked();
\
+ if (!b.IsUnchanged()) writable->b = std::move(b).ValueUnchecked();
\
+ return ffi::Unchanged();
\
+ } else {
\
+ auto copy = ffi::make_object<prim::Name##Node>(*node);
\
+ copy->a = std::move(a).ValueOrUnchanged(std::move(copy->a));
\
+ copy->b = std::move(b).ValueOrUnchanged(std::move(copy->b));
\
+ return ffi::Any(Expr(std::move(copy)));
\
+ }
\
+ }
+TVM_IR_BINARY_MUTATE_IMPL(Add)
+TVM_IR_BINARY_MUTATE_IMPL(Sub)
+TVM_IR_BINARY_MUTATE_IMPL(Mul)
+TVM_IR_BINARY_MUTATE_IMPL(Div)
+TVM_IR_BINARY_MUTATE_IMPL(Mod)
+TVM_IR_BINARY_MUTATE_IMPL(FloorDiv)
+TVM_IR_BINARY_MUTATE_IMPL(FloorMod)
+TVM_IR_BINARY_MUTATE_IMPL(Min)
+TVM_IR_BINARY_MUTATE_IMPL(Max)
+TVM_IR_BINARY_MUTATE_IMPL(EQ)
+TVM_IR_BINARY_MUTATE_IMPL(NE)
+TVM_IR_BINARY_MUTATE_IMPL(LT)
+TVM_IR_BINARY_MUTATE_IMPL(LE)
+TVM_IR_BINARY_MUTATE_IMPL(GT)
+TVM_IR_BINARY_MUTATE_IMPL(GE)
+TVM_IR_BINARY_MUTATE_IMPL(And)
+TVM_IR_BINARY_MUTATE_IMPL(Or)
+#undef TVM_IR_BINARY_MUTATE_IMPL
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::NotNode* node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, a,
+
MaybeInplaceMutateIfUniqueExpected(node->a, allow_inplace));
+ if (a.UnchangedOrSameAs(node->a)) return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::NotNode*>(node);
+ if (!a.IsUnchanged()) writable->a = std::move(a).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::NotNode>(*node);
+ copy->a = std::move(a).ValueOrUnchanged(std::move(copy->a));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::SelectNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<PrimExpr>, condition,
+ MaybeInplaceMutateIfUniqueExpected(node->condition, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<PrimExpr>, true_value,
+ MaybeInplaceMutateIfUniqueExpected(node->true_value, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<PrimExpr>, false_value,
+ MaybeInplaceMutateIfUniqueExpected(node->false_value, allow_inplace));
+ if (condition.UnchangedOrSameAs(node->condition) &&
+ true_value.UnchangedOrSameAs(node->true_value) &&
+ false_value.UnchangedOrSameAs(node->false_value))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::SelectNode*>(node);
+ if (!condition.IsUnchanged()) writable->condition =
std::move(condition).ValueUnchecked();
+ if (!true_value.IsUnchanged()) writable->true_value =
std::move(true_value).ValueUnchecked();
+ if (!false_value.IsUnchanged()) writable->false_value =
std::move(false_value).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::SelectNode>(*node);
+ copy->condition =
std::move(condition).ValueOrUnchanged(std::move(copy->condition));
+ copy->true_value =
std::move(true_value).ValueOrUnchanged(std::move(copy->true_value));
+ copy->false_value =
std::move(false_value).ValueOrUnchanged(std::move(copy->false_value));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::LetNode* node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<Var>, var, WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]
{
+ return MaybeInplaceMutateIfUniqueExpected(node->var, allow_inplace);
+ }));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, value,
+
MaybeInplaceMutateIfUniqueExpected(node->value, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, body,
+
MaybeInplaceMutateIfUniqueExpected(node->body, allow_inplace));
+ if (var.UnchangedOrSameAs(node->var) && value.UnchangedOrSameAs(node->value)
&&
+ body.UnchangedOrSameAs(node->body))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::LetNode*>(node);
+ if (!var.IsUnchanged()) writable->var = std::move(var).ValueUnchecked();
+ if (!value.IsUnchanged()) writable->value =
std::move(value).ValueUnchecked();
+ if (!body.IsUnchanged()) writable->body = std::move(body).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::LetNode>(*node);
+ copy->var = std::move(var).ValueOrUnchanged(std::move(copy->var));
+ copy->value = std::move(value).ValueOrUnchanged(std::move(copy->value));
+ copy->body = std::move(body).ValueOrUnchanged(std::move(copy->body));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::RampNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, base,
+
MaybeInplaceMutateIfUniqueExpected(node->base, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<PrimExpr>, stride,
+ MaybeInplaceMutateIfUniqueExpected(node->stride, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, lanes,
+
MaybeInplaceMutateIfUniqueExpected(node->lanes, allow_inplace));
+ if (base.UnchangedOrSameAs(node->base) &&
stride.UnchangedOrSameAs(node->stride) &&
+ lanes.UnchangedOrSameAs(node->lanes))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::RampNode*>(node);
+ if (!base.IsUnchanged()) writable->base = std::move(base).ValueUnchecked();
+ if (!stride.IsUnchanged()) writable->stride =
std::move(stride).ValueUnchecked();
+ if (!lanes.IsUnchanged()) writable->lanes =
std::move(lanes).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::RampNode>(*node);
+ copy->base = std::move(base).ValueOrUnchanged(std::move(copy->base));
+ copy->stride = std::move(stride).ValueOrUnchanged(std::move(copy->stride));
+ copy->lanes = std::move(lanes).ValueOrUnchanged(std::move(copy->lanes));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const
prim::BroadcastNode* node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, value,
+
MaybeInplaceMutateIfUniqueExpected(node->value, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<PrimExpr>, lanes,
+
MaybeInplaceMutateIfUniqueExpected(node->lanes, allow_inplace));
+ if (value.UnchangedOrSameAs(node->value) &&
lanes.UnchangedOrSameAs(node->lanes))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::BroadcastNode*>(node);
+ if (!value.IsUnchanged()) writable->value =
std::move(value).ValueUnchecked();
+ if (!lanes.IsUnchanged()) writable->lanes =
std::move(lanes).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::BroadcastNode>(*node);
+ copy->value = std::move(value).ValueOrUnchanged(std::move(copy->value));
+ copy->lanes = std::move(lanes).ValueOrUnchanged(std::move(copy->lanes));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const prim::ShuffleNode*
node,
+ bool allow_inplace) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ffi::Array<PrimExpr>>, vectors,
+ this->MaybeInplaceMutateIfUniqueExpected(node->vectors, allow_inplace));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
+ UnchangedOr<ffi::Array<PrimExpr>>, indices,
+ this->MaybeInplaceMutateIfUniqueExpected(node->indices, allow_inplace));
+ if (vectors.UnchangedOrSameAs(node->vectors) &&
indices.UnchangedOrSameAs(node->indices))
+ return ffi::Unchanged();
+ if (allow_inplace) {
+ auto* writable = const_cast<prim::ShuffleNode*>(node);
+ if (!vectors.IsUnchanged()) writable->vectors =
std::move(vectors).ValueUnchecked();
+ if (!indices.IsUnchanged()) writable->indices =
std::move(indices).ValueUnchecked();
+ return ffi::Unchanged();
+ } else {
+ auto copy = ffi::make_object<prim::ShuffleNode>(*node);
+ copy->vectors =
std::move(vectors).ValueOrUnchanged(std::move(copy->vectors));
+ copy->indices =
std::move(indices).ValueOrUnchanged(std::move(copy->indices));
+ return ffi::Any(Expr(std::move(copy)));
+ }
+}
+
+Expected<UnchangedOr<ffi::Any>> ExprMutator::Mutate_(const VarNode* node, bool
allow_inplace) {
+ if (node->ty.as<PrimTypeNode>()) return ffi::Unchanged();
+ if (TVM_FFI_PREDICT_TRUE(var_remap_.empty() && def_region_kind() ==
kTVMFFIDefRegionKindNone)) {
+ return ffi::Unchanged();
+ }
+ Expected<ffi::Any> remap_result = VarRemapGetExpected(ffi::AnyView(node));
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap_result);
+ if (ffi::details::ExpectedUnsafe::GetData(remap_result).type_index() !=
+ ffi::TypeIndex::kTVMFFINone) {
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Expr>, mapped,
std::move(remap_result));
+ return mapped;
+ }
+ if (def_region_kind() == kTVMFFIDefRegionKindNone) return ffi::Unchanged();
+ UnchangedOr<ffi::Any> result = ffi::Unchanged();
+ ffi::Any mapped_value = ffi::Unchanged();
+ // PrimType has no children; dynamic type fields inherit Pattern but are
visited outside Simple.
+ if (!node->ty.as<PrimTypeNode>()) {
+ Expected<UnchangedOr<ffi::Any>> mapped_ty_result =
+ def_region_kind() == kTVMFFIDefRegionKindSimple
+ ? WithDefRegionKind(
+ kTVMFFIDefRegionKindNone,
+ [&] { return
this->MaybeInplaceMutateIfUniqueExpected(node->ty, allow_inplace); })
+ : this->MaybeInplaceMutateIfUniqueExpected(node->ty,
allow_inplace);
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(UnchangedOr<Type>, mapped_ty,
std::move(mapped_ty_result));
+ if (!mapped_ty.UnchangedOrSameAs(node->ty)) {
+ if (allow_inplace) {
+ const_cast<VarNode*>(node)->ty = std::move(mapped_ty).ValueUnchecked();
+ mapped_value = ffi::Any(node);
+ } else {
+ auto copy = ffi::make_object<VarNode>(*node);
+ copy->ty = std::move(mapped_ty).ValueUnchecked();
+ mapped_value = ffi::Any(std::move(copy));
+ }
+ result = mapped_value;
+ }
+ }
+ if (!result.IsUnchanged() || def_region_kind() ==
kTVMFFIDefRegionKindPattern) {
+ auto set_result = VarRemapSetExpected(ffi::AnyView(node), mapped_value);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(set_result);
+ }
+ return result;
+}
+
+} // namespace tvm
diff --git a/tests/cpp/expr_functor_test.cc b/tests/cpp/expr_functor_test.cc
new file mode 100644
index 0000000000..06cb2c962e
--- /dev/null
+++ b/tests/cpp/expr_functor_test.cc
@@ -0,0 +1,95 @@
+/*
+ * 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/ffi/reflection/registry.h>
+#include <tvm/ir/expr_functor.h>
+
+#include <vector>
+
+namespace tvm {
+namespace {
+
+class PairExprNode : public ExprNode {
+ public:
+ Expr left;
+ Expr right;
+ static void RegisterReflection() {
+ ffi::reflection::ObjectDef<PairExprNode>()
+ .def_ro("left", &PairExprNode::left)
+ .def_ro("right", &PairExprNode::right);
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("test.NativePairExpr", PairExprNode,
ExprNode);
+};
+class PairExpr : public Expr {
+ public:
+ PairExpr(Expr left, Expr right) {
+ auto node = ffi::make_object<PairExprNode>();
+ node->left = std::move(left);
+ node->right = std::move(right);
+ data_ = std::move(node);
+ }
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PairExpr, Expr, PairExprNode);
+};
+TVM_FFI_STATIC_INIT_BLOCK() { PairExprNode::RegisterReflection(); }
+
+class Collect : public ExprVisitor {
+ public:
+ using ExprVisitor::Visit_;
+ std::vector<int64_t> values;
+ Expected<ffi::Optional<VisitInterrupt>> Visit_(const IntImmNode* node)
override {
+ values.push_back(node->value);
+ return std::nullopt;
+ }
+};
+
+TEST(ExprVisitor, StructuralFallback) {
+ auto visitor = ffi::make_object<Collect>();
+ PairExpr pair(IntImm::Int32(1), IntImm::Int32(2));
+ auto result = visitor->VisitExpected(pair);
+ ASSERT_TRUE(result.is_ok());
+ EXPECT_FALSE(result.value().has_value());
+ EXPECT_EQ(visitor->values, (std::vector<int64_t>{1, 2}));
+ EXPECT_EQ(pair->left.as<IntImmNode>()->value, 1);
+ EXPECT_EQ(pair->right.as<IntImmNode>()->value, 2);
+}
+
+class Rewrite : public ExprMutator {
+ public:
+ using ExprMutator::Mutate_;
+ Expected<UnchangedOr<ffi::Any>> Mutate_(const IntImmNode* node, bool
allow_inplace) override {
+ return ffi::Any(IntImm::Int32(node->value + 1));
+ }
+};
+
+TEST(ExprMutator, StructuralFallback) {
+ auto mutator = ffi::make_object<Rewrite>();
+ PairExpr pair(IntImm::Int32(1), IntImm::Int32(2));
+ auto result = mutator->MutateExpected(pair).value();
+ ASSERT_FALSE(result.IsUnchanged());
+ auto replacement = std::move(result).ValueUnchecked();
+ const auto* changed = replacement.as<PairExprNode>();
+ ASSERT_NE(changed, nullptr);
+ EXPECT_EQ(changed->left.as<IntImmNode>()->value, 2);
+ EXPECT_EQ(changed->right.as<IntImmNode>()->value, 3);
+ EXPECT_EQ(pair->left.as<IntImmNode>()->value, 1);
+ EXPECT_EQ(pair->right.as<IntImmNode>()->value, 2);
+}
+
+} // namespace
+} // namespace tvm
