MLIR中间表示和编译器框架

MLIR中间表示和编译器框架

TensorFlow生态系统包含许多在软件和硬件堆栈的多个级别上运行的编译器和优化器。作为TensorFlow的日常用户,使用不同类型的硬件(GPU,TPU,移动设备)时,此多级堆栈可能会表现为难以理解的编译器和运行时错误。

TensorFlow可以通过多种不同方式运行:

  • 发送到TensorFlow执行程序,该执行程序调用手写的操作内核
  • 转换为XLA高级优化器表示形式(XLA HLO),后者又可以调用用于CPU或GPU的LLVM编译器,或者继续将XLA用于TPU。(或两者的某种组合!)
  • 转换为TensorRTnGraph或另一种针对特定于硬件的指令集的编译器格式
  • 将图形转换为TensorFlow Lite格式,然后在TensorFlow Lite运行时内部执行,或者进一步转换为通过Android Neural Networks API(NNAPI)或相关技术在GPU或DSP上运行。

更复杂的路径,包括每层内的多轮优化,例如Grappler框架可优化TensorFlow中的张量布局和操作。

众多的编译器和表示实现大大提高了性能,但这种异构世界可能会给最终用户带来问题,例如在这些系统之间的边界处产生令人困惑的错误消息。同样,新的硬件和软件堆栈创建,必须为每个新路径重建优化和转换过程。

MLIR,或多级中间层表示。这是表示形式和编译器实用程序的库,位于模型表示和生成特定于硬件的代码的低级编译器/执行器之间。借助MLIR,在生产质量组件的支持下,实现优化编译器设计和实现方面的新颖探索。

MLIR引起许多团体的关注,包括:

  • 优化机器学习模型的性能和内存消耗的编译器研究和实施
  • 硬件制造商寻找一种将其硬件连接到TensorFlow的方法,例如TPU,手机便携式神经硬件以及其它定制ASIC。
  • 编写语言绑定时,利用优化编译器和硬件加速的优势。

什么是MLIR?

MLIR本质上是用于现代优化编译器的灵活基础架构。由一个中间表示(IR)规范和一个用于对该表示执行转换的代码工具包组成。(按照编译器的说法,从较高级别的表示形式转换为较低级别的表示形式时,这些转换可以称为“降低lowerings”。)

MLIR受LLVM的影响很大,重用了它的许多好思想。具有灵活的类型系统,允许在同一编译单元中结合多个抽象级别来表示,分析和转换图形。这些抽象包括TensorFlow操作,嵌套的多面循环区域,LLVM指令,以及固定的硬件操作和类型。

MLIR Dialects

为了分离不同的硬件和软件目标,MLIR提供了“方言”,其中包括:

  • TensorFlow IR,代表TensorFlow中的所有可能事物
  • XLA HLO IR,旨在利用XLA的编译功能(输出到TPU)
  • 实验仿射方言,重点关注多面体表示和优化
  • LLVM IR,与LLVM表示之间具有1:1映射,允许MLIR通过LLVM发出GPU和CPU代码
  • TensorFlow Lite,转换为在移动平台上运行的代码

每个方言由一组已定义的算子组成,这些算子具有不变的内容,例如:“一个二进制运算符,输入和输出具有相同的类型。”
添加到MLIR

MLIR没有固定/内置的全局已知算子列表(没有“内部”)。方言可以定义完全自定义的类型,这就是MLIR如何建模LLVM IR类型系统(具有一流的聚合),对ML优化的加速器(如量化类型)重要的域抽象,甚至是Swift或Clang类型系统(在以后围绕Swift / Clang声明节点构建)。

如果要连接新的低级编译器,则将创建一个新的方言以及TensorFlow Graph方言和方言之间的降低。这为硬件和编译器制造商铺平了道路。可以将同一模型中的不同级别的方言作为目标。较高级别的优化器将尊重IR的不熟悉部分,并等待较低级别的IR处理。

对于编译器研究和框架制造商来说,MLIR允许在各个层次上进行转换,甚至可以在IR中定义自己的算子和抽象-从而可以对要解决的问题领域进行最佳建模。这样,与LLVM相比,MLIR更像是纯编译器基础结构

虽然MLIR充当ML的编译器,使机器学习技术也可以在编译器中使用!这一点特别重要,因为开发数值库的工程师的缩放比例与ML模型或硬件的多样化不一样。MLIR的可扩展性促进了代码降低策略的探索和跨抽象的渐进式降低。

在TensorFlow 2.0中,图可以是隐式的;执行的操作可以单独,成组或作为完整图形(例如Keras顺序)运行。无论如何,这些图或图片段必须进行优化和执行。

 

posted @ 2021-04-21 05:51  吴建明wujianming  阅读(668)  评论(1编辑  收藏  举报