Writing unit tests for HLO passes

There are different ways to write unit test for HLO passes. This page describes the preferred method to ensure consistency and readability.

FileCheck with CHECK lines interleaved

Most HLO passes can be tested using FileCheck tests. Interleave CHECK lines in input HLO module texts, and make sure to use // CHECK instead of ; CHECK uniformly as the FileCheck delimiter.

For example, you can re-write the fusion cc_test for a priotity_fusion pass as follows:

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 runner and hlo-opt

Where feasible, use LIT runner and hlo-opt, and place CHECK lines locally next to the input IR they correspond to. Again, make sure to use // CHECK instead of ; CHECK as the delimiter.

For example, some GPU tests can be written as follows:

// RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/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
}

Automated CHECK-generation script

Writing test checks manually can be a lot of work, so it's often more practical to run an optimizer, read over the results to make sure they match expectations, and then convert the optimized HLO into CHECK directives. To simplify this process, you can use generate_hlo_test_checks.py to automatically insert generated CHECK directives above each test case in an HLO file.

(Don't) Graph traversal

Refrain from writing tests that travel leaf nodes of the result graph and match with expected op. These tests are tedious to write, difficult to quickly read, and more difficult to debug and fix. Use one of the above options instead.