Tianqi Chen
d883f5064f
## Summary
TVM-side cleanup that drops the `python/tvm/runtime/object.py` shim and
routes `tvm.runtime.Object` directly to `tvm_ffi.Object`. The
`tvm.runtime.Object` re-export is preserved (now a re-export of
`tvm_ffi.Object`) so external callers keep working.
The load-bearing `__object_repr__` install — which wires TVM IR objects
up to the rich C++ `ReprPrinter` registered through
`init_ffi_api("node", ...)` — moves into
`python/tvm/runtime/_ffi_node_api.py`.
That module is already imported as a side-effect-only module from
`python/tvm/runtime/__init__.py`, so the override fires at the right
time (after `init_ffi_api` registers the C++ printer).
`_ffi_node_api.AsRepr` itself is **kept**: `tvm_ffi`'s default repr is
primitive (`ClassName(ptr)`); TVM IR objects need the rich printer
registered via `init_ffi_api("node", ...)`. `AsRepr` is what bridges
that printer back into Python `repr(obj)` and is also the runtime-only
fallback when `libtvm.so` is unavailable.
The 7 in-tree importers of the deleted shim (plus one straggler in
`runtime/disco/session.py`) are switched to either
`from tvm.runtime import Object` or `from tvm_ffi import Object`,
depending on which pattern the file already uses.
## Test plan
- [x] `python -c "import tvm; print(repr(tvm.IRModule({})))"` produces
TVMScript-style output (rich repr preserved).
- [x] `pytest tests/python/all-platform-minimal-test/ -x` — 75 passed,
77 skipped (matches baseline).
- [x] `pytest tests/python/tirx-base/ -x` — 273 passed, 2 skipped.
- [x] `pre-commit run --files <changed files>` — all hooks pass.
- [ ] CI green.
109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
# 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.
|
|
# pylint: disable=invalid-name
|
|
"""Function definitions."""
|
|
|
|
from enum import IntEnum
|
|
|
|
import tvm_ffi
|
|
|
|
import tvm.runtime
|
|
from tvm.runtime import Object
|
|
|
|
from . import _ffi_api
|
|
from .attrs import DictAttrs
|
|
from .expr import RelaxExpr
|
|
|
|
|
|
class CallingConv(IntEnum):
|
|
"""Possible kinds of calling conventions."""
|
|
|
|
DEFAULT = 0
|
|
C_PACKED_FUNC = 1
|
|
DEVICE_KERNEL_LAUNCH = 2
|
|
|
|
|
|
@tvm_ffi.register_object("ir.BaseFunc")
|
|
class BaseFunc(RelaxExpr):
|
|
"""Base class of all functions."""
|
|
|
|
@property
|
|
def attrs(self):
|
|
"""Return the attrs member of the function."""
|
|
return _ffi_api.BaseFunc_Attrs(self)
|
|
|
|
def with_attr(self, attr_key_or_dict, attr_value=None) -> "BaseFunc":
|
|
"""Create a new copy of the function and update the attribute.
|
|
|
|
Parameters
|
|
----------
|
|
attr_key_or_dict : Union[str, dict]
|
|
The attribute key to use or a dict containing multiple key value pairs.
|
|
|
|
attr_value : Object
|
|
The new attribute value.
|
|
|
|
Returns
|
|
-------
|
|
func : BaseFunc
|
|
A new copy of the function
|
|
"""
|
|
# make sure we first copy so that we can safely do copy on write
|
|
# for multiple updates.
|
|
res = _ffi_api.BaseFuncCopy(self)
|
|
|
|
if isinstance(attr_key_or_dict, dict):
|
|
for key, val in attr_key_or_dict.items():
|
|
res = _ffi_api.BaseFuncWithAttr(res._move(), key, tvm.runtime.convert(val))
|
|
return res
|
|
|
|
return _ffi_api.BaseFuncWithAttr(
|
|
res._move(), attr_key_or_dict, tvm.runtime.convert(attr_value)
|
|
)
|
|
|
|
def with_attrs(self, attr_map: DictAttrs | dict[str, Object]) -> "BaseFunc":
|
|
"""Copy the IRModule and add the given attribute map to it.
|
|
Parameters
|
|
----------
|
|
attr_map: Union[DictAttrs, Dict[str, Object]]
|
|
The attribute map
|
|
Returns
|
|
-------
|
|
func : BaseFunc
|
|
A new copy of the function
|
|
"""
|
|
if isinstance(attr_map, tvm.ir.DictAttrs):
|
|
attr_map = attr_map._dict()
|
|
|
|
return _ffi_api.BaseFuncWithAttrs(self, attr_map)
|
|
|
|
def without_attr(self, attr_key: str) -> "BaseFunc":
|
|
"""Create a new copy of the function with an attribute without provided key.
|
|
|
|
Parameters
|
|
----------
|
|
attr_key : str
|
|
The attribute key to delete from the attrubte pairs.
|
|
|
|
|
|
Returns
|
|
-------
|
|
func : BaseFunc
|
|
A new copy of the function
|
|
"""
|
|
return _ffi_api.BaseFuncWithoutAttr(self, attr_key)
|