כתיבת בדיקות יחידה למעברים של 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 runner ו-hlo-opt

במקרים שבהם זה אפשרי, משתמשים ב-LIT runner וב-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 כדי להוסיף באופן אוטומטי הנחיות CHECK שנוצרו מעל כל מקרה בדיקה בקובץ HLO.

(Don't) Graph traversal

אל תכתבו בדיקות שמתבססות על צמתי עלים בגרף התוצאות ומתאימות לפעולה הצפויה. קשה לכתוב את הבדיקות האלה, קשה לקרוא אותן במהירות וקשה יותר לנפות באגים ולתקן אותן. במקום זאת, אפשר לבחור באחת מהאפשרויות שלמעלה.