타일 레이아웃

그림 1

그림 1은 배열 F32[3,5] 가 2x2 타일로 메모리에 배치되는 방식을 보여줍니다. 이 레이아웃이 포함된 도형은 F32[3,5]{1,0:T(2,2)}로 작성되며, 여기서 1,0은 차원의 물리적 순서 (레이아웃의 minor_to_major 필드)와 관련이 있는 반면 콜론 뒤의 (2,2)는 2x2 타일에 의한 물리적 크기의 타일식을 나타냅니다.

직관적으로는 타일이 도형을 덮도록 배치된 다음 각 카드 내에서 요소가 타일 없이 배치됩니다. 위의 예에서와 같이 예의 오른쪽 부분에서 메모리의 레이아웃을 보여줍니다. 여기에는 원래 배열 경계가 짝수가 아니더라도 2x2 타일을 완성하기 위해 추가된 흰색 패딩 요소가 포함됩니다.

패딩의 추가 요소는 특정 값을 포함할 필요가 없습니다.

도형과 타일이 지정된 타일 생성에 사용되는 선형 색인 수식

타일링을 하지 않으면 배열 경계가 d=(d_n, d_n-1, ... , d₁)인 배열의 요소 e=(e_n, e_n-1, ... , e₁) (d1이 가장 작은 차원임)이 위치에 큰 순서로 배치됩니다.

linear_index(e, d)
= linear_index((e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1...d₁ + e_n-2...d₁ + e_n-2...d_n-2...

이 문서에서는 표기를 쉽게 하기 위해 타일의 차원 수가 배열과 동일하다고 가정합니다. XLA의 타일 구현에서 이는 초기 메이저 치수를 변경하지 않고 가장 작은 치수에만 타일을 적용하여 치수가 더 적은 타일로 일반화됩니다. 그러면 지정된 타일이 타일로 연결되는 도형의 실제 치수 접미사를 언급합니다.

크기 타일 (t_n, t_n-1, ... , t₁)을 사용하면 색인이 있는 배열의 요소 (e_n, e_n-1, ... , e₁)가 최종 레이아웃에서 이 위치에 매핑됩니다.

t


_n

레이아웃은 (⌊e_n/t_n⌋, ..., ⌊e₁/t₁⌋) 두 부분으로 구성되는 것으로 볼 수 있습니다. 이는 크기(⌈d_n/t_n⌉, ... , mod_n⌉, ... , mod_n⌉, ... , ⌈d ceil 함수는 ⌈d_i/t_i⌉에 표시됩니다. 타일이 더 큰 배열의 경계를 초과하면 그림 1과 같이 패딩이 삽입되기 때문입니다. 카드 내의 카드와 요소는 모두 타일 없이 재귀적으로 배치됩니다.

그림 1의 예에서 요소 (2,3)에는 (1,1,0,1)의 결합된 좌표 벡터에 대한 타일 색인 (1,1)과 타일 내 색인 (0,1)이 있습니다. 카드 색인에는 경계 (2,3)가 있고 카드 자체는 (2,3,2,2)의 결합된 벡터의 경우 (2,2)입니다. 그러면 논리적 도형에서 색인 (2,3)이 있는 요소의 타일이 포함된 선형 색인은 다음과 같습니다.

linear_index_with_tile((2,3), (3,5), (2,2))
= linear_index((1,1,0,1), (2,3,2,2))
= linear_index((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index((0,1) ∙ 2 ∙ 2 + linear_index((0,1) ∙ 2 ∙ (2, 3), (3,5), (2,2))

패드-형태-트랜스포즈로 타일 지정

타일 기반 레이아웃은 다음과 같이 작동합니다.
크기 배열 (d_n, d_n-1, ... , d1)을 고려합니다(d1이 가장 작은 차원임). 타일은 크기(t_n, t_n-1, ... , t₁) (t₁이 가장 작은 차원임)의 타일로 배치될 때 다음과 같이 패드의 형태 변경 및 전치와 관련하여 설명할 수 있습니다.

1. 배열은 (⌈d_n/t_n⌉∙t_n, ... , ⌈d₁/t₁⌉∙t₁)에 패딩됩니다.
2. 각 차원 i는 (⌈d_i/ti⌉, t_i)로 나뉩니다. 즉, 배열의 형태가
   (⌈d_n/t_n⌉, t_n, ... , ⌈d₁/t₁⌉, t₁)로 변경됩니다.
   이 모양 변경 자체는 물리적 레이아웃 자체가 변경되지 않으므로 이 형태 변경은 비트캐스트입니다. 명시적으로 타일을 생각하지 않는 경우 이 재구성에서는 패딩 처리된 도형과 동일한 수의 요소를 사용하여 모든 도형을 표현할 수 있습니다. 여기에서 이러한 방식으로 카드를 표현하는 방법을 확인할 수 있습니다.
3. 전치는 t_n, ... , t₁을 상대 순서를 유지하면서 최단 차원으로 이동함으로써 발생합니다. 따라서 메이저 차원에서 가장 작은 차원 순으로 차원의 순서는
   이 됩니다(⌈d_n/t_n⌉, ... ,⌈d₁/t₁⌉, t_n).

최종 도형에는 각 차원의 타일 수를 나타내는 접두사
(⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉)가 있습니다. 배열의 요소 (e_n, ... , e₁)는 최종 형태의 이 요소에 매핑됩니다.
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀, e_n mod t_n, ... , mod₁). 요소의 선형 색인은 위 수식을 예상대로 따르기 쉽습니다.

반복되는 타일

XLA의 타일은 반복적으로 적용하면 더욱 유연해집니다.

그림 2

그림 2에서는 크기가 4x8인 배열이 두 수준의 타일 (처음 2x4, 2x1)에 의해 어떻게 타일화되는지를 보여줍니다. 이 반복되는 타일식을 (2,4)(2,1)로 표현합니다. 각 색상은 2x4 타일을 나타내고 각 빨간색 테두리 상자는 2x1 타일입니다. 숫자는 해당 요소의 메모리에 있는 선형 색인을 타일 형식으로 나타냅니다. 이 형식은 초기 타일이 더 크다는 점을 제외하면 (8,128)(2,1) TPU의 BF16에 사용된 형식과 일치합니다. (8,128)(2,1) 2x1로 두 번째 타일링하는 목적은 16비트 값 2개를 모아 TPU 아키텍처에 맞는 방식으로 16비트 값 1개를 형성하는 것입니다.

두 번째 이상의 타일은 이 예의 (8,128)(2,1)에서처럼 타일 내의 데이터를 재정렬하는 보조 타일 내 크기를 모두 참조할 수 있지만, 이전 타일의 주요 교차 타일 크기를 참조할 수도 있습니다.

타일을 사용하여 크기 결합

XLA의 타일은 치수 결합도 지원합니다. 예를 들어 F32[2,7,8,11,10]{4,3,2,1,0}의 차원을 (2,3)과 타일로 만들기 전에 먼저 F32[112,110]{1,0}으로 결합할 수 있습니다. 사용된 타일은 (**, *, 2, *, 3)입니다. 여기에서 카드의 별표는 해당 차원을 가져와서 다음 보조 차원과 결합함을 의미합니다. 인접한 여러 차원을 한 차원으로 함께 포함할 수 있습니다. 합산된 차원은 카드의 해당 치수에서 타일 값 -1로 표시되며, 이는 카드에서 차원 크기로 유효하지 않습니다.

더 정확히 설명하자면, 도형의 차원 i가 카드에서 별표를 통해 제거된 경우 이전에 타일링의 정의가 적용되기 전에 타일이 지정되는 도형과 타일 벡터에서 차원이 모두 삭제되고, 도형의 차원 i-1은 d_i-1에서 d_id_i-1로 배열 경계가 증가합니다. 이 단계는 타일 벡터의 각 별표에 대해 반복됩니다.