為 HLO 傳遞編寫單元測試

您可以透過不同方式為 HLO 傳遞編寫單元測試。本頁說明確保一致性和可讀性的偏好方法。

FileCheck,共 CHECK 行交錯

大多數 HLO 傳遞作業都可以使用 FileCheck 測試。在輸入 HLO 模組文字中交錯使用 CHECK 行,並確保統一使用 // CHECK 而非 ; CHECK 做為 FileCheck 分隔符。

舉例來說,您可以重新編寫 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,在 HLO 檔案中每個測試案例上方自動插入產生的 CHECK 指令。

(不要) 圖形遍歷

請勿編寫會遍歷結果圖葉節點並與預期作業相符的測試。這類測試編寫起來很繁瑣,難以快速解讀,而且更難以偵錯及修正。請改用上述任一選項。