InfiniOps/examples/runtime_api.h at e397d93f48544488854d4c36dccd524ebb153a75 · InfiniTensor/InfiniOps · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
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
#ifndef INFINI_OPS_EXAMPLES_RUNTIME_API_H_
#define INFINI_OPS_EXAMPLES_RUNTIME_API_H_

#ifdef WITH_NVIDIA
#include <cuda_runtime.h>
#define DEVICE_MALLOC cudaMalloc
#define DEVICE_FREE cudaFree
#define DEVICE_MEMCPY cudaMemcpy
#define DEVICE_MEMSET cudaMemset
#define DEVICE_MEMCPY_HOST_TO_DEVICE cudaMemcpyHostToDevice
#define DEVICE_MEMCPY_DEVICE_TO_HOST cudaMemcpyDeviceToHost
#define DEFAULT_DEVICE_TYPE Device::Type::kNvidia
#elif WITH_ILUVATAR
#include <cuda_runtime.h>
#define DEVICE_MALLOC cudaMalloc
#define DEVICE_FREE cudaFree
#define DEVICE_MEMCPY cudaMemcpy
#define DEVICE_MEMSET cudaMemset
#define DEVICE_MEMCPY_HOST_TO_DEVICE cudaMemcpyHostToDevice
#define DEVICE_MEMCPY_DEVICE_TO_HOST cudaMemcpyDeviceToHost
#define DEFAULT_DEVICE_TYPE Device::Type::kIluvatar
#elif WITH_HYGON
#include <cuda_runtime.h>
#define DEVICE_MALLOC cudaMalloc
#define DEVICE_FREE cudaFree
#define DEVICE_MEMCPY cudaMemcpy
#define DEVICE_MEMSET cudaMemset
#define DEVICE_MEMCPY_HOST_TO_DEVICE cudaMemcpyHostToDevice
#define DEVICE_MEMCPY_DEVICE_TO_HOST cudaMemcpyDeviceToHost
#define DEFAULT_DEVICE_TYPE Device::Type::kHygon
#elif WITH_METAX
#include <mcr/mc_runtime.h>
#define DEVICE_MALLOC mcMalloc
#define DEVICE_FREE mcFree
#define DEVICE_MEMCPY mcMemcpy
#define DEVICE_MEMSET mcMemset
#define DEVICE_MEMCPY_HOST_TO_DEVICE mcMemcpyHostToDevice
#define DEVICE_MEMCPY_DEVICE_TO_HOST mcMemcpyDeviceToHost
#define DEFAULT_DEVICE_TYPE Device::Type::kMetax
#elif WITH_CAMBRICON
#include <cnrt.h>
#define DEVICE_MALLOC cnrtMalloc
#define DEVICE_FREE cnrtFree
#define DEVICE_MEMCPY cnrtMemcpy
#define DEVICE_MEMSET cnrtMemset
#define DEVICE_MEMCPY_HOST_TO_DEVICE cnrtMemcpyHostToDev
#define DEVICE_MEMCPY_DEVICE_TO_HOST cnrtMemcpyDevToHost
#define DEFAULT_DEVICE_TYPE Device::Type::kCambricon
#elif WITH_MOORE
#include <musa_runtime_api.h>
#define DEVICE_MALLOC musaMalloc
#define DEVICE_FREE musaFree
#define DEVICE_MEMCPY musaMemcpy
#define DEVICE_MEMSET musaMemset
#define DEVICE_MEMCPY_HOST_TO_DEVICE musaMemcpyHostToDevice
#define DEVICE_MEMCPY_DEVICE_TO_HOST musaMemcpyDeviceToHost
#define DEFAULT_DEVICE_TYPE Device::Type::kMoore
#elif WITH_CPU
#include <cstdlib>
#include <cstring>
#define DEVICE_MALLOC(ptr, size) (*(ptr) = std::malloc(size))
#define DEVICE_FREE std::free
#define DEVICE_MEMCPY(dst, src, size, kind) std::memcpy(dst, src, size)
#define DEVICE_MEMSET std::memset
#define DEVICE_MEMCPY_HOST_TO_DEVICE 0
#define DEVICE_MEMCPY_DEVICE_TO_HOST 1
#define DEFAULT_DEVICE_TYPE Device::Type::kCpu
#endif

#endif