This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new bbd115877d [REFACTOR][IR] Consolidate expression operator overloading
into the base layer (#20248)
bbd115877d is described below
commit bbd115877dc9faf6bdc81b8d04332e468022ad35
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Sep 1 07:43:39 2026 -0400
[REFACTOR][IR] Consolidate expression operator overloading into the base
layer (#20248)
## Rationale
Python expression operators should have one shared base-layer protocol
instead of parallel TIRx and Relax mixins. A single surface keeps proxy
realization and type-directed dispatch consistent while leaving dialect
implementations behind their existing hooks.
## Changes
- expose public `ExprOperand` and `ExprWithOp` base-layer classes
- route TIRx and Relax expression nodes through the shared operator
surface
- preserve unresolved Relax operator chains and TE equality with
primitive expressions
- propagate `NotImplemented` for unsupported binary operators so Python
reflected dispatch remains available
- keep subscription and `SubscriptProxy` behavior outside this change
---
python/tvm/ir/__init__.py | 1 +
python/tvm/ir/expr.py | 94 +++++++++++------
python/tvm/ir/function.py | 7 +-
python/tvm/relax/expr.py | 136 +++++--------------------
python/tvm/relax/op/__init__.py | 25 ++++-
python/tvm/te/tensor.py | 4 +-
python/tvm/tirx/expr.py | 30 ++----
python/tvm/tirx/op.py | 4 +-
tests/python/relax/test_expr.py | 52 ++++++++++
tests/python/tirx-base/test_tir_constructor.py | 29 ++++++
10 files changed, 211 insertions(+), 171 deletions(-)
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 8a2d241cd1..412bd89061 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -38,6 +38,7 @@ from .expr import (
Call,
Expr,
ExprOperand,
+ ExprWithOp,
GlobalVar,
OpaqueExpr,
Range,
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 14d467263b..3cbd19c356 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -149,97 +149,113 @@ class ExprOperand:
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__add__(self, other)
- return _tensor_expr_overload.__add__(self, other)
+ result = _tensor_expr_overload.__add__(self, other)
+ return result
def __radd__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__radd__(self, other)
- return _tensor_expr_overload.__radd__(self, other)
+ result = _tensor_expr_overload.__radd__(self, other)
+ return result
def __sub__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__sub__(self, other)
- return _tensor_expr_overload.__sub__(self, other)
+ result = _tensor_expr_overload.__sub__(self, other)
+ return result
def __rsub__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rsub__(self, other)
- return _tensor_expr_overload.__rsub__(self, other)
+ result = _tensor_expr_overload.__rsub__(self, other)
+ return result
def __mul__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__mul__(self, other)
- return _tensor_expr_overload.__mul__(self, other)
+ result = _tensor_expr_overload.__mul__(self, other)
+ return result
def __rmul__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rmul__(self, other)
- return _tensor_expr_overload.__rmul__(self, other)
+ result = _tensor_expr_overload.__rmul__(self, other)
+ return result
def __div__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__div__(self, other)
- return _tensor_expr_overload.__div__(self, other)
+ result = _tensor_expr_overload.__div__(self, other)
+ return result
def __rdiv__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rdiv__(self, other)
- return _tensor_expr_overload.__rdiv__(self, other)
+ result = _tensor_expr_overload.__rdiv__(self, other)
+ return result
def __truediv__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__truediv__(self, other)
- return _tensor_expr_overload.__truediv__(self, other)
+ result = _tensor_expr_overload.__truediv__(self, other)
+ return result
def __rtruediv__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rtruediv__(self, other)
- return _tensor_expr_overload.__rtruediv__(self, other)
+ result = _tensor_expr_overload.__rtruediv__(self, other)
+ return result
def __floordiv__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__floordiv__(self, other)
- return _tensor_expr_overload.__floordiv__(self, other)
+ result = _tensor_expr_overload.__floordiv__(self, other)
+ return result
def __rfloordiv__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rfloordiv__(self, other)
- return _tensor_expr_overload.__rfloordiv__(self, other)
+ result = _tensor_expr_overload.__rfloordiv__(self, other)
+ return result
def __mod__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__mod__(self, other)
- return _tensor_expr_overload.__mod__(self, other)
+ result = _tensor_expr_overload.__mod__(self, other)
+ return result
def __rmod__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rmod__(self, other)
- return _tensor_expr_overload.__rmod__(self, other)
+ result = _tensor_expr_overload.__rmod__(self, other)
+ return result
def __pow__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return NotImplemented
- return _tensor_expr_overload.__pow__(self, other)
+ result = _tensor_expr_overload.__pow__(self, other)
+ return result
def __rpow__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return NotImplemented
- return _tensor_expr_overload.__rpow__(self, other)
+ result = _tensor_expr_overload.__rpow__(self, other)
+ return result
def __neg__(self):
self = _realize_operand(self)
@@ -250,7 +266,7 @@ class ExprOperand:
return result
result = _tensor_expr_overload.__neg__(self)
if result is NotImplemented:
- raise TypeError("Tensor expression overload negative is not
registered")
+ raise TypeError(f"Operator overloading is not supported for
expression type {self.ty}")
return result
def __lshift__(self, other):
@@ -320,19 +336,21 @@ class ExprOperand:
if result is NotImplemented:
raise TypeError("Primitive expression overload __invert__ is
not registered")
return result
- return NotImplemented
+ raise TypeError(f"Operator overloading is not supported for expression
type {self.ty}")
def __lt__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__lt__(self, other)
- return _tensor_expr_overload.__lt__(self, other)
+ result = _tensor_expr_overload.__lt__(self, other)
+ return result
def __le__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__le__(self, other)
- return _tensor_expr_overload.__le__(self, other)
+ result = _tensor_expr_overload.__le__(self, other)
+ return result
def __eq__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
@@ -350,13 +368,15 @@ class ExprOperand:
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__gt__(self, other)
- return _tensor_expr_overload.__gt__(self, other)
+ result = _tensor_expr_overload.__gt__(self, other)
+ return result
def __ge__(self, other):
self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__ge__(self, other)
- return _tensor_expr_overload.__ge__(self, other)
+ result = _tensor_expr_overload.__ge__(self, other)
+ return result
def __nonzero__(self):
raise ValueError(
@@ -369,6 +389,8 @@ class ExprOperand:
def equal(self, other, span=None):
self, other = _realize_operand(self), _realize_operand(other)
+ if not is_prim_expr(self):
+ raise TypeError(f"Operator overloading is not supported for
expression type {self.ty}")
result = _overload_prim_expr.equal(self, other, span)
if result is NotImplemented:
raise TypeError("Primitive expression overload equal is not
registered")
@@ -383,20 +405,26 @@ class ExprOperand:
return result
result = _tensor_expr_overload.astype(self, dtype, span)
if result is NotImplemented:
- raise TypeError("Tensor expression overload astype is not
registered")
+ raise TypeError(f"Operator overloading is not supported for
expression type {self.ty}")
return result
+
+class _ExprCallable:
+ """Function-call capability for expression operands that can denote
functions."""
+
+ __slots__ = ()
+
def __call__(self, *args, attrs=None):
self, args = _realize_operand(self), tuple(_realize_operand(arg) for
arg in args)
if is_prim_expr(self):
raise TypeError("A primitive-valued expression cannot be called")
result = _tensor_expr_overload.__call__(self, *args, attrs=attrs)
if result is NotImplemented:
- raise TypeError("Tensor expression overload __call__ is not
registered")
+ raise TypeError(f"Expression of type {self.ty} cannot be called")
return result
-class _ExprWithOp(ExprOperand, Expr, Scriptable):
+class ExprWithOp(ExprOperand, Expr, Scriptable):
"""Common type-directed operator behavior for core expressions."""
__hash__ = Expr.__hash__
@@ -405,7 +433,11 @@ class _ExprWithOp(ExprOperand, Expr, Scriptable):
return self
-class SubscriptProxy(ExprOperand, ObjectConvertible):
+class _CallableExprWithOp(_ExprCallable, ExprWithOp):
+ """Common operator behavior for expression nodes that support function
calls."""
+
+
+class SubscriptProxy(_ExprCallable, ExprOperand, ObjectConvertible):
"""An immutable, lazily-realized subscription of an :class:`Expr`.
Point subscriptions may be chained to accumulate dimensions. A proxy
@@ -480,7 +512,7 @@ class SubscriptProxy(ExprOperand, ObjectConvertible):
@tvm_ffi.register_object("ir.Tuple")
-class Tuple(_ExprWithOp):
+class Tuple(_CallableExprWithOp):
"""Tuple expression that groups several fields together.
Parameters
@@ -513,7 +545,7 @@ class Tuple(_ExprWithOp):
@tvm_ffi.register_object("ir.TupleGetItem")
-class TupleGetItem(_ExprWithOp):
+class TupleGetItem(_CallableExprWithOp):
"""Get the index-th item from a tuple.
Parameters
@@ -537,7 +569,7 @@ class TupleGetItem(_ExprWithOp):
@tvm_ffi.register_object("ir.TensorLoad")
-class TensorLoad(_ExprWithOp):
+class TensorLoad(_CallableExprWithOp):
"""An indexed load from an expression source.
TensorLoad objects are constructed by a dialect-specific helper that
@@ -555,7 +587,7 @@ class TensorLoad(_ExprWithOp):
@tvm_ffi.register_object("ir.Call")
-class Call(_ExprWithOp):
+class Call(_CallableExprWithOp):
"""Core function call node."""
op: Expr
@@ -594,7 +626,7 @@ class Call(_ExprWithOp):
@tvm_ffi.register_object("ir.Var")
-class Var(_ExprWithOp):
+class Var(_CallableExprWithOp):
"""A canonical local variable in the IR.
Parameters
diff --git a/python/tvm/ir/function.py b/python/tvm/ir/function.py
index 5bc7e368dc..7a58e4c576 100644
--- a/python/tvm/ir/function.py
+++ b/python/tvm/ir/function.py
@@ -26,7 +26,7 @@ from tvm.runtime import Object
from . import _ffi_api
from .attrs import DictAttrs
-from .expr import Expr
+from .expr import _CallableExprWithOp
class CallingConv(IntEnum):
@@ -38,9 +38,12 @@ class CallingConv(IntEnum):
@tvm_ffi.register_object("ir.BaseFunc")
-class BaseFunc(Expr):
+class BaseFunc(_CallableExprWithOp):
"""Base class of all functions."""
+ def __bool__(self) -> bool:
+ return True
+
@property
def attrs(self):
"""Return the attrs member of the function."""
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 51c63e21d1..fed8d1c3bc 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -32,6 +32,7 @@ import tvm.runtime
from tvm import DataType
from ..ir import BaseFunc, Node, Span
+from ..ir.expr import _CallableExprWithOp
from ..runtime import Scriptable
from . import _ffi_api
@@ -87,9 +88,15 @@ Type.is_base_of = _relax_type_is_base_of # type:
ignore[attr-defined]
_op_ffi_api = None # pylint: disable=invalid-name
-def _binary_op_helper(lhs: "ExprWithOp", rhs: "ExprWithOp", op: Callable) ->
"ExprWithOp":
+def _is_tensor_or_missing_type(ty: Type) -> bool:
+ return isinstance(ty, tvm.relax.TensorType) or ty.is_missing()
+
+
+def _binary_op_helper(lhs: Expr, rhs: Expr, op: Callable):
if not isinstance(lhs, Expr): # type: ignore
raise ValueError("lhs must be Expr")
+ if not _is_tensor_or_missing_type(lhs.ty):
+ return NotImplemented
if isinstance(rhs, Expr): # type: ignore
return op(lhs, rhs)
elif isinstance(rhs, Number):
@@ -98,117 +105,14 @@ def _binary_op_helper(lhs: "ExprWithOp", rhs:
"ExprWithOp", op: Callable) -> "Ex
raise TypeError(f"type {type(rhs)} not supported")
-def _binary_rhs_helper(rhs: "ExprWithOp") -> "ExprWithOp":
+def _binary_rhs_helper(rhs: Expr):
if isinstance(rhs, Number):
raise TypeError(f"Please convert {rhs} with `const` first")
raise TypeError(f"type {type(rhs)} not supported")
-class ExprWithOp(Expr, Scriptable):
- """Basetype of all relax expressions that defines op overloading."""
-
- def astype(self, dtype: str | DataType) -> "ExprWithOp":
- """Cast the content type of the current data to dtype.
-
- Parameters
- ----------
- dtype : str
- The target data type.
-
- Note
- ----
- This function only works for TensorType Exprs.
-
- Returns
- -------
- result : ExprWithOp
- The result expression.
- """
- return _op_ffi_api.astype(self, dtype) # type: ignore
-
- def __neg__(self) -> "ExprWithOp":
- return _op_ffi_api.negative(self) # type: ignore
-
- def __lt__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.less) # type: ignore
-
- def __gt__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.greater) # type:
ignore
-
- def __ge__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.greater_equal) #
type: ignore
-
- def __le__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.less_equal) # type:
ignore
-
- # NOTE: Cannot override __eq__ and __ne__, which will influence object
equal
-
- def __add__(self, other: Expr) -> "ExprWithOp":
- if isinstance(self.ty, tvm.relax.TupleType) and isinstance(other,
tuple):
- return tuple([*self, *other])
-
- return _binary_op_helper(self, other, _op_ffi_api.add) # type: ignore
-
- def __radd__(self, other: Expr) -> "ExprWithOp":
- return self.__add__(other)
-
- def __sub__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.subtract) # type:
ignore
-
- def __rsub__(self, other: Expr) -> "ExprWithOp":
- return _binary_rhs_helper(other)
-
- def __mul__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.multiply) # type:
ignore
-
- def __rmul__(self, other: Expr) -> "ExprWithOp":
- return self.__mul__(other)
-
- def __truediv__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.divide) # type:
ignore
-
- def __rtruediv__(self, other: Expr) -> "ExprWithOp":
- return _binary_rhs_helper(other)
-
- def __floordiv__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.floor_divide) #
type: ignore
-
- def __rfloordiv__(self, other: Expr) -> "ExprWithOp":
- return _binary_rhs_helper(other)
-
- def __mod__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.mod) # type: ignore
-
- def __rmod__(self, other: Expr) -> "ExprWithOp":
- return _binary_rhs_helper(other)
-
- def __pow__(self, other: Expr) -> "ExprWithOp":
- return _binary_op_helper(self, other, _op_ffi_api.power) # type:
ignore
-
- def __rpow__(self, other: Expr) -> "ExprWithOp":
- return _binary_rhs_helper(other)
-
- def __call__(self, *args: list[Expr], attrs: dict[str, Any] | None = None)
-> "ExprWithOp":
- """Call the variable (if it represents a function).
-
- Parameters
- ----------
- args: List[Expr]
- The arguments to the call.
-
- attr: Optional[Dict[str, object]]
- The additional attributes to the call.
-
- Returns
- -------
- call: ExprWithOp
- A call taking the variable as a function.
- """
- return tvm.ir.Call(self, args, attrs=attrs)
-
-
@tvm_ffi.register_object("relax.expr.If")
-class If(ExprWithOp):
+class If(_CallableExprWithOp):
"""A conditional expression in Relax.
Parameters
@@ -231,6 +135,9 @@ class If(ExprWithOp):
false_branch: Expr
span: Span | None
+ def __bool__(self) -> bool:
+ return True
+
def __init__(self, cond: Expr, true_branch: Expr, false_branch: Expr,
span: Span | None = None):
self.__init_handle_by_constructor__(
_ffi_api.If,
@@ -247,7 +154,7 @@ TupleGetItem = tvm.ir.TupleGetItem
@tvm_ffi.register_object("relax.expr.ShapeExpr")
-class ShapeExpr(ExprWithOp):
+class ShapeExpr(_CallableExprWithOp):
"""A shape expression which allows users to construct a shape containing
Expr.
Parameters
@@ -277,6 +184,9 @@ class ShapeExpr(ExprWithOp):
def __len__(self):
return len(self.values)
+ def __bool__(self) -> bool:
+ return len(self) != 0
+
def make_shape(shape: list[Any] | tuple[Any, ...]) -> ShapeExpr:
if isinstance(shape, list | tuple):
@@ -288,7 +198,7 @@ def make_shape(shape: list[Any] | tuple[Any, ...]) ->
ShapeExpr:
@tvm_ffi.register_object("relax.expr.Constant")
-class Constant(ExprWithOp):
+class Constant(_CallableExprWithOp):
"""Constant Tensor
Parameters
@@ -310,6 +220,9 @@ class Constant(ExprWithOp):
data: tvm.runtime.Tensor
span: Span | None
+ def __bool__(self) -> bool:
+ return True
+
def __init__(
self,
data: tvm.runtime.Tensor,
@@ -488,13 +401,16 @@ class DataflowBlock(BindingBlock):
@tvm_ffi.register_object("relax.expr.SeqExpr")
-class SeqExpr(ExprWithOp):
+class SeqExpr(_CallableExprWithOp):
"""A sequence of binding blocks followed by an expression."""
blocks: list[BindingBlock]
body: Expr
span: Span | None
+ def __bool__(self) -> bool:
+ return True
+
def __init__(self, blocks: list[BindingBlock], body: Expr, span: Span |
None = None) -> None:
self.__init_handle_by_constructor__(_ffi_api.SeqExpr, blocks, body,
span) # type: ignore
@@ -639,7 +555,7 @@ class Function(BaseFunc, Scriptable):
@tvm_ffi.register_object("relax.expr.ExternFunc")
-class ExternFunc(BaseFunc, ExprWithOp):
+class ExternFunc(BaseFunc):
"""extern function, which represents a PackedFunc."""
global_symbol: String
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 21a25e16b4..d8628c806d 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -182,14 +182,29 @@ def _register_op_make():
return tuple([*lhs, *rhs])
return expr._binary_op_helper(lhs, rhs, _ffi_api.add)
- def _rhs(_lhs, rhs):
+ def _rhs(lhs, rhs):
+ if not expr._is_tensor_or_missing_type(lhs.ty):
+ return NotImplemented
return expr._binary_rhs_helper(rhs)
- _tensor_expr_overload.astype = lambda lhs, dtype, _span=None:
_ffi_api.astype(lhs, dtype)
- _tensor_expr_overload.__call__ = lambda func, *args, attrs=None:
expr.tvm.ir.Call(
- func, args, attrs=attrs
+ def _unary(lhs, op):
+ if not expr._is_tensor_or_missing_type(lhs.ty):
+ return NotImplemented
+ return op(lhs)
+
+ def _call(func, *args, attrs=None):
+ if not (
+ isinstance(func.ty, expr.tvm.ir.FuncType | expr.tvm.relax.FuncType)
+ or func.ty.is_missing()
+ ):
+ return NotImplemented
+ return expr.tvm.ir.Call(func, args, attrs=attrs)
+
+ _tensor_expr_overload.astype = lambda lhs, dtype, _span=None: (
+ _ffi_api.astype(lhs, dtype) if expr._is_tensor_or_missing_type(lhs.ty)
else NotImplemented
)
- _tensor_expr_overload.__neg__ = lambda lhs: _ffi_api.negative(lhs)
+ _tensor_expr_overload.__call__ = _call
+ _tensor_expr_overload.__neg__ = lambda lhs: _unary(lhs, _ffi_api.negative)
_tensor_expr_overload.__lt__ = lambda lhs, rhs:
expr._binary_op_helper(lhs, rhs, _ffi_api.less)
_tensor_expr_overload.__le__ = lambda lhs, rhs: expr._binary_op_helper(
lhs, rhs, _ffi_api.less_equal
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index c124e7639c..de42bc714b 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -19,7 +19,7 @@
# pylint: disable=invalid-name
import tvm_ffi
-from tvm.ir import OpaqueExpr
+from tvm.ir import OpaqueExpr, is_prim_expr
from tvm.runtime import Object, ObjectConvertible, const
from tvm.tirx import expr as _expr
@@ -271,7 +271,7 @@ class Tensor(OpaqueExpr, TensorOpBase):
def __eq__(self, other):
if not isinstance(other, Tensor):
- if isinstance(other, _expr.ExprOp):
+ if isinstance(other, _expr.ExprOp) or is_prim_expr(other):
return _expr.EqualOp(self, other)
return False
if self.ndim == 0 and other.ndim == 0:
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index d46f231668..142fc1fee4 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -414,29 +414,21 @@ class IntImmEnum(ObjectConvertible):
return IntImm("int32", self.value, self.span) # type: ignore
-class ExprWithOp(ExprOp, Expr, Scriptable):
- """Helper base class to inherit from Expr."""
-
- # In Python3, We have to explicitly tell interpreter to retain __hash__ if
we overide __eq__
- # https://docs.python.org/3.1/reference/datamodel.html#object.__hash__
- __hash__ = Expr.__hash__
-
-
-class ConstExpr(ExprWithOp):
+class ConstExpr(ir.ExprWithOp):
pass
-class BinaryOpExpr(ExprWithOp):
+class BinaryOpExpr(ir.ExprWithOp):
a: Expr
b: Expr
-class CmpExpr(ExprWithOp):
+class CmpExpr(ir.ExprWithOp):
a: Expr
b: Expr
-class LogicalExpr(ExprWithOp):
+class LogicalExpr(ir.ExprWithOp):
pass
@@ -569,7 +561,7 @@ class CommReducer(Object, Scriptable):
@tvm_ffi.register_object("tirx.Reduce")
-class Reduce(ExprWithOp):
+class Reduce(ir.ExprWithOp):
"""Reduce node.
Parameters
@@ -738,7 +730,7 @@ class StringImm(ConstExpr):
@tvm_ffi.register_object("tirx.Cast")
-class Cast(ExprWithOp):
+class Cast(ir.ExprWithOp):
"""Cast expression.
Parameters
@@ -1124,7 +1116,7 @@ class Not(LogicalExpr):
@tvm_ffi.register_object("tirx.Select")
-class Select(ExprWithOp):
+class Select(ir.ExprWithOp):
"""Select node.
Note
@@ -1191,7 +1183,7 @@ def BufferLoad(buffer: Buffer, indices: list[Expr], span:
Span | None = None) ->
@tvm_ffi.register_object("tirx.Ramp")
-class Ramp(ExprWithOp):
+class Ramp(ir.ExprWithOp):
"""Ramp node.
Parameters
@@ -1224,7 +1216,7 @@ class Ramp(ExprWithOp):
@tvm_ffi.register_object("tirx.Broadcast")
-class Broadcast(ExprWithOp):
+class Broadcast(ir.ExprWithOp):
"""Broadcast node.
Parameters
@@ -1247,7 +1239,7 @@ class Broadcast(ExprWithOp):
@tvm_ffi.register_object("tirx.Shuffle")
-class Shuffle(ExprWithOp):
+class Shuffle(ir.ExprWithOp):
"""Shuffle node.
Parameters
@@ -1286,7 +1278,7 @@ class CallEffectKind:
@tvm_ffi.register_object("tirx.Let")
-class Let(ExprWithOp):
+class Let(ir.ExprWithOp):
"""Let node.
Parameters
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 5621e93195..2ce3402adc 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -24,7 +24,7 @@ from tvm_ffi import Array
import tvm
from tvm import tirx
-from tvm.ir import Call, Expr, Op, PointerType, PrimType, TensorLoad
+from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType, PrimType,
TensorLoad
from tvm.ir.base import Span
from tvm.ir.expr import _realize_operand
from tvm.ir.type import TensorMapType
@@ -32,7 +32,7 @@ from tvm.runtime import const
from . import _ffi_api
from .buffer import Buffer, buffer_data, is_buffer_var
-from .expr import BufferLoad, CommReducer, ExprOp, ExprWithOp, IntImm, Var
+from .expr import BufferLoad, CommReducer, ExprOp, IntImm, Var
tir = tirx # alias for backward compat with upstream tir.convert() calls
diff --git a/tests/python/relax/test_expr.py b/tests/python/relax/test_expr.py
index 12a650be02..fc53813afe 100644
--- a/tests/python/relax/test_expr.py
+++ b/tests/python/relax/test_expr.py
@@ -379,6 +379,58 @@ def test_call_raises_error_for_missing_operator():
rx.Call(None, [])
+def test_shared_operator_surface_preserves_relax_semantics():
+ tensor_ty = rx.TensorType([2], "float32")
+ x = rx.Var("x", tensor_ty)
+ y = rx.Var("y", tensor_ty)
+
+ assert (x == x) is True
+ assert (x != x) is False
+ assert (x == y) is False
+ assert (x != y) is True
+ assert isinstance(hash(x), int)
+
+ assert isinstance(x + y, tvm.ir.Call)
+ assert isinstance(-x, tvm.ir.Call)
+ assert isinstance(x.astype("float16"), tvm.ir.Call)
+
+ with pytest.raises(ValueError, match="Cannot use and"):
+ bool(x)
+
+
+def test_shared_operator_surface_rejects_non_tensor_relax_values():
+ shape = rx.ShapeExpr([1, 2])
+ message = "Operator overloading is not supported for expression type"
+
+ assert (shape == shape) is True
+ assert (shape != shape) is False
+ assert bool(shape)
+
+ with pytest.raises(TypeError, match="unsupported operand type"):
+ shape + shape
+
+ for operation in (
+ lambda: -shape,
+ lambda: shape.equal(shape),
+ lambda: shape.astype("float32"),
+ ):
+ with pytest.raises(TypeError, match=message):
+ operation()
+
+
+def test_shared_operator_surface_calls_function_typed_values():
+ tensor_ty = rx.TensorType([2], "float32")
+ arg = rx.Var("arg", tensor_ty)
+ func_ty = rx.FuncType([tensor_ty], tensor_ty)
+
+ for func in (rx.Var("func", func_ty), rx.ExternFunc("extern_func",
func_ty)):
+ call = func(arg)
+ assert isinstance(call, tvm.ir.Call)
+ assert call.op.same_as(func)
+ assert len(call.args) == 1
+ assert call.args[0].same_as(arg)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/tirx-base/test_tir_constructor.py
b/tests/python/tirx-base/test_tir_constructor.py
index 73b37f4b4d..65be8c07c9 100644
--- a/tests/python/tirx-base/test_tir_constructor.py
+++ b/tests/python/tirx-base/test_tir_constructor.py
@@ -214,6 +214,35 @@ def test_expr_constructor():
assert x.body == v
+def test_operator_base_categories_have_primitive_type():
+ var = tvm.tirx.Var("x", "int32")
+ buffer = tvm.tirx.decl_buffer([4], "float32")
+ expressions = [
+ tvm.tirx.IntImm("int32", 1),
+ tvm.tirx.Add(var, 1),
+ tvm.tirx.LT(var, 1),
+ tvm.tirx.And(var < 1, var < 2),
+ tvm.tirx.Reduce(
+ None,
+ [1],
+ [tvm.tirx.IterVar((0, 1), "i", tvm.tirx.IterVar.CommReduce)],
+ None,
+ 0,
+ ),
+ tvm.tirx.Cast("float32", var),
+ tvm.tirx.Select(var < 1, var, 1),
+ tvm.tirx.BufferLoad(buffer, [0]),
+ tvm.tirx.Ramp(0, 1, 4),
+ tvm.tirx.Broadcast(var, 4),
+ tvm.tirx.Shuffle([tvm.tirx.Broadcast(var, 2)], [0]),
+ tvm.tirx.Let(var, 1, var),
+ ]
+
+ for expression in expressions:
+ assert isinstance(expression, tvm.ir.ExprWithOp)
+ assert isinstance(expression.ty, tvm.ir.PrimType)
+
+
def test_stmt_constructor():
v = tvm.tirx.Var("aa", "int32")
nop = tvm.tirx.Evaluate(1)
