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) 在验证程序和形状函数中维护验证码
两者:
- verifiers:通过
Op::verify()
实现,并且 - 形状函数:通过
InferTypeOpInterface
(如Op::inferReturnTypes()
或Op::inferReturnTypeComponents
)实现
可能有验证码用于检查运算数/属性/结果。初始拆分可能如下所示:让验证程序检查运算数/属性,然后让形状函数仅计算推断的结果类型,并检查与实际结果类型的兼容性。但实际上,这种拆分也存在一些问题:
- 自动生成的
build()
函数可以调用形状函数,而无需先调用验证程序。因此,还必须在形状函数中验证相关输入。 - 重复代码:例如,在验证程序中,我们会对运算数进行一些处理,然后验证一些中间结果。然后,在形状函数中,这些中间结果有助于推断最终结果。这些中间结果必须计算两次。
- 维护负担:对操作的验证包含两种不同的方法。
解决方案如下:
对于没有区域的大多数操作(例如
PadOp
):将所有验证码放入形状函数中,然后彻底舍弃验证程序。对于带有区域的操作(例如
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
的单独测试文件。