infinirt.cc 9.56 KB
Newer Older
1
#include "infinirt.h"
2
3
#include "../utils.h"
#include "ascend/infinirt_ascend.h"
4
#include "bang/infinirt_bang.h"
5
6
#include "cpu/infinirt_cpu.h"
#include "cuda/infinirt_cuda.cuh"
zhangyue's avatar
zhangyue committed
7
#include "kunlun/infinirt_kunlun.h"
8
#include "metax/infinirt_metax.h"
9
#include "moore/infinirt_moore.h"
10
11
12

thread_local infiniDevice_t CURRENT_DEVICE_TYPE = INFINI_DEVICE_CPU;
thread_local int CURRENT_DEVICE_ID = 0;
13
14
thread_local infiniDevice_t PREVIOUS_NON_CPU_DEVICE_TYPE = INFINI_DEVICE_TYPE_COUNT;
thread_local int PREVIOUS_NON_CPU_DEVICE_ID = 0;
15

16
__INFINI_C infiniStatus_t infinirtInit() {
17
18
19
#ifdef ENABLE_ASCEND_API
    CHECK_STATUS(infinirt::ascend::init());
#endif
20
21
22
    return INFINI_STATUS_SUCCESS;
}

23
__INFINI_C infiniStatus_t infinirtGetAllDeviceCount(int *count_array) {
24
25
26
27
28
29
30
31
32
33
34
35
    if (count_array == nullptr) {
        return INFINI_STATUS_NULL_POINTER;
    }
    for (size_t i = 0; i < INFINI_DEVICE_TYPE_COUNT; i++) {
        auto status = infinirtGetDeviceCount(static_cast<infiniDevice_t>(i), &count_array[i]);
        if (status != INFINI_STATUS_SUCCESS) {
            return status;
        }
    }
    return INFINI_STATUS_SUCCESS;
}

36
__INFINI_C infiniStatus_t infinirtGetDevice(infiniDevice_t *device_ptr, int *device_id_ptr) {
37
38
39
40
41
42
43
44
45
46
47
48
49
    if (device_ptr == nullptr && device_id_ptr == nullptr) {
        return INFINI_STATUS_NULL_POINTER;
    }
    if (device_ptr != nullptr) {
        *device_ptr = CURRENT_DEVICE_TYPE;
    }
    if (device_id_ptr != nullptr) {
        *device_id_ptr = CURRENT_DEVICE_ID;
    }

    return INFINI_STATUS_SUCCESS;
}

50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
#define INFINIRT_CALL_DEVICE_API_AND(DEVICE_TYPE, API, PARAMS, ACTION) \
    {                                                                  \
        infiniStatus_t _status;                                        \
        switch (DEVICE_TYPE) {                                         \
        case INFINI_DEVICE_CPU:                                        \
            _status = infinirt::cpu::API PARAMS;                       \
            break;                                                     \
        case INFINI_DEVICE_NVIDIA:                                     \
            _status = infinirt::cuda::API PARAMS;                      \
            break;                                                     \
        case INFINI_DEVICE_CAMBRICON:                                  \
            _status = infinirt::bang::API PARAMS;                      \
            break;                                                     \
        case INFINI_DEVICE_ASCEND:                                     \
            _status = infinirt::ascend::API PARAMS;                    \
            break;                                                     \
        case INFINI_DEVICE_METAX:                                      \
67
            _status = infinirt::metax::API PARAMS;                     \
68
69
70
71
            break;                                                     \
        case INFINI_DEVICE_MOORE:                                      \
            _status = infinirt::musa::API PARAMS;                      \
            break;                                                     \
72
73
74
        case INFINI_DEVICE_KUNLUN:                                     \
            _status = infinirt::kunlun::API PARAMS;                    \
            break;                                                     \
zhangyue's avatar
zhangyue committed
75
        case INFINI_DEVICE_ILUVATAR:                                   \
76
            _status = infinirt::iluvatar::API PARAMS;                  \
zhangyue's avatar
zhangyue committed
77
            break;                                                     \
78
        case INFINI_DEVICE_QY:                                         \
79
            _status = infinirt::qy::API PARAMS;                        \
80
            break;                                                     \
81
        case INFINI_DEVICE_HYGON:                                      \
82
            _status = infinirt::hygon::API PARAMS;                     \
83
            break;                                                     \
wooway777's avatar
wooway777 committed
84
85
86
        case INFINI_DEVICE_ALI:                                        \
            _status = infinirt::ali::API PARAMS;                       \
            break;                                                     \
87
        default:                                                       \
zhangyue's avatar
zhangyue committed
88
            _status = INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;         \
89
90
91
        }                                                              \
        { ACTION; }                                                    \
        return _status;                                                \
92
93
    }

94
#define INFINIRT_CALL_DEVICE_API(API, PARAMS) INFINIRT_CALL_DEVICE_API_AND(CURRENT_DEVICE_TYPE, API, PARAMS, )
95

96
__INFINI_C infiniStatus_t infinirtGetDeviceCount(infiniDevice_t device, int *count) {
97
98
99
    if (count == nullptr) {
        return INFINI_STATUS_NULL_POINTER;
    }
100
    INFINIRT_CALL_DEVICE_API_AND(device, getDeviceCount, (count), {});
101
102
}

103
__INFINI_C infiniStatus_t infinirtSetDevice(infiniDevice_t device, int device_id) {
104
105
106
107
108
109
110
111
112
113
    bool skip_set = CURRENT_DEVICE_TYPE == INFINI_DEVICE_CPU && device == PREVIOUS_NON_CPU_DEVICE_TYPE && device_id == PREVIOUS_NON_CPU_DEVICE_ID;
    if (CURRENT_DEVICE_TYPE != INFINI_DEVICE_CPU) {
        PREVIOUS_NON_CPU_DEVICE_TYPE = CURRENT_DEVICE_TYPE;
        PREVIOUS_NON_CPU_DEVICE_ID = CURRENT_DEVICE_ID;
    }
    if (skip_set) {
        CURRENT_DEVICE_TYPE = device;
        CURRENT_DEVICE_ID = device_id;
        return INFINI_STATUS_SUCCESS;
    }
114
    INFINIRT_CALL_DEVICE_API_AND(device, setDevice, (device_id),
115
116
117
118
                                 { CURRENT_DEVICE_TYPE = device;
                                   CURRENT_DEVICE_ID = device_id; });
}

119
__INFINI_C infiniStatus_t infinirtDeviceSynchronize() {
120
121
122
    INFINIRT_CALL_DEVICE_API(deviceSynchronize, ());
}

123
__INFINI_C infiniStatus_t infinirtStreamCreate(infinirtStream_t *stream_ptr) {
124
125
126
    INFINIRT_CALL_DEVICE_API(streamCreate, (stream_ptr));
}

127
__INFINI_C infiniStatus_t infinirtStreamDestroy(infinirtStream_t stream) {
128
129
130
    INFINIRT_CALL_DEVICE_API(streamDestroy, (stream));
}

131
__INFINI_C infiniStatus_t infinirtStreamSynchronize(infinirtStream_t stream) {
132
133
134
    INFINIRT_CALL_DEVICE_API(streamSynchronize, (stream));
}

135
__INFINI_C infiniStatus_t infinirtStreamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) {
136
137
138
    INFINIRT_CALL_DEVICE_API(streamWaitEvent, (stream, event));
}

139
__INFINI_C infiniStatus_t infinirtEventCreate(infinirtEvent_t *event_ptr) {
140
141
142
    INFINIRT_CALL_DEVICE_API(eventCreate, (event_ptr));
}

143
__INFINI_C infiniStatus_t infinirtEventCreateWithFlags(infinirtEvent_t *event_ptr, uint32_t flags) {
thatPepe's avatar
thatPepe committed
144
145
146
    INFINIRT_CALL_DEVICE_API(eventCreateWithFlags, (event_ptr, flags));
}

147
__INFINI_C infiniStatus_t infinirtEventRecord(infinirtEvent_t event, infinirtStream_t stream) {
148
149
150
    INFINIRT_CALL_DEVICE_API(eventRecord, (event, stream));
}

151
__INFINI_C infiniStatus_t infinirtEventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr) {
152
153
154
    INFINIRT_CALL_DEVICE_API(eventQuery, (event, status_ptr));
}

155
__INFINI_C infiniStatus_t infinirtEventSynchronize(infinirtEvent_t event) {
156
157
158
    INFINIRT_CALL_DEVICE_API(eventSynchronize, (event));
}

159
__INFINI_C infiniStatus_t infinirtEventDestroy(infinirtEvent_t event) {
160
161
162
    INFINIRT_CALL_DEVICE_API(eventDestroy, (event));
}

163
__INFINI_C infiniStatus_t infinirtEventElapsedTime(float *ms_ptr, infinirtEvent_t start, infinirtEvent_t end) {
thatPepe's avatar
thatPepe committed
164
165
166
    INFINIRT_CALL_DEVICE_API(eventElapsedTime, (ms_ptr, start, end));
}

167
__INFINI_C infiniStatus_t infinirtMalloc(void **p_ptr, size_t size) {
168
169
170
    INFINIRT_CALL_DEVICE_API(mallocDevice, (p_ptr, size));
}

171
__INFINI_C infiniStatus_t infinirtMallocHost(void **p_ptr, size_t size) {
172
173
174
    INFINIRT_CALL_DEVICE_API(mallocHost, (p_ptr, size));
}

175
__INFINI_C infiniStatus_t infinirtFree(void *ptr) {
176
177
178
    INFINIRT_CALL_DEVICE_API(freeDevice, (ptr));
}

179
__INFINI_C infiniStatus_t infinirtFreeHost(void *ptr) {
180
181
182
    INFINIRT_CALL_DEVICE_API(freeHost, (ptr));
}

183
__INFINI_C infiniStatus_t infinirtMemcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind) {
184
185
186
    INFINIRT_CALL_DEVICE_API(memcpy, (dst, src, size, kind));
}

187
__INFINI_C infiniStatus_t infinirtMemcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream) {
188
189
190
    INFINIRT_CALL_DEVICE_API(memcpyAsync, (dst, src, size, kind, stream));
}

191
__INFINI_C infiniStatus_t infinirtMallocAsync(void **p_ptr, size_t size, infinirtStream_t stream) {
192
193
194
    INFINIRT_CALL_DEVICE_API(mallocAsync, (p_ptr, size, stream));
}

195
__INFINI_C infiniStatus_t infinirtFreeAsync(void *ptr, infinirtStream_t stream) {
196
197
    INFINIRT_CALL_DEVICE_API(freeAsync, (ptr, stream));
}
198

199
__INFINI_C infiniStatus_t infinirtStreamBeginCapture(infinirtStream_t stream, infinirtStreamCaptureMode_t mode) {
200
201
202
    INFINIRT_CALL_DEVICE_API(streamBeginCapture, (stream, mode));
}

203
__INFINI_C infiniStatus_t infinirtStreamEndCapture(infinirtStream_t stream, infinirtGraph_t *graph_ptr) {
204
205
206
    INFINIRT_CALL_DEVICE_API(streamEndCapture, (stream, graph_ptr));
}

207
__INFINI_C infiniStatus_t infinirtGraphDestroy(infinirtGraph_t graph) {
208
209
210
    INFINIRT_CALL_DEVICE_API(graphDestroy, (graph));
}

211
__INFINI_C infiniStatus_t infinirtGraphInstantiate(
212
213
214
215
216
217
218
219
    infinirtGraphExec_t *graph_exec_ptr,
    infinirtGraph_t graph,
    infinirtGraphNode_t *node_ptr,
    char *log_buffer,
    size_t buffer_size) {
    INFINIRT_CALL_DEVICE_API(graphInstantiate, (graph_exec_ptr, graph, node_ptr, log_buffer, buffer_size));
}

220
__INFINI_C infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exec) {
221
222
223
    INFINIRT_CALL_DEVICE_API(graphExecDestroy, (graph_exec));
}

224
__INFINI_C infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) {
225
226
    INFINIRT_CALL_DEVICE_API(graphLuanch, (graph_exec, stream));
}