diff --git a/3rdparty/tvm-ffi b/3rdparty/tvm-ffi index 5411a642d98b..21e30c3b1d44 160000 --- a/3rdparty/tvm-ffi +++ b/3rdparty/tvm-ffi @@ -1 +1 @@ -Subproject commit 5411a642d98bba14d4cb1e19a793c5b6758be019 +Subproject commit 21e30c3b1d4421f95fd69a6ba3fd0285c7b69e9c diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h index 6ed6ada0d230..a14135f3c855 100644 --- a/include/tvm/runtime/device_api.h +++ b/include/tvm/runtime/device_api.h @@ -25,6 +25,7 @@ #define TVM_RUNTIME_DEVICE_API_H_ #include +#include #include #include #include diff --git a/include/tvm/script/printer/ir_docsifier.h b/include/tvm/script/printer/ir_docsifier.h index e9c82265ff27..944e511457ae 100644 --- a/include/tvm/script/printer/ir_docsifier.h +++ b/include/tvm/script/printer/ir_docsifier.h @@ -19,6 +19,7 @@ #ifndef TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_ #define TVM_SCRIPT_PRINTER_IR_DOCSIFIER_H_ +#include #include #include #include diff --git a/python/tvm/ir/module.py b/python/tvm/ir/module.py index 1ce93e3700de..80cc6dfac297 100644 --- a/python/tvm/ir/module.py +++ b/python/tvm/ir/module.py @@ -40,6 +40,12 @@ class IRModule(Node, Scriptable): Map of global var to BaseFunc """ + # Opt into a per-instance ``__dict__``: ``IRModule`` stores the Python-only + # ``pyfuncs`` mapping. The base ``Object`` metaclass otherwise injects + # ``__slots__ = ()`` (required by the tvm-ffi lifetime model), which would + # forbid setting attributes. + __slots__ = ("__dict__",) + def __init__(self, functions=None, attrs=None, global_infos=None): if functions is None: functions = {} diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py index 41b49dd2c5d3..afa1f9b3a6a5 100644 --- a/python/tvm/relax/expr.py +++ b/python/tvm/relax/expr.py @@ -589,6 +589,12 @@ def __init__(self, blocks: list[BindingBlock], body: Expr, span: Span | None = N class Function(BaseFunc, Scriptable): """A Relax function.""" + # Opt into a per-instance ``__dict__``: callers attach ``__name__`` to the + # function (e.g. DataflowBlockRewrite.mutated_root_fn). The base ``Object`` + # metaclass otherwise injects ``__slots__ = ()`` (required by the tvm-ffi + # lifetime model), which would forbid setting attributes. + __slots__ = ("__dict__",) + params: list[Var] body: Expr ret_ty: Type diff --git a/python/tvm/runtime/_tensor.py b/python/tvm/runtime/_tensor.py index 6f1d60b01e2a..b53d88250e03 100644 --- a/python/tvm/runtime/_tensor.py +++ b/python/tvm/runtime/_tensor.py @@ -75,6 +75,12 @@ class Tensor(tvm_ffi.core.Tensor): how can we use TVM in existing project which might have their own array containers. """ + # Keep the same instance layout as ``tvm_ffi.core.Tensor``: the FFI runtime + # registers ``ffi.Tensor`` to its base wrapper and rejects re-registration + # with a larger wrapper (a subclass without ``__slots__`` would add + # ``__dict__``/``__weakref__`` and grow ``__basicsize__``). + __slots__ = () + def __setitem__(self, in_slice, value): """Set ndarray value""" if ( diff --git a/python/tvm/runtime/script_printer.py b/python/tvm/runtime/script_printer.py index 3f67e285de53..3863f43d20ae 100644 --- a/python/tvm/runtime/script_printer.py +++ b/python/tvm/runtime/script_printer.py @@ -118,6 +118,12 @@ def _relax_script(obj: Object, config: PrinterConfig) -> str: class Scriptable: """A base class that enables the script() and show() method.""" + # Empty slots so this mixin does not add ``__dict__``/``__weakref__`` to the + # FFI object subclasses that combine it with a slotted ``Object``/``Node``. + # The tvm-ffi lifetime model rejects registering a wrapper whose instance + # layout is larger than the base wrapper already registered for the type. + __slots__ = () + def script( self, *, diff --git a/python/tvm/te/tensor.py b/python/tvm/te/tensor.py index ae1fad55a4b6..ed045e2fc652 100644 --- a/python/tvm/te/tensor.py +++ b/python/tvm/te/tensor.py @@ -199,6 +199,12 @@ def astype(self, dtype, span=None): class TensorOpBase: """Operator overloads for whole TE Tensor values.""" + # Empty slots so this mixin does not add ``__dict__``/``__weakref__`` to the + # FFI object subclasses that combine it with a slotted ``Object``/``Node``. + # The tvm-ffi lifetime model rejects registering a wrapper whose instance + # layout is larger than the base wrapper already registered for the type. + __slots__ = () + def __add__(self, other): return _te_tensor_overload.__add__(self, other) diff --git a/python/tvm/tirx/expr.py b/python/tvm/tirx/expr.py index 3e4902c376a9..78638035f1e7 100644 --- a/python/tvm/tirx/expr.py +++ b/python/tvm/tirx/expr.py @@ -86,6 +86,12 @@ def _is_scalar_operand(value): class ExprOp: """Operator overloading for Expr like expressions.""" + # Empty slots so this mixin does not add ``__dict__``/``__weakref__`` to the + # FFI object subclasses that combine it with a slotted ``Object``/``Expr``. + # The tvm-ffi lifetime model rejects registering a wrapper whose instance + # layout is larger than the base wrapper already registered for the type. + __slots__ = () + # TODO(tkonolige): use inspect to add source information to these objects def expr_ty(self) -> ir.PrimType: diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py index ab7b6435a242..ce4cdf761b0b 100644 --- a/python/tvm/tirx/function.py +++ b/python/tvm/tirx/function.py @@ -60,6 +60,12 @@ class PrimFunc(BaseFunc, Scriptable): The location of this itervar in the source code. """ + # Opt into a per-instance ``__dict__``: the TVMScript parser attaches + # ``__name__`` to the parsed function. The base ``Object`` metaclass + # otherwise injects ``__slots__ = ()`` (required by the tvm-ffi lifetime + # model), which would forbid setting attributes. + __slots__ = ("__dict__",) + def __init__(self, params, body, ret_type=None, buffer_map=None, attrs=None, span=None): # Legacy compatibility: expand body-carrying leaf stmt wrappers # (e.g. DeclBuffer/AllocBuffer forms) into SeqStmt form. diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py index ba645985eaa7..4711d947a4f0 100644 --- a/python/tvm/tirx/stmt.py +++ b/python/tvm/tirx/stmt.py @@ -342,6 +342,12 @@ class AllocBuffer(Stmt): The location of this AllocBuffer in the source code. """ + # Opt into a per-instance ``__dict__``: the legacy ``body`` is carried on the + # Python side only. The base ``Object`` metaclass otherwise injects + # ``__slots__ = ()`` (required by the tvm-ffi lifetime model), which would + # forbid setting attributes. + __slots__ = ("__dict__",) + buffer: Buffer span: Span | None @@ -434,6 +440,12 @@ class DeclBuffer(Stmt): The location of this DeclBuffer in the source code. """ + # Opt into a per-instance ``__dict__``: the legacy ``body`` is carried on the + # Python side only. The base ``Object`` metaclass otherwise injects + # ``__slots__ = ()`` (required by the tvm-ffi lifetime model), which would + # forbid setting attributes. + __slots__ = ("__dict__",) + buffer: Buffer span: Span | None