StableHLO 最初是从 MHLO 方言引导的,它继承了类型推断的 MHLO 实现。实现进度会在 status.md 中跟踪。
下面建议的准则旨在确保为 StableHLO 操作实现高质量的验证程序和形状函数。
建议
这些方案适用于重新审视现有实现和实现新操作,直到全面覆盖。
(P1) 使用 StableHLO 规范作为可信来源
spec是 StableHLO 操作所有验证程序和形状函数的可信来源。每个操作的现有验证程序和形状函数都需要重新访问,以便与规范完全保持一致。请注意,规范文档会不断更新。如果没有可用的操作规范,则应改为使用 XLA 实现(包括 xla/service/shape_inference.cc 和 xla/service/hlo_verifier.cc)作为可信来源。XLA 实现不涵盖无限的动态特性,因此对于无限的动态特性,在动态 RFC 可用之前,我们会运用常识。
(P2) 充分利用 ODS
ODS 文件(如 StablehloOps.td)为每个操作数/属性/结果定义具有特征和类型的操作,并将执行验证。因此,对于已由 ODS 保证的属性,验证器或形状函数中不需要验证码。如果与 ODS 重复,请移除验证码,因为验证码永远不会被触发。
是否需要针对 ODS 中的限制条件添加测试?请参阅制定测试准则。
(P3) 在验证程序和 Shape 函数中维护验证码
两者:
- verifiers:通过
Op::verify()
实现,以及 - 形状函数:由
Op::inferReturnTypes()
或Op::inferReturnTypeComponents
等InferTypeOpInterface
实现
可能具有用于检查运算数/属性/结果的验证码。初始拆分可能如下所示:让验证程序检查运算数/属性,然后让形状函数仅计算推断的结果类型,并检查与实际结果类型的兼容性。但实际上,这种拆分存在几个问题:
- 形状函数可以由自动生成的
build()
函数调用,而无需先调用验证程序。因此,还必须在形状函数中验证相关输入。 - 重复的代码:例如,在验证程序中,我们会对运算数进行一些处理,然后验证一些中间结果。然后,在形状函数中,这些中间结果有助于推断最终结果。这些中间结果必须计算两次。
- 维护负担:操作验证包含两种不同的方法。
解决方案如下:
对于没有区域的大多数操作(例如
PadOp
):尝试将所有验证码都放入形状函数中,完全舍弃验证程序。如果由于无法推断返回值类型(例如使用ReshapeOp
或BroadcastInDimOp
)而导致无法做到这一点,请创建一个验证程序以包含必要的验证逻辑。通常,可推断的操作(如AddOp
)可能仍需要验证程序来执行其他验证,因为它们正在验证所提供的返回类型的约束条件,而该类型在类型/形状推断方法中无法访问。对于带有区域的操作(如
ReduceOp/IfOp
;如需查看完整列表,请点击此处):自动生成的构建器不将区域作为参数,因此,如果这些构建器涉及类型推断,系统将使用空区域调用形状函数(请参阅此示例)。如果类型推断不需要区域(如
ReduceOp
),请将与区域相关的验证逻辑放在验证程序中,而不是放在形状函数中。在不可避免的情况下,重复某些代码。如果类型推理 (
IfOp/CaseOp/MapOp
) 需要区域,则形状函数还必须明确验证区域是否不为空,即使 ODS 可能已经保证其存在于操作定义中。
(P4) 制定测试指南
对于 ODS 涵盖的验证,我们是否需要添加/维护测试?
我们没有。测试应侧重于验证器和形状函数,而对 ODS 的更改需要重新执行此操作。
但要注意缺失的部分:例如,如果操作包含特征 SameOperandsAndResultShape
(该特征仅检查形状,而不检查元素类型),则对运算数/结果的元素类型的验证仍然需要测试。
我们在哪里对验证程序和类型推断进行测试?
ops_stablehlo.mlir 包含 ops 的正例,以及每个验证错误(至少)1 项负例测试。此外,它还可以检查推断的返回值类型是否与实际结果类型兼容(不是!)。
infer_stablehlo.mlir 用于通过 hlo_test_infer.get_return_type_components"(%x):...
验证操作形状函数是否存在,并检查推断的类型是否与预期完全匹配。通常,每个操作一次测试为正例。
应对措施
在实现或重新访问操作的验证程序和/或形状函数时:
将所有正例和负例都放入 ops_stablehlo.mlir 中。
在 infer_stablehlo.mlir 中添加单个正向测试以测试接口。
(可选)如果操作很复杂并且可能包含大量测试,请考虑在同一文件夹中添加名为
verify_<op_name>.mlir
或verify_<your_topic>.mlir
的单独测试文件。