BladeDISC 简介

我们正在招聘！🔥🔥🔥

我们始终在寻找加入开发团队的人才。如果你符合以下任何一条，那么你就是我们一直在寻找的人：

• 🥷 如果你是编译器或人工智能领域的爱好者。
• ⭐️ 或者如果你在 CPU 和 GPU 上有优化经验。
• ⚙️ 或者如果你想构建一个统一且自动化的编译器来优化推理和训练工作负载。
• 🤿 或者如果你已经在生产或研究项目中使用 BladeDISC，并希望更深入地了解它。
• ✄ 或者你想在 AIGC 时代构建前沿的基础设施。

请通过页面底部的邮箱或钉钉联系我们。⬇️⬇️⬇️

• 最新动态
• 概述
  • 功能与路线图
    • 前端框架支持矩阵
    • 后端支持矩阵
    • 部署方案
  • 典型工作负载的性能数据
    • 在动态形状工作负载中的优势
• API 快速概览
  • 面向 TensorFlow 用户
  • 面向 PyTorch 用户
• 安装与示例
• 论文发表
• 开发者教程与文档
• 演讲与报告
• 如何贡献
• 构建状态
• 常见问题解答
  • 与 mlir-hlo 项目的路线图对比
  • 与 Torch-MLIR 项目的路线图对比
• 联系我们

最新动态

• [🔥 2023.03.17] BladeDISC v0.4.0：大规模性能与功能更新
• [2022.12.08] BladeDISC v0.3.0： 宣布支持 PyTorch 2.0 编译

概述

BladeDISC 是一个面向机器学习工作负载的端到端 动态形状编译器 项目，它是阿里巴巴 PAI-Blade 的关键组件之一。BladeDISC 为 GPGPU 和 CPU 后端上的 TensorFlow/PyTorch 工作负载提供通用、透明且易于使用的性能优化。该架构原生支持动态形状工作负载，并在静态和动态形状场景的性能上都进行了充分考虑。此外，它还支持多种灵活的部署方案，包括 TensorFlow/PyTorch 运行时内的插件模式，以及用于 AOT 独立执行的独立模式。该项目基于 MLIR 架构，并与 mlir-hlo 项目高度相关。

更多信息，请访问我们的官网 https://alibaba.github.io/BladeDISC/ ，其中包括安装教程、开发者指南、演示示例以及面向开发者的文档。

功能与路线图

前端框架支持矩阵

	TensorFlow [1]	PyTorch [2]
推理	是	是
训练	是 [3]	正在进行

[1] TensorFlow 1.12、1.15、2.4 和 2.5 已被支持并经过全面验证。对于其他版本，可能需要进行一些轻微的适配工作。

[2] PyTorch 版本 ≥ 1.6.0 已被完全验证。

[3] 尽管支持训练工作负载，但在算子覆盖方面仍有很大的改进空间。

后端支持矩阵

	状态
Nvidia GPU	是 [1]
AMD GPU	是
Hygon DCU	是
X86	是
AArch64	是

[1] 自 2022 年 8 月起，官方已不再支持 CUDA 11.0 以下的版本。

部署方案

• 插件模式 - BladeDISC 作为 TensorFlow 或 PyTorch 的插件运行。只有受支持的算子会被聚合并编译，而不支持的算子则会由原始的 TensorFlow 或 PyTorch 运行时执行。我们推荐大多数用户使用此模式，因为它透明且易于使用。

• 独立模式 - 在独立模式下，输入的工作负载会被编译成一个可以独立执行的二进制文件，即不依赖于 TensorFlow 或 PyTorch 运行时。在这种模式下，所有算子都必须被支持。

兰州典型工作负载的性能数据

通过对一组用于生产的典型机器学习工作负载进行评估，BladeDISC 相较于 PyTorch 最高可实现 6.95 倍的加速。此外，与静态优化编译器（如 XLA 和 TensorRT）相比，BladeDISC 表现出相当甚至更好的性能。

动态形状工作负载的优势

具体来说，在 T4 GPU 上进行 BERT large 推理时，我们在 示例 中展示了，静态编译器优化（XLA）由于其编译开销导致了严重的性能下降，而 BladeDISC 则实现了 1.75 倍的加速。

TensorFlow	XLA	BladeDISC
1.78 s	41.69s	1.02s
1X		1.75X

API 快速概览

面向 TensorFlow 用户

在原生 TensorFlow 程序中，只需添加两行代码即可启用 BladeDISC：

import numpy as np
import tensorflow as tf

## 在 TensorFlow 程序中启用 BladeDISC
import blade_disc_tf as disc
disc.enable()

## 构建 TensorFlow 图并运行
g = tf.Graph()
with g.as_default():
    ...
    with tf.session as sess:
        sess.run(...)

更多信息，请参阅 TensorFlow 用户快速入门。

面向 PyTorch 用户

PyTorch 用户只需几行代码即可启用 BladeDISC：

import torch_blade
# 构建 PyTorch 模块
class MyModule(nn.Module):
    ...

module = MyModule().eval()

with torch.no_grad():
    # blade_module 是由 BladeDISC 优化后的模块
    blade_module = torch_blade.optimize(module, allow_tracing=True, model_inputs=(x, y))

# 运行优化后的模块
blade_module(x, y)

torch_blade.optimize 接受一个 nn.Module 对象，并返回优化后的模块。更多信息，请参阅 PyTorch 用户快速入门。

安装与示例

• 如何从源码安装与构建
• TensorFlow 推理与训练用例
• PyTorch 推理用例

出版物

• 郑振, 潘在峰, 王大林, 朱凯, 赵文义, 郭天佑, 邱夏飞, 孙敏敏, 白俊杰, 张峰, 杜晓勇, 翟继东, 林伟。 BladeDISC：通过编译器方法优化动态形状机器学习工作负载。（SIGMOD'24）

• 郑振, 杨轩达, 赵鹏展, 龙国平, 朱凯, 朱飞文, 赵文义, 刘晓勇, 杨军, 翟继东, 宋帅文, 林伟。 AStitch：在现代SIMT架构上为内存密集型ML训练和推理开启新的多维优化空间。（ASPLOS'22）

针对开发者的教程与文档

• 教程：BladeDISC Pass流水线详解
• 运行时抽象层简介
• TorchBlade概览
• 教程：如何添加一个新的Torch算子

演示与演讲

• 基于编译器方法的动态形状AI工作负载性能优化实践
• 2022年7月31日 BladeDISC：动态形状深度学习编译器的实践（中文）
• 2022年7月7日 BladeDISC与Torch-MLIR路线图——Torch-MLIR社区演讲
• GTC22-S41073，利用云服务中的AI编译器实现通用且透明的AI优化解决方案
• GTC22-S41395，通过PAI-Blade使TensorRT更易用、更稳健
• 2023年2月17日 BladeDISC介绍（面向CPU厂商）
• 2023年3月10日 BladeDISC中基于Transform方言的代码生成

如何贡献

• 为BladeDISC做贡献

构建状态

框架	设备	状态
PyTorch Pre	GPU
PyTorch Pre	CPU
PyTorch2.0.0	GPU
PyTorch2.0.0	CPU
PyTorch2.0.0	Yitian
PyTorch1.13.0	GPU
PyTorch1.13.0	CPU
PyTorch1.13.0	Yitian
TensorFlow2.5	GPU
TensorFlow2.5	CPU
TensorFlow2.8	Yitian

常见问题解答

与mlir-hlo项目的路线图

BladeDISC与mlir-hlo项目保持着密切的合作关系。其中一部分基础组件，包括MHLO操作定义、TF到MHLO的转换以及一些通用的Pass，已经被合并到mlir-hlo仓库中。在未来较长的一段时间内，我们将继续与mlir-hlo项目保持紧密的合作关系。

与Torch-MLIR项目的路线图

BladeDISC基于Torch-MLIR来编译PyTorch工作loads。BladeDISC开发团队正与社区合作，计划将Torch到MHLO的转换功能加入到Torch-MLIR中，尤其是针对完全动态形状的支持。详情请参阅RFC：https://github.com/llvm/torch-mlir/issues/999。我们诚挚邀请对此感兴趣的社区开发者加入合作。

联系我们

• 邮件组：bladedisc-dev@list.alibaba-inc.com

• 支持与讨论的钉钉群：

BladeDISC 快速上手指南

BladeDISC 是阿里巴巴开源的端到端动态形状编译器，旨在为 TensorFlow 和 PyTorch 工作负载在 GPU 和 CPU 后端提供通用、透明且易用的性能优化。它原生支持动态形状场景，并显著优于静态编译器（如 XLA）。

1. 环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04/20.04)
• 硬件支持:
  • NVIDIA GPU (CUDA >= 11.0)
  • AMD GPU / Hygon DCU
  • x86_64 或 AArch64 CPU
• 软件依赖:
  • Python: 3.6+
  • 框架版本:
    • TensorFlow: 1.12, 1.15, 2.4, 2.5 (其他版本可能需要适配)
    • PyTorch: >= 1.6.0 (已全面验证)，支持 PyTorch 2.0
  • 编译工具: GCC, CMake, Git, Bazel (针对 TensorFlow 集成)

提示: 由于涉及底层编译和特定算子库，强烈建议使用 Docker 容器或参考官方提供的预构建镜像以避免环境配置冲突。

2. 安装步骤

BladeDISC 支持从源码构建，也提供部分预编译包。以下是通用的源码构建流程（以 PyTorch 为例，TensorFlow 类似但需额外配置 Bazel）：

步骤一：克隆代码库

git clone --recursive https://github.com/alibaba/BladeDISC.git
cd BladeDISC

步骤二：安装依赖

根据您使用的框架安装对应的 Python 依赖。

# 示例：安装 PyTorch 相关依赖 (请根据您的实际 CUDA 版本调整 torch 安装命令)
pip install -r requirements.txt

步骤三：构建与安装

BladeDISC 提供了便捷的构建脚本。

对于 PyTorch 用户:

# 构建并安装 torch_blade
python setup.py install
# 或者开发模式安装
python setup.py develop

对于 TensorFlow 用户: 需要确保 Bazel 环境配置正确，然后运行：

bazel build //tensorflow/compiler/mlir/disc/...
# 具体安装步骤请参考 docs/build_from_source.md 中关于 TF 的部分

注意: 完整构建过程可能需要较长时间。如果遇到网络问题，建议配置国内镜像源（如清华源、阿里源）加速 pip 和 git 依赖下载。

3. 基本使用

BladeDISC 的设计目标是“透明”和“易用”，通常只需极少代码改动即可启用优化。

场景一：PyTorch 用户

只需引入 torch_blade 并使用 optimize 接口包裹您的模型。

import torch
import torch.nn as nn
import torch_blade

# 1. 定义您的 PyTorch 模型
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.linear = nn.Linear(128, 128)

    def forward(self, x):
        return self.linear(x)

model = MyModule().eval()

# 2. 准备输入数据 (用于追踪模型结构)
x = torch.randn(1, 128).cuda()
y = torch.randn(1, 128).cuda()

# 3. 启用 BladeDISC 优化
# blade_module 是经过编译优化的新模块
with torch.no_grad():
    blade_module = torch_blade.optimize(model, allow_tracing=True, model_inputs=(x, y))

# 4. 运行优化后的模型
output = blade_module(x, y)

场景二：TensorFlow 用户

在原生 TensorFlow 程序中，仅需两行代码即可启用。

import numpy as np
import tensorflow as tf

# 1. 启用 BladeDISC
import blade_disc_tf as disc
disc.enable()

# 2. 构建并运行 TensorFlow Graph
g = tf.Graph()
with g.as_default():
    # 定义您的模型逻辑
    x = tf.placeholder(tf.float32, shape=[None, 128])
    y = tf.layers.dense(x, 128)
    
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        # 后续运行将自动经过 BladeDISC 优化
        input_data = np.random.rand(32, 128).astype(np.float32)
        result = sess.run(y, feed_dict={x: input_data})

部署模式说明

• 插件模式 (Plugin Mode): 上述示例均为此模式。BladeDISC 作为插件嵌入运行时，仅编译支持的算子簇，不支持的算子回退到原框架执行。推荐大多数用户直接使用此模式。
• 独立模式 (Standalone Mode): 将模型编译为不依赖 TensorFlow/PyTorch 运行时的独立二进制文件，适用于生产环境部署。