Como escrever testes de unidade para transmissões HLO

Há diferentes maneiras de escrever testes de unidade para transmissões de HLO. Nesta página, descrevemos o método preferido para garantir consistência e legibilidade.

FileCheck com CHECK linhas intercaladas

A maioria das transmissões HLO pode ser testada usando testes FileCheck. Intercale as linhas CHECK nos textos do módulo HLO de entrada e use // CHECK em vez de ; CHECK de maneira uniforme como o delimitador FileCheck.

Por exemplo, você pode reescrever o fusion cc_test para uma transmissão priotity_fusion da seguinte maneira:

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 e hlo-opt

Quando possível, use o executor LIT e hlo-opt, e coloque as linhas CHECK localmente ao lado da resposta de entrada correspondente. De novo, use // CHECK em vez de ; CHECK como delimitador.

Por exemplo, alguns testes de GPU podem ser escritos da seguinte maneira:

// 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
}

Script automatizado de geração de CHECK

Escrever verificações de teste manualmente pode dar muito trabalho. Por isso, geralmente é mais prático executar um otimizador, ler os resultados para garantir que eles correspondam às expectativas e converter o HLO otimizado em diretivas CHECK. Para simplificar esse processo, use generate_hlo_test_checks.py para inserir automaticamente as diretivas CHECK geradas acima de cada caso de teste em um arquivo HLO.

(Não) travessia de grafos

Evite escrever testes que percorram nós folha do gráfico de resultados e correspondam à operação esperada. Esses testes são tediosos de escrever, difíceis de ler rapidamente e mais difíceis de depurar e corrigir. Use uma das opções acima.