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.
93 lines
2.6 KiB
Python
93 lines
2.6 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
|
|
"""Testing utilities for relax VM"""
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np # type: ignore
|
|
|
|
import tvm
|
|
from tvm import relax
|
|
from tvm.runtime import Object
|
|
|
|
|
|
@tvm.register_global_func("test.vm.move")
|
|
def move(src):
|
|
return src
|
|
|
|
|
|
@tvm.register_global_func("test.vm.add")
|
|
def add(a, b):
|
|
ret = a.numpy() + b.numpy()
|
|
return tvm.runtime.tensor(ret)
|
|
|
|
|
|
@tvm.register_global_func("test.vm.mul")
|
|
def mul(a, b):
|
|
ret = a.numpy() * b.numpy()
|
|
return tvm.runtime.tensor(ret)
|
|
|
|
|
|
@tvm.register_global_func("test.vm.equal_zero")
|
|
def equal_zero(a):
|
|
ret = np.all(a.numpy() == 0)
|
|
return tvm.runtime.tensor(ret)
|
|
|
|
|
|
@tvm.register_global_func("test.vm.subtract_one")
|
|
def subtract_one(a):
|
|
ret = np.subtract(a.numpy(), 1)
|
|
return tvm.runtime.tensor(ret)
|
|
|
|
|
|
@tvm.register_global_func("test.vm.identity")
|
|
def identity_packed(a, b):
|
|
b[:] = tvm.runtime.tensor(a.numpy())
|
|
|
|
|
|
@tvm.register_global_func("test.vm.tile")
|
|
def tile_packed(a, b):
|
|
b[:] = tvm.runtime.tensor(np.tile(a.numpy(), (1, 2)))
|
|
|
|
|
|
@tvm.register_global_func("test.vm.add_scalar")
|
|
def add_scalar(a, b):
|
|
return a + b
|
|
|
|
|
|
@tvm.register_global_func("test.vm.get_device_id")
|
|
def get_device_id(device):
|
|
return device.index
|
|
|
|
|
|
def check_saved_func(vm: relax.VirtualMachine, func_name: str, *inputs: list[Any]) -> Object:
|
|
# uses save_function to create a closure with the given inputs
|
|
# and ensure the result is the same
|
|
# (assumes the functions return tensors and that they're idempotent)
|
|
saved_name = f"{func_name}_saved"
|
|
vm.save_function(func_name, saved_name, *inputs)
|
|
res1 = vm[func_name](*inputs)
|
|
res2 = vm[saved_name]()
|
|
tvm.testing.assert_allclose(res1.numpy(), res2.numpy(), rtol=1e-7, atol=1e-7)
|
|
return res1
|
|
|
|
|
|
@tvm.register_global_func("test.vm.check_if_defined")
|
|
def check_if_defined(obj: tvm.Object) -> tvm.tirx.IntImm:
|
|
return tvm.runtime.convert(obj is not None)
|