This is an automated email from the ASF dual-hosted git repository.
tqchen 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 33b01a8e80 [REFACTOR][Runtime] Clarify file helpers and VM parameter
ownership (#20679)
33b01a8e80 is described below
commit 33b01a8e8038fc5647a7c706e09763cb625814f3
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 17:49:35 2026 -0400
[REFACTOR][Runtime] Clarify file helpers and VM parameter ownership (#20679)
Move generic binary file I/O to inline support helpers named
ReadBytesFromFile and WriteBytesToFile. Place parameter serialization in
VM-owned C++ code and the Python runtime.vm.param module, with public VM
and parameter exports in the runtime.vm package.
Keep FunctionInfo reflection and JSON conversion with runtime metadata,
and use existing module serialization instead of source metadata
sidecars while preserving raw source and binary writes. Update RPC,
module, backend, tensor-cache, and amalgamated runtime consumers for the
new owners; parameter serialization retains its existing binary format.
---
apps/android_rpc/app/src/main/jni/tvm_runtime.h | 3 +-
apps/ios_rpc/tvmrpc/TVMRuntime.mm | 4 +-
docs/arch/relax_vm.rst | 4 +-
docs/reference/api/python/relax/relax.rst | 12 +-
docs/reference/api/python/runtime/vm.rst | 1 +
include/tvm/support/file_utils.h | 66 ++++++
python/tvm/runtime/__init__.py | 6 -
python/tvm/runtime/{ => vm}/__init__.py | 41 ++--
python/tvm/runtime/{params.py => vm/param.py} | 25 +-
python/tvm/runtime/{ => vm}/vm.py | 0
.../s_tir/meta_schedule/builder/local_builder.py | 3 +-
.../hexagon/codegen/llvm/codegen_hexagon.cc | 5 +-
src/backend/metal/runtime/metal_module.mm | 1 -
src/backend/opencl/runtime/opencl_common.h | 1 -
src/backend/webgpu/codegen/codegen_webgpu.cc | 1 -
src/runtime/extra/contrib/clml/clml_runtime.cc | 6 +-
src/runtime/extra/contrib/clml/clml_runtime.h | 2 +-
.../extra/contrib/tensorrt/tensorrt_runtime.cc | 14 +-
src/runtime/extra/disco/loader.cc | 6 +-
src/runtime/file_utils.cc | 259 ---------------------
src/runtime/file_utils.h | 177 --------------
src/runtime/metadata.cc | 66 ++++++
src/runtime/module.cc | 2 -
src/runtime/rpc/rpc_server_env.cc | 15 +-
src/runtime/static_library.cc | 6 +-
src/runtime/vm/executable.cc | 7 +-
src/runtime/vm/params.cc | 169 ++++++++++++++
src/runtime/vm/tensor_cache_support.cc | 8 +-
src/target/llvm/llvm_module.cc | 6 +-
src/target/source/source_module.cc | 13 +-
web/emcc/wasm_runtime.cc | 3 +-
web/emcc/webgpu_runtime.cc | 1 -
32 files changed, 394 insertions(+), 539 deletions(-)
diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
index fa08f82e31..fb45e61e56 100644
--- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -45,8 +45,8 @@
#include "../3rdparty/tvm-ffi/src/ffi/tensor.cc"
#include "../src/runtime/cpu_device_api.cc"
#include "../src/runtime/device_api.cc"
-#include "../src/runtime/file_utils.cc"
#include "../src/runtime/logging.cc"
+#include "../src/runtime/metadata.cc"
#include "../src/runtime/registry.cc"
#include "../src/runtime/rpc/rpc_channel.cc"
#include "../src/runtime/rpc/rpc_endpoint.cc"
@@ -60,6 +60,7 @@
#include "../src/runtime/thread_pool.cc"
#include "../src/runtime/threading_backend.cc"
#include "../src/runtime/timer.cc"
+#include "../src/runtime/vm/params.cc"
#include "../src/runtime/workspace_pool.cc"
#ifdef TVM_OPENCL_RUNTIME
diff --git a/apps/ios_rpc/tvmrpc/TVMRuntime.mm
b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
index 65d6d76303..f4866daa4d 100644
--- a/apps/ios_rpc/tvmrpc/TVMRuntime.mm
+++ b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
@@ -29,7 +29,7 @@
#include "RPCArgs.h"
// internal TVM header
-#include <../../../src/runtime/file_utils.h>
+#include <tvm/support/file_utils.h>
#if defined(USE_CUSTOM_DSO_LOADER) && USE_CUSTOM_DSO_LOADER == 1
// internal TVM header to achieve Library class
@@ -63,7 +63,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
})
.def_packed("tvm.rpc.server.load_module", [](ffi::PackedArgs args,
ffi::Any* rv) {
auto name = args[0].cast<std::string>();
- std::string fmt = GetFileFormat(name, "");
+ std::string fmt = support::GetFileFormat(name, "");
NSString* base;
if (fmt == "dylib") {
// only load dylib from frameworks.
diff --git a/docs/arch/relax_vm.rst b/docs/arch/relax_vm.rst
index 9bf3580786..22fa1ec4a9 100644
--- a/docs/arch/relax_vm.rst
+++ b/docs/arch/relax_vm.rst
@@ -409,7 +409,9 @@ Source Code Map
- Built-in operations (shape matching, allocation)
* - ``src/relax/backend/vm/codegen_vm.cc``
- CodeGenVM: Relax IR → bytecode
- * - ``python/tvm/runtime/vm.py``
+ * - ``python/tvm/runtime/vm/vm.py``
- Python VirtualMachine wrapper
+ * - ``python/tvm/runtime/vm/param.py``
+ - Parameter dictionary serialization
* - ``python/tvm/relax/vm_build.py``
- ``relax.build()`` and VMExecutable Python class
diff --git a/docs/reference/api/python/relax/relax.rst
b/docs/reference/api/python/relax/relax.rst
index 377e7bc99d..c6d71319d0 100644
--- a/docs/reference/api/python/relax/relax.rst
+++ b/docs/reference/api/python/relax/relax.rst
@@ -21,4 +21,14 @@ tvm.relax
.. automodule:: tvm.relax
:members:
:imported-members:
- :exclude-members: AnyType, ObjectType, BlockBuilder, Call, Tuple,
TupleGetItem, Var, Location, GlobalVar, SourceName, TupleType, Type, FuncType
+ :exclude-members: AnyType, ObjectType, BlockBuilder, Call, Tuple,
TupleGetItem, Var, Location, GlobalVar, SourceName, TupleType, Type, FuncType,
VirtualMachine, VMInstrumentReturnKind
+
+.. py:class:: VirtualMachine
+ :canonical: tvm.runtime.vm.VirtualMachine
+
+ Alias of :py:class:`tvm.runtime.vm.VirtualMachine`.
+
+.. py:class:: VMInstrumentReturnKind
+ :canonical: tvm.runtime.vm.VMInstrumentReturnKind
+
+ Alias of :py:class:`tvm.runtime.vm.VMInstrumentReturnKind`.
diff --git a/docs/reference/api/python/runtime/vm.rst
b/docs/reference/api/python/runtime/vm.rst
index 2d7b980f86..ce3237075f 100644
--- a/docs/reference/api/python/runtime/vm.rst
+++ b/docs/reference/api/python/runtime/vm.rst
@@ -19,3 +19,4 @@ tvm.runtime.vm
--------------------
.. automodule:: tvm.runtime.vm
:members:
+ :imported-members:
diff --git a/include/tvm/support/file_utils.h b/include/tvm/support/file_utils.h
new file mode 100644
index 0000000000..53fc33c53a
--- /dev/null
+++ b/include/tvm/support/file_utils.h
@@ -0,0 +1,66 @@
+/*
+ * 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/support/file_utils.h
+ * \brief Inline file format and binary I/O helpers.
+ */
+#ifndef TVM_SUPPORT_FILE_UTILS_H_
+#define TVM_SUPPORT_FILE_UTILS_H_
+
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/string.h>
+
+#include <fstream>
+#include <string>
+#include <utility>
+
+namespace tvm {
+namespace support {
+
+/*! \brief Get the explicit format, or the extension of the file path. */
+inline ffi::String GetFileFormat(const ffi::String& path, const ffi::String&
format) {
+ if (!format.empty()) return format;
+ std::string file_name = path;
+ size_t pos = file_name.find_last_of('.');
+ return pos == std::string::npos ? ffi::String() :
ffi::String(file_name.substr(pos + 1));
+}
+
+/*! \brief Read a file as binary bytes. */
+inline ffi::Bytes ReadBytesFromFile(const ffi::String& path) {
+ std::ifstream fs(path.c_str(), std::ios::in | std::ios::binary);
+ TVM_FFI_ICHECK(!fs.fail()) << "Cannot open " << path;
+ fs.seekg(0, std::ios::end);
+ size_t size = static_cast<size_t>(fs.tellg());
+ fs.seekg(0, std::ios::beg);
+ std::string data(size, '\0');
+ fs.read(data.data(), size);
+ return ffi::Bytes(std::move(data));
+}
+
+/*! \brief Write binary bytes to a file. */
+inline void WriteBytesToFile(const ffi::String& path, const ffi::Bytes& data) {
+ std::ofstream fs(path.c_str(), std::ios::out | std::ios::binary);
+ TVM_FFI_ICHECK(!fs.fail()) << "Cannot open " << path;
+ fs.write(data.data(), data.size());
+}
+
+} // namespace support
+} // namespace tvm
+#endif // TVM_SUPPORT_FILE_UTILS_H_
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/__init__.py
index 7b950012e4..035e1ed492 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/__init__.py
@@ -31,12 +31,6 @@ from .executable import Executable
from ._tensor import device, cpu, cuda, opencl, vulkan, metal
from ._tensor import vpi, rocm, ext_dev, from_dlpack
from .module import load_module, enabled, system_lib, load_static_library,
num_threads
-from .params import (
- save_param_dict,
- load_param_dict,
- save_param_dict_to_file,
- load_param_dict_from_file,
-)
try:
from . import disco
diff --git a/python/tvm/runtime/__init__.py b/python/tvm/runtime/vm/__init__.py
similarity index 51%
copy from python/tvm/runtime/__init__.py
copy to python/tvm/runtime/vm/__init__.py
index 7b950012e4..7350da0091 100644
--- a/python/tvm/runtime/__init__.py
+++ b/python/tvm/runtime/vm/__init__.py
@@ -1,4 +1,3 @@
-# isort: skip_file
# 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
@@ -15,34 +14,22 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""TVM runtime namespace."""
+# pylint: disable=invalid-name, redefined-builtin, no-else-return,
consider-using-dict-items
+"""Virtual machine execution and parameter serialization."""
-from tvm_ffi import convert, Object
-from tvm_ffi._dtype import dtype as DataType, DataTypeCode
-
-# class exposures
-from .object_generic import ObjectConvertible
-from .device import Device
-from ._tensor import Tensor, tensor, empty
-from .module import Module
-from .executable import Executable
-
-# function exposures
-from ._tensor import device, cpu, cuda, opencl, vulkan, metal
-from ._tensor import vpi, rocm, ext_dev, from_dlpack
-from .module import load_module, enabled, system_lib, load_static_library,
num_threads
-from .params import (
- save_param_dict,
+from .param import (
load_param_dict,
- save_param_dict_to_file,
load_param_dict_from_file,
+ save_param_dict,
+ save_param_dict_to_file,
)
+from .vm import VirtualMachine, VMInstrumentReturnKind
-try:
- from . import disco
-except (ImportError, ValueError):
- # disco C++ runtime is in libtvm_runtime_extra which may not be present.
- # Make the disco module optional.
- disco = None # type: ignore[assignment]
-
-from tvm_ffi import Shape as ShapeTuple
+__all__ = [
+ "VMInstrumentReturnKind",
+ "VirtualMachine",
+ "load_param_dict",
+ "load_param_dict_from_file",
+ "save_param_dict",
+ "save_param_dict_to_file",
+]
diff --git a/python/tvm/runtime/params.py b/python/tvm/runtime/vm/param.py
similarity index 77%
rename from python/tvm/runtime/params.py
rename to python/tvm/runtime/vm/param.py
index 1ea1c64262..88567d6542 100644
--- a/python/tvm/runtime/params.py
+++ b/python/tvm/runtime/vm/param.py
@@ -14,10 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=invalid-name
-"""Helper utility to save and load parameter dicts."""
+# pylint: disable=invalid-name, redefined-builtin, no-else-return,
consider-using-dict-items
+"""VM parameter dictionary serialization."""
-from . import Tensor, _ffi_api, tensor
+from tvm_ffi import get_global_func
+
+from .. import Tensor, tensor
def _to_tensor(params):
@@ -35,8 +37,7 @@ def _to_tensor(params):
def save_param_dict(params):
"""Save parameter dictionary to binary bytes.
- The result binary bytes can be loaded by the
- GraphModule with API "load_params".
+ The resulting bytes can be loaded with :func:`load_param_dict`.
Parameters
----------
@@ -55,12 +56,12 @@ def save_param_dict(params):
# set up the parameter dict
params = {"param0": arr0, "param1": arr1}
# save the parameters as byte array
- param_bytes = tvm.runtime.save_param_dict(params)
+ param_bytes = tvm.runtime.vm.save_param_dict(params)
# We can serialize the param_bytes and load it back later.
- # Pass in byte array to module to directly set parameters
- tvm.runtime.load_param_dict(param_bytes)
+ # Recover the parameter dictionary from the serialized bytes.
+ tvm.runtime.vm.load_param_dict(param_bytes)
"""
- return _ffi_api.SaveParams(_to_tensor(params))
+ return get_global_func("vm.SaveParams")(_to_tensor(params))
def save_param_dict_to_file(params, path):
@@ -74,7 +75,7 @@ def save_param_dict_to_file(params, path):
path: str
The path to the parameter file.
"""
- return _ffi_api.SaveParamsToFile(_to_tensor(params), path)
+ return get_global_func("vm.SaveParamsToFile")(_to_tensor(params), path)
def load_param_dict(param_bytes):
@@ -92,7 +93,7 @@ def load_param_dict(param_bytes):
"""
if isinstance(param_bytes, bytes | str):
param_bytes = bytearray(param_bytes)
- return _ffi_api.LoadParams(param_bytes)
+ return get_global_func("vm.LoadParams")(param_bytes)
def load_param_dict_from_file(path):
@@ -108,4 +109,4 @@ def load_param_dict_from_file(path):
params : dict of str to Tensor
The parameter dictionary.
"""
- return _ffi_api.LoadParamsFromFile(path)
+ return get_global_func("vm.LoadParamsFromFile")(path)
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm/vm.py
similarity index 100%
rename from python/tvm/runtime/vm.py
rename to python/tvm/runtime/vm/vm.py
diff --git a/python/tvm/s_tir/meta_schedule/builder/local_builder.py
b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
index 36b8825131..48edcba1c3 100644
--- a/python/tvm/s_tir/meta_schedule/builder/local_builder.py
+++ b/python/tvm/s_tir/meta_schedule/builder/local_builder.py
@@ -26,7 +26,8 @@ from tvm_ffi import register_global_func
from tvm.ir import IRModule
from tvm.ir.utils import derived_object
-from tvm.runtime import Module, Tensor, load_param_dict, save_param_dict
+from tvm.runtime import Module, Tensor
+from tvm.runtime.vm import load_param_dict, save_param_dict
from tvm.support.popen_pool import MapResult, PopenPoolExecutor, StatusKind
from tvm.target import Target
diff --git a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
index 129bedb211..ba856aa2b9 100644
--- a/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
+++ b/src/backend/hexagon/codegen/llvm/codegen_hexagon.cc
@@ -45,6 +45,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/function.h>
#include <tvm/runtime/logging.h>
+#include <tvm/support/file_utils.h>
#include <tvm/target/codegen.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/function.h>
@@ -58,7 +59,6 @@
#include <utility>
#include <vector>
-#include "../../../../runtime/file_utils.h"
#include "../../../../runtime/metadata.h"
#include "../../../../target/build_common.h"
#include "../../../../target/llvm/codegen_cpu.h"
@@ -567,8 +567,7 @@ ffi::Module BuildHexagon(IRModule mod, Target target) {
// The new design takes Bytes (not a filename); slurp the linked .so into
// memory. Receiver materializes the bytes back to a tempfile internally
// before dlopen.
- std::string so_data;
- runtime::LoadBinaryFromFile(so_name, &so_data);
+ std::string so_data = support::ReadBytesFromFile(so_name);
// Pack auxiliaries into the unified source map. Variant lets text blobs
// (asm/ll) and binary blobs (obj/bc) coexist in their natural form.
diff --git a/src/backend/metal/runtime/metal_module.mm
b/src/backend/metal/runtime/metal_module.mm
index c547b16a5f..713a34944a 100644
--- a/src/backend/metal/runtime/metal_module.mm
+++ b/src/backend/metal/runtime/metal_module.mm
@@ -37,7 +37,6 @@
#include <string>
#include <unordered_map>
#include <utility>
-#include "../../../runtime/file_utils.h"
#include "../../../runtime/launch_param.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
diff --git a/src/backend/opencl/runtime/opencl_common.h
b/src/backend/opencl/runtime/opencl_common.h
index b2b5ba2c54..0bf8596485 100644
--- a/src/backend/opencl/runtime/opencl_common.h
+++ b/src/backend/opencl/runtime/opencl_common.h
@@ -69,7 +69,6 @@
#include <unordered_map>
#include <vector>
-#include "../../../runtime/file_utils.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
#include "../../../runtime/workspace_pool.h"
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 152f10bb62..37618e86cc 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -44,7 +44,6 @@
#include <utility>
#include <vector>
-#include "../../../runtime/file_utils.h"
#include "../../../runtime/metadata.h"
#include "../../../runtime/thread_storage_scope.h"
#include "../../../support/bytes_io.h"
diff --git a/src/runtime/extra/contrib/clml/clml_runtime.cc
b/src/runtime/extra/contrib/clml/clml_runtime.cc
index e2dab689de..9c7d24635a 100644
--- a/src/runtime/extra/contrib/clml/clml_runtime.cc
+++ b/src/runtime/extra/contrib/clml/clml_runtime.cc
@@ -240,7 +240,7 @@ class CLMLRuntime : public JSONRuntimeBase {
if (!cws->is_tuning_run && cws->tuning_file) {
std::vector<unsigned char> tune_buffer;
std::string tune_blob;
- LoadBinaryFromFile(cws->tuning_file, &tune_blob);
+ tune_blob = support::ReadBytesFromFile(cws->tuning_file);
support::BytesInStream strm(tune_blob);
uint64_t header, reserve;
@@ -319,9 +319,9 @@ class CLMLRuntime : public JSONRuntimeBase {
dump_tensors.Set(node_name, narr);
}
- const auto f = tvm::ffi::Function::GetGlobal("runtime.SaveParams");
+ const auto f = tvm::ffi::Function::GetGlobal("vm.SaveParams");
if (f.has_value()) {
- std::string dump_bytes = (*f)(dump_tensors).cast<ffi::String>();
+ std::string dump_bytes = (*f)(dump_tensors).cast<ffi::Bytes>();
std::ostringstream oss;
/*TODO(Siva) HEX encoding doubles the size, look for better encode that
can cross the RPC. */
for (size_t i = 0; i < dump_bytes.size(); ++i) {
diff --git a/src/runtime/extra/contrib/clml/clml_runtime.h
b/src/runtime/extra/contrib/clml/clml_runtime.h
index 27286d5b22..db018a9583 100644
--- a/src/runtime/extra/contrib/clml/clml_runtime.h
+++ b/src/runtime/extra/contrib/clml/clml_runtime.h
@@ -34,6 +34,7 @@
#include <stdlib.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/tensor.h>
+#include <tvm/support/file_utils.h>
#include <fstream>
#include <map>
@@ -43,7 +44,6 @@
#include <vector>
#include "../../../../backend/opencl/runtime/opencl_common.h"
-#include "../../../file_utils.h"
#include "../../../thread_storage_scope.h"
#include "../json/json_node.h"
#include "../json/json_runtime.h"
diff --git a/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
b/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
index ab5548f72e..bb0a840a71 100644
--- a/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
+++ b/src/runtime/extra/contrib/tensorrt/tensorrt_runtime.cc
@@ -27,6 +27,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/runtime/tensor.h>
+#include <tvm/support/file_utils.h>
#include <cstdint>
#include <fstream>
@@ -38,7 +39,6 @@
#include <vector>
#include "../../../../support/env.h"
-#include "../../../file_utils.h"
#include "../json/json_node.h"
#include "../json/json_runtime.h"
@@ -606,8 +606,7 @@ class TensorRTRuntime : public JSONRuntimeBase {
if (!infile.good()) return false;
LOG(INFO) << "Loading cached TensorRT engine from " << path;
infile.close();
- std::string serialized_engine;
- LoadBinaryFromFile(path, &serialized_engine);
+ std::string serialized_engine = support::ReadBytesFromFile(path);
// Deserialize engine. TensorRT 10 dropped the trailing IPluginFactory*
argument and the runtime
// must outlive the engine, so it is owned by the cached
TensorRTEngineAndContext.
nvinfer1::IRuntime* runtime =
nvinfer1::createInferRuntime(GetTensorRTLogger());
@@ -626,8 +625,7 @@ class TensorRTRuntime : public JSONRuntimeBase {
// Load metadata
namespace json = ::tvm::ffi::json;
std::string meta_path = cache_dir + "/" + key + ".meta";
- std::string serialized_meta;
- LoadBinaryFromFile(meta_path, &serialized_meta);
+ std::string serialized_meta = support::ReadBytesFromFile(meta_path);
auto meta_obj = json::Parse(serialized_meta).cast<json::Object>();
int batch_size;
// Read inputs
@@ -667,8 +665,8 @@ class TensorRTRuntime : public JSONRuntimeBase {
// Serialize engine to disk
nvinfer1::IHostMemory* serialized_engine =
trt_engine_cache_[std::make_pair(symbol_name_,
batch_size)].engine->serialize();
- SaveBinaryToFile(path, std::string(static_cast<const
char*>(serialized_engine->data()),
- serialized_engine->size()));
+ support::WriteBytesToFile(path, std::string(static_cast<const
char*>(serialized_engine->data()),
+ serialized_engine->size()));
delete serialized_engine;
// Serialize metadata
namespace json = ::tvm::ffi::json;
@@ -689,7 +687,7 @@ class TensorRTRuntime : public JSONRuntimeBase {
}
meta_obj.Set(ffi::String("batch_size"), static_cast<int64_t>(batch_size));
std::string meta_path = cache_dir + "/" + key + ".meta";
- SaveBinaryToFile(meta_path, std::string(json::Stringify(meta_obj)));
+ support::WriteBytesToFile(meta_path,
std::string(json::Stringify(meta_obj)));
}
std::string GetSubgraphKey() {
diff --git a/src/runtime/extra/disco/loader.cc
b/src/runtime/extra/disco/loader.cc
index 0774b896fb..4e8e30b731 100644
--- a/src/runtime/extra/disco/loader.cc
+++ b/src/runtime/extra/disco/loader.cc
@@ -23,6 +23,7 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/vm/tensor_cache_support.h>
+#include <tvm/support/file_utils.h>
#include <functional>
#include <numeric>
@@ -30,7 +31,6 @@
#include <unordered_map>
#include <vector>
-#include "../../file_utils.h"
#include "./utils.h"
namespace tvm {
@@ -240,7 +240,7 @@ ffi::Tensor ShardLoaderObj::LoadParamOnWorker0(int
weight_index) const {
if (file != current_file_) {
current_file_ = file;
std::string file_name = GetSiblingPath(this->metadata_.path,
file->data_path);
- LoadBinaryFromFile(file_name, &this->current_file_stream_);
+ this->current_file_stream_ = support::ReadBytesFromFile(file_name);
}
return param->Load(device, &this->current_file_stream_);
};
@@ -287,7 +287,7 @@ ffi::Tensor ShardLoaderObj::LoadDirect(int weight_index)
const {
if (file != current_file_) {
current_file_ = file;
std::string file_name = GetSiblingPath(this->metadata_.path,
file->data_path);
- LoadBinaryFromFile(file_name, &this->current_file_stream_);
+ this->current_file_stream_ = support::ReadBytesFromFile(file_name);
}
return param->Load(device, &this->current_file_stream_);
}
diff --git a/src/runtime/file_utils.cc b/src/runtime/file_utils.cc
deleted file mode 100644
index ff2c15e8de..0000000000
--- a/src/runtime/file_utils.cc
+++ /dev/null
@@ -1,259 +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 file_utils.cc
- */
-#include "file_utils.h"
-
-#include <tvm/ffi/error.h>
-#include <tvm/ffi/extra/json.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/support/io.h>
-
-#include <fstream>
-#include <utility>
-#include <vector>
-
-#include "../support/bytes_io.h"
-
-namespace tvm {
-namespace runtime {
-
-std::string GetFileFormat(const std::string& file_name, const std::string&
format) {
- std::string fmt = format;
- if (fmt.length() == 0) {
- size_t pos = file_name.find_last_of(".");
- if (pos != std::string::npos) {
- return file_name.substr(pos + 1, file_name.length() - pos - 1);
- } else {
- return "";
- }
- } else {
- return format;
- }
-}
-
-std::string GetCacheDir() {
- char* env_cache_dir;
- if ((env_cache_dir = getenv("TVM_CACHE_DIR"))) return env_cache_dir;
- if ((env_cache_dir = getenv("XDG_CACHE_HOME"))) {
- return std::string(env_cache_dir) + "/tvm";
- }
- if ((env_cache_dir = getenv("HOME"))) {
- return std::string(env_cache_dir) + "/.cache/tvm";
- }
- return ".";
-}
-
-std::string GetFileBasename(const std::string& file_name) {
- size_t last_slash = file_name.find_last_of("/");
- if (last_slash == std::string::npos) return file_name;
- return file_name.substr(last_slash + 1);
-}
-
-std::string GetMetaFilePath(const std::string& file_name) {
- size_t pos = file_name.find_last_of(".");
- if (pos != std::string::npos) {
- return file_name.substr(0, pos) + ".tvm_meta.json";
- } else {
- return file_name + ".tvm_meta.json";
- }
-}
-
-void LoadBinaryFromFile(const std::string& file_name, std::string* data) {
- std::ifstream fs(file_name, std::ios::in | std::ios::binary);
- TVM_FFI_ICHECK(!fs.fail()) << "Cannot open " << file_name;
- // get its size:
- fs.seekg(0, std::ios::end);
- size_t size = static_cast<size_t>(fs.tellg());
- fs.seekg(0, std::ios::beg);
- data->resize(size);
- fs.read(&(*data)[0], size);
-}
-
-void SaveBinaryToFile(const std::string& file_name, const std::string& data) {
- std::ofstream fs(file_name, std::ios::out | std::ios::binary);
- TVM_FFI_ICHECK(!fs.fail()) << "Cannot open " << file_name;
- fs.write(&data[0], data.length());
-}
-
-void SaveMetaDataToFile(const std::string& file_name,
- const ffi::Map<ffi::String, FunctionInfo>& fmap) {
- namespace json = ::tvm::ffi::json;
- json::Object root;
- root.Set("tvm_version", ffi::String("0.1.0"));
- json::Object func_info;
- for (const auto& kv : fmap) {
- func_info.Set(kv.first, kv.second->SaveToJSON());
- }
- root.Set("func_info", std::move(func_info));
- std::ofstream fs(file_name.c_str());
- TVM_FFI_ICHECK(!fs.fail()) << "Cannot open file " << file_name;
- fs << std::string(json::Stringify(root, 2));
- fs.close();
-}
-
-void LoadMetaDataFromFile(const std::string& file_name, ffi::Map<ffi::String,
FunctionInfo>* fmap) {
- namespace json = ::tvm::ffi::json;
- std::ifstream fs(file_name.c_str());
- TVM_FFI_ICHECK(!fs.fail()) << "Cannot open file " << file_name;
- std::string content((std::istreambuf_iterator<char>(fs)),
std::istreambuf_iterator<char>());
- fs.close();
- auto root = json::Parse(content).cast<json::Object>();
- // tvm_version is ignored
- auto func_info_obj = root.at("func_info").cast<json::Object>();
- for (const auto& kv : func_info_obj) {
- auto info_node = ffi::make_object<FunctionInfoObj>();
- info_node->LoadFromJSON(kv.second.cast<json::Object>());
- fmap->Set(kv.first.cast<ffi::String>(), FunctionInfo(ffi::UnsafeInit{},
std::move(info_node)));
- }
-}
-
-void RemoveFile(const std::string& file_name) {
- // FIXME: This doesn't check the return code.
- std::remove(file_name.c_str());
-}
-
-void CopyFile(const std::string& src_file_name, const std::string&
dest_file_name) {
- std::ifstream src(src_file_name, std::ios::binary);
- TVM_FFI_ICHECK(src) << "Unable to open source file '" << src_file_name <<
"'";
-
- std::ofstream dest(dest_file_name, std::ios::binary | std::ios::trunc);
- TVM_FFI_ICHECK(dest) << "Unable to destination source file '" <<
src_file_name << "'";
-
- dest << src.rdbuf();
-
- src.close();
- dest.close();
-
- TVM_FFI_ICHECK(dest) << "File-copy operation failed."
- << " src='" << src_file_name << "'"
- << " dest='" << dest_file_name << "'";
-}
-
-ffi::Map<ffi::String, ffi::Tensor> LoadParams(const std::string& param_blob) {
- support::BytesInStream strm(param_blob);
- return LoadParams(&strm);
-}
-ffi::Map<ffi::String, ffi::Tensor> LoadParams(support::Stream* strm) {
- ffi::Map<ffi::String, ffi::Tensor> params;
- uint64_t header, reserved;
- TVM_FFI_ICHECK(strm->Read(&header)) << "Invalid parameters file format";
- TVM_FFI_ICHECK(header == kTVMTensorListMagic) << "Invalid parameters file
format";
- TVM_FFI_ICHECK(strm->Read(&reserved)) << "Invalid parameters file format";
-
- std::vector<std::string> names;
- TVM_FFI_ICHECK(strm->Read(&names)) << "Invalid parameters file format";
- uint64_t sz;
- strm->Read(&sz);
- size_t size = static_cast<size_t>(sz);
- TVM_FFI_ICHECK(size == names.size()) << "Invalid parameters file format";
- for (size_t i = 0; i < size; ++i) {
- // The data_entry is allocated on device, Tensor.load always load the
array into CPU.
- ffi::Tensor temp = LoadTensor(strm);
- params.Set(names[i], temp);
- }
- return params;
-}
-
-void SaveParams(support::Stream* strm, const ffi::Map<ffi::String,
ffi::Tensor>& params) {
- for (const auto& entry : params) {
- TVM_FFI_ICHECK(entry.second.defined());
- }
- std::vector<std::string> names;
- std::vector<const DLTensor*> arrays;
- for (auto& p : params) {
- names.push_back(p.first);
- arrays.push_back(p.second.GetDLTensorPtr());
- }
-
- uint64_t header = kTVMTensorListMagic, reserved = 0;
- strm->Write(header);
- strm->Write(reserved);
- strm->Write(names);
- {
- uint64_t sz = static_cast<uint64_t>(arrays.size());
- strm->Write(sz);
- for (size_t i = 0; i < sz; ++i) {
- tvm::runtime::SaveTensor(arrays[i], strm);
- }
- }
-}
-
-std::string SaveParams(const ffi::Map<ffi::String, ffi::Tensor>& params) {
- std::string result;
- support::BytesOutStream strm(&result);
- SaveParams(&strm, params);
- return result;
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<FunctionInfoObj>()
- .def_ro("name", &FunctionInfoObj::name)
- .def_ro("arg_types", &FunctionInfoObj::arg_types)
- .def_ro("launch_param_tags", &FunctionInfoObj::launch_param_tags)
- .def_ro("arg_extra_tags", &FunctionInfoObj::arg_extra_tags)
- .def_ro("cuda_launch_fields", &FunctionInfoObj::cuda_launch_fields,
- refl::DefaultValue(ffi::Array<ffi::String>()))
- .def_ro("cuda_kernel_attrs", &FunctionInfoObj::cuda_kernel_attrs,
- refl::DefaultValue(ffi::Map<ffi::String, int64_t>()));
- refl::GlobalDef()
- .def("runtime.SaveParams",
- [](const ffi::Map<ffi::String, ffi::Tensor>& params) {
- std::string s = ::tvm::runtime::SaveParams(params);
- return ffi::Bytes(std::move(s));
- })
- .def("runtime.SaveParamsToFile",
- [](const ffi::Map<ffi::String, ffi::Tensor>& params, const
ffi::String& path) {
- for (const auto& entry : params) {
- TVM_FFI_ICHECK(entry.second.defined());
- }
- tvm::runtime::SimpleBinaryFileStream strm(path, "wb");
- SaveParams(&strm, params);
- })
- .def("runtime.LoadParams", [](const ffi::Bytes& s) { return
::tvm::runtime::LoadParams(s); })
- .def("runtime.LoadParamsFromFile",
- [](const ffi::String& path) {
- tvm::runtime::SimpleBinaryFileStream strm(path, "rb");
- return LoadParams(&strm);
- })
- // Registry: "runtime.LoadMetaDataFromJSON" — parse a tvm_meta.json
- // string into Map<String, FunctionInfo>. Used by Python callers that
- // build FunctionInfo maps in-memory (e.g. tirx external_kernel) without
- // a disk round-trip.
- .def("runtime.LoadMetaDataFromJSON", [](const ffi::String& json_str) {
- namespace json = ::tvm::ffi::json;
- ffi::Map<ffi::String, FunctionInfo> fmap;
- auto root = json::Parse(std::string(json_str)).cast<json::Object>();
- auto func_info_obj = root.at("func_info").cast<json::Object>();
- for (const auto& kv : func_info_obj) {
- auto info_node = ffi::make_object<FunctionInfoObj>();
- info_node->LoadFromJSON(kv.second.cast<json::Object>());
- fmap.Set(kv.first.cast<ffi::String>(),
- FunctionInfo(ffi::UnsafeInit{}, std::move(info_node)));
- }
- return fmap;
- });
-}
-
-} // namespace runtime
-} // namespace tvm
diff --git a/src/runtime/file_utils.h b/src/runtime/file_utils.h
deleted file mode 100644
index 703b4947cc..0000000000
--- a/src/runtime/file_utils.h
+++ /dev/null
@@ -1,177 +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 file_utils.h
- * \brief Minimum file manipulation utils for runtime.
- */
-#ifndef TVM_RUNTIME_FILE_UTILS_H_
-#define TVM_RUNTIME_FILE_UTILS_H_
-
-#include <tvm/ffi/container/map.h>
-#include <tvm/ffi/string.h>
-
-#include <string>
-
-#include "metadata.h"
-
-namespace tvm {
-namespace runtime {
-/*!
- * \brief Get file format from given file name or format argument.
- * \param file_name The name of the file.
- * \param format The format of the file.
- */
-TVM_RUNTIME_DLL std::string GetFileFormat(const std::string& file_name, const
std::string& format);
-
-/*!
- * \return the directory in which TVM stores cached files.
- * May be set using TVM_CACHE_DIR; defaults to system locations.
- */
-std::string GetCacheDir();
-
-/*!
- * \brief Get meta file path given file name and format.
- * \param file_name The name of the file.
- */
-TVM_RUNTIME_DLL std::string GetMetaFilePath(const std::string& file_name);
-
-/*!
- * \brief Get file basename (i.e. without leading directories)
- * \param file_name The name of the file.
- * \return the base name
- */
-std::string GetFileBasename(const std::string& file_name);
-
-/*!
- * \brief Load binary file into a in-memory buffer.
- * \param file_name The name of the file.
- * \param data The data to be loaded.
- */
-TVM_RUNTIME_DLL void LoadBinaryFromFile(const std::string& file_name,
std::string* data);
-
-/*!
- * \brief Load binary file into a in-memory buffer.
- * \param file_name The name of the file.
- * \param data The binary data to be saved.
- */
-TVM_RUNTIME_DLL void SaveBinaryToFile(const std::string& file_name, const
std::string& data);
-
-/*!
- * \brief Save meta data to file.
- * \param file_name The name of the file.
- * \param fmap The function info map.
- */
-TVM_RUNTIME_DLL void SaveMetaDataToFile(const std::string& file_name,
- const ffi::Map<ffi::String,
FunctionInfo>& fmap);
-
-/*!
- * \brief Load meta data to file.
- * \param file_name The name of the file.
- * \param fmap The function info map.
- */
-TVM_RUNTIME_DLL void LoadMetaDataFromFile(const std::string& file_name,
- ffi::Map<ffi::String, FunctionInfo>*
fmap);
-
-/*!
- * \brief Copy the content of an existing file to another file.
- * \param src_file_name Path to the source file.
- * \param dest_file_name Path of the destination file. If this file already
exists,
- * replace its content.
- */
-void CopyFile(const std::string& src_file_name, const std::string&
dest_file_name);
-
-/*!
- * \brief Remove (unlink) a file.
- * \param file_name The file name.
- */
-void RemoveFile(const std::string& file_name);
-
-constexpr uint64_t kTVMTensorListMagic = 0xF7E58D4F05049CB7;
-/*!
- * \brief Load parameters from a string.
- * \param param_blob Serialized string of parameters.
- * \return Map of parameter name to parameter value.
- */
-ffi::Map<ffi::String, ffi::Tensor> LoadParams(const std::string& param_blob);
-/*!
- * \brief Load parameters from a stream.
- * \param strm Stream to load parameters from.
- * \return Map of parameter name to parameter value.
- */
-ffi::Map<ffi::String, ffi::Tensor> LoadParams(support::Stream* strm);
-/*!
- * \brief Serialize parameters to a byte array.
- * \param params Parameters to save.
- * \return ffi::String containing binary parameter data.
- */
-std::string SaveParams(const ffi::Map<ffi::String, ffi::Tensor>& params);
-/*!
- * \brief Serialize parameters to a stream.
- * \param strm Stream to write to.
- * \param params Parameters to save.
- */
-void SaveParams(support::Stream* strm, const ffi::Map<ffi::String,
ffi::Tensor>& params);
-
-/*!
- * \brief A dmlc stream which wraps standard file operations.
- */
-struct SimpleBinaryFileStream : public support::Stream {
- public:
- using Stream::Read;
- using Stream::Write;
-
- SimpleBinaryFileStream(const std::string& path, std::string mode) {
- const char* fname = path.c_str();
-
- TVM_FFI_ICHECK(mode == "wb" || mode == "rb") << "Only allowed modes are
'wb' and 'rb'";
- read_ = mode == "rb";
- fp_ = std::fopen(fname, mode.c_str());
- TVM_FFI_ICHECK(fp_ != nullptr) << "Unable to open file " << path;
- }
- virtual ~SimpleBinaryFileStream(void) { this->Close(); }
- virtual size_t Read(void* ptr, size_t size) {
- TVM_FFI_ICHECK(read_) << "File opened in write-mode, cannot read.";
- TVM_FFI_ICHECK(fp_ != nullptr) << "File is closed";
- return std::fread(ptr, 1, size, fp_);
- }
- virtual size_t Write(const void* ptr, size_t size) {
- TVM_FFI_ICHECK(!read_) << "File opened in read-mode, cannot write.";
- TVM_FFI_ICHECK(fp_ != nullptr) << "File is closed";
- size_t nwrite = std::fwrite(ptr, 1, size, fp_);
- int err = std::ferror(fp_);
-
- TVM_FFI_ICHECK_EQ(err, 0) << "SimpleBinaryFileStream.Write incomplete: "
<< std::strerror(err);
- return nwrite;
- }
- inline void Close(void) {
- if (fp_ != nullptr) {
- std::fclose(fp_);
- fp_ = nullptr;
- }
- }
-
- private:
- std::FILE* fp_ = nullptr;
- bool read_;
-}; // class SimpleBinaryFileStream
-
-} // namespace runtime
-} // namespace tvm
-#endif // TVM_RUNTIME_FILE_UTILS_H_
diff --git a/src/runtime/metadata.cc b/src/runtime/metadata.cc
new file mode 100644
index 0000000000..e509b661aa
--- /dev/null
+++ b/src/runtime/metadata.cc
@@ -0,0 +1,66 @@
+/*
+ * 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 metadata.cc
+ * \brief Runtime function metadata reflection and JSON conversion.
+ */
+#include "metadata.h"
+
+#include <tvm/ffi/reflection/registry.h>
+
+#include <string>
+#include <utility>
+
+namespace tvm {
+namespace runtime {
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<FunctionInfoObj>()
+ .def_ro("name", &FunctionInfoObj::name)
+ .def_ro("arg_types", &FunctionInfoObj::arg_types)
+ .def_ro("launch_param_tags", &FunctionInfoObj::launch_param_tags)
+ .def_ro("arg_extra_tags", &FunctionInfoObj::arg_extra_tags)
+ .def_ro("cuda_launch_fields", &FunctionInfoObj::cuda_launch_fields,
+ refl::DefaultValue(ffi::Array<ffi::String>()))
+ .def_ro("cuda_kernel_attrs", &FunctionInfoObj::cuda_kernel_attrs,
+ refl::DefaultValue(ffi::Map<ffi::String, int64_t>()));
+ refl::GlobalDef()
+ // Registry: "runtime.LoadMetaDataFromJSON" — parse a tvm_meta.json
+ // string into Map<String, FunctionInfo>. Used by Python callers that
+ // build FunctionInfo maps in-memory (e.g. tirx external_kernel) without
+ // a disk round-trip.
+ .def("runtime.LoadMetaDataFromJSON", [](const ffi::String& json_str) {
+ namespace json = ::tvm::ffi::json;
+ ffi::Map<ffi::String, FunctionInfo> fmap;
+ auto root = json::Parse(std::string(json_str)).cast<json::Object>();
+ auto func_info_obj = root.at("func_info").cast<json::Object>();
+ for (const auto& kv : func_info_obj) {
+ auto info_node = ffi::make_object<FunctionInfoObj>();
+ info_node->LoadFromJSON(kv.second.cast<json::Object>());
+ fmap.Set(kv.first.cast<ffi::String>(),
+ FunctionInfo(ffi::UnsafeInit{}, std::move(info_node)));
+ }
+ return fmap;
+ });
+}
+
+} // namespace runtime
+} // namespace tvm
diff --git a/src/runtime/module.cc b/src/runtime/module.cc
index 1bced00a39..7e12722e61 100644
--- a/src/runtime/module.cc
+++ b/src/runtime/module.cc
@@ -30,8 +30,6 @@
#include <cstring>
#include <unordered_set>
-#include "file_utils.h"
-
namespace tvm {
namespace runtime {
diff --git a/src/runtime/rpc/rpc_server_env.cc
b/src/runtime/rpc/rpc_server_env.cc
index b8c08b19c4..cab82b8b70 100644
--- a/src/runtime/rpc/rpc_server_env.cc
+++ b/src/runtime/rpc/rpc_server_env.cc
@@ -24,12 +24,20 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
+#include <tvm/support/file_utils.h>
-#include "../file_utils.h"
+#include <cstdio>
namespace tvm {
namespace runtime {
+namespace {
+void RemoveFile(const std::string& file_name) {
+ // Preserve the RPC removal operation's unchecked return code.
+ std::remove(file_name.c_str());
+}
+} // namespace
+
std::string RPCGetPath(const std::string& name) {
// do live lookup everytime as workpath can change.
const auto f = tvm::ffi::Function::GetGlobal("tvm.rpc.server.workpath");
@@ -44,13 +52,12 @@ TVM_FFI_STATIC_INIT_BLOCK() {
[](ffi::PackedArgs args, ffi::Any* rv) {
std::string file_name =
RPCGetPath(args[0].cast<std::string>());
auto data = args[1].cast<std::string>();
- SaveBinaryToFile(file_name, data);
+ support::WriteBytesToFile(file_name, data);
})
.def_packed("tvm.rpc.server.download",
[](ffi::PackedArgs args, ffi::Any* rv) {
std::string file_name =
RPCGetPath(args[0].cast<std::string>());
- std::string data;
- LoadBinaryFromFile(file_name, &data);
+ std::string data = support::ReadBytesFromFile(file_name);
LOG(INFO) << "Download " << file_name << "... nbytes=" <<
data.size();
*rv = ffi::Bytes(data);
})
diff --git a/src/runtime/static_library.cc b/src/runtime/static_library.cc
index 2ab201cdf1..97f77f09a4 100644
--- a/src/runtime/static_library.cc
+++ b/src/runtime/static_library.cc
@@ -30,13 +30,13 @@
#include <tvm/ffi/memory.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
+#include <tvm/support/file_utils.h>
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
#include <iostream>
#include "../support/bytes_io.h"
-#include "file_utils.h"
namespace tvm {
namespace runtime {
@@ -90,7 +90,7 @@ class StaticLibraryObj final : public ffi::ModuleObj {
void WriteToFile(const ffi::String& file_name, const ffi::String& format)
const final {
VLOG(0) << "Saving static library of " << data_.size() << " bytes
implementing " << FuncNames()
<< " to '" << file_name << "'";
- SaveBinaryToFile(file_name, data_);
+ support::WriteBytesToFile(file_name, data_);
}
/*! \brief Get the property of the runtime module .*/
@@ -128,7 +128,7 @@ class StaticLibraryObj final : public ffi::ModuleObj {
ffi::Module LoadStaticLibrary(const std::string& filename,
ffi::Array<ffi::String> func_names) {
auto node = ffi::make_object<StaticLibraryObj>();
- LoadBinaryFromFile(filename, &node->data_);
+ node->data_ = support::ReadBytesFromFile(filename);
node->func_names_ = std::move(func_names);
VLOG(0) << "Loaded static library from '" << filename << "' implementing "
<< node->FuncNames();
return ffi::Module(node);
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index 612935acaf..4cde076829 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -26,13 +26,13 @@
#include <tvm/runtime/logging.h>
#include <tvm/runtime/vm/executable.h>
#include <tvm/runtime/vm/vm.h>
+#include <tvm/support/file_utils.h>
#include <tvm/support/io.h>
#include <functional>
#include <sstream>
#include "../../support/bytes_io.h"
-#include "../file_utils.h"
#include "./module_utils.h"
namespace tvm {
@@ -203,7 +203,7 @@ ffi::Bytes VMExecutable::SaveToBytes() const {
}
void VMExecutable::WriteToFile(const ffi::String& file_name, const
ffi::String& format) const {
- runtime::SaveBinaryToFile(file_name, VMExecutable::SaveToBytes());
+ support::WriteBytesToFile(file_name, VMExecutable::SaveToBytes());
}
ffi::Module VMExecutable::LoadFromBytes(const ffi::Bytes& bytes) {
@@ -232,8 +232,7 @@ ffi::Module VMExecutable::LoadFromBytes(const ffi::Bytes&
bytes) {
}
ffi::Module VMExecutable::LoadFromFile(const ffi::String& file_name) {
- std::string data;
- runtime::LoadBinaryFromFile(file_name, &data);
+ std::string data = support::ReadBytesFromFile(file_name);
return VMExecutable::LoadFromBytes(ffi::Bytes(data));
}
diff --git a/src/runtime/vm/params.cc b/src/runtime/vm/params.cc
new file mode 100644
index 0000000000..80f22dedb3
--- /dev/null
+++ b/src/runtime/vm/params.cc
@@ -0,0 +1,169 @@
+/*
+ * 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 params.cc
+ * \brief VM parameter dictionary serialization.
+ */
+#include <tvm/ffi/container/map.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/string.h>
+#include <tvm/runtime/tensor.h>
+#include <tvm/support/io.h>
+
+#include <cstdio>
+#include <cstring>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "../../support/bytes_io.h"
+
+namespace tvm {
+namespace runtime {
+namespace vm {
+
+namespace {
+constexpr uint64_t kTVMTensorListMagic = 0xF7E58D4F05049CB7;
+
+struct SimpleBinaryFileStream : public support::Stream {
+ public:
+ using Stream::Read;
+ using Stream::Write;
+
+ SimpleBinaryFileStream(const std::string& path, std::string mode) {
+ const char* fname = path.c_str();
+
+ TVM_FFI_ICHECK(mode == "wb" || mode == "rb") << "Only allowed modes are
'wb' and 'rb'";
+ read_ = mode == "rb";
+ fp_ = std::fopen(fname, mode.c_str());
+ TVM_FFI_ICHECK(fp_ != nullptr) << "Unable to open file " << path;
+ }
+ virtual ~SimpleBinaryFileStream(void) { this->Close(); }
+ virtual size_t Read(void* ptr, size_t size) {
+ TVM_FFI_ICHECK(read_) << "File opened in write-mode, cannot read.";
+ TVM_FFI_ICHECK(fp_ != nullptr) << "File is closed";
+ return std::fread(ptr, 1, size, fp_);
+ }
+ virtual size_t Write(const void* ptr, size_t size) {
+ TVM_FFI_ICHECK(!read_) << "File opened in read-mode, cannot write.";
+ TVM_FFI_ICHECK(fp_ != nullptr) << "File is closed";
+ size_t nwrite = std::fwrite(ptr, 1, size, fp_);
+ int err = std::ferror(fp_);
+
+ TVM_FFI_ICHECK_EQ(err, 0) << "SimpleBinaryFileStream.Write incomplete: "
<< std::strerror(err);
+ return nwrite;
+ }
+ inline void Close(void) {
+ if (fp_ != nullptr) {
+ std::fclose(fp_);
+ fp_ = nullptr;
+ }
+ }
+
+ private:
+ std::FILE* fp_ = nullptr;
+ bool read_;
+}; // class SimpleBinaryFileStream
+
+} // namespace
+
+ffi::Map<ffi::String, ffi::Tensor> LoadParams(support::Stream* strm) {
+ ffi::Map<ffi::String, ffi::Tensor> params;
+ uint64_t header, reserved;
+ TVM_FFI_ICHECK(strm->Read(&header)) << "Invalid parameters file format";
+ TVM_FFI_ICHECK(header == kTVMTensorListMagic) << "Invalid parameters file
format";
+ TVM_FFI_ICHECK(strm->Read(&reserved)) << "Invalid parameters file format";
+
+ std::vector<std::string> names;
+ TVM_FFI_ICHECK(strm->Read(&names)) << "Invalid parameters file format";
+ uint64_t sz;
+ strm->Read(&sz);
+ size_t size = static_cast<size_t>(sz);
+ TVM_FFI_ICHECK(size == names.size()) << "Invalid parameters file format";
+ for (size_t i = 0; i < size; ++i) {
+ // The data_entry is allocated on device, Tensor.load always load the
array into CPU.
+ ffi::Tensor temp = LoadTensor(strm);
+ params.Set(names[i], temp);
+ }
+ return params;
+}
+
+ffi::Map<ffi::String, ffi::Tensor> LoadParams(const std::string& param_blob) {
+ support::BytesInStream strm(param_blob);
+ return LoadParams(&strm);
+}
+
+void SaveParams(support::Stream* strm, const ffi::Map<ffi::String,
ffi::Tensor>& params) {
+ for (const auto& entry : params) {
+ TVM_FFI_ICHECK(entry.second.defined());
+ }
+ std::vector<std::string> names;
+ std::vector<const DLTensor*> arrays;
+ for (auto& p : params) {
+ names.push_back(p.first);
+ arrays.push_back(p.second.GetDLTensorPtr());
+ }
+
+ uint64_t header = kTVMTensorListMagic, reserved = 0;
+ strm->Write(header);
+ strm->Write(reserved);
+ strm->Write(names);
+ {
+ uint64_t sz = static_cast<uint64_t>(arrays.size());
+ strm->Write(sz);
+ for (size_t i = 0; i < sz; ++i) {
+ tvm::runtime::SaveTensor(arrays[i], strm);
+ }
+ }
+}
+
+std::string SaveParams(const ffi::Map<ffi::String, ffi::Tensor>& params) {
+ std::string result;
+ support::BytesOutStream strm(&result);
+ SaveParams(&strm, params);
+ return result;
+}
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef()
+ .def("vm.SaveParams",
+ [](const ffi::Map<ffi::String, ffi::Tensor>& params) {
+ std::string s = ::tvm::runtime::vm::SaveParams(params);
+ return ffi::Bytes(std::move(s));
+ })
+ .def("vm.SaveParamsToFile",
+ [](const ffi::Map<ffi::String, ffi::Tensor>& params, const
ffi::String& path) {
+ for (const auto& entry : params) {
+ TVM_FFI_ICHECK(entry.second.defined());
+ }
+ SimpleBinaryFileStream strm(path, "wb");
+ SaveParams(&strm, params);
+ })
+ .def("vm.LoadParams", [](const ffi::Bytes& s) { return
::tvm::runtime::vm::LoadParams(s); })
+ .def("vm.LoadParamsFromFile", [](const ffi::String& path) {
+ SimpleBinaryFileStream strm(path, "rb");
+ return LoadParams(&strm);
+ });
+}
+
+} // namespace vm
+} // namespace runtime
+} // namespace tvm
diff --git a/src/runtime/vm/tensor_cache_support.cc
b/src/runtime/vm/tensor_cache_support.cc
index a1b91e330e..9bd9940fb8 100644
--- a/src/runtime/vm/tensor_cache_support.cc
+++ b/src/runtime/vm/tensor_cache_support.cc
@@ -35,16 +35,17 @@
* runtime builtin provide as in this file.
*/
#include <tvm/ffi/extra/json.h>
+#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/tensor.h>
#include <tvm/runtime/vm/tensor_cache_support.h>
+#include <tvm/support/file_utils.h>
#include <string>
#include <vector>
#include "../../support/utils.h"
-#include "../file_utils.h"
namespace tvm {
namespace runtime {
@@ -111,8 +112,7 @@ TensorCacheMetadata TensorCacheMetadata::LoadFromStr(const
std::string& json_str
}
TVM_RUNTIME_DLL TensorCacheMetadata TensorCacheMetadata::Load(const
std::string& path) {
- std::string json_str;
- LoadBinaryFromFile(path + "/tensor-cache.json", &json_str);
+ std::string json_str = support::ReadBytesFromFile(path +
"/tensor-cache.json");
ffi::String err;
json::Value json_info = json::Parse(json_str, &err);
if (!err.empty()) {
@@ -176,7 +176,7 @@ TVM_RUNTIME_DLL ffi::Array<ffi::Tensor>
TensorCacheMetadata::FileRecord::Load(
const std::string& path_prefix, //
std::string* raw_data_buffer, //
ffi::Optional<ffi::Tensor>* staging_buffer) const {
- LoadBinaryFromFile(path_prefix + "/" + this->data_path, raw_data_buffer);
+ *raw_data_buffer = support::ReadBytesFromFile(path_prefix + "/" +
this->data_path);
TVM_FFI_CHECK_EQ(this->format, "raw-shard", ValueError) << "Only `raw-shard`
format is supported";
TVM_FFI_CHECK_EQ(this->nbytes, raw_data_buffer->length(), ValueError)
<< "Encountered an corrupted parameter shard. It means it is not
downloaded "
diff --git a/src/target/llvm/llvm_module.cc b/src/target/llvm/llvm_module.cc
index 15b47b6f7d..586d6c45cb 100644
--- a/src/target/llvm/llvm_module.cc
+++ b/src/target/llvm/llvm_module.cc
@@ -58,6 +58,7 @@
#include <tvm/ir/module.h>
#include <tvm/ir/with_context.h>
#include <tvm/runtime/logging.h>
+#include <tvm/support/file_utils.h>
#include <tvm/target/codegen.h>
#include <tvm/target/target.h>
@@ -66,7 +67,6 @@
#include <system_error>
#include <utility>
-#include "../../runtime/file_utils.h"
#include "codegen_blob.h"
#include "codegen_cpu.h"
#include "codegen_llvm.h"
@@ -180,7 +180,7 @@ bool LLVMAddPassesToEmitFile(llvm::TargetMachine* tm,
llvm::legacy::PassManager*
void LLVMModuleObj::WriteToFile(const ffi::String& file_name_str, const
ffi::String& format) const {
// TVM_FFI_ICHECK(imports_.empty()) << "SaveToFile does not handle imported
modules";
std::string file_name = file_name_str;
- std::string fmt = runtime::GetFileFormat(file_name, format);
+ std::string fmt = support::GetFileFormat(file_name, format);
std::error_code ecode;
llvm::raw_fd_ostream dest(file_name, ecode, llvm_open_output_flag);
TVM_FFI_ICHECK_EQ(ecode.value(), 0)
@@ -214,7 +214,7 @@ ffi::Bytes LLVMModuleObj::SaveToBytes() const {
}
ffi::String LLVMModuleObj::InspectSource(const ffi::String& format) const {
- std::string fmt = runtime::GetFileFormat("", format);
+ std::string fmt = support::GetFileFormat("", format);
std::string type_str;
llvm::SmallString<256> str;
llvm::raw_svector_ostream rso(str);
diff --git a/src/target/source/source_module.cc
b/src/target/source/source_module.cc
index b9b0500a30..944ff2badf 100644
--- a/src/target/source/source_module.cc
+++ b/src/target/source/source_module.cc
@@ -27,6 +27,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/tensor.h>
+#include <tvm/support/file_utils.h>
#include <tvm/support/io.h>
#include <algorithm>
@@ -35,7 +36,6 @@
#include <unordered_set>
#include <vector>
-#include "../../runtime/file_utils.h"
#include "../../support/bytes_io.h"
#include "codegen_source_base.h"
@@ -47,9 +47,7 @@ using ffi::Function;
using ffi::PackedArgs;
using runtime::FunctionInfo;
-using runtime::GetFileFormat;
-using runtime::GetMetaFilePath;
-using runtime::SaveBinaryToFile;
+using support::GetFileFormat;
// Simulator function
class SourceModuleObj : public ffi::ModuleObj {
@@ -150,10 +148,9 @@ class CSourceModuleObj : public ffi::ModuleObj {
void WriteToFile(const ffi::String& file_name, const ffi::String& format)
const final {
std::string fmt = GetFileFormat(file_name, format);
- std::string meta_file = GetMetaFilePath(file_name);
if (fmt == "c" || fmt == "cc" || fmt == "cpp" || fmt == "cu") {
TVM_FFI_ICHECK_NE(code_.length(), 0);
- SaveBinaryToFile(file_name, code_);
+ support::WriteBytesToFile(file_name, code_);
} else {
TVM_FFI_ICHECK_EQ(fmt, fmt_) << "Can only save to format=" << fmt_;
}
@@ -231,9 +228,7 @@ class DeviceSourceModuleObj final : public ffi::ModuleObj {
void WriteToFile(const ffi::String& file_name, const ffi::String& format)
const final {
std::string fmt = GetFileFormat(file_name, format);
TVM_FFI_ICHECK_EQ(fmt, fmt_) << "Can only save to format=" << fmt_;
- std::string meta_file = GetMetaFilePath(file_name);
- SaveMetaDataToFile(meta_file, fmap_);
- SaveBinaryToFile(file_name, data_);
+ support::WriteBytesToFile(file_name, data_);
}
ffi::Bytes SaveToBytes() const final {
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index 6d3c1b7bb0..a56a55afcd 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -38,8 +38,8 @@
#include "src/runtime/cpu_device_api.cc"
#include "src/runtime/device_api.cc"
#include "src/runtime/extra/contrib/sort/sort.cc"
-#include "src/runtime/file_utils.cc"
#include "src/runtime/logging.cc"
+#include "src/runtime/metadata.cc"
#include "src/runtime/rpc/rpc_channel.cc"
#include "src/runtime/rpc/rpc_endpoint.cc"
#include "src/runtime/rpc/rpc_event_impl.cc"
@@ -48,6 +48,7 @@
#include "src/runtime/rpc/rpc_session.cc"
#include "src/runtime/tensor.cc"
#include "src/runtime/timer.cc"
+#include "src/runtime/vm/params.cc"
#include "src/runtime/workspace_pool.cc"
// relax setup
#include "3rdparty/tvm-ffi/src/ffi/backtrace.cc"
diff --git a/web/emcc/webgpu_runtime.cc b/web/emcc/webgpu_runtime.cc
index 83b8ebc03e..bced4636e6 100644
--- a/web/emcc/webgpu_runtime.cc
+++ b/web/emcc/webgpu_runtime.cc
@@ -37,7 +37,6 @@
#include <mutex>
#include <string>
-#include "../../src/runtime/file_utils.h"
#include "../../src/runtime/metadata.h"
#include "../../src/support/bytes_io.h"
