HLO パスの単体テストを作成する

HLO パスの単体テストを作成する方法はいくつかあります。このページでは、一貫性と読みやすさを確保するための推奨される方法について説明します。

FileCheck（CHECK 本の線が交互に配置されている）

ほとんどの HLO パスは FileCheck テストを使用してテストできます。入力 HLO モジュール テキストに CHECK 行をインターリーブし、FileCheck 区切り文字として ; CHECK ではなく // CHECK を一様に使用するようにします。

たとえば、priotity_fusion パスの fusion cc_test を次のように書き換えることができます。

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 ランナーと 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 を使用して、生成された CHECK ディレクティブを HLO ファイルの各テストケースの上に自動的に挿入します。

（非）グラフ走査

結果グラフのリーフノードをたどって、期待される op と照合するテストは記述しないでください。このようなテストは記述が面倒で、すぐに読み取ることが難しく、デバッグと修正がさらに難しくなります。代わりに、上記のいずれかの方法を使用してください。