device_api.h 6.2 KB
Newer Older
1 2
/*!
 *  Copyright (c) 2016 by Contributors
tqchen committed
3
 * \file tvm/runtime/device_api.h
4
 * \brief Abstract device memory management API
5 6 7 8
 */
#ifndef TVM_RUNTIME_DEVICE_API_H_
#define TVM_RUNTIME_DEVICE_API_H_

9
#include <string>
10 11
#include "./packed_func.h"
#include "./c_runtime_api.h"
12 13 14

namespace tvm {
namespace runtime {
15 16 17
/*!
 * \brief the query type into GetAttr
 */
18 19 20
enum DeviceAttrKind : int {
  kExist = 0,
  kMaxThreadsPerBlock = 1,
21
  kWarpSize = 2,
22 23
  kMaxSharedMemoryPerBlock = 3,
  kComputeVersion = 4,
24 25 26
  kDeviceName = 5,
  kMaxClockRate = 6,
  kMultiProcessorCount = 7
27
};
28 29 30 31

/*! \brief Number of bytes each allocation must align to */
constexpr int kAllocAlignment = 64;

32 33 34
/*! \brief Number of bytes each allocation must align to in temporary allocation */
constexpr int kTempAllocaAlignment = 64;

35 36 37
/*! \brief Maximum size that can be allocated on stack */
constexpr int kMaxStackAlloca = 1024;

38 39 40 41
/*!
 * \brief TVM Runtime Device API, abstracts the device
 *  specific interface for memory management.
 */
42 43 44 45 46
class DeviceAPI {
 public:
  /*! \brief virtual destructor */
  virtual ~DeviceAPI() {}
  /*!
47 48
   * \brief Set the environment device id to ctx
   * \param ctx The context to be set.
49
   */
50
  virtual void SetDevice(TVMContext ctx) = 0;
51 52
  /*!
   * \brief Get attribute of specified device.
53
   * \param ctx The device context
54 55
   * \param kind The result kind
   * \param rv The return value.
56
   * \sa DeviceAttrKind
57
   */
58
  virtual void GetAttr(TVMContext ctx, DeviceAttrKind kind, TVMRetValue* rv) = 0;
59
  /*!
60 61
   * \brief Allocate a data space on device.
   * \param ctx The device context to perform operation.
62
   * \param nbytes The number of bytes in memory.
63
   * \param alignment The alignment of the memory.
64 65 66
   * \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.
67
   */
68 69 70 71
  virtual void* AllocDataSpace(TVMContext ctx,
                               size_t nbytes,
                               size_t alignment,
                               TVMType type_hint) = 0;
72 73 74 75 76 77 78 79 80
  /*!
   * \brief Free a data space on device.
   * \param ctx The device context to perform operation.
   * \param ptr The data space.
   */
  virtual void FreeDataSpace(TVMContext ctx, void* ptr) = 0;
  /*!
   * \brief copy data from one place to another
   * \param from The source array.
81
   * \param from_offset The byte offeset in the from.
82
   * \param to The target array.
83
   * \param to_offset The byte offset in the to.
84
   * \param num_bytes The size of the memory in bytes
85 86
   * \param ctx_from The source context
   * \param ctx_to The target context
87 88
   * \param type_hint The type of elements, only neded by certain backends.
   *                  can be useful for cross device endian converison.
89
   * \param stream Optional stream object.
90 91
   */
  virtual void CopyDataFromTo(const void* from,
92
                              size_t from_offset,
93
                              void* to,
94
                              size_t to_offset,
95
                              size_t num_bytes,
96 97
                              TVMContext ctx_from,
                              TVMContext ctx_to,
98
                              TVMType type_hint,
99
                              TVMStreamHandle stream) = 0;
100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
    /*!
   * \brief Create a new stream of execution.
   *
   * \param ctx The context of allocation.
   */
  TVM_DLL virtual TVMStreamHandle CreateStream(TVMContext ctx);

  /*!
   * \brief Free a stream of execution
   *
   * \param ctx The context of the stream
   * \param stream The pointer to be freed.
   */
  TVM_DLL virtual void FreeStream(TVMContext ctx, TVMStreamHandle stream);

115 116 117 118 119 120
  /*!
   * \brief Synchronize the stream
   * \param ctx The context to perform operation.
   * \param stream The stream to be sync.
   */
  virtual void StreamSync(TVMContext ctx, TVMStreamHandle stream) = 0;
121
  /*!
122 123 124 125 126 127
   * \brief Set the stream
   * \param ctx The context to set stream.
   * \param stream The stream to be set.
   */
  virtual void SetStream(TVMContext ctx, TVMStreamHandle stream) {}
  /*!
128 129
   * \brief Synchronize 2 streams of execution.
   *
tqchen committed
130
   * An event is created in event_src stream that the second then
131 132 133 134 135 136 137 138 139 140 141 142
   * stream waits on.  Neither event_src or event_dst need to be of
   * the same device ID as the context, but they must be of the same
   * device type.
   *
   * \param ctx The context of the streams.
   * \param event_src The source stream to synchronize.
   * \param event_dst The destination stream to synchronize.
   */
  TVM_DLL virtual void SyncStreamFromTo(TVMContext ctx,
                                        TVMStreamHandle event_src,
                                        TVMStreamHandle event_dst);
  /*!
143 144 145 146 147 148 149 150 151 152 153
   * \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 ctx The context of allocation.
154 155 156
   * \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.
157
   */
158 159 160
  TVM_DLL virtual void* AllocWorkspace(TVMContext ctx,
                                       size_t nbytes,
                                       TVMType type_hint = {});
161 162 163 164 165 166
  /*!
   * \brief Free temporal workspace in backend execution.
   *
   * \param ctx The context of allocation.
   * \param ptr The pointer to be freed.
   */
167
  TVM_DLL virtual void FreeWorkspace(TVMContext ctx, void* ptr);
168

169
  /*!
170 171 172 173 174
   * \brief Get device API base don context.
   * \param ctx The context
   * \param allow_missing Whether allow missing
   * \return The corresponding device API.
   */
175
  TVM_DLL static DeviceAPI* Get(TVMContext ctx, bool allow_missing = false);
176
};
177

178 179
/*! \brief The device type bigger than this is RPC device */
constexpr int kRPCSessMask = 128;
180 181 182
}  // namespace runtime
}  // namespace tvm
#endif  // TVM_RUNTIME_DEVICE_API_H_