扩展硬件后端
背景知识
Tengine 在设计上将可扩展性作为第一优先级纳入考量,较早的版本注册机制依赖 GCC GNU 扩展,而 GNU 扩展并不是标准 C 的内容。当社区呼唤需要扩展支持到 Microsoft Visual Studio 上时,遇到了较多的困难。
在决定重新设计后,注册模块的易用性有了很大的提升。新的机制通过 CMake 额外的处理过程,取得类似遍历和注册的效果,完成模块的注册。具体的设计和改进可以参考架构详解中的重要模块介绍。
Tengine 在设计上将所有可以运行 CNN 的硬件单元均视为设备,CPU 就是一个典型的设备,在所有的编译选项里,CPU 设备都是默认包含的。如果描述一个新设备并注册,通常意义上这潜在上意味着要求编译的 Tengine 支持异构设备切图(相关内容可以阅读混合设备部分);如果注册的设备也描述了混合精度的接口,那么设备还支持混合精度。
Tengine 通过一个嵌套的结构体完成一个设备的描述:
/*!
* @struct nn_device_t
* @brief Abstract neural network runnable device description struct
*/
typedef struct device
{
const char* name;
struct interface* interface; //!< device scheduler operation interface
struct allocator* allocator; //!< device allocation operation interface
struct optimizer* optimizer; //!< device optimizer operation interface
struct scheduler* scheduler; //!< device scheduler
void* privacy; //!< device privacy data
} ir_device_t;
从结构体 ir_device_t 上可以看出,设计上将一个设备(device)分成 6 部分,第一部分 name 描述了设备的名字,设备名字不允许重复;interface 描述了设备接口;allocator描述了设备相关子图的操作;optimizer 描述了切图和混合精度的接口;scheduler 描述了设备独特的调度接口。
以上接口通常不需要全部填充,Tengine 提供一组丰富的示例指导如何自定义并添加用户自己的设备。
添加自定义设备
创建目录,编写 CMakeLists 文件
首先在source/device创建一个以用户设备命名的文件夹,文件夹可以是用户的设备缩写或其他用户认为比较酷的名字(这里假设起名为TPU,那么目录就是source/device/tpu),并从其他已经实现的 device/xxx 目录中复制一份 CMakeLists.txt 文件到当前文件夹;现在只需要对此 CMakeLists.txt 做些微的修改,而不需要从头创建。我们以从 source/device/acl/CMakeLists.txt 复制一份为例进行说明。该文件完整示例如下:
# 0. clear var
UNSET (_DEV_ACL_HEADER_PATH)
UNSET (_ACL_BASE_SOURCE)
UNSET (_ACL_OPS_SOURCE)
UNSET (_DEV_ACL_DEVICE_SOURCE)
UNSET (_DEV_ACL_COMPILER_DEFINES)
UNSET (_DEV_ACL_COMPILER_OPTIONS)
UNSET (_DEV_ACL_LINKER_OPTIONS)
UNSET (_DEV_ACL_LINK_LIBRARIES)
# 1. set source root path
SET(_ACL_ROOT ${CMAKE_SOURCE_DIR}/source/device/acl)
# 2. add header file path
LIST (APPEND _DEV_ACL_HEADER_PATH ${_ACL_ROOT})
LIST (APPEND _DEV_ACL_HEADER_PATH ${CMAKE_SOURCE_DIR}/3rdparty/acl/include)
# 3. add linking lib searching path
LIST (APPEND _DEV_ACL_LINK_PATH ${CMAKE_SOURCE_DIR}/3rdparty/acl/lib)
# 4. add source files
AUX_SOURCE_DIRECTORY("${_ACL_ROOT}" _ACL_BASE_SOURCE)
AUX_SOURCE_DIRECTORY("${_ACL_ROOT}/op" _ACL_OPS_SOURCE)
LIST (APPEND _DEV_ACL_DEVICE_SOURCE ${_ACL_BASE_SOURCE})
LIST (APPEND _DEV_ACL_DEVICE_SOURCE ${_ACL_OPS_SOURCE})
# 5. add build options for cpu device
# 5.1 is a gcc or clang like compiler
IF (TENGINE_COMPILER_GCC OR TENGINE_COMPILER_CLANG)
IF (TENGINE_COMPILER_GCC AND (${CMAKE_CXX_COMPILER_VERSION} VERSION_GREATER_EQUAL "6.1"))
LIST (APPEND _DEV_ACL_COMPILER_OPTIONS -Wno-ignored-attributes)
ENDIF()
ENDIF()
# 5.2 is Microsoft Visual C++
IF (TENGINE_COMPILER_MSVC)
ENDIF()
# 6. add link options
# 7. add link libs
LIST (APPEND _DEV_ACL_LINK_LIBRARIES arm_compute)
LIST (APPEND _DEV_ACL_LINK_LIBRARIES arm_compute_core)
# 8. set all to cmake cache
SET (TENGINE_ACL_HEADER_PATH ${_DEV_ACL_HEADER_PATH} CACHE INTERNAL "Tengine Arm Compute Library device header files searching path" FORCE)
SET (TENGINE_ACL_LINK_PATH ${_DEV_ACL_LINK_PATH} CACHE INTERNAL "Tengine Arm Compute Library device link libraries searching path" FORCE)
SET (TENGINE_ACL_DEVICE_SOURCE ${_DEV_ACL_DEVICE_SOURCE} CACHE INTERNAL "Tengine Arm Compute Library device main source files" FORCE)
SET (TENGINE_ACL_COMPILER_DEFINES ${_DEV_ACL_COMPILER_DEFINES} CACHE INTERNAL "Tengine Arm Compute Library about compiler defines" FORCE)
SET (TENGINE_ACL_COMPILER_OPTIONS ${_DEV_ACL_COMPILER_OPTIONS} CACHE INTERNAL "Tengine Arm Compute Library about compiler options" FORCE)
SET (TENGINE_ACL_LINKER_OPTIONS ${_DEV_ACL_LINKER_OPTIONS} CACHE INTERNAL "Tengine Arm Compute Library about linker options" FORCE)
SET (TENGINE_ACL_LINK_LIBRARIES ${_DEV_ACL_LINK_LIBRARIES} CACHE INTERNAL "Tengine Arm Compute Library about link libraries" FORCE)
# 9. install device option
INSTALL (FILES ${_ACL_ROOT}/acl_define.h DESTINATION include/tengine RENAME acl_device.h)
首先需要将使用的 CMake 变量的前缀进行修改,以避免潜在的变量冲突;将所有的 ACL 替换为TPU;然后修改模块的搜索根路径 _TPU_ROOT 为 source/device/tpu。
# 1. set source root path
SET(_TPU_ROOT ${CMAKE_SOURCE_DIR}/source/device/tpu)
自定义设备常常需要一些额外的3rdparty依赖,在 _DEV_TPU_HEADER_PATH 和 _DEV_TPU_LINK_PATH 中进行相应的修改;在 ACL 中,增加了 ACL 预编译库路径 ${CMAKE_SOURCE_DIR}/3rdparty/acl/lib。
# 2. add header file path
LIST (APPEND _DEV_TPU_HEADER_PATH ${_TPU_ROOT})
LIST (APPEND _DEV_TPU_HEADER_PATH ${CMAKE_SOURCE_DIR}/3rdparty/tpu/include)
# 3. add linking lib searching path
LIST (APPEND _DEV_TPU_LINK_PATH ${CMAKE_SOURCE_DIR}/3rdparty/tpu/lib)
源码搜集部分按实际情况修改即可。
# 4. add source files
AUX_SOURCE_DIRECTORY("${_TPU_ROOT}" _TPU_BASE_SOURCE)
AUX_SOURCE_DIRECTORY("${_TPU_ROOT}/op" _TPU_OPS_SOURCE)
LIST (APPEND _DEV_TPU_DEVICE_SOURCE ${_TPU_BASE_SOURCE})
LIST (APPEND _DEV_TPU_DEVICE_SOURCE ${_TPU_OPS_SOURCE})
接下来的部分是编译相关的选项,根据实际情况修改即可。Tengine 默认打开了 C/C++ 支持,并尝试打开标准到 C99/C++14,如果工具链不支持会降级为 C98/C++11;如果用户的代码有其他特殊要求可以根据情况调整 _DEV_TPU_COMPILER_DEFINES,_DEV_TPU_COMPILER_OPTIONS,_DEV_TPU_LINKER_OPTIONS 这 3 个变量。
# 5. add build options for cpu device
# 5.1 is a gcc or clang like compiler
IF (TENGINE_COMPILER_GCC OR TENGINE_COMPILER_CLANG)
ENDIF()
# 5.2 is Microsoft Visual C++
IF (TENGINE_COMPILER_MSVC)
ENDIF()
# 6. add link options
根据实际情况调整链接库情况,修改 _DEV_TPU_LINK_LIBRARIES 变量。
# 7. add link libs
LIST (APPEND _DEV_TPU_LINK_LIBRARIES tpu_runtime)
汇总一下临时变量到模块接口变量,接口变量设计为 cache 的,以便跨模块进行传递(另一方面这也是不同 device 不应重名的原因)。
SET (TENGINE_TPU_HEADER_PATH ${_DEV_TPU_HEADER_PATH} CACHE INTERNAL "Tengine TPU device header files searching path" FORCE)
SET (TENGINE_TPU_LINK_PATH ${_DEV_TPU_LINK_PATH} CACHE INTERNAL "Tengine TPU device link libraries searching path" FORCE)
SET (TENGINE_TPU_DEVICE_SOURCE ${_DEV_TPU_DEVICE_SOURCE} CACHE INTERNAL "Tengine TPU device main source files" FORCE)
SET (TENGINE_TPU_COMPILER_DEFINES ${_DEV_TPU_COMPILER_DEFINES} CACHE INTERNAL "Tengine TPU about compiler defines" FORCE)
SET (TENGINE_TPU_COMPILER_OPTIONS ${_DEV_TPU_COMPILER_OPTIONS} CACHE INTERNAL "Tengine TPU about compiler options" FORCE)
SET (TENGINE_TPU_LINKER_OPTIONS ${_DEV_TPU_LINKER_OPTIONS} CACHE INTERNAL "Tengine TPU about linker options" FORCE)
SET (TENGINE_TPU_LINK_LIBRARIES ${_DEV_TPU_LINK_LIBRARIES} CACHE INTERNAL "Tengine TPU about link libraries" FORCE)
如果设备有特殊选项,可以考虑将其插入到 install 阶段。
# 9. install device option
INSTALL (FILES ${_TPU_ROOT}/tpu_define.h DESTINATION include/tengine RENAME tpu_device.h)
在根目录下的 CMakeLists.txt 中添加 option 以便编译时条件打开。
OPTION (TENGINE_ENABLE_TPU "With Awesome TPU support" OFF)
还需要修改 source/device/CMakeLists.txt 添加 Option 相关的处理。
# Awesome TPU
IF (TENGINE_ENABLE_TPU)
ADD_SUBDIRECTORY (tpu)
LIST (APPEND _TENGINE_DEVICE_HEADER_PATH ${TENGINE_TPU_HEADER_PATH})
LIST (APPEND _TENGINE_DEVICE_LINK_PATH ${TENGINE_TPU_LINK_PATH})
LIST (APPEND _TENGINE_DEVICE_COMPILER_DEFINES ${TENGINE_TPU_COMPILER_DEFINES})
LIST (APPEND _TENGINE_DEVICE_COMPILER_OPTIONS ${TENGINE_TPU_COMPILER_OPTIONS})
LIST (APPEND _TENGINE_DEVICE_LINKER_OPTIONS ${TENGINE_TPU_LINKER_OPTIONS})
LIST (APPEND _TENGINE_DEVICE_LINK_LIBRARIES ${TENGINE_TPU_LINK_LIBRARIES})
LIST (APPEND _TENGINE_DEVICE_SOURCE ${TENGINE_TPU_DEVICE_SOURCE})
LIST (APPEND _REGISTER_DEVICE_LIST "${CMAKE_SOURCE_DIR}/source/device/tpu/tpu_device.cc")
ENDIF()
其中,_REGISTER_DEVICE_LIST 是设备注册的核心文件,需要根据实际情况进行填写。
填充结构体完成设备的注册
从某种意义上说,完成一个新设备的注册,只需要填充 ir_device_t 结构体,所有的其他代码工作都是围绕这个核心展开的。
/*!
* @struct nn_device_t
* @brief Abstract neural network runnable device description struct
*/
typedef struct device
{
const char* name;
struct interface* interface; //!< device scheduler operation interface
struct allocator* allocator; //!< device allocation operation interface
struct optimizer* optimizer; //!< device optimizer operation interface
struct scheduler* scheduler; //!< device scheduler
void* privacy; //!< device privacy data
} ir_device_t;
回顾 ir_device_t 结构体,struct interface 结构体描述了基本 API 接口:
/*!
* @struct ir_interface_t
* @brief Abstract neural network runnable device interface struct
*/
typedef struct interface
{
//!< interface of init this neural network device
int (*init)(struct device* device);
//!< interface of prepare runnable subgraph on device
int (*pre_run)(struct device* device, struct subgraph* subgraph, void* options);
//!< interface of run runnable subgraph on device
int (*run)(struct device* device, struct subgraph* subgraph);
//!< interface of post run runnable subgraph on device
int (*post_run)(struct device* device, struct subgraph* subgraph);
//!< interface of async run runnable subgraph on device
int (*async_run)(struct device* device, struct subgraph* subgraph);
//!< interface of async wait runnable subgraph on device
int (*async_wait)(struct device* device, struct subgraph* subgraph, int try_wait);
//!< interface of release runnable subgraph on device
int (*release_graph)(struct device* device, void* device_graph);
//!< interface of release this neural network device
int (*release_device)(struct device* device);
} ir_interface_t;
参考 ACL,一个可能的 TPU 的实现填充如下:
static struct interface tpu_interface = {
.init = tpu_dev_init,
.pre_run = tpu_dev_prerun,
.run = tpu_dev_run,
.post_run = tpu_dev_postrun,
.async_run = nullptr,
.async_wait = nullptr,
.release_graph = nullptr,
.release_device = tpu_dev_release,
};
tpu_dev_init() 是设备的全局初始化函数,注册设备时调用一次,反注册调用 release_device()。这个函数一般用来预申请设备内存作全局缓存,与设备驱动互操作初始化一些寄存器等。
tpu_dev_prerun() 是网络预处理部分,常见的处理包含申请 tensor 内存、转换数据 layout、创建设备运行图、编译设备 kernel 等。这部分申请的空间等需要在 tpu_release_graph() 中进行清理。
tpu_post_run() 与 tpu_release_graph() 可能会引发混淆,tpu_post_run() 常常用来只是清除运行一次的相关状态,与 tpu_dev_prerun() 相反,真正的释放工作可以放到 tpu_release_graph() 中进行。一个可能的场景是,运行一次分辨率的模型 tpu_dev_prerun() 后,换一个分辨率前运行 tpu_post_run(),然后再运行 tpu_dev_prerun()。当需要真正销毁时,运行 tpu_release_graph()。
ir_device_t 结构体中,struct interface 结构体描述了基本 API 接口, struct allocator 描述了设备能力上报接口和评估和调度的接口,struct optimizer 描述了切图和优化相关的接口,struct scheduler 描述了调度相关的接口。这几个接口的核心是 struct scheduler,设备并不总假设实现一个 struct scheduler,如果设备的这个接口描述是 nullptr,那么引擎会使用默认注册的 sync scheduler 运行网络,详情参考 source/scheduler/scheduler.c 中的 static ir_scheduler_t sync_scheduler。用户也可以实现一份自己的 struct scheduler 来完成特殊的任务;结合 struct allocator 和 struct optimizer 可以产生丰富的可能。下面的描述是假设用户不实现 struct scheduler 的情况下的逻辑。
static struct allocator tpu_allocator = {
.describe = tpu_describe,
.evaluation = tpu_evaluation,
.allocate = tpu_allocate,
.release = tpu_release,
};
在 tpu_allocator 中,tpu_describe() 上报模型的 OP 支持情况和精度支持情况,这里的 OP 和精度支持的描述并不会随网络变化而改变,潜在的含义是这种状态下总是假设用户特定设备是 OP 或精度 全场景支持的。以卷积为例,这意味着用户的设备支持所有模式的卷积,无论 pad、stride 、h 、w 、c 情况如何。如果设备实现确实需要在运行时评估,那么可以自定 struct scheduler 完成自定义过程。
tpu_evaluation() 用来运行前评估一下已经实现的设备子图是否可运行;这在需要编译 kernel 时特别有用。
tpu_allocate() 用来支持设备存储池的相关内容,在默认 scheduler 下无需填充这个入口。tpu_release() 是相反的释放过程。
static struct optimizer tpu_optimizer = {
.split_graph = tpu_split_graph,
.optimize_graph = nullptr,
};
在 tpu_optimizer 结构体中,tpu_split_graph() 用来实现切图,tpu_optimize_graph() 用来实现混合精度,其中 tpu_split_graph() 可以调用默认实现的 split_graph_node_to_sub_graph() 进行普通切图;如果有特殊需求可以结合其他结构体的不同字段形成组合。
最后,需要编写注册函数和反注册函数 int register_tpu_device() 和 int unregister_tpu_device(),需要注意的是注册函数和反注册函数的后半段就是文件名,需要和实际文件名匹配,CMake 会自动的完成注册函数的调用过程的链接。
总结
通过上文的描述,可以知道添加一个自定义设备的核心工作就是填充 ir_device_t 结构体,描述完成后,设备注册的所有工作就完成了。模块化的 device 使得 Tengine 非常易于扩展,并有足够的灵活性。
彩蛋
init_tengine(void) 函数中,当 operator prototype 完成注册后,注册的就是 serializer 和 devices,但在静态代码状态下函数并不会跳转,用户可以安装一款集成开发环境,比如 Microsoft Visual Studio 或 JetBrains CLion,打开文件夹后生成 CMake 过程后即可进行跳转。