HLO graf (High-Level Optimizer graph) je interní reprezentace výpočetního grafu, kterou využívá kompilátor XLA (Accelerated Linear Algebra). Slouží jako klíčový most mezi programovacím jazykem (např. JAX či TensorFlow) a cílovým strojovým kódem pro hardware jako CPU, GPU nebo TPU.

Když v JAXu označíte funkci dekorátorem `jax.jit`, XLA nepracuje přímo s Python kódem. Proces probíhá v několika krocích:

1. Trasa (Tracing): JAX provede funkci s abstraktními hodnotami (tzv. *ShapedArray*), čímž zaznamená posloupnost všech volaných operací lineární algebry.
2. Generování HLO: Z tohoto záznamu se vytvoří strukturovaný výpočetní graf v HLO formátu.
3. Optimalizace a kompilace: HLO graf projde řadou optimalizačních průchodů (passes) XLA kompilátoru a následně se přeloží do binárního kódu.

HLO je navrženo tak, aby umožňovalo nezávislé optimalizace na vysoké úrovni:

• Nezávislost na hardwaru: HLO reprezentuje operace (jako `dot`, `add`, `reshape`, `slice`) abstraktně, bez ohledu na to, zda poběží na jádru CPU, grafické kartě NVIDIA nebo čipu Google TPU.
• Optimizovatelnost (Fúze operací): Kompilátor analyzuje sousední uzly v HLO grafu a hledá příležitosti pro sloučení (fúzi) – např. spojení element-wise operací do jednoho bloku, což šetří paměťovou propustnost.
• Statické tvary (Static Shapes): HLO operuje s pevně danými rozměry polí, což kompilátoru umožňuje přesně alokovat paměť a eliminovat dynamickou režii.

XLA umožňuje exportovat a zobrazit HLO graf v textové podobě. Ukázka jednoduché funkce (`f(x) = x * 2 + 1`) může vypadat přibližně takto:

HloModule jit_f.4
 
ENTRY %jit_f.4 (parameter.1: f32[10]) -> f32[10] {
  %parameter.1 = f32[10] parameter(0)
  %constant.1 = f32[] constant(2)
  %broadcast.1 = f32[10] broadcast(%constant.1), dimensions={}
  %multiply.1 = f32[10] multiply(%parameter.1, %broadcast.1)
  %constant.2 = f32[] constant(1)
  %broadcast.2 = f32[10] broadcast(%constant.2), dimensions={}
  ROOT %add.1 = f32[10] add(%multiply.1, %broadcast.2)
}

Při vývoji nebo optimalizaci výkonu v JAXu či TensorFlow je často nutné HLO graf prozkoumat:

• XLA Dump Flags: Pomocí proměnných prostředí (např. `XLA_FLAGS=„–xla_dump_to=/tmp/xla_dump“`) lze nechat XLA vygenerovat HLO soubory pro každou fázi kompilace.
• Visualization Tools: Vygenerované grafy lze vizualizovat pomocí webového nástroje XLA Visualizer, což pomáhá odhalit neefektivní místa ve výpočtu nebo ověřit, zda proběhla správná fúze operací.

HLO graf představuje srdce optimalizačního procesu XLA. Umožňuje kompilátoru nahlédnout „pod pokličku“ matematických operací, globálně je přeuspořádat a vygenerovat maximálně efektivní strojový kód.

• Architektura a HLO dokumentace XLA
• Oficiální dokumentace JAX