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 a771de82bd [REFACTOR][TIRx] Promote LambdaExpr to Expr (#20546)
a771de82bd is described below
commit a771de82bdc9517ca4596727dcf0c4c15cacab31
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Oct 4 21:29:12 2026 -0400
[REFACTOR][TIRx] Promote LambdaExpr to Expr (#20546)
Promote `LambdaExpr` to `Expr` in C++ and Python, retaining the
inherited `MissingType` and existing construction and application
behavior.
Preserve lexical parameter scope during structural traversal and
mutation, and recognize lambda expressions in TIR path traversal.
---
include/tvm/tirx/tile_primitive.h | 8 ++--
python/tvm/tirx/tile_primitive.py | 2 +-
src/tirx/ir/lambdaexpr.cc | 83 ++++++++++++++++++++++++++++++++++--
src/tirx/ir/tir_visitor_with_path.cc | 11 +++++
src/tirx/ir/tir_visitor_with_path.h | 4 ++
5 files changed, 99 insertions(+), 9 deletions(-)
diff --git a/include/tvm/tirx/tile_primitive.h
b/include/tvm/tirx/tile_primitive.h
index 879c9b5e72..c41c38c2b4 100644
--- a/include/tvm/tirx/tile_primitive.h
+++ b/include/tvm/tirx/tile_primitive.h
@@ -42,7 +42,7 @@ namespace tirx {
* At lowering time the dispatch substitutes ``vars`` with the concrete
* instruction axes via ``Apply``.
*/
-class LambdaExprNode : public ffi::Object {
+class LambdaExprNode : public ExprNode {
public:
explicit LambdaExprNode(ffi::UnsafeInit tag) : pred(tag) {}
@@ -64,18 +64,18 @@ class LambdaExprNode : public ffi::Object {
}
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind =
kTVMFFISEqHashKindTreeNode;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.LambdaExpr", LambdaExprNode,
ffi::Object);
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.LambdaExpr", LambdaExprNode,
ExprNode);
};
/*!
* \brief Managed reference to LambdaExprNode.
* \sa LambdaExprNode
*/
-class LambdaExpr : public ffi::ObjectRef {
+class LambdaExpr : public Expr {
public:
explicit LambdaExpr(Array<Var> vars, PrimExpr pred);
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(LambdaExpr, ffi::ObjectRef,
LambdaExprNode);
+ TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(LambdaExpr, Expr,
LambdaExprNode);
};
/*!
diff --git a/python/tvm/tirx/tile_primitive.py
b/python/tvm/tirx/tile_primitive.py
index 3118130d01..8adbd3abe5 100644
--- a/python/tvm/tirx/tile_primitive.py
+++ b/python/tvm/tirx/tile_primitive.py
@@ -39,7 +39,7 @@ from .stmt import Stmt
@register_object("tirx.LambdaExpr")
-class LambdaExpr(Object):
+class LambdaExpr(Expr):
"""A reified Python lambda: bound variables and a body over them.
Used by tile primitive ops that take a per-element expression over the
diff --git a/src/tirx/ir/lambdaexpr.cc b/src/tirx/ir/lambdaexpr.cc
index 789bf9dc0e..b73e6caab1 100644
--- a/src/tirx/ir/lambdaexpr.cc
+++ b/src/tirx/ir/lambdaexpr.cc
@@ -22,13 +22,88 @@
* \brief Implementation of LambdaExpr, a reified lambda used by tile
primitive ops.
*/
-#include "tvm/ffi/extra/structural_mutate.h"
-#include "tvm/tirx/tile_primitive.h"
+#include <tvm/ffi/extra/structural_mutate.h>
+#include <tvm/ffi/extra/structural_visit.h>
+#include <tvm/tirx/tile_primitive.h>
+
+#include <utility>
+#include <vector>
namespace tvm {
namespace tirx {
-TVM_FFI_STATIC_INIT_BLOCK() { LambdaExprNode::RegisterReflection(); }
+namespace {
+
+ffi::Expected<ffi::Optional<ffi::VisitInterrupt>> LambdaExprVisit(
+ ffi::StructuralVisitorObj* visitor, ffi::AnyView value) noexcept {
+ const auto* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
LambdaExprNode>(value);
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->VisitExpected(self->ty));
+ TVM_FFI_S_VISIT_MAYBE_EARLY_RETURN(visitor->WithDefRegionKind(
+ kTVMFFIDefRegionKindSimple, [&]() { return
visitor->VisitExpected(self->vars); }));
+ return visitor->VisitExpected(self->pred);
+}
+
+ffi::Expected<ffi::UnchangedOr<ffi::Any>>
LambdaExprMutate(ffi::StructuralMutatorObj* mutator,
+ ffi::AnyView value)
noexcept {
+ const auto* self =
+ ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
LambdaExprNode>(value);
+ // Parameters shadow substitutions from the surrounding expression. Save and
restore
+ // their remaps so a binder rewrite stays local to this lambda's body.
+ std::vector<ffi::Any> saved;
+ for (const Var& var : self->vars) {
+ auto previous = mutator->VarRemapGetExpected(var);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(previous);
+ saved.push_back(std::move(previous).value());
+ }
+ auto result = [&]() -> ffi::Expected<ffi::UnchangedOr<ffi::Any>> {
+ for (const Var& var : self->vars) {
+ auto cleared = mutator->VarRemapSetExpected(var, nullptr);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(cleared);
+ }
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ty,
+ mutator->MutateExpected(self->ty));
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Array<Var>>,
mapped_vars,
+
mutator->WithDefRegionKind(kTVMFFIDefRegionKindSimple, [&]() {
+ return
mutator->MutateExpected(self->vars);
+ }));
+ auto vars = std::move(mapped_vars).ValueOrUnchanged(self->vars);
+ for (size_t i = 0; i < vars.size(); ++i) {
+ auto remap = mutator->VarRemapSetExpected(self->vars[i], vars[i]);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(remap);
+ }
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<PrimExpr>, mapped_pred,
+ mutator->MutateExpected(self->pred));
+ if (vars.same_as(self->vars) && mapped_pred.UnchangedOrSameAs(self->pred)
&&
+ mapped_ty.UnchangedOrSameAs(self->ty)) {
+ return ffi::Unchanged();
+ }
+ auto copy = ffi::make_object<LambdaExprNode>(*self);
+ copy->vars = std::move(vars);
+ copy->pred = std::move(mapped_pred).ValueOrUnchanged(self->pred);
+ copy->ty = std::move(mapped_ty).ValueOrUnchanged(self->ty);
+ return ffi::Any(std::move(copy));
+ }();
+ for (size_t i = 0; i < self->vars.size(); ++i) {
+ auto restored = mutator->VarRemapSetExpected(self->vars[i], saved[i]);
+ TVM_FFI_S_MUTATE_MAYBE_EARLY_RETURN(restored);
+ }
+ return result;
+}
+
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ LambdaExprNode::RegisterReflection();
+ refl::TypeAttrDef<LambdaExprNode>()
+ .attr(refl::type_attr::kStructuralVisit,
+ ffi::FStructuralVisit::FromNative<&LambdaExprVisit>())
+ .attr(refl::type_attr::kStructuralMutate,
+ ffi::FStructuralMutate::FromNative<&LambdaExprMutate>())
+ .attr(refl::type_attr::kStructuralMaybeInplaceMutate,
+ ffi::FStructuralMutate::FromNative<&LambdaExprMutate>());
+}
PrimExpr LambdaExprNode::Apply(const ffi::Array<PrimExpr>& indices) const {
TVM_FFI_ICHECK_EQ(indices.size(), vars.size());
@@ -46,7 +121,7 @@ PrimExpr LambdaExprNode::Apply(const ffi::Array<PrimExpr>&
indices) const {
.as_or_throw<PrimExpr>();
}
-LambdaExpr::LambdaExpr(ffi::Array<Var> vars, PrimExpr pred) {
+LambdaExpr::LambdaExpr(ffi::Array<Var> vars, PrimExpr pred) :
Expr(ffi::UnsafeInit{}) {
auto n = ffi::make_object<LambdaExprNode>(std::move(pred));
n->vars = std::move(vars);
data_ = std::move(n);
diff --git a/src/tirx/ir/tir_visitor_with_path.cc
b/src/tirx/ir/tir_visitor_with_path.cc
index a1844b0536..c0402f184f 100644
--- a/src/tirx/ir/tir_visitor_with_path.cc
+++ b/src/tirx/ir/tir_visitor_with_path.cc
@@ -256,6 +256,8 @@ void TIRVisitorWithPath::Dispatch_(const
tirx::TilePrimitiveCallNode* op, Access
}
if (auto buf_region = op->args[i].as<TensorRegion>()) {
Visit(buf_region.value(), path->Attr("args")->ArrayItem(i));
+ } else if (auto lambda = op->args[i].as<LambdaExpr>()) {
+ Visit(lambda.value(), path->Attr("args")->ArrayItem(i));
} else if (auto expr = op->args[i].as<PrimExpr>()) {
Visit(expr.value(), path->Attr("args")->ArrayItem(i));
} else if (auto stmt = op->args[i].as<Stmt>()) {
@@ -297,6 +299,15 @@ void TIRVisitorWithPath::Dispatch_(const TensorRegionNode*
op, AccessPath path)
void TIRVisitorWithPath::Dispatch_(const OpaqueExprNode* op, AccessPath path)
{}
+void TIRVisitorWithPath::VisitLambda(const LambdaExprNode* op, AccessPath
path) {
+ Visit(op->ty, path->Attr("ty"));
+ std::vector<DefContext<Var>> context;
+ for (size_t i = 0; i < op->vars.size(); ++i) {
+ context.push_back(WithDef(op->vars[i], path->Attr("vars")->ArrayItem(i)));
+ }
+ Visit(op->pred, path->Attr("pred"));
+}
+
void TIRVisitorWithPath::Dispatch_(const TupleNode* op, AccessPath path) {
Visit(op->fields, path->Attr("fields"));
}
diff --git a/src/tirx/ir/tir_visitor_with_path.h
b/src/tirx/ir/tir_visitor_with_path.h
index 1115c4754b..c4da3d05b2 100644
--- a/src/tirx/ir/tir_visitor_with_path.h
+++ b/src/tirx/ir/tir_visitor_with_path.h
@@ -30,6 +30,7 @@
#include <tvm/runtime/logging.h>
#include <tvm/tirx/expr_functor.h>
#include <tvm/tirx/stmt_functor.h>
+#include <tvm/tirx/tile_primitive.h>
#include <exception>
#include <optional>
@@ -80,6 +81,8 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
Dispatch_(tuple_get_item, path);
} else if (auto* buffer_region = obj.as<TensorRegionNode>()) {
Dispatch_(buffer_region, path);
+ } else if (auto* lambda = obj.as<LambdaExprNode>()) {
+ VisitLambda(lambda, path);
} else if (obj.as<OpaqueExprNode>()) {
Dispatch(obj, path);
} else {
@@ -99,6 +102,7 @@ class TIRVisitorWithPath : protected ExprFunctor<void(const
Expr&, ffi::reflecti
// Visit type metadata through its reflected fields, preserving source
access paths.
virtual void Visit(const Type& obj, ffi::reflection::AccessPath path);
virtual void Visit(ffi::AnyView obj, ffi::reflection::AccessPath path);
+ void VisitLambda(const LambdaExprNode* op, ffi::reflection::AccessPath path);
// Visitors for TIR constructs that are neither PrimExpr nor Stmt
virtual void Visit(const IRModule& obj, ffi::reflection::AccessPath path);
