140 lines
5.7 KiB
C
140 lines
5.7 KiB
C
/****************************************************************************
|
|
*
|
|
* Copyright (c) 2017 - 2023 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_CUSTOM_OP_H
|
|
#define _RKNN_CUSTOM_OP_H
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
#include "rknn_api.h"
|
|
|
|
#include <stdint.h>
|
|
|
|
#define RKNN_CUSTOM_OP_MAX_STR_LEN 64
|
|
#define RKNN_CUSTOM_OP_MAX_VALUE_LEN 32
|
|
#define RKNN_CUSTOM_OP_EXPORT __attribute__((visibility("default")))
|
|
|
|
#ifdef __arm__
|
|
typedef uint32_t rknn_custom_op_interal_context;
|
|
#else
|
|
typedef uint64_t rknn_custom_op_interal_context;
|
|
#endif
|
|
/*
|
|
the backend execution device of custom operator.
|
|
*/
|
|
typedef enum _rknn_target_type
|
|
{
|
|
RKNN_TARGET_TYPE_CPU = 1, /* backend device is cpu */
|
|
RKNN_TARGET_TYPE_GPU = 2, /* backend device is gpu */
|
|
RKNN_TARGET_TYPE_MAX
|
|
} rknn_target_type;
|
|
|
|
typedef struct _rknn_gpu_op_context
|
|
{
|
|
void* cl_context;
|
|
void* cl_command_queue;
|
|
void* cl_kernel;
|
|
|
|
} rknn_gpu_op_context;
|
|
|
|
typedef struct _rknn_custom_op_context
|
|
{
|
|
rknn_target_type target; /* custom op backend target */
|
|
rknn_custom_op_interal_context internal_ctx; /* the context of custom op*/
|
|
rknn_gpu_op_context gpu_ctx; /* the gpu context of custom op */
|
|
void* priv_data; /* the private data managed by user */
|
|
} rknn_custom_op_context;
|
|
|
|
typedef struct _rknn_custom_op_tensor
|
|
{
|
|
rknn_tensor_attr attr; /* the attribute of tensor buffer. */
|
|
rknn_tensor_mem mem; /* the memory information of tensor. */
|
|
} rknn_custom_op_tensor;
|
|
|
|
typedef struct _rknn_custom_op_attr
|
|
{
|
|
char name[RKNN_MAX_NAME_LEN]; /* the name of operator atrributes. */
|
|
rknn_tensor_type dtype; /* the data type of operator attributes, indicate the 'array' type. */
|
|
uint32_t n_elems; /* the number of 'array'. */
|
|
void* data; /* the array pointer of operator attributes, the data type of each element is determined by type. */
|
|
} rknn_custom_op_attr;
|
|
|
|
/*
|
|
the information of custom operator to add to the rknn_context.
|
|
*/
|
|
typedef struct _rknn_custom_op
|
|
{
|
|
uint32_t version; /* custom op version */
|
|
rknn_target_type target; /* custom op backend target */
|
|
char op_type[RKNN_MAX_NAME_LEN]; /* custom op type */
|
|
|
|
char cl_kernel_name[RKNN_MAX_NAME_LEN]; /* the opencl kernel name used by custom op */
|
|
char* cl_kernel_source; /* if cl_source_size > 0, pointer to the cl kernel source string, if cl_source_size = 0,
|
|
filepath to the cl kernel file. */
|
|
uint64_t cl_source_size; /* the size of cl_kernel_source */
|
|
char cl_build_options[RKNN_MAX_NAME_LEN]; /* the options for opencl to build clProgram used by custom op */
|
|
|
|
/**
|
|
* The callback function sets that the users need to code
|
|
*/
|
|
int (*init)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
|
|
rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [optional] custom op kernel init falllback function*/
|
|
int (*prepare)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
|
|
rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [optional] custom op kernel prepare falllback function*/
|
|
int (*compute)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
|
|
rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [required] custom op kernel compute falllback function */
|
|
int (*compute_native)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs,
|
|
rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [optional] custom op kernel compute with native attribute falllback function */
|
|
int (*destroy)(rknn_custom_op_context* op_ctx); /* [optional] custom op kernel compute falllback function */
|
|
|
|
} rknn_custom_op;
|
|
|
|
/**
|
|
* dlopen custom op with so required this function
|
|
*/
|
|
typedef rknn_custom_op* (*get_custom_op_func)();
|
|
|
|
/* rknn_register_custom_ops
|
|
|
|
Register custom operators to rknn_context.
|
|
Steps to use a custom op:
|
|
1. Create a rknn_custom_op structure array and fill in it.
|
|
2. Setup prepare/compute/compute_native/destroy callback function and add them to the
|
|
rknn_custom_op.(compute is required and other function is optional, compute_native is not supported now, set it
|
|
to nullptr)
|
|
3. Call rknn_register_custom_ops to register the op type after rknn_init.
|
|
input:
|
|
rknn_context ctx the handle of context.
|
|
rknn_custom_op* op the custom operator array, each of which contains operator information and calllback function.
|
|
uint32_t custom_op_num the length of rknn_custom_op array.
|
|
return:
|
|
int error code.
|
|
*/
|
|
int rknn_register_custom_ops(rknn_context ctx, rknn_custom_op* op, uint32_t custom_op_num);
|
|
|
|
/* rknn_custom_op_get_op_attr
|
|
|
|
input:
|
|
rknn_custom_op_context* op_ctx the handle of custom op context.
|
|
const char* attr_name the attribute name of operator.
|
|
rknn_custom_op_attr* op_attr the data and information of operator attributes.
|
|
*/
|
|
void rknn_custom_op_get_op_attr(rknn_custom_op_context* op_ctx, const char* attr_name, rknn_custom_op_attr* op_attr);
|
|
|
|
#ifdef __cplusplus
|
|
} // extern "C"
|
|
#endif
|
|
|
|
#endif //_RKNN_CUSTOM_OP_H
|