类型推断

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) 在验证程序和 Shape 函数中维护验证码

两者:

  • verifiers:通过 Op::verify() 实现,以及
  • 形状函数:由 Op::inferReturnTypes()Op::inferReturnTypeComponentsInferTypeOpInterface 实现

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

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

解决方案如下:

  1. 对于没有区域的大多数操作(例如 PadOp):尝试将所有验证码都放入形状函数中,完全舍弃验证程序。如果由于无法推断返回值类型(例如使用 ReshapeOpBroadcastInDimOp)而导致无法做到这一点,请创建一个验证程序以包含必要的验证逻辑。通常,可推断的操作(如 AddOp)可能仍需要验证程序来执行其他验证,因为它们正在验证所提供的返回类型的约束条件,而该类型在类型/形状推断方法中无法访问。

  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 的单独测试文件。