device_api.h 3.42 KB
Newer Older
1 2
/*!
 *  Copyright (c) 2016 by Contributors
3
 * \file device_api.h
4 5 6 7 8 9
 * \brief Device specific API
 */
#ifndef TVM_RUNTIME_DEVICE_API_H_
#define TVM_RUNTIME_DEVICE_API_H_

#include <tvm/base.h>
10
#include <tvm/runtime/c_runtime_api.h>
11 12 13 14

namespace tvm {
namespace runtime {
/*!
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29
 * \brief Initialize the device.
 * \param option_keys Additional option  keys to pass.
 * \param option_vals Additional option values to pass
 * \param num_options Number of options to be passed into it.
 * \return 0 if success, 1: if already initialized
 * \tparam xpu The device mask.
 */
template<TVMDeviceMask xpu>
inline bool DeviceInit(const char** option_keys,
                       const char** option_vals,
                       int num_options) {
  return true;
}

/*!
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111
 * \brief Whether ctx is enabled.
 * \param ctx The device context to perform operation.
 * \tparam xpu The device mask.
 */
template<TVMDeviceMask xpu>
inline bool CheckEnabled(TVMContext ctx) {
  return true;
}

/*!
 * \brief Allocate a data space on device.
 * \param ctx The device context to perform operation.
 * \param size The size of the memory
 * \param alignment The alignment of the memory.
 * \return The allocated device pointer
 * \tparam xpu The device mask.
 */
template<TVMDeviceMask xpu>
inline void* AllocDataSpace(TVMContext ctx, size_t size, size_t alignment);

/*!
 * \brief Free a data space on device.
 * \param ctx The device context to perform operation.
 * \param ptr The data space.
 * \tparam xpu The device mask.
 */
template<TVMDeviceMask xpu>
inline void FreeDataSpace(TVMContext ctx, void* ptr);

/*!
 * \brief copy data from one place to another
 * \param dev The device to perform operation.
 * \param from The source array.
 * \param to The target array.
 * \param size The size of the memory
 * \param ctx_from The source context
 * \param ctx_to The target context
 * \tparam xpu The device mask.
 */
template<TVMDeviceMask xpu>
inline void CopyDataFromTo(const void* from,
                           void* to,
                           size_t size,
                           TVMContext ctx_from,
                           TVMContext ctx_to,
                           TVMStreamHandle stream);
/*!
 * \brief Synchronize the stream
 * \param ctx The context to perform operation.
 * \param stream The stream to be sync.
 * \tparam xpu The device mask.
 */
template<TVMDeviceMask xpu>
inline void StreamSync(TVMContext ctx, TVMStreamHandle stream);

// macro to run cuda related code
#if TVM_CUDA_RUNTIME
#define TVM_RUN_CUDA(OP) { const TVMDeviceMask xpu = kGPU; OP; }
#else
#define TVM_RUN_CUDA(OP) LOG(FATAL) << "CUDA is not enabled";
#endif

// macro to run opencl related code
#if TVM_OPENCL_RUNTIME
#define TVM_RUN_OPENCL(OP) { const TVMDeviceMask xpu = kOpenCL; OP; }
#else
#define TVM_RUN_OPENCL(OP) LOG(FATAL) << "OpenCL is not enabled";
#endif

// macro to switch options between devices
#define TVM_DEVICE_SWITCH(ctx, OP)                                  \
  switch (ctx.dev_mask) {                                           \
    case kCPU: { const TVMDeviceMask xpu = kCPU; OP; break; }       \
    case kGPU: TVM_RUN_CUDA(OP); break;                             \
    case kOpenCL: TVM_RUN_OPENCL(OP); break;                        \
    default: LOG(FATAL) << "unknown device_mask " << ctx.dev_mask;  \
  }

}  // namespace runtime
}  // namespace tvm

#include "./device_api_cpu.h"
112 113
#include "./cuda/device_api_cuda.h"
#include "./opencl/device_api_opencl.h"
114 115

#endif  // TVM_RUNTIME_DEVICE_API_H_