【从零开始学深度学习编译器】十七，MLIR ODS要点总结下篇

最新推荐文章于 2025-03-23 17:09:51 发布

原创

最新推荐文章于 2025-03-23 17:09:51 发布 · 841 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #计算机视觉

前言

这一节在【从零开始学深度学习编译器】十六，MLIR ODS要点总结上篇的基础上补充完整了ODS的要点。约束和属性的定义都是MLIR中相当重要的元素，至于类型的定义个人认为了解即可，等到我们需要自定义类型的时候再仔细研究。最后MLIR的语法比较晦涩，初学者可以借助mlir-tblgen来辅助debug。

在这两篇文章里，我跟着MLIR的ODS规范完整走了一遍并总结了14个要点，对于每一个要点我都在OneFlow MLIR的Op定义中进行了对照，并给出了一些示例代码和位置。希望对读者入门MLIR有帮助。

11. 约束（这个很重要）

约束(Constraint)是表驱动Operation定义中的一个核心概念：Operation验证和图Operation匹配都是基于约束来做的。因此，Operation定义和重写规则都直接涉及写入约束。MLIR在OpBase.td(https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/OpBase.td)中定义了Constraint基类。一个Operation的约束可以覆盖不同的范围，可能是：

仅关注单个属性（例如大于 5 的 32 位整数）
多个操作数和结果（例如，第一个结果的形状必须与第一个操作数（可理解为Tensor）相同）
操作本身固有的。（例如没有副作用，参考Transpose Op消除那个案例）

我们将它们分别称为单实体约束、多实体约束和特征。这里的概念了解下即可，我觉得写新的约束是最重要的。

单体约束。单体约束作用域为单个操作数，属性或结果的约束在实体的声明位置进行指定，如Operation arguments 和 Operation results 中（在【从零开始学深度学习编译器】十六，MLIR ODS要点总结上篇中总结了Operation arguments和Operation results需要注意的知识）。
多实体约束。多实体约束在https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/OpBase.td中被建模为PredOpTrait类（是OpTrait的一个子类）。查看OpBase.td获取完整列表。
特征。特征是Operation的内在属性，例如是否具有副作用、可交换与否、是否是终止符等。这些约束应指定为 Op 类模板参数，如【从零开始学深度学习编译器】十六，MLIR ODS要点总结上篇中第三节的Op的特征和约束(Operation traits and constraints) 所示。特征在https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/OpBase.td中被建模成一个NativeOpTrait类（OpTrait的一个子类）。它们得到支持并将被翻译成相应的 C++ mlir::OpTrait 类。
如何指定新的约束？要写一个新的约束，我们必须为它提供一个谓词并指定一个描述名。使用Pred类建模的谓词是构成约束的核心。约束的谓词通常以嵌套的方式构建，有两种类型的谓词：1.CPred：原始的叶子节点谓词。2.复合谓词：由使用谓词组合器的子谓词组成的谓词（conjunction: And, disjunction: Or, negation: Neg, substitution: SubstLeaves, concatenation: Concat）。CPred 是构成更复杂谓词的基础。它是TableGen 视角下的“原子”谓词，是TableGen 与C++ 之间的“接口”。里面已经是 C++ 代码了，它会被当作不透明的字符串来处理，并带有特殊的占位符来替换。我们可以将任何返回布尔值的 C++ 代码放在 CPred 中，包括计算表达式、调用函数、调用类方法等。

为了帮助与 C++ 环境交互，提供了一些特殊的占位符来引用使用该谓词的上下文中的实体。它们充当封闭环境的“钩子”。这包括 $_builder、$_op 和 $_self：

$_builder会被替换成一个mlir::Builder实例，以便我们可以访问常见的构建方法。
$_op 会被当前的Operation替换，以便我们可以访问当前Operation的信息。
$_self 会被替换为该谓词所附加的实体。例如，BoolAttr 是一个包含 CPred<"$_self.isa<BoolAttr>()"> 的属性约束。那么对于 BoolAttr:$attr，$_self 将被 $attr 替换。对于类型约束，它有点特殊，因为我们希望每个类型定义的约束自然读取，并且我们希望将类型约束直接附加到操作数/结果，$_self 将被操作数/结果的类型替换。例如，对于 F32:$operand 中的 F32，它的 $_self 将被扩展为operand(...).getType()。

例如，要写一个属性 attr 是一个 IntegerAttr，在 C++ 中我们可以调用 attr.isa<IntegerAttr>()来实现。这行代码也可以作为 $_self.isa<IntegerAttr>() 包装在 CPred 中，其中 $_self 作为特殊占位符，在扩展时由当前属性 attr 替换来实现相同的功能（指在Tablegen中）。

对于更复杂的谓词，我们可以将其包装在单个 CPred 中，也可以使用谓词组合器将它们组合起来。例如，要写出属性 attr 是 32 位或 64 位整数的约束，可以将其写为：

And<[
  CPred<"$_self.isa<IntegerAttr>()">,
  Or<[
    CPred<"$_self.cast<IntegerAttr>().getType().isInteger(32)">,
    CPred<"$_self.cast<IntegerAttr>().getType().isInteger(64)">
  ]>
]>

（注意，上面只是用一个熟悉的例子来展示如何使用CPred和谓词组合器来编写复杂的谓词。具体来说，对于整数属性，OpBase.td已经定义了I32Attr和I64Attr。所以我们实际上可以重用它们来编写它 Or<[I32Attr.predicate, I64Attr.predicate]>.)

这里再以OneFlow的一个例子来讲解一下，我们定义了一个IsGPU的约束：

def IsGPU: Constraint<CPred<"$0.getValue().equals(\"gpu\")">, "is GPU device">;

然后OneFlow在Transformer部分做了一个定制优化，就是将Scale和Tril这两个连续的Kernel融合成一个大的Kernel，这样可以省掉一部分内存读写的时间。但这个融合的kernel只在GPU的情况下生效，所以这个时候就需要判断当前计算图检测到的Scale和Tril这两个Operation的device是否是GPU的，就需要这个约束。FusedScaleTrilPattern这个Pass的实现如下，可以看到在最后使用了IsGPU这个约束。

def FusedScaleTrilPattern : Pat<
  (
    OneFlow_TrilOp
    (
      OneFlow_ScalarMulOp
        $x,
        $scale_op_name,
        $scale_trainable,
        $scale_device_tag,
        $scale_device_name,
        $scale_scope_symbol_id,
        $scale_hierarchy,
        $has_int_operand,
        $has_float_operand,
        $int_operand,
        $float_operand
    ),
    $tril_op_name,
    $tril_trainable,
    $tril_device_tag,
    $tril_device_name,
    $tril_scope_symbol_id,
    $tril_hierarchy,
    $diagonal,
    $floating_fill_value,
    $integer_fill_value,
    $is_floating_fill_value
  ),
  (OneFlow_FusedScaleTrilOp $x,
    $tril_op_name,
    $tril_trainable,
    $tril_device_tag,
    $tril_device_name,
    $tril_scope_symbol_id,
    $tril_hierarchy,
    $diagonal,
    $floating_fill_value,
    $integer_fill_value,
    $is_floating_fill_value,
    $float_operand,
    $int_operand,
    $has_float_operand
  ),
  [
    (IsGPU $tril_device_tag),
    (IsGPU $scale_device_tag)
  ]
>;

这个Pass的功能就是检测到连续的Scale+Tril Operation就将这两个Operation融合成一个FusedScaleTril Operation。

如果谓词用 CPred 和谓词组合器一起编写非常复杂，我们也可以将其编写为普通的 C++ 函数，并使用 CPred 作为“调用”函数的一种方式。例如，要验证属性 attr 是否具有某些属性，我们可以编写一个 C++ 函数，如：

bool HasSomeProperty(Attribute attr) {
   
    ... }

然后定义Op如下：

def HasSomeProperty : AttrConstraint<CPred<"HasSomeProperty($_self)">,
                                     "has some property">;

def MyOp : Op<...> {
   
   
  let arguments = (ins
    ...
    HasSomeProperty:$attr
  );
}

至于我们是否应该使用单个 CPred 包装整个表达式、多个带有谓词组合器的 CPreds 或单个 CPred “调用”一个函数来定义谓词，没有明确的标准。使用 CPred 和谓词组合器进行定义是可取的，因为它将更多信息（而不是隐藏 C++ 函数背后的所有逻辑）公开到操作定义规范中，以便它可以潜在地驱动更多的自动生成案例。但它需要一个很好的通用谓词库作为构建块，以避免重复，目前正在研究中。

12. 属性定义（很重要+1）

属性是编译期就知道的Operation的常量。ODS 在 C++ 属性类上提供属性包装器。 MLIR 的核心 IR 库中定义了一些常见的 C++ 属性类（https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Attributes.h）。ODS 允许在 TableGen 中使用这些属性来定义Operation，可能具有更细粒度的约束。比如StrAttr直接映射到StringAttr； F32Attr/F64Attr 要求 FloatAttr 额外具有一定的位宽。 ODS属性被定义为具有存储类型（对应于存储属性的mlir::Attribute类），返回类型（对应于生成的getters帮助函数的C++返回类型）以及在内部存储类型和帮助函数进行互转的方法。

属性装饰器。有一些重要的属性适配器/装饰器/修饰符可以应用于 ODS 属性以指定常见的附加属性，如可选性、默认值等。

DefaultValuedAttr：为一个属性指定默认值。
OptionalAttr：将一个属性指定为可选的。
Confined：Confined作为一种通用机制被提供，以帮助对值类型带来的属性约束进行进一步建模。可以通过Confined将较为原始的约束组合成为复杂约束。举个例子，一个32bit的整型最小值为10，可以被表示为Confined<I32Attr, [IntMinValue<10>]>。还有一些其它例子，比如IntMinValue<N>：指定一个大于等于N的整型属性等等。

枚举属性 。某些属性只能从预定义的enum获取值，例如，比较op的比较类型。为了定义这些属性，ODS 提供了几种机制：StrEnumAttr、IntEnumAttr 和 BitEnumAttr。

StrEnumAttr：每个enum case 都是一个字符串，属性在op中存储为 StringAttr。
IntEnumAttr：每个enum case 都是一个整数，属性在op中存储为 IntegerType。
BitEnumAttr：每个 enum case 都是一个位，属性在 op 中存储为 IntegerAttr。

所有这些 *EnumAttr 属性都需要通过其对应的 *EnumAttrCase 完全指定所有允许的情况。有了这个，ODS 能够生成额外的验证以只接受允许的案例。为了促进 *EnumAttrs 和它们的 C++ 使用者之间的交互，EnumsGen(https://github.com/llvm/llvm-project/blob/main/mlir/tools/mlir-tblgen/EnumsGen.cpp) TableGen 后端可以生成一些常见的实用程序：C++ 枚举类、用于枚举类的 llvm::DenseMapInfo、从/到字符串的转换函数。这是通过 mlir-tblgen 的 -gen-enum-decls 和 -gen-enum-defs 命令行选项控制的。

例如，给定下面的EnumAttr：

def Case15: I32EnumAttrCase<"Case15", 15>;
def Case20: I32EnumAttrCase<"Case20", 20>;

def MyIntEnum: I32EnumAttr<"MyIntEnum", "An example int enum",
                           [Case15, Case20]> {
   
   
  let cppNamespace = "Outer::Inner";
  let stringToSymbolFnName = "ConvertToEnum";
  let symbolToStringFnName = "ConvertToString";
}

以下代码将通过 mlir-tblgen -gen-enum-decls 生成：

namespace Outer {
   
   
namespace Inner {
   
   
// An example int enum
enum class MyIntEnum : uint32_t {
   
   
  Case15 = 15,
  Case20 = 20,
};

llvm::Optional<MyIntEnum> symbolizeMyIntEnum(uint32_t);
llvm::StringRef ConvertToString(MyIntEnum);
llvm::Optional<MyIntEnum> ConvertToEnum(llvm::StringRef);<

最低0.47元/天解锁文章