This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tvm-direct-json-graph-serialization-apis in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 6a6f6639a5271df9803b421cae9c64be39bfefdc Author: Tianqi Chen <[email protected]> AuthorDate: Wed Jun 3 14:37:29 2026 +0000 [FFI][IR] Route JSON serialization through tvm-ffi TVM can rely on tvm-ffi's JSON graph serialization helpers directly instead of routing through TVM-side node.SaveJSON/node.LoadJSON registry entries. This changes tvm.ir save/load to call tvm_ffi.serialization with tvm_version metadata, removes the C++ registry wrapper, and moves the disco debug object path to ffi::ToJSONGraph/FromJSONGraph plus JSON parse/stringify. The disco Python wrappers now declare Python attribute storage explicitly for DRef and Session so DPackedFunc/DModule and method caches continue to work with the current tvm-ffi object model. The socket address helper also normalizes localhost consistently across constructors so the disco socket debug round-trip can bind an IPv4 socket when localhost resolves to IPv6 first. --- python/tvm/ir/base.py | 8 +++--- python/tvm/runtime/disco/session.py | 4 +++ src/ir/serialization.cc | 47 --------------------------------- src/runtime/extra/disco/protocol.h | 12 +++------ src/support/socket.h | 3 +++ tests/python/ir/test_node_reflection.py | 8 ++++++ 6 files changed, 24 insertions(+), 58 deletions(-) diff --git a/python/tvm/ir/base.py b/python/tvm/ir/base.py index cff43bb8c1..ceccb401f4 100644 --- a/python/tvm/ir/base.py +++ b/python/tvm/ir/base.py @@ -18,9 +18,11 @@ import tvm_ffi from tvm_ffi import get_global_func, register_object +from tvm_ffi.serialization import from_json_graph_str, to_json_graph_str -from tvm.runtime import Object, _ffi_node_api +from tvm.runtime import Object +from ..base import __version__ from . import _ffi_api, json_compact @@ -141,7 +143,7 @@ def load_json(json_str) -> Object: """ json_str = json_compact.upgrade_json(json_str) - return _ffi_node_api.LoadJSON(json_str) + return from_json_graph_str(json_str) def save_json(node) -> str: @@ -157,7 +159,7 @@ def save_json(node) -> str: json_str : str Saved json string. """ - return _ffi_node_api.SaveJSON(node) + return to_json_graph_str(node, {"tvm_version": __version__}) def structural_equal(lhs, rhs, map_free_vars=False): diff --git a/python/tvm/runtime/disco/session.py b/python/tvm/runtime/disco/session.py index 08afbbfc80..1615f07412 100644 --- a/python/tvm/runtime/disco/session.py +++ b/python/tvm/runtime/disco/session.py @@ -40,6 +40,8 @@ class DRef(Object): to each object, and the worker process uses this id to refer to the object residing on itself. """ + __slots__ = ("__dict__",) + def debug_get_from_remote(self, worker_id: int) -> Any: """Get the value of a DRef from a remote worker. It is only used for debugging purposes. @@ -103,6 +105,8 @@ class Session(Object): """A Disco interactive session. It allows users to interact with the Disco command queue with various PackedFunc calling convention.""" + __slots__ = ("__dict__",) + def _get_cached_method(self, name: str) -> Callable: if "_cache" not in self.__dict__: cache = self._cache = {} # pylint: disable=attribute-defined-outside-init diff --git a/src/ir/serialization.cc b/src/ir/serialization.cc deleted file mode 100644 index fc080802a3..0000000000 --- a/src/ir/serialization.cc +++ /dev/null @@ -1,47 +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. - */ - -/*! - * \file src/ir/serialization.cc - * \brief Utilities to serialize TVM AST/IR objects. - */ -#include <tvm/ffi/extra/json.h> -#include <tvm/ffi/extra/serialization.h> -#include <tvm/ffi/reflection/registry.h> -#include <tvm/runtime/base.h> - -namespace tvm { - -static std::string SaveJSON(ffi::Any n) { - int indent = 2; - ffi::json::Object metadata{{"tvm_version", TVM_VERSION}}; - ffi::json::Value jgraph = ffi::ToJSONGraph(n, metadata); - return ffi::json::Stringify(jgraph, indent); -} - -static ffi::Any LoadJSON(std::string json_str) { - ffi::json::Value jgraph = ffi::json::Parse(json_str); - return ffi::FromJSONGraph(jgraph); -} - -TVM_FFI_STATIC_INIT_BLOCK() { - namespace refl = tvm::ffi::reflection; - refl::GlobalDef().def("node.SaveJSON", SaveJSON).def("node.LoadJSON", LoadJSON); -} -} // namespace tvm diff --git a/src/runtime/extra/disco/protocol.h b/src/runtime/extra/disco/protocol.h index 89905b7ce1..25662051dc 100644 --- a/src/runtime/extra/disco/protocol.h +++ b/src/runtime/extra/disco/protocol.h @@ -19,7 +19,8 @@ #ifndef TVM_RUNTIME_DISCO_PROTOCOL_H_ #define TVM_RUNTIME_DISCO_PROTOCOL_H_ -#include <tvm/ffi/function.h> +#include <tvm/ffi/extra/json.h> +#include <tvm/ffi/extra/serialization.h> #include <tvm/runtime/base.h> #include <tvm/runtime/disco/session.h> #include <tvm/support/io.h> @@ -233,10 +234,7 @@ inline std::string DiscoDebugObject::SaveToStr() const { return result; } else if (auto opt_obj = this->data.as<ffi::ObjectRef>()) { ffi::ObjectRef obj = opt_obj.value(); - const auto f = tvm::ffi::Function::GetGlobal("node.SaveJSON"); - TVM_FFI_CHECK(f.has_value(), ValueError) - << "Cannot serialize object in non-debugging mode: " << obj->GetTypeKey(); - std::string result = (*f)(obj).cast<std::string>(); + std::string result = ffi::json::Stringify(ffi::ToJSONGraph(obj)); result.push_back('0'); return result; } @@ -251,9 +249,7 @@ inline ffi::ObjectPtr<DiscoDebugObject> DiscoDebugObject::LoadFromStr(std::strin json_str.pop_back(); ffi::ObjectPtr<DiscoDebugObject> result = ffi::make_object<DiscoDebugObject>(); if (control_bit == '0') { - const auto f = tvm::ffi::Function::GetGlobal("node.LoadJSON"); - TVM_FFI_CHECK(f.has_value(), ValueError) << "Cannot deserialize object in non-debugging mode"; - result->data = (*f)(json_str); + result->data = ffi::FromJSONGraph(ffi::json::Parse(json_str)); } else if (control_bit == '1') { support::BytesInStream mstrm(json_str); support::Base64InStream b64strm(&mstrm); diff --git a/src/support/socket.h b/src/support/socket.h index 3647054497..813dcebc06 100644 --- a/src/support/socket.h +++ b/src/support/socket.h @@ -133,6 +133,9 @@ struct SockAddr { * \param port the port of address */ void Set(const char* host, int port) { + if (strcmp(host, "localhost") == 0) { + host = "127.0.0.1"; + } addrinfo hints; memset(&hints, 0, sizeof(hints)); hints.ai_family = PF_UNSPEC; diff --git a/tests/python/ir/test_node_reflection.py b/tests/python/ir/test_node_reflection.py index 6cfff9d184..dce3fbffeb 100644 --- a/tests/python/ir/test_node_reflection.py +++ b/tests/python/ir/test_node_reflection.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. # ruff: noqa: E712, F401, F841 +import json import sys import numpy as np @@ -37,6 +38,13 @@ def test_const_saveload_json(): tvm.ir.assert_structural_equal(zz, z, map_free_vars=True) +def test_save_json_metadata_version(): + obj = tvm.runtime.convert([1, 2]) + json_str = tvm.ir.save_json(obj) + assert json.loads(json_str)["metadata"]["tvm_version"] == tvm.__version__ + assert list(tvm.ir.load_json(json_str)) == [1, 2] + + def _test_infinity_value(value, dtype): x = tvm.tirx.const(value, dtype) json_str = tvm.ir.save_json(x)
