BladeDISC

BladeDISC 项目简介

概述

BladeDISC 是一个用于机器学习任务的动态形状编译器项目，是阿里巴巴 [PAI-Blade] 的重要组成部分。该项目特别针对 TensorFlow 和 PyTorch 的工作负载提供了一种全面、透明且易于使用的性能优化方案。其架构原生支持动态形状工作负载，并在静态和动态形状场景下具有优异的性能表现。此外，BladeDISC 提供了灵活的部署解决方案，包括插件模式和独立模式，可满足不同用户的需求。

BladeDISC 基于 MLIR 项目进行开发，与 mlir-hlo 项目紧密相关，为 TensorFlow 和 PyTorch 提供了支持，并且支持在 CPU 和 GPGPU 上进行优化。

功能和路线图

前端框架支持矩阵

BladeDISC 支持以下框架：

	TensorFlow	PyTorch
推理	是	是
训练	支持	正在进行中

后端支持矩阵

BladeDISC 支持多种硬件后端，包括：

	支持状况
Nvidia GPU	是
AMD GPU	是
海光 DCU	是
X86	是
AArch64	是

部署解决方案

插件模式：BladeDISC 可作为 TensorFlow 或 PyTorch 的插件使用，仅支持的算子会被聚合和编译，未支持的算子将由原始框架运行。此模式使用便捷，推荐给大多数用户。
独立模式：在独立模式下，输入的工作负载会被编译成一个可独立执行的二进制文件，不依赖于 TensorFlow 或 PyTorch 的运行时。

典型工作负载的数据表现

BladeDISC 在一组典型的机器学习工作负载中表现出相较于 PyTorch 高达 6.95 倍的提速。相比其他静态优化编译器（如 XLA 和 TensorRT），BladeDISC 还展示了可比甚至更优的性能。

动态形状工作负载的优势

在 T4 GPU 上进行大型 BERT 推理时，BladeDISC 提供了高达 1.75 倍的速度提升，而静态编译优化（XLA）由于高编译开销表现不佳。

TensorFlow	XLA	BladeDISC
1.78 秒	41.69 秒	1.02 秒
1 倍		1.75 倍

API 快速查看

对于 TensorFlow 用户

只需简单的两行代码即可在 TensorFlow 程序中启用 BladeDISC：

import blade_disc_tf as disc
disc.enable()

# 构建并运行 TensorFlow 图

对于 PyTorch 用户

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

import torch_blade

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

# 使用 BladeDISC 进行优化
blade_module = torch_blade.optimize(module, allow_tracing=True, model_inputs=(x, y))

# 运行优化后的模块