chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,332 @@
|
||||
/*
|
||||
* 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/serializer.h
|
||||
* \brief Serializer<T> specializations for tvm::support::Stream.
|
||||
*
|
||||
* Built-in support for:
|
||||
* - Arithmetic types (endian-aware)
|
||||
* - Enum types (via underlying arithmetic type)
|
||||
* - std::string, std::vector<T>, std::pair<A,B>, std::unordered_map<K,V>
|
||||
* - DLDataType, DLDevice
|
||||
*
|
||||
* Custom types with Save(Stream*)/Load(Stream*) methods should define
|
||||
* Serializer<Type> specializations directly.
|
||||
*/
|
||||
#ifndef TVM_SUPPORT_SERIALIZER_H_
|
||||
#define TVM_SUPPORT_SERIALIZER_H_
|
||||
|
||||
#include <dlpack/dlpack.h>
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/container/map.h>
|
||||
#include <tvm/ffi/endian.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/support/io.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <type_traits>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace support {
|
||||
|
||||
// ---- Arithmetic types ----
|
||||
template <typename T>
|
||||
struct Serializer<T, std::enable_if_t<std::is_arithmetic_v<T>>> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const T& data) {
|
||||
if constexpr (TVM_FFI_IO_NO_ENDIAN_SWAP) {
|
||||
strm->Write(&data, sizeof(T));
|
||||
} else {
|
||||
T copy = data;
|
||||
ffi::ByteSwap(©, sizeof(T), 1);
|
||||
strm->Write(©, sizeof(T));
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, T* data) {
|
||||
bool ok = strm->Read(data, sizeof(T)) == sizeof(T);
|
||||
if constexpr (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
|
||||
ffi::ByteSwap(data, sizeof(T), 1);
|
||||
}
|
||||
return ok;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- Enum types (delegate to underlying arithmetic type) ----
|
||||
template <typename T>
|
||||
struct Serializer<T, std::enable_if_t<std::is_enum_v<T>>> {
|
||||
static constexpr bool enabled = true;
|
||||
using U = std::underlying_type_t<T>;
|
||||
|
||||
static void Write(Stream* strm, const T& data) {
|
||||
Serializer<U>::Write(strm, static_cast<U>(data));
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, T* data) {
|
||||
U val;
|
||||
if (!Serializer<U>::Read(strm, &val)) return false;
|
||||
*data = static_cast<T>(val);
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- std::string ----
|
||||
template <>
|
||||
struct Serializer<std::string> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const std::string& data) {
|
||||
uint64_t sz = static_cast<uint64_t>(data.size());
|
||||
Serializer<uint64_t>::Write(strm, sz);
|
||||
if (sz != 0) {
|
||||
strm->Write(data.data(), data.size());
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, std::string* data) {
|
||||
uint64_t sz;
|
||||
if (!Serializer<uint64_t>::Read(strm, &sz)) return false;
|
||||
data->resize(static_cast<size_t>(sz));
|
||||
if (sz != 0) {
|
||||
size_t nbytes = static_cast<size_t>(sz);
|
||||
return strm->Read(&(*data)[0], nbytes) == nbytes;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- std::vector<T> ----
|
||||
template <typename T>
|
||||
struct Serializer<std::vector<T>> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const std::vector<T>& vec) {
|
||||
uint64_t sz = static_cast<uint64_t>(vec.size());
|
||||
Serializer<uint64_t>::Write(strm, sz);
|
||||
if constexpr (std::is_arithmetic_v<T> && TVM_FFI_IO_NO_ENDIAN_SWAP) {
|
||||
if (sz != 0) {
|
||||
strm->Write(vec.data(), sizeof(T) * vec.size());
|
||||
}
|
||||
} else {
|
||||
for (const auto& v : vec) {
|
||||
Serializer<T>::Write(strm, v);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, std::vector<T>* vec) {
|
||||
uint64_t sz;
|
||||
if (!Serializer<uint64_t>::Read(strm, &sz)) return false;
|
||||
vec->resize(static_cast<size_t>(sz));
|
||||
if constexpr (std::is_arithmetic_v<T> && TVM_FFI_IO_NO_ENDIAN_SWAP) {
|
||||
if (sz != 0) {
|
||||
size_t nbytes = sizeof(T) * static_cast<size_t>(sz);
|
||||
return strm->Read(vec->data(), nbytes) == nbytes;
|
||||
}
|
||||
return true;
|
||||
} else {
|
||||
for (size_t i = 0; i < static_cast<size_t>(sz); ++i) {
|
||||
if (!Serializer<T>::Read(strm, &(*vec)[i])) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// ---- std::pair<A, B> ----
|
||||
template <typename A, typename B>
|
||||
struct Serializer<std::pair<A, B>> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const std::pair<A, B>& data) {
|
||||
Serializer<A>::Write(strm, data.first);
|
||||
Serializer<B>::Write(strm, data.second);
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, std::pair<A, B>* data) {
|
||||
return Serializer<A>::Read(strm, &data->first) && Serializer<B>::Read(strm, &data->second);
|
||||
}
|
||||
};
|
||||
|
||||
// ---- std::unordered_map<K, V> ----
|
||||
template <typename K, typename V>
|
||||
struct Serializer<std::unordered_map<K, V>> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const std::unordered_map<K, V>& data) {
|
||||
std::vector<std::pair<K, V>> vec(data.begin(), data.end());
|
||||
Serializer<std::vector<std::pair<K, V>>>::Write(strm, vec);
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, std::unordered_map<K, V>* data) {
|
||||
std::vector<std::pair<K, V>> vec;
|
||||
if (!Serializer<std::vector<std::pair<K, V>>>::Read(strm, &vec)) return false;
|
||||
data->clear();
|
||||
data->insert(vec.begin(), vec.end());
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- ffi::String ----
|
||||
template <>
|
||||
struct Serializer<ffi::String> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const ffi::String& data) {
|
||||
uint64_t sz = static_cast<uint64_t>(data.size());
|
||||
Serializer<uint64_t>::Write(strm, sz);
|
||||
if (sz != 0) {
|
||||
strm->Write(data.data(), data.size());
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, ffi::String* data) {
|
||||
std::string s;
|
||||
if (!Serializer<std::string>::Read(strm, &s)) return false;
|
||||
*data = ffi::String(std::move(s));
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- ffi::Bytes (binary payload, same wire format as std::string) ----
|
||||
template <>
|
||||
struct Serializer<ffi::Bytes> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const ffi::Bytes& data) {
|
||||
uint64_t sz = static_cast<uint64_t>(data.size());
|
||||
Serializer<uint64_t>::Write(strm, sz);
|
||||
if (sz != 0) {
|
||||
strm->Write(data.data(), data.size());
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, ffi::Bytes* data) {
|
||||
std::string s;
|
||||
if (!Serializer<std::string>::Read(strm, &s)) return false;
|
||||
*data = ffi::Bytes(std::move(s));
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- ffi::Array<T> (binary-compatible with std::vector<T>) ----
|
||||
template <typename T>
|
||||
struct Serializer<ffi::Array<T>> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const ffi::Array<T>& arr) {
|
||||
uint64_t sz = static_cast<uint64_t>(arr.size());
|
||||
Serializer<uint64_t>::Write(strm, sz);
|
||||
for (size_t i = 0; i < arr.size(); ++i) {
|
||||
Serializer<T>::Write(strm, arr[i]);
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, ffi::Array<T>* arr) {
|
||||
uint64_t sz;
|
||||
if (!Serializer<uint64_t>::Read(strm, &sz)) return false;
|
||||
*arr = ffi::Array<T>();
|
||||
for (uint64_t i = 0; i < sz; ++i) {
|
||||
T val;
|
||||
if (!Serializer<T>::Read(strm, &val)) return false;
|
||||
arr->push_back(std::move(val));
|
||||
}
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- ffi::Map<K, V> (binary format: uint64_t count, then key/value pairs) ----
|
||||
template <typename K, typename V>
|
||||
struct Serializer<ffi::Map<K, V>> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const ffi::Map<K, V>& data) {
|
||||
uint64_t sz = static_cast<uint64_t>(data.size());
|
||||
Serializer<uint64_t>::Write(strm, sz);
|
||||
for (const auto& kv : data) {
|
||||
Serializer<K>::Write(strm, kv.first);
|
||||
Serializer<V>::Write(strm, kv.second);
|
||||
}
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, ffi::Map<K, V>* data) {
|
||||
uint64_t sz;
|
||||
if (!Serializer<uint64_t>::Read(strm, &sz)) return false;
|
||||
*data = ffi::Map<K, V>();
|
||||
for (uint64_t i = 0; i < sz; ++i) {
|
||||
K key;
|
||||
V val;
|
||||
if (!Serializer<K>::Read(strm, &key)) return false;
|
||||
if (!Serializer<V>::Read(strm, &val)) return false;
|
||||
data->Set(std::move(key), std::move(val));
|
||||
}
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- DLDataType ----
|
||||
template <>
|
||||
struct Serializer<DLDataType> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const DLDataType& dtype) {
|
||||
Serializer<uint8_t>::Write(strm, dtype.code);
|
||||
Serializer<uint8_t>::Write(strm, dtype.bits);
|
||||
Serializer<uint16_t>::Write(strm, dtype.lanes);
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, DLDataType* dtype) {
|
||||
if (!Serializer<uint8_t>::Read(strm, &dtype->code)) return false;
|
||||
if (!Serializer<uint8_t>::Read(strm, &dtype->bits)) return false;
|
||||
if (!Serializer<uint16_t>::Read(strm, &dtype->lanes)) return false;
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
// ---- DLDevice ----
|
||||
template <>
|
||||
struct Serializer<DLDevice> {
|
||||
static constexpr bool enabled = true;
|
||||
|
||||
static void Write(Stream* strm, const DLDevice& dev) {
|
||||
int32_t device_type = static_cast<int32_t>(dev.device_type);
|
||||
Serializer<int32_t>::Write(strm, device_type);
|
||||
Serializer<int32_t>::Write(strm, dev.device_id);
|
||||
}
|
||||
|
||||
static bool Read(Stream* strm, DLDevice* dev) {
|
||||
int32_t device_type = 0;
|
||||
if (!Serializer<int32_t>::Read(strm, &device_type)) return false;
|
||||
dev->device_type = static_cast<DLDeviceType>(device_type);
|
||||
if (!Serializer<int32_t>::Read(strm, &dev->device_id)) return false;
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace support
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_SUPPORT_SERIALIZER_H_
|
||||
Reference in New Issue
Block a user