021-伍华林.pdf

1、伍华林8/21/24Terapines Confidential1兆松科技（武汉）有限公司面向RISC-V异构AI芯片的“大编译器”设计和实现目录C O N T E N T S8/21/24Terapines Confidential2动机02“大编译器”设计“大编译器”实现总结010304动机 软件1.0时代-软件2.0时代 新模型，AI加速器层不穷 模型 x 框架 x 加速器：指数级增 算库维护成本，CUDA兼容问题8/21/24Terapines Confidential3动机 LLVM对加速器的持成熟且容易 LLVM擅做标量优化 MLIR提供了多层抽象，适合做LLVM不擅长的优化 ML

2、IR后端直接成LLVM IR方言8/21/24Terapines Confidential4“大编译器”设计 传统编译器：针对某款特定指令 集的芯片的编译器 编译器 对控制器和加速器的计算和任务程序，以及两者之间的通信进行抽象 自动生成异构计算芯片的所有代码（包括驱动，加速器代码）异构计算系统案例-OpenCL Platform Model8/21/24Terapines Confidential5“大编译器”设计8/21/24Terapines Confidential6“大编译器”设计 MLIR前端 基于IREE/MLIR的些核心方言 持Triton/Mojo*编程语言 支持多种模型文件导

3、入 支持Dataflow芯片计算图生成 统一的算子和计算图抽象层 支持RTL计算图到Dataflow计算图映射*未来支持8/21/24Terapines Confidential7“大编译器”设计 LLVM后端 无缝衔接MLIR 支持多种硬件平台 处理标量相关的优化 指令选择，寄存器分配，指令 排序等 快速适配各类RISC-V DSAGlobalISelLLVM IRLLVM IR DialectMachineIRMC IR8/21/24Terapines Confidential8“大编译器”实现 模型导入 基于IREE的模型导入 支持导入PyTorch,JAX/Tensorflow,ONN

4、X Dialect:TorchIR,StableHLO,TOSA 统一下降到LinalgExt,Arith,Flow etc8/21/24Terapines Confidential9“大编译器”实现 算子库 Triton 算子库-MLIR Dialect TTIR 纯算法，跨平台的开源算子库 MK(Magic Kernel)抽象不同硬件的算子库 通过MK映射到各硬件平台的算子库 自动生成Kernel runtime8/21/24Terapines Confidential10“大编译器”实现 运行时抽象 和平台无关的运行时抽象 自动生成控制器代码 自动生成平台相关算子 算子库编写者无需考虑硬

5、件平台API 硬件平台API通过编译器隐藏和自动生成8/21/24Terapines Confidential11“大编译器”实现 计算图 Flow AI模型计算图抽象 Stream 硬件流水线抽象 MG(Magic Graph)通计算图抽象 图切割算法 动态计算时间模型 数据通信模型 全局变量重排 缩减def-use链8/21/24Terapines Confidential12Example-简化的 GEMV in Triton#Y=A X,where A is a matrix of M x N,X is a vector of N.triton.jitdef gemv_kernel(Y

6、,A,X,M,N,stride_am,BLOCK_SIZE_M:tl.constexpr,BLOCK_SIZE_N:tl.constexpr):start_m=tl.program_id(0)rm=start_m*BLOCK_SIZE_M+tl.arange(0,BLOCK_SIZE_M)rn=tl.arange(0,BLOCK_SIZE_N)A=A+(rm:,None*stride_am+rnNone,:)X=X+rnacc=tl.zeros(BLOCK_SIZE_M,),