chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,78 @@
|
||||
/*
|
||||
* 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/cuda/nvtx.h
|
||||
* \brief NVTX scoped range utility (header-only).
|
||||
*
|
||||
* Provides NVTXScopedRange: a lightweight RAII wrapper over
|
||||
* nvtxRangePush/Pop. When TVM_NVTX_ENABLED is not defined or is 0,
|
||||
* all methods are no-ops compiled away by the optimizer.
|
||||
*/
|
||||
#ifndef TVM_SUPPORT_CUDA_NVTX_H_
|
||||
#define TVM_SUPPORT_CUDA_NVTX_H_
|
||||
|
||||
#include <string>
|
||||
|
||||
#ifndef TVM_NVTX_ENABLED
|
||||
#define TVM_NVTX_ENABLED 0
|
||||
#endif
|
||||
|
||||
#if TVM_NVTX_ENABLED
|
||||
#include <nvtx3/nvToolsExt.h>
|
||||
#endif // TVM_NVTX_ENABLED
|
||||
|
||||
namespace tvm {
|
||||
namespace support {
|
||||
|
||||
/*!
|
||||
* \brief A class to create a NVTX range. No-op if TVM is not built against NVTX.
|
||||
*/
|
||||
class NVTXScopedRange {
|
||||
public:
|
||||
/*! \brief Enter an NVTX scoped range */
|
||||
#if TVM_NVTX_ENABLED
|
||||
explicit NVTXScopedRange(const char* name) { nvtxRangePush(name); }
|
||||
#else
|
||||
explicit NVTXScopedRange(const char* name) {}
|
||||
#endif // TVM_NVTX_ENABLED
|
||||
/*! \brief Enter an NVTX scoped range */
|
||||
explicit NVTXScopedRange(const std::string& name) : NVTXScopedRange(name.c_str()) {}
|
||||
/*! \brief Exit an NVTX scoped range */
|
||||
#if TVM_NVTX_ENABLED
|
||||
~NVTXScopedRange() { nvtxRangePop(); }
|
||||
#else
|
||||
~NVTXScopedRange() {}
|
||||
#endif // TVM_NVTX_ENABLED
|
||||
NVTXScopedRange(const NVTXScopedRange& other) = delete;
|
||||
NVTXScopedRange(NVTXScopedRange&& other) = delete;
|
||||
NVTXScopedRange& operator=(const NVTXScopedRange& other) = delete;
|
||||
NVTXScopedRange& operator=(NVTXScopedRange&& other) = delete;
|
||||
};
|
||||
|
||||
#ifdef _MSC_VER
|
||||
#define TVM_NVTX_FUNC_SCOPE() ::tvm::support::NVTXScopedRange _nvtx_func_scope_(__FUNCSIG__);
|
||||
#else
|
||||
#define TVM_NVTX_FUNC_SCOPE() \
|
||||
::tvm::support::NVTXScopedRange _nvtx_func_scope_(__PRETTY_FUNCTION__);
|
||||
#endif
|
||||
|
||||
} // namespace support
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_SUPPORT_CUDA_NVTX_H_
|
||||
@@ -0,0 +1,129 @@
|
||||
/*
|
||||
* 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/io.h
|
||||
* \brief Binary stream I/O interface.
|
||||
*/
|
||||
#ifndef TVM_SUPPORT_IO_H_
|
||||
#define TVM_SUPPORT_IO_H_
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace tvm {
|
||||
namespace support {
|
||||
|
||||
/*!
|
||||
* \brief Primary Serializer template. Specialize for each serializable type.
|
||||
*
|
||||
* A valid specialization must provide:
|
||||
* static constexpr bool enabled = true;
|
||||
* static void Write(Stream* strm, const T& data);
|
||||
* static bool Read(Stream* strm, T* data);
|
||||
*/
|
||||
template <typename T, typename = void>
|
||||
struct Serializer {
|
||||
static constexpr bool enabled = false;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Abstract binary stream for serialization.
|
||||
*
|
||||
* Subclasses implement the raw Read/Write byte methods.
|
||||
* The template Write/Read methods delegate to Serializer
|
||||
* for endian-aware, type-safe binary I/O.
|
||||
*
|
||||
* \note Subclasses that override Read(void*, size_t) or Write(const void*, size_t)
|
||||
* must add `using Stream::Read;` and `using Stream::Write;` to make the
|
||||
* template overloads visible (C++ name-hiding rule).
|
||||
*/
|
||||
class Stream {
|
||||
public:
|
||||
/*!
|
||||
* \brief Read raw bytes from the stream.
|
||||
* \param ptr Destination buffer.
|
||||
* \param size Number of bytes to read.
|
||||
* \return Number of bytes actually read.
|
||||
*/
|
||||
virtual size_t Read(void* ptr, size_t size) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Write raw bytes to the stream.
|
||||
* \param ptr Source buffer.
|
||||
* \param size Number of bytes to write.
|
||||
* \return Number of bytes actually written.
|
||||
*/
|
||||
virtual size_t Write(const void* ptr, size_t size) = 0;
|
||||
|
||||
/*! \brief Virtual destructor. */
|
||||
virtual ~Stream() = default;
|
||||
|
||||
/*!
|
||||
* \brief Write a typed value using Serializer<T>.
|
||||
* \tparam T The data type (must have a Serializer<T> specialization).
|
||||
* \param data The value to write.
|
||||
*/
|
||||
template <typename T>
|
||||
void Write(const T& data) {
|
||||
static_assert(Serializer<T>::enabled, "No Serializer<T> specialization for this type");
|
||||
Serializer<T>::Write(this, data);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Read a typed value using Serializer<T>.
|
||||
* \tparam T The data type (must have a Serializer<T> specialization).
|
||||
* \param out_data Pointer to receive the value.
|
||||
* \return true on success.
|
||||
*/
|
||||
template <typename T>
|
||||
bool Read(T* out_data) {
|
||||
static_assert(Serializer<T>::enabled, "No Serializer<T> specialization for this type");
|
||||
return Serializer<T>::Read(this, out_data);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Write an array of typed values element by element.
|
||||
* \param data Pointer to the first element.
|
||||
* \param num_elems Number of elements.
|
||||
*/
|
||||
template <typename T>
|
||||
void WriteArray(const T* data, size_t num_elems) {
|
||||
for (size_t i = 0; i < num_elems; ++i) Write<T>(data[i]);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Read an array of typed values element by element.
|
||||
* \param data Pointer to the first element.
|
||||
* \param num_elems Number of elements.
|
||||
* \return true on success.
|
||||
*/
|
||||
template <typename T>
|
||||
bool ReadArray(T* data, size_t num_elems) {
|
||||
for (size_t i = 0; i < num_elems; ++i) {
|
||||
if (!Read<T>(data + i)) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace support
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_SUPPORT_IO_H_
|
||||
@@ -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