Написание модульных тестов для проходов 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 runner и hlo-opt и размещайте строки CHECK локально рядом с входным IR, которому они соответствуют. Опять же, обязательно используйте // CHECK вместо ; CHECK в качестве разделителя.

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

// RUN: hlo-opt %s --platform=gpu --stage=llvm-before-optimizations --xla_gpu_target_config_filename=%S/../../../tools/hlo_opt/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 для автоматической вставки сгенерированных директив CHECK над каждым тестовым примером в файле HLO.

(Не делайте) Обход графа

Воздержитесь от написания тестов, которые перемещаются по конечным узлам графа результатов и соответствуют ожидаемым операциям. Эти тесты утомительно писать, их трудно быстро прочитать, а также сложнее отлаживать и исправлять. Вместо этого используйте один из вышеперечисленных вариантов.