Как писать модульные тесты для проходов HLO

Существуют разные способы написания модульных тестов для проходов HLO. На этой странице описан предпочтительный способ обеспечения единообразия и удобочитаемости.

FileCheck с CHECK чередующимися строками

Большинство проходов HLO можно протестировать с помощью тестов FileCheck. Чередуйте строки CHECK в текстах входных модулей HLO и используйте // CHECK вместо ; CHECK в качестве разделителя FileCheck.

Например, вы можете переписать fusion cc_test для priotity_fusion следующим образом:

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.

Например, некоторые тесты графического процессора можно написать следующим образом:

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

Writing test checks manually can be a lot of work, so it's often more practical to run an optimizer, read over the results to make sure they match expectations, and then convert the optimized HLO into CHECK directives. Чтобы упростить этот процесс, вы можете использовать generate_hlo_test_checks.py для автоматической вставки сгенерированных директив CHECK над каждым тестовым примером в файле HLO.

Обход графа (не рекомендуется)

Не пишите тесты, которые проходят по конечным узлам графа результатов и сопоставляют их с ожидаемыми операциями. Такие тесты сложно писать, читать, отлаживать и исправлять. Используйте один из вариантов выше.