MLIR 方言

概述

为了区分不同的硬件和软件目标,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 计算图方言与您的方言之间的降级。这为硬件和编译器制造商铺平了道路。您甚至可以将同一模型中不同级别的方言作为目标;较高级别的优化器将注重 IR 的不熟悉部分,并等待较低级别的优化器对它进行处理。