File size: 6,056 Bytes
477da44 |
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 |
/****************************************************************************
*
* 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>
/*
Error code returned by the RKNN Custom Operator API.
*/
#define RKNN_WARNING_SKIP_CUSTOM_OP_COMPUTE -14 /* if custom op init callback funtion return this code and op type is supported by RKNN, it will use RKNN implementation. */
#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
|