BladeDISC 开源项目教程

项目介绍

BladeDISC 是一个端到端的动态形状编译器项目，专门为机器学习工作负载设计。它是阿里巴巴 PAI-Blade 的关键组件之一，提供对 TensorFlow 和 PyTorch 工作负载在 GPGPU 和 CPU 后端上的通用、透明且易于使用的性能优化。BladeDISC 原生支持动态形状工作负载，并在性能方面进行了多方面的考虑。

项目快速启动

对于 TensorFlow 用户

以下是启用 BladeDISC 的 TensorFlow 程序示例：

import blade_disc_tf as disc
disc.enable()

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

更多信息请参考 QuickStart for TensorFlow Users。

对于 PyTorch 用户

以下是启用 BladeDISC 的 PyTorch 程序示例：

import torch_blade
from torch import nn

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

module = MyModule()
module.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)

更多信息请参考 QuickStart for PyTorch Users。

应用案例和最佳实践

BladeDISC 在生产环境中显示出高达 6.95 倍的加速比，相比于 PyTorch。与静态优化编译器（如 XLA 和 TensorRT）相比，BladeDISC 显示出可比或更好的性能。特别是在动态形状工作负载中，例如 BERT large 推理在 T4 GPU 上，BladeDISC 显示出 1.75 倍的加速。

典型生态项目

BladeDISC 与多个生态项目紧密结合，包括 TensorFlow、PyTorch 以及 MLIR。这些项目共同构成了一个强大的机器学习优化生态系统，支持从模型开发到部署的全过程优化。

• TensorFlow: BladeDISC 提供了对 TensorFlow 工作负载的优化支持。
• PyTorch: BladeDISC 提供了对 PyTorch 工作负载的优化支持。
• MLIR: BladeDISC 利用 MLIR 进行中间表示的优化和转换。

通过这些生态项目的结合，BladeDISC 能够为各种机器学习工作负载提供高效的性能优化解决方案。