tvm
device_api.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
24 #ifndef TVM_RUNTIME_DEVICE_API_H_
25 #define TVM_RUNTIME_DEVICE_API_H_
26 
27 #include <tvm/ffi/any.h>
28 #include <tvm/ffi/optional.h>
29 #include <tvm/runtime/base.h>
30 #include <tvm/runtime/logging.h>
31 
32 #include <string>
37 typedef void* TVMStreamHandle;
38 
39 namespace tvm {
40 
41 // alias DLDevice
42 using Device = DLDevice;
43 
44 namespace runtime {
45 
69 #ifdef __cplusplus
70 typedef enum : int32_t {
71 #else
72 typedef enum {
73 #endif
74  // To help avoid accidental conflicts between `DLDeviceType`
75  // and this enumeration, start numbering the new enumerators
76  // a little higher than (currently) seems necessary.
77  TVMDeviceExtType_End = 36, // sentinel value
79 
83 enum DeviceAttrKind : int {
84  kExist = 0,
86  kWarpSize = 2,
94  kGcnArch = 10,
101 };
102 
103 #ifdef TVM_KALLOC_ALIGNMENT
105 constexpr int kAllocAlignment = TVM_KALLOC_ALIGNMENT;
106 
108 constexpr int kTempAllocaAlignment = TVM_KALLOC_ALIGNMENT;
109 #else
111 constexpr int kAllocAlignment = 64;
112 
114 constexpr int kTempAllocaAlignment = 64;
115 #endif // TVM_KALLOC_ALIGNMENT
116 
118 constexpr int kMaxStackAlloca = 1024;
119 
122 constexpr int kDefaultWorkspaceAlignment = 1;
123 
128 class TVM_DLL DeviceAPI {
129  public:
131  virtual ~DeviceAPI() {}
136  virtual void SetDevice(Device dev) = 0;
144  virtual void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) = 0;
145 
152  virtual size_t GetDataSize(const DLTensor& arr,
153  ffi::Optional<ffi::String> mem_scope = std::nullopt);
154 
161  virtual void GetTargetProperty(Device dev, const std::string& property, ffi::Any* rv) {}
162 
172  virtual void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment,
173  DLDataType type_hint) = 0;
183  virtual void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
184  ffi::Optional<ffi::String> mem_scope = std::nullopt);
190  virtual void FreeDataSpace(Device dev, void* ptr) = 0;
201  virtual void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream);
208 
215  virtual void FreeStream(Device dev, TVMStreamHandle stream);
216 
222  virtual void StreamSync(Device dev, TVMStreamHandle stream) = 0;
228  virtual void SetStream(Device dev, TVMStreamHandle stream);
247  virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst);
264  virtual void* AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint = {});
271  virtual void FreeWorkspace(Device dev, void* ptr);
272 
279  static DeviceAPI* Get(Device dev, bool allow_missing = false);
280 
286  static bool NeedSetDevice(int device_type) { return device_type != kDLCPU; }
287 
291  virtual bool SupportsDevicePointerArithmeticsOnHost() { return false; }
292 
293  protected:
307  virtual void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset,
308  size_t num_bytes, Device dev_from, Device dev_to,
309  DLDataType type_hint, TVMStreamHandle stream);
310 };
311 
317 inline const char* DLDeviceType2Str(int type) {
318  switch (type) {
319  case kDLCPU:
320  return "cpu";
321  case kDLCUDA:
322  return "cuda";
323  case kDLCUDAHost:
324  return "cuda_host";
325  case kDLCUDAManaged:
326  return "cuda_managed";
327  case kDLOpenCL:
328  return "opencl";
329  case kDLVulkan:
330  return "vulkan";
331  case kDLMetal:
332  return "metal";
333  case kDLVPI:
334  return "vpi";
335  case kDLROCM:
336  return "rocm";
337  case kDLROCMHost:
338  return "rocm_host";
339  case kDLExtDev:
340  return "ext_dev";
341  case kDLOneAPI:
342  return "oneapi";
343  case kDLWebGPU:
344  return "webgpu";
345  case kDLHexagon:
346  return "hexagon";
347  default:
348  LOG(FATAL) << "unknown type = " << type;
349  }
350  throw;
351 }
352 
354 constexpr int kRPCSessMask = 128;
355 static_assert(kRPCSessMask >= TVMDeviceExtType_End);
356 
360 inline bool IsRPCSessionDevice(Device dev) { return (dev.device_type / kRPCSessMask) > 0; }
361 
366 inline int GetRPCSessionIndex(Device dev) {
367  ICHECK(IsRPCSessionDevice(dev)) << "GetRPCSessionIndex: dev has no RPC session";
368  return dev.device_type / kRPCSessMask - 1;
369 }
370 
379  dev.device_type = static_cast<DLDeviceType>(dev.device_type % kRPCSessMask);
380  return dev;
381 }
382 
383 inline std::ostream& operator<<(std::ostream& os, DLDevice dev) { // NOLINT(*)
385  os << "remote[" << tvm::runtime::GetRPCSessionIndex(dev) << "]-";
387  }
388  os << tvm::runtime::DLDeviceType2Str(static_cast<int>(dev.device_type)) << ":" << dev.device_id;
389  return os;
390 }
391 
399 inline Device AddRPCSessionMask(Device dev, int session_table_index) {
400  CHECK(!IsRPCSessionDevice(dev)) << "AddRPCSessionMask: dev already non-zero RPCSessionIndex: "
401  << dev;
402  dev.device_type =
403  static_cast<DLDeviceType>(dev.device_type | (kRPCSessMask * (session_table_index + 1)));
404  return dev;
405 }
406 
407 } // namespace runtime
408 } // namespace tvm
409 
410 #endif // TVM_RUNTIME_DEVICE_API_H_
TVM Runtime Device API, abstracts the device specific interface for memory management.
Definition: device_api.h:128
static bool NeedSetDevice(int device_type)
Whether a certian device type requires set device device before launching the kernel function.
Definition: device_api.h:286
virtual void GetTargetProperty(Device dev, const std::string &property, ffi::Any *rv)
Query the device for specified properties.
Definition: device_api.h:161
static DeviceAPI * Get(Device dev, bool allow_missing=false)
Get device API based on device.
virtual size_t GetDataSize(const DLTensor &arr, ffi::Optional< ffi::String > mem_scope=std::nullopt)
Get the physical memory size required.
virtual void CopyDataFromTo(DLTensor *from, DLTensor *to, TVMStreamHandle stream)
copy data from one place to another
virtual TVMStreamHandle CreateStream(Device dev)
Create a new stream of execution.
virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst)
Synchronize 2 streams of execution.
virtual void FreeWorkspace(Device dev, void *ptr)
Free temporal workspace in backend execution.
virtual bool SupportsDevicePointerArithmeticsOnHost()
Whether pointer arithmetics on a device owned pointer may be performed on the host.
Definition: device_api.h:291
virtual void SetDevice(Device dev)=0
Set the environment device id to device.
virtual void FreeStream(Device dev, TVMStreamHandle stream)
Free a stream of execution.
virtual TVMStreamHandle GetCurrentStream(Device dev)
Get the current stream.
virtual void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any *rv)=0
Get attribute of specified device.
virtual void * AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint={})
Allocate temporal workspace for backend execution.
virtual void StreamSync(Device dev, TVMStreamHandle stream)=0
Synchronize the stream.
virtual void FreeDataSpace(Device dev, void *ptr)=0
Free a data space on device.
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
virtual void SetStream(Device dev, TVMStreamHandle stream)
Set the stream.
virtual void * AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint)=0
Allocate a data space on device.
virtual void * AllocDataSpace(Device dev, int ndim, const int64_t *shape, DLDataType dtype, ffi::Optional< ffi::String > mem_scope=std::nullopt)
Allocate a data space on device with memory scope support.
virtual ~DeviceAPI()
virtual destructor
Definition: device_api.h:131
void * TVMStreamHandle
The stream that is specific to device can be NULL, which indicates the default one.
Definition: device_api.h:37
const char * DLDeviceType2Str(int type)
The name of DLDeviceType.
Definition: device_api.h:317
TVMDeviceExtType
Extension device types in TVM.
Definition: device_api.h:72
@ TVMDeviceExtType_End
Definition: device_api.h:77
constexpr int kMaxStackAlloca
Maximum size that can be allocated on stack.
Definition: device_api.h:118
DeviceAttrKind
the query type into GetAttr
Definition: device_api.h:83
@ kDeviceName
Definition: device_api.h:89
@ kDriverVersion
Definition: device_api.h:96
@ kMaxThreadsPerBlock
Definition: device_api.h:85
@ kMultiProcessorCount
Definition: device_api.h:91
@ kMaxThreadDimensions
Definition: device_api.h:92
@ kApiVersion
Definition: device_api.h:95
@ kImagePitchAlignment
Definition: device_api.h:100
@ kMaxClockRate
Definition: device_api.h:90
@ kWarpSize
Definition: device_api.h:86
@ kTotalGlobalMemory
Definition: device_api.h:98
@ kAvailableGlobalMemory
Definition: device_api.h:99
@ kMaxRegistersPerBlock
Definition: device_api.h:93
@ kComputeVersion
Definition: device_api.h:88
@ kGcnArch
Definition: device_api.h:94
@ kMaxSharedMemoryPerBlock
Definition: device_api.h:87
@ kExist
Definition: device_api.h:84
@ kL2CacheSizeBytes
Definition: device_api.h:97
constexpr int kRPCSessMask
The device type bigger than this is RPC device.
Definition: device_api.h:354
constexpr int kDefaultWorkspaceAlignment
Number of bytes each allocation must align to by default in the workspace buffer to service intermedi...
Definition: device_api.h:122
std::ostream & operator<<(std::ostream &os, const DataType &dtype)
Definition: data_type.h:453
int GetRPCSessionIndex(Device dev)
Return the RPCSessTable index of the RPC Session that owns this device.
Definition: device_api.h:366
constexpr int kTempAllocaAlignment
Number of bytes each allocation must align to in temporary allocation.
Definition: device_api.h:114
bool IsRPCSessionDevice(Device dev)
Return true if a Device is owned by an RPC session.
Definition: device_api.h:360
Device AddRPCSessionMask(Device dev, int session_table_index)
Add a RPC session mask to a Device. RPC clients typically do this when decoding a Device received fro...
Definition: device_api.h:399
constexpr int kAllocAlignment
Number of bytes each allocation must align to.
Definition: device_api.h:111
Device RemoveRPCSessionMask(Device dev)
Remove the RPC session mask from a Device. RPC clients typically do this when encoding a Device for t...
Definition: device_api.h:378
constexpr const char * device_type
The device type.
Definition: stmt.h:1092
Tensor shape(const Tensor &src, DataType dtype, const std::string name="T_shape", const std::string tag=kInjective)
Get the shape of input tensor.
Definition: transform.h:1945
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
DLDevice Device
Definition: device_api.h:42