261 lines
8.0 KiB
C
261 lines
8.0 KiB
C
/****************************************************************************
|
||
*
|
||
* 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
|