This is an automated email from the ASF dual-hosted git repository.
tlopex 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 f7ebbdb176 [REFACTOR] Further cleanup node redirections (#18844)
f7ebbdb176 is described below
commit f7ebbdb17666859430819965647b6317ccd5dbc1
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Feb 28 10:13:54 2026 -0500
[REFACTOR] Further cleanup node redirections (#18844)
## Summary
Remove node/serialization indirection headers and redirect to direct ffi
API calls.
## Changes
- Replace tvm::SaveJSON/LoadJSON wrappers with direct
ffi::ToJSONGraph/FromJSONGraph calls
- Remove C++ MakeNode, redirect Python make_node to
ffi.MakeObjectFromPackedArgs
- Move attr_registry.h to its logical home under src/ir/
---
include/tvm/ir/serialization.h | 55 --------------------
include/tvm/node/serialization.h | 30 -----------
include/tvm/runtime/packed_func.h | 41 ---------------
python/tvm/ir/attrs.py | 8 ++-
src/contrib/msc/core/ir/plugin.h | 2 +-
src/{node => ir}/attr_registry.h | 8 +--
src/ir/op.cc | 2 +-
src/ir/serialization.cc | 4 +-
src/node/reflection.cc | 59 ----------------------
.../backend/adreno/annotate_custom_storage.cc | 1 -
.../backend/adreno/fold_vdevice_scope_change.cc | 1 -
src/relax/transform/alter_op_impl.cc | 4 +-
src/relax/transform/convert_layout.cc | 4 +-
.../specialize_primfunc_based_on_callsite.cc | 1 -
.../contrib/cutlass/fp16_group_gemm_sm100.cu | 1 -
.../cutlass/fp8_groupwise_scaled_gemm_sm100.cu | 1 -
.../cutlass/fp8_groupwise_scaled_gemm_sm90.cu | 1 -
.../fp8_groupwise_scaled_group_gemm_sm100.cu | 1 -
src/runtime/hexagon/rpc/simulator/rpc_server.cc | 1 -
src/s_tir/meta_schedule/database/database.cc | 5 +-
src/s_tir/meta_schedule/database/database_utils.cc | 5 +-
.../mutator/mutate_compute_location.cc | 12 +++--
src/s_tir/meta_schedule/mutator/mutate_parallel.cc | 10 ++--
.../meta_schedule/mutator/mutate_thread_binding.cc | 4 +-
src/s_tir/meta_schedule/utils.h | 7 ++-
src/s_tir/schedule/primitive/blockize_tensorize.cc | 2 +-
.../schedule/primitive/layout_transformation.cc | 12 +++--
src/s_tir/schedule/trace.cc | 2 +-
src/s_tir/schedule/utils.h | 5 +-
src/script/printer/utils.h | 10 +++-
src/target/tag.cc | 2 +-
src/target/target_kind.cc | 2 +-
32 files changed, 69 insertions(+), 234 deletions(-)
diff --git a/include/tvm/ir/serialization.h b/include/tvm/ir/serialization.h
deleted file mode 100644
index 59bdb87067..0000000000
--- a/include/tvm/ir/serialization.h
+++ /dev/null
@@ -1,55 +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 tvm/ir/serialization.h
- * \brief Utility functions for serialization.
- *
- * This is a thin forwarding header to ffi/extra/serialization.h.
- * Prefer using ffi::ToJSONGraph / ffi::FromJSONGraph directly.
- */
-#ifndef TVM_IR_SERIALIZATION_H_
-#define TVM_IR_SERIALIZATION_H_
-
-#include <tvm/ffi/extra/json.h>
-#include <tvm/ffi/extra/serialization.h>
-#include <tvm/runtime/base.h>
-
-#include <string>
-
-namespace tvm {
-
-/*!
- * \brief Save the node as well as all the node it depends on as json.
- * This can be used to serialize any TVM object.
- *
- * \return the string representation of the node.
- */
-TVM_DLL std::string SaveJSON(ffi::Any node);
-
-/*!
- * \brief Load tvm Node object from json and return a shared_ptr of Node.
- * \param json_str The json string to load from.
- *
- * \return The shared_ptr of the Node.
- */
-TVM_DLL ffi::Any LoadJSON(std::string json_str);
-
-} // namespace tvm
-#endif // TVM_IR_SERIALIZATION_H_
diff --git a/include/tvm/node/serialization.h b/include/tvm/node/serialization.h
deleted file mode 100644
index 892ff5fdf1..0000000000
--- a/include/tvm/node/serialization.h
+++ /dev/null
@@ -1,30 +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 tvm/node/serialization.h
- * \brief Forwarding header. Use tvm/ir/serialization.h instead.
- */
-#ifndef TVM_NODE_SERIALIZATION_H_
-#define TVM_NODE_SERIALIZATION_H_
-
-// This header has moved to tvm/ir/serialization.h
-#include <tvm/ir/serialization.h>
-
-#endif // TVM_NODE_SERIALIZATION_H_
diff --git a/include/tvm/runtime/packed_func.h
b/include/tvm/runtime/packed_func.h
deleted file mode 100644
index 3f8ec66bc1..0000000000
--- a/include/tvm/runtime/packed_func.h
+++ /dev/null
@@ -1,41 +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 tvm/runtime/packed_func.h
- * \brief Type-erased function used across TVM API.
- */
-#ifndef TVM_RUNTIME_PACKED_FUNC_H_
-#define TVM_RUNTIME_PACKED_FUNC_H_
-
-#include <tvm/ffi/any.h>
-#include <tvm/ffi/function.h>
-
-namespace tvm {
-namespace runtime {
-
-#define TVM_DLL_EXPORT_TYPED_FUNC TVM_FFI_DLL_EXPORT_TYPED_FUNC
-
-using ffi::Any;
-using ffi::AnyView;
-
-} // namespace runtime
-} // namespace tvm
-
-#endif // TVM_RUNTIME_PACKED_FUNC_H_
diff --git a/python/tvm/ir/attrs.py b/python/tvm/ir/attrs.py
index 82a526ca7f..01b6fe20c1 100644
--- a/python/tvm/ir/attrs.py
+++ b/python/tvm/ir/attrs.py
@@ -17,8 +17,8 @@
"""TVM Attribute module, which is mainly used for defining attributes of
operators."""
import tvm_ffi
+import tvm_ffi._ffi_api as _tvm_ffi_api
-import tvm.runtime._ffi_node_api
from tvm.runtime import Object
from . import _ffi_api
@@ -149,7 +149,11 @@ def make_node(type_key, **kwargs):
assert isinstance(x, tvm.tir.IntImm)
assert x.value == 10
"""
+ if type_key == "ir.DictAttrs":
+ # DictAttrs stores kwargs as a key-value dict, not as named fields.
+ # MakeObjectFromPackedArgs would look for a field named "__dict__".
+ return _tvm_ffi_api.MakeObjectFromPackedArgs("ir.DictAttrs",
"__dict__", kwargs)
args = [type_key]
for k, v in kwargs.items():
args += [k, v]
- return tvm.runtime._ffi_node_api.MakeNode(*args)
+ return _tvm_ffi_api.MakeObjectFromPackedArgs(*args)
diff --git a/src/contrib/msc/core/ir/plugin.h b/src/contrib/msc/core/ir/plugin.h
index 2d2441353f..46b41d4fe6 100644
--- a/src/contrib/msc/core/ir/plugin.h
+++ b/src/contrib/msc/core/ir/plugin.h
@@ -32,7 +32,7 @@
#include <unordered_map>
#include <vector>
-#include "../../../../node/attr_registry.h"
+#include "../../../../ir/attr_registry.h"
#include "../utils.h"
namespace tvm {
diff --git a/src/node/attr_registry.h b/src/ir/attr_registry.h
similarity index 97%
rename from src/node/attr_registry.h
rename to src/ir/attr_registry.h
index b956cb5616..1d1fdb2742 100644
--- a/src/node/attr_registry.h
+++ b/src/ir/attr_registry.h
@@ -18,11 +18,11 @@
*/
/*!
- * \file tvm/node/attr_registry.h
+ * \file tvm/ir/attr_registry.h
* \brief Common global registry for objects that also have additional attrs.
*/
-#ifndef TVM_NODE_ATTR_REGISTRY_H_
-#define TVM_NODE_ATTR_REGISTRY_H_
+#ifndef TVM_IR_ATTR_REGISTRY_H_
+#define TVM_IR_ATTR_REGISTRY_H_
#include <tvm/ffi/function.h>
#include <tvm/node/attr_registry_map.h>
@@ -169,4 +169,4 @@ class AttrRegistry {
};
} // namespace tvm
-#endif // TVM_NODE_ATTR_REGISTRY_H_
+#endif // TVM_IR_ATTR_REGISTRY_H_
diff --git a/src/ir/op.cc b/src/ir/op.cc
index 825437c216..508d4eafcd 100644
--- a/src/ir/op.cc
+++ b/src/ir/op.cc
@@ -30,7 +30,7 @@
#include <memory>
-#include "../node/attr_registry.h"
+#include "attr_registry.h"
namespace tvm {
diff --git a/src/ir/serialization.cc b/src/ir/serialization.cc
index 4d9074e98c..fc080802a3 100644
--- a/src/ir/serialization.cc
+++ b/src/ir/serialization.cc
@@ -28,14 +28,14 @@
namespace tvm {
-std::string SaveJSON(Any n) {
+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);
}
-Any LoadJSON(std::string json_str) {
+static ffi::Any LoadJSON(std::string json_str) {
ffi::json::Value jgraph = ffi::json::Parse(json_str);
return ffi::FromJSONGraph(jgraph);
}
diff --git a/src/node/reflection.cc b/src/node/reflection.cc
deleted file mode 100644
index a5f4ffb84d..0000000000
--- a/src/node/reflection.cc
+++ /dev/null
@@ -1,59 +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.
- */
-
-/*!
- * Reflection utilities.
- * \file node/reflection.cc
- */
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/accessor.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/attrs.h>
-
-namespace tvm {
-
-using ffi::Any;
-using ffi::Function;
-using ffi::PackedArgs;
-
-// API function to make node.
-// args format:
-// key1, value1, ..., key_n, value_n
-void MakeNode(const ffi::PackedArgs& args, ffi::Any* rv) {
- // TODO(tvm-team): consider further simplify by removing DictAttrsNode
special handling
- ffi::String type_key = args[0].cast<ffi::String>();
- int32_t type_index;
- TVMFFIByteArray type_key_array = TVMFFIByteArray{type_key.data(),
type_key.size()};
- TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeKeyToIndex(&type_key_array, &type_index));
- if (type_index == DictAttrsNode::RuntimeTypeIndex()) {
- ObjectPtr<DictAttrsNode> attrs = ffi::make_object<DictAttrsNode>();
- attrs->InitByPackedArgs(args.Slice(1), false);
- *rv = ObjectRef(attrs);
- } else {
- auto fcreate_object =
ffi::Function::GetGlobalRequired("ffi.MakeObjectFromPackedArgs");
- fcreate_object.CallPacked(args, rv);
- }
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def_packed("node.MakeNode", MakeNode);
-}
-
-} // namespace tvm
diff --git a/src/relax/backend/adreno/annotate_custom_storage.cc
b/src/relax/backend/adreno/annotate_custom_storage.cc
index 396e5f9cbb..eda0148779 100644
--- a/src/relax/backend/adreno/annotate_custom_storage.cc
+++ b/src/relax/backend/adreno/annotate_custom_storage.cc
@@ -236,7 +236,6 @@
*
*/
-#include <tvm/ir/serialization.h>
#include <tvm/relax/attrs/op.h>
#include <tvm/relax/backend/adreno/transform.h>
#include <tvm/relax/dataflow_matcher.h>
diff --git a/src/relax/backend/adreno/fold_vdevice_scope_change.cc
b/src/relax/backend/adreno/fold_vdevice_scope_change.cc
index af34a3ac10..f5898021de 100644
--- a/src/relax/backend/adreno/fold_vdevice_scope_change.cc
+++ b/src/relax/backend/adreno/fold_vdevice_scope_change.cc
@@ -23,7 +23,6 @@
* store into global scope avoiding unnecessary device copy.
*/
-#include <tvm/ir/serialization.h>
#include <tvm/relax/attrs/op.h>
#include <tvm/relax/backend/adreno/transform.h>
#include <tvm/relax/dataflow_matcher.h>
diff --git a/src/relax/transform/alter_op_impl.cc
b/src/relax/transform/alter_op_impl.cc
index 93db755059..f324af5a42 100644
--- a/src/relax/transform/alter_op_impl.cc
+++ b/src/relax/transform/alter_op_impl.cc
@@ -24,9 +24,9 @@
* true.
*/
#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/attrs.h>
-#include <tvm/ir/serialization.h>
#include <tvm/relax/analysis.h>
#include <tvm/relax/attrs/manipulate.h>
#include <tvm/relax/expr_functor.h>
@@ -59,7 +59,7 @@ static ffi::Array<PrimExpr> GetShapeFromTensor(const Expr&
expr) {
}
static IndexMap DeepCopyIndexMap(const IndexMap& index_map) {
- return Downcast<IndexMap>(LoadJSON(SaveJSON(index_map)));
+ return Downcast<IndexMap>(ffi::FromJSONGraph(ffi::ToJSONGraph(index_map)));
}
/*! \brief Checks if the \p transform is bijective on the shape of \p expr */
diff --git a/src/relax/transform/convert_layout.cc
b/src/relax/transform/convert_layout.cc
index 888554df67..62d029a334 100644
--- a/src/relax/transform/convert_layout.cc
+++ b/src/relax/transform/convert_layout.cc
@@ -21,8 +21,8 @@
* \brief Automatic layout conversion pass, especially for axis swapping.
*/
+#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/ir/serialization.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/nested_msg.h>
#include <tvm/relax/op_attr_types.h>
@@ -130,7 +130,7 @@ class LayoutConvertMutator : public ExprMutator {
ObjectPtr<LayoutTransformAttrs> attrs =
ffi::make_object<LayoutTransformAttrs>();
ffi::Array<IntImm> axis_separator;
ffi::Array<IntImm> input_axis_separator;
- attrs->index_map = Downcast<IndexMap>(LoadJSON(SaveJSON(index_map)));
+ attrs->index_map =
Downcast<IndexMap>(ffi::FromJSONGraph(ffi::ToJSONGraph(index_map)));
attrs->axis_separators = std::move(axis_separator);
attrs->input_axis_separators = std::move(input_axis_separator);
const Op& layout_transform_op_ = Op::Get("relax.layout_transform");
diff --git a/src/relax/transform/specialize_primfunc_based_on_callsite.cc
b/src/relax/transform/specialize_primfunc_based_on_callsite.cc
index 8a38baedd7..a0dcdaf4a1 100644
--- a/src/relax/transform/specialize_primfunc_based_on_callsite.cc
+++ b/src/relax/transform/specialize_primfunc_based_on_callsite.cc
@@ -21,7 +21,6 @@
* \brief Update PrimFunc buffers based on updated scope (or structure) info.
*/
-#include <tvm/ir/serialization.h>
#include <tvm/relax/attrs/op.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/nested_msg.h>
diff --git a/src/runtime/contrib/cutlass/fp16_group_gemm_sm100.cu
b/src/runtime/contrib/cutlass/fp16_group_gemm_sm100.cu
index 0c9fe0fff1..e7ae7bba75 100644
--- a/src/runtime/contrib/cutlass/fp16_group_gemm_sm100.cu
+++ b/src/runtime/contrib/cutlass/fp16_group_gemm_sm100.cu
@@ -21,7 +21,6 @@
#include <float.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/tensor.h>
#include "fp16_group_gemm.cuh"
diff --git a/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm100.cu
b/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm100.cu
index e8035c172a..be8a2156da 100644
--- a/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm100.cu
+++ b/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm100.cu
@@ -21,7 +21,6 @@
#include <float.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/tensor.h>
#include "../cublas/cublas_utils.h"
diff --git a/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm90.cu
b/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm90.cu
index 3c326e3143..c8490b542c 100644
--- a/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm90.cu
+++ b/src/runtime/contrib/cutlass/fp8_groupwise_scaled_gemm_sm90.cu
@@ -21,7 +21,6 @@
#include <float.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/tensor.h>
#include "../cublas/cublas_utils.h"
diff --git
a/src/runtime/contrib/cutlass/fp8_groupwise_scaled_group_gemm_sm100.cu
b/src/runtime/contrib/cutlass/fp8_groupwise_scaled_group_gemm_sm100.cu
index 4f5dd1e1c7..955a01765c 100644
--- a/src/runtime/contrib/cutlass/fp8_groupwise_scaled_group_gemm_sm100.cu
+++ b/src/runtime/contrib/cutlass/fp8_groupwise_scaled_group_gemm_sm100.cu
@@ -22,7 +22,6 @@
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/packed_func.h>
#include <tvm/runtime/tensor.h>
#include "fp8_groupwise_scaled_group_gemm_runner_sm100.cuh"
diff --git a/src/runtime/hexagon/rpc/simulator/rpc_server.cc
b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
index fba9253e53..2cef7f9c71 100644
--- a/src/runtime/hexagon/rpc/simulator/rpc_server.cc
+++ b/src/runtime/hexagon/rpc/simulator/rpc_server.cc
@@ -33,7 +33,6 @@
#include "../../profiler/prof_utils.h"
#include "hexagon_sim_proto.h"
#include "tvm/ffi/function.h"
-#include "tvm/runtime/packed_func.h"
namespace tvm {
namespace runtime {
diff --git a/src/s_tir/meta_schedule/database/database.cc
b/src/s_tir/meta_schedule/database/database.cc
index c3cd0aca0e..33bf54b429 100644
--- a/src/s_tir/meta_schedule/database/database.cc
+++ b/src/s_tir/meta_schedule/database/database.cc
@@ -43,7 +43,8 @@ Workload::Workload(IRModule mod, Workload::THashCode shash) {
ObjectRef WorkloadNode::AsJSON() const {
// Convert `this->mod` to JSON
- std::string json_mod = tvm::SaveJSON(this->mod);
+ std::string json_mod = ffi::json::Stringify(
+ ffi::ToJSONGraph(this->mod, ffi::json::Object{{"tvm_version",
TVM_VERSION}}), /*indent=*/2);
// Dump the JSON string to base64
std::string b64_mod = Base64Encode(json_mod);
// Output
@@ -62,7 +63,7 @@ Workload Workload::FromJSON(const ObjectRef& json_obj) {
{
ffi::String b64_mod = json_array->at(1).cast<ffi::String>();
std::string json_mod = Base64Decode(b64_mod);
- mod = LoadJSON(json_mod).cast<IRModule>();
+ mod = ffi::FromJSONGraph(ffi::json::Parse(json_mod)).cast<IRModule>();
std::stringstream(str_shash) >> shash;
}
} catch (const std::runtime_error& e) { // includes tvm::Error and
dmlc::Error
diff --git a/src/s_tir/meta_schedule/database/database_utils.cc
b/src/s_tir/meta_schedule/database/database_utils.cc
index 120c380aa2..6ef1f254ae 100644
--- a/src/s_tir/meta_schedule/database/database_utils.cc
+++ b/src/s_tir/meta_schedule/database/database_utils.cc
@@ -82,7 +82,10 @@ void JSONDumps(Any json_obj, std::ostringstream& os) {
}
os << "}";
} else if (json_obj.as<tir::IndexMapNode>()) {
- JSONDumps(ffi::String(SaveJSON(json_obj)), os);
+ JSONDumps(ffi::String(ffi::json::Stringify(
+ ffi::ToJSONGraph(json_obj, ffi::json::Object{{"tvm_version",
TVM_VERSION}}),
+ /*indent=*/2)),
+ os);
} else {
TVM_FFI_THROW(TypeError) << "Unsupported type in JSON object: " <<
json_obj.GetTypeKey();
}
diff --git a/src/s_tir/meta_schedule/mutator/mutate_compute_location.cc
b/src/s_tir/meta_schedule/mutator/mutate_compute_location.cc
index 67fd867e48..86f3c5c1c1 100644
--- a/src/s_tir/meta_schedule/mutator/mutate_compute_location.cc
+++ b/src/s_tir/meta_schedule/mutator/mutate_compute_location.cc
@@ -44,7 +44,9 @@ class MutateComputeLocationNode : public MutatorNode {
public:
// Inherit from `MutatorNode`
void InitializeWithTuneContext(const TuneContext& context) final {
- this->json_mod_ = SaveJSON(context->mod.value());
+ this->json_mod_ = ffi::json::Stringify(
+ ffi::ToJSONGraph(context->mod.value(),
ffi::json::Object{{"tvm_version", TVM_VERSION}}),
+ /*indent=*/2);
}
// Inherit from `MutatorNode`
ffi::Optional<Trace> Apply(const Trace& trace, TRandState* rand_state) final;
@@ -76,10 +78,10 @@ class MutateComputeLocationNode : public MutatorNode {
*/
std::vector<MutateComputeLocationNode::Candidate>
MutateComputeLocationNode::FindCandidates(
const Trace& trace, TRandState* rand_state) {
- s_tir::Schedule sch = s_tir::Schedule::Traced( //
- /*mod=*/LoadJSON(this->json_mod_).cast<IRModule>(), //
- /*rand_state=*/ForkSeed(rand_state), //
- /*debug_mode=*/0, //
+ s_tir::Schedule sch = s_tir::Schedule::Traced(
//
+
/*mod=*/ffi::FromJSONGraph(ffi::json::Parse(this->json_mod_)).cast<IRModule>(),
//
+ /*rand_state=*/ForkSeed(rand_state),
//
+ /*debug_mode=*/0,
//
/*error_render_level=*/s_tir::ScheduleErrorRenderLevel::kNone);
static InstructionKind inst_sample_compute_location =
diff --git a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc
b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc
index fa540b7231..3365b7f64f 100644
--- a/src/s_tir/meta_schedule/mutator/mutate_parallel.cc
+++ b/src/s_tir/meta_schedule/mutator/mutate_parallel.cc
@@ -189,7 +189,9 @@ class MutateParallelNode : public MutatorNode {
void InitializeWithTuneContext(const TuneContext& context) final {
Target target = context->target.value();
this->max_parallel_extent_ = GetTargetNumCores(target) *
this->max_jobs_per_core;
- this->json_mod_ = SaveJSON(context->mod.value());
+ this->json_mod_ = ffi::json::Stringify(
+ ffi::ToJSONGraph(context->mod.value(),
ffi::json::Object{{"tvm_version", TVM_VERSION}}),
+ /*indent=*/2);
}
// Inherit from `MutatorNode`
ffi::Optional<Trace> Apply(const Trace& trace, TRandState* rand_state) final;
@@ -258,9 +260,9 @@ ffi::Optional<Trace> MutateParallelNode::Apply(const Trace&
trace, TRandState* r
return std::nullopt;
}
// Step 2. Replay the instructions to recover loop extents
- s_tir::Schedule sch = s_tir::Schedule::Traced( //
- /*mod=*/LoadJSON(this->json_mod_).cast<IRModule>(), //
- /*rand_state=*/ForkSeed(rand_state), //
+ s_tir::Schedule sch = s_tir::Schedule::Traced(
//
+
/*mod=*/ffi::FromJSONGraph(ffi::json::Parse(this->json_mod_)).cast<IRModule>(),
//
+ /*rand_state=*/ForkSeed(rand_state),
//
/*debug_mode=*/0,
/*error_render_level=*/s_tir::ScheduleErrorRenderLevel::kNone);
trace->ApplyToSchedule(sch, /*remove_postproc=*/true);
diff --git a/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc
b/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc
index d1310907b0..9386d106c1 100644
--- a/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc
+++ b/src/s_tir/meta_schedule/mutator/mutate_thread_binding.cc
@@ -44,7 +44,9 @@ class MutateThreadBindingNode : public MutatorNode {
public:
// Inherit from `MutatorNode`
void InitializeWithTuneContext(const TuneContext& context) final {
- this->json_mod_ = SaveJSON(context->mod.value());
+ this->json_mod_ = ffi::json::Stringify(
+ ffi::ToJSONGraph(context->mod.value(),
ffi::json::Object{{"tvm_version", TVM_VERSION}}),
+ /*indent=*/2);
}
// Inherit from `MutatorNode`
ffi::Optional<Trace> Apply(const Trace& trace, TRandState* rand_state) final;
diff --git a/src/s_tir/meta_schedule/utils.h b/src/s_tir/meta_schedule/utils.h
index d5569d07ec..f2938c9068 100644
--- a/src/s_tir/meta_schedule/utils.h
+++ b/src/s_tir/meta_schedule/utils.h
@@ -20,8 +20,9 @@
#define TVM_S_TIR_META_SCHEDULE_UTILS_H_
#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/serialization.h>
#include <tvm/ffi/optional.h>
-#include <tvm/ir/serialization.h>
#include <tvm/node/cast.h>
#include <tvm/runtime/object.h>
#include <tvm/s_tir/meta_schedule/arg_info.h>
@@ -268,7 +269,9 @@ inline
std::vector<support::LinearCongruentialEngine::TRandState> ForkSeed(
* \param mod The IRModule to make a deep copy.
* \return The deep copy of the IRModule.
*/
-inline IRModule DeepCopyIRModule(IRModule mod) { return
LoadJSON(SaveJSON(mod)).cast<IRModule>(); }
+inline IRModule DeepCopyIRModule(IRModule mod) {
+ return ffi::FromJSONGraph(ffi::ToJSONGraph(mod)).cast<IRModule>();
+}
/*!
* \brief Concatenate strings
diff --git a/src/s_tir/schedule/primitive/blockize_tensorize.cc
b/src/s_tir/schedule/primitive/blockize_tensorize.cc
index f9631b96d4..5c6c152925 100644
--- a/src/s_tir/schedule/primitive/blockize_tensorize.cc
+++ b/src/s_tir/schedule/primitive/blockize_tensorize.cc
@@ -39,7 +39,7 @@ Range RangeFromExtent(const PrimExpr& extent) {
template <class T>
T DeepCopy(const T& stmt) {
- return Downcast<T>(LoadJSON(SaveJSON(stmt)));
+ return Downcast<T>(ffi::FromJSONGraph(ffi::ToJSONGraph(stmt)));
}
/*!
diff --git a/src/s_tir/schedule/primitive/layout_transformation.cc
b/src/s_tir/schedule/primitive/layout_transformation.cc
index 2d8629c06f..08d4de6f82 100644
--- a/src/s_tir/schedule/primitive/layout_transformation.cc
+++ b/src/s_tir/schedule/primitive/layout_transformation.cc
@@ -1605,7 +1605,9 @@ struct TransformLayoutTraits : public
UnpackedInstTraits<TransformLayoutTraits>
attrs_record.push_back(attrs[0]);
attrs_record.push_back(attrs[1]);
if (attrs[2] != nullptr) {
- attrs_record.push_back(ffi::String(::tvm::SaveJSON(attrs[2])));
+ attrs_record.push_back(ffi::String(ffi::json::Stringify(
+ ffi::ToJSONGraph(attrs[2], ffi::json::Object{{"tvm_version",
TVM_VERSION}}),
+ /*indent=*/2)));
} else {
attrs_record.push_back(attrs[2]);
}
@@ -1619,7 +1621,7 @@ struct TransformLayoutTraits : public
UnpackedInstTraits<TransformLayoutTraits>
attrs.push_back(attrs_record[0]);
attrs.push_back(attrs_record[1]);
if (attrs_record[2] != nullptr) {
- attrs.push_back(::tvm::LoadJSON(Downcast<ffi::String>(attrs_record[2])));
+
attrs.push_back(ffi::FromJSONGraph(ffi::json::Parse(Downcast<ffi::String>(attrs_record[2]))));
} else {
attrs.push_back(attrs_record[2]);
}
@@ -1656,14 +1658,16 @@ struct TransformBlockLayoutTraits : public
UnpackedInstTraits<TransformBlockLayo
static ObjectRef AttrsAsJSON(const ffi::Array<Any>& attrs) {
ffi::Array<Any> attrs_record;
attrs_record.reserve(kNumAttrs);
- attrs_record.push_back(ffi::String(::tvm::SaveJSON(attrs[0])));
+ attrs_record.push_back(ffi::String(ffi::json::Stringify(
+ ffi::ToJSONGraph(attrs[0], ffi::json::Object{{"tvm_version",
TVM_VERSION}}),
+ /*indent=*/2)));
return attrs_record;
}
static ffi::Array<Any> AttrsFromJSON(const ObjectRef& attrs_record_) {
ffi::Array<Any> attrs_record = Downcast<ffi::Array<Any>>(attrs_record_);
ffi::Array<Any> attrs;
- attrs.push_back(::tvm::LoadJSON(Downcast<ffi::String>(attrs_record[0])));
+
attrs.push_back(ffi::FromJSONGraph(ffi::json::Parse(Downcast<ffi::String>(attrs_record[0]))));
return attrs;
}
diff --git a/src/s_tir/schedule/trace.cc b/src/s_tir/schedule/trace.cc
index 24e4636dae..6114a275a5 100644
--- a/src/s_tir/schedule/trace.cc
+++ b/src/s_tir/schedule/trace.cc
@@ -188,7 +188,7 @@ ffi::Array<Any> TranslateInputRVs(const ffi::Array<Any>&
inputs,
const char* name = (*opt_str).data();
int64_t size = (*opt_str).size();
if (name[0] == '{' && name[size - 1] == '}') {
- Any obj = LoadJSON(name);
+ Any obj = ffi::FromJSONGraph(ffi::json::Parse(name));
// Case 6. IndexMap
if (obj.as<IndexMapNode>()) {
IndexMap index_map = Downcast<IndexMap>(obj);
diff --git a/src/s_tir/schedule/utils.h b/src/s_tir/schedule/utils.h
index d8aebb2f6d..b1c3903eb0 100644
--- a/src/s_tir/schedule/utils.h
+++ b/src/s_tir/schedule/utils.h
@@ -22,7 +22,8 @@
#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_set.h>
#include <tvm/arith/iter_affine_map.h>
-#include <tvm/ir/serialization.h>
+#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/serialization.h>
#include <tvm/s_tir/schedule/instruction.h>
#include <tvm/s_tir/schedule/schedule.h>
#include <tvm/s_tir/schedule/state.h>
@@ -39,7 +40,7 @@
#include <utility>
#include "../../arith/pattern_match.h"
-#include "../../node/attr_registry.h"
+#include "../../ir/attr_registry.h"
#include "../../runtime/thread_storage_scope.h"
#include "../../support/array.h"
#include "../../support/nd_int_set.h"
diff --git a/src/script/printer/utils.h b/src/script/printer/utils.h
index c0dbd2e46c..ddeab5e754 100644
--- a/src/script/printer/utils.h
+++ b/src/script/printer/utils.h
@@ -19,7 +19,9 @@
#ifndef TVM_SCRIPT_PRINTER_UTILS_H_
#define TVM_SCRIPT_PRINTER_UTILS_H_
-#include <tvm/ir/serialization.h>
+#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/serialization.h>
+#include <tvm/runtime/base.h>
#include <tvm/script/printer/ir_docsifier.h>
#include <string>
@@ -72,7 +74,11 @@ inline std::string Docsify(const ObjectRef& obj, const
IRDocsifier& d, const Fra
if (d->cfg->show_meta) {
os << "metadata = tvm.ir.load_json(\"\"\""
<< support::StrEscape(
- SaveJSON(ffi::Map<ffi::String, ffi::Any>(d->metadata.begin(),
d->metadata.end())),
+ ffi::json::Stringify(
+ ffi::ToJSONGraph(
+ ffi::Map<ffi::String, ffi::Any>(d->metadata.begin(),
d->metadata.end()),
+ ffi::json::Object{{"tvm_version", TVM_VERSION}}),
+ 2),
false, false)
<< "\"\"\")\n";
} else {
diff --git a/src/target/tag.cc b/src/target/tag.cc
index c41d077084..74fa65b0e6 100644
--- a/src/target/tag.cc
+++ b/src/target/tag.cc
@@ -28,7 +28,7 @@
#include <tvm/target/tag.h>
#include <tvm/target/target.h>
-#include "../node/attr_registry.h"
+#include "../ir/attr_registry.h"
namespace tvm {
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 46e36d1019..8c328dd9cf 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -30,7 +30,7 @@
#include <algorithm>
-#include "../node/attr_registry.h"
+#include "../ir/attr_registry.h"
#include "../support/utils.h"
#include "./canonicalizer/llvm/canonicalize.h"
