Sintesi Tecnica: Verso i Modelli di Mondo per i Compilatori

Definizione del Problema

L'ottimizzazione dei programmi tensoriali è essenziale per i moderni sistemi di machine learning, tuttavia l'identificazione del programma più efficiente è ostacolata da uno spazio di ricerca combinatorio massiccio. Per evitare il costo proibitivo delle misurazioni on-device, gli auto-scheduler esistenti (ad es., Ansor, AutoTVM) si affidano a modelli di costo appresi per classificare i programmi candidati. Tuttavia, questi modelli trattano tipicamente i candidati come istantanee di codice statico (ad es., il testo finale di TensorIR).

Questa prospettiva statica introduce due criticità fondamentali:

1. Perdita del Contesto di Trasformazione: Ignora il processo di trasformazione passo dopo passo che ha generato il programma. Di conseguenza, gli valutatori non possono comprendere come una decisione di scheduling in un passaggio successivo dipenda dal contesto stabilito dalle decisioni precedenti.
2. Sensibilità al Rumore Sintattico: Gli encoder statici sono facilmente fuorviati da variazioni superficiali del codice. Ad esempio, due programmi strutturalmente equivalenti (ad es., moltiplicazione di matrici con dimensioni dei loop $N=1024$ vs $N=1023$) possono produrre testi TensorIR significativamente diversi a causa della gestione dei bordi (loop peeling), causando una classificazione errata del potenziale di performance da parte del modello nonostante l'identico comportamento hardware.

Metodologia

Gli autori propongono un framework di Modello di Mondo per i Compilatori (Compiler World Model) che inquadra la valutazione dello scheduling come dinamiche latenti condizionate dalle azioni su stati di programma. Invece di valutare un prodotto finale di codice statico, il framework simula l'evoluzione dello stato del programma attraverso uno spazio latente continuo.

Il framework consiste di tre componenti apprese istanziate all'interno dell'AutoScheduler di TVM:

1. Rappresentazione dello Stato del Programma Tensoriale (Encoder)

• Input: Testo TensorIR (stati sia pre-scheduling che post-scheduling).
• Architettura: Un'architettura basata su CodeBERT addestinabile e adattata per input lunghi. Utilizza una finestra scorrevole per tokenizzare e codificare chunk sovrapposti, aggregandoli tramite masked mean pooling in una rappresentazione vettoriale densa $z = E(s)$.
• Addestramento: Utilizza l'apprendimento contrastivo per mappare stati TensorIR simili (provenienti dallo stesso workload sorgente) vicini nello spazio latente, allontanando al contempo quelli dissimili. Ciò crea uno spazio di stato in cui l'equivalenza semantica è preservata nonostante le variazioni sintattiche.

2. Transizione di Stato Multi-step Condizionata dalle Azioni

• Obiettivo: Predire la rappresentazione dello stato terminale indotta da una sequenza di azioni di scheduling senza materializzare esplicitamente gli AST intermedi.
• Meccanismo: Partendo da un embedding di stato iniziale $z_0$, il modello esegue il rollout dell'effetto di una sequenza di azioni di scheduling $a_{1:T}$ passo dopo passo.
  • Le azioni vengono serializzate e codificate in embedding densi.
  • Un modello di transizione $F(\cdot)$, implementato tramite un'architettura basata su TransH (originariamente derivata dall'apprendimento di rappresentazioni di grafi di conoscenza), predice il successivo stato latente: $z_t = F(z_{t-1}, a_t)$.
  • Questo modella le azioni di scheduling come trasformazioni geometriche in un iperpiano specifico per la relazione, tracciando efficacementmente la provenienza delle decisioni e risolvendo le ambiguità referenziali (ad es., cosa rappresenta una variabile di loop dopo una fusione).
• Addestramento: Utilizza una perdita di rollout multi-step che combina un termine di corrispondenza dello stato (minimizzando la distanza rispetto ai futuri stati reali) e un termine di ranking basato su margine (assicurando che lo stato predetto sia più vicino alla verità fondamentale rispetto ai campioni negativi).

3. Valutazione dei Candidati Basata sul Ranking

• Input: La rappresentazione dello stato terminale predetta $\hat{z}_T$, feature strutturate delle azioni (ad es., conteggio delle azioni, lunghezza della sequenza) e opzionalmente feature hardware.
• Output: Un punteggio scalare per la priorità relativa dei candidati all'interno dello stesso workload.
• Architettura: Un modello di ranking XGBoost addestrato con un obiettivo LambdaRank.
• Strategia di Addestramento: Il modello non predice il tempo di esecuzione assoluto. Inveve, impara a classificare i candidati all'interno di un gruppo di workload basandosi sui tempi di esecuzione misurati, trattando il problema come un compito di learning-to-rank.

Contributi Chiave

1. Prospettiva del Modello di Mondo per i Compilatori: Il documento adatta il paradigma del modello di mondo all'ottimizzazione tensoriale, stabilendo un framework in cui i candidati di scheduling sono modellati attraverso l'evoluzione dello stato latente condizionata dalle azioni piuttosto che come istantanee statiche.
2. Dataset di Traiettoria Azione-Stato: Costruzione di un dataset basato su TVM/TenSet che organizza i log di tuning in traiettorie azione-stato allineate (stato pre-scheduling, sequenza di azioni, stati intermedi, stato post-scheduling) per la predizione multi-step.
3. Integrazione Efficiente della Ricerca: Implementazione del framework in TVM AutoScheduler, dimostrando che la modellazione delle dinamiche latenti riduce significativamente i costi di autotuning mantenendo un'alta qualità della ricerca.

Risultati Sperimentali

Il metodo è stato valutato su una CPU Intel Xeon Gold 6430 e una GPU NVIDIA RTX 4090 su sette modelli di reti neurali (inclusi ResNet, BERT, GPT e OPT) e 103 sottografi estratti.

• Performance vs. Ansor: Sotto un budget di 64 prove, il metodo migliora la latenza dei sottografi rappresentativi rispetto ad Ansor di 1.37× su GPU e 1.54× su CPU.
• Efficienza del Campionamento: Il metodo raggiunge una qualità di ricerca comparabile a Ansor-10K (10.000 prove) utilizzando solo 1.000 prove (10× meno misurazioni). Nello specifico, supera leggermente Ansor-10K nella media aritmetica e si avvicina ad essa entro il 2.2% nella media geometrica sui sottografi rappresentativi.
• Speedup End-to-End: Rispetto a PyTorch/PyTorch-opt (cuDNN), il metodo fornisce speedup della media geometrica di 4.61× (CPU) e 3.67× (GPU), con picchi di velocità che raggiungono i 58.61×.
• Studi di Ablazione: Rimuovere il modulo di transizione appreso (tornando alla codifica dello stato iniziale o finale statico) ha causato aumenti significativi della latenza (mediana ~56–69%), confermando che la modellazione dell'evoluzione dello stato intermedio è la componente più critica.

Significato e Limitazioni

Significato:
Il documento sostiene che trattare l'ottimizzazione del compilatore come un processo dinamico di transizione di stato, piuttosto che come un problema di valutazione di codice statico, permette agli valutatori di catturare le dipendenze semantiche tra le decisioni di scheduling. Questo approccio sopprime efficacementamente il rumore sintattico e risolve le ambiguità referenziali, portando a una ricerca più efficiente con meno misurazioni hardware.

Limitazioni (come dichiarato dagli autori):

1. Dipendenza dallo Spazio di Ricerca: Il metodo migliora l'valutazione dei candidati, non la loro generazione. La sua efficacia è limitata dalla capacità dello scheduler sottostante di proporre schedule validi; un valutatore migliore non può recuperare schedule che non vengono mai proposti.
2. Accumulo di Errori: Gli errori di predizione possono accumularsi in traiettorie di scheduling lunghe o complesse nello spazio latente, riducendo potenzialmente l'affidabilità in casi difficili.
3. Ranking vs. Predizione: Il modello è addestrato per il ranking all'interno del workload, non per la predizione della latenza assoluta. Ciò lo rende ideale per la prioritizzazione online, ma meno adatto per contesti che richiedono stime di performance calibrate tra diversi workload.