Unittests für HLO-Passes schreiben

Es gibt verschiedene Möglichkeiten, Unittests für HLO-Passes zu schreiben. Auf dieser Seite wird die bevorzugte Methode beschrieben, um Konsistenz und Lesbarkeit zu gewährleisten.

FileCheck mit verschachtelten CHECK-Zeilen

Die meisten HLO-Passes können mit FileCheck-Tests getestet werden. Verschachteln Sie CHECK-Zeilen in den Eingabe-HLO-Modultexten und verwenden Sie einheitlich // CHECK anstelle von ; CHECK als FileCheck-Trennzeichen.

Sie können beispielsweise die fusion cc_test für eine priotity_fusion-Karte so umschreiben:

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

Verwenden Sie nach Möglichkeit den LIT-Runner und hlo-opt und platzieren Sie CHECK-Zeilen lokal neben der Eingabe-IR, der sie entsprechen. Verwenden Sie auch hier // CHECK anstelle von ; CHECK als Trennzeichen.

Einige GPU-Tests können beispielsweise so geschrieben werden:

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

Automatisiertes Skript zur CHECK-Generierung

Das manuelle Schreiben von Testprüfungen kann sehr aufwendig sein. Daher ist es oft praktischer, einen Optimierer auszuführen, die Ergebnisse zu prüfen, um sicherzustellen, dass sie den Erwartungen entsprechen, und dann das optimierte HLO in CHECK-Anweisungen zu konvertieren. Um diesen Prozess zu vereinfachen, können Sie generate_hlo_test_checks.py verwenden, um automatisch generierte CHECK-Anweisungen über jedem Testlauf in einer HLO-Datei einzufügen.

(Kein) Knotenbasierter Durchlauf mit Teilindexierung (Graph Traversal)

Schreiben Sie keine Tests, die Blattknoten des Ergebnisgraphen durchlaufen und mit dem erwarteten Vorgang übereinstimmen. Diese Tests sind mühsam zu schreiben, schwer zu lesen und noch schwieriger zu debuggen und zu korrigieren. Verwenden Sie stattdessen eine der oben genannten Optionen.