BladeDISC 简介

我们正在招聘！🔥🔥🔥

我们一直在寻找志同道合的开发者加入开发团队。如果你符合以下任一条件，就是我们寻找已久的人：
* 🥷 如果你是编译器或 AI 爱好者。
* ⭐️ 或者你在 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 项目高度相关。

请访问我们的网站获取更多信息，包括安装教程、开发者指南、演示示例和开发者文档。

功能特性与路线图

前端框架支持矩阵

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

[2] PyTorch 版本 >= 1.6.0 已得到充分验证。

[3] 虽然支持训练，但在训练工作负载的算子覆盖范围上仍有很大改进空间。

后端支持矩阵

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

部署方案

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

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

典型工作负载性能数据

通过使用一组用于生产目的的典型机器学习工作负载评估 BladeDISC，与 PyTorch 相比，BladeDISC 显示出高达 6.95 倍的加速。此外，与静态优化编译器（即 XLA 和 TensorRT）相比，BladeDISC 显示出相当甚至更好的性能。

动态形状工作负载的优势

具体来说，对于 T4 GPU 上的 BERT large 推理（示例中提供），静态编译器优化（XLA）由于其编译开销而显示出严重的性能下降，而 BladeDISC 则显示出 1.75 倍的加速。

API 快速概览

TensorFlow 用户

在原生 TensorFlow 程序中只需添加以下两行代码：

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 推理用例

相关论文

• Zhen Zheng, Zaifeng Pan, Dalin Wang, Kai Zhu, Wenyi Zhao, Tianyou Guo, Xiafei Qiu, Minmin Sun, Junjie Bai, Feng Zhang, Xiaoyong Du, Jidong Zhai, Wei Lin.
  BladeDISC: Optimizing Dynamic Shape Machine Learning Workloads via Compiler Approach. (SIGMOD'24)

• Zhen Zheng, Xuanda Yang, Pengzhan Zhao, Guoping Long, Kai Zhu, Feiwen Zhu, Wenyi Zhao, Xiaoyong Liu, Jun Yang, Jidong Zhai, Shuaiwen Leon Song, Wei Lin.
  AStitch: Enabling a New Multi-dimensional Optimization Space for Memory-Intensive ML Training and Inference on Modern SIMT Architectures. (ASPLOS'22)

开发者教程与文档

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

演讲与报告

• 通过基于编译器的方案优化动态形状 AI 工作负载的性能实践
• 2022/07/31 BladeDISC：动态 Shape 深度学习编译器实践（中文）
• 2022/07/07 在 Torch-MLIR 社区关于 BladeDISC 和 Torch-MLIR 路线图的演讲
• GTC22-S41073，通过云服务的 AI 编译器提供通用且透明的 AI 优化解决方案
• GTC22-S41395，通过 PAI-Blade 实现更易用、更鲁棒的 TensorRT
• 2023/2/17 BladeDISC 介绍（面向 CPU 厂商）
• 2023/3/10 BladeDISC 中基于 Transform Dialect 的代码生成

如何贡献

• 为 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 工作负载。BladeDISC 开发团队正在与社区合作，为 Torch-MLIR 添加 Torch 到 Mhlo 的转换，特别是完全动态形状的特性。
参见 RFC：https://github.com/llvm/torch-mlir/issues/999。
我们呼吁有兴趣的社区开发者加入。

联系我们

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

• 钉钉支持与讨论群：