오류 코드: E2002

카테고리: 컴파일 시간: 모자이크 입력/출력 정렬되지 않은 블록 및 타일링

이 오류는 커널 입력 또는 출력의 블록 모양이 사용 중인 특정 TPU 하드웨어의 데이터 유형 기본 타일링과 정렬되지 않을 때 발생합니다.

샘플 오류 메시지:

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.

XLA 백엔드: TPU

개요

TPU의 텐서 코어 (TC)에는 2차원 벡터 레지스터가 있습니다. 두 측정기준을 하위 차선차선이라고 합니다. TPU 컴퓨팅 단위 (예: MXU)는 벡터 레지스터의 세부사항으로 작동하며 XLA 배열은 TPU 메모리에 타일로 배치됩니다.

이 타일링 (예: 8x128) 컴퓨팅 단위에 공급할 때 데이터 변환을 최소화합니다. 정확한 타일링 크기 (서브레인 × 레인)는 하드웨어 세대와 데이터 유형에 따라 달라집니다. 예를 들어 대부분의 유형에 대한 일반적인 타일링은 8×128입니다.

컴파일 시간에 XLA는 각 커널 입력/출력의 부 버전 및 두 번째 부 버전 측정기준에 다음 제약 조건을 적용합니다.

  1. 나누기 가능성: 블록 차원은 기본 텐서의 타일 차원의 배수여야 합니다.
  2. 전체 모양 예외: 블록 차원이 나누어지지 않으면 기본 텐서에서 해당 차원의 전체 크기와 같아야 합니다.

이 오류는 블록이 두 조건을 모두 위반하는 경우에 트리거됩니다. 예를 들어 (8, 128) 타일링이 있는 하드웨어에서 (8, 1024) 모양의 입력으로부터 (8, 100) 모양의 블록을 로드하는 것은 100128100 != 1024로 나눌 수 없기 때문에 실패합니다. 하지만 입력 모양이 (32, 100)이면 허용됩니다.

디버깅

이 오류를 해결하려면 커널의 블록 모양이 현재 하드웨어 타일링과 일치해야 합니다. 필요한 타일링의 배수가 되도록 블록 크기를 정렬하도록 커널 코드를 수정합니다.

  • 예: 오류에 타일링이 (16, 128)라고 표시되어 있지만 블록 모양이 (8, 128)인 경우 모양이 (16, 128)와 일치하도록 블록 사양을 변경합니다.