CINN - 神经网络的编译器基础设施 [飞桨]

伊织code2025-01-06 17:34
c 复制代码 
                        ___                    ___          ___     
                       /__\                  /\  \        /\  \   
                      /:/  /       ___        \:\  \       \:\  \   
                     /:/  /       /__\        \:\  \       \:\  \  
                    /:/  /  ___  /:/__/    _____\:\  \  _____\:\  \ 
                   /:/__/  /__\/::\  \   /::::::::__\/::::::::__\
                   \:\  \ /:/  /\/\:\  __\:\~~\~~\/__/\:\~~\~~\/__/
                    \:\  /:/  /    \:\/__\\:\  \       \:\  \      
                     \:\/:/  /      \::/  / \:\  \       \:\  \     
                      \::/  /       /:/  /   \:__\       \:__\   
                       \/__/        \/__/     \/__/        \/__/

文章目录

一、关于 CINN

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

该项目正在积极开发中。

它是如何工作的

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::Program，该程序有助于定义机器学习计算，
  • hlir::framework::Tensor，多维数组有助于管理内存缓冲区。
  • hlir::framework::Program，运行时的最终可执行程序。它包含许多基本的可执行元素。
  • 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

python 复制代码 
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

c 复制代码 
# 打开组合算子
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 打印示例：

c 复制代码 
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 得到的结果

c 复制代码 
// 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) 迭代空间对齐

c 复制代码 
// 所有 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

c 复制代码 
// 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

c 复制代码 
// 例如 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

c 复制代码 
// 为了简洁，此处省略 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，融合外层循环乘积相同的循环，并且保证不破坏图级别依赖（规则负责）和元素级别依赖（原语负责）

c 复制代码 
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 轴

c 复制代码 
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_ 的代码如下：

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

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

c 复制代码 
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 如下：

c 复制代码 
// 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加载桨模型。

python 复制代码 
# 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都是一个桨操作符。

python 复制代码 
# 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实现

c 复制代码 
#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源代码，如：

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（四）

相关推荐
NAGNIP11 小时前
一文搞懂深度学习中的通用逼近定理！
人工智能·算法·面试
冬奇Lab13 小时前
一天一个开源项目（第36篇）：EverMemOS - 跨 LLM 与平台的长时记忆 OS，让 Agent 会记忆更会推理
人工智能·开源·资讯
冬奇Lab13 小时前
OpenClaw 源码深度解析（一）：Gateway——为什么需要一个"中枢"
人工智能·开源·源码阅读
AngelPP16 小时前
OpenClaw 架构深度解析：如何把 AI 助手搬到你的个人设备上
人工智能
宅小年16 小时前
Claude Code 换成了Kimi K2.5后，我再也回不去了
人工智能·ai编程·claude
九狼17 小时前
Flutter URL Scheme 跨平台跳转
人工智能·flutter·github
ZFSS17 小时前
Kimi Chat Completion API 申请及使用
前端·人工智能
天翼云开发者社区18 小时前
春节复工福利就位！天翼云息壤2500万Tokens免费送，全品类大模型一键畅玩！
人工智能·算力服务·息壤
知识浅谈18 小时前
教你如何用 Gemini 将课本图片一键转为精美 PPT
人工智能
Ray Liang19 小时前
被低估的量化版模型，小身材也能干大事
人工智能·ai·ai助手·mindx
热门推荐
01GitHub 镜像站点02【OpenClaw 本地实战 Ep.3】突破瓶颈：强制修改 openclaw.json 解锁 32k 上下文记忆03OpenClaw 使用和管理 MCP 完全指南04Clawdbot部署教程：解决‘gateway token missing’授权问题的完整步骤05OpenClaw + 飞书（Feishu）环境搭建指南06Window 10部署openclaw报错node.exe : npm error code 12807Claude Code + GLM4.7 避坑指南：解决 Unable to connect to Anthropic services08AI 规范驱动开发“三剑客”深度对比：Spec-Kit、Kiro 与 OpenSpec 实战指南09OpenClaw优化飞书API 额度已耗尽问题10AI Agent 平台横评：ZeroClaw vs OpenClaw vs Nanobot