Categoria:tempo de compilação: alinhamento de acesso à memória não comprovado do Mosaic
Esse erro ocorre quando o compilador analisa uma operação de acesso à memória (como
vector.load, vector.store, tpu.load ou tpu.store) e não consegue
provar estaticamente que o índice dinâmico usado para uma dimensão específica é um
múltiplo do tamanho de mosaico necessário.
Exemplos de mensagens de erro:
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>
Back-ends do XLA:TPU
Visão geral
Quando o kernel carrega ou armazena um vetor, o endereço de memória (calculado com base no ponteiro base mais o índice dinâmico) precisa estar alinhado com o tamanho de mosaico do vetor no hardware. Por exemplo, se uma dimensão for segmentada por 128 elementos, o índice dinâmico usado para acessá-la precisará ser 0, 128, 256 etc. Muitas operações (como cargas e armazenamentos de vetores) não têm esses requisitos para índices estáticos.
O compilador impõe esse requisito usando a análise estática. Ele rastreia o histórico da variável de índice nas operações aritméticas que a produziram (por exemplo, multiplicações, adições). Se o compilador não puder garantir (no momento da compilação) que o valor resultante sempre será divisível pelo tamanho do bloco, ele vai gerar esse erro.
O compilador trata "desalinhamento comprovado" e "alinhamento desconhecido" de forma idêntica.
Portanto, se você usar um índice que tem garantia matemática de desalinhamento (por exemplo,
i * 128 + 32), o compilador vai gerar o mesmo erro.
Esse erro pode ocorrer quando
- Você usa uma variável de tempo de execução (índice dinâmico) para acessar a memória.
- A lógica de cálculo do índice é muito complexa para o compilador analisar.
- O índice é matematicamente válido, mas não tem uma prova explícita no código.
- A análise estática determina o "desalinhamento comprovado".
Depuração
Para resolver esse erro, você tem as seguintes opções:
1. Declarar o alinhamento explicitamente
Se você sabe que o índice é válido, mas o compilador não consegue provar isso, use a operação tpu.assume_multiple. Isso funciona como uma promessa ao compilador de que um valor é divisível por um fator específico.
2. Usar "Cargas alinhadas" e "Girar"
Em cenários em que o desalinhamento é intencional, em vez de carregar um segmento de vetor pequeno e desalinhado:
- Carregue um bloco maior e totalmente alinhado e gire os valores em uma quantidade dinâmica para mover os dados desejados para a posição (já que não há suporte para intervalos de vetores com índices de início dinâmicos). ou
- Remodele ou faça padding no tensor para que os dados comecem no índice 0 e a passada entre os acessos corresponda ao alinhamento de hardware.
- Exemplo: se você iterar sobre blocos de tamanho 32 começando no deslocamento 1, seus deslocamentos serão 1, 33, 65... (não alinhado).
- Correção: recompacte os dados em um novo tensor em que o primeiro bloco está em 0 e a dimensão é preenchida com 128. Seus deslocamentos se tornam 0, 128, 256..., o que atende ao requisito de alinhamento.
Esses métodos consomem mais memória, mas geralmente simplificam a lógica do kernel e eliminam a necessidade de declarações de alinhamento manual.