도형 및 레이아웃

XLA 작업의 구조

다음 HLO를 살펴보세요.

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

다음 구성요소로 구성됩니다.

  • 작업 이름: add.936
    • 작업의 고유한 이름입니다.
  • 모양: bf16[8,1,1280,16384]
    • 이는 작업의 출력 모양입니다. 여기에서 dtype은 bf16이고 모양은 [8,1,1280,16384]입니다.
  • 레이아웃 (타일링 사용): 3,2,0,1:T(8,128)(2,1)
    • 이는 배열이 메모리에 저장되는 방식을 설명합니다. 3,2,0,1는 메모리의 축 순서 (예: 열 우선, 행 우선 등)를 나타내고 T(8,128)(2,1)는 사용된 타일링 및 패딩을 나타냅니다.
    • 레이아웃은 선택사항입니다. 지정하지 않으면 타일링이 없으며 차원은 가장 주요한 것부터 가장 사소한 것 순으로 정렬된 것으로 간주됩니다.
  • 작업: add
    • 실행 중인 작업입니다. 여기서는 Add이며, 이는 작업 이름에도 언급되어 있습니다.
  • 인수: exponential.183, broadcast.3115
    • 이 작업은 고유한 이름으로 지정된 두 개의 인수를 사용합니다.

융합 작업의 또 다른 예를 살펴보겠습니다.

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

앞에서 설명한 구성요소 외에도 다음이 포함됩니다.

  • 속성: kind 및 calls
    • 이러한 값은 수행되는 작업(이 경우 융합)에 관한 추가 정보를 제공합니다.
  • 메모리 위치 (메모리 공간 식별자): S(1)
    • 배열이 저장된 메모리 공간/위치를 나타냅니다. S(1)는 이 배열이 VMEM (TPU에 있음)에 있음을 나타냅니다.
  • %fusion.32 입력 인수의 모양 및 레이아웃 세부정보

다음 섹션에서는 모양, 레이아웃, 메모리 공간 식별자를 설명합니다. 타일링에 관한 자세한 내용은 타일 레이아웃을 참고하세요.

도형

XLA ShapeProto 프로토(xla_data.proto)는 N차원 배열 (간단히 배열)의 차원 수, 크기, 데이터 유형을 설명합니다.

용어, 표기법, 규칙

  • 배열의 실제 차원 수는 크기가 1보다 큰 차원 수입니다.

  • N 차원 배열의 경우 측정기준은 0부터 N-1까지 번호가 매겨집니다. 측정기준의 크기는 음이 아닌 정수입니다. 특히 크기 0이 유효합니다. 측정기준 번호는 편의를 위해 임의로 지정된 라벨입니다. 이러한 차원 번호의 순서는 모양 레이아웃의 특정 부전공/전공 순서를 의미하지 않습니다. 레이아웃은 LayoutProto 프로토에 의해 결정됩니다.

  • 관례에 따라 측정기준은 측정기준 번호의 오름차순으로 나열됩니다. 예를 들어 크기가 [A x B x C]인 3차원 배열의 경우 차원 0의 크기는 A이고, 차원 1의 크기는 B이며, 차원 2의 크기는 C입니다.

    XLA의 일부 유틸리티는 Python과 유사한 음수 색인도 지원합니다. 측정기준 -1은 마지막 측정기준입니다 (N 차원 배열의 경우 N-1와 동일). 예를 들어 위에 설명된 3차원 배열의 경우 -1차원의 크기는 C이고 -2차원의 크기는 B입니다.

  • 2차원, 3차원, 4차원 배열에는 측정기준과 관련된 특정 문자가 있는 경우가 많습니다. 예를 들어 2D 배열의 경우 다음과 같습니다.

    • 측정기준 0: y
    • 측정기준 1: x

    3D 배열의 경우:

    • 측정기준 0: z
    • 측정기준 1: y
    • 측정기준 2: x

    4D 배열의 경우:

    • 측정기준 0: p
    • 측정기준 1: z
    • 측정기준 2: y
    • 측정기준 3: x

  • 차원을 사용하는 XLA API의 함수는 차원 번호의 오름차순으로 차원을 사용합니다. 이는 측정기준을 initializer_list로 전달할 때 사용되는 순서와 일치합니다.예를 들면 다음과 같습니다.

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    차원 크기 배열이 [A, B, C, D] 시퀀스로 구성된 도형을 만듭니다.

레이아웃

LayoutProto 프로토콜은 메모리에 배열이 표현되는 방식을 설명합니다. 여기에는 다음 필드가 포함됩니다.

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

하위 측정기준에서 상위 측정기준으로의 순서 지정

유일한 필수 필드는 minor_to_major입니다. 이 필드는 셰이프 내에서 측정기준의 사소한 순서부터 중요한 순서까지를 설명합니다. minor_to_major의 값은 배열의 차원 순서입니다 (N 차원 배열의 경우 0~N-1). 첫 번째 값은 가장 작은 차원이고 마지막 값은 가장 큰 차원입니다. 가장 작은 차원은 선형 메모리에 배치된 배열의 요소를 단계별로 실행할 때 가장 빠르게 변경되는 차원입니다.

예를 들어 크기가 [2 x 3]인 다음 2D 배열을 고려해 보세요.

a b c
d e f

여기서 0 차원은 크기가 2이고 1 차원은 크기가 3입니다. 레이아웃의 minor_to_major 필드가 [0, 1]이면 측정기준 0가 가장 작은 측정기준이고 측정기준 1가 가장 큰 측정기준입니다. 이는 선형 메모리의 다음 레이아웃에 해당합니다.

a d b e c f

0부터 N-1까지의 하위-상위 측정기준 순서는 열 우선(2차원의 경우)과 유사합니다. 측정기준의 단조 순서를 가정하면 코드에서 이 레이아웃을 참조하는 또 다른 방법은 단순히 'dim 0이 마이너'입니다.

반면 레이아웃의 minor_to_major 필드가 [1, 0]이면 선형 메모리의 레이아웃은 다음과 같습니다.

a b c d e f

N 차원 배열의 N-1에서 0까지의 부-주 측정기준 순서는 행 우선 (2차원의 경우)과 유사합니다. 차원의 단조 순서를 가정하면 코드에서 이 레이아웃을 참조하는 또 다른 방법은 단순히 'dim 0이 주요'입니다.

기본 마이너-메이저 순서

새로 생성된 모양의 기본 레이아웃은 '차원 순서가 메이저에서 마이너로' (즉, [N-1, ..., 0])입니다.

패딩

tail_padding_alignment_in_elements 필드는 요소 수 측면에서 타일식 배열의 정렬을 정의합니다. 타일링을 적용한 후 패딩된 요소가 전체 요소 수가 이 값의 배수가 될 때까지 레이아웃 끝에 추가됩니다.

배열 색인 생성

index_util.h의 IndexUtil 클래스는 모양과 레이아웃이 지정된 경우 다차원 색인과 선형 색인 간 변환을 위한 유틸리티를 제공합니다. 다차원 색인에는 각 차원의 int64 색인이 포함됩니다. 선형 색인은 배열을 보유하는 버퍼로 색인을 생성하는 단일 int64 값입니다. 모양과 레이아웃의 생성과 조작을 간소화하는 유틸리티는 동일한 디렉터리의 shape_util.h 및 layout_util.h을 참고하세요.

메모리 공간 식별자

HLO에서 각 배열은 S(n)으로 작성된 메모리 공간 식별자로 주석 처리될 수 있습니다.

  • S(0) (생략되는 경우가 많음)는 기기 고대역폭 메모리 (HBM)를 나타냅니다.
  • S(1)는 기기 내 가상 메모리 (VMEM)를 나타냅니다.
  • S(2), S(3) 등은 추가 기기별 메모리 공간에 해당합니다.
  • S(5)은 호스트 메모리를 나타냅니다.