ToB企服应用市场:ToB评测及商务社交产业平台

标题: CINN - 神经网络的编译器根本办法 [飞桨] [打印本页]

作者: 温锦文欧普厨电及净水器总代理 时间: 2025-1-3 06:37
标题: CINN - 神经网络的编译器根本办法 [飞桨]

___ ___ ___
/__\ /\ \ /\ \
/:/ / ___ \:\ \ \:\ \
/:/ / /__\ \:\ \ \:\ \
/:/ / ___ /:/__/ _____\:\ \ _____\:\ \
/:/__/ /__\/::\ \ /::::::::__\/::::::::__\
\:\ \ /:/ /\/\:\ __\:\~~\~~\/__/\:\~~\~~\/__/
\:\ /:/ / \:\/__\\:\ \ \:\ \
\:\/:/ / \::/ / \:\ \ \:\ \
\::/ / /:/ / \:__\ \:__\
\/__/ \/__/ \/__/ \/__/

复制代码

一、关于 CINN

CINN项目是多硬件后端的机器学习编译器和实行器。它旨在提供多层API，使张量计算更容易定义、实行更快，并且更方便地与硬件后端一起扩展。目前，它的目标是x86处理惩罚器和英伟达图形处理惩罚器。
该项目正在积极开发中。

github : https://github.com/PaddlePaddle/CINN | License：Apache 2.0许可证
官方文档：https://paddlepaddle.github.io/CINN/
官方先容：https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/paddle_v3_features/cinn_cn.html
CINN 算子开发教程（一）
https://aistudio.baidu.com/aistudio/projectdetail/5129681
手把手教你为神经网络编译器CINN增加One-Hot算子
https://www.paddlepaddle.org.cn/support/news?action=detail&id=3258

它是怎样工作的

CINN将传统的DNN模型低落为两级中心表示（IR），即高级IR（HLIR）和CINN IR。
HLIR有助于定义一些特定领域的计算并对IR图实行一些整体优化；CINN IR有助于表示一些计算语义，末了低落到硬件后端。
两个级别的IR都具有相似的SSA图、分析和优化工具。
CINN基于多面体编译，因此很容易通过更多的循环优化进行扩展。调度变换应用于从HLIR到CINN IR的低落之间。
整体架构如下

概念简介

深度学习编译器是一种专门为深度学习模型优化和部署而设计的工具，用于提高模型的计算服从、低落内存占用、加速训练推理过程。其功能是将高层次的深度学习模型转换为低层次的、高效的、底层硬件可实行的代码。简朴来说，深度学习编译器在深度学习框架和底层硬件之间充当了“翻译”的脚色，能够将用户定义的神经网络模型形貌转化为底层硬件能够理解和实行的指令。编译器在实现这种转换的过程中，应用了一系列优化技术，以提高模型在各种硬件平台上（如 CPU、GPU）的实行服从。深度学习编译器的主要功能包罗：

模型转换：将高层次的深度学习模型转换为得当目标硬件的中心表示（IR）。
优化：应用各种编译优化技术，如图优化、内存优化、算子融合等，以提高实行服从。
代码天生：天生得当目标硬件的可实行代码。

背景与动机

深度学习模型的训练和推理过程涉及大量的计算，对硬件性能要求很高。飞桨框架虽然提供了高级的编程接口和丰富的算子库，但在实行服从和模型部署方面另有很大的优化空间。利用深度学习编译器的主要动机包罗：

1. 优化性能与资源利用率

深度学习模型每每需要处理惩罚大量的数据和复杂的计算，直接在高层次框架上实行大概无法充分利用底层硬件的本领。深度学习编译器能够深入硬件特性，应用多种优化技术，提高计算服从，低落延迟。并且通过优化模型的计算图和内存利用，深度学习编译器也能够明显低落模型的内存和 IO 资源的斲丧，进而提高计算性能。

2. 硬件多样性支持

不同的硬件平台有不同的特性和优化需求。在现有机制下，新的异构硬件设备接入深度学习框架需要手工实现几百个算子对应的硬件 Kernel 代码，开发的工作量非常大。如果利用深度学习编译器，理论上仅需实现新硬件 IR 层面的对接，以及相应的硬件 IR 优化计谋就能完成与深度学习框架的对接，相比于实现几百个硬件 Kernel，开发的工作量会大幅减少。

3. 提升开发服从

深度学习编译器可以主动化很多优化过程，减少手动调优的工作量。开发者只需关注模型的设计和训练，而不必深入了解底层硬件优化细节，从而提高开发服从。

二、入门

编译并实行代码

请参阅安装指南并遵照指南。

概念

CINN中的API有两个级别，较高级别是HLIR，较低级别是CINN IR，都包罗一些概念。
在HLIR

Primitive Emitter（PE），封装了不同基于张量的算法的计算，
frontend::Interpreter，实行模型（PaddlePaddle）的容器，
frontend:rogram，该步伐有助于定义机器学习计算，
hlir::framework::Tensor，多维数组有助于管理内存缓冲区。
hlir::framework:rogram，运行时的最终可实行步伐。它包罗很多基本的可实行元素。
hlir::framework::Graph，表示模型结构的图。图中的每个节点代表一个运算符（con2d、relu、mul等）。
hlir::framework::GraphCompiler，将模型的图形表示（hlir::框架::Graph）转换为可实行步伐（hlir::框架::步伐）的编译器。

在CINN IR

Compute，定义计算的方法，
Lower，将计算低落到相应IR的方法，
LoweredFunc，CINN IR中定义的函数，
Var，一个标量变量，
Expr，表达式表示任何CINN IR节点（没有指定的语句节点），
Stage，包罗张量的一些时间表细节，

参考API利用情况

阅读测试中的代码
对于Python API，引用python/tests中的代码。
C++API位于cinn/*/*_test.cc，高级API位于hlir/frontend，低级API位于cinn/cinn.h。

三、利用示例：

飞桨框架编译器（CINN, Compiler Infrastructure for Neural Networks）利用时仅需在原先的模型动转静或推理流程下打开编译器相干 FLAGS 即可，无需对模型代码做任何改动。以下是一个利用样例：
示例代码文件：run_net.py

import paddle
from paddle import nn
from paddle.static import InputSpec
# 定义神经网络
class RMSNorm(nn.Layer):
def __init__(self):
super().__init__()
paddle.seed(2024)
self.hidden_size = 768
self.weight = paddle.randn([self.hidden_size], dtype="float32")
self.variance_epsilon = 1e-6
def forward(self, hidden_states):
variance = (hidden_states * hidden_states).sum(-1, keepdim=True) / 768
hidden_states = (
paddle.rsqrt(variance + self.variance_epsilon) * hidden_states
)
return hidden_states * self.weight
def run_net(input_data):
net = RMSNorm()
# 指定输入变量的维度、数据类型等信息，具体接口可参考：
# https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/jit/basic_usage_cn.html#inputspec
input_spec = [
InputSpec(shape=[1, None, 768], dtype='float32'),
]
net = paddle.jit.to_static(
net,
input_spec=input_spec,
full_graph=True,
)
# 使用 eval 模式
net.eval()
# 执行计算图
out = net(input_data)
return out
# 创建输入数据
input_data = paddle.randn([1, 2048, 768], dtype="float32")
# 运行神经网络
out = run_net(input_data)
print(out)

复制代码

脚本实行：run.sh

# 打开组合算子
export FLAGS_prim_enable_dynamic=true && export FLAGS_prim_all=true
# 打开 CINN 编译器相关 FLAG
export FLAGS_use_cinn=true
export FLAGS_cinn_new_group_scheduler=true
export FLAGS_group_schedule_tiling_first=true
export FLAGS_cinn_bucket_compile=true
# 打开 PIR 模式
export FLAGS_enable_pir_api=true
# 是否打印 Program IR 信息
export FLAGS_print_ir=false
python run_net.py

复制代码

上述代码示例中我们创建了一个简朴的rms_norm计算子图，利用飞桨的动转静流程将子图转为静态图并调用编译器 CINN 进行优化和实行。经过性能对比测试，在 A100 GPU 环境中上述子图利用 CINN 可以取得 3 倍左右的性能提升（该性能数据仅供学习参考，在现实应用模型中能够取得的性能提升效果一样寻常会低于该数据）。
注：由于飞桨的编译器仍旧处在快速迭代开发阶段，我们设置了较多 FLAGS 进行分支的选择和调试，因此现阶段在利用 CINN 时需要对如下 FLAGS（FLAGS_prim_enable_dynamic、 FLAGS_cinn_new_group_scheduler、 FLAGS_group_schedule_tiling_first、 FLAGS_cinn_bucket_compile、 FLAGS_enable_pir_api）进行手动设置，待后续相干功能完备后这些 FLAGS 会默认开启，无需再手动设置。

四、设计架构

图 1 CINN 整体架构
飞桨框架编译器（CINN, Compiler Infrastructure for Neural Networks）整体架构如上图所示，大体可以分为三个模块，分别是编译器前端、编译器后端和实行器部分。

1. 编译器前端

一样寻常来说编译器前端需要将不同框架和格式的深度学习模型转换为编译器的内部 IR 并进行图级别的优化，CINN 作为飞桨框架原生编译器，可以直接利用飞桨框架提供的模型加载和中心表示（Paddle IR，简称 PIR）组件，因此 CINN 前端的主要功能是基于 PIR 进行图层级别的优化，并对子图进行分别为后端高性能 Kernel 代码天生提供支持。CINN 前端关键的流程可分为三部分：

a. 组合算子拆分

飞桨框架中将算子分别为根本算子（也称作原子算子，语义上该算子无法更进一步拆分成其他算子。根本算子语义上可以通过重组等价实现组合算子的逻辑）和非根本算子两类大，由于非根本算子数量较多，并且在编译器中较难辨认和处理惩罚，因此我们利用组合算子拆分的方式将非根本算子拆分为等价的根本算子组合，原始计算图经过组合算子拆分后可以大幅提升性能的可优化空间。

b. 图优化 Pass

在计算图层级进行 PIR 的 Pass 优化，常见的图优化 Pass 包罗：常量折叠、死代码消除（DCE）、公共子表达式消除（CSE）、冗余算子消除、算子计算归并等。

c. 算子融合

算子融合是编译器前端非常紧张的一个功能，主要是将多个算子打包到一个子图中（对应为一个 FusionOp），交给编译器后端天生一个高效的硬件相干计算 Kernel。算子融合的本质是通过 IO 优化加速访存麋集算子，如果我们将两个一连 Kernel 归并为一个 Kernel 调用，我们会减少中心变量的读写开销，因此在访存麋集型的 2 个 Op 上，融合可以获取更高的性能。举个例子，如下图：

图 2 算子融合示例
我们有两个算子 Relu 和 Scale，由于两个算子都是 IO 麋集型算子（计算复杂度不高）。正常情况下我们需要读取 A 和 B 一次，写 B 和 C 一次。但是对于融合之后的 Kernel（右图）而言，我们只需要读取 A 和写 C 一次，这样我们通过算子融合可以取得更少的访存次数，在 IO 麋集算子而言，可以极大提高性能。具体的算子融合计谋实现非常复杂，这里不做展开先容，感兴趣的读者可以阅读相干源码 #cinn_group_cluster_pass。

2. 编译器后端

编译器后端主要负责将前端处理惩罚后的 IR 转换为目标硬件可实行的代码或硬件形貌。主要功能包罗基于硬件特性的 IR 优化、高效内存管理和代码天生等。

2.1. CINN AST IR

AST IR 打印示例：

ScheduleBlock(root)
{
serial for (i, 0, 32)
{
serial for (j_0, 0, 64)
{
serial for (j_1, 0, 128)
{
ScheduleBlock(A)
{
vi, vj = axis.bind(i, j_0 * 64 + j_1) // tensor 下标与循环变量的仿射变换
A[vi, vj] = X[vi, vj] * 2
}
}
}
}
}

复制代码

CINN AST IR 中包罗了以下信息，但集合和映射并不显示利用某种数据结构进行存储。
集合：语句实例 & 内存单元**
映射：
访存关系：语句实例 <—> 内存单元
依赖关系：语句实例 <—> 语句实例
实行顺序：语句实例 -----> 语句实例

实行顺序 = 语句实例的先后关系
语句实例集合范围 = 循环界限 + 循环步长 ------ 循环构成一个带约束的整数空间，即迭代空间，迭代空间决定了语句实例，语句实例充满了迭代空间。

2.2. 基于 AST IR 的 Schedule

Schedule 为定义在 CINN AST IR 上的优化计谋，常见的 Schedule 包罗：LoopAlignment, Tile, Inline, Vectorize, Unroll 等。
以一个组合算子为例模仿大概的 AST 变换过程：
[S1, S2, 1024] ==E=> [S1, S2, 1024] ==R=> [S1, S2] ==E=> [S1, S2] ==B=> [S1, S2, 1024] ==E=> [S1, S2, 1024]

(1) LowerToAst 得到的结果

// Elemenwise-1
serial for (i, 0, S1)
serial for (j, 0, S2)
serial for (k, 0, 1024)
ScheduleBlock(A)
vi, vj, vk = axis.bind(i, j, k)
A[vi, vj, vk] = X[vi, vj, vk] * 2
// Elemenwise-2
serial for (i, 0, S1)
serial for (j, 0, S2)
serial for (k, 0, 1024)
ScheduleBlock(B)
vi, vj, vk = axis.bind(i, j, k)
B[vi, vj, vk] = A[vi, vj, vk] + 1
// Reduce-1
serial for (i, 0, S1)
serial for (j, 0, S2)
ScheduleBlock(C__reduce_init)
vi, vj = axis.bind(i, j)
C_init[vi, vj] = 0
serial for (i, 0, S1)
serial for (j, 0, S2)
serial for (k, 0, 1024) // Reduce
ScheduleBlock(C)
vi, vj, vk = axis.bind(i, j, k)
C[vi, vj] = C[vi, vj] + B[vi, vj, vk]
// Elemenwise-3
serial for (i, 0, S1)
serial for (j, 0, S2)
ScheduleBlock(D)
vi, vj = axis.bind(i, j)
D[vi, vj] = C[vi, vj] * 2
// Broadcast-1
serial for (i, 0, S1)
serial for (j, 0, S2)
serial for (k, 0, 1024) // Broadcast
ScheduleBlock(E)
vi, vj, vk = axis.bind(i, j, k)
E[vi, vj, vk] = D[vi, vj]
// Elemenwise-4
serial for (i, 0, S1)
serial for (j, 0, S2)
serial for (k, 0, 1024)
ScheduleBlock(F)
vi, vj, vk = axis.bind(i, j, k)
F[vi, vj, vk] = E[vi, vj, vk] + 1

复制代码

(2) 迭代空间对齐

// 所有 ScheduleBlock 的 loop nest 都变为以下 2 种格式中的一种
// 1
serial for (sp, 0, S1 * S2) // pure_spatial_iter
serial for (rb, 0, 1024) // impure_spatial_iter
ScheduleBlock(XXX)
vsp1, vsp2, vrb = axis.bind(sp / S2, sp % S2, rb)
XXX = XXXXXX
// 2
serial for (sp, 0, S1 * S2) // pure_spatial_iter
ScheduleBlock(XXX)
vsp1, vsp2 = axis.bind(sp / S2, sp % S2)
XXX = XXXXXX

复制代码

(3) Tile: 对全部 ScheduleBlock 的 loop nest 做相同的 Tile

// pure_spatial 轴 Tile 为：-1 * 16 * 64 Tile size 可为参数传入
serial for (sp1, 0, S1 * S2 / 1024)
serial for (sp2, 0, 16)
serial for (sp3, 0, 64) // S1 * S2 / 16 / 64, predicate: sp1 * 1024 + sp2 * 16 + sp3 < S1 * S2
XXXXXX
// impure_spatial_iter 轴 Tile 为 32
serial for (sp1, 0, S1 * S2 / 1024)
serial for (sp2, 0, 16)
serial for (sp3, 0, 64)
serial for (rb1, 0, 32)
serial for (rb2, 0, 32)
ScheduleBlock(XXX)
predicate = sp1 * 1024 + sp2 * 16 + sp3 < S1 * S2
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
vrb = axis.bind(rb1 * 32 + rb2)
XXX = XXXXX

复制代码

(4) ComputeInline

// 例如 ScheduleBlock(A) inline 到 ScheduleBlock(B)
serial for (sp1, 0, S1 * S2 / 1024)
serial for (sp2, 0, 16)
serial for (sp3, 0, 64)
serial for (rb1, 0, 32)
serial for (rb2, 0, 32)
ScheduleBlock(A)
predicate = sp1 * 1024 + sp2 * 16 + sp3 < S1 * S2
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
vrb = axis.bind(rb1 * 32 + rb2)
B[vsp1, vsp2, vrb] = (X[vsp1, vsp2, vrb] * 2) + 1

复制代码

(5) Reduce 优化: two step reduce & 绑定部分 reduce 轴到 cuda

// 为了简洁，此处省略 reduce_init Block 和 predicate
serial for (sp1, 0, S1 * S2 / 1024)
serial for (sp2, 0, 16)
serial for (sp3, 0, 64)
CudaBind[ThreadIdx.x] for (rb1, 0, 32)
serial for (rb2, 0, 32)
ScheduleBlock(C_rf)
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
vrb1 = axis.bind(rb1)
vrb2 = axis.bind(rb2)
C_rf[vsp1, vsp2, vrb1] = C_rf[vsp1, vsp2, vrb1] + B[vsp1, vsp2, vrb1 * 32 + vrb2]
ScheduleBlock(C)
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
vrb1 = axis.bind(rb1)
C[vsp1, vsp2] = C[vsp1, vsp2] + C_rf[vsp1, vsp2, vrb1]

复制代码

(6) 循环融合: ComputeAt && SimpleComputeAt，融合外层循环乘积相同的循环，并且包管不粉碎图级别依赖（规则负责）和元素级别依赖（原语负责）

serial for (sp1, 0, S1 * S2 / 1024)
serial for (sp2, 0, 16)
serial for (sp3, 0, 64)
ScheduleBlock(D)
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
D[vsp1, vsp2] = C[vsp1, vsp2] * 2
serial for (rb1, 0, 32)
serial for (rb2, 0, 32)
ScheduleBlock(E)
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
vrb = axis.bind(rb1 * 32 + rb2)
E[vsp1, vsp2, vrb] = D[vsp1, vsp2]
ScheduleBlock(F)
vsp1 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) / S2)
vsp2 = axis.bind((sp1 * 1024 + sp2 * 16 + sp3) % S2)
vrb = axis.bind(rb1 * 32 + rb2)
F[vsp1, vsp2, vrb] = E[vsp1, vsp2, vrb] + 1

复制代码

(7) Bind Cuda 轴：在第二步中，全部 ScheduleBlock 对应的循环要 bind 到同一 Cuda 轴

serial for (sp1, 0, S1 * S2 / 1024)
CudaBind[BlockIdx.x] for (sp2, 0, 16)
CudaBind[ThreadIdx.y] for (sp3, 0, 64)
CudaBind[ThreadIdx.x] for (rb1, 0, 32)
serial for (rb2, 0, 32)
ScheduleBlock(XXX)

复制代码

2.3. Kernel 代码天生与编译

Codegen 在 CINN IR AST 上做前序遍历，打印出对应硬件的指令，并通过硬件相对应的编译器（如 llvm、nvcc 等）进行编译得到可运行的函数指针，该指针会被封装到 `JitKernelOp`` 中用于后续实行器的解析实行。
a. 以函数定义为例子，cuda kernel func 和 x86 kernel func 的不同的是，cuda kernel func 会在函数名前增加 __global__
针对 x86 硬件，转义 ir::_LoweredFunc_ 的代码如下：

void CodeGenC::Visit(const ir::_LoweredFunc_ *op) {
PrintFunctionDeclaration(op); // 前序遍历继续转义函数名、函数参数等
str_ += "\n";
...
...
}

复制代码

在 NV GPU 上的转义代码如下：

void CodeGenCUDA_Dev::Visit(const ir::_LoweredFunc_ *op) {
str_ += "__global__\n"; // 和 x86 的不同，增加 __global__
PrintFunctionDeclaration(op); // 前序遍历继续转义函数名、函数参数等
str_ += "\n";
...
...
}

复制代码

b. 在动态外形场景下，还会 codegen 出 infer shape function, infer shape function 的 CINN IR 会在 Bucket Lowering 中得到，转义过程复用的 x86 硬件的 codegen。infer shape kernel 如下：

// infer shape 函数名字的组成：kernel_name + "infer_shape"
// 函数参数：
// kernel_args: 指针数组，和 kernel func args 一致
// kernel_args_num: kernel_args 的长度
// tensor_shape_args: 指针数组，存储输出 tensor 的 shape
function fn_exp_0_subtract_0_infer_shape (kernel_args, kernel_args_num, tensor_shape_args)
{
int64 S0 = cinn_get_value_in_cuda_kernel_args(kernel_args, 2)
{
// CINN IR 暂时不支持数据索引的语法，暂时用函数调用实现，下面 2 条语句等价于
// tensor_shape_args[0] = {S0, 256ll};
// 即第 0 个出 tensor 的 shape 为{S0, 256ll};
infer_shape_set_value(0, 0, S0, tensor_shape_args)
infer_shape_set_value(0, 1, 256ll, tensor_shape_args)
}
}

复制代码

3. 实行器

编译器天生的 Kernel 代码需要与深度学习框架实行器完成交互和集成才能最终运行起来，因此需要基于实行器的运行调度接口对编译器天生的 Kernel 进行封装。
接入实行器后在运行时对于经过编译器处理惩罚的子图将实行 CINN 天生的 Kernel, 否则将实行通例的 PHI 算子 K

五、更多示例

让我们以python API为例，相应的C++API可用且差别不大。

加载PaddlePaddle模型

您可以直接利用CINN加载桨模型。

# Load Model to CINN
computation = Computation.compile_paddle_model(
target = DefaultHostTarget(), model_dir = "./ResNet50", input_tensors = ['inputs'], input_sapes = [[1, 3, 224, 224]], params_combined = True)
# Get input tensor and set input data
a_t = computation.get_tensor(input_tensor)
a_t.from_numpy(np.random.random(x_shape).astype("float32"), target)

复制代码

通过NetBuilder构建网络

您可以利用NetBuilder API构建和运行模型。每个NetBuilder API都是一个桨操纵符。

# Define the NetBuilder.
builder = frontend.NetBuilder(name="network")
# Define the input variables of the model
a = builder.create_input(type=common.Float(32), shape=(1, 3, 224, 224), id_hint="A")
b = builder.create_input(type=common.Float(32), shape=(1, 3, 224, 224), id_hint="B")
# Build the model using NetBuilder API
y = builder.add(a, b)
res = builder.relu(y)
# Specify target and generate the computation
target = common.DefaultHostTarget()
computation = Computation.build_and_compile(target, builder)

复制代码

利用CINN低级DSL定义一些计算并实行

以下是利用CINN DSL的简朴matrix-multiplication实现

#include "cinn/cinn.h"
using namespace cinn;
// Declare constants
Expr M(10), N(20), K(30);
// Declare the inputs
auto A = Placeholder<float>("A", {M, K});
auto B = Placeholder<float>("B", {K, N});
auto k1 = Var(K.as_int32(), "k1");
auto C = Compute(
{M, N}, [&](Var i, Var j) { return ReduceSum(A(i, k1) * B(k1, j), {k1}); }, "C");
Target target = common::DefaultHostTarget();
int block_size = 32;
// The stages holds all the schedules for each tensors.
auto stages = CreateStages({C});
// Blocking optimization by loop tiling stragety.
auto [i_outer, i_inner, j_outer, j_inner] = stages[C]->Tile(0, 1, bn, bn);
auto [k_outer, k_inner] = stages[C]->Split("k0", 4);
stages[C]->Reorder({i_outer, j_outer, k_outer, k_inner, i_inner, j_inner});
// Generate C source code:
Module::Builder builder("module_block", target);
auto func = Lower("matmul_block", stages, {A, B, C});
builder.AddFunction(func);
CodeGenCX86 compiler(target, CodeGenCX86::Feature::AVX512);
Outputs outputs;
outputs = outputs.c_header("./test02_matmul_block.h").c_source("./test02_matmul_block.cc");
compiler.Compile(builder.Build(), outputs);

复制代码

这可以天生优化的C源代码，如：

void matmul_block(void* _args, int32_t num_args)
{
const cinn_buffer_t* _A = cinn_pod_value_to_buffer_p(&(((cinn_pod_value_t*)(_args))[0]));
const cinn_buffer_t* _B = cinn_pod_value_to_buffer_p(&(((cinn_pod_value_t*)(_args))[1]));
cinn_buffer_t* _C = cinn_pod_value_to_buffer_p(&(((cinn_pod_value_t*)(_args))[2]));
cinn_buffer_malloc((void*)(0), _C);
const float* A = ((const float*)(_A->memory));
const float* B = ((const float*)(_B->memory));
float* C = ((float*)(_C->memory));
float* C__reduce_init = ((float*)(_C->memory));
for (int32_t i = 0; i < 1024; i += 1) {
for (int32_t j = 0; j < 1024; j += 1) {
C__reduce_init[((1024 * i) + j)] = 0;
};
};
for (int32_t i_outer = 0; i_outer < 32; i_outer += 1) {
for (int32_t j_outer = 0; j_outer < 32; j_outer += 1) {
for (int32_t k0_outer = 0; k0_outer < 256; k0_outer += 1) {
for (int32_t k0_inner = 0; k0_inner < 4; k0_inner += 1) {
for (int32_t i_inner = 0; i_inner < 32; i_inner += 1) {
for (int32_t j_inner = 0; j_inner < 32; j_inner += 1) {
C[((1024 * i_inner) + ((32768 * i_outer) + ((32 * j_outer) + j_inner)))] = (C[((1024 * i_inner) + ((32768 * i_outer) + ((32 * j_outer) + j_inner)))] + (A[((1024 * i_inner) + ((32768 * i_outer) + ((4 * k0_outer) + k0_inner)))] * B[((32 * j_outer) + ((1024 * k0_inner) + ((4096 * k0_outer) + j_inner)))]));
};
};
};
};
};
};
cinn_buffer_free((void*)(0), _C);
}

复制代码

将CodeGenCX86的用法更改为CodeGenLLVM，它将天生一个可以实时调用的LLVM JIT编译函数。

致谢

CINN从以下项目中学到了很多：

Halide：参考了大多数红外节点的设计，
TVM：我们学习了很多想法，包罗一些时间表原语的语义学、TOPI、NNVM等等，
tiramisu：isl用法、多面体编译、调度原语实现等，
TensorFlow/xla：引用了原始操纵的语义学。

伊织 2024-01-02（四）

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。

欢迎光临 ToB企服应用市场:ToB评测及商务社交产业平台 (https://dis.qidao123.com/)