BladeDISC: 支持动态Shape的深度学习编译器实践

邱侠斐

中文演讲 2022-07-30 16:00 GMT+8  #ai

随着深度学习的不断发展,模型结构在不断演进的同时,底层计算硬件技术更是层出不穷。对于广大开发者来说,不仅要考虑如何在复杂多变的场景下有效的将算力发挥出来,还要应对计算框架的持续迭代。深度编译器是业界在实践中探索出的解决以上问题的技术方向,它让开发者仅需专注于上层模型开发,降低手工优化的人力成本,并进一步压榨硬件性能。

业内已经出现了以TensorFlow XLA、Apache TVM为代表的一批深度学习编译器,但是随着深度学习的发展,也来越来越多的模型呈现出了动态Tensor Shape的特性,例如,1)可以支持多种尺寸图片输入的CNN模型;2)以Bert为代表了NLP模型在BatchSize和SequenceLength维度均有动态性。然而前述编译器最初是面向静态Shape场景设计,需要输入及中间Tensor每个维度上具有固定的尺寸,因此并不能很好地适配动态Shape的场景。面临业务中的强需求,深度学习社区也出现了TVM Relay VM的工作。本演讲主要介绍阿里云PAI团队以BladeDISC为中心,在动态Shape编译器上的工作,主要包括如下内容:

  • BladeDISC的主要架构:为何以及如何基于MLIR框架搭建支持动态Shape的编译器,为什么BladeDISC选择MHLO作为接入层Dialect,以及BladeDISC在为MHLO Dialect做了哪些改造。相较于Apache TVM,在技术路径上有哪些差异以及背后的设计考量。
  • 动态Shape带来的挑战:许多静态Shape语义下比较确定性的问题,例如指令层的向量化,codegen模版选择,是否需要implicit broadcast等等在动态shape场景下都会面临更大的复杂性。部分在编译时期的决策需要移动到运行时进行。
  • 大粒度算子融合:在动态Shape的语义下,如何实现通过低访存开销的硬件(Shared memory in GPU, Memroy Cache in CPU)实现更大尺度的算子融合。以及如何在未知Shape的情况下,推到Tensor之间在Shape上的约束关系进一步实现性能提升。
  • 计算密集型算子:如何利用厂商库、Apache TVM算子生成、手写算子等手段提升计算密集型算子效率,以及非编译生成的算子如何在架构上实现与编译器的统一与互补。
  • 如何支持多种深度学习框架,降低终端用户的使用门槛,以及BladeDISC在阿里云业务中的应用。

Speakers:


邱侠斐: 阿里云计算有限公司, 高级技术专家, 阿里云PAI团队是负责阿里集团内部和外部AI基础设施的团队,模型系统优化一直是团队的重点技术方向之一。编译器作为系统优化的重要手段,经过内部多年的沉淀打磨,目前已经在GitHub开源(https://github.com/alibaba/BladeDISC)。

团队其他分享包括: