이 문서에서는 HLO 작업의 색인 생성 맵을 기호 방식으로 계산할 수 있는 HLO 색인 생성 분석을 설명합니다. 색인 생성 맵은 한 텐서의 색인을 다른 텐서의 색인에 매핑하는 함수입니다(예: HLO 명령 출력의 색인을 HLO 명령 입력의 색인에 매핑하거나 그 반대의 경우도 마찬가지).
예
tensor<20xf32>~tensor<10x20x30xf32>의 방송
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
출력에서 입력까지의 색인 생성 맵은 \((i, j, k) \mapsto (j)\) for $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$입니다.
동기
XLA GPU는 여러 맞춤형 솔루션을 사용하여 병합, 피연산자 사용률, 타일링 체계를 추론합니다 (자세한 내용은 아래 참고). 색인 생성 분석의 목표는 이러한 사용 사례에 재사용 가능한 구성요소를 제공하는 것입니다. 색인 생성 분석은 MLIR의 아핀 지도 인프라를 기반으로 하며 HLO 의미 체계를 추가합니다.
합병
출력 요소를 계산하기 위해 어떤 입력 요소/슬라이스를 읽었는지 알 수 있는 중요한 사례에서 메모리 합병에 관한 추론이 가능해집니다.
피연산자 사용률
XLA의 피연산자 사용률은 출력이 완전히 사용되었다는 가정하에 명령어의 각 입력이 얼마나 사용되었는지를 나타냅니다. 현재 사용률은 일반적인 사례에 대해서도 계산되지 않습니다. 색인 생성 분석을 통해 사용률을 정확하게 계산할 수 있습니다.
타일식:
타일/슬라이스는 오프셋, 크기, 스트라이드로 매개변수화된 텐서의 직사각형 하위 집합입니다. 타일 전파는 작업 자체의 타일 매개변수를 사용하여 작업의 생산자/소비자의 카드 매개변수를 계산하는 방법입니다. 소프트맥스 및 닷을 위해 이 작업을 실행하는 라이브러리가 이미 있습니다. 타일 전파는 색인 생성 맵을 통해 표현한다면 더 일반적이고 강력할 수 있습니다.
함수 및 도메인
색인 생성 맵은 \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\)텐서의 \(\boldsymbol{d}\) 다중 색인을 \(A\) 텐서의 요소/범위에 \(B\)매핑하는 함수입니다. \(\boldsymbol{s}\) 매개변수는 텐서 \(B\)에는 있지만 텐서에는 없는 차원의 색인 범위를 참조합니다. \(A\).
예를 들어 tensor<2x4x8x16xf32>에서 tensor<4x8xf32>로 감소한 경우 2D 출력에서 4D 입력으로의 색인 생성 맵은\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\)입니다. 여기서 \(d_i\) 는 출력 텐서의 색인에 해당하는 측정기준 매개변수입니다. 매개변수는 \(s_j\)
여러 값을 인코딩합니다. 즉, 출력의 \((d_0, d_1)\) 요소를 계산하려면 입력의 \((s_0, d_0, d_1, s_1)\) 요소가 필요합니다(여기서 \(s_0 \in [0, 2)\) 및\(s_1 \in [0, 16)\)).
이 매핑은 HLO 명령어의 속성으로 구성되거나, 통합되지 않은 명령어의 매핑을 구성하여 퓨전의 색인을 생성할 수 있습니다. 매핑에는 매핑이 존재하는 텐서의 요소를 지정하는 도메인도 있습니다.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
재계산을 최소화하려면 심볼릭 계산을 위한 라이브러리가 필요합니다. XLA는 이미 MLIR에 종속되므로 기호화된 산술 라이브러리를 작성하는 대신 mlir::AffineMap을 사용합니다.
일반적인 AffineMap 모습
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
AffineMap에는 편리하게 \(\boldsymbol d\) 및 \(\boldsymbol s\) 에 사용할 수 있는 크기와 기호라는 두 가지 유형의 매개변수가 있습니다.
AffineMap에는 측정기준의 범위에 대한 메타데이터가 없으므로 이 데이터는 Google에서 직접 제공해야 합니다.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges는 색인 생성 맵의 측정기준 매개변수에 대한 포괄적 상자 제약 조건을 인코딩합니다. 이는 일반적으로 전치, 감소, 요소별, 점과 같은 작업의 출력 텐서 형태와 일치하지만 HloConcatenateInstruction와 같은 몇 가지 예외가 있습니다. \(\boldsymbol{d}\)
symbol_ranges는 매개변수가 취할 수 있는 \(\boldsymbol {s}\) 가능한 값을 인코딩합니다.
위의 모든 내용이 실제로 무엇을 의미하는지 알아보기 위해 예를 들어 설명해 보겠습니다.
통합되지 않은 작업을 위한 지도 색인 생성
요소별
요소별 작업의 경우 색인 생성 맵은 ID입니다.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
입력 맵의 출력은 다음과 같습니다.
- output -> input_0: \((d_0, d_1) \mapsto (d_0, d_1)\) for $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(출력)$
- output -> input_1: \((d_0, d_1) \mapsto (d_0, d_1)\) for $\boldsymbol{d} \in {\rm Dom} (출력)$
입력-출력 맵은
- input_i -> output: \((d_0, d_1) \mapsto (d_0, d_1)\) for $\boldsymbol{d} \in {\rm Dom}(출력)$
방송
브로드캐스트는 출력을 입력에 매핑할 때 일부 측정기준이 삭제되고 입력을 출력에 매핑할 때 추가된다는 의미입니다.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: \((d_0, d_1, d_2) \mapsto (d_1)\) for $\boldsymbol{d} \in {\rm Dom}(출력)$
입력-출력 맵
- input -> output: \((d_0) \mapsto (s_0, d_1, s_1)\) for $\boldsymbol{d} \in {\rm Dom}(출력)\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$.
이제 입력-출력 매핑을 위한 \(\boldsymbol s\) 가 오른쪽에 있습니다. 이러한 기호는 값 범위를 나타내는 기호입니다. 예를 들어 이 특정 사례에서는 색인이 \(d_0\) 인 입력의 모든 요소가 출력의 10x1x30 슬라이스에 매핑됩니다.
Constant 및 Iota
편리하게도 입력 매개변수가 없으므로 색인 생성을 계산할 필요가 없습니다.
전치행렬
전치를 위한 지도 색인 생성은 입력/출력 크기의 순열입니다.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) (\(\boldsymbol{d} \in {\rm Dom}(output)\)의 경우)
입력과 출력 매핑은 다음과 같습니다.
- input -> output: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) (\(\boldsymbol{d} \in {\rm Dom}(input)\)의 경우)
리버스
역방향 사용을 위한 맵의 색인을 생성하면 되돌린 크기가 $upper_bound(d_i) - d_i$로 변경됩니다.
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(출력)$
입력과 출력 매핑은 다음과 같습니다.
- input -> output: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(입력)$
(Variadic) 줄이기
가변적인 축소는 여러 개의 입력과 여러 개의 초기화를 가지며, 출력에서 입력으로의 매핑은 축소된 차원을 추가합니다. 따라서 어떤 면에서는 방송의 역으로 작동합니다.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
입력 맵의 출력은 다음과 같습니다.
- output -> input_j: \((d_0) \mapsto (s_0, d_0)\) for $\boldsymbol{d} \in {\rm Dom}(출력)\( and \)\boldsymbol{s} \in [0, 9]$
- output -> init_j: \((d_0) \mapsto ()\) for $\boldsymbol{d} \in {\rm Dom}(output)$
입력과 출력 맵은 다음과 같습니다.
- input_i -> output_j: \((d_0, d_1) \mapsto (d_1)\) for $\boldsymbol{d} \in {\rm Dom}(입력)$
- init_i -> output_j: \(() \mapsto (s_0)\) for \(\boldsymbol{s} \in [0, 9]\)
\(i, j = 0, \ldots, INPUT\\_COUNT\)
슬라이스
슬라이스의 출력에서 입력으로 색인을 생성하면 출력의 모든 요소에 유효한 스트라이드 색인 생성 맵이 생성됩니다. 입력에서 출력으로의 매핑은 입력 요소의 스트라이딩 범위로 제한됩니다.
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) (\(\boldsymbol{d} \in {\rm Dom}(output)\)의 경우)
입력과 출력 매핑은 다음과 같습니다.
- input -> output: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) for \(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) with strides $[1, 7, 2]$.
미정: 입력-출력 색인 생성
재구성
형태 변경에는 여러 가지 맛이 있습니다.
도형 접기
이는 N-D에서 1D로 셰이프를 '선형화'하는 것입니다.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) for $\boldsymbol{d} \in {\rm Dom}(output)$
입력과 출력 매핑은 다음과 같습니다.
- input -> output: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) for $\boldsymbol{d} \in {\rm Dom}(input)$.
도형 펼치기
이는 역 '셰이프 축소' 작업으로, 1차원 입력을 N-D 출력으로 재구성합니다.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) for $\boldsymbol{d} \in {\rm Dom}(출력)$
입력과 출력 매핑은 다음과 같습니다.
- input -> output: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) for $\boldsymbol{d} \in {\rm Dom}(input)$.
일반 재구성
단일 펼치기 또는 접기로 표현할 수 없는 형태 변경 작업입니다. 확장 또는 축소 도형이 2개 이상 있는 컴포지션으로만 표현될 수 있습니다.
예시 1: 선형화-비선형화.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
이러한 형태 변경은 tensor<4x8xf32>의 축소 모양에서 tensor<32xf32>로, 그런 다음 tensor<2x4x4xf32>로 확장되는 도형의 합성으로 표현될 수 있습니다.
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
\(\boldsymbol{d} \in {\rm Dom}(output)\)동안
입력과 출력 매핑은 다음과 같습니다.
- 입력 -> 출력: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
\(\boldsymbol{d} \in {\rm Dom}(input)\)
예 2: 펼친 하위 도형 및 접힌 하위 도형
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
이 재구성은 두 가지 재구성의 컴포지션으로 표현할 수 있습니다. 첫 번째는 가장 바깥쪽 치수 tensor<4x8x12xf32>를 tensor<32x12xf32>로 축소하고 두 번째는 가장 안쪽 차원 tensor<32x12xf32>를 tensor<32x3x4xf32>로 확장합니다.
입력 맵의 출력은 다음과 같습니다.
- 출력 -> input: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) for \(\boldsymbol{d} \in {\rm Dom}(output)\)
입력과 출력 매핑은 다음과 같습니다.
- input -> output: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\) ( \(\boldsymbol{d} \in {\rm Dom}(input)\)의 경우)
비트캐스트
비트캐스트 작업은 전치-재형-전치 시퀀스로 표현할 수 있습니다. 따라서 색인 생성 맵은 이 시퀀스의 색인 생성 맵을 조합한 것에 불과합니다.
Concatenate
concat의 출력-입력 매핑은 모든 입력에 대해 정의되지만 도메인이 겹치지 않습니다. 즉, 한 번에 하나의 입력만 사용됩니다.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
입력 맵의 출력은 다음과 같습니다.
- 출력 -> 입력 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) 대상: \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- 출력 -> 입력 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) $\boldsymbol{d} \in [0, 2] \times [50, 79]$
출력 매핑은 다음과 같습니다.
- input 1 -> output: \((d_0, d_1) \mapsto (d_0, d_1)\) for $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- input 2 -> output: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) for $\boldsymbol{d} \in {\rm Dom}(input_2)$.
점 (출력-입력 구현됨
점의 색인 생성 맵은 감소의 맵과 매우 유사합니다.
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
입력 매핑의 출력은 다음과 같습니다.
- 출력 -> input_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) for \(\boldsymbol{d} \in {\rm Dom}(output)\) and \(\boldsymbol{s} \in [0, 255]\)
- output -> input_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) for \(\boldsymbol{d} \in {\rm Dom}(output)\) 및 \(\boldsymbol{s} \in [0, 255]\)
출력 맵의 입력은 다음과 같습니다.
- input_1 -> output: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) (\(\boldsymbol{d} \in {\rm Dom}(input_1)\) 및 \(\boldsymbol{s} \in [0, 63]\))
- input_2 -> output: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) (\(\boldsymbol{d} \in {\rm Dom}(input_2)\) 및 \(\boldsymbol{s} \in [0, 127]\))
기간 축소 (미정)
패드 (미정)
퓨전용 지도 색인 생성
퓨전 작업의 색인 생성 맵은 클러스터의 모든 작업에 대한 색인 생성 맵을 조합한 것입니다. 일부 입력은 서로 다른 액세스 패턴으로 여러 번 읽힐 수 있습니다.
하나의 입력, 여러 개의 인덱싱 맵
다음은 \(p_0 + p_0^T\)의 예입니다.
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
p0의 출력-입력 색인 생성 맵은 $(d_0, d_1) \mapsto (d_0,
d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$입니다. 즉, 출력의 한 요소를 계산하려면 입력 매개변수를 두 번 읽어야 할 수도 있습니다.
입력 1개, 중복 삭제된 색인 생성 맵
즉시 명확하지는 않지만 색인 생성 맵이 실제로 동일한 경우가 있습니다.
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
이 경우 p0의 출력-입력 색인 생성 맵은 $(d_0, d_1, d_2) \mapsto (d_2, d_0, d_1)$입니다.
소프트맥스
소프트맥스의 parameter 0의 출력-입력 색인 생성 매핑은 다음과 같습니다.
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
for \(\boldsymbol{d} \in {\rm Dom}(output)\) 및 \(\boldsymbol{s} \in [0, 124]\)는 입력의 가장 안쪽에 있는 차원을 나타냅니다.
지도 단순화 기능 색인 생성
mlir::AffineMap 업스트림의 기본 단순화자는 치수/기호 범위를 가정할 수 없습니다. 따라서 mod 및 div를 사용한 표현식을 효율적으로 단순화할 수 없습니다.
아핀 맵에 있는 하위 표현식의 하한 및 상한에 관한 정보를 활용하여 훨씬 더 단순화할 수 있습니다.
단순화자는 다음 표현식을 다시 작성할 수 있습니다.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \mod 10) [\( becomes \)_0, \mod 10) [\( becomes \)_0, \\( for \)
- \\( for \)\( becomes \)
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) for $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.
지도 단순화의 색인을 생성하면 HLO에서 연결된 재구성의 일부가 서로를 상쇄한다는 것을 파악할 수 있습니다.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
색인 생성 맵의 구성과 단순화를 완료하면 다음과 같은 결과를 얻을 수 있습니다.
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).