카테고리: 컴파일 시간: 모자이크 입증되지 않은 메모리 액세스 정렬
이 오류는 컴파일러가 메모리 액세스 작업 (예: vector.load, vector.store, tpu.load 또는 tpu.store)을 분석할 때 특정 차원에 사용된 동적 색인이 필요한 타일링 크기의 배수임을 정적으로 증명할 수 없는 경우에 발생합니다.
샘플 오류 메시지:
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>
XLA 백엔드: TPU
개요
커널이 벡터를 로드하거나 저장할 때 메모리 주소(기본 포인터와 동적 색인에서 계산됨)가 하드웨어의 벡터 타일링 크기와 정렬되어야 합니다. 예를 들어 차원이 128개 요소로 타일링된 경우 액세스하는 데 사용되는 동적 색인은 0, 128, 256 등이어야 합니다. 벡터 로드 및 저장과 같은 많은 작업에는 정적 색인에 이러한 요구사항이 없습니다.
컴파일러는 정적 분석을 사용하여 이 요구사항을 적용합니다. 색인 변수의 기록을 생성한 산술 연산 (예: 곱셈, 덧셈)을 통해 추적합니다. 컴파일러가 결과 값이 항상 타일링 크기로 나누어질 수 있다고 컴파일 시간에 보장할 수 없는 경우 이 오류가 발생합니다.
컴파일러는 '증명된 정렬 불일치'와 '알 수 없는 정렬'을 동일하게 취급합니다.
따라서 수학적으로 정렬되지 않는 것이 보장된 색인 (예:
i * 128 + 32)와 동일한 오류가 발생합니다.
따라서 이 오류는 다음과 같은 경우에 발생할 수 있습니다.
- 런타임 변수 (동적 색인)를 사용하여 메모리에 액세스합니다.
- 컴파일러가 분석하기에는 색인 계산 로직이 너무 복잡합니다.
- 색인은 수학적으로 유효하지만 코드에 명시적인 증명이 없습니다.
- 정적 분석을 통해 '입증된 불일치'가 확인됩니다.
디버깅
이 오류를 해결하려면 다음 옵션을 사용하세요.
1. 정렬 명시적으로 어설션
색인이 유효하지만 컴파일러가 이를 증명할 수 없는 경우 tpu.assume_multiple 작업을 사용하세요. 이는 값이 특정 요소로 나눌 수 있다는 컴파일러에 대한 약속 역할을 합니다.
2. 정렬된 로드 및 회전 사용
의도적으로 정렬되지 않은 시나리오에서는 정렬되지 않은 작은 벡터 세그먼트를 로드하는 대신 다음을 실행합니다.
- 완전히 정렬된 더 큰 타일을 로드한 다음 동적 양만큼 값을 회전하여 원하는 데이터를 위치로 이동합니다 (동적 시작 색인이 있는 벡터 슬라이스는 지원되지 않음). 또는
- 액세스 간의 스트라이드가 하드웨어 정렬과 일치하도록 텐서를 재구성하거나 패딩합니다.
- 예: 오프셋 1에서 시작하여 크기가 32인 청크를 반복하는 경우 오프셋은 1, 33, 65...입니다. (정렬되지 않음)
- 수정: 첫 번째 청크가 0에 있고 차원이 128로 패딩된 새 텐서로 데이터를 다시 패킹합니다. 오프셋은 0, 128, 256이 되므로 정렬 요구사항을 충족합니다.
이러한 메서드는 메모리를 더 많이 소비하지만 커널 로직을 간소화하고 수동 정렬 어설션의 필요성을 없애는 경우가 많습니다.