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 00cabd2a6d [REFACTOR][PYTHON] Use the native FFI object protocol
(#20726)
00cabd2a6d is described below
commit 00cabd2a6df5118a5c57651ab5d2613674a03cae
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 13:34:14 2026 -0400
[REFACTOR][PYTHON] Use the native FFI object protocol (#20726)
Deferred expressions, TE tensor slices, and lazy enum values now expose
`__tvm_ffi_object__`, the native FFI object protocol. Script builders
and Relax type-argument helpers use that hook, and the obsolete
`tvm.runtime.ObjectConvertible` export and `object_generic.py` shim are
removed.
This migrates TVM consumers before removal of the legacy protocol in
TVM-FFI.
---
python/tvm/ir/_overload_prim_expr.py | 10 +-
python/tvm/relax/op/base.py | 10 +-
python/tvm/relax/script/ir_builder/op.py | 17 ++-
python/tvm/runtime/__init__.py | 1 -
python/tvm/runtime/object_generic.py | 21 ----
python/tvm/script/ir_builder/stmt.py | 6 +-
python/tvm/te/tensor.py | 80 ++++++-------
python/tvm/tirx/expr.py | 5 +-
python/tvm/tirx/script/ir_builder/op.py | 26 ++--
python/tvm/tirx/tile/instruction.py | 4 +-
.../python/runtime/test_runtime_object_protocol.py | 132 +++++++++++++++++++++
tests/python/script/test_basic_usage.py | 2 +-
tests/python/sym/test_sym_rewrite_simplify.py | 6 +-
tests/python/te/test_te_tensor.py | 2 +-
14 files changed, 213 insertions(+), 109 deletions(-)
diff --git a/python/tvm/ir/_overload_prim_expr.py
b/python/tvm/ir/_overload_prim_expr.py
index d0fe6df594..f30d3c52f2 100644
--- a/python/tvm/ir/_overload_prim_expr.py
+++ b/python/tvm/ir/_overload_prim_expr.py
@@ -18,7 +18,7 @@
from tvm.ir.location import UNKNOWN_LOC, Location
-from ..runtime import DataTypeCode, ObjectConvertible
+from ..runtime import DataTypeCode
from ._constant import const
from .expr import Expr, TensorRegion, is_prim_expr
from .prim import _ffi_api as _prim_ffi_api
@@ -246,7 +246,7 @@ class ExprOp:
return _prim_ffi_api._cast(self, dtype, loc) # type: ignore
-class EqualOp(ObjectConvertible, ExprOp):
+class EqualOp(ExprOp):
"""Deferred equal operator.
This is used to support sugar that a == b can either
@@ -280,7 +280,7 @@ class EqualOp(ObjectConvertible, ExprOp):
def __bool__(self) -> bool:
return self.__nonzero__()
- def asobject(self) -> Expr:
+ def __tvm_ffi_object__(self) -> Expr:
"""Convert object."""
return _prim_ffi_api._OpEQ(self.a, self.b, self.loc) # type: ignore
@@ -292,7 +292,7 @@ class EqualOp(ObjectConvertible, ExprOp):
return f"EqualOp({self.a!r}, {self.b!r})"
-class NotEqualOp(ObjectConvertible, ExprOp):
+class NotEqualOp(ExprOp):
"""Deferred NE operator.
This is used to support sugar that a != b can either
@@ -326,7 +326,7 @@ class NotEqualOp(ObjectConvertible, ExprOp):
def __bool__(self) -> bool:
return self.__nonzero__()
- def asobject(self) -> Expr:
+ def __tvm_ffi_object__(self) -> Expr:
"""Convert object."""
return _prim_ffi_api._OpNE(self.a, self.b, self.loc) # type: ignore
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
index 688aa08117..8f0125c7cd 100644
--- a/python/tvm/relax/op/base.py
+++ b/python/tvm/relax/op/base.py
@@ -27,7 +27,7 @@ from tvm.ir import Attrs, Call, Op
from tvm.ir.attrs import make_node as _make_attrs
from tvm.ir.location import UNKNOWN_LOC, Location
from tvm.ir.op import _make_op_api
-from tvm.runtime import DataTypeCode, Object, ObjectConvertible
+from tvm.runtime import DataTypeCode, Object
from ..expr import Expr, ExternFunc, GlobalVar, Var
from ..type import Type
@@ -356,8 +356,8 @@ def call_builtin_with_ctx(
ty_args = [
value()
if callable(value) and not isinstance(value, Type)
- else value.asobject()
- if isinstance(value, ObjectConvertible)
+ else value.__tvm_ffi_object__()
+ if hasattr(value, "__tvm_ffi_object__")
else value
for value in ty_args
]
@@ -779,8 +779,8 @@ def call_pure_packed(
(
ty()
if callable(ty) and not isinstance(ty, Type)
- else ty.asobject()
- if isinstance(ty, ObjectConvertible)
+ else ty.__tvm_ffi_object__()
+ if hasattr(ty, "__tvm_ffi_object__")
else ty
)
for ty in ty_args
diff --git a/python/tvm/relax/script/ir_builder/op.py
b/python/tvm/relax/script/ir_builder/op.py
index fea36c5faf..0c7370501b 100644
--- a/python/tvm/relax/script/ir_builder/op.py
+++ b/python/tvm/relax/script/ir_builder/op.py
@@ -24,8 +24,6 @@ import numbers as _numbers
from types import SimpleNamespace as _Namespace
from typing import Any
-import tvm_ffi as _ffi
-
import tvm
from tvm import ir as _ir
from tvm import relax as _relax
@@ -236,7 +234,6 @@ from tvm.relax.op import to_vdevice as _to_vdevice
from tvm.relax.op.builtin import stop_lift_params
from tvm.relax.type import Type
from tvm.relax.utils import convert_to_expr
-from tvm.runtime import ObjectConvertible
from tvm.script.ir_builder import base as _base
from . import distributed as dist # noqa: F401
@@ -316,8 +313,8 @@ def call_packed(
(
ty()
if callable(ty) and not isinstance(ty, _ir.Type)
- else ty.asobject()
- if isinstance(ty, ObjectConvertible)
+ else ty.__tvm_ffi_object__()
+ if hasattr(ty, "__tvm_ffi_object__")
else ty
)
for ty in ty_args
@@ -487,15 +484,17 @@ def select(condition, true_value, false_value):
def if_then_else_(condition: Any, true_value: Any, false_value: Any) -> Any:
"""Implements
:func:`tvm.script.ir_builder.parser_protocol.if_then_else_`."""
- if isinstance(condition, _ffi.ObjectConvertible):
- condition = condition.asobject()
+ if hasattr(condition, "__tvm_ffi_object__"):
+ condition = condition.__tvm_ffi_object__()
if not isinstance(condition, _ir.Expr):
return true_value if condition else false_value
true_value = (
- true_value.asobject() if isinstance(true_value,
_ffi.ObjectConvertible) else true_value
+ true_value.__tvm_ffi_object__() if hasattr(true_value,
"__tvm_ffi_object__") else true_value
)
false_value = (
- false_value.asobject() if isinstance(false_value,
_ffi.ObjectConvertible) else false_value
+ false_value.__tvm_ffi_object__()
+ if hasattr(false_value, "__tvm_ffi_object__")
+ else false_value
)
if _ir.is_prim_expr(condition) and all(
_ir.is_prim_expr(value)
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py
index 552dfd426a..b57d3af1b2 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -21,7 +21,6 @@ from tvm_ffi import convert, Object
from tvm_ffi._dtype import dtype as DataType, DataTypeCode
# class exposures
-from .object_generic import ObjectConvertible
from .device import Device
from .memory_scope import MemoryScope
from ._tensor import Tensor, tensor, empty
diff --git a/python/tvm/runtime/object_generic.py
b/python/tvm/runtime/object_generic.py
deleted file mode 100644
index 14f54dc82d..0000000000
--- a/python/tvm/runtime/object_generic.py
+++ /dev/null
@@ -1,21 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-# ruff: noqa: F401
-"""Common implementation of object generic related logic"""
-
-# pylint: disable=unused-import, invalid-name
-from tvm_ffi import ObjectConvertible
diff --git a/python/tvm/script/ir_builder/stmt.py
b/python/tvm/script/ir_builder/stmt.py
index f99043ee69..ca1265b41d 100644
--- a/python/tvm/script/ir_builder/stmt.py
+++ b/python/tvm/script/ir_builder/stmt.py
@@ -28,8 +28,6 @@ from collections.abc import Sequence
from functools import partial as _partial
from typing import Any
-import tvm_ffi as _ffi
-
import tvm
from tvm import ir as _ir
from tvm.ir import Expr, TensorRegion, Type, Var
@@ -44,8 +42,8 @@ _Loc = _base.LocationEntry | _ir.Location | None
def _as_expr(value):
- if isinstance(value, _ffi.ObjectConvertible):
- value = value.asobject()
+ if hasattr(value, "__tvm_ffi_object__"):
+ value = value.__tvm_ffi_object__()
if isinstance(value, _ir.Expr):
return value
if isinstance(value, str):
diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py
index 2899f50f62..47253d2a79 100644
--- a/python/tvm/te/tensor.py
+++ b/python/tvm/te/tensor.py
@@ -19,19 +19,19 @@
# pylint: disable=invalid-name
import tvm_ffi
-from tvm.ir import OpaqueExpr, SeqStmt, const, is_prim_expr
+from tvm.ir import Expr, OpaqueExpr, SeqStmt, const, is_prim_expr
from tvm.ir.location import UNKNOWN_LOC, Location
-from tvm.runtime import Object, ObjectConvertible
+from tvm.runtime import Object
from tvm.tirx import expr as _expr
from . import _ffi_api, _te_tensor_overload
def _as_scalar_operand(value):
- return value.asobject() if isinstance(value, TensorSlice) else value
+ return value.__tvm_ffi_object__() if isinstance(value, TensorSlice) else
value
-class TensorSlice(ObjectConvertible):
+class TensorSlice:
"""Auxiliary data structure for enable slicing syntax from tensor."""
def __init__(self, tensor, indices):
@@ -45,7 +45,7 @@ class TensorSlice(ObjectConvertible):
indices = (indices,)
return TensorSlice(self.tensor, self.indices + indices)
- def asobject(self):
+ def __tvm_ffi_object__(self) -> Expr:
"""Convert slice to object."""
return self.tensor.__call__(*self.indices)
@@ -62,139 +62,139 @@ class TensorSlice(ObjectConvertible):
result = _te_tensor_overload.__add__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__add__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__add__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __radd__(self, other):
result = _te_tensor_overload.__radd__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__radd__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__radd__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __sub__(self, other):
result = _te_tensor_overload.__sub__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__sub__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__sub__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rsub__(self, other):
result = _te_tensor_overload.__rsub__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__rsub__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rsub__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __mul__(self, other):
result = _te_tensor_overload.__mul__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__mul__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__mul__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rmul__(self, other):
result = _te_tensor_overload.__rmul__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__rmul__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rmul__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __div__(self, other):
result = _te_tensor_overload.__div__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__div__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__div__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rdiv__(self, other):
result = _te_tensor_overload.__rdiv__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__rdiv__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rdiv__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __truediv__(self, other):
result = _te_tensor_overload.__truediv__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__truediv__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__truediv__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rtruediv__(self, other):
result = _te_tensor_overload.__rtruediv__(self, other)
if result is not NotImplemented:
return result
- return _expr.ExprOp.__rtruediv__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rtruediv__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __floordiv__(self, other):
- return _expr.ExprOp.__floordiv__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__floordiv__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rfloordiv__(self, other):
- return _expr.ExprOp.__rfloordiv__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rfloordiv__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __mod__(self, other):
- return _expr.ExprOp.__mod__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__mod__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rmod__(self, other):
- return _expr.ExprOp.__rmod__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rmod__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __neg__(self):
- return _expr.ExprOp.__neg__(self.asobject())
+ return _expr.ExprOp.__neg__(self.__tvm_ffi_object__())
def __lshift__(self, other):
- return _expr.ExprOp.__lshift__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__lshift__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rlshift__(self, other):
- return _expr.ExprOp.__rlshift__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rlshift__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rshift__(self, other):
- return _expr.ExprOp.__rshift__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rshift__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rrshift__(self, other):
- return _expr.ExprOp.__rrshift__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rrshift__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __and__(self, other):
- return _expr.ExprOp.__and__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__and__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rand__(self, other):
- return _expr.ExprOp.__rand__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rand__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __or__(self, other):
- return _expr.ExprOp.__or__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__or__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __ror__(self, other):
- return _expr.ExprOp.__ror__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__ror__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __xor__(self, other):
- return _expr.ExprOp.__xor__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__xor__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __rxor__(self, other):
- return _expr.ExprOp.__rxor__(self.asobject(),
_as_scalar_operand(other))
+ return _expr.ExprOp.__rxor__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __invert__(self):
- return _expr.ExprOp.__invert__(self.asobject())
+ return _expr.ExprOp.__invert__(self.__tvm_ffi_object__())
def __lt__(self, other):
- return _expr.ExprOp.__lt__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__lt__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __le__(self, other):
- return _expr.ExprOp.__le__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__le__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __eq__(self, other):
- return _expr.ExprOp.__eq__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__eq__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __ne__(self, other):
- return _expr.ExprOp.__ne__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__ne__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __gt__(self, other):
- return _expr.ExprOp.__gt__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__gt__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __ge__(self, other):
- return _expr.ExprOp.__ge__(self.asobject(), _as_scalar_operand(other))
+ return _expr.ExprOp.__ge__(self.__tvm_ffi_object__(),
_as_scalar_operand(other))
def __nonzero__(self):
- return _expr.ExprOp.__nonzero__(self.asobject())
+ return _expr.ExprOp.__nonzero__(self.__tvm_ffi_object__())
def __bool__(self):
return self.__nonzero__()
def equal(self, other, loc: Location = UNKNOWN_LOC):
- return _expr.ExprOp.equal(self.asobject(), _as_scalar_operand(other),
loc)
+ return _expr.ExprOp.equal(self.__tvm_ffi_object__(),
_as_scalar_operand(other), loc)
def astype(self, dtype, loc: Location = UNKNOWN_LOC):
- return _expr.ExprOp.astype(self.asobject(), dtype, loc)
+ return _expr.ExprOp.astype(self.__tvm_ffi_object__(), dtype, loc)
class TensorOpBase:
diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py
index 47b234d9b6..12ebc8e596 100644
--- a/python/tvm/tirx/expr.py
+++ b/python/tvm/tirx/expr.py
@@ -78,10 +78,9 @@ from tvm.ir.prim.expr import ( # noqa: F401
Shuffle,
Sub,
)
-from tvm.runtime import ObjectConvertible
-class IntImmEnum(ObjectConvertible):
+class IntImmEnum:
"""Lazily evaluate an IntImm in case
the constructor is not available in runtime.
@@ -98,7 +97,7 @@ class IntImmEnum(ObjectConvertible):
self.value = value
self.loc = loc
- def asobject(self) -> "IntImm":
+ def __tvm_ffi_object__(self) -> "IntImm":
"""Convert object."""
return IntImm("int32", self.value, self.loc) # type: ignore
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 6f36e164b3..2322c1b0e2 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -21,10 +21,6 @@ from __future__ import annotations
import functools
from typing import Any
-import tvm_ffi as _ffi
-
-# isort: off
-# isort: on
from tvm import ir
from tvm import ir as _ir
from tvm import tirx as _tir
@@ -608,7 +604,7 @@ def logical_and(*values):
Parameters
----------
values : Expr or Python value
- One or more operands. Object-convertible values are normalized first.
+ One or more operands. Values implementing ``__tvm_ffi_object__`` are
normalized first.
Host pairs follow Python logical operations; IR pairs use scalar
logical
or vector bitwise operations according to their types.
@@ -625,7 +621,8 @@ def logical_and(*values):
if not values:
raise TypeError("logical_and requires at least one operand")
values = [
- value.asobject() if isinstance(value, _ffi.ObjectConvertible) else
value for value in values
+ value.__tvm_ffi_object__() if hasattr(value, "__tvm_ffi_object__")
else value
+ for value in values
]
result = values[0]
for value in values[1:]:
@@ -643,7 +640,7 @@ def logical_or(*values):
Parameters
----------
values : Expr or Python value
- One or more operands. Object-convertible values are normalized first.
+ One or more operands. Values implementing ``__tvm_ffi_object__`` are
normalized first.
Host pairs follow Python logical operations; IR pairs use scalar
logical
or vector bitwise operations according to their types.
@@ -660,7 +657,8 @@ def logical_or(*values):
if not values:
raise TypeError("logical_or requires at least one operand")
values = [
- value.asobject() if isinstance(value, _ffi.ObjectConvertible) else
value for value in values
+ value.__tvm_ffi_object__() if hasattr(value, "__tvm_ffi_object__")
else value
+ for value in values
]
result = values[0]
for value in values[1:]:
@@ -678,7 +676,7 @@ def logical_not(value):
Parameters
----------
value : Expr or Python value
- Operand to negate. Object-convertible values are normalized first;
+ Operand to negate. Values implementing ``__tvm_ffi_object__`` are
normalized first;
IR expressions use primitive Not and host values use Python not.
Returns
@@ -686,8 +684,8 @@ def logical_not(value):
result : Expr or bool
The logical negation without testing an IR expression as a Python bool.
"""
- if isinstance(value, _ffi.ObjectConvertible):
- value = value.asobject()
+ if hasattr(value, "__tvm_ffi_object__"):
+ value = value.__tvm_ffi_object__()
return _tir.Not(value) if isinstance(value, _ir.Expr) else not value
@@ -697,7 +695,7 @@ def select(condition, true_value, false_value):
Parameters
----------
condition : Expr or Python value
- Scalar IR condition or host truth value. Object-convertible conditions
+ Scalar IR condition or host truth value. Conditions implementing
``__tvm_ffi_object__``
are normalized before selection.
true_value : Expr or Python value
Value selected when the condition is true.
@@ -714,8 +712,8 @@ def select(condition, true_value, false_value):
Both Python value arguments are constructed before this call. The generated
IR conditional evaluates only its selected arm at runtime.
"""
- if isinstance(condition, _ffi.ObjectConvertible):
- condition = condition.asobject()
+ if hasattr(condition, "__tvm_ffi_object__"):
+ condition = condition.__tvm_ffi_object__()
if not isinstance(condition, _ir.Expr):
return true_value if condition else false_value
return _tir.if_then_else(condition, true_value, false_value)
diff --git a/python/tvm/tirx/tile/instruction.py
b/python/tvm/tirx/tile/instruction.py
index 90fc08d7ff..c24b358c6c 100644
--- a/python/tvm/tirx/tile/instruction.py
+++ b/python/tvm/tirx/tile/instruction.py
@@ -204,8 +204,8 @@ def _expr(value):
return Tuple([_expr(v) for v in value])
if isinstance(value, str):
return StringImm(value)
- if hasattr(value, "asobject"):
- return value.asobject()
+ if hasattr(value, "__tvm_ffi_object__"):
+ return value.__tvm_ffi_object__()
return const(value)
diff --git a/tests/python/runtime/test_runtime_object_protocol.py
b/tests/python/runtime/test_runtime_object_protocol.py
new file mode 100644
index 0000000000..1f07859b4a
--- /dev/null
+++ b/tests/python/runtime/test_runtime_object_protocol.py
@@ -0,0 +1,132 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Native FFI object conversion for TVM wrappers and script helpers."""
+
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import ir, relax, te, tirx
+from tvm.relax.script.ir_builder import op as relax_op
+from tvm.script.ir_builder.stmt import _as_expr
+from tvm.tirx.script.ir_builder import op as tir_op
+from tvm.tirx.tile.instruction import _expr
+
+
+class ObjectWrapper:
+ def __init__(self, value):
+ self.value = value
+
+ def __tvm_ffi_object__(self) -> tvm_ffi.Object:
+ return self.value
+
+
+def test_object_protocol_identity():
+ value = ir.Var("x", "int32")
+ wrapped = ObjectWrapper(value)
+ assert tvm_ffi.convert(wrapped) is value
+ echo = tvm_ffi.convert_func(lambda x: x)
+ assert echo(wrapped).same_as(value)
+ assert tvm_ffi.convert([wrapped])[0].same_as(value)
+ assert tvm_ffi.convert({"value": [wrapped]})["value"][0].same_as(value)
+ assert _as_expr(wrapped) is value
+ assert _expr(wrapped) is value
+
+
[email protected]("unequal", [False, True])
+def test_deferred_comparison_protocol(unequal):
+ x, y = ir.Var("x", "int32"), ir.Var("y", "int32")
+ comparison = x != y if unequal else x == y
+ assert bool(comparison) is unequal
+ assert bool(x != x if unequal else x == x) is not unequal
+ result = tvm_ffi.convert(comparison)
+ assert isinstance(result, tirx.NE if unequal else tirx.EQ)
+ assert result.a.same_as(x)
+ assert result.b.same_as(y)
+ echo = tvm_ffi.convert_func(lambda x: x)
+ tvm.ir.assert_structural_equal(echo(comparison), result)
+ tvm.ir.assert_structural_equal(tvm_ffi.convert([comparison])[0], result)
+ assert isinstance(tir_op.logical_not(comparison), tirx.Not)
+ assert isinstance(tir_op.logical_and(comparison, True), tirx.And)
+ assert isinstance(tir_op.logical_or(comparison, False), tirx.Or)
+ tvm.ir.assert_structural_equal(_as_expr(comparison), result)
+ tvm.ir.assert_structural_equal(_expr(comparison), result)
+
+
+def test_tensor_slice_protocol():
+ tensor = te.placeholder((2, 3), name="A")
+ sliced = tensor[0][1]
+ result = tvm_ffi.convert(sliced)
+ assert result.op.same_as(ir.Op.get("te.tensor_load"))
+ assert result.args[0].same_as(tensor)
+ assert list(result.args[1:]) == [0, 1]
+ echo = tvm_ffi.convert_func(lambda x: x)
+ tvm.ir.assert_structural_equal(echo(sliced), result)
+ tvm.ir.assert_structural_equal(tvm_ffi.convert([sliced])[0], result)
+
+
+def test_enum_protocol():
+ enum = tirx.CallEffectKind.Pure
+ result = tvm_ffi.convert(enum)
+ assert isinstance(result, tirx.IntImm)
+ assert result.value == 1
+ assert result.ty == ir.PrimType("int32")
+ echo = tvm_ffi.convert_func(lambda x: x)
+ tvm.ir.assert_structural_equal(echo(enum), result)
+ tvm.ir.assert_structural_equal(tvm_ffi.convert([enum])[0], result)
+
+
+def test_builder_protocol_preserves_host_values():
+ host = object()
+ assert tir_op.select(True, host, None) is host
+ assert relax_op.if_then_else_(True, host, None) is host
+ assert tir_op.logical_and(True, host) is host
+ assert tir_op.logical_or(False, host) is host
+ x = ir.Var("x", "int32")
+ condition = ObjectWrapper(tvm_ffi.convert(x == 0))
+ true_value = ObjectWrapper(x)
+ tvm.ir.assert_structural_equal(
+ tir_op.select(condition, true_value, 0),
tir_op.select(condition.value, x, 0)
+ )
+ tvm.ir.assert_structural_equal(
+ relax_op.if_then_else_(condition, true_value, 0),
+ relax_op.if_then_else_(condition.value, x, 0),
+ )
+
+
[email protected](
+ "make_call", [relax.op.call_builtin_with_ctx, relax.op.call_pure_packed,
relax_op.call_packed]
+)
+def test_wrapped_type_arguments(make_call):
+ ty = ir.PrimType("int32")
+ call = make_call("custom", (), ty_args=[ObjectWrapper(ty)])
+ assert call.ty_args[0].same_as(ty)
+ call = make_call("custom", (), ty_args=[lambda: ty])
+ assert call.ty_args[0].same_as(ty)
+
+
+def test_invalid_object_protocol_result():
+ echo = tvm_ffi.convert_func(lambda x: x)
+ with pytest.raises(TypeError):
+ echo(ObjectWrapper(1))
+ with pytest.raises(TypeError):
+ tvm_ffi.convert([ObjectWrapper(1)])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/script/test_basic_usage.py
b/tests/python/script/test_basic_usage.py
index 78fb654bcd..3600886350 100644
--- a/tests/python/script/test_basic_usage.py
+++ b/tests/python/script/test_basic_usage.py
@@ -471,7 +471,7 @@ def
test_host_equality_result_is_not_arbitrarily_converted(primitive_language):
M = primitive_language.M
class CustomEquality(EqualOp):
- def asobject(self):
+ def __tvm_ffi_object__(self):
raise AssertionError("custom equality subclasses are host values")
result = CustomEquality(ir.Var("x", "int32"), 0)
diff --git a/tests/python/sym/test_sym_rewrite_simplify.py
b/tests/python/sym/test_sym_rewrite_simplify.py
index 6cebaff321..ded9372e92 100644
--- a/tests/python/sym/test_sym_rewrite_simplify.py
+++ b/tests/python/sym/test_sym_rewrite_simplify.py
@@ -34,9 +34,9 @@ from tvm.tirx import truncmod as tmod
class TestCase:
def __init__(self, before, expected, preconditions=None):
if isinstance(before, tirx.expr.EqualOp):
- before = before.asobject()
+ before = before.__tvm_ffi_object__()
if isinstance(expected, tirx.expr.EqualOp):
- expected = expected.asobject()
+ expected = expected.__tvm_ffi_object__()
self.before = self._convert(before)
self.expected = self._convert(expected)
@@ -45,7 +45,7 @@ class TestCase:
@staticmethod
def _convert(expr):
if isinstance(expr, tirx.expr.EqualOp):
- return expr.asobject()
+ return expr.__tvm_ffi_object__()
elif isinstance(expr, int):
return T.int32(expr)
elif isinstance(expr, float):
diff --git a/tests/python/te/test_te_tensor.py
b/tests/python/te/test_te_tensor.py
index 73dc5d0926..555c9b6790 100644
--- a/tests/python/te/test_te_tensor.py
+++ b/tests/python/te/test_te_tensor.py
@@ -41,7 +41,7 @@ def test_tensor():
assert T.op.output(0).__hash__() == T.__hash__()
d = {T.op.output(0): 1}
assert d[T] == 1
- load = T[0][0][0].asobject()
+ load = T[0][0][0].__tvm_ffi_object__()
assert isinstance(load, tvm.ir.Call)
assert load.op.same_as(tvm.ir.Op.get("te.tensor_load"))
assert load.args[0].same_as(T)
