类型推断

StableHLO 最初是从 MHLO 方言引导的,它继承了类型推断的 MHLO 实现。实现进度会在 status.md 中跟踪。

下面建议的准则旨在确保为 StableHLO 操作实现高质量的验证程序和形状函数。

建议

这些方案适用于重新审视现有实现和实现新操作,直到全面覆盖。

(P1) 使用 StableHLO 规范作为可信来源

spec是 StableHLO 操作的所有验证程序和形状函数的可信来源。您需要重新访问每个操作的现有验证程序和形状函数,以便与规范完全保持一致。请注意,规范文档会不断更新。如果没有操作规范,则应改用 XLA 实现(包括 xla/service/shape_inference.ccxla/service/hlo_verifier.cc)作为可信来源。XLA 实现不涵盖无限动态,因此对于无限动态,我们将在动态 RFC 可用之前运用常识。

(P2) 充分利用 ODS

ODS 文件(例如 StablehloOps.td)为每个操作数/属性/结果定义具有特征和类型的操作,并将执行验证。因此,对于已由 ODS 保证的属性,验证器或形状函数中不需要验证码。如果验证码与 ODS 重复,请移除该验证码,因为验证码永远不会被触发。

我们需要针对 ODS 约束条件添加测试吗?请参阅制定测试准则

(P3) 在验证程序和形状函数中维护验证码

两者:

  • verifiers:通过 Op::verify() 实现,并且
  • 形状函数:通过 InferTypeOpInterface(如 Op::inferReturnTypes()Op::inferReturnTypeComponents)实现

可能有验证码用于检查运算数/属性/结果。初始拆分可能如下所示:让验证程序检查运算数/属性,然后让形状函数仅计算推断的结果类型,并检查与实际结果类型的兼容性。但实际上,这种拆分也存在一些问题:

  • 自动生成的 build() 函数可以调用形状函数,而无需先调用验证程序。因此,还必须在形状函数中验证相关输入。
  • 重复代码:例如,在验证程序中,我们会对运算数进行一些处理,然后验证一些中间结果。然后,在形状函数中,这些中间结果有助于推断最终结果。这些中间结果必须计算两次。
  • 维护负担:对操作的验证包含两种不同的方法。

解决方案如下:

  1. 对于没有区域的大多数操作(例如 PadOp):将所有验证码放入形状函数中,然后彻底舍弃验证程序。

  2. 对于带有区域的操作(例如 ReduceOp/IfOp;点击此处可查看完整列表):自动生成的构建器不会将区域作为参数,因此,如果这些构建器涉及类型推断,则系统会使用空区域调用形状函数(请参阅此示例)。

    1. 如果类型推断不需要区域(如 ReduceOp),请将与区域相关的验证逻辑放在验证程序中,而不是放在形状函数中。如果不可避免,请复制某些代码。

    2. 如果类型推断 (IfOp/CaseOp/MapOp) 需要区域,则此外形状函数必须明确验证区域是否不为空,即使 ODS 可能已经保证它存在于操作定义中。

(P4) 制定测试指南

对于 ODS 涵盖的验证,我们是否需要添加/维护测试?

我们没有。测试应侧重于验证程序和形状函数,而对 ODS 的更改需要重新检查此操作。

但要注意缺失部分:例如,如果操作包含特征 SameOperandsAndResultShape,该特征仅检查形状,而不检查元素类型,则对运算数/结果的元素类型的验证仍然需要测试。

我们在哪里对验证程序和类型推断进行测试?

ops_stablehlo.mlir 包含 Ops 的正例,以及每个验证错误(至少)1 项负例测试。此外,它还可以检查推断的返回值类型是否与实际结果类型兼容(并非相同!)。

infer_stablehlo.mlir 使用 hlo_test_infer.get_return_type_components"(%x):... 逐行验证操作的形状函数是否存在,并检查推断的类型是否与预期完全匹配。通常,每个操作一次测试为正例。

需要采取的行动

在实现或重新访问操作的验证程序和/或形状函数时:

  1. 将所有正例和负例放入 ops_stablehlo.mlir 中。

  2. infer_stablehlo.mlir 中添加一项正向测试以测试接口。

  3. (可选)如果操作很复杂并且可能包含大量测试,请考虑在同一文件夹中添加名为 verify_<op_name>.mlirverify_<your_topic>.mlir 的单独测试文件。