From fe78ec00a561e600a9cb820556a12f5dcd662ced Mon Sep 17 00:00:00 2001 From: tqchen Date: Sat, 29 Aug 2026 16:51:43 +0000 Subject: [PATCH 01/12] [REFACTOR][TIRX] Restore BufferRegion as typed expression --- include/tvm/tirx/expr_functor.h | 5 +++ include/tvm/tirx/stmt.h | 33 +++++++++++++++---- include/tvm/tirx/stmt_functor.h | 2 ++ python/tvm/tirx/__init__.py | 2 +- python/tvm/tirx/expr.py | 8 +---- python/tvm/tirx/expr_functor.py | 27 +++++++++++++++ python/tvm/tirx/stmt.py | 12 +++++-- src/tirx/ir/expr_functor.cc | 19 +++++++++++ src/tirx/ir/stmt.cc | 28 ++++++---------- src/tirx/ir/stmt_functor.cc | 20 +++++++++++ tests/python/tirx-base/test_tir_buffer.py | 17 ++++++++++ .../python/tirx-base/test_tir_stmt_functor.py | 2 ++ tests/python/tirx/test_op.py | 9 +++++ .../tirx/transform/test_stmt_functor.py | 2 ++ .../tirx/transform/test_tirx_expr_functor.py | 29 ++++++++++++++++ 15 files changed, 181 insertions(+), 34 deletions(-) diff --git a/include/tvm/tirx/expr_functor.h b/include/tvm/tirx/expr_functor.h index 193bdcadc4b6..9152535cc834 100644 --- a/include/tvm/tirx/expr_functor.h +++ b/include/tvm/tirx/expr_functor.h @@ -27,6 +27,7 @@ #include #include +#include #include @@ -117,6 +118,7 @@ class ExprFunctor { virtual R VisitExpr_(const VarNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; virtual R VisitExpr_(const BufferLoadNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; virtual R VisitExpr_(const OpaqueExprNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; + virtual R VisitExpr_(const BufferRegionNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; virtual R VisitExpr_(const TupleNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; virtual R VisitExpr_(const TupleGetItemNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; virtual R VisitExpr_(const LetNode* op, Args... args) EXPR_FUNCTOR_DEFAULT; @@ -161,6 +163,7 @@ class ExprFunctor { IR_EXPR_FUNCTOR_DISPATCH(VarNode); IR_EXPR_FUNCTOR_DISPATCH(BufferLoadNode); IR_EXPR_FUNCTOR_DISPATCH(OpaqueExprNode); + IR_EXPR_FUNCTOR_DISPATCH(BufferRegionNode); IR_EXPR_FUNCTOR_DISPATCH(TupleNode); IR_EXPR_FUNCTOR_DISPATCH(TupleGetItemNode); IR_EXPR_FUNCTOR_DISPATCH(LetNode); @@ -213,6 +216,7 @@ class TVM_DLL ExprVisitor : public ExprFunctor { void VisitExpr_(const VarNode* op) override; void VisitExpr_(const BufferLoadNode* op) override; void VisitExpr_(const OpaqueExprNode* op) override; + void VisitExpr_(const BufferRegionNode* op) override; void VisitExpr_(const TupleNode* op) override; void VisitExpr_(const TupleGetItemNode* op) override; void VisitExpr_(const LetNode* op) override; @@ -261,6 +265,7 @@ class TVM_DLL ExprMutator : protected ExprFunctor { Expr VisitExpr_(const VarNode* op) override; Expr VisitExpr_(const BufferLoadNode* op) override; Expr VisitExpr_(const OpaqueExprNode* op) override; + Expr VisitExpr_(const BufferRegionNode* op) override; Expr VisitExpr_(const TupleNode* op) override; Expr VisitExpr_(const TupleGetItemNode* op) override; Expr VisitExpr_(const LetNode* op) override; diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index d95c1af50102..b03aa6cb66e6 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -770,10 +770,33 @@ class Continue : public Stmt { TVM_DEFINE_OBJECT_REF_COW_METHOD(ContinueNode); }; +/*! + * \brief The type of a multi-dimensional buffer region expression. + */ +class BufferRegionTypeNode : public TypeNode { + public: + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef(); + } + + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType", BufferRegionTypeNode, TypeNode); +}; + +/*! + * \brief Managed reference to BufferRegionTypeNode. + */ +class BufferRegionType : public Type { + public: + TVM_DLL BufferRegionType(Span span = Span()); + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, Type, BufferRegionTypeNode); +}; + /*! * \brief Representing the region of multi-dimensional buffer access. */ -class BufferRegionNode : public PrimExprConvertibleNode { +class BufferRegionNode : public ExprNode { public: /*! \brief The buffer of the buffer region. */ BufferVar buffer; @@ -787,17 +810,15 @@ class BufferRegionNode : public PrimExprConvertibleNode { .def_ro("region", &BufferRegionNode::region); } - TVM_DLL PrimExpr ToPrimExpr() const final; - static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode, PrimExprConvertibleNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode, ExprNode); }; /*! * \brief Managed reference to BufferRegionNode. * \sa BufferRegionNode */ -class BufferRegion : public PrimExprConvertible { +class BufferRegion : public Expr { public: TVM_DLL explicit BufferRegion(BufferVar buffer, ffi::Array region); @@ -816,7 +837,7 @@ class BufferRegion : public PrimExprConvertible { */ TVM_DLL static BufferRegion FromPoint(BufferVar buffer, ffi::Array indices); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, PrimExprConvertible, BufferRegionNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, Expr, BufferRegionNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode); }; diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h index 5f8562161c32..d4d67625fc0e 100644 --- a/include/tvm/tirx/stmt_functor.h +++ b/include/tvm/tirx/stmt_functor.h @@ -343,6 +343,7 @@ class TVM_DLL StmtExprVisitor : public ExprVisitor, public StmtVisitor { void VisitExpr(const Expr& e) override { return ExprVisitor::VisitExpr(e); } void VisitExpr_(const BufferLoadNode* op) override; + void VisitExpr_(const BufferRegionNode* op) override; }; /*! @@ -362,6 +363,7 @@ class TVM_DLL StmtExprMutator : public ExprMutator, public StmtMutator { Expr VisitExpr(const Expr& e) override { return ExprMutator::VisitExpr(e); } Expr VisitExpr_(const VarNode* op) override; Expr VisitExpr_(const BufferLoadNode* op) override; + Expr VisitExpr_(const BufferRegionNode* op) override; }; /*! diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py index ae228360aeca..48712407b884 100644 --- a/python/tvm/tirx/__init__.py +++ b/python/tvm/tirx/__init__.py @@ -51,7 +51,7 @@ from .stmt import SeqStmt from .stmt import IfThenElse, Evaluate, stmt_seq, stmt_list -from .stmt import BufferRegion, MatchBufferRegion, SBlock, SBlockRealize +from .stmt import BufferRegion, BufferRegionType, MatchBufferRegion, SBlock, SBlockRealize from .stmt import ScopeIdDefStmt from .tile_primitive import DispatchContext, LambdaExpr, TilePrimitiveCall diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index 6e29d9444499..99d2ceb17fe2 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -74,13 +74,7 @@ def _dtype_is_float(value): def _is_scalar_operand(value): - if isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value): - return True - - # BufferRegion is a C++ PrimExprConvertible, but its Python wrapper is not an ExprOp. - from .stmt import BufferRegion # pylint: disable=import-outside-toplevel - - return isinstance(value, BufferRegion) + return isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value) class ExprOp: diff --git a/python/tvm/tirx/expr_functor.py b/python/tvm/tirx/expr_functor.py index def3b18bda90..101f8ef8b3cc 100644 --- a/python/tvm/tirx/expr_functor.py +++ b/python/tvm/tirx/expr_functor.py @@ -51,6 +51,7 @@ def __init__(self): self._dispatch_map = { "tirx.Var": self.visit_var_, "tirx.BufferLoad": self.visit_buffer_load_, + "tirx.BufferRegion": self.visit_buffer_region_, "tirx.Tuple": self.visit_tuple_, "tirx.TupleGetItem": self.visit_tuple_get_item_, "tirx.Let": self.visit_let_, @@ -123,6 +124,11 @@ def visit_buffer_load_(self, op): def visit_opaque_expr_(self, op): """Default visitor for an opaque construction-time expression.""" + + return self.visit_expr_default_(op) + + def visit_buffer_region_(self, op): + """Default visitor for BufferRegion node.""" return self.visit_expr_default_(op) def visit_tuple_(self, op): @@ -292,6 +298,12 @@ def visit_opaque_expr_(self, op): """Visitor implementation for an opaque construction-time expression.""" pass + def visit_buffer_region_(self, op): + """Visitor implementation for BufferRegion.""" + for region in op.region: + self.visit_expr(region.min) + self.visit_expr(region.extent) + def visit_tuple_(self, op): """Visitor implementation for Tuple.""" _visit_array(op.fields, self.visit_expr) @@ -478,6 +490,21 @@ def visit_opaque_expr_(self, op): """Mutator implementation for an opaque construction-time expression.""" return op + def visit_buffer_region_(self, op): + """Mutator implementation for BufferRegion.""" + + def mutate_range(old): + new_min = self.visit_expr(old.min) + new_extent = self.visit_expr(old.extent) + if new_min is old.min and new_extent is old.extent: + return old + return Range.from_min_extent(new_min, new_extent) + + region = [mutate_range(r) for r in op.region] + if all(old is new for old, new in zip(op.region, region)): + return op + return tvm.tirx.BufferRegion(op.buffer, region) + def visit_tuple_(self, op): """Mutator implementation for Tuple.""" fields = [self.visit_expr(field) for field in op.fields] diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index fe4c98a4025a..6398c9062c8a 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -33,7 +33,7 @@ import tvm_ffi -from tvm.ir import Expr, Range, Span, is_prim_expr +from tvm.ir import Expr, Range, Span, Type, is_prim_expr from tvm.runtime import Object, Scriptable, const from tvm.tirx import IntImm @@ -615,8 +615,16 @@ def __init__(self, value: Expr, span: Span | None = None) -> None: self.__init_handle_by_constructor__(_ffi_api.Evaluate, value, span) # type: ignore +@tvm_ffi.register_object("tirx.BufferRegionType") +class BufferRegionType(Type): + """The structural type of a :class:`BufferRegion` expression.""" + + def __init__(self, span: Span | None = None) -> None: + self.__init_handle_by_constructor__(_ffi_api.BufferRegionType, span) # type: ignore + + @tvm_ffi.register_object("tirx.BufferRegion") -class BufferRegion(Object, Scriptable): +class BufferRegion(Expr, Scriptable): """BufferRegion node. Parameters diff --git a/src/tirx/ir/expr_functor.cc b/src/tirx/ir/expr_functor.cc index 9a73caf2c828..4a3b335600fe 100644 --- a/src/tirx/ir/expr_functor.cc +++ b/src/tirx/ir/expr_functor.cc @@ -36,6 +36,13 @@ void ExprVisitor::VisitExpr_(const BufferLoadNode* op) { void ExprVisitor::VisitExpr_(const OpaqueExprNode* op) {} +void ExprVisitor::VisitExpr_(const BufferRegionNode* op) { + VisitArray(op->region, [this](const Range& range) { + this->VisitExpr(range->min); + this->VisitExpr(range->extent); + }); +} + void ExprVisitor::VisitExpr_(const TupleNode* op) { VisitArray(op->fields, [this](const Expr& e) { this->VisitExpr(e); }); } @@ -130,6 +137,18 @@ Expr ExprMutator::VisitExpr_(const BufferLoadNode* op) { Expr ExprMutator::VisitExpr_(const OpaqueExprNode* op) { return ffi::GetRef(op); } +Expr ExprMutator::VisitExpr_(const BufferRegionNode* op) { + ffi::Array region = op->region.Map([this](const Range& range) { + PrimExpr min = this->VisitPrimExpr(range->min); + PrimExpr extent = this->VisitPrimExpr(range->extent); + return min.same_as(range->min) && extent.same_as(range->extent) + ? range + : Range::FromMinExtent(std::move(min), std::move(extent)); + }); + return region.same_as(op->region) ? ffi::GetRef(op) + : BufferRegion(op->buffer, std::move(region)); +} + Expr ExprMutator::VisitExpr_(const TupleNode* op) { ffi::Array fields = op->fields.Map([this](const Expr& field) { return this->VisitExpr(field); }); diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index a35791c60af3..2837bb79b926 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -51,6 +51,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { ReturnNode::RegisterReflection(); BreakNode::RegisterReflection(); ContinueNode::RegisterReflection(); + BufferRegionTypeNode::RegisterReflection(); BufferRegionNode::RegisterReflection(); MatchBufferRegionNode::RegisterReflection(); SBlockNode::RegisterReflection(); @@ -510,21 +511,10 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // BufferRegion -PrimExpr BufferRegionNode::ToPrimExpr() const { - // Auto convert to PrimExpr if it is a single point load - ffi::Array indices; - indices.reserve(this->region.size()); - for (const Range& r : this->region) { - if (tvm::tirx::is_one(r->extent)) { - indices.push_back(r->min); - } else if (r->extent.as()) { - indices.push_back(tirx::Ramp(r->min, IntImm(r->min.ty(), 1), r->extent)); - } else { - TVM_FFI_THROW(ValueError) << "Cannot convert to BufferLoad: " - << ffi::GetRef(this); - } - } - return tirx::BufferLoad(this->buffer, indices); +BufferRegionType::BufferRegionType(Span span) : Type(ffi::UnsafeInit{}) { + ffi::ObjectPtr node = ffi::make_object(); + node->span = std::move(span); + data_ = std::move(node); } BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region) { @@ -532,6 +522,7 @@ BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region) { << "The dimension between " << buffer << " and region " << region << " mismatched, the buffer is " << buffer; ffi::ObjectPtr node = ffi::make_object(); + node->ty = BufferRegionType(); node->buffer = std::move(buffer); node->region = std::move(region); data_ = std::move(node); @@ -560,9 +551,10 @@ BufferRegion BufferRegion::FromPoint(BufferVar buffer, ffi::Array indi TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("tirx.BufferRegion", [](BufferVar buffer, ffi::Array region) { - return BufferRegion(buffer, region); - }); + refl::GlobalDef() + .def("tirx.BufferRegionType", [](Span span) { return BufferRegionType(span); }) + .def("tirx.BufferRegion", + [](BufferVar buffer, ffi::Array region) { return BufferRegion(buffer, region); }); } // MatchBufferRegion diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc index be7364fb5329..710cc350ad5c 100644 --- a/src/tirx/ir/stmt_functor.cc +++ b/src/tirx/ir/stmt_functor.cc @@ -95,6 +95,11 @@ void StmtExprVisitor::VisitExpr_(const BufferLoadNode* op) { ExprVisitor::VisitExpr_(op); } +void StmtExprVisitor::VisitExpr_(const BufferRegionNode* op) { + this->VisitBufferUse(op->buffer); + ExprVisitor::VisitExpr_(op); +} + void StmtVisitor::VisitStmt_(const AllocBufferNode* op) { this->VisitBufferDef(op->buffer, /*alloc_data=*/true); } @@ -464,6 +469,21 @@ Expr StmtExprMutator::VisitExpr_(const BufferLoadNode* op) { return expr; } +Expr StmtExprMutator::VisitExpr_(const BufferRegionNode* op) { + BufferVar new_buf = this->VisitBufferUse(op->buffer); + ffi::Array new_region = op->region.Map([this](const Range& range) { + PrimExpr min = this->VisitPrimExpr(range->min); + PrimExpr extent = this->VisitPrimExpr(range->extent); + return min.same_as(range->min) && extent.same_as(range->extent) + ? range + : Range::FromMinExtent(std::move(min), std::move(extent)); + }); + if (new_buf.same_as(op->buffer) && new_region.same_as(op->region)) { + return ffi::GetRef(op); + } + return BufferRegion(std::move(new_buf), std::move(new_region)); +} + Stmt StmtMutator::VisitStmt_(const AllocBufferNode* op) { BufferVar new_buf = this->VisitBufferDef(op->buffer, /*alloc_data=*/true); diff --git a/tests/python/tirx-base/test_tir_buffer.py b/tests/python/tirx-base/test_tir_buffer.py index 9db789c1484b..a8bceef88342 100644 --- a/tests/python/tirx-base/test_tir_buffer.py +++ b/tests/python/tirx-base/test_tir_buffer.py @@ -16,6 +16,8 @@ # under the License. # ruff: noqa: E741, F401, F841 +import pickle + import numpy as np import pytest @@ -40,6 +42,21 @@ def test_buffer(): assert not tvm.tirx.is_buffer_var(m) +def test_buffer_region_is_typed_expr_and_call_argument(): + buffer = tvm.tirx.decl_buffer((16,), "float32") + region = buffer[2:10] + + assert isinstance(region, tvm.ir.Expr) + assert isinstance(region.ty, tvm.tirx.BufferRegionType) + + call = tvm.ir.Call(tvm.ir.GlobalVar("consume_region"), [region]) + assert call.args[0].same_as(region) + + restored = pickle.loads(pickle.dumps(call)) + tvm.ir.assert_structural_equal(restored, call, map_free_vars=True) + assert isinstance(restored.args[0].ty, tvm.tirx.BufferRegionType) + + def test_buffer_compatibility_alias_and_global_var_properties(): scalar = tvm.ir.Var("scalar", tvm.ir.PrimType("int32")) buffer = tvm.tirx.decl_buffer((8,), "float32") diff --git a/tests/python/tirx-base/test_tir_stmt_functor.py b/tests/python/tirx-base/test_tir_stmt_functor.py index e3862b062b67..0951d2683316 100644 --- a/tests/python/tirx-base/test_tir_stmt_functor.py +++ b/tests/python/tirx-base/test_tir_stmt_functor.py @@ -364,6 +364,8 @@ def visit_expr(self, expr): if a is expr.a and b is expr.b: return expr return tir.GT(a, b) + elif isinstance(expr, tir.BufferRegion): + return self.visit_buffer_region_(expr) else: self.log.add(f"Expr::{type(expr).__name__}") return expr diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py index 8d6e6326e189..6ee7af745ade 100644 --- a/tests/python/tirx/test_op.py +++ b/tests/python/tirx/test_op.py @@ -73,6 +73,15 @@ def test_tile_primitive_call_pickle_roundtrip(): assert_structural_equal(restored.scope, call.scope) +def test_compose_op_retains_statement_arguments(): + buffer = decl_buffer((16,), "float32", scope="local") + inner = _test("fill", buffer[:], 1.0) + composed = _test("compose_op", inner) + + assert isinstance(composed.args[0], TilePrimitiveCall) + assert composed.args[0].same_as(inner) + + def test_buffer_replacer_no_shared_default(): """Regression test for F4: BufferReplacer default dicts must not be shared.""" from tvm.tirx.transform.common import BufferReplacer diff --git a/tests/python/tirx/transform/test_stmt_functor.py b/tests/python/tirx/transform/test_stmt_functor.py index bf845c65163f..5224967c8a84 100644 --- a/tests/python/tirx/transform/test_stmt_functor.py +++ b/tests/python/tirx/transform/test_stmt_functor.py @@ -375,6 +375,8 @@ def visit_expr(self, expr): if a is expr.a and b is expr.b: return expr return tir.GT(a, b) + elif isinstance(expr, tir.BufferRegion): + return self.visit_buffer_region_(expr) else: self.log.add(f"Expr::{type(expr).__name__}") return expr diff --git a/tests/python/tirx/transform/test_tirx_expr_functor.py b/tests/python/tirx/transform/test_tirx_expr_functor.py index 38845fe6f184..3c24be57e6f6 100644 --- a/tests/python/tirx/transform/test_tirx_expr_functor.py +++ b/tests/python/tirx/transform/test_tirx_expr_functor.py @@ -65,6 +65,35 @@ class BasicVisitor(ExprVisitor): """Default ExprVisitor""" +def test_buffer_region_expr_functor_traversal_and_mutation(): + buffer = tir.decl_buffer((16,), "float32") + begin = tir.Var("begin", "int32") + replacement = tir.Var("replacement", "int32") + region = tir.BufferRegion(buffer, [tvm.ir.Range.from_min_extent(begin, 4)]) + call = tvm.ir.Call(tvm.ir.GlobalVar("consume_region"), [region]) + + class VarCollector(ExprVisitor): + def __init__(self): + super().__init__() + self.vars = [] + + def visit_var_(self, op): + self.vars.append(op) + + collector = VarCollector() + collector(call) + assert collector.vars == [begin] + + class ReplaceBegin(ExprMutator): + def visit_var_(self, op): + return replacement if op.same_as(begin) else op + + updated = ReplaceBegin()(call) + assert isinstance(updated.args[0], tir.BufferRegion) + assert isinstance(updated.args[0].ty, tir.BufferRegionType) + assert updated.args[0].region[0].min.same_as(replacement) + + class ASTLog: """Helper class to log AST""" From d4ab387ddba8708b53f4eb02f4979ba7d2295cb6 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sat, 29 Aug 2026 21:23:04 +0000 Subject: [PATCH 02/12] [FIX][TIRX] Materialize buffer regions at primitive boundaries --- include/tvm/tirx/stmt.h | 3 +++ python/tvm/tirx/op.py | 11 ++++++++++- python/tvm/tirx/script/builder/ir.py | 1 + src/tirx/ir/stmt.cc | 19 ++++++++++++++++++- 4 files changed, 32 insertions(+), 2 deletions(-) diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index b03aa6cb66e6..ffdaebef2b45 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -822,6 +822,9 @@ class BufferRegion : public Expr { public: TVM_DLL explicit BufferRegion(BufferVar buffer, ffi::Array region); + /*! \brief Materialize this region as a scalar or vector buffer load. */ + TVM_DLL PrimExpr ToBufferLoad() const; + /*! * \brief Create a BufferRegion which is full region of the given buffer. * \param buffer The buffer to generate full BufferRegion. diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 39939803a667..db9b421d2dcb 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -32,9 +32,17 @@ from . import _ffi_api from .buffer import Buffer, buffer_data, is_buffer_var from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var +from .stmt import BufferRegion tir = tirx # alias for backward compat with upstream tir.convert() calls + +def _convert_to_prim_expr(value): + if isinstance(value, BufferRegion): + return _ffi_api.BufferRegionToBufferLoad(value) # type: ignore + return value + + # Insertion order matters: a longer prefix has to be tried before the shorter # one it starts with, or `ptx_legacy_mma` would strip as `ptx` + `legacy_mma`. _DEVICE_INTRIN_PREFIX_TO_NAMESPACE = { @@ -253,6 +261,7 @@ def call_intrin(dtype: str | tvm.ir.Type, func_name, *args, attrs=None, span=Non """ if isinstance(func_name, str): func_name = _canonical_device_intrin_name(func_name) + args = tuple(_convert_to_prim_expr(arg) for arg in args) return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype) @@ -1375,7 +1384,7 @@ def reinterpret(dtype, value, span: Span | None = None) -> Expr: dtype = ( PointerType(tvm.ir.PrimType("void")) if dtype == "handle" else tvm.ir.PrimType(dtype) ) - return _ffi_api.reinterpret(dtype, value, span) # type: ignore + return _ffi_api.reinterpret(dtype, _convert_to_prim_expr(value), span) # type: ignore def exp(x): diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index d7897647fdbb..6978d1774a55 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -2376,6 +2376,7 @@ def buffer_store( expr_indices.append(index) if isinstance(value, bool) and buffer.ty.dtype == "bool": value = IntImm("bool", value) + value = _tir_op._convert_to_prim_expr(value) # pylint: disable=protected-access return _ffi_api.BufferStore( # type: ignore[attr-defined] # pylint: disable=no-member buffer, value, expr_indices, predicate ) diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 2837bb79b926..b491fd439ad2 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -517,6 +517,21 @@ BufferRegionType::BufferRegionType(Span span) : Type(ffi::UnsafeInit{}) { data_ = std::move(node); } +PrimExpr BufferRegion::ToBufferLoad() const { + ffi::Array indices; + indices.reserve((*this)->region.size()); + for (const Range& r : (*this)->region) { + if (tirx::is_one(r->extent)) { + indices.push_back(r->min); + } else if (r->extent.as()) { + indices.push_back(Ramp(r->min, IntImm(r->min.ty(), 1), r->extent)); + } else { + TVM_FFI_THROW(ValueError) << "Cannot convert to BufferLoad: " << *this; + } + } + return BufferLoad((*this)->buffer, indices); +} + BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region) { TVM_FFI_ICHECK_EQ(buffer->shape.size(), region.size()) << "The dimension between " << buffer << " and region " << region @@ -554,7 +569,9 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef() .def("tirx.BufferRegionType", [](Span span) { return BufferRegionType(span); }) .def("tirx.BufferRegion", - [](BufferVar buffer, ffi::Array region) { return BufferRegion(buffer, region); }); + [](BufferVar buffer, ffi::Array region) { return BufferRegion(buffer, region); }) + .def("tirx.BufferRegionToBufferLoad", + [](BufferRegion region) { return region.ToBufferLoad(); }); } // MatchBufferRegion From 9d8fe85d250d1d56eae59afa7640431e91443f7a Mon Sep 17 00:00:00 2001 From: tqchen Date: Sat, 29 Aug 2026 23:00:13 +0000 Subject: [PATCH 03/12] [FIX][TIRX] Materialize buffer regions in expression operators --- python/tvm/tirx/expr.py | 82 +++++++++++++++++++++++------------------ python/tvm/tirx/op.py | 17 +++++---- python/tvm/tirx/stmt.py | 5 ++- 3 files changed, 60 insertions(+), 44 deletions(-) diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index 99d2ceb17fe2..c307f0e990f2 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -77,6 +77,18 @@ def _is_scalar_operand(value): return isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value) +def _convert_to_prim_expr(value): + from .stmt import BufferRegion # pylint: disable=import-outside-toplevel + + if isinstance(value, BufferRegion): + return _ffi_api.BufferRegionToBufferLoad(value) # type: ignore + return value + + +def _binary_prim_expr_op(op, lhs, rhs, span=None): + return op(_convert_to_prim_expr(lhs), _convert_to_prim_expr(rhs), span) + + class ExprOp: """Operator overloading for Expr like expressions.""" @@ -84,7 +96,7 @@ class ExprOp: def expr_ty(self) -> ir.PrimType: """Return the compile-time primitive type for expression operators.""" - ty = getattr(self, "ty", None) + ty = getattr(_convert_to_prim_expr(self), "ty", None) if isinstance(ty, ir.PrimType): return ty raise TypeError(f"Cannot determine PrimType for {type(self).__name__}") @@ -92,129 +104,129 @@ def expr_ty(self) -> ir.PrimType: def __add__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _ffi_api._OpAdd(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpAdd, self, other) # type: ignore def __radd__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _ffi_api._OpAdd(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpAdd, other, self) # type: ignore def __sub__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _ffi_api._OpSub(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpSub, self, other) # type: ignore def __rsub__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _ffi_api._OpSub(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpSub, other, self) # type: ignore def __mul__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _ffi_api._OpMul(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpMul, self, other) # type: ignore def __rmul__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _ffi_api._OpMul(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpMul, other, self) # type: ignore def __div__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _ffi_api._OpDiv(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpDiv, self, other) # type: ignore def __rdiv__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _ffi_api._OpDiv(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpDiv, other, self) # type: ignore def __truediv__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _ffi_api._OpDiv(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpDiv, self, other) # type: ignore def __rtruediv__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _ffi_api._OpDiv(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpDiv, other, self) # type: ignore def __floordiv__(self, other: Expr) -> Expr: - return _ffi_api._OpFloorDiv(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpFloorDiv, self, other) # type: ignore def __rfloordiv__(self, other: Expr) -> Expr: - return _ffi_api._OpFloorDiv(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpFloorDiv, other, self) # type: ignore def __mod__(self, other: Expr) -> Expr: - return _ffi_api._OpFloorMod(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpFloorMod, self, other) # type: ignore def __rmod__(self, other: Expr) -> Expr: - return _ffi_api._OpFloorMod(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpFloorMod, other, self) # type: ignore def __neg__(self) -> Expr: neg_one = const(-1, self.expr_ty().dtype) - return self.__mul__(neg_one) + return _ffi_api._OpMul(_convert_to_prim_expr(self), neg_one, None) # type: ignore def __lshift__(self, other: Expr) -> Expr: - return _ffi_api.left_shift(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.left_shift, self, other) # type: ignore def __rlshift__(self, other: Expr) -> Expr: - return _ffi_api.left_shift(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.left_shift, other, self) # type: ignore def __rshift__(self, other: Expr) -> Expr: - return _ffi_api.right_shift(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.right_shift, self, other) # type: ignore def __rrshift__(self, other: Expr) -> Expr: - return _ffi_api.right_shift(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.right_shift, other, self) # type: ignore def __and__(self, other: Expr) -> Expr: - return _ffi_api.bitwise_and(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.bitwise_and, self, other) # type: ignore def __rand__(self, other: Expr) -> Expr: - return _ffi_api.bitwise_and(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.bitwise_and, other, self) # type: ignore def __or__(self, other: Expr) -> Expr: - return _ffi_api.bitwise_or(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.bitwise_or, self, other) # type: ignore def __ror__(self, other: Expr) -> Expr: - return _ffi_api.bitwise_or(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.bitwise_or, other, self) # type: ignore def __xor__(self, other: Expr) -> Expr: - return _ffi_api.bitwise_xor(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.bitwise_xor, self, other) # type: ignore def __rxor__(self, other: Expr) -> Expr: - return _ffi_api.bitwise_xor(other, self, None) # type: ignore + return _binary_prim_expr_op(_ffi_api.bitwise_xor, other, self) # type: ignore def __invert__(self) -> Expr: if _dtype_is_float(self): raise RuntimeError("Cannot use ~ operator on float type Expr.") - return _ffi_api.bitwise_not(self, None) # type: ignore + return _ffi_api.bitwise_not(_convert_to_prim_expr(self), None) # type: ignore def __lt__(self, other: Expr) -> Expr: - return _ffi_api._OpLT(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpLT, self, other) # type: ignore def __le__(self, other: Expr) -> Expr: - return _ffi_api._OpLE(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpLE, self, other) # type: ignore def __eq__(self, other: Expr) -> Expr: - return EqualOp(self, other) + return EqualOp(_convert_to_prim_expr(self), _convert_to_prim_expr(other)) def __ne__(self, other: Expr) -> Expr: - return NotEqualOp(self, other) + return NotEqualOp(_convert_to_prim_expr(self), _convert_to_prim_expr(other)) def __gt__(self, other: Expr) -> Expr: - return _ffi_api._OpGT(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpGT, self, other) # type: ignore def __ge__(self, other: Expr) -> Expr: - return _ffi_api._OpGE(self, other, None) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpGE, self, other) # type: ignore def __nonzero__(self): raise ValueError( @@ -241,7 +253,7 @@ def equal(self, other: Expr, span: Span | None = None) -> bool: ret : Expr The equality expression. """ - return _ffi_api._OpEQ(self, other, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpEQ, self, other, span) # type: ignore def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> Expr: """Cast the expression to other type. @@ -259,7 +271,7 @@ def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> Expr: expr : Expr Expression with new type """ - return _ffi_api._cast(dtype, self, span) # type: ignore + return _ffi_api._cast(dtype, _convert_to_prim_expr(self), span) # type: ignore _overload_prim_expr.__add__ = ExprOp.__add__ diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index db9b421d2dcb..9d25f50b87c3 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -31,18 +31,19 @@ from . import _ffi_api from .buffer import Buffer, buffer_data, is_buffer_var -from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var -from .stmt import BufferRegion +from .expr import ( + BufferLoad, + CommReducer, + ExprOp, + ExprWithOp, + IntImm, + Var, + _convert_to_prim_expr, +) tir = tirx # alias for backward compat with upstream tir.convert() calls -def _convert_to_prim_expr(value): - if isinstance(value, BufferRegion): - return _ffi_api.BufferRegionToBufferLoad(value) # type: ignore - return value - - # Insertion order matters: a longer prefix has to be tried before the shorter # one it starts with, or `ptx_legacy_mma` would strip as `ptx` + `legacy_mma`. _DEVICE_INTRIN_PREFIX_TO_NAMESPACE = { diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 6398c9062c8a..c45a8eac5002 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -36,6 +36,7 @@ from tvm.ir import Expr, Range, Span, Type, is_prim_expr from tvm.runtime import Object, Scriptable, const from tvm.tirx import IntImm +from tvm.tirx.expr import ExprOp from . import _ffi_api from .buffer import Buffer @@ -624,7 +625,7 @@ def __init__(self, span: Span | None = None) -> None: @tvm_ffi.register_object("tirx.BufferRegion") -class BufferRegion(Expr, Scriptable): +class BufferRegion(ExprOp, Expr, Scriptable): """BufferRegion node. Parameters @@ -636,6 +637,8 @@ class BufferRegion(Expr, Scriptable): The region array of the buffer region """ + __hash__ = Expr.__hash__ + buffer: Buffer region: list[Range] From 00e1fe8e6ef60c1adada1be110ad8c640c3220f1 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 00:37:59 +0000 Subject: [PATCH 04/12] [FIX][TIRX] Preserve buffer region identity comparisons --- python/tvm/tirx/stmt.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index c45a8eac5002..338e7a2b1210 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -645,6 +645,12 @@ class BufferRegion(ExprOp, Expr, Scriptable): def __init__(self, buffer: Buffer, region: list[Range]) -> None: self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer, region) # type: ignore + def __eq__(self, other) -> bool: + return Object.__eq__(self, other) + + def __ne__(self, other) -> bool: + return Object.__ne__(self, other) + def __getitem__(self, indices): from ..arith import Analyzer From 6155d3690bf9af3b33bbc7bab024d917acba8678 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 01:11:47 +0000 Subject: [PATCH 05/12] [REFACTOR][TIRX] Localize buffer region primitive conversion --- python/tvm/ir/__init__.py | 12 ++- python/tvm/ir/expr.py | 94 ++++++++++++---------- python/tvm/ir/type.py | 4 + python/tvm/tirx/expr.py | 23 ++++-- python/tvm/tirx/op.py | 43 +++++----- python/tvm/tirx/script/parser/operation.py | 3 +- python/tvm/tirx/stmt.py | 15 +--- 7 files changed, 111 insertions(+), 83 deletions(-) diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py index ca83d21f3aac..729c90465f40 100644 --- a/python/tvm/ir/__init__.py +++ b/python/tvm/ir/__init__.py @@ -33,10 +33,19 @@ # Register Type before Expr. Expr's reflected ``ty`` field otherwise creates # an auto-generated Type wrapper before the concrete Python class is available. -from .type import FuncType, OpaqueType, PointerType, PrimType, TupleType, Type +from .type import ( + FuncType, + OpaqueType, + PointerType, + PrimExprConvertibleType, + PrimType, + TupleType, + Type, +) from .expr import ( Call, Expr, + ExprWithOp, GlobalVar, OpaqueExpr, Range, @@ -44,6 +53,7 @@ TupleGetItem, Var, is_prim_expr, + is_prim_expr_convertible, is_prim_var, ) from .function import BaseFunc, CallingConv diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py index 393e825646b3..4ba3e478faca 100644 --- a/python/tvm/ir/expr.py +++ b/python/tvm/ir/expr.py @@ -25,6 +25,7 @@ from ..runtime import Object, Scriptable from . import _ffi_api, _overload_prim_expr, _tensor_expr_overload from .base import Node, Span +from .type import PrimExprConvertibleType @tvm_ffi.register_object("ir.Expr") @@ -50,6 +51,15 @@ def is_prim_var(value: object) -> bool: return isinstance(value, Var) and type(value) is Var and is_prim_expr(value) +def is_prim_expr_convertible(value: object) -> bool: + """Return whether an expression's type opts into primitive operators.""" + return isinstance(value, Expr) and isinstance(value.ty, PrimExprConvertibleType) + + +def _supports_prim_expr_ops(value: object) -> bool: + return is_prim_expr(value) or is_prim_expr_convertible(value) + + @tvm_ffi.register_object("ir.GlobalVar") class GlobalVar(Expr): """A global variable in the IR. @@ -100,99 +110,97 @@ def is_tir_arg(x): raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}") -class _ExprWithOp(Expr, Scriptable): +class ExprWithOp(Expr, Scriptable): """Common type-directed operator behavior for core expressions.""" __hash__ = Expr.__hash__ def expr_ty(self): - """Return this expression's primitive result type.""" - if is_prim_expr(self): - return self.ty - raise TypeError(f"Expected a primitive-valued expression, but result type is {self.ty}") + """Return this expression's result type.""" + return self.ty def __add__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__add__(self, other) return _tensor_expr_overload.__add__(self, other) def __radd__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__radd__(self, other) return _tensor_expr_overload.__radd__(self, other) def __sub__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__sub__(self, other) return _tensor_expr_overload.__sub__(self, other) def __rsub__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rsub__(self, other) return _tensor_expr_overload.__rsub__(self, other) def __mul__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__mul__(self, other) return _tensor_expr_overload.__mul__(self, other) def __rmul__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rmul__(self, other) return _tensor_expr_overload.__rmul__(self, other) def __div__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__div__(self, other) return _tensor_expr_overload.__div__(self, other) def __rdiv__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rdiv__(self, other) return _tensor_expr_overload.__rdiv__(self, other) def __truediv__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__truediv__(self, other) return _tensor_expr_overload.__truediv__(self, other) def __rtruediv__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rtruediv__(self, other) return _tensor_expr_overload.__rtruediv__(self, other) def __floordiv__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__floordiv__(self, other) return _tensor_expr_overload.__floordiv__(self, other) def __rfloordiv__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rfloordiv__(self, other) return _tensor_expr_overload.__rfloordiv__(self, other) def __mod__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__mod__(self, other) return _tensor_expr_overload.__mod__(self, other) def __rmod__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rmod__(self, other) return _tensor_expr_overload.__rmod__(self, other) def __pow__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return NotImplemented return _tensor_expr_overload.__pow__(self, other) def __rpow__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return NotImplemented return _tensor_expr_overload.__rpow__(self, other) def __neg__(self): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): result = _overload_prim_expr.__neg__(self) if result is NotImplemented: raise TypeError("Primitive expression overload __neg__ is not registered") @@ -203,57 +211,57 @@ def __neg__(self): return result def __lshift__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__lshift__(self, other) return NotImplemented def __rlshift__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rlshift__(self, other) return NotImplemented def __rshift__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rshift__(self, other) return NotImplemented def __rrshift__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rrshift__(self, other) return NotImplemented def __and__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__and__(self, other) return NotImplemented def __rand__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rand__(self, other) return NotImplemented def __or__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__or__(self, other) return NotImplemented def __ror__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__ror__(self, other) return NotImplemented def __xor__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__xor__(self, other) return NotImplemented def __rxor__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__rxor__(self, other) return NotImplemented def __invert__(self): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): result = _overload_prim_expr.__invert__(self) if result is NotImplemented: raise TypeError("Primitive expression overload __invert__ is not registered") @@ -261,12 +269,12 @@ def __invert__(self): return NotImplemented def __lt__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__lt__(self, other) return _tensor_expr_overload.__lt__(self, other) def __le__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__le__(self, other) return _tensor_expr_overload.__le__(self, other) @@ -281,12 +289,12 @@ def __ne__(self, other): return Object.__ne__(self, other) def __gt__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__gt__(self, other) return _tensor_expr_overload.__gt__(self, other) def __ge__(self, other): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): return _overload_prim_expr.__ge__(self, other) return _tensor_expr_overload.__ge__(self, other) @@ -306,7 +314,7 @@ def equal(self, other, span=None): return result def astype(self, dtype, span=None): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): result = _overload_prim_expr.astype(self, dtype, span) if result is NotImplemented: raise TypeError("Primitive expression overload astype is not registered") @@ -317,7 +325,7 @@ def astype(self, dtype, span=None): return result def __call__(self, *args, attrs=None): - if is_prim_expr(self): + if _supports_prim_expr_ops(self): raise TypeError("A primitive-valued expression cannot be called") result = _tensor_expr_overload.__call__(self, *args, attrs=attrs) if result is NotImplemented: @@ -334,7 +342,7 @@ def __getitem__(self, index): @tvm_ffi.register_object("ir.Tuple") -class Tuple(_ExprWithOp): +class Tuple(ExprWithOp): """Tuple expression that groups several fields together. Parameters @@ -367,7 +375,7 @@ def __len__(self) -> int: @tvm_ffi.register_object("ir.TupleGetItem") -class TupleGetItem(_ExprWithOp): +class TupleGetItem(ExprWithOp): """Get the index-th item from a tuple. Parameters @@ -391,7 +399,7 @@ def __init__(self, tuple_value: Expr, index: int, span: Span | None = None): @tvm_ffi.register_object("ir.Call") -class Call(_ExprWithOp): +class Call(ExprWithOp): """Core function call node.""" op: Expr @@ -430,7 +438,7 @@ def __init__( @tvm_ffi.register_object("ir.Var") -class Var(_ExprWithOp): +class Var(ExprWithOp): """A canonical local variable in the IR. Parameters diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py index 015232963e1b..cf71f8f7f23f 100644 --- a/python/tvm/ir/type.py +++ b/python/tvm/ir/type.py @@ -54,6 +54,10 @@ def same_as(self, other): return self.is_(other) +class PrimExprConvertibleType(Type): + """Marker for non-primitive expressions accepted by primitive operators.""" + + @tvm_ffi.register_object("ir.OpaqueType") class OpaqueType(Type): """Type marker for opaque values that must be removed from finished IR.""" diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index c307f0e990f2..ee56f4e81a94 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -54,6 +54,7 @@ def div_ambiguity_error() -> RuntimeError: def _dtype_is_int(value): + value = _convert_to_prim_expr(value) if isinstance(value, int): return True if isinstance(value, ExprOp): @@ -64,6 +65,7 @@ def _dtype_is_int(value): def _dtype_is_float(value): + value = _convert_to_prim_expr(value) if isinstance(value, float): return True if isinstance(value, ExprOp): @@ -74,7 +76,11 @@ def _dtype_is_float(value): def _is_scalar_operand(value): - return isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value) + return ( + isinstance(value, ExprOp | int | float) + or ir.is_prim_expr(value) + or ir.is_prim_expr_convertible(value) + ) def _convert_to_prim_expr(value): @@ -94,12 +100,12 @@ class ExprOp: # TODO(tkonolige): use inspect to add source information to these objects - def expr_ty(self) -> ir.PrimType: - """Return the compile-time primitive type for expression operators.""" - ty = getattr(_convert_to_prim_expr(self), "ty", None) - if isinstance(ty, ir.PrimType): + def expr_ty(self) -> ir.Type: + """Return the expression's compile-time type.""" + ty = getattr(self, "ty", None) + if isinstance(ty, ir.Type): return ty - raise TypeError(f"Cannot determine PrimType for {type(self).__name__}") + raise TypeError(f"Cannot determine Expr type for {type(self).__name__}") def __add__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): @@ -172,8 +178,9 @@ def __rmod__(self, other: Expr) -> Expr: return _binary_prim_expr_op(_ffi_api._OpFloorMod, other, self) # type: ignore def __neg__(self) -> Expr: - neg_one = const(-1, self.expr_ty().dtype) - return _ffi_api._OpMul(_convert_to_prim_expr(self), neg_one, None) # type: ignore + value = _convert_to_prim_expr(self) + neg_one = const(-1, value.ty.dtype) + return _ffi_api._OpMul(value, neg_one, None) # type: ignore def __lshift__(self, other: Expr) -> Expr: return _binary_prim_expr_op(_ffi_api.left_shift, self, other) # type: ignore diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 9d25f50b87c3..57a5c5887d1e 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -38,6 +38,7 @@ ExprWithOp, IntImm, Var, + _binary_prim_expr_op, _convert_to_prim_expr, ) @@ -72,6 +73,7 @@ def _canonical_device_intrin_name(func_name: str) -> str: def _primexpr_ty(expr): """Return the runtime primitive type of an expression.""" + expr = _convert_to_prim_expr(expr) if isinstance(expr, tvm.ir.PrimType): return expr ty = getattr(expr, "ty", None) @@ -1940,7 +1942,7 @@ def bitwise_and(x, y, span=None): res : Expr The result. """ - return _ffi_api.bitwise_and(x, y, span) + return _binary_prim_expr_op(_ffi_api.bitwise_and, x, y, span) def bitwise_not(x, span=None): @@ -1959,7 +1961,7 @@ def bitwise_not(x, span=None): res : Expr The result. """ - return _ffi_api.bitwise_not(x, span) + return _ffi_api.bitwise_not(_convert_to_prim_expr(x), span) def bitwise_or(x, y, span=None): @@ -1981,7 +1983,7 @@ def bitwise_or(x, y, span=None): res : Expr The result. """ - return _ffi_api.bitwise_or(x, y, span) + return _binary_prim_expr_op(_ffi_api.bitwise_or, x, y, span) def bitwise_xor(x, y, span=None): @@ -2003,7 +2005,7 @@ def bitwise_xor(x, y, span=None): res : Expr The result. """ - return _ffi_api.bitwise_xor(x, y, span) + return _binary_prim_expr_op(_ffi_api.bitwise_xor, x, y, span) def round(x, span=None): @@ -2282,7 +2284,7 @@ def power(x, y, span=None): z : Expr The result. """ - return _ffi_api._OpPow(x, y, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpPow, x, y, span) # type: ignore def pow(x, y, span=None): @@ -2304,7 +2306,7 @@ def pow(x, y, span=None): z : Expr The result. """ - return _ffi_api._OpPow(x, y, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpPow, x, y, span) # type: ignore def popcount(x): @@ -2415,7 +2417,7 @@ def shift_left(x, y, span=None): z : Expr The result. """ - return _ffi_api.left_shift(x, y, span) + return _binary_prim_expr_op(_ffi_api.left_shift, x, y, span) def shift_right(x, y, span=None): @@ -2434,7 +2436,7 @@ def shift_right(x, y, span=None): z : Expr The result. """ - return _ffi_api.right_shift(x, y, span) + return _binary_prim_expr_op(_ffi_api.right_shift, x, y, span) def fmod(x, y): @@ -2487,7 +2489,12 @@ def if_then_else(cond, t, f, span=None): Unlike Select, if_then_else cannot be vectorized if some lanes in the vector have different conditions. """ - return _ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore + return _ffi_api._OpIfThenElse( + _convert_to_prim_expr(cond), + _convert_to_prim_expr(t), + _convert_to_prim_expr(f), + span, + ) # type: ignore def div(a, b, span=None): @@ -2512,7 +2519,7 @@ def div(a, b, span=None): ---- When operands are integers, returns truncdiv(a, b, span). """ - return _ffi_api._OpDiv(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpDiv, a, b, span) # type: ignore def indexdiv(a, b, span=None): @@ -2540,7 +2547,7 @@ def indexdiv(a, b, span=None): This function may take advantage of operands' non-negativeness. """ - return _ffi_api._OpIndexDiv(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpIndexDiv, a, b, span) # type: ignore def indexmod(a, b, span=None): @@ -2568,7 +2575,7 @@ def indexmod(a, b, span=None): This function may take advantage of operands' non-negativeness. """ - return _ffi_api._OpIndexMod(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpIndexMod, a, b, span) # type: ignore def truncdiv(a, b, span=None): @@ -2594,7 +2601,7 @@ def truncdiv(a, b, span=None): ---- This is the default integer division behavior in C. """ - return _ffi_api._OpTruncDiv(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpTruncDiv, a, b, span) # type: ignore def truncmod(a, b, span=None): @@ -2620,7 +2627,7 @@ def truncmod(a, b, span=None): ---- This is the default integer division behavior in C. """ - return _ffi_api._OpTruncMod(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpTruncMod, a, b, span) # type: ignore def floordiv(a, b, span=None): @@ -2642,7 +2649,7 @@ def floordiv(a, b, span=None): res : Expr The result expression. """ - return _ffi_api._OpFloorDiv(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpFloorDiv, a, b, span) # type: ignore def logaddexp(a, b, span=None): @@ -2664,7 +2671,7 @@ def logaddexp(a, b, span=None): res : Expr The result expression. """ - return _ffi_api._OpLogAddExp(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpLogAddExp, a, b, span) # type: ignore def floormod(a, b, span=None): @@ -2686,7 +2693,7 @@ def floormod(a, b, span=None): res : Expr The result expression. """ - return _ffi_api._OpFloorMod(a, b, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpFloorMod, a, b, span) # type: ignore def ceildiv(lhs, rhs, span=None): @@ -2706,7 +2713,7 @@ def ceildiv(lhs, rhs, span=None): op : tvm.Expr The result Expr of ceildiv operaton. """ - return _ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore + return _binary_prim_expr_op(_ffi_api._OpCeilDiv, lhs, rhs, span) # type: ignore def comm_reducer(fcombine, fidentity, name="reduce"): diff --git a/python/tvm/tirx/script/parser/operation.py b/python/tvm/tirx/script/parser/operation.py index fd67d6f12591..eacf15ed2e9d 100644 --- a/python/tvm/tirx/script/parser/operation.py +++ b/python/tvm/tirx/script/parser/operation.py @@ -22,13 +22,14 @@ from tvm.runtime import DataTypeCode from tvm.script.parser._core import OpMethod, doc, register_op from tvm.tirx import IntImm -from tvm.tirx.expr import FloatImm +from tvm.tirx.expr import FloatImm, _convert_to_prim_expr def _register_expr_op(ty: type): # pylint: disable=invalid-name ty._dispatch_type = ty # pylint: disable=protected-access def _expr_ty(expr): + expr = _convert_to_prim_expr(expr) ty = expr.ty if tvm.ir.is_prim_expr(expr) else None if not isinstance(ty, PrimType): ty = expr.expr_ty() diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 338e7a2b1210..6f0f21871029 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -33,10 +33,9 @@ import tvm_ffi -from tvm.ir import Expr, Range, Span, Type, is_prim_expr +from tvm.ir import Expr, ExprWithOp, PrimExprConvertibleType, Range, Span, is_prim_expr from tvm.runtime import Object, Scriptable, const from tvm.tirx import IntImm -from tvm.tirx.expr import ExprOp from . import _ffi_api from .buffer import Buffer @@ -617,7 +616,7 @@ def __init__(self, value: Expr, span: Span | None = None) -> None: @tvm_ffi.register_object("tirx.BufferRegionType") -class BufferRegionType(Type): +class BufferRegionType(PrimExprConvertibleType): """The structural type of a :class:`BufferRegion` expression.""" def __init__(self, span: Span | None = None) -> None: @@ -625,7 +624,7 @@ def __init__(self, span: Span | None = None) -> None: @tvm_ffi.register_object("tirx.BufferRegion") -class BufferRegion(ExprOp, Expr, Scriptable): +class BufferRegion(ExprWithOp): """BufferRegion node. Parameters @@ -637,20 +636,12 @@ class BufferRegion(ExprOp, Expr, Scriptable): The region array of the buffer region """ - __hash__ = Expr.__hash__ - buffer: Buffer region: list[Range] def __init__(self, buffer: Buffer, region: list[Range]) -> None: self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer, region) # type: ignore - def __eq__(self, other) -> bool: - return Object.__eq__(self, other) - - def __ne__(self, other) -> bool: - return Object.__ne__(self, other) - def __getitem__(self, indices): from ..arith import Analyzer From 7bb4f86d0866d499571c2e7c276b3ae86cf23ac3 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 01:47:34 +0000 Subject: [PATCH 06/12] [REFACTOR][TIRX] Convert primitive operands at FFI boundaries --- include/tvm/ir/base_expr.h | 29 ++++++- include/tvm/tirx/stmt.h | 12 ++- python/tvm/ir/type.py | 1 + python/tvm/tirx/expr.py | 97 +++++++++------------- python/tvm/tirx/op.py | 54 +++++------- python/tvm/tirx/script/builder/ir.py | 1 - python/tvm/tirx/script/parser/operation.py | 9 +- src/ir/expr.cc | 16 ++++ src/ir/type.cc | 1 + src/tirx/ir/stmt.cc | 10 ++- src/tirx/op/op.cc | 1 + 11 files changed, 125 insertions(+), 106 deletions(-) diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h index 1e39240a7f16..a0f46309e117 100644 --- a/include/tvm/ir/base_expr.h +++ b/include/tvm/ir/base_expr.h @@ -37,6 +37,9 @@ namespace tvm { +class Expr; +class PrimExpr; + /*! * \brief Type is the base type of all types. * @@ -85,6 +88,29 @@ class Type : public ffi::ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef, TypeNode); }; +/*! + * \brief Base type for expressions that can be converted to PrimExpr at typed FFI boundaries. + */ +class PrimExprConvertibleTypeNode : public TypeNode { + public: + virtual PrimExpr ConvertToPrimExpr(Expr expr) const = 0; + + static void RegisterReflection() { + namespace refl = tvm::ffi::reflection; + refl::ObjectDef(); + } + + static constexpr const uint32_t _type_child_slots = 1; + TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertibleType", PrimExprConvertibleTypeNode, TypeNode); +}; + +/*! \brief Managed reference to PrimExprConvertibleTypeNode. */ +class PrimExprConvertibleType : public Type { + public: + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimExprConvertibleType, Type, + PrimExprConvertibleTypeNode); +}; + /*! * \brief Type marker for opaque construction-time expressions. * @@ -537,10 +563,11 @@ struct TypeTraits using Base::GetMismatchTypeInfo; using Base::MoveFromAnyAfterCheck; using Base::MoveToAny; - using Base::TryCastFromAnyView; using Base::TypeSchema; using Base::TypeStr; + TVM_DLL static std::optional TryCastFromAnyView(const TVMFFIAny* src); + TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value); TVM_DLL static PrimExpr ConvertFallbackValue(int64_t value); TVM_DLL static PrimExpr ConvertFallbackValue(double value); diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index ffdaebef2b45..7129f33ba12f 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -773,24 +773,28 @@ class Continue : public Stmt { /*! * \brief The type of a multi-dimensional buffer region expression. */ -class BufferRegionTypeNode : public TypeNode { +class BufferRegionTypeNode : public PrimExprConvertibleTypeNode { public: + TVM_DLL PrimExpr ConvertToPrimExpr(Expr expr) const final; + static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType", BufferRegionTypeNode, TypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType", BufferRegionTypeNode, + PrimExprConvertibleTypeNode); }; /*! * \brief Managed reference to BufferRegionTypeNode. */ -class BufferRegionType : public Type { +class BufferRegionType : public PrimExprConvertibleType { public: TVM_DLL BufferRegionType(Span span = Span()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, Type, BufferRegionTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, PrimExprConvertibleType, + BufferRegionTypeNode); }; /*! diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py index cf71f8f7f23f..c4d0c29c18ab 100644 --- a/python/tvm/ir/type.py +++ b/python/tvm/ir/type.py @@ -54,6 +54,7 @@ def same_as(self, other): return self.is_(other) +@tvm_ffi.register_object("ir.PrimExprConvertibleType") class PrimExprConvertibleType(Type): """Marker for non-primitive expressions accepted by primitive operators.""" diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index ee56f4e81a94..7533e1d2b901 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -54,24 +54,18 @@ def div_ambiguity_error() -> RuntimeError: def _dtype_is_int(value): - value = _convert_to_prim_expr(value) if isinstance(value, int): return True - if isinstance(value, ExprOp): - return value.expr_ty().matches_code(DataTypeCode.INT) - if ir.is_prim_expr(value): - return value.ty.matches_code(DataTypeCode.INT) + if isinstance(value, ExprOp) or ir.is_prim_expr(value) or ir.is_prim_expr_convertible(value): + return _ffi_api._PrimExprType(value).matches_code(DataTypeCode.INT) # type: ignore return False def _dtype_is_float(value): - value = _convert_to_prim_expr(value) if isinstance(value, float): return True - if isinstance(value, ExprOp): - return value.expr_ty().matches_code(DataTypeCode.FLOAT) - if ir.is_prim_expr(value): - return value.ty.matches_code(DataTypeCode.FLOAT) + if isinstance(value, ExprOp) or ir.is_prim_expr(value) or ir.is_prim_expr_convertible(value): + return _ffi_api._PrimExprType(value).matches_code(DataTypeCode.FLOAT) # type: ignore return False @@ -83,18 +77,6 @@ def _is_scalar_operand(value): ) -def _convert_to_prim_expr(value): - from .stmt import BufferRegion # pylint: disable=import-outside-toplevel - - if isinstance(value, BufferRegion): - return _ffi_api.BufferRegionToBufferLoad(value) # type: ignore - return value - - -def _binary_prim_expr_op(op, lhs, rhs, span=None): - return op(_convert_to_prim_expr(lhs), _convert_to_prim_expr(rhs), span) - - class ExprOp: """Operator overloading for Expr like expressions.""" @@ -110,130 +92,129 @@ def expr_ty(self) -> ir.Type: def __add__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _binary_prim_expr_op(_ffi_api._OpAdd, self, other) # type: ignore + return _ffi_api._OpAdd(self, other, None) # type: ignore def __radd__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _binary_prim_expr_op(_ffi_api._OpAdd, other, self) # type: ignore + return _ffi_api._OpAdd(other, self, None) # type: ignore def __sub__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _binary_prim_expr_op(_ffi_api._OpSub, self, other) # type: ignore + return _ffi_api._OpSub(self, other, None) # type: ignore def __rsub__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _binary_prim_expr_op(_ffi_api._OpSub, other, self) # type: ignore + return _ffi_api._OpSub(other, self, None) # type: ignore def __mul__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _binary_prim_expr_op(_ffi_api._OpMul, self, other) # type: ignore + return _ffi_api._OpMul(self, other, None) # type: ignore def __rmul__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented - return _binary_prim_expr_op(_ffi_api._OpMul, other, self) # type: ignore + return _ffi_api._OpMul(other, self, None) # type: ignore def __div__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _binary_prim_expr_op(_ffi_api._OpDiv, self, other) # type: ignore + return _ffi_api._OpDiv(self, other, None) # type: ignore def __rdiv__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _binary_prim_expr_op(_ffi_api._OpDiv, other, self) # type: ignore + return _ffi_api._OpDiv(other, self, None) # type: ignore def __truediv__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _binary_prim_expr_op(_ffi_api._OpDiv, self, other) # type: ignore + return _ffi_api._OpDiv(self, other, None) # type: ignore def __rtruediv__(self, other: Expr) -> Expr: if not _is_scalar_operand(other): return NotImplemented if _dtype_is_int(self) and _dtype_is_int(other): raise div_ambiguity_error() - return _binary_prim_expr_op(_ffi_api._OpDiv, other, self) # type: ignore + return _ffi_api._OpDiv(other, self, None) # type: ignore def __floordiv__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpFloorDiv, self, other) # type: ignore + return _ffi_api._OpFloorDiv(self, other, None) # type: ignore def __rfloordiv__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpFloorDiv, other, self) # type: ignore + return _ffi_api._OpFloorDiv(other, self, None) # type: ignore def __mod__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpFloorMod, self, other) # type: ignore + return _ffi_api._OpFloorMod(self, other, None) # type: ignore def __rmod__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpFloorMod, other, self) # type: ignore + return _ffi_api._OpFloorMod(other, self, None) # type: ignore def __neg__(self) -> Expr: - value = _convert_to_prim_expr(self) - neg_one = const(-1, value.ty.dtype) - return _ffi_api._OpMul(value, neg_one, None) # type: ignore + neg_one = const(-1, _ffi_api._PrimExprType(self).dtype) # type: ignore + return _ffi_api._OpMul(self, neg_one, None) # type: ignore def __lshift__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.left_shift, self, other) # type: ignore + return _ffi_api.left_shift(self, other, None) # type: ignore def __rlshift__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.left_shift, other, self) # type: ignore + return _ffi_api.left_shift(other, self, None) # type: ignore def __rshift__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.right_shift, self, other) # type: ignore + return _ffi_api.right_shift(self, other, None) # type: ignore def __rrshift__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.right_shift, other, self) # type: ignore + return _ffi_api.right_shift(other, self, None) # type: ignore def __and__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.bitwise_and, self, other) # type: ignore + return _ffi_api.bitwise_and(self, other, None) # type: ignore def __rand__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.bitwise_and, other, self) # type: ignore + return _ffi_api.bitwise_and(other, self, None) # type: ignore def __or__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.bitwise_or, self, other) # type: ignore + return _ffi_api.bitwise_or(self, other, None) # type: ignore def __ror__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.bitwise_or, other, self) # type: ignore + return _ffi_api.bitwise_or(other, self, None) # type: ignore def __xor__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.bitwise_xor, self, other) # type: ignore + return _ffi_api.bitwise_xor(self, other, None) # type: ignore def __rxor__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api.bitwise_xor, other, self) # type: ignore + return _ffi_api.bitwise_xor(other, self, None) # type: ignore def __invert__(self) -> Expr: if _dtype_is_float(self): raise RuntimeError("Cannot use ~ operator on float type Expr.") - return _ffi_api.bitwise_not(_convert_to_prim_expr(self), None) # type: ignore + return _ffi_api.bitwise_not(self, None) # type: ignore def __lt__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpLT, self, other) # type: ignore + return _ffi_api._OpLT(self, other, None) # type: ignore def __le__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpLE, self, other) # type: ignore + return _ffi_api._OpLE(self, other, None) # type: ignore def __eq__(self, other: Expr) -> Expr: - return EqualOp(_convert_to_prim_expr(self), _convert_to_prim_expr(other)) + return EqualOp(self, other) def __ne__(self, other: Expr) -> Expr: - return NotEqualOp(_convert_to_prim_expr(self), _convert_to_prim_expr(other)) + return NotEqualOp(self, other) def __gt__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpGT, self, other) # type: ignore + return _ffi_api._OpGT(self, other, None) # type: ignore def __ge__(self, other: Expr) -> Expr: - return _binary_prim_expr_op(_ffi_api._OpGE, self, other) # type: ignore + return _ffi_api._OpGE(self, other, None) # type: ignore def __nonzero__(self): raise ValueError( @@ -260,7 +241,7 @@ def equal(self, other: Expr, span: Span | None = None) -> bool: ret : Expr The equality expression. """ - return _binary_prim_expr_op(_ffi_api._OpEQ, self, other, span) # type: ignore + return _ffi_api._OpEQ(self, other, span) # type: ignore def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> Expr: """Cast the expression to other type. @@ -278,7 +259,7 @@ def astype(self, dtype: str | ir.PrimType, span: Span | None = None) -> Expr: expr : Expr Expression with new type """ - return _ffi_api._cast(dtype, _convert_to_prim_expr(self), span) # type: ignore + return _ffi_api._cast(dtype, self, span) # type: ignore _overload_prim_expr.__add__ = ExprOp.__add__ diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 57a5c5887d1e..199f44c09480 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -34,12 +34,9 @@ from .expr import ( BufferLoad, CommReducer, - ExprOp, ExprWithOp, IntImm, Var, - _binary_prim_expr_op, - _convert_to_prim_expr, ) tir = tirx # alias for backward compat with upstream tir.convert() calls @@ -73,15 +70,9 @@ def _canonical_device_intrin_name(func_name: str) -> str: def _primexpr_ty(expr): """Return the runtime primitive type of an expression.""" - expr = _convert_to_prim_expr(expr) if isinstance(expr, tvm.ir.PrimType): return expr - ty = getattr(expr, "ty", None) - if isinstance(ty, tvm.ir.PrimType): - return ty - if isinstance(expr, ExprOp): - return expr.expr_ty() - raise TypeError(f"Cannot determine Expr type for {type(expr).__name__}") + return _ffi_api._PrimExprType(expr) # type: ignore def _primexpr_dtype(expr): @@ -264,7 +255,6 @@ def call_intrin(dtype: str | tvm.ir.Type, func_name, *args, attrs=None, span=Non """ if isinstance(func_name, str): func_name = _canonical_device_intrin_name(func_name) - args = tuple(_convert_to_prim_expr(arg) for arg in args) return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype) @@ -1387,7 +1377,7 @@ def reinterpret(dtype, value, span: Span | None = None) -> Expr: dtype = ( PointerType(tvm.ir.PrimType("void")) if dtype == "handle" else tvm.ir.PrimType(dtype) ) - return _ffi_api.reinterpret(dtype, _convert_to_prim_expr(value), span) # type: ignore + return _ffi_api.reinterpret(dtype, value, span) # type: ignore def exp(x): @@ -1942,7 +1932,7 @@ def bitwise_and(x, y, span=None): res : Expr The result. """ - return _binary_prim_expr_op(_ffi_api.bitwise_and, x, y, span) + return _ffi_api.bitwise_and(x, y, span) def bitwise_not(x, span=None): @@ -1961,7 +1951,7 @@ def bitwise_not(x, span=None): res : Expr The result. """ - return _ffi_api.bitwise_not(_convert_to_prim_expr(x), span) + return _ffi_api.bitwise_not(x, span) def bitwise_or(x, y, span=None): @@ -1983,7 +1973,7 @@ def bitwise_or(x, y, span=None): res : Expr The result. """ - return _binary_prim_expr_op(_ffi_api.bitwise_or, x, y, span) + return _ffi_api.bitwise_or(x, y, span) def bitwise_xor(x, y, span=None): @@ -2005,7 +1995,7 @@ def bitwise_xor(x, y, span=None): res : Expr The result. """ - return _binary_prim_expr_op(_ffi_api.bitwise_xor, x, y, span) + return _ffi_api.bitwise_xor(x, y, span) def round(x, span=None): @@ -2284,7 +2274,7 @@ def power(x, y, span=None): z : Expr The result. """ - return _binary_prim_expr_op(_ffi_api._OpPow, x, y, span) # type: ignore + return _ffi_api._OpPow(x, y, span) # type: ignore def pow(x, y, span=None): @@ -2306,7 +2296,7 @@ def pow(x, y, span=None): z : Expr The result. """ - return _binary_prim_expr_op(_ffi_api._OpPow, x, y, span) # type: ignore + return _ffi_api._OpPow(x, y, span) # type: ignore def popcount(x): @@ -2417,7 +2407,7 @@ def shift_left(x, y, span=None): z : Expr The result. """ - return _binary_prim_expr_op(_ffi_api.left_shift, x, y, span) + return _ffi_api.left_shift(x, y, span) def shift_right(x, y, span=None): @@ -2436,7 +2426,7 @@ def shift_right(x, y, span=None): z : Expr The result. """ - return _binary_prim_expr_op(_ffi_api.right_shift, x, y, span) + return _ffi_api.right_shift(x, y, span) def fmod(x, y): @@ -2490,9 +2480,9 @@ def if_then_else(cond, t, f, span=None): if some lanes in the vector have different conditions. """ return _ffi_api._OpIfThenElse( - _convert_to_prim_expr(cond), - _convert_to_prim_expr(t), - _convert_to_prim_expr(f), + cond, + t, + f, span, ) # type: ignore @@ -2519,7 +2509,7 @@ def div(a, b, span=None): ---- When operands are integers, returns truncdiv(a, b, span). """ - return _binary_prim_expr_op(_ffi_api._OpDiv, a, b, span) # type: ignore + return _ffi_api._OpDiv(a, b, span) # type: ignore def indexdiv(a, b, span=None): @@ -2547,7 +2537,7 @@ def indexdiv(a, b, span=None): This function may take advantage of operands' non-negativeness. """ - return _binary_prim_expr_op(_ffi_api._OpIndexDiv, a, b, span) # type: ignore + return _ffi_api._OpIndexDiv(a, b, span) # type: ignore def indexmod(a, b, span=None): @@ -2575,7 +2565,7 @@ def indexmod(a, b, span=None): This function may take advantage of operands' non-negativeness. """ - return _binary_prim_expr_op(_ffi_api._OpIndexMod, a, b, span) # type: ignore + return _ffi_api._OpIndexMod(a, b, span) # type: ignore def truncdiv(a, b, span=None): @@ -2601,7 +2591,7 @@ def truncdiv(a, b, span=None): ---- This is the default integer division behavior in C. """ - return _binary_prim_expr_op(_ffi_api._OpTruncDiv, a, b, span) # type: ignore + return _ffi_api._OpTruncDiv(a, b, span) # type: ignore def truncmod(a, b, span=None): @@ -2627,7 +2617,7 @@ def truncmod(a, b, span=None): ---- This is the default integer division behavior in C. """ - return _binary_prim_expr_op(_ffi_api._OpTruncMod, a, b, span) # type: ignore + return _ffi_api._OpTruncMod(a, b, span) # type: ignore def floordiv(a, b, span=None): @@ -2649,7 +2639,7 @@ def floordiv(a, b, span=None): res : Expr The result expression. """ - return _binary_prim_expr_op(_ffi_api._OpFloorDiv, a, b, span) # type: ignore + return _ffi_api._OpFloorDiv(a, b, span) # type: ignore def logaddexp(a, b, span=None): @@ -2671,7 +2661,7 @@ def logaddexp(a, b, span=None): res : Expr The result expression. """ - return _binary_prim_expr_op(_ffi_api._OpLogAddExp, a, b, span) # type: ignore + return _ffi_api._OpLogAddExp(a, b, span) # type: ignore def floormod(a, b, span=None): @@ -2693,7 +2683,7 @@ def floormod(a, b, span=None): res : Expr The result expression. """ - return _binary_prim_expr_op(_ffi_api._OpFloorMod, a, b, span) # type: ignore + return _ffi_api._OpFloorMod(a, b, span) # type: ignore def ceildiv(lhs, rhs, span=None): @@ -2713,7 +2703,7 @@ def ceildiv(lhs, rhs, span=None): op : tvm.Expr The result Expr of ceildiv operaton. """ - return _binary_prim_expr_op(_ffi_api._OpCeilDiv, lhs, rhs, span) # type: ignore + return _ffi_api._OpCeilDiv(lhs, rhs, span) # type: ignore def comm_reducer(fcombine, fidentity, name="reduce"): diff --git a/python/tvm/tirx/script/builder/ir.py b/python/tvm/tirx/script/builder/ir.py index 6978d1774a55..d7897647fdbb 100644 --- a/python/tvm/tirx/script/builder/ir.py +++ b/python/tvm/tirx/script/builder/ir.py @@ -2376,7 +2376,6 @@ def buffer_store( expr_indices.append(index) if isinstance(value, bool) and buffer.ty.dtype == "bool": value = IntImm("bool", value) - value = _tir_op._convert_to_prim_expr(value) # pylint: disable=protected-access return _ffi_api.BufferStore( # type: ignore[attr-defined] # pylint: disable=no-member buffer, value, expr_indices, predicate ) diff --git a/python/tvm/tirx/script/parser/operation.py b/python/tvm/tirx/script/parser/operation.py index eacf15ed2e9d..169b3ced7c13 100644 --- a/python/tvm/tirx/script/parser/operation.py +++ b/python/tvm/tirx/script/parser/operation.py @@ -21,18 +21,15 @@ from tvm.ir import PrimType from tvm.runtime import DataTypeCode from tvm.script.parser._core import OpMethod, doc, register_op -from tvm.tirx import IntImm -from tvm.tirx.expr import FloatImm, _convert_to_prim_expr +from tvm.tirx import IntImm, _ffi_api +from tvm.tirx.expr import FloatImm def _register_expr_op(ty: type): # pylint: disable=invalid-name ty._dispatch_type = ty # pylint: disable=protected-access def _expr_ty(expr): - expr = _convert_to_prim_expr(expr) - ty = expr.ty if tvm.ir.is_prim_expr(expr) else None - if not isinstance(ty, PrimType): - ty = expr.expr_ty() + ty = _ffi_api._PrimExprType(expr) # type: ignore if not isinstance(ty, PrimType): raise TypeError(f"Expected a PrimType expression, but got {ty}") return ty diff --git a/src/ir/expr.cc b/src/ir/expr.cc index 5f80e20cbdf4..6374e2740ae0 100644 --- a/src/ir/expr.cc +++ b/src/ir/expr.cc @@ -108,6 +108,22 @@ PrimExpr PrimExpr::ConvertFallbackValue(ffi::String value) { return tirx::String namespace ffi { +std::optional TypeTraits::TryCastFromAnyView(const TVMFFIAny* src) { + if (auto value = Base::TryCastFromAnyView(src)) { + return value; + } + if (src->type_index < TypeIndex::kTVMFFIStaticObjectBegin || + !details::IsObjectInstance(src->type_index)) { + return std::nullopt; + } + Expr expr = details::ObjectUnsafe::ObjectRefFromObjectPtr( + details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); + if (const auto* type = expr->ty.as()) { + return type->ConvertToPrimExpr(std::move(expr)); + } + return std::nullopt; +} + PrimExpr TypeTraits::ConvertFallbackValue(StrictBool value) { return IntImm::Bool(value); } diff --git a/src/ir/type.cc b/src/ir/type.cc index ea4fa4e15675..c7a636f2ad7f 100644 --- a/src/ir/type.cc +++ b/src/ir/type.cc @@ -70,6 +70,7 @@ ffi::ObjectPtr GetCachedPrimTypeNode(DLDataType dtype) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; TypeNode::RegisterReflection(); + PrimExprConvertibleTypeNode::RegisterReflection(); OpaqueTypeNode::RegisterReflection(); PrimTypeNode::RegisterReflection(); refl::TypeAttrDef() diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index b491fd439ad2..aede24005618 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -511,12 +511,16 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // BufferRegion -BufferRegionType::BufferRegionType(Span span) : Type(ffi::UnsafeInit{}) { +BufferRegionType::BufferRegionType(Span span) : PrimExprConvertibleType(ffi::UnsafeInit{}) { ffi::ObjectPtr node = ffi::make_object(); node->span = std::move(span); data_ = std::move(node); } +PrimExpr BufferRegionTypeNode::ConvertToPrimExpr(Expr expr) const { + return std::move(expr).as_or_throw().ToBufferLoad(); +} + PrimExpr BufferRegion::ToBufferLoad() const { ffi::Array indices; indices.reserve((*this)->region.size()); @@ -569,9 +573,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { refl::GlobalDef() .def("tirx.BufferRegionType", [](Span span) { return BufferRegionType(span); }) .def("tirx.BufferRegion", - [](BufferVar buffer, ffi::Array region) { return BufferRegion(buffer, region); }) - .def("tirx.BufferRegionToBufferLoad", - [](BufferRegion region) { return region.ToBufferLoad(); }); + [](BufferVar buffer, ffi::Array region) { return BufferRegion(buffer, region); }); } // MatchBufferRegion diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc index 68230f419a53..2a4bcd879873 100644 --- a/src/tirx/op/op.cc +++ b/src/tirx/op/op.cc @@ -1297,6 +1297,7 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("tirx.trunc", tvm::trunc) .def("tirx._cast", [](PrimType dtype, PrimExpr value, Span span) { return tvm::cast(dtype, value, span); }) + .def("tirx._PrimExprType", [](PrimExpr value) { return value.ty(); }) .def("tirx.reinterpret", [](Type dtype, Expr value, Span span) { return tvm::reinterpret(dtype, value, span); }); } From 90f37085522e2aab32a46437780c705e12665de0 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 02:17:17 +0000 Subject: [PATCH 07/12] [REFACTOR][TIRX] Keep primitive-convertible type as marker --- include/tvm/ir/base_expr.h | 7 ++----- include/tvm/tirx/stmt.h | 2 -- src/ir/expr.cc | 8 ++++++-- src/tirx/ir/stmt.cc | 7 +++---- 4 files changed, 11 insertions(+), 13 deletions(-) diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h index a0f46309e117..ea4a713f3925 100644 --- a/include/tvm/ir/base_expr.h +++ b/include/tvm/ir/base_expr.h @@ -37,9 +37,6 @@ namespace tvm { -class Expr; -class PrimExpr; - /*! * \brief Type is the base type of all types. * @@ -93,8 +90,6 @@ class Type : public ffi::ObjectRef { */ class PrimExprConvertibleTypeNode : public TypeNode { public: - virtual PrimExpr ConvertToPrimExpr(Expr expr) const = 0; - static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); @@ -104,6 +99,8 @@ class PrimExprConvertibleTypeNode : public TypeNode { TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertibleType", PrimExprConvertibleTypeNode, TypeNode); }; +inline constexpr const char* kPrimExprConversionTypeAttr = "__tvm_ffi_to_prim_expr__"; + /*! \brief Managed reference to PrimExprConvertibleTypeNode. */ class PrimExprConvertibleType : public Type { public: diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index 7129f33ba12f..156fc82a699e 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -775,8 +775,6 @@ class Continue : public Stmt { */ class BufferRegionTypeNode : public PrimExprConvertibleTypeNode { public: - TVM_DLL PrimExpr ConvertToPrimExpr(Expr expr) const final; - static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); diff --git a/src/ir/expr.cc b/src/ir/expr.cc index 6374e2740ae0..08dadfa41c6b 100644 --- a/src/ir/expr.cc +++ b/src/ir/expr.cc @@ -118,8 +118,12 @@ std::optional TypeTraits::TryCastFromAnyView(const TVMFFIAny } Expr expr = details::ObjectUnsafe::ObjectRefFromObjectPtr( details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); - if (const auto* type = expr->ty.as()) { - return type->ConvertToPrimExpr(std::move(expr)); + if (!expr->ty.as()) { + return std::nullopt; + } + static const reflection::TypeAttrColumn converters(kPrimExprConversionTypeAttr); + if (auto converter = converters[src->type_index].try_cast()) { + return (*converter)(std::move(expr)).cast(); } return std::nullopt; } diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index aede24005618..3e28f6f3b369 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -35,6 +35,7 @@ namespace tvm { namespace tirx { TVM_FFI_STATIC_INIT_BLOCK() { + namespace refl = tvm::ffi::reflection; StmtNode::RegisterReflection(); BindNode::RegisterReflection(); @@ -53,6 +54,8 @@ TVM_FFI_STATIC_INIT_BLOCK() { ContinueNode::RegisterReflection(); BufferRegionTypeNode::RegisterReflection(); BufferRegionNode::RegisterReflection(); + refl::TypeAttrDef().def( + kPrimExprConversionTypeAttr, [](BufferRegion region) { return region.ToBufferLoad(); }); MatchBufferRegionNode::RegisterReflection(); SBlockNode::RegisterReflection(); SBlockRealizeNode::RegisterReflection(); @@ -517,10 +520,6 @@ BufferRegionType::BufferRegionType(Span span) : PrimExprConvertibleType(ffi::Uns data_ = std::move(node); } -PrimExpr BufferRegionTypeNode::ConvertToPrimExpr(Expr expr) const { - return std::move(expr).as_or_throw().ToBufferLoad(); -} - PrimExpr BufferRegion::ToBufferLoad() const { ffi::Array indices; indices.reserve((*this)->region.size()); From 60548191d93ea7ee3dd93ba17c4ddf097f4ca30c Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 02:54:49 +0000 Subject: [PATCH 08/12] [REFACTOR][TIRX] Reuse primitive-convertible expression base --- include/tvm/ir/base_expr.h | 38 +++++++++----------------------------- include/tvm/tirx/stmt.h | 23 ++++++++++------------- include/tvm/tirx/var.h | 10 +--------- python/tvm/ir/__init__.py | 2 +- python/tvm/ir/expr.py | 10 +++++++--- python/tvm/ir/type.py | 5 ----- python/tvm/tirx/stmt.py | 6 +++--- src/ir/expr.cc | 20 -------------------- src/ir/type.cc | 1 - src/tirx/ir/stmt.cc | 16 +++++++--------- 10 files changed, 38 insertions(+), 93 deletions(-) diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h index ea4a713f3925..f3900652b7c0 100644 --- a/include/tvm/ir/base_expr.h +++ b/include/tvm/ir/base_expr.h @@ -85,29 +85,6 @@ class Type : public ffi::ObjectRef { TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Type, ffi::ObjectRef, TypeNode); }; -/*! - * \brief Base type for expressions that can be converted to PrimExpr at typed FFI boundaries. - */ -class PrimExprConvertibleTypeNode : public TypeNode { - public: - static void RegisterReflection() { - namespace refl = tvm::ffi::reflection; - refl::ObjectDef(); - } - - static constexpr const uint32_t _type_child_slots = 1; - TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertibleType", PrimExprConvertibleTypeNode, TypeNode); -}; - -inline constexpr const char* kPrimExprConversionTypeAttr = "__tvm_ffi_to_prim_expr__"; - -/*! \brief Managed reference to PrimExprConvertibleTypeNode. */ -class PrimExprConvertibleType : public Type { - public: - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(PrimExprConvertibleType, Type, - PrimExprConvertibleTypeNode); -}; - /*! * \brief Type marker for opaque construction-time expressions. * @@ -458,21 +435,24 @@ class PrimExpr : public TypedExpr { * This is useful for the FFI to convert the expressions to PrimExpr. * \sa PrimExpr */ -class PrimExprConvertibleNode : public ffi::Object { +class PrimExprConvertibleNode : public ExprNode { public: virtual ~PrimExprConvertibleNode() {} virtual PrimExpr ToPrimExpr() const = 0; - TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertible", PrimExprConvertibleNode, ffi::Object); + static constexpr const uint32_t _type_child_slots = 2; + TVM_FFI_DECLARE_OBJECT_INFO("ir.PrimExprConvertible", PrimExprConvertibleNode, ExprNode); }; /*! * \brief Managed reference to PrimExprConvertibleNode. * \sa PrimExprConvertibleNode */ -class PrimExprConvertible : public ffi::ObjectRef { +class PrimExprConvertible : public Expr { public: - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExprConvertible, ffi::ObjectRef, - PrimExprConvertibleNode); + bool operator==(const PrimExprConvertible& other) const { return this->same_as(other); } + bool operator!=(const PrimExprConvertible& other) const { return !(*this == other); } + + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(PrimExprConvertible, Expr, PrimExprConvertibleNode); }; namespace ffi { @@ -563,7 +543,7 @@ struct TypeTraits using Base::TypeSchema; using Base::TypeStr; - TVM_DLL static std::optional TryCastFromAnyView(const TVMFFIAny* src); + using Base::TryCastFromAnyView; TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value); TVM_DLL static PrimExpr ConvertFallbackValue(int64_t value); diff --git a/include/tvm/tirx/stmt.h b/include/tvm/tirx/stmt.h index 156fc82a699e..0b0b5d5d4a28 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -773,32 +773,30 @@ class Continue : public Stmt { /*! * \brief The type of a multi-dimensional buffer region expression. */ -class BufferRegionTypeNode : public PrimExprConvertibleTypeNode { +class BufferRegionTypeNode : public TypeNode { public: static void RegisterReflection() { namespace refl = tvm::ffi::reflection; refl::ObjectDef(); } - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType", BufferRegionTypeNode, - PrimExprConvertibleTypeNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegionType", BufferRegionTypeNode, TypeNode); }; /*! * \brief Managed reference to BufferRegionTypeNode. */ -class BufferRegionType : public PrimExprConvertibleType { +class BufferRegionType : public Type { public: TVM_DLL BufferRegionType(Span span = Span()); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, PrimExprConvertibleType, - BufferRegionTypeNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(BufferRegionType, Type, BufferRegionTypeNode); }; /*! * \brief Representing the region of multi-dimensional buffer access. */ -class BufferRegionNode : public ExprNode { +class BufferRegionNode : public PrimExprConvertibleNode { public: /*! \brief The buffer of the buffer region. */ BufferVar buffer; @@ -812,21 +810,20 @@ class BufferRegionNode : public ExprNode { .def_ro("region", &BufferRegionNode::region); } + TVM_DLL PrimExpr ToPrimExpr() const final; + static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; - TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode, ExprNode); + TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.BufferRegion", BufferRegionNode, PrimExprConvertibleNode); }; /*! * \brief Managed reference to BufferRegionNode. * \sa BufferRegionNode */ -class BufferRegion : public Expr { +class BufferRegion : public PrimExprConvertible { public: TVM_DLL explicit BufferRegion(BufferVar buffer, ffi::Array region); - /*! \brief Materialize this region as a scalar or vector buffer load. */ - TVM_DLL PrimExpr ToBufferLoad() const; - /*! * \brief Create a BufferRegion which is full region of the given buffer. * \param buffer The buffer to generate full BufferRegion. @@ -842,7 +839,7 @@ class BufferRegion : public Expr { */ TVM_DLL static BufferRegion FromPoint(BufferVar buffer, ffi::Array indices); - TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, Expr, BufferRegionNode); + TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(BufferRegion, PrimExprConvertible, BufferRegionNode); TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferRegionNode); }; diff --git a/include/tvm/tirx/var.h b/include/tvm/tirx/var.h index 1886aaca51eb..6e9fac41e9f7 100644 --- a/include/tvm/tirx/var.h +++ b/include/tvm/tirx/var.h @@ -162,12 +162,6 @@ class IterVarNode : public PrimExprConvertibleNode { * set this if this is bound already to a known thread tag. */ ffi::String thread_tag; - /*! - * \brief Span that points to the original source code. - * Reserved debug information. - */ - mutable Span span; - PrimExpr ToPrimExpr() const final { return var; } static void RegisterReflection() { @@ -176,9 +170,7 @@ class IterVarNode : public PrimExprConvertibleNode { .def_ro("dom", &IterVarNode::dom) .def_ro("var", &IterVarNode::var, refl::AttachFieldFlag::SEqHashDefRecursive()) .def_ro("iter_type", &IterVarNode::iter_type) - .def_ro("thread_tag", &IterVarNode::thread_tag) - .def_ro("span", &IterVarNode::span, refl::DefaultValue(Span()), - refl::AttachFieldFlag::SEqHashIgnore()); + .def_ro("thread_tag", &IterVarNode::thread_tag); } static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode; diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py index 729c90465f40..a8b6d6bcc433 100644 --- a/python/tvm/ir/__init__.py +++ b/python/tvm/ir/__init__.py @@ -37,7 +37,6 @@ FuncType, OpaqueType, PointerType, - PrimExprConvertibleType, PrimType, TupleType, Type, @@ -48,6 +47,7 @@ ExprWithOp, GlobalVar, OpaqueExpr, + PrimExprConvertible, Range, Tuple, TupleGetItem, diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py index 4ba3e478faca..faf25c24cf0e 100644 --- a/python/tvm/ir/expr.py +++ b/python/tvm/ir/expr.py @@ -25,7 +25,6 @@ from ..runtime import Object, Scriptable from . import _ffi_api, _overload_prim_expr, _tensor_expr_overload from .base import Node, Span -from .type import PrimExprConvertibleType @tvm_ffi.register_object("ir.Expr") @@ -52,8 +51,8 @@ def is_prim_var(value: object) -> bool: def is_prim_expr_convertible(value: object) -> bool: - """Return whether an expression's type opts into primitive operators.""" - return isinstance(value, Expr) and isinstance(value.ty, PrimExprConvertibleType) + """Return whether an expression supports conversion at primitive FFI boundaries.""" + return isinstance(value, PrimExprConvertible) def _supports_prim_expr_ops(value: object) -> bool: @@ -341,6 +340,11 @@ def __getitem__(self, index): return result +@tvm_ffi.register_object("ir.PrimExprConvertible") +class PrimExprConvertible(ExprWithOp): + """Expression that converts to PrimExpr at typed FFI boundaries.""" + + @tvm_ffi.register_object("ir.Tuple") class Tuple(ExprWithOp): """Tuple expression that groups several fields together. diff --git a/python/tvm/ir/type.py b/python/tvm/ir/type.py index c4d0c29c18ab..015232963e1b 100644 --- a/python/tvm/ir/type.py +++ b/python/tvm/ir/type.py @@ -54,11 +54,6 @@ def same_as(self, other): return self.is_(other) -@tvm_ffi.register_object("ir.PrimExprConvertibleType") -class PrimExprConvertibleType(Type): - """Marker for non-primitive expressions accepted by primitive operators.""" - - @tvm_ffi.register_object("ir.OpaqueType") class OpaqueType(Type): """Type marker for opaque values that must be removed from finished IR.""" diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index 6f0f21871029..5a55544c61bd 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -33,7 +33,7 @@ import tvm_ffi -from tvm.ir import Expr, ExprWithOp, PrimExprConvertibleType, Range, Span, is_prim_expr +from tvm.ir import Expr, PrimExprConvertible, Range, Span, Type, is_prim_expr from tvm.runtime import Object, Scriptable, const from tvm.tirx import IntImm @@ -616,7 +616,7 @@ def __init__(self, value: Expr, span: Span | None = None) -> None: @tvm_ffi.register_object("tirx.BufferRegionType") -class BufferRegionType(PrimExprConvertibleType): +class BufferRegionType(Type): """The structural type of a :class:`BufferRegion` expression.""" def __init__(self, span: Span | None = None) -> None: @@ -624,7 +624,7 @@ def __init__(self, span: Span | None = None) -> None: @tvm_ffi.register_object("tirx.BufferRegion") -class BufferRegion(ExprWithOp): +class BufferRegion(PrimExprConvertible): """BufferRegion node. Parameters diff --git a/src/ir/expr.cc b/src/ir/expr.cc index 08dadfa41c6b..5f80e20cbdf4 100644 --- a/src/ir/expr.cc +++ b/src/ir/expr.cc @@ -108,26 +108,6 @@ PrimExpr PrimExpr::ConvertFallbackValue(ffi::String value) { return tirx::String namespace ffi { -std::optional TypeTraits::TryCastFromAnyView(const TVMFFIAny* src) { - if (auto value = Base::TryCastFromAnyView(src)) { - return value; - } - if (src->type_index < TypeIndex::kTVMFFIStaticObjectBegin || - !details::IsObjectInstance(src->type_index)) { - return std::nullopt; - } - Expr expr = details::ObjectUnsafe::ObjectRefFromObjectPtr( - details::ObjectUnsafe::ObjectPtrFromUnowned(src->v_obj)); - if (!expr->ty.as()) { - return std::nullopt; - } - static const reflection::TypeAttrColumn converters(kPrimExprConversionTypeAttr); - if (auto converter = converters[src->type_index].try_cast()) { - return (*converter)(std::move(expr)).cast(); - } - return std::nullopt; -} - PrimExpr TypeTraits::ConvertFallbackValue(StrictBool value) { return IntImm::Bool(value); } diff --git a/src/ir/type.cc b/src/ir/type.cc index c7a636f2ad7f..ea4fa4e15675 100644 --- a/src/ir/type.cc +++ b/src/ir/type.cc @@ -70,7 +70,6 @@ ffi::ObjectPtr GetCachedPrimTypeNode(DLDataType dtype) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; TypeNode::RegisterReflection(); - PrimExprConvertibleTypeNode::RegisterReflection(); OpaqueTypeNode::RegisterReflection(); PrimTypeNode::RegisterReflection(); refl::TypeAttrDef() diff --git a/src/tirx/ir/stmt.cc b/src/tirx/ir/stmt.cc index 3e28f6f3b369..d4a0478a9050 100644 --- a/src/tirx/ir/stmt.cc +++ b/src/tirx/ir/stmt.cc @@ -35,7 +35,6 @@ namespace tvm { namespace tirx { TVM_FFI_STATIC_INIT_BLOCK() { - namespace refl = tvm::ffi::reflection; StmtNode::RegisterReflection(); BindNode::RegisterReflection(); @@ -54,8 +53,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { ContinueNode::RegisterReflection(); BufferRegionTypeNode::RegisterReflection(); BufferRegionNode::RegisterReflection(); - refl::TypeAttrDef().def( - kPrimExprConversionTypeAttr, [](BufferRegion region) { return region.ToBufferLoad(); }); MatchBufferRegionNode::RegisterReflection(); SBlockNode::RegisterReflection(); SBlockRealizeNode::RegisterReflection(); @@ -514,25 +511,26 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // BufferRegion -BufferRegionType::BufferRegionType(Span span) : PrimExprConvertibleType(ffi::UnsafeInit{}) { +BufferRegionType::BufferRegionType(Span span) : Type(ffi::UnsafeInit{}) { ffi::ObjectPtr node = ffi::make_object(); node->span = std::move(span); data_ = std::move(node); } -PrimExpr BufferRegion::ToBufferLoad() const { +PrimExpr BufferRegionNode::ToPrimExpr() const { ffi::Array indices; - indices.reserve((*this)->region.size()); - for (const Range& r : (*this)->region) { + indices.reserve(this->region.size()); + for (const Range& r : this->region) { if (tirx::is_one(r->extent)) { indices.push_back(r->min); } else if (r->extent.as()) { indices.push_back(Ramp(r->min, IntImm(r->min.ty(), 1), r->extent)); } else { - TVM_FFI_THROW(ValueError) << "Cannot convert to BufferLoad: " << *this; + TVM_FFI_THROW(ValueError) << "Cannot convert to BufferLoad: " + << ffi::GetRef(this); } } - return BufferLoad((*this)->buffer, indices); + return BufferLoad(this->buffer, indices); } BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region) { From 7288510922dbaf2c96c87bfcf60ad54ebafe98bc Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 04:10:48 +0000 Subject: [PATCH 09/12] [FIX][TIRX] Convert primitive intrinsic operands at FFI --- python/tvm/tirx/op.py | 10 ++++++++++ src/tirx/op/op.cc | 18 ++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 199f44c09480..a87dd4e30abd 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -255,6 +255,14 @@ def call_intrin(dtype: str | tvm.ir.Type, func_name, *args, attrs=None, span=Non """ if isinstance(func_name, str): func_name = _canonical_device_intrin_name(func_name) + if any(tvm.ir.is_prim_expr_convertible(arg) for arg in args): + if isinstance(func_name, str): + func_name = Op.get(func_name) + if not isinstance(dtype, tvm.ir.Type): + dtype = PointerType(PrimType("void")) if dtype == "handle" else PrimType(dtype) + if attrs is not None and isinstance(attrs, dict): + attrs = tvm.ir.DictAttrs(attrs) + return _ffi_api._CallPrimExpr(dtype, func_name, args, attrs, span) # type: ignore return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype) @@ -1377,6 +1385,8 @@ def reinterpret(dtype, value, span: Span | None = None) -> Expr: dtype = ( PointerType(tvm.ir.PrimType("void")) if dtype == "handle" else tvm.ir.PrimType(dtype) ) + if tvm.ir.is_prim_expr_convertible(value): + return _ffi_api._reinterpret_prim(dtype, value, span) # type: ignore return _ffi_api.reinterpret(dtype, value, span) # type: ignore diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc index 2a4bcd879873..1633abdd8309 100644 --- a/src/tirx/op/op.cc +++ b/src/tirx/op/op.cc @@ -1298,6 +1298,24 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("tirx._cast", [](PrimType dtype, PrimExpr value, Span span) { return tvm::cast(dtype, value, span); }) .def("tirx._PrimExprType", [](PrimExpr value) { return value.ty(); }) + .def("tirx._CallPrimExpr", + [](Type ret_ty, Expr op, ffi::Array args, Attrs attrs, Span span) { + ffi::Array expr_args; + expr_args.reserve(args.size()); + for (const ffi::Any& arg : args) { + if (auto prim_arg = arg.try_cast()) { + expr_args.push_back(std::move(prim_arg.value())); + } else { + expr_args.push_back(arg.cast()); + } + } + return Call(std::move(ret_ty), std::move(op), std::move(expr_args), std::move(attrs), + {}, std::move(span)); + }) + .def("tirx._reinterpret_prim", + [](PrimType dtype, PrimExpr value, Span span) { + return tvm::reinterpret(std::move(dtype), std::move(value), std::move(span)); + }) .def("tirx.reinterpret", [](Type dtype, Expr value, Span span) { return tvm::reinterpret(dtype, value, span); }); } From 1f0546118a60cc34d149bc42003345f904aaaad9 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 13:43:13 +0000 Subject: [PATCH 10/12] [REFACTOR][TIRX] Make primitive conversion explicit --- include/tvm/ir/base_expr.h | 3 +- python/tvm/ir/__init__.py | 11 +----- python/tvm/ir/expr.py | 23 ++++++------- python/tvm/tirx/expr.py | 19 ++++++---- python/tvm/tirx/op.py | 40 ++++++++-------------- python/tvm/tirx/script/parser/operation.py | 6 ++-- src/ir/expr.cc | 2 ++ src/tirx/op/op.cc | 19 ---------- 8 files changed, 46 insertions(+), 77 deletions(-) diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h index f3900652b7c0..d3698d81a814 100644 --- a/include/tvm/ir/base_expr.h +++ b/include/tvm/ir/base_expr.h @@ -540,11 +540,10 @@ struct TypeTraits using Base::GetMismatchTypeInfo; using Base::MoveFromAnyAfterCheck; using Base::MoveToAny; + using Base::TryCastFromAnyView; using Base::TypeSchema; using Base::TypeStr; - using Base::TryCastFromAnyView; - TVM_DLL static PrimExpr ConvertFallbackValue(StrictBool value); TVM_DLL static PrimExpr ConvertFallbackValue(int64_t value); TVM_DLL static PrimExpr ConvertFallbackValue(double value); diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py index a8b6d6bcc433..dd7d897ef3c9 100644 --- a/python/tvm/ir/__init__.py +++ b/python/tvm/ir/__init__.py @@ -33,18 +33,10 @@ # Register Type before Expr. Expr's reflected ``ty`` field otherwise creates # an auto-generated Type wrapper before the concrete Python class is available. -from .type import ( - FuncType, - OpaqueType, - PointerType, - PrimType, - TupleType, - Type, -) +from .type import FuncType, OpaqueType, PointerType, PrimType, TupleType, Type from .expr import ( Call, Expr, - ExprWithOp, GlobalVar, OpaqueExpr, PrimExprConvertible, @@ -53,7 +45,6 @@ TupleGetItem, Var, is_prim_expr, - is_prim_expr_convertible, is_prim_var, ) from .function import BaseFunc, CallingConv diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py index faf25c24cf0e..a985e6d32b64 100644 --- a/python/tvm/ir/expr.py +++ b/python/tvm/ir/expr.py @@ -50,13 +50,8 @@ def is_prim_var(value: object) -> bool: return isinstance(value, Var) and type(value) is Var and is_prim_expr(value) -def is_prim_expr_convertible(value: object) -> bool: - """Return whether an expression supports conversion at primitive FFI boundaries.""" - return isinstance(value, PrimExprConvertible) - - def _supports_prim_expr_ops(value: object) -> bool: - return is_prim_expr(value) or is_prim_expr_convertible(value) + return is_prim_expr(value) or isinstance(value, PrimExprConvertible) @tvm_ffi.register_object("ir.GlobalVar") @@ -109,7 +104,7 @@ def is_tir_arg(x): raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}") -class ExprWithOp(Expr, Scriptable): +class _ExprWithOp(Expr, Scriptable): """Common type-directed operator behavior for core expressions.""" __hash__ = Expr.__hash__ @@ -341,12 +336,16 @@ def __getitem__(self, index): @tvm_ffi.register_object("ir.PrimExprConvertible") -class PrimExprConvertible(ExprWithOp): +class PrimExprConvertible(_ExprWithOp): """Expression that converts to PrimExpr at typed FFI boundaries.""" + def to_prim_expr(self): + """Convert this expression to its primitive representation.""" + return _ffi_api.PrimExprConvertibleToPrimExpr(self) + @tvm_ffi.register_object("ir.Tuple") -class Tuple(ExprWithOp): +class Tuple(_ExprWithOp): """Tuple expression that groups several fields together. Parameters @@ -379,7 +378,7 @@ def __len__(self) -> int: @tvm_ffi.register_object("ir.TupleGetItem") -class TupleGetItem(ExprWithOp): +class TupleGetItem(_ExprWithOp): """Get the index-th item from a tuple. Parameters @@ -403,7 +402,7 @@ def __init__(self, tuple_value: Expr, index: int, span: Span | None = None): @tvm_ffi.register_object("ir.Call") -class Call(ExprWithOp): +class Call(_ExprWithOp): """Core function call node.""" op: Expr @@ -442,7 +441,7 @@ def __init__( @tvm_ffi.register_object("ir.Var") -class Var(ExprWithOp): +class Var(_ExprWithOp): """A canonical local variable in the IR. Parameters diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index 7533e1d2b901..3d9ebaa4eb62 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -56,16 +56,20 @@ def div_ambiguity_error() -> RuntimeError: def _dtype_is_int(value): if isinstance(value, int): return True - if isinstance(value, ExprOp) or ir.is_prim_expr(value) or ir.is_prim_expr_convertible(value): - return _ffi_api._PrimExprType(value).matches_code(DataTypeCode.INT) # type: ignore + if isinstance(value, ir.PrimExprConvertible): + value = value.to_prim_expr() + if isinstance(value, ExprOp) or ir.is_prim_expr(value): + return value.expr_ty().matches_code(DataTypeCode.INT) return False def _dtype_is_float(value): if isinstance(value, float): return True - if isinstance(value, ExprOp) or ir.is_prim_expr(value) or ir.is_prim_expr_convertible(value): - return _ffi_api._PrimExprType(value).matches_code(DataTypeCode.FLOAT) # type: ignore + if isinstance(value, ir.PrimExprConvertible): + value = value.to_prim_expr() + if isinstance(value, ExprOp) or ir.is_prim_expr(value): + return value.expr_ty().matches_code(DataTypeCode.FLOAT) return False @@ -73,7 +77,7 @@ def _is_scalar_operand(value): return ( isinstance(value, ExprOp | int | float) or ir.is_prim_expr(value) - or ir.is_prim_expr_convertible(value) + or isinstance(value, ir.PrimExprConvertible) ) @@ -160,8 +164,9 @@ def __rmod__(self, other: Expr) -> Expr: return _ffi_api._OpFloorMod(other, self, None) # type: ignore def __neg__(self) -> Expr: - neg_one = const(-1, _ffi_api._PrimExprType(self).dtype) # type: ignore - return _ffi_api._OpMul(self, neg_one, None) # type: ignore + value = self.to_prim_expr() if isinstance(self, ir.PrimExprConvertible) else self + neg_one = const(-1, value.expr_ty().dtype) + return _ffi_api._OpMul(value, neg_one, None) # type: ignore def __lshift__(self, other: Expr) -> Expr: return _ffi_api.left_shift(self, other, None) # type: ignore diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index a87dd4e30abd..9b1cccee0570 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -31,17 +31,10 @@ from . import _ffi_api from .buffer import Buffer, buffer_data, is_buffer_var -from .expr import ( - BufferLoad, - CommReducer, - ExprWithOp, - IntImm, - Var, -) +from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var tir = tirx # alias for backward compat with upstream tir.convert() calls - # Insertion order matters: a longer prefix has to be tried before the shorter # one it starts with, or `ptx_legacy_mma` would strip as `ptx` + `legacy_mma`. _DEVICE_INTRIN_PREFIX_TO_NAMESPACE = { @@ -72,7 +65,14 @@ def _primexpr_ty(expr): """Return the runtime primitive type of an expression.""" if isinstance(expr, tvm.ir.PrimType): return expr - return _ffi_api._PrimExprType(expr) # type: ignore + if isinstance(expr, tvm.ir.PrimExprConvertible): + expr = expr.to_prim_expr() + ty = getattr(expr, "ty", None) + if isinstance(ty, tvm.ir.PrimType): + return ty + if isinstance(expr, ExprOp): + return expr.expr_ty() + raise TypeError(f"Cannot determine primitive expression type for {type(expr).__name__}") def _primexpr_dtype(expr): @@ -255,14 +255,9 @@ def call_intrin(dtype: str | tvm.ir.Type, func_name, *args, attrs=None, span=Non """ if isinstance(func_name, str): func_name = _canonical_device_intrin_name(func_name) - if any(tvm.ir.is_prim_expr_convertible(arg) for arg in args): - if isinstance(func_name, str): - func_name = Op.get(func_name) - if not isinstance(dtype, tvm.ir.Type): - dtype = PointerType(PrimType("void")) if dtype == "handle" else PrimType(dtype) - if attrs is not None and isinstance(attrs, dict): - attrs = tvm.ir.DictAttrs(attrs) - return _ffi_api._CallPrimExpr(dtype, func_name, args, attrs, span) # type: ignore + args = [ + arg.to_prim_expr() if isinstance(arg, tvm.ir.PrimExprConvertible) else arg for arg in args + ] return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype) @@ -1385,8 +1380,8 @@ def reinterpret(dtype, value, span: Span | None = None) -> Expr: dtype = ( PointerType(tvm.ir.PrimType("void")) if dtype == "handle" else tvm.ir.PrimType(dtype) ) - if tvm.ir.is_prim_expr_convertible(value): - return _ffi_api._reinterpret_prim(dtype, value, span) # type: ignore + if isinstance(value, tvm.ir.PrimExprConvertible): + value = value.to_prim_expr() return _ffi_api.reinterpret(dtype, value, span) # type: ignore @@ -2489,12 +2484,7 @@ def if_then_else(cond, t, f, span=None): Unlike Select, if_then_else cannot be vectorized if some lanes in the vector have different conditions. """ - return _ffi_api._OpIfThenElse( - cond, - t, - f, - span, - ) # type: ignore + return _ffi_api._OpIfThenElse(cond, t, f, span) # type: ignore def div(a, b, span=None): diff --git a/python/tvm/tirx/script/parser/operation.py b/python/tvm/tirx/script/parser/operation.py index 169b3ced7c13..3cce1ad13360 100644 --- a/python/tvm/tirx/script/parser/operation.py +++ b/python/tvm/tirx/script/parser/operation.py @@ -21,7 +21,7 @@ from tvm.ir import PrimType from tvm.runtime import DataTypeCode from tvm.script.parser._core import OpMethod, doc, register_op -from tvm.tirx import IntImm, _ffi_api +from tvm.tirx import IntImm from tvm.tirx.expr import FloatImm @@ -29,7 +29,9 @@ def _register_expr_op(ty: type): # pylint: disable=invalid-name ty._dispatch_type = ty # pylint: disable=protected-access def _expr_ty(expr): - ty = _ffi_api._PrimExprType(expr) # type: ignore + if isinstance(expr, tvm.ir.PrimExprConvertible): + expr = expr.to_prim_expr() + ty = expr.ty if tvm.ir.is_prim_expr(expr) else expr.expr_ty() if not isinstance(ty, PrimType): raise TypeError(f"Expected a PrimType expression, but got {ty}") return ty diff --git a/src/ir/expr.cc b/src/ir/expr.cc index 5f80e20cbdf4..3e12170abdf8 100644 --- a/src/ir/expr.cc +++ b/src/ir/expr.cc @@ -90,6 +90,8 @@ TupleGetItem::TupleGetItem(Expr tuple, int index, Span span) { TVM_FFI_STATIC_INIT_BLOCK() { namespace refl = tvm::ffi::reflection; refl::GlobalDef() + .def("ir.PrimExprConvertibleToPrimExpr", + [](PrimExprConvertible value) { return value->ToPrimExpr(); }) .def("ir.Tuple", [](ffi::Array fields, Span span) { return Tuple(fields, span); }) .def("ir.TupleGetItem", [](Expr tuple, int index, Span span) { return TupleGetItem(tuple, index, span); }) diff --git a/src/tirx/op/op.cc b/src/tirx/op/op.cc index 1633abdd8309..68230f419a53 100644 --- a/src/tirx/op/op.cc +++ b/src/tirx/op/op.cc @@ -1297,25 +1297,6 @@ TVM_FFI_STATIC_INIT_BLOCK() { .def("tirx.trunc", tvm::trunc) .def("tirx._cast", [](PrimType dtype, PrimExpr value, Span span) { return tvm::cast(dtype, value, span); }) - .def("tirx._PrimExprType", [](PrimExpr value) { return value.ty(); }) - .def("tirx._CallPrimExpr", - [](Type ret_ty, Expr op, ffi::Array args, Attrs attrs, Span span) { - ffi::Array expr_args; - expr_args.reserve(args.size()); - for (const ffi::Any& arg : args) { - if (auto prim_arg = arg.try_cast()) { - expr_args.push_back(std::move(prim_arg.value())); - } else { - expr_args.push_back(arg.cast()); - } - } - return Call(std::move(ret_ty), std::move(op), std::move(expr_args), std::move(attrs), - {}, std::move(span)); - }) - .def("tirx._reinterpret_prim", - [](PrimType dtype, PrimExpr value, Span span) { - return tvm::reinterpret(std::move(dtype), std::move(value), std::move(span)); - }) .def("tirx.reinterpret", [](Type dtype, Expr value, Span span) { return tvm::reinterpret(dtype, value, span); }); } From 5b2b19764726993e9d93d1089e7e17bb8f2b5f0c Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 15:31:27 +0000 Subject: [PATCH 11/12] [REFACTOR][TIRX] Keep type queries shallow --- python/tvm/tirx/expr.py | 16 ++++++---------- python/tvm/tirx/op.py | 2 -- python/tvm/tirx/script/parser/operation.py | 6 +++--- 3 files changed, 9 insertions(+), 15 deletions(-) diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index 3d9ebaa4eb62..c824661bd77e 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -35,7 +35,7 @@ from tvm import ir from tvm.ir import Expr from tvm.ir.base import Span -from tvm.runtime import DataTypeCode, Object, ObjectConvertible, Scriptable, const +from tvm.runtime import DataTypeCode, Object, ObjectConvertible, Scriptable from . import _ffi_api from .buffer import Buffer @@ -56,20 +56,18 @@ def div_ambiguity_error() -> RuntimeError: def _dtype_is_int(value): if isinstance(value, int): return True - if isinstance(value, ir.PrimExprConvertible): - value = value.to_prim_expr() if isinstance(value, ExprOp) or ir.is_prim_expr(value): - return value.expr_ty().matches_code(DataTypeCode.INT) + ty = value.expr_ty() + return isinstance(ty, ir.PrimType) and ty.matches_code(DataTypeCode.INT) return False def _dtype_is_float(value): if isinstance(value, float): return True - if isinstance(value, ir.PrimExprConvertible): - value = value.to_prim_expr() if isinstance(value, ExprOp) or ir.is_prim_expr(value): - return value.expr_ty().matches_code(DataTypeCode.FLOAT) + ty = value.expr_ty() + return isinstance(ty, ir.PrimType) and ty.matches_code(DataTypeCode.FLOAT) return False @@ -164,9 +162,7 @@ def __rmod__(self, other: Expr) -> Expr: return _ffi_api._OpFloorMod(other, self, None) # type: ignore def __neg__(self) -> Expr: - value = self.to_prim_expr() if isinstance(self, ir.PrimExprConvertible) else self - neg_one = const(-1, value.expr_ty().dtype) - return _ffi_api._OpMul(value, neg_one, None) # type: ignore + return _ffi_api._OpMul(self, -1, None) # type: ignore def __lshift__(self, other: Expr) -> Expr: return _ffi_api.left_shift(self, other, None) # type: ignore diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 9b1cccee0570..610087dd8c95 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -65,8 +65,6 @@ def _primexpr_ty(expr): """Return the runtime primitive type of an expression.""" if isinstance(expr, tvm.ir.PrimType): return expr - if isinstance(expr, tvm.ir.PrimExprConvertible): - expr = expr.to_prim_expr() ty = getattr(expr, "ty", None) if isinstance(ty, tvm.ir.PrimType): return ty diff --git a/python/tvm/tirx/script/parser/operation.py b/python/tvm/tirx/script/parser/operation.py index 3cce1ad13360..fd67d6f12591 100644 --- a/python/tvm/tirx/script/parser/operation.py +++ b/python/tvm/tirx/script/parser/operation.py @@ -29,9 +29,9 @@ def _register_expr_op(ty: type): # pylint: disable=invalid-name ty._dispatch_type = ty # pylint: disable=protected-access def _expr_ty(expr): - if isinstance(expr, tvm.ir.PrimExprConvertible): - expr = expr.to_prim_expr() - ty = expr.ty if tvm.ir.is_prim_expr(expr) else expr.expr_ty() + ty = expr.ty if tvm.ir.is_prim_expr(expr) else None + if not isinstance(ty, PrimType): + ty = expr.expr_ty() if not isinstance(ty, PrimType): raise TypeError(f"Expected a PrimType expression, but got {ty}") return ty From 4c7fdaa0c2ef397451fd1a243a7feefdcd7c69fc Mon Sep 17 00:00:00 2001 From: tqchen Date: Sun, 30 Aug 2026 16:40:43 +0000 Subject: [PATCH 12/12] [REFACTOR][TIRX] Localize primitive expression conversion --- python/tvm/s_tir/tensor_intrin/rocm.py | 4 ++-- python/tvm/tirx/op.py | 9 ++++++--- tests/python/tirx/test_op.py | 9 --------- 3 files changed, 8 insertions(+), 14 deletions(-) diff --git a/python/tvm/s_tir/tensor_intrin/rocm.py b/python/tvm/s_tir/tensor_intrin/rocm.py index 8573c45304da..5a5e667e214a 100644 --- a/python/tvm/s_tir/tensor_intrin/rocm.py +++ b/python/tvm/s_tir/tensor_intrin/rocm.py @@ -363,8 +363,8 @@ def mfma_sync_impl_integer(a: T.handle, b: T.handle, c: T.handle) -> None: C[tx, 0:local_size_out] = T.call_llvm_pure_intrin( T.llvm_lookup_intrinsic_id(mfma_intrin), - T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size]), - T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size]), + T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size].to_prim_expr()), + T.call_intrin("int32", "tirx.reinterpret", A[tx, 0:local_size].to_prim_expr()), C[tx, 0:local_size_out], T.int32(0), T.int32(0), diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py index 610087dd8c95..416574b38b99 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -253,9 +253,6 @@ def call_intrin(dtype: str | tvm.ir.Type, func_name, *args, attrs=None, span=Non """ if isinstance(func_name, str): func_name = _canonical_device_intrin_name(func_name) - args = [ - arg.to_prim_expr() if isinstance(arg, tvm.ir.PrimExprConvertible) else arg for arg in args - ] return Call(func_name, args, attrs=attrs, span=span, ret_ty=dtype) @@ -360,6 +357,9 @@ def call_llvm_intrin(dtype, name, *args, span=None): llvm_id = name if llvm_id == 0: raise ValueError(f"Unknown llvm intrinsic function {name}") + args = tuple( + arg.to_prim_expr() if isinstance(arg, tvm.ir.PrimExprConvertible) else arg for arg in args + ) return call_intrin( dtype, Op.get("tirx.call_llvm_intrin"), @@ -402,6 +402,9 @@ def call_llvm_pure_intrin(dtype, name, *args, span=None): llvm_id = name if llvm_id == 0: raise ValueError(f"Unknown llvm intrinsic function {name}") + args = tuple( + arg.to_prim_expr() if isinstance(arg, tvm.ir.PrimExprConvertible) else arg for arg in args + ) return call_intrin( dtype, Op.get("tirx.call_llvm_pure_intrin"), diff --git a/tests/python/tirx/test_op.py b/tests/python/tirx/test_op.py index 6ee7af745ade..8d6e6326e189 100644 --- a/tests/python/tirx/test_op.py +++ b/tests/python/tirx/test_op.py @@ -73,15 +73,6 @@ def test_tile_primitive_call_pickle_roundtrip(): assert_structural_equal(restored.scope, call.scope) -def test_compose_op_retains_statement_arguments(): - buffer = decl_buffer((16,), "float32", scope="local") - inner = _test("fill", buffer[:], 1.0) - composed = _test("compose_op", inner) - - assert isinstance(composed.args[0], TilePrimitiveCall) - assert composed.args[0].same_as(inner) - - def test_buffer_replacer_no_shared_default(): """Regression test for F4: BufferReplacer default dicts must not be shared.""" from tvm.tirx.transform.common import BufferReplacer