英语 | 简体中文

AutoKernel

简介

神经网络如今已被广泛应用于各种场景。在不同设备上高效地执行神经网络对于这些应用至关重要。然而，面对深度学习算法的快速演进，能够针对不同硬件平台手写优化底层内核的专业程序员却十分有限。因此，利用自动优化工具生成高性能实现方案，成为一种极具前景的解决方案。

AutoKernel 最初是 OPEN AI LAB 的一个研究项目，现已开源。AutoKernel 是一款算子优化工具，用于为多种硬件后端自动生成高性能的底层代码。其目标是在包括专用加速器在内的各类硬件上加速高性能算子的开发。

AutoKernel 架构

AutoKernel 由三个模块组成：

• 算子生成器：

  该模块使用开源项目 Halide。Halide 是一种嵌入式 C++ 的领域特定语言（DSL），旨在简化现代机器上高性能图像处理代码的编写。Halide 将算法描述与调度分离。该模块的输入是算子的算法描述，输出则是针对相应后端编译优化后的汇编代码或目标文件。

• 自动搜索

  自动搜索模块利用多种优化算法（贪心算法、强化学习、机器学习等）自动为 Halide 算子搜索最优调度。它支持在 CPU 和 GPU 上进行优化调度搜索，并生成可在不同平台（x86 或 ARM）上运行的代码文件。该模块目前仍在开发中。

• AutoKernel 插件：

  AutoKernel 插件实现了将自动生成的优化算子代码一键集成到 Tengine 中，而无需修改 Tengine 的核心代码库。通过此插件，可以实现自动生成算子实现的一键部署。

特性

• 自动化
• 高效
• 用户友好

Docker

我们提供了安装了 Halide 和 Tengine 的以下 Docker 镜像：

• CPU：openailab/autokernel
• CUDA：openailab/autokernel:cuda
• OpenCL：openailab/autokernel:opencl

详细的 Dockerfile 可参见 Dockerfiles。

[注意]： 如果使用 CUDA 镜像，需使用 nvidia-docker 而不是 docker。以下是 nvidia-docker 安装指南：

nvidia-docker pull openailab/autokernel:cuda
nvidia-docker run -it openailab/autokernel:cuda /bin/bash

许可证

• Apache 2.0

讨论

• Github Issues
• QQ 群：829565581

AutoKernel 快速上手指南

AutoKernel 是一款由 OPEN AI LAB 开源的算子优化工具，旨在通过自动化方式生成针对不同硬件后端（如 CPU、GPU、专用加速器）的高性能底层代码。它基于 Halide DSL 分离算法与调度，并支持一键集成到 Tengine 推理框架中。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04)
• 硬件架构: x86_64 或 ARM64
• GPU 支持 (可选): NVIDIA GPU (需安装对应驱动) 或 OpenCL 设备

前置依赖

为了简化环境配置，强烈推荐使用官方提供的 Docker 镜像，其中已预装 Halide、Tengine 及相关编译工具链。

• Docker: 版本 19.03+
• NVIDIA Container Toolkit: 若使用 CUDA 镜像，必须安装 nvidia-docker (参考 安装指南)

安装步骤

方式一：使用 Docker（推荐）

根据目标硬件平台选择对应的镜像：

1. CPU 版本

   docker pull openailab/autokernel
   docker run -it openailab/autokernel /bin/bash
   

2. CUDA (GPU) 版本 注意：需使用 nvidia-docker 运行

   nvidia-docker pull openailab/autokernel:cuda
   nvidia-docker run -it openailab/autokernel:cuda /bin/bash
   

3. OpenCL 版本

   docker pull openailab/autokernel:opencl
   docker run -it openailab/autokernel:opencl /bin/bash
   

方式二：源码编译（进阶）

若需在本地环境构建，请克隆仓库并参考 Dockerfiles 中的依赖安装步骤手动配置 Halide 和 Tengine 环境。

git clone https://github.com/OAID/AutoKernel.git
cd AutoKernel
# 后续需手动编译 Halide 及配置环境变量

基本使用

AutoKernel 的核心工作流是通过定义 Halide 算法描述，自动搜索最优调度策略并生成优化后的算子代码。

1. 编写算子描述

在容器内，创建一个 Halide 算法文件（例如 my_operator.cpp），定义输入输出及计算逻辑。Halide 将算法描述与调度策略分离。

2. 生成优化代码

利用 AutoKernel 的 Operator Generator 模块编译该描述。以下是一个简化的调用示例（具体命令依实际项目结构可能略有不同）：

# 假设当前目录包含 my_operator.cpp
# 使用 g++ 链接 Halide 运行时进行编译
g++ -std=c++17 -I$HALIDE_INCLUDE my_operator.cpp -L$HALIDE_LIB -lHalide -o my_operator_gen

# 运行生成器，指定目标后端 (如 x86-64-linux-avx2)
./my_operator_gen -o optimized_kernel.o -target=x86-64-linux-avx2

3. 集成到 Tengine

使用 AutoKernel Plugin 可实现无需修改 Tengine 核心代码的一键部署。将生成的 .o 文件或动态库放置在 Tengine 插件加载路径下，并在初始化时启用插件即可自动加载优化算子。

// Tengine C++ 伪代码示例
#include "tengine/c_api.h"

int main() {
    init_tengine();
    
    // 加载 AutoKernel 插件
    load_plugin("libautokernel_plugin.so"); 
    
    // 创建图并运行，自动匹配优化后的算子实现
    graph_t graph = create_graph(nullptr, "tmfile", "model.tmfile");
    run_graph(graph, 1);
    
    release_graph_tensor(graph);
    destroy_tengine();
    return 0;
}

提示：详细的算子定义语法、自动搜索模块（AutoSearch）的使用方法及插件配置参数，请参阅 中文文档。