-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathrknn_matmul_api.h
261 lines (229 loc) · 7.98 KB
/
rknn_matmul_api.h
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
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
/****************************************************************************
*
* Copyright (c) 2017 - 2018 by Rockchip Corp. All rights reserved.
*
* The material in this file is confidential and contains trade secrets
* of Rockchip Corporation. This is proprietary information owned by
* Rockchip Corporation. No part of this work may be disclosed,
* reproduced, copied, transmitted, or used in any way for any purpose,
* without the express written permission of Rockchip Corporation.
*
*****************************************************************************/
#ifndef _RKNN_MATMUL_API_H
#define _RKNN_MATMUL_API_H
#ifdef __cplusplus
extern "C" {
#endif
#include "rknn_api.h"
typedef rknn_context rknn_matmul_ctx;
typedef struct _rknn_matmul_tensor_attr
{
char name[RKNN_MAX_NAME_LEN];
// indicate A(M, K) or B(K, N) or C(M, N)
uint32_t n_dims;
uint32_t dims[RKNN_MAX_DIMS];
// matmul tensor size
uint32_t size;
// matmul tensor data type
// int8 : A, B
// int32: C
rknn_tensor_type type;
} rknn_matmul_tensor_attr;
typedef struct _rknn_matmul_io_attr
{
// indicate A(M, K) or B(K, N) or C(M, N)
rknn_matmul_tensor_attr A;
rknn_matmul_tensor_attr B;
rknn_matmul_tensor_attr C;
} rknn_matmul_io_attr;
/*
matmul information struct
*/
typedef struct rknn_matmul_info_t
{
int32_t M;
int32_t K; // limit: rk356x: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte;
// rk3588: int8 type must be aligned with 32byte, float16 type must be aligned with 32byte;
int32_t N; // limit: rk356x: int8 type must be aligned with 16byte, float16 type must be aligned with 8byte;
// rk3588: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte;
// matmul data type
// int8: int8(A) x int8(B) -> int32(C)
// float16: float16(A) x float16(B) -> float32(C)
rknn_tensor_type type;
// matmul native layout for B
// 0: normal layout
// 1: native layout
int32_t native_layout;
// matmul perf layout for A and C
// 0: normal layout
// 1: perf layout
int32_t perf_layout;
} rknn_matmul_info;
/* rknn_matmul_create
params:
rknn_matmul_ctx *ctx the handle of context.
rknn_matmul_info *info the matmal information.
rknn_matmul_io_attr *io_attr inputs/output attribute
return:
int error code
*/
int rknn_matmul_create(rknn_matmul_ctx* ctx, rknn_matmul_info* info, rknn_matmul_io_attr* io_attr);
/* rknn_matmul_set_io_mem
params:
rknn_matmul_ctx ctx the handle of context.
rknn_tensor_mem *mem the pointer of tensor memory information.
rknn_matmul_tensor_attr *attr the attribute of input or output tensor buffer.
return:
int error code.
formula:
C = A * B,
limit:
K <= 4096
K limit: rk356x: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte;
rk3588: int8 type must be aligned with 32byte, float16 type must be aligned with 32byte;
N limit: rk356x: int8 type must be aligned with 16byte, float16 type must be aligned with 8byte;
rk3588: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte;
A shape: M x K
normal layout: (M, K)
[M1K1, M1K2, ..., M1Kk,
M2K1, M2K2, ..., M2Kk,
...
MmK1, MmK2, ..., MmKk]
for rk356x：
int8:
perf layout: (K / 8, M, 8)
[K1M1, K2M1, ..., K8M1,
K9M2, K10M2, ..., K16M2,
...
K(k-7)Mm, K(k-6)Mm, ..., KkMm]
float16:
perf layout: (K / 4, M, 4)
[K1M1, K2M1, ..., K4M1,
K9M2, K10M2, ..., K8M2,
...
K(k-3)Mm, K(k-2)Mm, ..., KkMm]
for rk3588：
int8:
perf layout: (K / 16, M, 16)
[K1M1, K2M1, ..., K16M1,
K9M2, K10M2, ..., K32M2,
...
K(k-15)Mm, K(k-14)Mm, ..., KkMm]
float16:
perf layout: (K / 8, M, 8)
[K1M1, K2M1, ..., K8M1,
K9M2, K10M2, ..., K16M2,
...
K(k-7)Mm, K(k-6)Mm, ..., KkMm]
B shape: K x N
normal layout: (K, N)
[K1N1, K1N2, ..., K1Nn,
K2N1, K2N2, ..., K2Nn,
...
KkN1, KkN2, ..., KkNn]
for rk356x：
int8:
native layout: (N / 16, K / 32, 16, 32)
[K1N1, K2N1, ..., K32N1,
K1N2, K2N2, ..., K32N2,
...
K1N16, K2N16, ..., K32N16,
K33N1, K34N1, ..., K64N1,
K33N2, K34N2, ..., K64N2,
...
K(k-31)N16, K(k-30)N16, ..., KkN16,
K1N17, K2N17, ..., K32N17,
K1N18, K2N18, ..., K32N18,
...
K(k-31)Nn, K(k-30)Nn, ..., KkNn]
float16:
native layout: (N / 8, K / 16, 8, 16)
[K1N1, K2N1, ..., K16N1,
K1N2, K2N2, ..., K16N2,
...
K1N8, K2N8, ..., K16N8,
K17N1, K18N1, ..., K32N1,
K17N2, K18N2, ..., K32N2,
...
K(k-15)N8, K(k-30)N8, ..., KkN8,
K1N9, K2N9, ..., K16N9,
K1N10, K2N10, ..., K16N10,
...
K(k-15)Nn, K(k-14)Nn, ..., KkNn]
for rk3588：
int8:
native layout: (N / 32, K / 32, 32, 32)
[K1N1, K2N1, ..., K32N1,
K1N2, K2N2, ..., K32N2,
...
K1N32, K2N32, ..., K32N32,
K33N1, K34N1, ..., K64N1,
K33N2, K34N2, ..., K64N2,
...
K(k-31)N32, K(k-30)N32, ..., KkN32,
K1N33, K2N33, ..., K32N33,
K1N34, K2N34, ..., K32N34,
...
K(k-31)Nn, K(k-30)Nn, ..., KkNn]
float16:
native layout: (N / 16, K / 32, 16, 32)
[K1N1, K2N1, ..., K32N1,
K1N2, K2N2, ..., K32N2,
...
K1N16, K2N16, ..., K32N16,
K33N1, K34N1, ..., K64N1,
K33N2, K34N2, ..., K64N2,
...
K(k-31)N16, K(k-30)N16, ..., KkN16,
K1N17, K2N17, ..., K32N17,
K1N18, K2N18, ..., K32N18,
...
K(k-31)Nn, K(k-30)Nn, ..., KkNn]
C shape: M x N
normal layout: (M, N)
[M1N1, M1N2, ..., M1Nn,
M2N1, M2N2, ..., M2Nn,
...
MmN1, MmN2, ..., MmNn]
perf layout: (N / 4, M, 4)
[N1M1, N2M1, ..., N4M1,
N5M2, N6M2, ..., N8M2,
...
N(n-3)Mm, N(n-2)Mm, ..., NnMm]
*/
int rknn_matmul_set_io_mem(rknn_matmul_ctx ctx, rknn_tensor_mem* mem, rknn_matmul_tensor_attr* attr);
/* rknn_matmul_set_core_mask
set rknn core mask.(only support rk3588 in current)
RKNN_NPU_CORE_AUTO: auto mode, default value
RKNN_NPU_CORE_0: core 0 mode
RKNN_NPU_CORE_1: core 1 mode
RKNN_NPU_CORE_2: core 2 mode
RKNN_NPU_CORE_0_1: combine core 0/1 mode
RKNN_NPU_CORE_0_1_2: combine core 0/1/2 mode
input:
rknn_matmul_ctx context the handle of context.
rknn_core_mask core_mask the core mask.
return:
int error code.
*/
int rknn_matmul_set_core_mask(rknn_matmul_ctx context, rknn_core_mask core_mask);
/* rknn_matmul_run
run the matmul in blocking mode
params:
rknn_matmul_ctx ctx the handle of context.
return:
int error code.
*/
int rknn_matmul_run(rknn_matmul_ctx ctx);
/* rknn_matmul_destroy
destroy the matmul context
params:
rknn_matmul_ctx ctx the handle of context.
return:
int error code.
*/
int rknn_matmul_destroy(rknn_matmul_ctx ctx);
#ifdef __cplusplus
} // extern "C"
#endif
#endif // _RKNN_MATMUL_API_H