HLO 패스 단위 테스트 작성

HLO 패스의 단위 테스트를 작성하는 방법은 여러 가지가 있습니다. 이 페이지에서는 일관성과 가독성을 보장하는 기본 방법을 설명합니다.

CHECK개의 선이 인터리브된 FileCheck

대부분의 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 파일의 각 테스트 사례 위에 자동으로 삽입하면 됩니다.

그래프 순회

결과 그래프의 리프 노드를 이동하고 예상 작업과 일치하는 테스트를 작성하지 마세요. 이러한 테스트는 작성하기 지루하고, 빠르게 읽기 어려우며, 디버그하고 수정하기가 더 어렵습니다. 대신 위의 옵션 중 하나를 사용하세요.