Código do erro: E2002

Categoria:tempo de compilação: bloco e mosaico de entrada/saída desalinhados

Esse erro ocorre quando o formato de bloco de uma entrada ou saída de kernel não se alinha com o bloco padrão do tipo de dados no hardware de TPU específico que está sendo usado.

Exemplos de mensagens de erro:

UNIMPLEMENTED: Mosaic failed to compile TPU kernel: Failed to set window params
for input 0: Operand of shape (..., 256, 8192) has tiling (16, 128), but its
block shape (..., 8, 8192) is not divisible by tiling evenly nor matches the
full shape.

Back-ends do XLA:TPU

Visão geral

Os TensorCores (TCs) nas TPUs têm registros vetoriais bidimensionais. As duas dimensões são chamadas de subpista e pista. Como as unidades de computação de TPU (por exemplo, MXU) operam na granularidade dos registros de vetor, e as matrizes XLA são dispostas na memória da TPU em blocos.

Esse mosaico (por exemplo, 8x128) minimiza as transformações de dados ao alimentar as unidades de computação. As dimensões exatas de mosaico (subvias × vias) dependem da geração de hardware e do tipo de dados. Por exemplo, um bloco comum para a maioria dos tipos é 8×128.

No momento da compilação, a XLA aplica as seguintes restrições às dimensões secundárias e secundárias de cada entrada/saída do kernel:

  1. Divisibilidade:a dimensão do bloco precisa ser um múltiplo da dimensão do bloco no tensor subjacente ou
  2. Exceção de forma completa:se a dimensão do bloco não for divisível, ela precisará ser igual ao tamanho completo dessa dimensão no tensor subjacente.

Esse erro é acionado quando um bloco viola as duas condições. Por exemplo, o carregamento de um bloco de formato (8, 100) de uma entrada de formato (8, 1024) em hardware com bloco de formato (8, 128) falha porque 100 não é divisível por 128 e 100 != 1024. No entanto, seria permitido se o formato de entrada fosse (32, 100).

Depuração

Para resolver esse erro, verifique se os formatos de bloco do kernel estão alinhados com o tiling de hardware atual. Modifique o código do kernel para alinhar o tamanho do bloco de forma que ele seja um múltiplo do tiling necessário.

  • Exemplo:se o erro indicar que o revestimento é (16, 128), mas o formato do bloco for (8, 128), mude a especificação para que o formato corresponda a (16, 128).