wehub-resource-sync
428 lines
14 KiB
C++
428 lines
14 KiB
C++
/*
|
|
* 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_
|