为 HLO 传递编写单元测试

有多种方法可以为 HLO 传递编写单元测试。本页介绍了确保一致性和可读性的首选方法。

FileCheck,其中交织了 CHECK

大多数 HLO 传递都可以使用 FileCheck 测试进行测试。在输入 HLO 模块文本中交错使用 CHECK 行,并确保统一使用 // CHECK 而不是 ; CHECK 作为 FileCheck 分隔符。

例如,您可以将 fusion cc_test(针对 priotity_fusion 传递)重写为如下形式:

TEST_F(PriorityFusionTest, FuseBroadcastIntoBitcastConsumers) {
  absl::string_view kHlo = R"(
    HloModule test_module

    // CHECK: ENTRY main
    ENTRY main {
      // CHECK-NEXT: %[[PARAM:.*]] = f32[96]{0} parameter(0)
      param_0 = f32[96]{0} parameter(0)
      broadcast = f32[8,96,128,7]{3,2,1,0} broadcast(param_0), dimensions={1}
      bitcast.6079.2 = f32[8,24,4,128,7]{4,3,2,1,0} bitcast(broadcast)
      // CHECK-NEXT: ROOT %{ {.*} } fusion(%[[PARAM]]) { {.*} }
      ROOT transpose.1990.2 = f32[8,24,128,7,4]{4,3,2,1,0} transpose(bitcast.6079.2), dimensions={0,1,3,4,2}
    }
  )";
  RunAndFilecheckHloRewrite(kHlo, std::move(priority_fusion_));
}

LIT 跑者和 hlo-opt

在可行的情况下,使用 LIT runner 和 hlo-opt,并将 CHECK 行放置在本地,紧邻其对应的输入 IR。同样,请务必使用 // CHECK 而不是 ; CHECK 作为分隔符。

例如,一些 GPU 测试可以按如下方式编写:

// RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../../../backends/gpu/specs/%{GPU}.txtpb | FileCheck --check-prefixes=CHECK-%{PTX} %s

HloModule Test, is_scheduled=true
fused_computation {
  param_0 = f32[100,200]{1,0} parameter(0)
  ROOT b.1 = f32[200,100]{1,0} transpose(f32[100,200]{1,0} param_0), dimensions={1,0}
}
ENTRY main {
  a = f32[100, 200]{1,0} parameter(0)
  // CHECK-PTX:         call void @llvm.nvvm.barrier0
  // CHECK-GCN:         call void @llvm.amdgcn.s.barrier
  ROOT wrapped_b = f32[200,100]{1,0} fusion(f32[100,200]{1,0} a), kind=kInput, calls=fused_computation
}

自动 CHECK 生成脚本

手动编写测试检查可能需要大量工作,因此通常更实用的做法是运行优化器,检查结果以确保其符合预期,然后将优化后的 HLO 转换为 CHECK 指令。为了简化此流程,您可以使用 generate_hlo_test_checks.py 自动在 HLO 文件中每个测试用例上方插入生成的 CHECK 指令。

(不)图形遍历

请勿编写遍历结果图的叶节点并与预期操作匹配的测试。这些测试编写起来很繁琐,难以快速阅读,并且更难调试和修复。请改为使用上述任一选项。