BladeDISC: 阿里巴巴开源的动态形状深度学习编译器

BladeDISC简介

BladeDISC (Blade Dynamic Shape Compiler)是阿里巴巴开源的一个端到端动态形状编译器项目,专门用于优化机器学习工作负载。它是阿里云机器学习平台PAI-Blade的核心组件之一,为TensorFlow和PyTorch等深度学习框架提供通用、透明且易用的性能优化。

BladeDISC的主要特点包括:

• 支持多种前端框架:TensorFlow和PyTorch
• 支持多种后端硬件:CUDA、ROCM和x86 CPU
• 原生支持动态形状语义编译
• 支持推理和训练
• 轻量级API,对用户透明
• 支持插件模式嵌入宿主框架和独立部署模式

设计原理

BladeDISC的设计初衷是解决XLA和TVM等静态形状编译器的局限性,实现完全支持动态形状语义的深度学习编译器。它基于MLIR (Multi-Level Intermediate Representation)基础设施构建,采用了以下关键技术:

1. 动态形状编译

BladeDISC采用host-device联合代码生成的方案,在编译时生成形状计算代码,运行时进行实时形状推导。这使得它能够完整支持动态计算图语义,避免了静态编译器需要指定输入形状或捕获实际输入形状组合的限制。

2. 性能优化

为了在动态形状场景下实现良好的性能,BladeDISC采用了以下策略:

• 区分计算密集和内存访问密集部分,分别采用不同的优化方法
• 将部分优化决策从编译时下沉到运行时,生成多个版本的kernel实现
• 引入形状约束抽象,保留更多优化机会

3. 大粒度算子融合

BladeDISC实现了基于fusion-stitching的大粒度算子融合技术,能够将不同调度的计算子图缝合到同一个kernel中,显著提高融合粒度。

4. 多前端框架支持

BladeDISC通过实现轻量级转换器,将不同前端框架的计算图转换为统一的中间表示(IR),实现了对多种框架的支持。

5. 运行时环境适配

BladeDISC引入了运行时抽象层(RAL),支持在不同运行环境中执行编译结果,如独立运行、TensorFlow子图和PyTorch子图等。

应用场景

BladeDISC在阿里巴巴内部和阿里云上已广泛应用,覆盖了NLP、机器翻译、语音识别、图像检测等多种AI应用场景。它主要有两类典型应用:

1. 作为通用透明的性能优化工具,减少用户部署AI任务的人力负担,提高模型迭代效率。

2. 帮助新硬件在AI场景中实现适配和对接。

在推理场景中,BladeDISC与TensorRT等厂商提供的推理优化工具形成了良好的技术互补,在以下方面具有优势:

• 完整支持动态形状语义
• 基于编译器技术路线,对非标准模型具有更好的性能优势
• 插件形式支持前端框架,部署模式更灵活

在新硬件支持方面,BladeDISC的编译器技术路线对硬件后端具有一定的通用性,能够帮助新兴AI硬件厂商快速实现软件栈的适配,发挥硬件算力。

开源生态

BladeDISC已在GitHub上开源(https://github.com/alibaba/BladeDISC),欢迎社区开发者参与贡献。开源的主要目的是:

1. 回馈社区,分享在AI编译器领域的经验和理解。
2. 获得更多用户在业务场景中的反馈,持续改进和迭代。

未来的roadmap包括:

• 持续提升鲁棒性和性能
• x86后端支持计算密集算子
• GPGPU上基于Stitching的大粒度自动代码生成
• 支持AMD ROCM GPU后端
• 支持PyTorch训练

BladeDISC团队也在探索更多前沿方向,如新硬件架构支持、动态形状下的全局布局优化、稀疏子图优化等,欢迎社区提供反馈和建议,共同推动AI编译器技术的发展。