diff --git a/include/tvm/ir/base_expr.h b/include/tvm/ir/base_expr.h index 1e39240a7f16..d3698d81a814 100644 --- a/include/tvm/ir/base_expr.h +++ b/include/tvm/ir/base_expr.h @@ -435,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 { 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..0b0b5d5d4a28 100644 --- a/include/tvm/tirx/stmt.h +++ b/include/tvm/tirx/stmt.h @@ -770,6 +770,29 @@ 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. */ 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/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 ca83d21f3aac..dd7d897ef3c9 100644 --- a/python/tvm/ir/__init__.py +++ b/python/tvm/ir/__init__.py @@ -39,6 +39,7 @@ Expr, GlobalVar, OpaqueExpr, + PrimExprConvertible, Range, Tuple, TupleGetItem, diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py index 393e825646b3..a985e6d32b64 100644 --- a/python/tvm/ir/expr.py +++ b/python/tvm/ir/expr.py @@ -50,6 +50,10 @@ def is_prim_var(value: object) -> bool: return isinstance(value, Var) and type(value) is Var and is_prim_expr(value) +def _supports_prim_expr_ops(value: object) -> bool: + return is_prim_expr(value) or isinstance(value, PrimExprConvertible) + + @tvm_ffi.register_object("ir.GlobalVar") class GlobalVar(Expr): """A global variable in the IR. @@ -106,93 +110,91 @@ class _ExprWithOp(Expr, Scriptable): __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 +205,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 +263,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 +283,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 +308,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 +319,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: @@ -333,6 +335,15 @@ def __getitem__(self, index): return result +@tvm_ffi.register_object("ir.PrimExprConvertible") +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): """Tuple expression that groups several fields together. 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/__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..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,31 +56,27 @@ def div_ambiguity_error() -> RuntimeError: def _dtype_is_int(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): + 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, 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): + ty = value.expr_ty() + return isinstance(ty, ir.PrimType) and ty.matches_code(DataTypeCode.FLOAT) return False 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) + or isinstance(value, ir.PrimExprConvertible) + ) class ExprOp: @@ -88,12 +84,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.""" + def expr_ty(self) -> ir.Type: + """Return the expression's compile-time type.""" ty = getattr(self, "ty", None) - if isinstance(ty, ir.PrimType): + 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): @@ -166,8 +162,7 @@ def __rmod__(self, other: Expr) -> Expr: return _ffi_api._OpFloorMod(other, self, None) # type: ignore def __neg__(self) -> Expr: - neg_one = const(-1, self.expr_ty().dtype) - return self.__mul__(neg_one) + 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/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/op.py b/python/tvm/tirx/op.py index 39939803a667..416574b38b99 100644 --- a/python/tvm/tirx/op.py +++ b/python/tvm/tirx/op.py @@ -70,7 +70,7 @@ def _primexpr_ty(expr): return ty if isinstance(expr, ExprOp): return expr.expr_ty() - raise TypeError(f"Cannot determine Expr type for {type(expr).__name__}") + raise TypeError(f"Cannot determine primitive expression type for {type(expr).__name__}") def _primexpr_dtype(expr): @@ -357,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"), @@ -399,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"), @@ -1375,6 +1381,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 isinstance(value, tvm.ir.PrimExprConvertible): + value = value.to_prim_expr() return _ffi_api.reinterpret(dtype, value, span) # type: ignore diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index fe4c98a4025a..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, 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 @@ -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(PrimExprConvertible): """BufferRegion node. Parameters 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/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..d4a0478a9050 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,26 @@ TVM_FFI_STATIC_INIT_BLOCK() { } // BufferRegion +BufferRegionType::BufferRegionType(Span span) : Type(ffi::UnsafeInit{}) { + ffi::ObjectPtr node = ffi::make_object(); + node->span = std::move(span); + data_ = std::move(node); +} + 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)) { + if (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)); + indices.push_back(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); + return BufferLoad(this->buffer, indices); } BufferRegion::BufferRegion(BufferVar buffer, ffi::Array region) { @@ -532,6 +538,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 +567,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/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"""