This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 1382707e8f [FFI][IR] Route JSON serialization through tvm-ffi (#19662)
1382707e8f is described below
commit 1382707e8fde785872283d517c5790d1e9f5cfcb
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Jun 3 15:58:07 2026 -0400
[FFI][IR] Route JSON serialization through tvm-ffi (#19662)
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`/`ffi::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.
Validated locally in an isolated worktree build with `ninja -C build
tvm_compiler tvm_runtime_extra`, targeted IR/target tests,
`tests/python/disco/test_session.py::test_string_obj`, import smoke, and
touched-file pre-commit.
---
python/tvm/ir/base.py | 8 +++---
src/ir/serialization.cc | 47 ---------------------------------
src/runtime/extra/disco/protocol.h | 12 +++------
tests/python/ir/test_node_reflection.py | 8 ++++++
4 files changed, 17 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/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/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)
