chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,96 @@
|
||||
/*
|
||||
* 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/base.h
|
||||
* \brief base macros
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_BASE_H_
|
||||
#define TVM_RUNTIME_BASE_H_
|
||||
|
||||
// TVM runtime fully relies on TVM FFI C API
|
||||
// we will avoid defining extra C APIs here
|
||||
#include <tvm/ffi/c_api.h>
|
||||
|
||||
// TVM version. Overridable at build time via -DTVM_VERSION="..." (scikit-build-core
|
||||
// passes the setuptools_scm-resolved version through CMake). The literal below is the
|
||||
// fallback for a bare build with no override.
|
||||
#ifndef TVM_VERSION
|
||||
#define TVM_VERSION "0.26.dev0"
|
||||
#endif
|
||||
|
||||
// TVM ships two shared libraries: libtvm_compiler and libtvm_runtime.
|
||||
// Each exposes its own DLL macro pair. The two families are defined
|
||||
// independently so that each can be overridden separately by downstream
|
||||
// embedders who need custom visibility on only one of the two libraries.
|
||||
//
|
||||
// TVM_DLL / TVM_DLL_EXPORT: symbols in libtvm_compiler.
|
||||
// - TVM_DLL is dllexport when TVM_EXPORTS is defined (compiler build),
|
||||
// dllimport otherwise (downstream consumers, runtime TUs).
|
||||
// - TVM_DLL_EXPORT is always dllexport.
|
||||
//
|
||||
// TVM_RUNTIME_DLL / TVM_RUNTIME_DLL_EXPORT: symbols in libtvm_runtime.
|
||||
// - TVM_RUNTIME_DLL is dllexport when TVM_RUNTIME_EXPORTS is defined
|
||||
// (runtime build), dllimport otherwise.
|
||||
// - TVM_RUNTIME_DLL_EXPORT is always dllexport.
|
||||
//
|
||||
// On non-MSVC platforms the import/export decision is made by the dynamic
|
||||
// loader, so all four macros expand to visibility("default"). Under
|
||||
// Emscripten they expand to EMSCRIPTEN_KEEPALIVE.
|
||||
#ifdef __EMSCRIPTEN__
|
||||
#include <emscripten/emscripten.h>
|
||||
#endif
|
||||
|
||||
// --- TVM_DLL family (libtvm_compiler) ---
|
||||
#if !defined(TVM_DLL) && defined(__EMSCRIPTEN__)
|
||||
#define TVM_DLL EMSCRIPTEN_KEEPALIVE
|
||||
#define TVM_DLL_EXPORT EMSCRIPTEN_KEEPALIVE
|
||||
#endif
|
||||
#if !defined(TVM_DLL) && defined(_MSC_VER)
|
||||
#ifdef TVM_EXPORTS
|
||||
#define TVM_DLL __declspec(dllexport)
|
||||
#else
|
||||
#define TVM_DLL __declspec(dllimport)
|
||||
#endif
|
||||
#define TVM_DLL_EXPORT __declspec(dllexport)
|
||||
#endif
|
||||
#ifndef TVM_DLL
|
||||
#define TVM_DLL __attribute__((visibility("default")))
|
||||
#define TVM_DLL_EXPORT __attribute__((visibility("default")))
|
||||
#endif
|
||||
|
||||
// --- TVM_RUNTIME_DLL family (libtvm_runtime) ---
|
||||
#if !defined(TVM_RUNTIME_DLL) && defined(__EMSCRIPTEN__)
|
||||
#define TVM_RUNTIME_DLL EMSCRIPTEN_KEEPALIVE
|
||||
#define TVM_RUNTIME_DLL_EXPORT EMSCRIPTEN_KEEPALIVE
|
||||
#endif
|
||||
#if !defined(TVM_RUNTIME_DLL) && defined(_MSC_VER)
|
||||
#ifdef TVM_RUNTIME_EXPORTS
|
||||
#define TVM_RUNTIME_DLL __declspec(dllexport)
|
||||
#else
|
||||
#define TVM_RUNTIME_DLL __declspec(dllimport)
|
||||
#endif
|
||||
#define TVM_RUNTIME_DLL_EXPORT __declspec(dllexport)
|
||||
#endif
|
||||
#ifndef TVM_RUNTIME_DLL
|
||||
#define TVM_RUNTIME_DLL __attribute__((visibility("default")))
|
||||
#define TVM_RUNTIME_DLL_EXPORT __attribute__((visibility("default")))
|
||||
#endif
|
||||
|
||||
#endif // TVM_RUNTIME_BASE_H_
|
||||
@@ -0,0 +1,109 @@
|
||||
/*
|
||||
* 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/c_backend_api.h
|
||||
* \brief TVM runtime backend API.
|
||||
*
|
||||
* The functions defined in this header are intended to be
|
||||
* used by compiled tvm operators, usually user do not need to use these
|
||||
* function directly.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_C_BACKEND_API_H_
|
||||
#define TVM_RUNTIME_C_BACKEND_API_H_
|
||||
|
||||
#include <tvm/runtime/base.h>
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
/*!
|
||||
* \brief Backend function to allocate temporal workspace.
|
||||
*
|
||||
* \note The result allocated space is ensured to be aligned to kTempAllocaAlignment.
|
||||
*
|
||||
* \param nbytes The size of the space requested.
|
||||
* \param device_type The device type which the space will be allocated.
|
||||
* \param device_id The device id which the space will be allocated.
|
||||
* \param dtype_code_hint The type code of the array elements. Only used in
|
||||
* certain backends such as OpenGL.
|
||||
* \param dtype_bits_hint The type bits of the array elements. Only used in
|
||||
* certain backends such as OpenGL.
|
||||
* \return nullptr when error is thrown, a valid ptr if success
|
||||
*/
|
||||
TVM_RUNTIME_DLL void* TVMBackendAllocWorkspace(int device_type, int device_id, uint64_t nbytes,
|
||||
int dtype_code_hint, int dtype_bits_hint);
|
||||
|
||||
/*!
|
||||
* \brief Backend function to free temporal workspace.
|
||||
*
|
||||
* \param ptr The result allocated space pointer.
|
||||
* \param device_type The device type which the space will be allocated.
|
||||
* \param device_id The device id which the space will be allocated.
|
||||
* \return 0 when no error is thrown, -1 when failure happens
|
||||
*
|
||||
* \sa TVMBackendAllocWorkspace
|
||||
*/
|
||||
TVM_RUNTIME_DLL int TVMBackendFreeWorkspace(int device_type, int device_id, void* ptr);
|
||||
|
||||
/*!
|
||||
* \brief Environment for TVM parallel task.
|
||||
*/
|
||||
typedef struct {
|
||||
/*!
|
||||
* \brief Auxiliary used for synchronization
|
||||
*/
|
||||
void* sync_handle;
|
||||
/*! \brief total amount of task */
|
||||
int32_t num_task;
|
||||
} TVMParallelGroupEnv;
|
||||
|
||||
/*!
|
||||
* \brief The callback function to execute a parallel lambda
|
||||
* \param task_id the task id of the function.
|
||||
* \param penv The parallel environment backs the execution.
|
||||
* \param cdata The supporting closure data.
|
||||
*/
|
||||
typedef int (*FTVMParallelLambda)(int task_id, TVMParallelGroupEnv* penv, void* cdata);
|
||||
|
||||
/*!
|
||||
* \brief Backend function for running parallel jobs.
|
||||
*
|
||||
* \param flambda The parallel function to be launched.
|
||||
* \param cdata The closure data.
|
||||
* \param num_task Number of tasks to launch, can be 0, means launch
|
||||
* with all available threads.
|
||||
*
|
||||
* \return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
TVM_RUNTIME_DLL int TVMBackendParallelLaunch(FTVMParallelLambda flambda, void* cdata, int num_task);
|
||||
|
||||
/*!
|
||||
* \brief BSP barrrier between parallel threads
|
||||
* \param task_id the task id of the function.
|
||||
* \param penv The parallel environment backs the execution.
|
||||
* \return 0 when no error is thrown, -1 when failure happens
|
||||
*/
|
||||
TVM_RUNTIME_DLL int TVMBackendParallelBarrier(int task_id, TVMParallelGroupEnv* penv);
|
||||
|
||||
#ifdef __cplusplus
|
||||
} // TVM_EXTERN_C
|
||||
#endif
|
||||
#endif // TVM_RUNTIME_C_BACKEND_API_H_
|
||||
@@ -0,0 +1,427 @@
|
||||
/*
|
||||
* 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/device_api.h
|
||||
* \brief Abstract device memory management API
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_DEVICE_API_H_
|
||||
#define TVM_RUNTIME_DEVICE_API_H_
|
||||
|
||||
#include <tvm/ffi/any.h>
|
||||
#include <tvm/ffi/error.h>
|
||||
#include <tvm/ffi/optional.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/base.h>
|
||||
|
||||
#include <string>
|
||||
/*!
|
||||
* \brief The stream that is specific to device
|
||||
* can be NULL, which indicates the default one.
|
||||
*/
|
||||
typedef void* TVMStreamHandle;
|
||||
|
||||
namespace tvm {
|
||||
|
||||
// alias DLDevice
|
||||
using Device = DLDevice;
|
||||
|
||||
namespace runtime {
|
||||
|
||||
/*! \brief Extension device types in TVM
|
||||
*
|
||||
* Additional enumerators to supplement those provided by
|
||||
* DLPack's `DLDeviceType` enumeration.
|
||||
*
|
||||
* MAINTAINERS NOTE #1: We need to ensure that the two devices
|
||||
* are identified by the same integer.
|
||||
* Currently this requires manual verification.
|
||||
* Discussed here: https://github.com/dmlc/dlpack/issues/111
|
||||
* As of DLPack v0.7, the highest-valued enumerator in
|
||||
* `DLDeviceType` is kDLHexagon = 16.
|
||||
*
|
||||
* MAINTAINERS NOTE #2: As of DLPack v0.7, the definition for
|
||||
* `DLDeviceType` specifies an underlying storage type of
|
||||
* `int32_t`. That guarantees a variable of type
|
||||
* `DLDeviceType` is capable of holding any integers provided
|
||||
* by *either* of these enumerations.
|
||||
*
|
||||
* However, the `int32_t` specification only applies when the
|
||||
* header file is compiled as C++, and this header file is also
|
||||
* meant to work as C code. So the unspecified storage type
|
||||
* could be a latent bug when compiled as C.
|
||||
*/
|
||||
#ifdef __cplusplus
|
||||
typedef enum : int32_t {
|
||||
#else
|
||||
typedef enum {
|
||||
#endif
|
||||
// To help avoid accidental conflicts between `DLDeviceType`
|
||||
// and this enumeration, start numbering the new enumerators
|
||||
// a little higher than (currently) seems necessary.
|
||||
TVMDeviceExtType_End = 36, // sentinel value
|
||||
} TVMDeviceExtType;
|
||||
|
||||
/*!
|
||||
* \brief the query type into GetAttr
|
||||
*/
|
||||
enum DeviceAttrKind : int {
|
||||
kExist = 0,
|
||||
kMaxThreadsPerBlock = 1,
|
||||
kWarpSize = 2,
|
||||
kMaxSharedMemoryPerBlock = 3,
|
||||
kComputeVersion = 4,
|
||||
kDeviceName = 5,
|
||||
kMaxClockRate = 6,
|
||||
kMultiProcessorCount = 7,
|
||||
kMaxThreadDimensions = 8,
|
||||
kMaxRegistersPerBlock = 9,
|
||||
kGcnArch = 10,
|
||||
kApiVersion = 11,
|
||||
kDriverVersion = 12,
|
||||
kL2CacheSizeBytes = 13,
|
||||
kTotalGlobalMemory = 14,
|
||||
kAvailableGlobalMemory = 15,
|
||||
kImagePitchAlignment = 16,
|
||||
};
|
||||
|
||||
#ifdef TVM_KALLOC_ALIGNMENT
|
||||
/*! \brief Number of bytes each allocation must align to */
|
||||
constexpr int kAllocAlignment = TVM_KALLOC_ALIGNMENT;
|
||||
|
||||
/*! \brief Number of bytes each allocation must align to in temporary allocation */
|
||||
constexpr int kTempAllocaAlignment = TVM_KALLOC_ALIGNMENT;
|
||||
#else
|
||||
/*! \brief Number of bytes each allocation must align to */
|
||||
constexpr int kAllocAlignment = 64;
|
||||
|
||||
/*! \brief Number of bytes each allocation must align to in temporary allocation */
|
||||
constexpr int kTempAllocaAlignment = 64;
|
||||
#endif // TVM_KALLOC_ALIGNMENT
|
||||
|
||||
/*! \brief Maximum size that can be allocated on stack */
|
||||
constexpr int kMaxStackAlloca = 1024;
|
||||
|
||||
/*! \brief Number of bytes each allocation must align to by default in the workspace buffer to
|
||||
* service intermediate tensors */
|
||||
constexpr int kDefaultWorkspaceAlignment = 1;
|
||||
|
||||
/*!
|
||||
* \brief TVM Runtime Device API, abstracts the device
|
||||
* specific interface for memory management.
|
||||
*/
|
||||
class TVM_RUNTIME_DLL DeviceAPI {
|
||||
public:
|
||||
/*! \brief virtual destructor */
|
||||
virtual ~DeviceAPI() {}
|
||||
/*!
|
||||
* \brief Set the environment device id to device
|
||||
* \param dev The device to be set.
|
||||
*/
|
||||
virtual void SetDevice(Device dev) = 0;
|
||||
/*!
|
||||
* \brief Get attribute of specified device.
|
||||
* \param dev The device device
|
||||
* \param kind The result kind
|
||||
* \param rv The return value.
|
||||
* \sa DeviceAttrKind
|
||||
*/
|
||||
virtual void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Get the physical memory size required.
|
||||
* \param arr the tensor object.
|
||||
* \param mem_scope the memory scope if any
|
||||
* \return the memory size.
|
||||
*/
|
||||
virtual size_t GetDataSize(const DLTensor& arr,
|
||||
ffi::Optional<ffi::String> mem_scope = std::nullopt);
|
||||
|
||||
/*!
|
||||
* \brief Query the device for specified properties.
|
||||
*
|
||||
* This is used to expand "-from_device=N" in the target string to
|
||||
* all properties that can be determined from that device.
|
||||
*/
|
||||
virtual void GetTargetProperty(Device dev, const std::string& property, ffi::Any* rv) {}
|
||||
|
||||
/*!
|
||||
* \brief Allocate a data space on device.
|
||||
* \param dev The device device to perform operation.
|
||||
* \param nbytes The number of bytes in memory.
|
||||
* \param alignment The alignment of the memory.
|
||||
* \param type_hint The type of elements. Only needed by certain backends such
|
||||
* as OpenGL, as nbytes & alignment are sufficient for most backends.
|
||||
* \return The allocated device pointer.
|
||||
*/
|
||||
virtual void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment,
|
||||
DLDataType type_hint) = 0;
|
||||
/*!
|
||||
* \brief Allocate a data space on device with memory scope support.
|
||||
* \param dev The device device to perform operation.
|
||||
* \param ndim The number of dimension of allocated tensor.
|
||||
* \param shape The shape of allocated tensor.
|
||||
* \param dtype The type of elements.
|
||||
* \param mem_scope The memory scope of allocated tensor.
|
||||
* \return The allocated device pointer.
|
||||
*/
|
||||
virtual void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
|
||||
ffi::Optional<ffi::String> mem_scope = std::nullopt);
|
||||
/*!
|
||||
* \brief Free a data space on device.
|
||||
* \param dev The device device to perform operation.
|
||||
* \param ptr The data space.
|
||||
*/
|
||||
virtual void FreeDataSpace(Device dev, void* ptr) = 0;
|
||||
/*!
|
||||
* \brief copy data from one place to another
|
||||
* \note This API is designed to support special memory with shape dependent layout.
|
||||
* We pass in DLTensor* with shape information to support these cases.
|
||||
* \param from The source array.
|
||||
* \param to The target array.
|
||||
* \param stream Optional stream object.
|
||||
* \note The copy may happen asynchronously if it involves a GPU context.
|
||||
* Call StreamSync to ensure the copy completes from host's pov.
|
||||
*/
|
||||
virtual void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream);
|
||||
/*!
|
||||
* \brief Create a new stream of execution.
|
||||
*
|
||||
* \param dev The device of allocation.
|
||||
*/
|
||||
virtual TVMStreamHandle CreateStream(Device dev);
|
||||
|
||||
/*!
|
||||
* \brief Free a stream of execution
|
||||
*
|
||||
* \param dev The device of the stream
|
||||
* \param stream The pointer to be freed.
|
||||
*/
|
||||
virtual void FreeStream(Device dev, TVMStreamHandle stream);
|
||||
|
||||
/*!
|
||||
* \brief Synchronize the stream
|
||||
* \param dev The device to perform operation.
|
||||
* \param stream The stream to be sync.
|
||||
*/
|
||||
virtual void StreamSync(Device dev, TVMStreamHandle stream) = 0;
|
||||
/*!
|
||||
* \brief Set the stream
|
||||
* \param dev The device to set stream.
|
||||
* \param stream The stream to be set.
|
||||
*/
|
||||
virtual void SetStream(Device dev, TVMStreamHandle stream);
|
||||
/*!
|
||||
* \brief Get the current stream
|
||||
* \param dev The device to get stream.
|
||||
* \return The current stream of the device.
|
||||
*/
|
||||
virtual TVMStreamHandle GetCurrentStream(Device dev);
|
||||
/*!
|
||||
* \brief Synchronize 2 streams of execution.
|
||||
*
|
||||
* An event is created in event_src stream that the second then
|
||||
* stream waits on. Neither event_src or event_dst need to be of
|
||||
* the same device ID as the device, but they must be of the same
|
||||
* device type.
|
||||
*
|
||||
* \param dev The device of the streams.
|
||||
* \param event_src The source stream to synchronize.
|
||||
* \param event_dst The destination stream to synchronize.
|
||||
*/
|
||||
virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst);
|
||||
/*!
|
||||
* \brief Allocate temporal workspace for backend execution.
|
||||
*
|
||||
* \note We have the following assumption about backend temporal
|
||||
* workspace allocation, and backend will optimize for such assumption:
|
||||
*
|
||||
* - Only a few allocation will happen, and space will be released after use.
|
||||
* - The release order is usually in reverse order of allocate (stack style).
|
||||
* - Repeative pattern of same allocations over different runs.
|
||||
* - Workspace should not overlap between different threads(i.e. be threadlocal)
|
||||
*
|
||||
* \param dev The device of allocation.
|
||||
* \param nbytes The size to be allocated.
|
||||
* \param type_hint The type of elements. Only needed by certain backends such
|
||||
* as OpenGL, as nbytes is sufficient for most backends.
|
||||
*/
|
||||
virtual void* AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint = {});
|
||||
/*!
|
||||
* \brief Free temporal workspace in backend execution.
|
||||
*
|
||||
* \param dev The device of allocation.
|
||||
* \param ptr The pointer to be freed.
|
||||
*/
|
||||
virtual void FreeWorkspace(Device dev, void* ptr);
|
||||
|
||||
/*!
|
||||
* \brief Get device API based on device.
|
||||
* \param dev The device
|
||||
* \param allow_missing Whether allow missing
|
||||
* \return The corresponding device API.
|
||||
*/
|
||||
static DeviceAPI* Get(Device dev, bool allow_missing = false);
|
||||
|
||||
/*!
|
||||
* \brief Whether a certian device type requires set device device
|
||||
* before launching the kernel function.
|
||||
* \param device_type The device type.
|
||||
*/
|
||||
static bool NeedSetDevice(int device_type) { return device_type != kDLCPU; }
|
||||
|
||||
/*!
|
||||
* \brief Whether pointer arithmetics on a device owned pointer may be performed on the host.
|
||||
*/
|
||||
virtual bool SupportsDevicePointerArithmeticsOnHost() { return false; }
|
||||
|
||||
protected:
|
||||
/*!
|
||||
* \brief copy data from one place to another
|
||||
* \param from The source array.
|
||||
* \param from_offset The byte offeset in the from.
|
||||
* \param to The target array.
|
||||
* \param to_offset The byte offset in the to.
|
||||
* \param num_bytes The size of the memory in bytes
|
||||
* \param dev_from The source device
|
||||
* \param dev_to The target device
|
||||
* \param type_hint The type of elements, only neded by certain backends.
|
||||
* can be useful for cross device endian converison.
|
||||
* \param stream Optional stream object.
|
||||
*/
|
||||
virtual void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset,
|
||||
size_t num_bytes, Device dev_from, Device dev_to,
|
||||
DLDataType type_hint, TVMStreamHandle stream);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The name of DLDeviceType.
|
||||
* \param type The device type.
|
||||
* \return the device name.
|
||||
*/
|
||||
inline const char* DLDeviceType2Str(int type) {
|
||||
switch (type) {
|
||||
case kDLCPU:
|
||||
return "cpu";
|
||||
case kDLCUDA:
|
||||
return "cuda";
|
||||
case kDLCUDAHost:
|
||||
return "cuda_host";
|
||||
case kDLCUDAManaged:
|
||||
return "cuda_managed";
|
||||
case kDLOpenCL:
|
||||
return "opencl";
|
||||
case kDLVulkan:
|
||||
return "vulkan";
|
||||
case kDLMetal:
|
||||
return "metal";
|
||||
case kDLVPI:
|
||||
return "vpi";
|
||||
case kDLROCM:
|
||||
return "rocm";
|
||||
case kDLROCMHost:
|
||||
return "rocm_host";
|
||||
case kDLExtDev:
|
||||
return "ext_dev";
|
||||
case kDLOneAPI:
|
||||
return "oneapi";
|
||||
case kDLWebGPU:
|
||||
return "webgpu";
|
||||
case kDLHexagon:
|
||||
return "hexagon";
|
||||
case kDLTrn:
|
||||
return "trn";
|
||||
default:
|
||||
TVM_FFI_THROW(InternalError) << "unknown type = " << type;
|
||||
}
|
||||
throw;
|
||||
}
|
||||
|
||||
/*! \brief The device type bigger than this is RPC device */
|
||||
constexpr int kRPCSessMask = 128;
|
||||
static_assert(kRPCSessMask >= TVMDeviceExtType_End);
|
||||
|
||||
/*!
|
||||
* \brief Return true if a Device is owned by an RPC session.
|
||||
*/
|
||||
inline bool IsRPCSessionDevice(Device dev) { return (dev.device_type / kRPCSessMask) > 0; }
|
||||
|
||||
/*!
|
||||
* \brief Return the RPCSessTable index of the RPC Session that owns this device.
|
||||
* \return the table index.
|
||||
*/
|
||||
inline int GetRPCSessionIndex(Device dev) {
|
||||
TVM_FFI_ICHECK(IsRPCSessionDevice(dev)) << "GetRPCSessionIndex: dev has no RPC session";
|
||||
return dev.device_type / kRPCSessMask - 1;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Remove the RPC session mask from a Device.
|
||||
* RPC clients typically do this when encoding a Device for transmission to an RPC remote.
|
||||
* On the wire, RPCdevice are expected to be valid on the server without interpretation.
|
||||
* \param dev A Device with non-zero RPC Session mask, valid on the RPC client.
|
||||
* \return A Device without any RPC Session mask, valid on the RPC server.
|
||||
*/
|
||||
inline Device RemoveRPCSessionMask(Device dev) {
|
||||
dev.device_type = static_cast<DLDeviceType>(dev.device_type % kRPCSessMask);
|
||||
return dev;
|
||||
}
|
||||
|
||||
inline std::ostream& operator<<(std::ostream& os, DLDevice dev) { // NOLINT(*)
|
||||
if (tvm::runtime::IsRPCSessionDevice(dev)) {
|
||||
os << "remote[" << tvm::runtime::GetRPCSessionIndex(dev) << "]-";
|
||||
dev = tvm::runtime::RemoveRPCSessionMask(dev);
|
||||
}
|
||||
os << tvm::runtime::DLDeviceType2Str(static_cast<int>(dev.device_type)) << ":" << dev.device_id;
|
||||
return os;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Add a RPC session mask to a Device.
|
||||
* RPC clients typically do this when decoding a Device received from a RPC remote.
|
||||
* \param dev A Device without any RPC Session mask, valid on the RPC server.
|
||||
* \param session_table_index Numeric index of the RPC session in the session table.
|
||||
* \return A Device with RPC session mask added, valid on the RPC client.
|
||||
*/
|
||||
inline Device AddRPCSessionMask(Device dev, int session_table_index) {
|
||||
TVM_FFI_ICHECK(!IsRPCSessionDevice(dev))
|
||||
<< "AddRPCSessionMask: dev already non-zero RPCSessionIndex: " << dev;
|
||||
dev.device_type =
|
||||
static_cast<DLDeviceType>(dev.device_type | (kRPCSessMask * (session_table_index + 1)));
|
||||
return dev;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Check if runtime module is enabled for target.
|
||||
* \param target The target module name.
|
||||
* \return Whether runtime is enabled.
|
||||
*/
|
||||
TVM_RUNTIME_DLL bool RuntimeEnabled(const ffi::String& target);
|
||||
|
||||
/*! \brief namespace for constant symbols */
|
||||
namespace symbol {
|
||||
constexpr const char* tvm_global_barrier_state = "__tvm_global_barrier_state";
|
||||
/*! \brief global function to set device */
|
||||
constexpr const char* tvm_set_device = "__tvm_set_device";
|
||||
} // namespace symbol
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_DEVICE_API_H_
|
||||
@@ -0,0 +1,154 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_DISCO_BUILTIN_H_
|
||||
#define TVM_RUNTIME_DISCO_BUILTIN_H_
|
||||
|
||||
#include <tvm/ffi/dtype.h>
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <string>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
/*!
|
||||
* \brief Possible kinds of reduction operations.
|
||||
*/
|
||||
enum class ReduceKind : int32_t {
|
||||
kSum = 0,
|
||||
kProd = 1,
|
||||
kMin = 2,
|
||||
kMax = 3,
|
||||
kAvg = 4,
|
||||
};
|
||||
|
||||
/*! \brief Converts `ReduceKind` to string */
|
||||
inline std::string ReduceKind2String(ReduceKind kind) {
|
||||
switch (kind) {
|
||||
case ReduceKind::kSum:
|
||||
return "kSum";
|
||||
case ReduceKind::kProd:
|
||||
return "kProd";
|
||||
case ReduceKind::kMin:
|
||||
return "kMin";
|
||||
case ReduceKind::kMax:
|
||||
return "kMax";
|
||||
case ReduceKind::kAvg:
|
||||
return "kAvg";
|
||||
}
|
||||
TVM_FFI_THROW(ValueError) << "Unknown ReduceKind: " << static_cast<int>(kind);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Load a runtime Module, then create and initialize a RelaxVM
|
||||
* \param path The path to the runtime Module (a DSO file) to be loaded
|
||||
* \param device The default device used to initialize the RelaxVM
|
||||
* \return The RelaxVM as a runtime Module
|
||||
*/
|
||||
TVM_RUNTIME_DLL ffi::Module LoadVMModule(std::string path, ffi::Optional<Device> device);
|
||||
/*!
|
||||
* \brief Create an uninitialized empty Tensor
|
||||
* \param shape The shape of the Tensor
|
||||
* \param dtype The dtype of the Tensor
|
||||
* \param device The device the Tensor is created on. If None, use the thread local default device
|
||||
* \return The Tensor created
|
||||
*/
|
||||
TVM_RUNTIME_DLL Tensor DiscoEmptyTensor(ffi::Shape shape, DLDataType dtype,
|
||||
ffi::Optional<Device> device);
|
||||
/*!
|
||||
* \brief Perform an allreduce operation using the underlying communication library
|
||||
* \param send The array send to perform allreduce on
|
||||
* \param reduce_kind The kind of reduction operation (e.g. sum, avg, min, max)
|
||||
* \param in_group Whether the allreduce operation performs globally or in group as default.
|
||||
* \param recv The array receives the outcome of allreduce
|
||||
*/
|
||||
TVM_RUNTIME_DLL void AllReduce(Tensor send, ReduceKind reduce_kind, bool in_group, Tensor recv);
|
||||
/*!
|
||||
* \brief Perform an allgather operation using the underlying communication library
|
||||
* \param send The array send to perform allgather on
|
||||
* \param in_group Whether the allgather operation performs globally or in group as default.
|
||||
* \param recv The array receives the outcome of allgather
|
||||
*/
|
||||
TVM_RUNTIME_DLL void AllGather(Tensor send, bool in_group, Tensor recv);
|
||||
/*!
|
||||
* \brief Perform a broadcast operation from worker-0
|
||||
* \param send The buffer to be broadcasted
|
||||
* \param in_group Whether the broadcast operation performs globally or in group as default.
|
||||
* \param recv The buffer receives the broadcasted array
|
||||
*/
|
||||
TVM_RUNTIME_DLL void BroadcastFromWorker0(Tensor send, bool in_group, Tensor recv);
|
||||
/*!
|
||||
* \brief Perform a scatter operation from worker-0, chunking the given buffer into equal parts.
|
||||
* \param send For worker-0, it must be provided, and otherwise, the buffer must be None.
|
||||
* The buffer will be divided into equal parts and sent to each worker accordingly.
|
||||
* \param in_group Whether the scatter operation performs globally or in group as default.
|
||||
* \param recv The receiving buffer, which must not be None.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void ScatterFromWorker0(ffi::Optional<Tensor> send, bool in_group, Tensor recv);
|
||||
/*!
|
||||
* \brief Perform a gather operation to worker-0.
|
||||
* \param send The sending buffer, which must not be None.
|
||||
* \param in_group Whether the gather operation performs globally or in group as default.
|
||||
* \param recv For worker-0, it must be provided, and otherwise, the buffer must be None. The
|
||||
* receiving buffer will be divided into equal parts and receive from each worker accordingly.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void GatherToWorker0(Tensor send, bool in_group, ffi::Optional<Tensor> recv);
|
||||
/*!
|
||||
* \brief Receive a buffer from worker-0. No-op if the current worker is worker-0.
|
||||
* \param buffer The buffer to be received
|
||||
*/
|
||||
TVM_RUNTIME_DLL void RecvFromWorker0(Tensor buffer);
|
||||
/*!
|
||||
* \brief Send a buffer to the corresponding worker in the next group.
|
||||
* An error is thrown if the worker is already in the last group.
|
||||
* \param buffer The sending buffer.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void SendToNextGroup(Tensor buffer);
|
||||
/*!
|
||||
* \brief Receive a buffer from the corresponding worker in the previous group.
|
||||
* An error is thrown if the worker is already in the first group.
|
||||
* \param buffer The receiving buffer.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void RecvFromPrevGroup(Tensor buffer);
|
||||
/*!
|
||||
* \brief Send a buffer to the target receiver worker (globally across all groups).
|
||||
* \param buffer The sending buffer.
|
||||
* \param receiver_id The global receiver worker id.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void SendToWorker(Tensor buffer, int receiver_id);
|
||||
/*!
|
||||
* \brief Receive a buffer from the target sender worker (globally across all groups).
|
||||
* \param buffer The receiving buffer.
|
||||
* \param sender_id The global sender worker id.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void RecvFromWorker(Tensor buffer, int sender_id);
|
||||
/*! \brief Get the local worker id */
|
||||
TVM_RUNTIME_DLL int WorkerId();
|
||||
/*!
|
||||
* \brief Called by the worker thread. Waiting until the worker completes all its tasks.
|
||||
* As a specific example, on a CUDA worker, it blocks until all kernels are launched and
|
||||
* cudaStreamSynchronize is complete.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void SyncWorker();
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_DISCO_BUILTIN_H_
|
||||
@@ -0,0 +1,99 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#ifndef TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
|
||||
#define TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
|
||||
|
||||
#include <tvm/runtime/base.h>
|
||||
#include <tvm/runtime/memory/memory_manager.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace cuda_ipc {
|
||||
|
||||
/*!
|
||||
* \brief The CUDA IPC (interprocess communication) memory object,
|
||||
* which internally contains data pointers to CUDA IPC memory.
|
||||
* It is be useful for efficient all-reduce implementation.
|
||||
* \note Right now the class members are closely tied with customized
|
||||
* all-reduce kernel. They may also be extended for other uses in
|
||||
* the future.
|
||||
*/
|
||||
class CUDAIPCMemoryObj : public ffi::Object {
|
||||
public:
|
||||
/*! \brief The number of GPU workers. */
|
||||
int num_workers;
|
||||
/*! \brief The worker id corresponding to this IPC memory object. */
|
||||
int worker_id;
|
||||
/*!
|
||||
* \brief The data pointers of all all-reduce inputs.
|
||||
* It has "num_workers" pointers. The i-th pointer is the data pointer on worker i.
|
||||
* If "i != worker_id", the pointer is an IPC data pointer.
|
||||
* Otherwise, the pointer is a local CUDA data pointer.
|
||||
*/
|
||||
std::vector<void*> remote_data;
|
||||
|
||||
// We introduce the barrier helper data below per CUDAIPCMemory object
|
||||
// so that they can be used by custom collective operations and allow
|
||||
// fine-grained synchronization on each buffer. These barriers have
|
||||
// low overhead, and can potentially enable concurrent execution of
|
||||
// kernels in future.
|
||||
/*!
|
||||
* \brief The pointers to input barrier signals of all workers for all-reduce.
|
||||
* It has "num_workers" pointers, and the pointer arrangement is the same as "remote_data".
|
||||
*/
|
||||
std::vector<void*> barrier_in;
|
||||
/*!
|
||||
* \brief The pointers to output barrier signals of all workers for all-reduce.
|
||||
* It has "num_workers" pointers, and the pointer arrangement is the same as "remote_data".
|
||||
*/
|
||||
std::vector<void*> barrier_out;
|
||||
/*! \brief The integer buffer flag for all-reduce. */
|
||||
int barrier_flag;
|
||||
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("tvm.runtime.disco.cuda_ipc_memory", CUDAIPCMemoryObj, ffi::Object);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to CUDAIPCMemoryObj.
|
||||
* \sa CUDAIPCMemory
|
||||
*/
|
||||
class CUDAIPCMemory : public ffi::ObjectRef {
|
||||
public:
|
||||
/*! \brief Get the global singleton CUDAIPCMemory allocator. */
|
||||
TVM_RUNTIME_DLL static memory::Allocator* GlobalAllocator();
|
||||
/*!
|
||||
* \brief Given a local CUDA data pointer, return the CUDAIPCMemory object of the pointer.
|
||||
* \note The pointer's CUDAIPCMemory is expected to have been allocated
|
||||
* through global function "cuda_ipc.alloc_storage". Or otherwise this
|
||||
* function will raise exception.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static CUDAIPCMemory GetIPCMemoryFromDevicePtr(void* ptr);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(CUDAIPCMemory, ffi::ObjectRef, CUDAIPCMemoryObj);
|
||||
};
|
||||
|
||||
} // namespace cuda_ipc
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
|
||||
@@ -0,0 +1,118 @@
|
||||
/*
|
||||
* 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 disco_worker.h
|
||||
* \brief This file defines a worker in Disco. A worker can be launched in a separate thread or
|
||||
* process as long as the channel supports bi-directional communication in-between the worker and
|
||||
* the controler.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_DISCO_DISCO_WORKER_H_
|
||||
#define TVM_RUNTIME_DISCO_DISCO_WORKER_H_
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/disco/session.h>
|
||||
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
/*!
|
||||
* \brief A worker in Disco. It takes a channel to communication with the controler.
|
||||
* The worker can be run in a separate thread or process as long as the channel supports
|
||||
* bi-directional communication in-between.
|
||||
*/
|
||||
class DiscoWorker {
|
||||
public:
|
||||
/*!
|
||||
* \brief Construct a worker.
|
||||
* \param worker_id The id of the worker.
|
||||
* \param num_workers The number of the workers.
|
||||
* \param num_groups The number of the worker groups.
|
||||
* \param worker_zero_data The data shared between worker-0 and the controler. It's a nullptr if
|
||||
* the worker is not worker-0.
|
||||
* \param channel The communication channel between the worker and the controler.
|
||||
*/
|
||||
explicit DiscoWorker(int worker_id, int num_workers, int num_groups,
|
||||
WorkerZeroData* worker_zero_data, DiscoChannel* channel)
|
||||
: worker_id(worker_id),
|
||||
local_worker_id(worker_id),
|
||||
num_workers(num_workers),
|
||||
num_groups(num_groups),
|
||||
default_device(Device{DLDeviceType::kDLCPU, 0}),
|
||||
worker_zero_data(worker_zero_data),
|
||||
channel(channel),
|
||||
register_file{} {}
|
||||
|
||||
/*! \brief Main loop of the worker */
|
||||
void MainLoop();
|
||||
/*! \brief Get the worker instance on the current thread */
|
||||
TVM_RUNTIME_DLL static DiscoWorker* ThreadLocal();
|
||||
/*! \brief Set the specific register to a specific value */
|
||||
void SetRegister(int reg_id, ffi::AnyView value);
|
||||
|
||||
/*! \brief The id of the worker.*/
|
||||
int worker_id;
|
||||
/*! \brief The local id of the worker. This can be different from worker_id if the session is
|
||||
* consisted with multiple sub-sessions. */
|
||||
int local_worker_id;
|
||||
/*! \brief Total number of workers */
|
||||
int num_workers;
|
||||
/*! \brief Total number of workers */
|
||||
int num_groups;
|
||||
/*! \brief The default device to allocate data if not specified */
|
||||
Device default_device;
|
||||
/*! \brief The name of the underlying collective communication library. */
|
||||
ffi::String ccl;
|
||||
/*!
|
||||
* \brief The data shared between worker-0 and the controler. It's a nullptr if
|
||||
* the worker is not worker-0.
|
||||
* \note This data structure is owned by the controler.
|
||||
*/
|
||||
WorkerZeroData* worker_zero_data;
|
||||
/*!
|
||||
* \brief The communication channel between the worker and the controler.
|
||||
* \note This data structure is owned by the controler.
|
||||
*/
|
||||
DiscoChannel* channel;
|
||||
/*! \brief The registers in the worker */
|
||||
std::vector<ffi::Any> register_file;
|
||||
|
||||
struct Impl;
|
||||
friend struct DiscoWorker::Impl;
|
||||
};
|
||||
/*!
|
||||
* \brief A threadlocal wrapper of DiscoWorker.
|
||||
*/
|
||||
struct ThreadLocalDiscoWorker {
|
||||
/*! \brief The Disco worker */
|
||||
DiscoWorker* worker;
|
||||
|
||||
/*!
|
||||
* \brief Get the threadlocal Disco worker.
|
||||
*/
|
||||
static ThreadLocalDiscoWorker* Get() {
|
||||
thread_local static ThreadLocalDiscoWorker worker;
|
||||
return &worker;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
#endif // TVM_RUNTIME_DISCO_DISCO_WORKER_H_
|
||||
@@ -0,0 +1,385 @@
|
||||
/*
|
||||
* 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 session.h
|
||||
* \brief This file serves as the entry point of Disco and defines key data structures and
|
||||
* interfaces.
|
||||
*
|
||||
* Disco is a distributed runtime that consists of a controler and a cluster of workers. The
|
||||
* controler is responsible for managing the workers by broadcasting commands to all the workers
|
||||
* together, and the workers are responsible for executing the commands and. The controler and
|
||||
* workers communicate with each other through a bi-directional channel.
|
||||
*
|
||||
* Different from a generic system, Disco is designed to as "single-program-multiple-data" (SPMD)
|
||||
* runtime, which means that all the workers execute the same instruction at the same time, but the
|
||||
* data they are working on may be different. For example, in data parallelism, each worker may
|
||||
* work on a different batches of the data, but they all execute the same set of instructions.
|
||||
* Therefore, imagine there is a virtual machine that executes the program, the structures of
|
||||
* workers' register files could be considered as "identical" (single program) although the values
|
||||
* may differ (multiple data).
|
||||
*
|
||||
*
|
||||
* **DRef.** Following the design above, consider the program in SPMD in a virtual ISA, then each
|
||||
* worker is a virtual machine instance to execute the ISA maintaining its own register file.
|
||||
* The controler denotes each of their register files with a unique integer "register id",
|
||||
* and the workers use this id to refer to the register file that resides on itself.
|
||||
* DRef is a control-side object backed by such a register id. The data it contains is not assumed
|
||||
* to be directly accessible by the controler, with an exception for worker-0, which is a special
|
||||
* worker that is always co-located with the controler.
|
||||
*
|
||||
* **Worker-0.** Worker-0 is a special worker that is always co-located with the controler.
|
||||
* It is assumed that the controler can synchronize with and access the registers of worker-0.
|
||||
* The Disco session provides multiple APIs to interact specifically with the worker-0.
|
||||
* To shared data with other workers, a common paradigm in Disco is to copy data from the
|
||||
* controler-side Tensor to the worker-0, and then copy it to other workers using primitives on
|
||||
* the data plane, for example, `broadcast` and `send`.
|
||||
*
|
||||
* **Control plane.** The controler broadcasts commands to all the workers as control signals.
|
||||
* For example, the control may ask all workers to load a library or call a function respectively.
|
||||
* Common control signals include: shutdown, retrievel a global ffi::Function, call packed function,
|
||||
* etc. The controler is assumed to keep a message channel to each worker to implement the broadcast
|
||||
* behavior, and the message channel may vary depends on usecases.
|
||||
*
|
||||
* **Data plane.** The data channel is usually used to exchange data between workers, especially for
|
||||
* tensor data which is usually large. For example, performing an allreduce operator for sharded
|
||||
* matrix multiplication, or broadcasting for an input tensor. For efficiency, the data channel is
|
||||
* usually backed by NCCL on NVIDIA GPUs, RCCL on AMD GPUs, or MPI on CPUs.
|
||||
*
|
||||
* **Session.** A Disco session is a primary interface to interact with the Disco runtime, serving
|
||||
* as a global context that manages the control and workers. It could be implemented as a
|
||||
* multi-threaded with a pool of workers for single-node multi-gpu scenarios, or TCP sockets for
|
||||
* workloads that span over a cluster of nodes.
|
||||
*
|
||||
* **Channel.** Disco channel is a bi-directional communication channel between the controler and
|
||||
* workers for exchanging control signals. It is no different from a generic RPC channel, but
|
||||
* adopts TVM's ffi::Function calling convention to support polymorphic and variadic arguments.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_DISCO_SESSION_H_
|
||||
#define TVM_RUNTIME_DISCO_SESSION_H_
|
||||
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <mutex>
|
||||
#include <queue>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
/*!
|
||||
* \brief Static FFI type index for `runtime::disco::DRef`.
|
||||
*
|
||||
* Allocated within the [kTVMFFIDynObjectBegin - 16, kTVMFFIDynObjectBegin)
|
||||
* custom-static slot range. The sibling constant `kRuntimeRPCObjectRef`
|
||||
* lives in `src/runtime/rpc/rpc_session.h` and uses `... - 13`; values must
|
||||
* remain disjoint across this small reserved block.
|
||||
*/
|
||||
constexpr int32_t kRuntimeDiscoDRef = TVMFFITypeIndex::kTVMFFIDynObjectBegin - 14;
|
||||
|
||||
static_assert(kRuntimeDiscoDRef >= TVMFFITypeIndex::kTVMFFIStaticObjectEnd &&
|
||||
kRuntimeDiscoDRef < TVMFFITypeIndex::kTVMFFIDynObjectBegin,
|
||||
"kRuntimeDiscoDRef must live in the static custom-index slot range");
|
||||
|
||||
/*!
|
||||
* \brief All possible kinds of Disco commands.
|
||||
*/
|
||||
enum class DiscoAction : int32_t {
|
||||
kShutDown = 0,
|
||||
kKillReg = 1,
|
||||
kGetGlobalFunc = 2,
|
||||
kCallPacked = 3,
|
||||
kSyncWorker = 4,
|
||||
kCopyFromWorker0 = 5,
|
||||
kCopyToWorker0 = 6,
|
||||
kDebugGetFromRemote = 7,
|
||||
kDebugSetRegister = 8,
|
||||
};
|
||||
|
||||
/*! \brief Converts the enum class `DiscoAction` to string */
|
||||
inline std::string DiscoAction2String(DiscoAction action) {
|
||||
switch (action) {
|
||||
case DiscoAction::kShutDown:
|
||||
return "kShutDown";
|
||||
case DiscoAction::kKillReg:
|
||||
return "kKillReg";
|
||||
case DiscoAction::kGetGlobalFunc:
|
||||
return "kGetGlobalFunc";
|
||||
case DiscoAction::kCallPacked:
|
||||
return "kCallPacked";
|
||||
case DiscoAction::kSyncWorker:
|
||||
return "kSyncWorker";
|
||||
case DiscoAction::kCopyFromWorker0:
|
||||
return "kCopyFromWorker0";
|
||||
case DiscoAction::kCopyToWorker0:
|
||||
return "kCopyToWorker0";
|
||||
case DiscoAction::kDebugGetFromRemote:
|
||||
return "kDebugGetFromRemote";
|
||||
case DiscoAction::kDebugSetRegister:
|
||||
return "kDebugSetRegister";
|
||||
}
|
||||
TVM_FFI_THROW(ValueError) << "Unknown DiscoAction: " << static_cast<int>(action);
|
||||
}
|
||||
|
||||
class SessionObj;
|
||||
|
||||
/*!
|
||||
* \brief An object that exists on all workers.
|
||||
*
|
||||
* The controler assigns a unique "register id" to each object, and the worker uses this id to
|
||||
* refer to the object residing on itself.
|
||||
*/
|
||||
class DRefObj : public ffi::Object {
|
||||
public:
|
||||
/*!\ brief Send dellocation command for `reg_id` */
|
||||
inline ~DRefObj();
|
||||
/*!
|
||||
* \brief Get the value of a DRef from a remote worker.
|
||||
* \param worker_id The id of the worker to be fetched from.
|
||||
* \return The value of the register.
|
||||
*/
|
||||
inline ffi::Any DebugGetFromRemote(int worker_id);
|
||||
/*!
|
||||
* \brief Copy from the Tensor provided to a remote worker.
|
||||
* \param worker_id The id of the worker to be copied to.
|
||||
* \param source The Tensor to be copied.
|
||||
*/
|
||||
inline void DebugCopyFrom(int worker_id, ffi::AnyView source);
|
||||
|
||||
static constexpr const uint32_t _type_index = kRuntimeDiscoDRef;
|
||||
static const constexpr bool _type_final = true;
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_STATIC("runtime.disco.DRef", DRefObj, ffi::Object);
|
||||
|
||||
/*! \brief The id of the register */
|
||||
int64_t reg_id;
|
||||
/*! \brief Back-pointer to the host controler session */
|
||||
ffi::ObjectRef session{nullptr};
|
||||
|
||||
private:
|
||||
inline SessionObj* GetSession();
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to DRefObj.
|
||||
* \sa DRefObj
|
||||
* \note No public constructor is provided as it is not supposed to be directly created by users.
|
||||
*/
|
||||
class DRef : public ffi::ObjectRef {
|
||||
public:
|
||||
explicit DRef(ffi::ObjectPtr<DRefObj> data) : ffi::ObjectRef(data) {
|
||||
TVM_FFI_ICHECK(data != nullptr);
|
||||
}
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(DRef, ffi::ObjectRef, DRefObj);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief A Disco interactive session. It allows users to interact with the Disco command queue with
|
||||
* various ffi::Function calling convention.
|
||||
*/
|
||||
class SessionObj : public ffi::Object {
|
||||
public:
|
||||
virtual ~SessionObj() = default;
|
||||
/*!
|
||||
* \brief Call a ffi::Function on workers providing variadic arguments.
|
||||
* \tparam Args In the variadic arguments, the supported types include:
|
||||
* - integers and floating point numbers;
|
||||
* - DataType;
|
||||
* - Device;
|
||||
* - std::string;
|
||||
* - DRef.
|
||||
* Examples of unsupported types:
|
||||
* - Tensor, DLTensor;
|
||||
* - TVM Objects, including ffi::Function, Module and String;
|
||||
* \param func The function to be called.
|
||||
* \param args The variadic arguments.
|
||||
* \return The return value of function call
|
||||
*/
|
||||
template <typename... Args>
|
||||
TVM_FFI_INLINE DRef CallPacked(const DRef& func, Args&&... args);
|
||||
/*!
|
||||
* \brief Call packed function on each worker using a packed sequence. The calling convention:
|
||||
* The first element must be DiscoAction::kCallPacked,
|
||||
* The second element must be 0, which will later be updated by the session to return reg_id
|
||||
* The thirtd element is the function to be called.
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual DRef CallWithPacked(const ffi::PackedArgs& args) = 0;
|
||||
/*! \brief Get the number of workers in the session. */
|
||||
TVM_RUNTIME_DLL virtual int64_t GetNumWorkers() = 0;
|
||||
/*! \brief Get a global functions on workers. */
|
||||
TVM_RUNTIME_DLL virtual DRef GetGlobalFunc(const std::string& name) = 0;
|
||||
/*!
|
||||
* \brief Copy an Tensor from worker-0 to the controler-side Tensor
|
||||
* \param host_array The array to be copied to worker-0
|
||||
* \param remote_array The Tensor on worker-0
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual void CopyFromWorker0(const Tensor& host_array,
|
||||
const DRef& remote_array) = 0;
|
||||
/*!
|
||||
* \brief Copy the controler-side Tensor to worker-0
|
||||
* \param host_array The array to be copied to worker-0
|
||||
* \param remote_array The Tensor on worker-0
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual void CopyToWorker0(const Tensor& host_array,
|
||||
const DRef& remote_array) = 0;
|
||||
/*!
|
||||
* \brief Synchrnoize the controler with a worker, and it will wait until worker finishes
|
||||
* executing this instruction.
|
||||
* \param worker_id The id of the worker to be synced with.
|
||||
* \note This function is usually used for worker-0, because it is the only worker that is
|
||||
* assumed to collocate with the controler. Syncing with other workers may not be supported.
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual void SyncWorker(int worker_id) = 0;
|
||||
/*! \brief Signal all the workers to shutdown */
|
||||
TVM_RUNTIME_DLL virtual void Shutdown() = 0;
|
||||
/*!
|
||||
* \brief Initialize the data plane between workers.
|
||||
* \param ccl The name of the communication backend, e.g., nccl, rccl, mpi.
|
||||
* \param device_ids The device ids of the workers.
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual void InitCCL(ffi::String ccl, ffi::Shape device_ids) = 0;
|
||||
/*!
|
||||
* \brief Get the value of a register from a remote worker.
|
||||
* \param reg_id The id of the register to be fetched.
|
||||
* \param worker_id The id of the worker to be fetched from.
|
||||
* \return The value of the register.
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual ffi::Any DebugGetFromRemote(int64_t reg_id, int worker_id) = 0;
|
||||
/*!
|
||||
* \brief Set the value of a register on a remote worker.
|
||||
* \param reg_id The id of the register to be set.
|
||||
* \param value The value to be set.
|
||||
* \param worker_id The id of the worker to be set.
|
||||
*/
|
||||
TVM_RUNTIME_DLL virtual void DebugSetRegister(int64_t reg_id, ffi::AnyView value,
|
||||
int worker_id) = 0;
|
||||
|
||||
struct FFI;
|
||||
friend struct SessionObj::FFI;
|
||||
friend class DRefObj;
|
||||
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("runtime.disco.Session", SessionObj, ffi::Object);
|
||||
|
||||
protected:
|
||||
/*! \brief Deallocate a register id, kill it on all workers, and append it to `free_regs_`. */
|
||||
virtual void DeallocReg(int reg_id) = 0;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Managed reference to SessionObj
|
||||
* \sa SessionObj
|
||||
*/
|
||||
class Session : public ffi::ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Create a session backed by a thread pool of workers
|
||||
* \param num_workers The number of workers.
|
||||
* \param num_groups The number of worker groups.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static Session ThreadedSession(int num_workers, int num_groups);
|
||||
/*!
|
||||
* \brief Create a session backed by pipe-based multiprocessing
|
||||
* \param num_workers The number of workers.
|
||||
* \param num_groups The number of worker groups.
|
||||
* \param process_pool_creator The name of a global function that takes `num_workers` as an input,
|
||||
* and returns a ffi::Function, which takes an integer `worker_id` as the input and returns None.
|
||||
* When `worker-id` is 0, it shuts down the process pool; Otherwise, it retursn a tuple
|
||||
* (read_fd, writefd) used to communicate with the corresponding worker.
|
||||
* \param entrypoint The entrypoint of DiscoWorker main worker function.
|
||||
* \note Worker-0 is always co-located with the controler as a separate thread, and therefore
|
||||
* worker-0 does not exist in the process pool.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static Session ProcessSession(int num_workers, int num_groups,
|
||||
ffi::String process_pool_creator,
|
||||
ffi::String entrypoint);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Session, ffi::ObjectRef, SessionObj);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief A bi-directional channel for controler-worker communication.
|
||||
* This channel is primarily used to transfer control messages but not data.
|
||||
*/
|
||||
class DiscoChannel {
|
||||
public:
|
||||
virtual ~DiscoChannel() = default;
|
||||
/*! \brief Send a packed sequence to the receiver */
|
||||
virtual void Send(const ffi::PackedArgs& args) = 0;
|
||||
/*! \brief Receive a packed sequence from worker */
|
||||
virtual ffi::PackedArgs Recv() = 0;
|
||||
/*! \brief Reply a packed sequence to the sender */
|
||||
virtual void Reply(const ffi::PackedArgs& args) = 0;
|
||||
/*! \brief Receive a reply from the worker */
|
||||
virtual ffi::PackedArgs RecvReply() = 0;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief A special communication channel between controler and worker-0,
|
||||
* assuming they are always collocated in the same process.
|
||||
*/
|
||||
class WorkerZeroData {
|
||||
public:
|
||||
/*!
|
||||
* \brief The host-side arrays to passed to worker-0 for special uses, for example,
|
||||
* copy-to-worker0 and copy-from-worker0
|
||||
*/
|
||||
std::queue<Tensor> host_arrays;
|
||||
/*! \brief The mutex that guards `host_arrays` */
|
||||
std::mutex queue_mutex_;
|
||||
};
|
||||
|
||||
// Implementation details
|
||||
|
||||
inline SessionObj* DRefObj::GetSession() {
|
||||
return const_cast<SessionObj*>(static_cast<const SessionObj*>(session.get()));
|
||||
}
|
||||
|
||||
DRefObj::~DRefObj() {
|
||||
if (this->session.defined()) {
|
||||
GetSession()->DeallocReg(reg_id);
|
||||
}
|
||||
}
|
||||
|
||||
ffi::Any DRefObj::DebugGetFromRemote(int worker_id) {
|
||||
return GetSession()->DebugGetFromRemote(this->reg_id, worker_id);
|
||||
}
|
||||
|
||||
void DRefObj::DebugCopyFrom(int worker_id, ffi::AnyView value) {
|
||||
return GetSession()->DebugSetRegister(this->reg_id, value, worker_id);
|
||||
}
|
||||
|
||||
template <typename... Args>
|
||||
DRef SessionObj::CallPacked(const DRef& func, Args&&... args) {
|
||||
constexpr int offset = 3;
|
||||
constexpr int kNumArgs = offset + sizeof...(Args);
|
||||
ffi::AnyView packed_args[kNumArgs];
|
||||
ffi::PackedArgs::Fill(packed_args,
|
||||
/*.0=*/static_cast<int>(DiscoAction::kCallPacked), // action
|
||||
/*.1=*/0, // reg_id, which will be updated by this->CallWithPacked
|
||||
/*.2=*/func, // the function to be called
|
||||
std::forward<Args>(args)...);
|
||||
return this->CallWithPacked(ffi::PackedArgs(packed_args, kNumArgs));
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
#endif // TVM_RUNTIME_DISCO_SESSION_H_
|
||||
@@ -0,0 +1,402 @@
|
||||
/*
|
||||
* 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/logging.h
|
||||
* \brief logging utilities
|
||||
*
|
||||
* We use the following facilities from tvm/ffi/error.h for
|
||||
* error handling and checking:
|
||||
*
|
||||
* - TVM_FFI_THROW(ErrorKind) << "msg";
|
||||
* - TVM_FFI_CHECK(cond, ErrorKind) << "msg";
|
||||
* - TVM_FFI_CHECK_EQ(x, y, ErrorKind) << "msg";
|
||||
* - TVM_FFI_ICHECK(x) << "msg"; // InternalError
|
||||
* - TVM_FFI_ICHECK_EQ(x, y) << "msg";
|
||||
* - TVM_FFI_DCHECK(x) << "msg"; // Debug-only InternalError
|
||||
*
|
||||
* LOG(INFO), LOG(WARNING), LOG(ERROR) are kept for logging.
|
||||
* LOG(FATAL) is kept for completeness, it throws InternalError.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_LOGGING_H_
|
||||
#define TVM_RUNTIME_LOGGING_H_
|
||||
|
||||
#include <tvm/ffi/error.h>
|
||||
#include <tvm/runtime/base.h>
|
||||
|
||||
#include <ctime>
|
||||
#include <iomanip>
|
||||
#include <iostream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
/*!
|
||||
* \brief Whether or not customize the logging output.
|
||||
* If log customize is enabled, the user must implement
|
||||
* tvm::runtime::detail::LogFatalImpl and tvm::runtime::detail::LogMessageImpl.
|
||||
*/
|
||||
#ifndef TVM_LOG_CUSTOMIZE
|
||||
#define TVM_LOG_CUSTOMIZE 0
|
||||
#endif
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
/*! \brief Internal implementation */
|
||||
namespace detail {
|
||||
// Provide support for customized logging.
|
||||
#if TVM_LOG_CUSTOMIZE
|
||||
/*!
|
||||
* \brief Custom implementations of LogFatal.
|
||||
*
|
||||
* \sa TVM_LOG_CUSTOMIZE
|
||||
*/
|
||||
[[noreturn]] TVM_RUNTIME_DLL void LogFatalImpl(const std::string& file, int lineno,
|
||||
const std::string& message);
|
||||
|
||||
/*!
|
||||
* \brief Custom implementations of LogMessage.
|
||||
*
|
||||
* \sa TVM_LOG_CUSTOMIZE
|
||||
*/
|
||||
TVM_RUNTIME_DLL void LogMessageImpl(const std::string& file, int lineno, int level,
|
||||
const std::string& message);
|
||||
|
||||
/*!
|
||||
* \brief Class to accumulate an error message and throw it. Do not use
|
||||
* directly, instead use LOG(FATAL).
|
||||
*/
|
||||
class LogFatal {
|
||||
public:
|
||||
LogFatal(const std::string& file, int lineno) : file_(file), lineno_(lineno) {}
|
||||
#ifdef _MSC_VER
|
||||
#pragma warning(push)
|
||||
#pragma warning(disable : 4722)
|
||||
#endif
|
||||
[[noreturn]] ~LogFatal() noexcept(false) { LogFatalImpl(file_, lineno_, stream_.str()); }
|
||||
#ifdef _MSC_VER
|
||||
#pragma warning(pop)
|
||||
#endif
|
||||
std::ostringstream& stream() { return stream_; }
|
||||
|
||||
private:
|
||||
std::ostringstream stream_;
|
||||
std::string file_;
|
||||
int lineno_;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Class to accumulate an log message. Do not use directly, instead use
|
||||
* LOG(INFO), LOG(WARNING), LOG(ERROR).
|
||||
*/
|
||||
class LogMessage {
|
||||
public:
|
||||
LogMessage(const std::string& file, int lineno, int level)
|
||||
: file_(file), lineno_(lineno), level_(level) {}
|
||||
~LogMessage() { LogMessageImpl(file_, lineno_, level_, stream_.str()); }
|
||||
std::ostringstream& stream() { return stream_; }
|
||||
|
||||
private:
|
||||
std::string file_;
|
||||
int lineno_;
|
||||
int level_;
|
||||
std::ostringstream stream_;
|
||||
};
|
||||
|
||||
#else
|
||||
|
||||
/*!
|
||||
* \brief Class to accumulate an error message and throw it. Do not use
|
||||
* directly, instead use LOG(FATAL).
|
||||
* \note The `LogFatal` class is designed to be an empty class to reduce stack size usage.
|
||||
* To play this trick, we use the thread-local storage to store its internal data.
|
||||
*/
|
||||
class LogFatal {
|
||||
public:
|
||||
TVM_FFI_NO_INLINE LogFatal(const char* file, int lineno) { GetEntry().Init(file, lineno); }
|
||||
#ifdef _MSC_VER
|
||||
#pragma warning(push)
|
||||
#pragma warning(disable : 4722)
|
||||
#endif
|
||||
[[noreturn]] ~LogFatal() noexcept(false) {
|
||||
GetEntry().Finalize();
|
||||
throw;
|
||||
}
|
||||
#ifdef _MSC_VER
|
||||
#pragma warning(pop)
|
||||
#endif
|
||||
std::ostringstream& stream() { return GetEntry().stream_; }
|
||||
|
||||
private:
|
||||
struct Entry {
|
||||
void Init(const char* file, int lineno) {
|
||||
this->stream_.str("");
|
||||
this->file_ = file;
|
||||
this->lineno_ = lineno;
|
||||
}
|
||||
[[noreturn]] TVM_FFI_NO_INLINE ffi::Error Finalize() noexcept(false) {
|
||||
ffi::Error error("InternalError", stream_.str(),
|
||||
TVMFFIBacktrace(file_.c_str(), lineno_, "", 0));
|
||||
throw error;
|
||||
}
|
||||
std::ostringstream stream_;
|
||||
std::string file_;
|
||||
int lineno_;
|
||||
};
|
||||
|
||||
TVM_FFI_NO_INLINE TVM_RUNTIME_DLL static Entry& GetEntry();
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Class to accumulate an log message. Do not use directly, instead use
|
||||
* LOG(INFO), LOG(WARNING), LOG(ERROR).
|
||||
*/
|
||||
class LogMessage {
|
||||
public:
|
||||
LogMessage(const std::string& file, int lineno, int level) {
|
||||
// Use inline constexpr to avoid ODR-issues with hidden static members
|
||||
// when libtvm_runtime.so is built with -fvisibility=hidden.
|
||||
static constexpr const char* kLevelStrings[] = {
|
||||
": Debug: ", // TVM_LOG_LEVEL_DEBUG
|
||||
": ", // TVM_LOG_LEVEL_INFO
|
||||
": Warning: ", // TVM_LOG_LEVEL_WARNING
|
||||
": Error: ", // TVM_LOG_LEVEL_ERROR
|
||||
};
|
||||
std::time_t t = std::time(nullptr);
|
||||
stream_ << "[" << std::put_time(std::localtime(&t), "%H:%M:%S") << "] " << file << ":" << lineno
|
||||
<< kLevelStrings[level];
|
||||
}
|
||||
TVM_FFI_NO_INLINE ~LogMessage() { std::cerr << stream_.str() << std::endl; }
|
||||
std::ostringstream& stream() { return stream_; }
|
||||
|
||||
private:
|
||||
std::ostringstream stream_;
|
||||
};
|
||||
|
||||
#endif
|
||||
|
||||
// This class is used to explicitly ignore values in the conditional
|
||||
// logging macros. This avoids compiler warnings like "value computed
|
||||
// is not used" and "statement has no effect".
|
||||
class LogMessageVoidify {
|
||||
public:
|
||||
LogMessageVoidify() {}
|
||||
// This has to be an operator with a precedence lower than << but
|
||||
// higher than "?:". See its usage.
|
||||
void operator&(std::ostream&) {}
|
||||
};
|
||||
|
||||
/*! \brief Captures the state of the \p TVM_LOG_DEBUG environment flag. */
|
||||
class TVM_RUNTIME_DLL TvmLogDebugSettings {
|
||||
public:
|
||||
/*!
|
||||
* \brief Parses the \p TVM_LOG_DEBUG environment flag as per the specification given by
|
||||
* \p DebugLoggingEnabled and \p VerboseLoggingEnabled, and caches the result.
|
||||
*/
|
||||
inline static const TvmLogDebugSettings& FromFlag() {
|
||||
// Parse and cache the verbosity level map.
|
||||
static const auto* settings =
|
||||
new TvmLogDebugSettings(TvmLogDebugSettings::ParseSpec(std::getenv("TVM_LOG_DEBUG")));
|
||||
return *settings;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Parses \p opt_spec as per specification for \p TVM_LOG_DEBUG given by
|
||||
* \p DebugLoggingEnabled and \p VerboseLoggingEnabled. Throws if specification is ill-formed.
|
||||
*/
|
||||
static TvmLogDebugSettings ParseSpec(const char* opt_spec);
|
||||
|
||||
/*!
|
||||
* \brief Implements \p VerboseLoggingEnabled below w.r.t. the already parsed \p TVM_LOG_DEBUG
|
||||
* environment variable.
|
||||
*/
|
||||
inline bool VerboseEnabled(const char* opt_filename, int level) const {
|
||||
if (opt_filename == nullptr || level < 0 || vlog_level_map_.empty()) {
|
||||
return false;
|
||||
}
|
||||
return VerboseEnabledImpl(opt_filename, level);
|
||||
}
|
||||
|
||||
/*! \brief Returns true if \p DLOG statements should be executed. */
|
||||
bool dlog_enabled() const { return dlog_enabled_; }
|
||||
|
||||
private:
|
||||
// Slow path for VerboseEnabled.
|
||||
bool VerboseEnabledImpl(const std::string& filename, int level) const;
|
||||
|
||||
/*! \brief If true, DLOG statements are enabled. */
|
||||
bool dlog_enabled_ = false;
|
||||
/*!
|
||||
* \brief A map from canonicalized filenames to the maximum VLOG verbosity level for that file.
|
||||
* May also contain the 'wildcard' entry \p "DEFAULT" representing the level for all other files.
|
||||
*/
|
||||
std::unordered_map<std::string, int> vlog_level_map_;
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Returns true if a DLOG statement is enabled by the \p TVM_LOG_DEBUG environment
|
||||
* variable. Requires:
|
||||
* \code
|
||||
* TVM_LOG_DEBUG=1
|
||||
* \endcode
|
||||
* or a valid setting as described by \p VerboseLoggingEnabled below.
|
||||
*/
|
||||
inline bool DebugLoggingEnabled() {
|
||||
static int state = 0;
|
||||
if (state == 0) {
|
||||
state = TvmLogDebugSettings::FromFlag().dlog_enabled() ? 1 : -1;
|
||||
}
|
||||
return state == 1;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Returns true if a VLOG statement in \p filename is enabled by the \p TVM_LOG_DEBUG
|
||||
* environment variable for logging at verbosity \p level. Levels should be non-negative.
|
||||
*
|
||||
* Filenames are canonicalized to be w.r.t. the src/ dir of the TVM tree. (VLOG's should not
|
||||
* appear under include/).
|
||||
*
|
||||
* To enable file \p ir/bar.cc for level 0 only set:
|
||||
* \code
|
||||
* TVM_LOG_DEBUG="ir/bar.cc=0"
|
||||
* \endcode
|
||||
*
|
||||
* To enable all files up to level 3 but disable \p ir/bar.cc set:
|
||||
* \code
|
||||
* TVM_LOG_DEBUG="DEFAULT=2,ir/bar.cc=-1"
|
||||
* \endcode
|
||||
*
|
||||
* Any of these settings will also enable DLOG statements.
|
||||
*/
|
||||
inline bool VerboseLoggingEnabled(const char* opt_filename, int level) {
|
||||
return TvmLogDebugSettings::FromFlag().VerboseEnabled(opt_filename, level);
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief A stack of VLOG context messages.
|
||||
*
|
||||
* For use by \p VLOG_CONTEXT macro only.
|
||||
*/
|
||||
class VLogContext {
|
||||
public:
|
||||
void Push(std::stringstream* stream) { context_stack_.push_back(stream); }
|
||||
void Pop() {
|
||||
if (!context_stack_.empty()) {
|
||||
context_stack_.pop_back();
|
||||
}
|
||||
}
|
||||
|
||||
std::string str() const;
|
||||
|
||||
private:
|
||||
std::vector<std::stringstream*> context_stack_;
|
||||
};
|
||||
|
||||
/*! \brief Get thread local \p VLogContext for tracking a stack of VLOG context messages. */
|
||||
inline VLogContext* ThreadLocalVLogContext() {
|
||||
static thread_local VLogContext inst;
|
||||
return &inst;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief A RAII class to push/pos a VLOG context message onto the thread-local stack.
|
||||
*
|
||||
* For use by \p VLOG_CONTEXT macro only.
|
||||
*/
|
||||
class VLogContextEntry {
|
||||
public:
|
||||
VLogContextEntry() { ThreadLocalVLogContext()->Push(&sstream_); }
|
||||
~VLogContextEntry() { ThreadLocalVLogContext()->Pop(); }
|
||||
std::ostream& stream() { return sstream_; }
|
||||
|
||||
private:
|
||||
std::stringstream sstream_;
|
||||
};
|
||||
|
||||
} // namespace detail
|
||||
|
||||
#define TVM_LOG_LEVEL_DEBUG 0
|
||||
#define TVM_LOG_LEVEL_INFO 1
|
||||
#define TVM_LOG_LEVEL_WARNING 2
|
||||
#define TVM_LOG_LEVEL_ERROR 3
|
||||
#define TVM_LOG_LEVEL_FATAL 4
|
||||
#define LOG(level) LOG_##level
|
||||
#define LOG_DEBUG \
|
||||
::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_DEBUG).stream()
|
||||
#define LOG_FATAL ::tvm::runtime::detail::LogFatal(__FILE__, __LINE__).stream()
|
||||
#define LOG_INFO ::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_INFO).stream()
|
||||
#define LOG_ERROR \
|
||||
::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_ERROR).stream()
|
||||
#define LOG_WARNING \
|
||||
::tvm::runtime::detail::LogMessage(__FILE__, __LINE__, TVM_LOG_LEVEL_WARNING).stream()
|
||||
|
||||
#define LOG_IF(severity, condition) \
|
||||
!(condition) ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(severity)
|
||||
|
||||
#if TVM_LOG_DEBUG
|
||||
|
||||
#define LOG_DFATAL LOG_FATAL
|
||||
#define DFATAL FATAL
|
||||
#define DLOG(severity) LOG_IF(severity, ::tvm::runtime::detail::DebugLoggingEnabled())
|
||||
#define DLOG_IF(severity, condition) \
|
||||
LOG_IF(severity, ::tvm::runtime::detail::DebugLoggingEnabled() && (condition))
|
||||
|
||||
/*!
|
||||
* \brief If the \p TVM_LOG_DEBUG build flag is enabled, push a context message onto an internal
|
||||
* stack. All VLOG messages will include this stack in their prefix to help with debugging. E.g.:
|
||||
* \code
|
||||
* VLOG_CONTEXT << "my context";
|
||||
* VLOG(1) << "my log message";
|
||||
* \endcode
|
||||
* Thread safe. No-op with no execution overhead if the \p TVM_LOG_DEBUG build flag is not enabled.
|
||||
*/
|
||||
#define VLOG_CONTEXT \
|
||||
::tvm::runtime::detail::VLogContextEntry vlog_entry_; \
|
||||
vlog_entry_.stream()
|
||||
|
||||
#else
|
||||
|
||||
#define LOG_DFATAL LOG_ERROR
|
||||
#define DFATAL ERROR
|
||||
#define DLOG(severity) true ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(severity)
|
||||
#define DLOG_IF(severity, condition) \
|
||||
(true || !(condition)) ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(severity)
|
||||
#define VLOG_CONTEXT true ? (void)0 : ::tvm::runtime::detail::LogMessageVoidify() & LOG(INFO)
|
||||
|
||||
#endif
|
||||
|
||||
/*!
|
||||
* \brief If the \p TVM_LOG_DEBUG build flag is enabled, and the containing file has been enabled
|
||||
* at \p level or greater in the \p TVM_LOG_DEBUG environment variable, then log a message at
|
||||
* \p INFO severity.
|
||||
*
|
||||
* See \p VerboseLoggingEnabled for the format of the \p TVM_LOG_DEBUG environment variable.
|
||||
* Thread safe. No-op with no execution overhead if the \p TVM_LOG_DEBUG build flag is not enabled.
|
||||
* No-op with some execution overhead if the \p TVM_LOG_DEBUG build flag is enabled but the
|
||||
* containing file is not enabled.
|
||||
*/
|
||||
#define VLOG(level) \
|
||||
DLOG_IF(INFO, ::tvm::runtime::detail::VerboseLoggingEnabled(__FILE__, (level))) \
|
||||
<< ::tvm::runtime::detail::ThreadLocalVLogContext()->str()
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
#endif // TVM_RUNTIME_LOGGING_H_
|
||||
@@ -0,0 +1,199 @@
|
||||
/*
|
||||
* 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/memory/memory_manager.h
|
||||
* \brief Abstract device memory management API
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
|
||||
#define TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
|
||||
|
||||
#include <tvm/runtime/base.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace memory {
|
||||
|
||||
enum AllocatorType {
|
||||
kNaive = 1,
|
||||
kPooled,
|
||||
};
|
||||
|
||||
struct Buffer {
|
||||
/*! \brief The pointer to the allocated block of memory. */
|
||||
void* data{nullptr};
|
||||
/*! \brief The size of the block. */
|
||||
size_t size{0};
|
||||
/*! \brief The context of the allocated buffers. */
|
||||
Device device;
|
||||
/*! \brief The allocator that created this buffer. */
|
||||
AllocatorType alloc_type;
|
||||
};
|
||||
|
||||
class TVM_RUNTIME_DLL Allocator {
|
||||
public:
|
||||
explicit Allocator(AllocatorType type) : type_(type) {}
|
||||
virtual ~Allocator() = default;
|
||||
/*! \brief Allocate an empty Tensor using from the allocator.
|
||||
* \param shape The shape of the Tensor.
|
||||
* \param dtype The datatype of the Tensor.
|
||||
* \param dev The device where the array is allocated.
|
||||
* \param mem_scope The device memory scope hint.
|
||||
* \return The empty Tensor.
|
||||
*/
|
||||
Tensor Empty(ffi::Shape shape, DLDataType dtype, Device dev,
|
||||
ffi::Optional<ffi::String> mem_scope = std::nullopt);
|
||||
/*! \brief Return the allocator type. */
|
||||
inline AllocatorType type() const { return type_; }
|
||||
/*! \brief Allocate a buffer given a size, alignment and type.
|
||||
* \param dev The device where the array is allocated.
|
||||
* \param nbytes The size of the buffer.
|
||||
* \param alignment The alignment of the buffer.
|
||||
* \param type_hint A type hint to the allocator.
|
||||
* \return A sized allocation in the form of a buffer.
|
||||
*/
|
||||
virtual Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) = 0;
|
||||
/*! \brief Allocate a buffer given a shape and type.
|
||||
* \param dev The device where the array is allocated.
|
||||
* \param shape The shape of the tensor.
|
||||
* \param type_hint A type hint to the allocator.
|
||||
* \param mem_scope A memory scope of the buffer.
|
||||
* \return A sized allocation in the form of a buffer.
|
||||
*/
|
||||
virtual Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
|
||||
const std::string& mem_scope = "");
|
||||
|
||||
/*! \brief Create a view for the buffer given a shape, type and scope.
|
||||
* \param buffer The existing buffer upon which we need to create a view.
|
||||
* \param shape The shape of the view.
|
||||
* \param type_hint A type hint to the view.
|
||||
* \param mem_scope A memory scope of the view.
|
||||
* \return A device pointer to the created view.
|
||||
*/
|
||||
virtual void* CreateView(const Buffer& buffer, ffi::Shape shape, DLDataType type_hint,
|
||||
const std::string& mem_scope = "global") {
|
||||
return buffer.data;
|
||||
}
|
||||
|
||||
/*! \brief Release the view .
|
||||
* \param dev is the device where this view is created
|
||||
* \param data The view pointer to be freed.
|
||||
*/
|
||||
virtual void FreeView(Device dev, void* data) {}
|
||||
|
||||
/*! \brief Free a buffer allocated by the allocator.
|
||||
* \param buffer The buffer to free.
|
||||
*/
|
||||
virtual void Free(const Buffer& buffer) = 0;
|
||||
/*! \brief Clear the allocated memory. */
|
||||
virtual void Clear();
|
||||
/*! \brief The amount of memory currently allocated.
|
||||
* \return The amount of memory currently allocated.
|
||||
*/
|
||||
virtual size_t UsedMemory() const = 0;
|
||||
|
||||
protected:
|
||||
/*! \brief Check if the given memory scope is allowed to allocate by the allocator. */
|
||||
virtual bool AllowMemoryScope(const std::string& mem_scope) const;
|
||||
|
||||
private:
|
||||
AllocatorType type_;
|
||||
};
|
||||
|
||||
class MemoryManager {
|
||||
public:
|
||||
TVM_RUNTIME_DLL static MemoryManager* Global();
|
||||
/*!
|
||||
* \brief Get or create an allocator given the context and allocator type.
|
||||
* \param dev The TVM device
|
||||
* \param type The allocator type
|
||||
* \return The memory allocator.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static Allocator* GetOrCreateAllocator(Device dev, AllocatorType type);
|
||||
/*!
|
||||
* \brief Get an allocator given the context.
|
||||
* \param dev The TVM device
|
||||
* \param type The allocator type
|
||||
* \return The memory allocator.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static Allocator* GetAllocator(Device dev, AllocatorType type);
|
||||
/*! \brief Clear the allocators. */
|
||||
static void Clear();
|
||||
|
||||
private:
|
||||
MemoryManager() {}
|
||||
|
||||
protected:
|
||||
std::mutex mu_;
|
||||
std::unordered_map<Device, std::unordered_map<AllocatorType, std::unique_ptr<Allocator>>>
|
||||
allocators_;
|
||||
};
|
||||
|
||||
/*! \brief An object representing a storage allocation. */
|
||||
class StorageObj : public ffi::Object {
|
||||
public:
|
||||
/*! \brief The index into the VM function table. */
|
||||
Buffer buffer;
|
||||
/*! \brief The allocator where the storage buffer is allocated from. */
|
||||
Allocator* allocator = nullptr;
|
||||
|
||||
/*! \brief Allocate an Tensor from a given piece of storage. */
|
||||
TVM_RUNTIME_DLL Tensor AllocTensor(int64_t offset, ffi::Shape shape, DLDataType dtype);
|
||||
|
||||
/*! \brief Allocate an Tensor with memory scope from a given piece of storage. */
|
||||
TVM_RUNTIME_DLL Tensor AllocTensorScoped(int64_t offset, ffi::Shape shape, DLDataType dtype,
|
||||
ffi::String scope = "global");
|
||||
|
||||
~StorageObj() {
|
||||
if (allocator) {
|
||||
allocator->Free(buffer);
|
||||
}
|
||||
}
|
||||
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("vm.Storage", StorageObj, ffi::Object);
|
||||
};
|
||||
|
||||
/*! \brief reference to storage. */
|
||||
class Storage : public ffi::ObjectRef {
|
||||
public:
|
||||
TVM_RUNTIME_DLL explicit Storage(Buffer buffer, Allocator* allocator);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Storage, ffi::ObjectRef, StorageObj);
|
||||
};
|
||||
|
||||
} // namespace memory
|
||||
|
||||
using memory::Allocator;
|
||||
using memory::AllocatorType;
|
||||
using memory::MemoryManager;
|
||||
using memory::StorageObj;
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
|
||||
@@ -0,0 +1,360 @@
|
||||
/*
|
||||
* 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/tensor.h
|
||||
* \brief A device-independent managed Tensor abstraction.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_TENSOR_H_
|
||||
#define TVM_RUNTIME_TENSOR_H_
|
||||
|
||||
#include <tvm/ffi/container/shape.h>
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
#include <tvm/ffi/dtype.h>
|
||||
#include <tvm/ffi/optional.h>
|
||||
#include <tvm/ffi/string.h>
|
||||
#include <tvm/runtime/base.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
#include <tvm/support/io.h>
|
||||
#include <tvm/support/serializer.h>
|
||||
|
||||
#include <atomic>
|
||||
#include <functional>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
/*!
|
||||
* \brief Managed Tensor.
|
||||
* The array is backed by reference counted blocks.
|
||||
*/
|
||||
class Tensor : public tvm::ffi::Tensor {
|
||||
public:
|
||||
Tensor() = default;
|
||||
/*!
|
||||
* \brief constructor.
|
||||
* \param data ffi::ObjectPtr to the data container.
|
||||
*/
|
||||
explicit Tensor(ffi::ObjectPtr<ffi::TensorObj> data) : tvm::ffi::Tensor(data) {}
|
||||
explicit Tensor(ffi::UnsafeInit tag) : tvm::ffi::Tensor(tag) {}
|
||||
Tensor(ffi::Tensor&& other) : tvm::ffi::Tensor(std::move(other)) {} // NOLINT(*)
|
||||
Tensor(const ffi::Tensor& other) : tvm::ffi::Tensor(other) {} // NOLINT(*)
|
||||
|
||||
ffi::ShapeView Shape() const { return this->shape(); }
|
||||
DLDataType DataType() const { return this->dtype(); }
|
||||
|
||||
// DLPack handling
|
||||
static Tensor FromDLPack(DLManagedTensor* tensor) {
|
||||
return tvm::ffi::Tensor::FromDLPack(tensor, kAllocAlignment, true);
|
||||
}
|
||||
|
||||
static Tensor FromDLPackVersioned(DLManagedTensorVersioned* tensor) {
|
||||
return tvm::ffi::Tensor::FromDLPackVersioned(tensor, kAllocAlignment, true);
|
||||
}
|
||||
inline const DLTensor* operator->() const { return this->get(); }
|
||||
/*!
|
||||
* \brief Copy data content from another array.
|
||||
* \param other The source array to be copied from.
|
||||
* \note The copy may happen asynchronously if it involves a GPU context.
|
||||
* TVMSynchronize is necessary.
|
||||
*/
|
||||
inline void CopyFrom(const DLTensor* other);
|
||||
inline void CopyFrom(const Tensor& other);
|
||||
/*!
|
||||
* \brief Copy data content from a byte buffer.
|
||||
* \param data The source bytes to be copied from.
|
||||
* \param nbytes The size of the buffer in bytes
|
||||
* Must be equal to the size of the Tensor.
|
||||
* \note The copy always triggers a TVMSynchronize.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void CopyFromBytes(const void* data, size_t nbytes);
|
||||
/*!
|
||||
* \brief Copy data content into another array.
|
||||
* \param other The source array to be copied from.
|
||||
* \note The copy may happen asynchronously if it involves a GPU context.
|
||||
* TVMSynchronize is necessary.
|
||||
*/
|
||||
inline void CopyTo(DLTensor* other) const;
|
||||
inline void CopyTo(const Tensor& other) const;
|
||||
/*!
|
||||
* \brief Copy data content into another array.
|
||||
* \param data The source bytes to be copied from.
|
||||
* \param nbytes The size of the data buffer.
|
||||
* Must be equal to the size of the Tensor.
|
||||
* \note The copy always triggers a TVMSynchronize.
|
||||
*/
|
||||
TVM_RUNTIME_DLL void CopyToBytes(void* data, size_t nbytes) const;
|
||||
/*!
|
||||
* \brief Copy the data to another device.
|
||||
* \param dev The target device.
|
||||
* \param mem_scope The memory scope of the target array.
|
||||
* \return The array under another device.
|
||||
* \note The copy always triggers a TVMSynchronize.
|
||||
*/
|
||||
TVM_RUNTIME_DLL Tensor CopyTo(const Device& dev,
|
||||
ffi::Optional<ffi::String> mem_scope = std::nullopt) const;
|
||||
/*!
|
||||
* \brief Load Tensor from stream
|
||||
* \param stream The input data stream
|
||||
* \return Whether load is successful
|
||||
*/
|
||||
inline bool Load(support::Stream* stream);
|
||||
/*!
|
||||
* \brief Save Tensor to stream
|
||||
* \param stream The output data stream
|
||||
*/
|
||||
inline void Save(support::Stream* stream) const;
|
||||
|
||||
/*!
|
||||
* \brief Create a Tensor that shares the data memory with the current one.
|
||||
*
|
||||
* \param shape The shape of the new array.
|
||||
*
|
||||
* \param dtype The data type of the new array.
|
||||
*
|
||||
* \param relative_byte_offset The offset of the output Tensor,
|
||||
* relative to the current byte offset.
|
||||
*
|
||||
* By default, the offset of the view is the same as the offset
|
||||
* of the current array.
|
||||
*
|
||||
* \note The new array must not allow access of addresses which
|
||||
* would be out of bounds in the current array. If the new
|
||||
* array is larger than the current array, or if the
|
||||
* `relative_byte_offset` would place the end of the new array
|
||||
* outside the bounds of the current array, this function will
|
||||
* raise an exception.
|
||||
*/
|
||||
TVM_RUNTIME_DLL Tensor CreateView(ffi::Shape shape, DLDataType dtype,
|
||||
uint64_t relative_byte_offset = 0) const;
|
||||
/*!
|
||||
* \brief Create an empty Tensor.
|
||||
* \param shape The shape of the new array.
|
||||
* \param dtype The data type of the new array.
|
||||
* \param dev The device of the array.
|
||||
* \param mem_scope The memory scope of the array.
|
||||
* \return The created Array
|
||||
*/
|
||||
TVM_RUNTIME_DLL static Tensor Empty(ffi::Shape shape, DLDataType dtype, Device dev,
|
||||
ffi::Optional<ffi::String> mem_scope = std::nullopt);
|
||||
/*!
|
||||
* \brief Function to copy data from one array to another.
|
||||
* \param from The source array.
|
||||
* \param to The target array.
|
||||
* \param stream The stream used in copy.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static void CopyFromTo(const DLTensor* from, DLTensor* to,
|
||||
TVMStreamHandle stream = nullptr);
|
||||
|
||||
/*!
|
||||
* \brief Function to copy data from one array to a byte buffer.
|
||||
* \param from The source array.
|
||||
* \param to The target byte buffer.
|
||||
* \param nbytes The size of the data buffer.
|
||||
* \param stream The stream used in copy.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static void CopyToBytes(const DLTensor* from, void* to, size_t nbytes,
|
||||
TVMStreamHandle stream = nullptr);
|
||||
|
||||
/*!
|
||||
* \brief Function to copy data from one array to a byte buffer.
|
||||
* \param from The source array.
|
||||
* \param to The target byte buffer.
|
||||
* \param nbytes The size of the data buffer.
|
||||
* \param stream The stream used in copy.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static void CopyFromBytes(const DLTensor* to, void* from, size_t nbytes,
|
||||
TVMStreamHandle stream = nullptr);
|
||||
|
||||
/*!
|
||||
* \brief Check if two tensors share the same underlying storage.
|
||||
*
|
||||
* This detects runtime storage aliasing (e.g. views from CreateView, etc.) but does
|
||||
* not imply either tensor was created by CreateView.
|
||||
*
|
||||
* \param a The first tensor.
|
||||
* \param b The second tensor.
|
||||
* \return True if the tensors share the same storage.
|
||||
*/
|
||||
TVM_RUNTIME_DLL static bool IsStorageShared(const DLTensor* a, const DLTensor* b);
|
||||
|
||||
/*!
|
||||
* \brief Tensor overload of IsStorageShared.
|
||||
* \param a The first tensor.
|
||||
* \param b The second tensor.
|
||||
* \return True if the tensors share the same storage.
|
||||
*/
|
||||
static bool IsStorageShared(const Tensor& a, const Tensor& b);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Save a DLTensor to stream
|
||||
* \param strm The output stream
|
||||
* \param tensor The tensor to be saved.
|
||||
*/
|
||||
inline bool SaveDLTensor(support::Stream* strm, const DLTensor* tensor);
|
||||
|
||||
inline void Tensor::CopyFrom(const DLTensor* other) {
|
||||
TVM_FFI_ICHECK(data_ != nullptr);
|
||||
CopyFromTo(other, get_mutable());
|
||||
}
|
||||
|
||||
inline void Tensor::CopyFrom(const Tensor& other) {
|
||||
TVM_FFI_ICHECK(data_ != nullptr);
|
||||
TVM_FFI_ICHECK(other.data_ != nullptr);
|
||||
CopyFromTo(other.get_mutable(), get_mutable());
|
||||
}
|
||||
|
||||
inline void Tensor::CopyTo(DLTensor* other) const {
|
||||
TVM_FFI_ICHECK(data_ != nullptr);
|
||||
CopyFromTo(get_mutable(), other);
|
||||
}
|
||||
|
||||
inline void Tensor::CopyTo(const Tensor& other) const {
|
||||
TVM_FFI_ICHECK(data_ != nullptr);
|
||||
TVM_FFI_ICHECK(other.data_ != nullptr);
|
||||
CopyFromTo(get_mutable(), other.get_mutable());
|
||||
}
|
||||
|
||||
/*! \brief Magic number for Tensor file */
|
||||
constexpr uint64_t kTVMTensorMagic = 0xDD5E40F096B4A13F;
|
||||
|
||||
inline bool SaveDLTensor(support::Stream* strm, const DLTensor* tensor) {
|
||||
uint64_t header = kTVMTensorMagic, reserved = 0;
|
||||
strm->Write(header);
|
||||
strm->Write(reserved);
|
||||
// Always save data as CPU context
|
||||
//
|
||||
// Parameters that get serialized should be in CPU by default.
|
||||
// So even the array's context is GPU, it will be stored as CPU array.
|
||||
// This is used to prevent case when another user loads the parameters
|
||||
// back on machine that do not have GPU or related context.
|
||||
//
|
||||
// We can always do array.CopyTo(target_dev) to get a corresponding
|
||||
// array in the target context.
|
||||
Device cpu_dev;
|
||||
cpu_dev.device_type = kDLCPU;
|
||||
cpu_dev.device_id = 0;
|
||||
strm->Write(cpu_dev);
|
||||
strm->Write(tensor->ndim);
|
||||
strm->Write(tensor->dtype);
|
||||
int ndim = tensor->ndim;
|
||||
strm->WriteArray(tensor->shape, ndim);
|
||||
int type_bytes = (tensor->dtype.bits + 7) / 8;
|
||||
int64_t num_elems = 1;
|
||||
for (int i = 0; i < ndim; ++i) {
|
||||
num_elems *= tensor->shape[i];
|
||||
}
|
||||
int64_t data_byte_size = type_bytes * num_elems;
|
||||
strm->Write(data_byte_size);
|
||||
|
||||
if (TVM_FFI_IO_NO_ENDIAN_SWAP && tensor->device.device_type == kDLCPU &&
|
||||
ffi::IsContiguous(*tensor) && tensor->byte_offset == 0) {
|
||||
// quick path
|
||||
strm->Write(tensor->data, data_byte_size);
|
||||
} else {
|
||||
std::vector<uint8_t> bytes(data_byte_size);
|
||||
Tensor::CopyToBytes(const_cast<DLTensor*>(tensor), bytes.data(), data_byte_size);
|
||||
if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
|
||||
ffi::ByteSwap(bytes.data(), type_bytes, num_elems);
|
||||
}
|
||||
strm->Write(bytes.data(), data_byte_size);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
inline void Tensor::Save(support::Stream* strm) const { SaveDLTensor(strm, operator->()); }
|
||||
|
||||
inline bool Tensor::Load(support::Stream* strm) {
|
||||
uint64_t header, reserved;
|
||||
TVM_FFI_ICHECK(strm->Read(&header)) << "Invalid DLTensor file format";
|
||||
TVM_FFI_ICHECK(strm->Read(&reserved)) << "Invalid DLTensor file format";
|
||||
TVM_FFI_ICHECK(header == kTVMTensorMagic) << "Invalid DLTensor file format";
|
||||
Device dev;
|
||||
int ndim;
|
||||
DLDataType dtype;
|
||||
TVM_FFI_ICHECK(strm->Read(&dev)) << "Invalid DLTensor file format";
|
||||
TVM_FFI_ICHECK(strm->Read(&ndim)) << "Invalid DLTensor file format";
|
||||
TVM_FFI_ICHECK(strm->Read(&dtype)) << "Invalid DLTensor file format";
|
||||
TVM_FFI_ICHECK_EQ(dev.device_type, kDLCPU)
|
||||
<< "Invalid DLTensor device: can only save as CPU tensor";
|
||||
std::vector<int64_t> shape(ndim);
|
||||
if (ndim != 0) {
|
||||
TVM_FFI_ICHECK(strm->ReadArray(&shape[0], ndim)) << "Invalid DLTensor file format";
|
||||
}
|
||||
Tensor ret = Tensor::Empty(ffi::Shape(shape), dtype, dev);
|
||||
int64_t num_elems = 1;
|
||||
int elem_bytes = (ret->dtype.bits + 7) / 8;
|
||||
for (int i = 0; i < ret->ndim; ++i) {
|
||||
num_elems *= ret->shape[i];
|
||||
}
|
||||
int64_t data_byte_size;
|
||||
TVM_FFI_ICHECK(strm->Read(&data_byte_size)) << "Invalid DLTensor file format";
|
||||
TVM_FFI_ICHECK(data_byte_size == num_elems * elem_bytes) << "Invalid DLTensor file format";
|
||||
auto read_ret = strm->Read(ret->data, data_byte_size);
|
||||
// Only check non-empty data
|
||||
if (ndim > 0 && shape[0] != 0) {
|
||||
TVM_FFI_ICHECK(read_ret) << "Invalid DLTensor file format";
|
||||
}
|
||||
if (!TVM_FFI_IO_NO_ENDIAN_SWAP) {
|
||||
ffi::ByteSwap(ret->data, elem_bytes, num_elems);
|
||||
}
|
||||
*this = ret;
|
||||
return true;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Get the preferred host device from the input device.
|
||||
* - For CUDA and ROCm, CUDAHost and ROCMHost will be returned for pinned memory,
|
||||
* since pinned memory reduces copy overhead.
|
||||
* - For other devices, CPU is returned as a fallback.
|
||||
*/
|
||||
inline Device GetPreferredHostDevice(Device device) {
|
||||
if (device.device_type == DLDeviceType::kDLCUDA) {
|
||||
return Device{DLDeviceType::kDLCUDAHost, 0};
|
||||
} else if (device.device_type == DLDeviceType::kDLROCM) {
|
||||
return Device{DLDeviceType::kDLROCMHost, 0};
|
||||
} else {
|
||||
// Fallback to CPU.
|
||||
return Device{DLDeviceType::kDLCPU, 0};
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
namespace std {
|
||||
template <>
|
||||
struct hash<tvm::Device> {
|
||||
std::size_t operator()(const tvm::Device& dev) const {
|
||||
return ((dev.device_id << 8) | dev.device_type);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct equal_to<tvm::Device> {
|
||||
bool operator()(const tvm::Device& lhs, const tvm::Device& rhs) const {
|
||||
return (lhs.device_type == rhs.device_type && lhs.device_id == rhs.device_id);
|
||||
}
|
||||
};
|
||||
} // namespace std
|
||||
|
||||
#endif // TVM_RUNTIME_TENSOR_H_
|
||||
@@ -0,0 +1,194 @@
|
||||
/*
|
||||
* 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 include/tvm/runtime/timer.h
|
||||
* \brief Runtime timer primitives: Timer, TimerNode, WrapTimeEvaluator.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_TIMER_H_
|
||||
#define TVM_RUNTIME_TIMER_H_
|
||||
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/base.h>
|
||||
#include <tvm/runtime/device_api.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
/*! \brief Base class for all timer implementations.
|
||||
*
|
||||
* New implementations of this interface should make sure that `Start` and `Stop`
|
||||
* are as lightweight as possible. Expensive state synchronization should be
|
||||
* done in `SyncAndGetElapsedNanos`.
|
||||
*/
|
||||
class TimerNode : public ffi::Object {
|
||||
public:
|
||||
/*! \brief Start the timer.
|
||||
*
|
||||
* Note: this function should only be called once per object.
|
||||
*/
|
||||
virtual void Start() = 0;
|
||||
/*! \brief Stop the timer.
|
||||
*
|
||||
* Note: this function should only be called once per object.
|
||||
*/
|
||||
virtual void Stop() = 0;
|
||||
/*! \brief Synchronize timer state and return elapsed time between `Start` and `Stop`.
|
||||
* \return The time in nanoseconds between `Start` and `Stop`.
|
||||
*
|
||||
* This function is necessary because we want to avoid timing the overhead of
|
||||
* doing timing. When using multiple timers, it is recommended to stop all of
|
||||
* them before calling `SyncAndGetElapsedNanos` on any of them.
|
||||
*
|
||||
* Note: this function should be only called once per object. It may incur
|
||||
* a large synchronization overhead (for example, with GPUs).
|
||||
*/
|
||||
virtual int64_t SyncAndGetElapsedNanos() = 0;
|
||||
|
||||
virtual ~TimerNode() {}
|
||||
|
||||
static constexpr const bool _type_mutable = true;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("runtime.TimerNode", TimerNode, ffi::Object);
|
||||
};
|
||||
|
||||
/*! \brief Timer for a specific device.
|
||||
*
|
||||
* This is a managed reference to a TimerNode.
|
||||
*
|
||||
* \sa TimerNode
|
||||
*/
|
||||
class Timer : public ffi::ObjectRef {
|
||||
public:
|
||||
/*!
|
||||
* \brief Get a device specific timer.
|
||||
* \param dev The device to time.
|
||||
* \return A `Timer` that has already been started.
|
||||
*
|
||||
* Use this function to time runtime of arbitrary regions of code on a specific
|
||||
* device. The code that you want to time should be running on the device
|
||||
* otherwise the timer will not return correct results. This is a lower level
|
||||
* interface than TimeEvaluator and only runs the timed code once
|
||||
* (TimeEvaluator runs the code multiple times).
|
||||
*
|
||||
* A default timer is used if a device specific one does not exist. This
|
||||
* timer performs synchronization between the device and CPU, which can lead
|
||||
* to overhead in the reported results.
|
||||
*
|
||||
* Example usage:
|
||||
* \code{.cpp}
|
||||
* Timer t = Timer::Start(Device::cpu());
|
||||
* my_long_running_function();
|
||||
* t->Stop();
|
||||
* ... // some more computation
|
||||
* int64_t nanosecs = t->SyncAndGetElapsedNanos() // elapsed time in nanoseconds
|
||||
* \endcode
|
||||
*
|
||||
* To add a new device-specific timer, register a new function
|
||||
* "runtime.timer.my_device" (where `my_device` is the `DeviceName` of your
|
||||
* device). This function should accept a `Device` and return a new `Timer`
|
||||
* that has already been started.
|
||||
*
|
||||
* For example, this is how the CPU timer is implemented:
|
||||
* \code{.cpp}
|
||||
* class CPUTimerNode : public TimerNode {
|
||||
* public:
|
||||
* virtual void Start() { start_ = std::chrono::high_resolution_clock::now(); }
|
||||
* virtual void Stop() { duration_ = std::chrono::high_resolution_clock::now() - start_; }
|
||||
* virtual int64_t SyncAndGetElapsedNanos() { return duration_.count(); }
|
||||
* virtual ~CPUTimerNode() {}
|
||||
*
|
||||
* static constexpr const char* _type_key = "runtime.CPUTimerNode";
|
||||
* TVM_FFI_DECLARE_OBJECT_INFO_FINAL(CPUTimerNode, TimerNode);
|
||||
*
|
||||
* private:
|
||||
* std::chrono::high_resolution_clock::time_point start_;
|
||||
* std::chrono::duration<int64_t, std::nano> duration_;
|
||||
* };
|
||||
*
|
||||
*
|
||||
* TVM_FFI_STATIC_INIT_BLOCK() {
|
||||
* namespace refl = tvm::ffi::reflection;
|
||||
* refl::GlobalDef().def("runtime.timer.cpu", [](Device dev) {
|
||||
* return Timer(ffi::make_object<CPUTimerNode>());
|
||||
* });
|
||||
* }
|
||||
* \endcode
|
||||
*/
|
||||
static TVM_RUNTIME_DLL Timer Start(Device dev);
|
||||
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Timer, ffi::ObjectRef, TimerNode);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Wrap a timer function to measure the time cost of a given packed function.
|
||||
*
|
||||
* Approximate implementation:
|
||||
* \code{.py}
|
||||
* f() // warmup
|
||||
* for i in range(repeat)
|
||||
* f_preproc()
|
||||
* while True:
|
||||
* start = time()
|
||||
* for j in range(number):
|
||||
* f()
|
||||
* duration_ms = time() - start
|
||||
* if duration_ms >= min_repeat_ms:
|
||||
* break
|
||||
* else:
|
||||
* number = (min_repeat_ms / (duration_ms / number) + 1
|
||||
* if cooldown_interval_ms and i % repeats_to_cooldown == 0:
|
||||
* sleep(cooldown_interval_ms)
|
||||
* \endcode
|
||||
*
|
||||
* \param f The function argument.
|
||||
* \param dev The device.
|
||||
* \param number The number of times to run this function for taking average.
|
||||
* We call these runs as one `repeat` of measurement.
|
||||
* \param repeat The number of times to repeat the measurement.
|
||||
* In total, the function will be invoked (1 + number x repeat) times,
|
||||
* where the first one is warm up and will be discarded.
|
||||
* The returned result contains `repeat` costs,
|
||||
* each of which is an average of `number` costs.
|
||||
* \param min_repeat_ms The minimum duration of one `repeat` in milliseconds.
|
||||
* By default, one `repeat` contains `number` runs. If this parameter is set,
|
||||
* the parameters `number` will be dynamically adjusted to meet the
|
||||
* minimum duration requirement of one `repeat`.
|
||||
* i.e., When the run time of one `repeat` falls below this time,
|
||||
* the `number` parameter will be automatically increased.
|
||||
* \param limit_zero_time_iterations The maximum number of repeats when
|
||||
* measured time is equal to 0. It helps to avoid hanging during measurements.
|
||||
* \param cooldown_interval_ms The cooldown interval in milliseconds between the number of repeats
|
||||
* defined by `repeats_to_cooldown`.
|
||||
* \param repeats_to_cooldown The number of repeats before the
|
||||
* cooldown is activated.
|
||||
* \param cache_flush_bytes The number of bytes to flush from cache before
|
||||
* \param f_preproc The function to be executed before we execute time
|
||||
* evaluator.
|
||||
* \return f_timer A timer function.
|
||||
*/
|
||||
ffi::Function WrapTimeEvaluator(ffi::Function f, Device dev, int number, int repeat,
|
||||
int min_repeat_ms, int limit_zero_time_iterations,
|
||||
int cooldown_interval_ms, int repeats_to_cooldown,
|
||||
int cache_flush_bytes = 0, ffi::Function f_preproc = nullptr);
|
||||
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_TIMER_H_
|
||||
@@ -0,0 +1,89 @@
|
||||
/*
|
||||
* 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/vm/builtin.h
|
||||
* \brief Builtin runtime APIs.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_VM_BUILTIN_H_
|
||||
#define TVM_RUNTIME_VM_BUILTIN_H_
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace vm {
|
||||
|
||||
/*!
|
||||
* \brief Op code used in built-in match-shape function.
|
||||
*
|
||||
* The function takes the following signature:
|
||||
|
||||
* MatchShape(input_shape, shape_heap, n, c[0], r[0], c[1], r[1], ... c[n], r[n], err_ctx)
|
||||
*
|
||||
* This function provides runtime shape population and checking support for match-cast.
|
||||
* When a shape variable appears in the first time, we should load the shape and
|
||||
* populate the variable. When a shape variable already appears, we should
|
||||
* assert that it already equals an existing shape value.
|
||||
*
|
||||
* NOTE: It is OK to pass nullptr shape_heap if all code are AssertEqualToImm.
|
||||
*/
|
||||
enum class MatchShapeCode : int {
|
||||
/*!
|
||||
* \brief Perform an assertion that shape equals immediate.
|
||||
*
|
||||
* assert input_shape[i] == r[i]
|
||||
*/
|
||||
kAssertEqualToImm = 0,
|
||||
/*!
|
||||
* \brief This is the first time we see a symbolic shape variable, store to heap.
|
||||
*
|
||||
* shape_heap[r[i]] = input_shape[i]
|
||||
*/
|
||||
kStoreToHeap = 1,
|
||||
/*!
|
||||
* \brief skip and do not do anything.
|
||||
*/
|
||||
kNoOp = 2,
|
||||
/*!
|
||||
* \brief Peform an assertion that the shape equals a loaded value.
|
||||
*
|
||||
* assert input_shape[i] == shape_heap[r[i]]
|
||||
*/
|
||||
kAssertEqualToLoad = 3,
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Op code used in builtin function MakeShape.
|
||||
*
|
||||
* MakeShape(shape_heap, n, c[0], r[0], c[1], r[1], ... c[n], r[n]).
|
||||
*
|
||||
* \note It is OK to pass nullptr to shape_heap if all code are UseImm.
|
||||
*/
|
||||
enum class MakeShapeCode : int {
|
||||
/*! \brief Use the following r[i] as immediate shape value. */
|
||||
kUseImm = 0,
|
||||
/*!
|
||||
* \brief Load shape value from the shape_heap[[r[i]].
|
||||
*/
|
||||
kLoadShape = 1,
|
||||
};
|
||||
|
||||
} // namespace vm
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
#endif // TVM_RUNTIME_VM_BUILTIN_H_
|
||||
@@ -0,0 +1,246 @@
|
||||
/*
|
||||
* 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/vm/bytecode.h
|
||||
* \brief The bytecode for the virtual machine.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_VM_BYTECODE_H_
|
||||
#define TVM_RUNTIME_VM_BYTECODE_H_
|
||||
|
||||
#include <tvm/ffi/dtype.h>
|
||||
#include <tvm/ffi/error.h>
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace vm {
|
||||
|
||||
/*!
|
||||
* \brief The storage type for the bytecode in the VM.
|
||||
*/
|
||||
using ExecWord = int64_t;
|
||||
|
||||
/*! \brief A register name. */
|
||||
using RegName = ExecWord;
|
||||
|
||||
/*!
|
||||
* \brief An alias for the integer type used ubiquitously in the VM.
|
||||
*/
|
||||
using Index = ExecWord;
|
||||
|
||||
/*!
|
||||
* \brief An enumeration of Relax's opcodes.
|
||||
*
|
||||
* The opcode is used to implement instruction
|
||||
* as a tagged union.
|
||||
*/
|
||||
enum class Opcode {
|
||||
Call = 1U,
|
||||
Ret = 2U,
|
||||
Goto = 3U,
|
||||
If = 4U,
|
||||
};
|
||||
|
||||
/*! \brief A single virtual machine instruction.
|
||||
*
|
||||
* The representation of the instruction is as
|
||||
* a tagged union.
|
||||
*
|
||||
* The first field represents which instruction,
|
||||
* and by extension which field of the union
|
||||
* is active.
|
||||
*/
|
||||
struct Instruction {
|
||||
/*! \brief The number of bit for storing value. */
|
||||
static constexpr ExecWord kKindBit = 8;
|
||||
/*! \brief The number of bit for storing value. */
|
||||
static constexpr ExecWord kValueBit = sizeof(ExecWord) * 8 - kKindBit;
|
||||
/*! \brief The bit mask of the value part. */
|
||||
static constexpr ExecWord kValueMask = (static_cast<ExecWord>(1) << kValueBit) - 1;
|
||||
/*! \brief Maximum possible value, use 1 bit for sign. */
|
||||
static constexpr ExecWord kValueMaxLimit = (static_cast<ExecWord>(1) << (kValueBit - 1)) - 1;
|
||||
/*! \brief Minimum possible value, remove 1 slot to keep things symmetric. */
|
||||
static constexpr ExecWord kValueMinLimit = -kValueMaxLimit;
|
||||
/*! \brief Beginning of special register section. */
|
||||
static constexpr RegName kBeginSpecialReg = static_cast<ExecWord>(1) << 54;
|
||||
/*! \brief Random magic number that represents void argument, indicate null value */
|
||||
static constexpr RegName kVoidRegister = kBeginSpecialReg + 0;
|
||||
/*! \brief Random magic number that represents the VM context */
|
||||
static constexpr RegName kVMRegister = kBeginSpecialReg + 1;
|
||||
/*!
|
||||
* \brief The kind of instruction's argument.
|
||||
*/
|
||||
enum class ArgKind : int { kRegister = 0, kImmediate = 1, kConstIdx = 2, kFuncIdx = 3 };
|
||||
|
||||
friend std::ostream& operator<<(std::ostream& os, const ArgKind& kind) {
|
||||
switch (kind) {
|
||||
case ArgKind::kRegister:
|
||||
os << "kRegister";
|
||||
break;
|
||||
case ArgKind::kImmediate:
|
||||
os << "kImmediate";
|
||||
break;
|
||||
case ArgKind::kConstIdx:
|
||||
os << "kConstIdx";
|
||||
break;
|
||||
case ArgKind::kFuncIdx:
|
||||
os << "kFuncIdx";
|
||||
break;
|
||||
default:
|
||||
TVM_FFI_THROW(InternalError)
|
||||
<< "Internal error: "
|
||||
<< "Invalid ArgKind with integer value " << static_cast<int>(kind);
|
||||
}
|
||||
return os;
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief The auxiliary data structure for instruction argument.
|
||||
*/
|
||||
class Arg {
|
||||
public:
|
||||
/*! \brief Construct a void argument. */
|
||||
Arg() : data_(Instruction::kVoidRegister) {}
|
||||
/*!
|
||||
* \brief construct Arg from data.
|
||||
* \param data The data repr.
|
||||
*/
|
||||
static Arg FromData(ExecWord data) { return Arg(data); }
|
||||
/*!
|
||||
* \brief construct a register Arg.
|
||||
* \param reg The register number.
|
||||
* \return The constructed arg.
|
||||
*/
|
||||
static Arg Register(RegName reg) { return Arg(ArgKind::kRegister, reg); }
|
||||
/*!
|
||||
* \brief construct a ConstIdx arg.
|
||||
* \param index The constant index.
|
||||
* \return The constructed arg.
|
||||
*/
|
||||
static Arg ConstIdx(Index index) { return Arg(ArgKind::kConstIdx, index); }
|
||||
/*!
|
||||
* \brief construct a immediate arg.
|
||||
* \param imm_value The immediate value.
|
||||
* \return The constructed arg.
|
||||
*/
|
||||
static Arg Immediate(int64_t imm_value) { return Arg(ArgKind::kImmediate, imm_value); }
|
||||
/*!
|
||||
* \brief construct a FuncIdx arg.
|
||||
* \param index The func index in the function table.
|
||||
* \return The constructed arg.
|
||||
*/
|
||||
static Arg FuncIdx(Index index) { return Arg(ArgKind::kFuncIdx, index); }
|
||||
/*!
|
||||
* \brief Get the kind of argument.
|
||||
* \return The kind of argument.
|
||||
*/
|
||||
ArgKind kind() const {
|
||||
uint8_t kind = (data_ >> kValueBit) & 0xFF;
|
||||
return Instruction::ArgKind(kind);
|
||||
}
|
||||
/*!
|
||||
* \brief Get the value of argument.
|
||||
* \return The value of argument.
|
||||
* \note We store both positive and negative values by sign extension.
|
||||
*/
|
||||
ExecWord value() const { return ((data_ & kValueMask) << kKindBit) >> kKindBit; }
|
||||
/*!
|
||||
* \brief Get the raw data repr of the arg.
|
||||
* \return The raw data.
|
||||
*/
|
||||
ExecWord data() const { return data_; }
|
||||
|
||||
private:
|
||||
/*! \brief Construct from the data. */
|
||||
explicit Arg(ExecWord data) : data_(data) {}
|
||||
/*! \brief Construct from the kind and value. */
|
||||
Arg(ArgKind kind, Index value) {
|
||||
TVM_FFI_ICHECK_LE(value, kValueMaxLimit);
|
||||
TVM_FFI_ICHECK_GE(value, kValueMinLimit);
|
||||
data_ = (static_cast<ExecWord>(kind) << kValueBit) | (value & kValueMask);
|
||||
}
|
||||
/*! \brief The underlying stored data. */
|
||||
ExecWord data_;
|
||||
};
|
||||
/*! \brief The instruction opcode. */
|
||||
Opcode op;
|
||||
union {
|
||||
struct /* Call */ {
|
||||
/*! \brief The destination register. */
|
||||
RegName dst;
|
||||
/*! \brief The index into the packed function table. */
|
||||
Index func_idx;
|
||||
/*! \brief The number of arguments to the packed function. */
|
||||
Index num_args;
|
||||
/*! \brief The arguments of the packed function. */
|
||||
Arg* args;
|
||||
};
|
||||
struct /* Ret */ {
|
||||
/*! \brief The return result. */
|
||||
RegName result;
|
||||
};
|
||||
struct /* Goto */ {
|
||||
/*! \brief The jump offset. */
|
||||
Index pc_offset;
|
||||
};
|
||||
struct /* If */ {
|
||||
/*! \brief The register containing the cond value. */
|
||||
RegName cond;
|
||||
/*! \brief The program counter offset for the false branch. */
|
||||
Index false_offset;
|
||||
};
|
||||
};
|
||||
/*!
|
||||
* \brief Construct a Call instruction.
|
||||
* \param func_idx The index of the function to call.
|
||||
* \param num_args The number of arguments.
|
||||
* \param args The input arguments.
|
||||
* \param dst The destination register.
|
||||
* \return The call instruction.
|
||||
*/
|
||||
static Instruction Call(Index func_idx, Index num_args, Arg* args, RegName dst);
|
||||
/*!
|
||||
* \brief Construct a return instruction.
|
||||
* \param result The register containing the return value.
|
||||
* \return The return instruction.
|
||||
*/
|
||||
static Instruction Ret(RegName result);
|
||||
/*!
|
||||
* \brief Construct a goto instruction.
|
||||
* \param pc_offset The register containing the jump offset.
|
||||
* \return The goto instruction.
|
||||
*/
|
||||
static Instruction Goto(RegName pc_offset);
|
||||
/*!
|
||||
* \brief Construct an If instruction.
|
||||
* \param cond The register containing the cond value.
|
||||
* \param false_offset The program counter offset for the false branch.
|
||||
* \return The If instruction.
|
||||
*/
|
||||
static Instruction If(RegName cond, Index false_offset);
|
||||
};
|
||||
|
||||
} // namespace vm
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_VM_BYTECODE_H_
|
||||
@@ -0,0 +1,242 @@
|
||||
/*
|
||||
* 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/vm/executable.h
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_VM_EXECUTABLE_H_
|
||||
#define TVM_RUNTIME_VM_EXECUTABLE_H_
|
||||
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/support/io.h>
|
||||
#include <tvm/support/serializer.h>
|
||||
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "./bytecode.h"
|
||||
|
||||
// Convention: this version should set to minimum TVM version it support
|
||||
// NOTE: this file only changes if we change relax vm format
|
||||
// for example if relax vm format do not change in 0.15, this should remain as 0.14
|
||||
// if it changes in 0.16, we will change it to 0.16
|
||||
#define VM_VERSION "0.14"
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace vm {
|
||||
|
||||
/*!
|
||||
* \brief Information entry in executable function table.
|
||||
*
|
||||
* Contains metadata about the compiled function, as
|
||||
* well as the compiled VM instructions.
|
||||
*/
|
||||
struct VMFuncInfo {
|
||||
/*! \brief kind of the function. */
|
||||
enum class FuncKind : int {
|
||||
/*! \brief system level packed function */
|
||||
kPackedFunc = 0,
|
||||
/*! \brief VM function. */
|
||||
kVMFunc = 1,
|
||||
/*! \brief VMTIR function. */
|
||||
kVMTIRFunc = 2,
|
||||
};
|
||||
/*! \brief The kind of function. */
|
||||
FuncKind kind;
|
||||
/*! \brief The function's name, global symbol */
|
||||
std::string name;
|
||||
/*! \brief The start instruction index of the function. */
|
||||
Index start_instr = 0;
|
||||
/*! \brief The end instruction index of the function. */
|
||||
Index end_instr = 0;
|
||||
/*! \brief The number of arguments of the function. */
|
||||
Index num_args = 0;
|
||||
/*! \brief The register file size of the function. */
|
||||
Index register_file_size = 0;
|
||||
/*! \brief The function parameter names.*/
|
||||
std::vector<std::string> param_names;
|
||||
|
||||
// defined customized loader save
|
||||
void Save(support::Stream* writer) const;
|
||||
bool Load(support::Stream* reader);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The virtual machine executable emitted by the VM compiler.
|
||||
*
|
||||
* The executable contains information (e.g. data in different memory regions)
|
||||
* to run in a virtual machine.
|
||||
*/
|
||||
class TVM_RUNTIME_DLL VMExecutable : public ffi::ModuleObj {
|
||||
public:
|
||||
/*! \brief Get the property of the runtime module .*/
|
||||
int GetPropertyMask() const final { return ffi::Module::kBinarySerializable; };
|
||||
|
||||
/*!
|
||||
* \brief Print the detailed statistics of the given code, i.e. number of
|
||||
* globals and constants, etc.
|
||||
* \return The statistics represented by a string.
|
||||
*/
|
||||
std::string Stats() const;
|
||||
/*!
|
||||
* \brief Get the i-th instruction from the executable.
|
||||
* \param i The index of the instruction to be fetched.
|
||||
* \return The instruction.
|
||||
*/
|
||||
Instruction GetInstruction(Index i) const;
|
||||
/*!
|
||||
* \brief Set j-th byte data of i-th instruction to val.
|
||||
* \param i The index of the instruction to be updated.
|
||||
* \param j The index of the byte data of the instruction to be updated.
|
||||
* \param val The value to be set
|
||||
*/
|
||||
void SetInstructionData(Index i, Index j, ExecWord val);
|
||||
/*!
|
||||
* \brief Print the instructions as text format.
|
||||
* \return The text format of the instructions.
|
||||
*/
|
||||
ffi::String AsText() const;
|
||||
/*!
|
||||
* \brief Print the instructions as python program.
|
||||
* \return The python program of the instructions, represented by a string.
|
||||
*/
|
||||
ffi::String AsPython() const;
|
||||
/*!
|
||||
* \brief Write the VMExecutable to the binary stream in serialized form.
|
||||
* \return The binary bytes that save the executable to.
|
||||
*/
|
||||
ffi::Bytes SaveToBytes() const final;
|
||||
/*!
|
||||
* \brief Load VMExecutable from the binary stream in serialized form.
|
||||
* \param bytes The binary bytes that load the executable from.
|
||||
* \return The loaded executable, in the form of a `runtime::Module`.
|
||||
*/
|
||||
static ffi::Module LoadFromBytes(const ffi::Bytes& bytes);
|
||||
/*!
|
||||
* \brief Write the VMExecutable to the provided path as a file containing its serialized content.
|
||||
* \param file_name The name of the file to write the serialized data to.
|
||||
* \param format The target format of the saved file.
|
||||
*/
|
||||
void WriteToFile(const ffi::String& file_name, const ffi::String& format) const final;
|
||||
/*! \brief Create a Relax virtual machine and load `this` as the executable. */
|
||||
ffi::Module VMLoadExecutable() const;
|
||||
/*! \brief Check if the VMExecutable contains a specific function. */
|
||||
bool HasFunction(const ffi::String& name) const;
|
||||
/*!
|
||||
* \brief Load VMExecutable from the file.
|
||||
* \param file_name The path of the file that load the executable from.
|
||||
* \return The loaded executable, in the form of a `runtime::Module`.
|
||||
*/
|
||||
static ffi::Module LoadFromFile(const ffi::String& file_name);
|
||||
|
||||
/*! \brief The virtual machine's function table. */
|
||||
std::vector<VMFuncInfo> func_table;
|
||||
/*! \brief A map from globals (as strings) to their index in the function map. */
|
||||
std::unordered_map<std::string, Index> func_map;
|
||||
/*! \brief The global constant pool. */
|
||||
std::vector<ffi::Any> constants;
|
||||
/*! \brief The VDevice memory scopes */
|
||||
std::unordered_map<Index, std::string> memory_scopes;
|
||||
/*! \brief The offset of instruction. */
|
||||
std::vector<Index> instr_offset;
|
||||
/*! \brief The byte data of instruction. */
|
||||
std::vector<ExecWord> instr_data;
|
||||
|
||||
virtual ~VMExecutable() {}
|
||||
|
||||
/*! \brief Module type key. */
|
||||
const char* kind() const final;
|
||||
/*!
|
||||
* \brief Look up an exported function by name.
|
||||
* \param name The function name.
|
||||
* \return The function if found, otherwise std::nullopt.
|
||||
*/
|
||||
ffi::Optional<ffi::Function> GetFunction(const ffi::String& name) override;
|
||||
|
||||
private:
|
||||
/*!
|
||||
* \brief Save the globals.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void SaveGlobalSection(support::Stream* strm) const;
|
||||
/*!
|
||||
* \brief Save the memory scopes.
|
||||
* \param strm The output stream.
|
||||
*/
|
||||
void SaveMemoryScopeSection(support::Stream* strm) const;
|
||||
/*!
|
||||
* \brief Save the constant pool.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void SaveConstantSection(support::Stream* strm) const;
|
||||
/*!
|
||||
* \brief Save the instructions.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void SaveCodeSection(support::Stream* strm) const;
|
||||
/*!
|
||||
* \brief Save the packed functions.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void SavePackedFuncNames(support::Stream* strm) const;
|
||||
/*!
|
||||
* \brief Load the globals.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void LoadGlobalSection(support::Stream* strm);
|
||||
/*!
|
||||
* \brief Load the memory scopes.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void LoadMemoryScopeSection(support::Stream* strm);
|
||||
/*!
|
||||
* \brief Load the constant pool.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void LoadConstantSection(support::Stream* strm);
|
||||
/*!
|
||||
* \brief Load the instructions.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void LoadCodeSection(support::Stream* strm);
|
||||
/*!
|
||||
* \brief Save the packed functions.
|
||||
* \param strm The input stream.
|
||||
*/
|
||||
void LoadPackedFuncNames(support::Stream* strm);
|
||||
};
|
||||
|
||||
} // namespace vm
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
namespace tvm {
|
||||
namespace support {
|
||||
template <>
|
||||
struct Serializer<::tvm::runtime::vm::VMFuncInfo> {
|
||||
static constexpr bool enabled = true;
|
||||
static void Write(Stream* strm, const ::tvm::runtime::vm::VMFuncInfo& data) { data.Save(strm); }
|
||||
static bool Read(Stream* strm, ::tvm::runtime::vm::VMFuncInfo* data) { return data->Load(strm); }
|
||||
};
|
||||
} // namespace support
|
||||
} // namespace tvm
|
||||
#endif // TVM_RUNTIME_VM_EXECUTABLE_H_
|
||||
@@ -0,0 +1,97 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
|
||||
#define TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
|
||||
|
||||
#include <tvm/ffi/container/array.h>
|
||||
#include <tvm/ffi/function.h>
|
||||
#include <tvm/runtime/tensor.h>
|
||||
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
namespace vm {
|
||||
|
||||
/*!
|
||||
* \brief Metadata for Tensor cache, which by default, is named as "tensor-cache.json".
|
||||
*/
|
||||
struct TensorCacheMetadata {
|
||||
/*! \brief Each shard of Tensor cache, which by default, is named as "params_shard_x.bin". */
|
||||
struct FileRecord {
|
||||
/*! \brief Metadata of each parameter */
|
||||
struct ParamRecord {
|
||||
/*!
|
||||
* \brief Load the parameter from raw data.
|
||||
* \param device The device to load the parameter onto.
|
||||
* \param raw_data The raw data stream
|
||||
* \param staging_buffer The buffer to be used to avoid extra OpenCL copies. Pass in a nullptr
|
||||
* in other cases
|
||||
*/
|
||||
TVM_RUNTIME_DLL Tensor Load(Device device, const std::string* raw_data,
|
||||
ffi::Optional<Tensor>* staging_buffer = nullptr) const;
|
||||
|
||||
/*! \brief Name of the parameter */
|
||||
std::string name;
|
||||
/*! \brief Shape of the parameter */
|
||||
ffi::Shape shape;
|
||||
/*! \brief Data type of the parameter */
|
||||
DLDataType dtype;
|
||||
/*! \brief Format of the parameter */
|
||||
std::string format;
|
||||
/*! \brief Number of bytes */
|
||||
int64_t nbytes;
|
||||
/*! \brief Offset from the raw stream */
|
||||
int64_t byte_offset;
|
||||
};
|
||||
|
||||
/*! \brief Load a FileRecord into memory */
|
||||
TVM_RUNTIME_DLL ffi::Array<Tensor> Load(Device device, //
|
||||
const std::string& path_prefix, //
|
||||
std::string* raw_data_buffer, //
|
||||
ffi::Optional<Tensor>* staging_buffer = nullptr) const;
|
||||
|
||||
/*! \brief Relative path to the bin file */
|
||||
std::string data_path;
|
||||
/*! \brief Format of the file */
|
||||
std::string format;
|
||||
/*! \brief Size of the file */
|
||||
int64_t nbytes;
|
||||
/*! \brief The parameters in the file */
|
||||
std::vector<ParamRecord> records;
|
||||
};
|
||||
/*! \brief The files in the Tensor cache */
|
||||
std::vector<FileRecord> records;
|
||||
/*! \brief The path to the `tensor-cache.json` file */
|
||||
std::string path;
|
||||
|
||||
/*! \brief Load the metadata from a specific directory */
|
||||
TVM_RUNTIME_DLL static TensorCacheMetadata Load(const std::string& path);
|
||||
/*! \brief Load the metadata from a given JSON string */
|
||||
TVM_RUNTIME_DLL static TensorCacheMetadata LoadFromStr(const std::string& json_str,
|
||||
const std::string& path);
|
||||
};
|
||||
|
||||
} // namespace vm
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_VM_TENSOR_CACHE_SUPPORT_H_
|
||||
@@ -0,0 +1,226 @@
|
||||
/*
|
||||
* 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/vm/vm.h
|
||||
*/
|
||||
#ifndef TVM_RUNTIME_VM_VM_H_
|
||||
#define TVM_RUNTIME_VM_VM_H_
|
||||
|
||||
#include <tvm/ffi/extra/module.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "../memory/memory_manager.h"
|
||||
#include "./bytecode.h"
|
||||
#include "./executable.h"
|
||||
|
||||
namespace tvm {
|
||||
namespace runtime {
|
||||
|
||||
using memory::Allocator;
|
||||
using memory::AllocatorType;
|
||||
using memory::MemoryManager;
|
||||
using memory::Storage;
|
||||
using memory::StorageObj;
|
||||
|
||||
namespace vm {
|
||||
|
||||
/*!
|
||||
* \brief Possible instrument actions.
|
||||
*/
|
||||
enum class VMInstrumentReturnKind : int {
|
||||
/*! \brief Running as normal. */
|
||||
kNoOp = 0,
|
||||
/*! \brief Skip the following run, only valid in before. */
|
||||
kSkipRun = 1,
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief An object representing a vm closure.
|
||||
*/
|
||||
class VMClosureObj : public ffi::Object {
|
||||
public:
|
||||
/*!
|
||||
* \brief The function name. The function could be any
|
||||
* function object that is compatible to the VM runtime.
|
||||
*/
|
||||
ffi::String func_name;
|
||||
|
||||
/*!
|
||||
* \brief The implementation of the Closure.
|
||||
* \note This function takes context pointer(VirtualMachine*)
|
||||
* as the first argument. The rest of arguments follows
|
||||
* the same arguments as the normal function call.
|
||||
*/
|
||||
ffi::Function impl;
|
||||
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("relax.vm.Closure", VMClosureObj, ffi::Object);
|
||||
};
|
||||
|
||||
/*! \brief reference to closure. */
|
||||
class VMClosure : public ffi::ObjectRef {
|
||||
public:
|
||||
VMClosure(ffi::String func_name, ffi::Function impl);
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMClosure, ffi::ObjectRef, VMClosureObj);
|
||||
|
||||
/*!
|
||||
* \brief Create another ffi::Function with last arguments already bound to last_args.
|
||||
*
|
||||
* This is a helper function to create captured closures.
|
||||
* \param func The input func, can be a VMClosure or ffi::Function.
|
||||
* \param last_args The arguments to bound to in the end of the function.
|
||||
* \note The new function takes in arguments and append the last_args in the end.
|
||||
*/
|
||||
static ffi::Function BindLastArgs(ffi::Function func, std::vector<ffi::Any> last_args);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief Represent a VM extension.
|
||||
* A VM extension allows the user to extend the VM with target specific functionalities.
|
||||
* The VM holds the reference of the extensions to ensure the extensions have the same lifetime
|
||||
* as the VM.
|
||||
*
|
||||
* This is the base class for all VM extensions and should not be used directly.
|
||||
*/
|
||||
class VMExtensionNode : public ffi::Object {
|
||||
protected:
|
||||
TVM_FFI_DECLARE_OBJECT_INFO("runtime.VMExtension", VMExtensionNode, ffi::Object);
|
||||
};
|
||||
|
||||
/*! \brief Managed reference to VM extension. */
|
||||
class VMExtension : public ffi::ObjectRef {
|
||||
public:
|
||||
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(VMExtension, ffi::ObjectRef, VMExtensionNode);
|
||||
};
|
||||
|
||||
/*!
|
||||
* \brief The virtual machine.
|
||||
*
|
||||
* The virtual machine contains all the current execution state,
|
||||
* as well as the executable.
|
||||
*
|
||||
* The goal is to have a single self-contained object,
|
||||
* enabling one to easily pass around VMs, execute them on
|
||||
* multiple threads, or serialize them to disk or over the
|
||||
* wire.
|
||||
*/
|
||||
class VirtualMachine : public ffi::ModuleObj {
|
||||
public:
|
||||
/*!
|
||||
* \brief Initialize the virtual machine for a set of devices.
|
||||
* \param devices The set of TVM devices.
|
||||
* \param alloc_types The allocator types for each device.
|
||||
*/
|
||||
virtual void Init(const std::vector<Device>& devices,
|
||||
const std::vector<AllocatorType>& alloc_types) = 0;
|
||||
/*!
|
||||
* \brief Load the executable for the virtual machine.
|
||||
* \param exec The executable.
|
||||
*/
|
||||
virtual void LoadExecutable(ffi::ObjectPtr<VMExecutable> exec) = 0;
|
||||
/*!
|
||||
* \brief Get global function in the VM.
|
||||
* \param func_name The name of the function.
|
||||
* \return The closure
|
||||
*/
|
||||
virtual VMClosure GetClosure(const ffi::String& func_name) = 0;
|
||||
/*!
|
||||
* \brief Invoke closure or packed function using ffi::Function convention.
|
||||
* \param closure_or_packedfunc A VM closure or a packed_func.
|
||||
* \param args The input arguments.
|
||||
* \param rv The return value.
|
||||
*/
|
||||
virtual void InvokeClosurePacked(const ffi::ObjectRef& closure_or_packedfunc,
|
||||
ffi::PackedArgs args, ffi::Any* rv) = 0;
|
||||
/*!
|
||||
* \brief Set an instrumentation function.
|
||||
*
|
||||
* If instrument is present, the function will be called
|
||||
* before/after each Call instruction.
|
||||
*
|
||||
* bool instrument(func, func_symbol, before_run, args...)
|
||||
*
|
||||
* - func: Union[VMClosure, ffi::Function], the function object.
|
||||
* - func_symbol: string, the symbol of the function.
|
||||
* - before_run: bool, whether it is before or after call.
|
||||
* - ret_value: Only valid in after run, otherwise it is null.
|
||||
* - args: the arguments being passed to call.
|
||||
*
|
||||
* instrument can return an int which corresponds to the action value.
|
||||
* \sa VMInstrumentAction
|
||||
*
|
||||
* \param instrument The instrument function.
|
||||
*/
|
||||
virtual void SetInstrument(ffi::Function instrument) = 0;
|
||||
|
||||
/*!
|
||||
* \brief Get or create a VM extension. Once created, the extension will be stored in the VM
|
||||
* and held until the VM is destructed.
|
||||
*
|
||||
* \tparam T The type of the extension
|
||||
* \return The extension instance
|
||||
*/
|
||||
template <typename T, typename = std::enable_if_t<std::is_base_of<VMExtension, T>::value>>
|
||||
T GetOrCreateExtension() {
|
||||
using ContainerType = typename T::ContainerType;
|
||||
uint32_t key = ContainerType::RuntimeTypeIndex();
|
||||
if (auto it = extensions.find(key); it != extensions.end()) {
|
||||
ffi::Any value = (*it).second;
|
||||
return value.cast<T>();
|
||||
}
|
||||
auto [it, _] = extensions.emplace(key, T::Create());
|
||||
ffi::Any value = (*it).second;
|
||||
return value.cast<T>();
|
||||
}
|
||||
|
||||
/*!
|
||||
* \brief Create a specific instance of VM.
|
||||
* \return Created VM
|
||||
*/
|
||||
static ffi::ObjectPtr<VirtualMachine> Create();
|
||||
/*!
|
||||
* \brief Helper function for vm closure functions to get the context ptr
|
||||
* \param arg The argument value.
|
||||
*/
|
||||
static VirtualMachine* GetContextPtr(ffi::AnyView arg) {
|
||||
return static_cast<VirtualMachine*>(arg.cast<void*>());
|
||||
}
|
||||
|
||||
~VirtualMachine() {}
|
||||
|
||||
//--------------------------------------------------------------------------
|
||||
// The following section contains states that other builtin can depend on
|
||||
//--------------------------------------------------------------------------
|
||||
/*! \brief The memory allocators. */
|
||||
std::vector<Allocator*> allocators;
|
||||
/*! \brief Runtime physical device list. */
|
||||
std::vector<Device> devices;
|
||||
/*! \brief The VM extensions. Mapping from the type index of the extension to the extension
|
||||
* instance. */
|
||||
std::unordered_map<uint32_t, Any> extensions;
|
||||
};
|
||||
|
||||
} // namespace vm
|
||||
} // namespace runtime
|
||||
} // namespace tvm
|
||||
|
||||
#endif // TVM_RUNTIME_VM_VM_H_
|
||||
Reference in New Issue
Block a user