Categoria: Tempo di compilazione: allineamento dell'accesso alla memoria non verificato di Mosaic
Questo errore si verifica quando il compilatore analizza un'operazione di accesso alla memoria (ad esempio
vector.load, vector.store, tpu.load o tpu.store) e non può
dimostrare staticamente che l'indice dinamico utilizzato per una dimensione specifica è un
multiplo della dimensione di tiling richiesta.
Messaggi di errore di esempio:
INTERNAL: Mosaic failed to compile TPU kernel: cannot statically prove that index in dimension 1 is a multiple of 128
at location: ...
The MLIR operation involved:
%14372 = "vector.load"(%14371, %93, %14363) : (memref<4x256xf32, #tpu.memory_space<vmem>>, index, index) -> vector<1x32xf32>
Backend XLA: TPU
Panoramica
Quando il kernel carica o memorizza un vettore, l'indirizzo di memoria
(calcolato dal puntatore di base più l'indice dinamico) deve essere allineato con le
dimensioni del tiling del vettore sull'hardware. Ad esempio, se una dimensione è suddivisa in riquadri
da 128 elementi, l'indice dinamico utilizzato per accedervi deve essere 0, 128, 256,
e così via. Tieni presente che molte operazioni (come i caricamenti e gli archiviamenti di vettori) non hanno questi
requisiti per gli indici statici.
Il compilatore applica questo requisito utilizzando l'analisi statica. Traccia la cronologia della variabile indice attraverso le operazioni aritmetiche che l'hanno prodotta (ad es. moltiplicazioni, addizioni). Se il compilatore non può garantire (in fase di compilazione) che il valore risultante sarà sempre divisibile per le dimensioni del riquadro, genera questo errore.
Il compilatore tratta "allineamento dimostrato" e "allineamento sconosciuto" in modo identico.
Pertanto, se utilizzi un indice che è matematicamente garantito che non sia allineato (ad es.
i * 128 + 32), il compilatore genererà lo stesso errore.
Pertanto, questo errore può verificarsi quando
- Utilizzi una variabile runtime (indice dinamico) per accedere alla memoria.
- La logica di calcolo dell'indice è troppo complessa da analizzare per il compilatore.
- L'indice è matematicamente valido, ma manca una dimostrazione esplicita nel codice.
- L'analisi statica determina un "disallineamento dimostrato".
Debug
Per risolvere questo errore, hai le seguenti opzioni:
1. Assert Alignment Explicitly
Se sai che l'indice è valido, ma il compilatore non riesce a dimostrarlo, utilizza l'operazione
tpu.assume_multiple. che funge da promessa al compilatore che un valore è divisibile per un fattore specifico.
2. Utilizzare Carichi allineati e Ruota
Negli scenari in cui il disallineamento è intenzionale, anziché caricare un segmento vettoriale piccolo e non allineato:
- Carica un riquadro più grande e completamente allineato, quindi ruota i valori di un importo dinamico per spostare i dati desiderati nella posizione corretta (poiché le sezioni vettoriali con indici di inizio dinamici non sono supportate). oppure
- Rimodella o riempi il tensore in modo che i dati inizino dall'indice 0 e lo stride
tra gli accessi corrisponda all'allineamento hardware.
- Esempio: se esegui l'iterazione su blocchi di dimensione 32 a partire dall'offset 1, gli offset sono 1, 33, 65... (non allineato).
- Correzione: ricomponi i dati in un nuovo tensore in cui il primo blocco è a 0 e la dimensione è riempita fino a 128. Gli offset diventano 0, 128, 256..., il che soddisfa il requisito di allineamento.
Questi metodi consumano più memoria, ma spesso semplificano la logica del kernel ed eliminano la necessità di asserzioni di allineamento manuale.