Écrire des tests unitaires pour les passes HLO

Il existe différentes manières d'écrire des tests unitaires pour les passes HLO. Cette page décrit la méthode recommandée pour assurer la cohérence et la lisibilité.

FileCheck avec CHECK lignes entrelacées

La plupart des passes HLO peuvent être testées à l'aide de tests FileCheck. Entrelacez les lignes CHECK dans les textes des modules HLO d'entrée et assurez-vous d'utiliser // CHECK au lieu de ; CHECK de manière uniforme comme délimiteur FileCheck.

Par exemple, vous pouvez réécrire fusion cc_test pour un pass priotity_fusion comme suit :

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

Dans la mesure du possible, utilisez le runner LIT et hlo-opt, et placez les lignes CHECK en local à côté de l'IR d'entrée à laquelle elles correspondent. Encore une fois, veillez à utiliser // CHECK au lieu de ; CHECK comme délimiteur.

Par exemple, certains tests de GPU peuvent être écrits comme suit :

// 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 de génération automatique de CHECK

Écrire manuellement des vérifications de test peut être beaucoup de travail. Il est donc souvent plus pratique d'exécuter un optimiseur, de lire les résultats pour s'assurer qu'ils correspondent aux attentes, puis de convertir le HLO optimisé en directives CHECK. Pour simplifier ce processus, vous pouvez utiliser generate_hlo_test_checks.py pour insérer automatiquement les directives CHECK générées au-dessus de chaque cas de test dans un fichier HLO.

(Ne pas) parcourir le graphique

Évitez d'écrire des tests qui parcourent les nœuds feuilles du graphique de résultats et correspondent à l'opération attendue. Ces tests sont fastidieux à écrire, difficiles à lire rapidement et encore plus difficiles à déboguer et à corriger. Utilisez plutôt l'une des options ci-dessus.