Cómo escribir pruebas de unidades para los pases de HLO

Existen diferentes formas de escribir pruebas de unidades para los pases de HLO. En esta página, se describe el método preferido para garantizar la coherencia y la legibilidad.

FileCheck con CHECK líneas intercaladas

La mayoría de los pases de HLO se pueden probar con pruebas de FileCheck. Intercala líneas de CHECK en los textos de los módulos de HLO de entrada y asegúrate de usar // CHECK en lugar de ; CHECK de manera uniforme como delimitador de FileCheck.

Por ejemplo, puedes volver a escribir fusion cc_test para un pase de priotity_fusion de la siguiente manera:

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 corredor y hlo-opt

Cuando sea posible, usa el ejecutor LIT y hlo-opt, y coloca las líneas CHECK de forma local junto a la RI de entrada a la que corresponden. Nuevamente, asegúrate de usar // CHECK en lugar de ; CHECK como delimitador.

Por ejemplo, algunas pruebas de GPU se pueden escribir de la siguiente manera:

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

Secuencia de comandos automatizada para generar CHECK

Escribir verificaciones de prueba de forma manual puede ser mucho trabajo, por lo que suele ser más práctico ejecutar un optimizador, leer los resultados para asegurarse de que coincidan con las expectativas y, luego, convertir el HLO optimizado en directivas CHECK. Para simplificar este proceso, puedes usar generate_hlo_test_checks.py para insertar automáticamente directivas CHECK generadas sobre cada caso de prueba en un archivo de HLO.

Recorrido de grafos (no)

Evita escribir pruebas que recorran los nodos hoja del gráfico de resultados y coincidan con la operación esperada. Estas pruebas son tediosas de escribir, difíciles de leer rápidamente y más difíciles de depurar y corregir. En su lugar, usa una de las opciones anteriores.