Class DeviceAPI

Class DeviceAPI#

Class Documentation#

class DeviceAPI#

TVM Runtime Device API, abstracts the device specific interface for memory management.

Public Functions

inline virtual ~DeviceAPI()#

virtual destructor

virtual void SetDevice(Device dev) = 0#

Set the environment device id to device.

Parameters:

dev – The device to be set.

virtual int GetDeviceIndex()#

Get the currently selected runtime device index on the calling host thread.

The device type is determined by this DeviceAPI instance. The result is a runtime device index (DLDevice::device_id), not a device type or a target configuration index. Implementations must query the current selection without changing it. The default implementation throws NotImplementedError; support for SetDevice does not imply support for this query. RPC devices are excluded.

Returns:

The selected device index, on success.

virtual ffi::Any GetAttr(DLDevice dev, const ffi::String &name) = 0#

Query an attribute of a device.

Parameters:

  • dev – The device to query.

  • name – Attribute name. “exist” returns a boolean; unsupported optional attributes return None. Errors querying a supported attribute are propagated.

Returns:

The attribute value.

virtual ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice device, ffi::Optional<MemoryScope> mem_scope = std::nullopt) = 0#

Create an owning, compact tensor with uninitialized contents.

Parameters:

  • shape – Logical shape, copied into the returned tensor.

  • dtype – Element datatype.

  • device – Device on which to allocate storage.

  • mem_scope – Optional backend memory scope.

Returns:

A tensor whose deleter owns the backend allocation and its resources.

virtual void CopyDataFromTo(DLTensor *from, DLTensor *to, TVMStreamHandle stream)#

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.

Note

The copy may happen asynchronously if it involves a GPU context. Call StreamSync to ensure the copy completes from host’s pov.

Parameters:

  • from – The source array.

  • to – The target array.

  • stream – Optional stream object.

virtual TVMStreamHandle CreateStream(Device dev)#

Create a new stream of execution.

Parameters:

dev – The device of allocation.

virtual void FreeStream(Device dev, TVMStreamHandle stream)#

Free a stream of execution.

Parameters:

  • dev – The device of the stream

  • stream – The pointer to be freed.

virtual void StreamSync(Device dev, TVMStreamHandle stream) = 0#

Synchronize the stream.

Parameters:

  • dev – The device to perform operation.

  • stream – The stream to be sync.

virtual void SetStream(Device dev, TVMStreamHandle stream)#

Set the stream.

Parameters:

  • dev – The device to set stream.

  • stream – The stream to be set.

virtual TVMStreamHandle GetCurrentStream(Device dev)#

Get the current stream.

Parameters:

dev – The device to get stream.

Returns:

The current stream of the device.

virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst)#

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.

Parameters:

  • dev – The device of the streams.

  • event_src – The source stream to synchronize.

  • event_dst – The destination stream to synchronize.

virtual void *AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint = {})#

Allocate temporal workspace for backend execution.

  • 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)

Note

We have the following assumption about backend temporal workspace allocation, and backend will optimize for such assumption:

Parameters:

  • dev – The device of allocation.

  • nbytes – The size to be allocated.

  • type_hint – The type of elements. Only needed by certain backends such as OpenGL, as nbytes is sufficient for most backends.

virtual void FreeWorkspace(Device dev, void *ptr)#

Free temporal workspace in backend execution.

Parameters:

  • dev – The device of allocation.

  • ptr – The pointer to be freed.

inline virtual bool SupportsDevicePtrArithmeticOnHost()#

Whether pointer arithmetics on a device owned pointer may be performed on the host.

inline virtual void *GetMappedHostPtrIfAvailable(void *data)#

Query an existing host mapping of a device allocation, if available.

The returned pointer is borrowed from the allocation and remains subject to its mapping lifetime and synchronization requirements. This query does not map, unmap, or synchronize memory, and does not apply a tensor byte offset.

Parameters:

data – The backend’s device allocation handle.

Returns:

The base of the existing host mapping, or nullptr if unavailable.

Public Static Functions

static DeviceAPI *Get(Device dev, bool allow_missing = false)#

Get device API based on device.

Parameters:

  • dev – The device

  • allow_missing – Whether allow missing

Returns:

The corresponding device API.

static inline bool NeedSetDevice(int device_type)#

Whether a certain device type requires setting the device before launching the kernel function.

Parameters:

device_type – The device type.

Protected Functions

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)#

copy data from one place to another

Parameters:

  • from – The source array.

  • from_offset – The byte offset in the from.

  • to – The target array.

  • to_offset – The byte offset in the to.

  • num_bytes – The size of the memory in bytes

  • dev_from – The source device

  • dev_to – The target device

  • type_hint – The type of elements, only neded by certain backends. can be useful for cross device endian converison.

  • stream – Optional stream object.